import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, MinMaxScaler

cols = list(pd.read_csv("avocado.csv", nrows=1))
# print("###\n", cols, "\n###")
avocados = pd.read_csv("avocado.csv", usecols=cols[1:])
avocados.describe(include="all")

float_cols = ['AveragePrice','Total Volume','4046','4225','4770','Total Bags','Small Bags','Large Bags','XLarge Bags']

avocados.loc[:, float_cols] = StandardScaler().fit_transform(avocados.loc[:, float_cols])
print(avocados.head())

# avocados.loc[:, float_cols] = MinMaxScaler().fit_transform(avocados.loc[:, float_cols])
# print(avocados.head())

avocado_train, avocado_test = train_test_split(avocados, test_size=2000, random_state=3337)
avocado_train, avocado_valid = train_test_split(avocado_train, test_size=2249, random_state=3337)

print("Train\n", avocado_train.describe(include="all"), "\n")
print("Valid\n", avocado_valid.describe(include="all"), "\n")
print("Test\n", avocado_test.describe(include="all"))

avocado_train.to_csv("avocado.data.train", index=False)
avocado_valid.to_csv("avocado.data.valid", index=False)
avocado_test.to_csv("avocado.data.test", index=False)