# -- Code Cell --
import pandas as pd
train = pd.read_csv('train.csv')
test = pd.read_csv('test.csv')

# -- Code Cell --
train.head()

# -- Code Cell --
train['engine'].value_counts().nunique()

# -- Code Cell --
train.info()

# -- Code Cell --
# cat_cols2 = train.select_dtypes(include=object).columns.to_list()

# -- Code Cell --
train = train.dropna(subset=['price'])

# -- Code Cell --
num_cols = ['cylinders','mileage','doors']
cat_cols = ['engine','fuel','transmission','trim','body','exterior_color','interior_color']
from sklearn.impute import SimpleImputer
imp_num = SimpleImputer(strategy='mean')
train[num_cols] = imp_num.fit_transform(train[num_cols])
cat_num = SimpleImputer(strategy='most_frequent')
train[cat_cols] = cat_num.fit_transform(train[cat_cols])

# -- Code Cell --
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.linear_model import LinearRegression
from sklearn.metrics import r2_score
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
from catboost import CatBoostRegressor

X = train.drop(columns=['price','id'])
combined = pd.concat([train.drop(columns=['price','id']), test.drop(columns=['id'])], axis=0)
enc = OneHotEncoder(sparse_output=False)
enc.fit(combined) 
X = enc.transform(X)
y = train['price']
X_train, X_test, y_train, y_test = train_test_split(X,y,test_size=0.2, random_state=42)
scaler = StandardScaler()
X_train_sc = scaler.fit_transform(X_train)
X_test_sc = scaler.transform(X_test)
# model = LinearRegression()
model = RandomForestRegressor(random_state=42)
model.fit(X_train_sc,y_train)
acc = r2_score(y_test, model.predict(X_test_sc))
acc

# -- Code Cell --
# model = CatBoostRegressor(cat_features=cat_cols2,random_state=42)
# X2 = train.drop(columns = ['price','id'])
# y2 = train['price']
# X_train2, X_test2, y_train2,y_test2 =train_test_split(X2,y2, test_size=0.2, random_state=42)
# model.fit(X_train2, y_train2)
# preds = model.predict(X_test2)
# acc2 = r2_score(y_test2, preds)
# model.fit(X,y)

# -- Code Cell --


# -- Code Cell --
len(X)

# -- Code Cell --
test.head()

# -- Code Cell --
ids = test['id']
test = test.drop(columns='id')
test[num_cols] = imp_num.transform(test[num_cols])
test[cat_cols] = cat_num.transform(test[cat_cols])
test_enc = enc.transform(test)
test_sc = scaler.transform(test_enc)

pred = model.predict(test_sc)

# -- Code Cell --
pred

# -- Code Cell --
ids

# -- Code Cell --
t = train['price'].groupby(train['cylinders']).mean().max()
t

# -- Code Cell --
rows = []
rows.append({'subtaskID':1, 'datapointID':1, 'answer':t})
for idx, row in test.iterrows():
    rows.append({'subtaskID':2,'datapointID':ids[idx],'answer':pred[idx]})
pd.DataFrame(rows).to_csv('submission.csv',index=False)

# -- Code Cell --
