#!/usr/bin/env python
import pandas as pd
from sklearn.model_selection import train_test_split

def preprocess():
    df_all_data = pd.read_csv('/opt/ml/input/data/raw/001.csv')
    X = df_all_data[['floors']]
    y = df_all_data[['price']]
    
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=0)
    
    X_train.to_csv('/opt/ml/input/data/train/X_train.csv', index=False)
    y_train.to_csv('/opt/ml/input/data/train/y_train.csv', index=False)
    X_test.to_csv('/opt/ml/input/data/test/X_test.csv', index=False)
    y_test.to_csv('/opt/ml/input/data/test/y_test.csv', index=False)
    
if __name__ == "__main__":
    preprocess()