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

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

# -- Code Cell --
answer1 = train[(train['native_country'] != 'United-States') & (train['income'] == '>50K')]['native_country'].value_counts().idxmax()

# -- Code Cell --
answer2 = train.groupby('occupation')['income'].apply(lambda x: (x == '>50K').mean()).idxmax()

# -- Code Cell --
tr = train.drop(columns=['profile_description']).copy()
te = test.drop(columns=['profile_description']).copy()
#SCOATEM FEATURE-ul DE TEXT

# -- Code Cell --
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import OneHotEncoder

cat_cols = tr.select_dtypes(include='object').columns 
num_cols = [c for c in tr.columns if c not in cat_cols]  
cat_cols = [c for c in cat_cols if c != 'income']

imp = SimpleImputer(strategy='most_frequent')
tr[cat_cols] = imp.fit_transform(tr[cat_cols])
te[cat_cols] = imp.transform(te[cat_cols])

ohe = OneHotEncoder(sparse_output=False, handle_unknown='ignore')
X_train_cat = ohe.fit_transform(tr[cat_cols])
X_test_cat = ohe.transform(te[cat_cols])

X_train = np.hstack([tr[num_cols].values, X_train_cat])
X_test = np.hstack([te[num_cols].values, X_test_cat])
y_train = tr['income']

# -- Code Cell --
#Sau inainte de direct fit pe tot train, se poate face train_test_split

# -- Code Cell --
from sklearn.ensemble import GradientBoostingClassifier
model = GradientBoostingClassifier(n_estimators=200, max_depth=5, random_state=42)
model.fit(X_train, y_train)
pred3 = model.predict(X_test)

# -- Code Cell --
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.cluster import KMeans

tfidf = TfidfVectorizer(max_features=200)
X_train_tfidf = tfidf.fit_transform(train['profile_description'])
X_test_tfidf = tfidf.transform(test['profile_description'])

km = KMeans(n_clusters=5, random_state=42, n_init=10)
km.fit(X_train_tfidf)
pred4 = km.predict(X_test_tfidf)

# -- Code Cell --
rows = []
rows.append({'subtaskID': 1, 'datapointID': 1, 'answer': answer1})
rows.append({'subtaskID': 2, 'datapointID': 2, 'answer': answer2})

for idx, row in test.iterrows():
    rows.append({'subtaskID': 3, 'datapointID': row['sampleid'], 'answer': pred3[idx]})
    rows.append({'subtaskID': 4, 'datapointID': row['sampleid'], 'answer': str(pred4[idx])})

submission = pd.DataFrame(rows)
submission.to_csv('submission.csv', index=False)