import numpy as np
import pandas as pd
import torch
from sklearn.preprocessing import StandardScaler
from torch.utils.data import Dataset, DataLoader
from transformers import AutoTokenizer, AutoModel
from tqdm.auto import tqdm
from sklearn.ensemble import AdaBoostClassifier
from sklearn.tree import DecisionTreeClassifier
from sklearn.metrics import f1_score, classification_report
from sklearn.model_selection import train_test_split

train = pd.read_csv("/kaggle/input/tst-day-3-upsolving/train.csv")
test = pd.read_csv("/kaggle/input/tst-day-3-upsolving/test.csv")

train_df, val_df = train_test_split(
    train,
    test_size=0.2,
    random_state=42,
    stratify=train["label"]
)

print("Train :", train_df.shape)
print("Validation :", val_df.shape)

MODEL_NAME = "FacebookAI/roberta-base"
MAX_LENGTH = 256
BATCH_SIZE = 8

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
encoder = AutoModel.from_pretrained(MODEL_NAME).to(device)
encoder.eval()

class CodeDataset(Dataset):
    def __init__(self, texts, tokenizer, max_length=256):
        self.texts = texts.fillna("").astype(str).tolist()
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        return self.tokenizer(
            self.texts[idx],
            truncation=True,
            padding="max_length",
            max_length=self.max_length,
            return_tensors="pt"
        )

def extract_embeddings(texts):
    dataset = CodeDataset(texts, tokenizer, MAX_LENGTH)
    loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=False)
    all_embeddings = []

    with torch.no_grad():
        for batch in tqdm(loader):
            input_ids = batch["input_ids"].squeeze(1).to(device)
            attention_mask = batch["attention_mask"].squeeze(1).to(device)

            outputs = encoder(
                input_ids=input_ids,
                attention_mask=attention_mask
            )

            hidden = outputs.last_hidden_state
            mask = attention_mask.unsqueeze(-1).float()
            summed = (hidden * mask).sum(dim=1)
            counts = mask.sum(dim=1).clamp(min=1)
            embeddings = summed / counts

            all_embeddings.append(embeddings.cpu().numpy())

    return np.concatenate(all_embeddings, axis=0)

sc = StandardScaler()

X_train_emb = extract_embeddings(train_df["code"])
X_train_emb = sc.fit_transform(X_train_emb)

X_val_emb = extract_embeddings(val_df["code"])
X_val_emb = sc.transform(X_val_emb)

X_test_emb = extract_embeddings(test["code"])
X_test_emb = sc.transform(X_test_emb)

y_train = train_df["label"].to_numpy()
y_val = val_df["label"].to_numpy()

print("Train embeddings:", X_train_emb.shape)
print("Validation embeddings:", X_val_emb.shape)
print("Test embeddings:", X_test_emb.shape)

ada = AdaBoostClassifier(
    estimator=DecisionTreeClassifier(
        max_depth=2,
        min_samples_leaf=5,
        class_weight="balanced",
        random_state=42
    ),
    n_estimators=200,
    learning_rate=0.05,
    random_state=42
)

ada.fit(X_train_emb, y_train)

val_pred = ada.predict(X_val_emb)

print("Macro F1:", f1_score(
    y_val,
    val_pred,
    average="macro"
))

print(classification_report(
    y_val,
    val_pred,
    labels=[0, 1, 2],
    target_names=["easy", "medium", "hard"],
    zero_division=0
))

test_pred = ada.predict(X_test_emb).astype(int)

id2label = {
    0: "easy",
    1: "medium",
    2: "hard"
}

submission = pd.DataFrame({
    "id": test["id"],
    "difficulty": [id2label[int(p)] for p in test_pred]
})

submission.to_csv(
    "submission_roberta_ada.csv",
    index=False
)