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

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

# -- Code Cell --
from sentence_transformers import SentenceTransformer
embeddings = []
model = SentenceTransformer('all-mpnet-base-v2')
for i in range(len(train)):
    embedding = model.encode(train['text'][i])
    embeddings.append(embedding)

# -- Code Cell --
len(embeddings)

# -- Code Cell --
from sklearn.metrics import f1_score
from sklearn.linear_model import LogisticRegression
from catboost import CatBoostClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split

X = embeddings
y = train['label']
X_train, X_test, y_train, y_test = train_test_split(X,y,test_size=0.2, random_state=42)
model_ml = CatBoostClassifier()
model_ml.fit(X_train, y_train)
pred = model_ml.predict(X_test)
acc = f1_score(y_test,pred)
acc

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

# -- Code Cell --
embeddings_test = []
for i in range(len(test[test['subtaskID'] == 1])):
    embedding = model.encode(test[test['subtaskID'] == 1]['text'][i])
    embeddings_test.append(embedding)

# -- Code Cell --
model_ml.fit(X,y)
pred_final = model_ml.predict(embeddings_test)

# -- Code Cell --
label_texts = {
    "SCIENCE": "Scientific research and discovery in physics, biology, chemistry, medicine, and technology",
    "BUSINESS": "Business news about companies, markets, finance, economics, stocks, and corporate strategy",
    "CRIME": "Criminal activity including murder, theft, robbery, arrest, police investigation, and court trials",
    "RELIGION": "Religious beliefs, faith, worship, God, church, prayer, and spiritual practices"
}

label_names = list(label_texts.keys())
label_emb = model.encode(list(label_texts.values()))

test2 = test[test['subtaskID'] == 2]
test2_emb = model.encode(test2['text'].tolist())

from sentence_transformers import util
scores = util.cos_sim(test2_emb, label_emb) 
preds = [label_names[i] for i in scores.argmax(dim=1)]

# -- Code Cell --
rows = []
for idx, row in test[test['subtaskID'] == 1].iterrows():
    rows.append({"subtaskID": 1, "datapointID": row['ID'], "answer": pred_final[idx]})

for idx, row in test[test['subtaskID'] == 2].iterrows():
    rows.append({"subtaskID": 2, "datapointID": row['ID'], "answer": preds[idx]})

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

# -- Code Cell --
