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

# -- Code Cell --
IDs = test['Patient_ID']

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

# -- Code Cell --
train['Gender'].value_counts()

# -- Code Cell --
gender_map = {'Female':0, 'Male':1}
gender_map2 = {'Female':0, 'Male':1,'Other':2}
train['Gender'] = train['Gender'].astype(str).map(gender_map2)
test['Gender'] = test['Gender'].astype(str).map(gender_map2)

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

# -- Code Cell --
train = train.drop(columns='Patient_ID')
test = test.drop(columns='Patient_ID')

# -- Code Cell --
train['Symptom_Count'].value_counts()

# -- Code Cell --
#MAX = 6

# -- Code Cell --
def add_symptom_columns(df, col="Symptoms", n=6, sep=","):

    parts = (
        df[col]
        .fillna("")
        .astype(str)
        .str.split(sep)
        .apply(lambda xs: [x.strip() for x in xs if x.strip()])  
        .apply(lambda xs: xs[:n])  
    )

    parts_df = pd.DataFrame(parts.tolist(), index=df.index)
    parts_df = parts_df.reindex(columns=range(n), fill_value=0)
    parts_df.columns = [f"Symptoms{i+1}" for i in range(n)]

    return df.join(parts_df)

train = add_symptom_columns(train)
test= add_symptom_columns(test)

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

# -- Code Cell --
train = train.drop(columns = ['Symptoms'])
test = test.drop(columns = ['Symptoms'])

# -- Code Cell --
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
from catboost import CatBoostClassifier

symptom_cols = [f"Symptoms{i}" for i in range(1, 7)]
feature_cols = ['Gender', 'Age'] + symptom_cols

X = train[feature_cols]
X[symptom_cols] = X[symptom_cols].fillna("missing").astype(str)
y = train['Disease']

cat_features = [feature_cols.index(col) for col in symptom_cols]

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

model = CatBoostClassifier(
    iterations=1200,
    depth=8,
    learning_rate=0.05,
    loss_function="MultiClass",
    verbose=False
)

model.fit(X_train, y_train, cat_features=cat_features)

y_pred = model.predict(X_test)
acc = accuracy_score(y_test, y_pred)

acc


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

# -- Code Cell --
test= test.drop(columns = 'Symptom_Count')

# -- Code Cell --
import numpy as np
X_pred = test[feature_cols]
X_pred[symptom_cols] = X_pred[symptom_cols].fillna("missing").astype(str)

y_pred_final = model.predict(X_pred)

# make it 1D
y_pred_final = np.array(y_pred_final).ravel()   

submission = pd.DataFrame({
    "Patient_ID": IDs.values,          
    "Disease": y_pred_final.astype(str)
})

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

# -- Code Cell --
