

import pandas as pd
import numpy as np

train = pd.read_csv("train.csv")
test  = pd.read_csv("test.csv")

TARGET = "Diagnosis"

X = train.drop(columns=[TARGET])
y = train[TARGET]
print(y.value_counts(normalize=True)) #vedem dacă datele sunt dezechilibrate
from sklearn.model_selection import train_test_split

X_train, X_val, y_train, y_val = train_test_split(
    X, y,
    test_size=0.2,
    stratify=y,
    random_state=42
)
!pip install catboost
from catboost import CatBoostClassifier, Pool
from sklearn.metrics import roc_auc_score
train_pool = Pool(X_train, y_train)
val_pool   = Pool(X_val, y_val)
model = CatBoostClassifier(
    loss_function="Logloss",   # probabilități corecte
    eval_metric="AUC",         # ce optimizează
    iterations=5000,
    learning_rate=0.03,
    depth=6,
    l2_leaf_reg=3,
    random_seed=42,
    od_type="Iter",
    od_wait=200,
    auto_class_weights="Balanced",
    verbose=200
)
model.fit(train_pool, eval_set=val_pool, use_best_model=True)
val_proba = model.predict_proba(X_val)[:, 1]
auc = roc_auc_score(y_val, val_proba)

print("Validation ROC AUC:", auc)
full_pool = Pool(X, y)

final_model = CatBoostClassifier(
    loss_function="Logloss",
    eval_metric="AUC",
    iterations=int(model.get_best_iteration()),
    learning_rate=0.03,
    depth=6,
    l2_leaf_reg=3,
    random_seed=42,
    auto_class_weights="Balanced",
    verbose=200
)

final_model.fit(full_pool)
#TASK 3:
probs_ml = final_model.predict_proba(test)[:, 1]
#TASK 1:
age_counts = train["Age"].value_counts()
#TASK 2:
age_smoker_pct = train.groupby("Age")["Smoking"].mean() * 100
subtask1_answer = {}
subtask2_answer = {}

for _, row in test.iterrows():
    pid = row["PatientID"]
    age = row["Age"]

    subtask1_answer[pid] = int(age_counts.get(age, 0))
    subtask2_answer[pid] = float(age_smoker_pct.get(age, 0.0))
submission_rows = []

for i, row in test.iterrows():
    pid = row["PatientID"]

    submission_rows.append([1, pid, subtask1_answer[pid]])
    submission_rows.append([2, pid, subtask2_answer[pid]])
    submission_rows.append([3, pid, float(probs_ml[i])])

submission = pd.DataFrame(
    submission_rows,
    columns=["subtaskID", "datapointID", "answer"]
)

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