# %%
import pandas as pd

# %%
train = pd.read_csv('train_data.csv')
test = pd.read_csv('test_data.csv')

# %%
train.head()

# %%
import numpy as np

# %%
np.random.seed(42)
n = len(train)
X1 = np.random.normal(0, 1, n)
X2 = np.random.uniform(-5, 5, n)
X345 = np.random.multivariate_normal(mean=[1,-1,2], cov=[[1.0,0.8,-0.3],[0.8,1,0.1],[-0.3,0.1,1]], size=n)
X3, X4, X5 = X345[:,0], X345[:,1], X345[:,2]
X6 = np.random.choice(['Copper','Aluminum','Titanium','Tungsten'], size=n, p=[0.1,0.2,0.3,0.4])
X7 = 2*X2**2 - 3*X2 + 1
coin = np.random.random(n)
X8 = np.where(coin < 0.7, np.random.normal(-3,1,n), np.random.normal(3,1,n))
X9 = np.where(X6 != 'Copper', X1*X2, np.nan)

# %%
target = []
for z in range(len(train)):
  for i in train.values[z]:
    target.append(i)

# %%
d = {"X1": X1, "X2": X2,"X3": X3, "X4": X4,"X5": X5, "X6": X6,"X7": X7, "X8": X8,"X9": X9, 'Y':target}

# %%
df = pd.DataFrame(data=d)

# %%
df

# %%
ss = df.to_csv("df_train.csv",index=False)

# %%


# %%


# %%
df.info()

# %%
# from sklearn.impute import KNNImputer,SimpleImputer
# imp = SimpleImputer().set_output(transform="pandas")
# ore_map = {"Titanium":0, "Copper":1, "Tungsten":2, "Aluminum":3}
# df['X6'] =df['X6'].map(ore_map)
# df_cleaned = imp.fit_transform(df[tag_cols])

# %%


# %%
df

# %%

from sklearn.metrics import root_mean_squared_error as RMSE
from sklearn.linear_model import LinearRegression
from sklearn.ensemble import RandomForestRegressor
from catboost import CatBoostRegressor
from lightgbm import LGBMRegressor
from sklearn.svm import SVR
from xgboost import XGBRegressor
from sklearn.model_selection import train_test_split
X = df.drop(columns = "Y")
y = df['Y']
X_train, X_test, y_train, y_test = train_test_split(X,y,test_size=0.2,random_state=42)
# from sklearn.decomposition import PCA
# pca = PCA(n_components=2).set_output(transform="pandas")
# X_train = pca.fit_transform(X_train)
# X_test = pca.transform(X_test)
X_train["X6"] = X_train["X6"].astype("category")
X_test["X6"] = X_test["X6"].astype("category")
model = CatBoostRegressor(random_state=42, cat_features=['X6'])
model.fit(X_train,y_train)
pred = model.predict(X_test)
acc = RMSE(y_test,pred)
acc

# %%
test.info()

# %%
test['X6'] =test['X6'].astype('category')

# %%
test.head()

# %%
preds = model.predict(test)

# %%
subs = pd.DataFrame({
    "subtaskID":1,
    'datapointID':range(len(preds)),
    "answer":preds
}).to_csv('subs.csv',index=False)

# %%



