# -- Code Cell --
import pandas as pd

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

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

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

# -- Code Cell --
train_task1 = train[(train['subtaskID'] == 1)]
train_task1

# -- Code Cell --
train_task1_query = train_task1['datapointID']
train_task1 = train_task1.drop(columns='datapointID')

# -- Code Cell --
from sklearn.impute import SimpleImputer,KNNImputer

imp= SimpleImputer(strategy='median').set_output(transform="pandas")
train_task1 = imp.fit_transform(train_task1)

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

# -- Code Cell --
!pip install lightgbm

# -- Code Cell --
from sklearn.model_selection import train_test_split
from sklearn.metrics import f1_score
from catboost import CatBoostClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestClassifier
from lightgbm import LGBMClassifier
X = train_task1.drop(columns = 'Target')
y = train_task1['Target']
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 = LGBMClassifier()
model.fit(X_train,y_train)
pred = model.predict(X_test)
acc = f1_score(y_test,pred)
acc

# -- Code Cell --
test_task1 = test[(test['subtaskID'] == 1)]
test_task1

# -- Code Cell --
test_task1_query = test_task1['datapointID']
test_task1 = test_task1.drop(columns = 'datapointID')

# -- Code Cell --
model.fit(X,y)
pred_final = model.predict(test_task1)

# -- Code Cell --
train_task2 = train[train['subtaskID'] == 2]
train_task2 = train_task2.drop(columns='datapointID')

# -- Code Cell --
X = train_task2.drop(columns = 'Target')
y = train_task2['Target']
X_train, X_test, y_train, y_test = train_test_split(X,y, test_size=0.2, random_state=42)
model = CatBoostClassifier()
model.fit(X_train,y_train)
pred = model.predict(X_test)
acc = f1_score(y_test,pred)
acc

# -- Code Cell --
test_task2 = test[test['subtaskID']==2]
test_task2_querys = test_task2['datapointID']
test_task2 = test_task2.drop(columns='datapointID')

# -- Code Cell --
model.fit(X,y)
pred_final_2 = model.predict(test_task2)

# -- Code Cell --
re_index_array_task1_querry = []
for i in test_task1_query:
    re_index_array_task1_querry.append(i)
len(re_index_array_task1_querry)

# -- Code Cell --
re_index_array_task1_querry[14974]

# -- Code Cell --
re_index_array_task2_querry = []
for i in test_task2_querys:
    re_index_array_task2_querry.append(i)
len(re_index_array_task2_querry)

# -- Code Cell --
sda = pd.read_csv('sample_output.csv')
len(sda[sda['subtaskID'] == 1])

# -- Code Cell --
rows = []
for querry in test['datapointID']:
    for idx,elem in enumerate(re_index_array_task1_querry):   
        if querry == elem:
            rows.append({"subtaskID":1,'datapointID':querry,'answer':int(pred_final[idx])})
    for idx,elem in enumerate(re_index_array_task2_querry):   
        if querry == elem:
            rows.append({"subtaskID":2,'datapointID':querry,'answer':int(pred_final_2[idx])})

# -- Code Cell --
rows

# -- Code Cell --
for row in rows:
    row['answer'] = int(row['answer'])

# -- Code Cell --
pd.DataFrame(rows).to_csv('subs.csv',index=False)

# -- Code Cell --
