# -- Code Cell --
# %% [ONE CELL] CatBoost with native categorical handling + submission file
import pandas as pd
import numpy as np
from catboost import CatBoostClassifier
from sklearn.model_selection import StratifiedKFold

# -----------------------------
# Load
# -----------------------------
train = pd.read_csv("train.csv")
test  = pd.read_csv("test.csv")

# -----------------------------
# Subtask 1: label ratio (least common / most common)
# -----------------------------
vc = train["HadHeartAttack"].astype(str).str.strip().value_counts()
ratio = (vc.min() / vc.max())

# -----------------------------
# Prep data (NO manual mapping of categoricals)
# - CatBoost can handle NaNs and categoricals directly.
# - Convert object cols to string to ensure CatBoost treats them as categorical.
# -----------------------------
ID_COL = "ID"
TARGET = "HadHeartAttack"
DROP_COLS = [ID_COL]
if "State" in train.columns: DROP_COLS.append("State")  # keep if you want; remove this line if you want State as feature

X = train.drop(columns=[TARGET] + DROP_COLS, errors="ignore")
y_raw = train[TARGET].astype(str).str.strip()

# map target to 0/1 robustly (handles Yes/No or 1/0)
y = y_raw.map({"Yes": 1, "No": 0, "1": 1, "0": 0}).astype(int)

X_test = test.drop(columns=DROP_COLS, errors="ignore")
test_ids = test[ID_COL].copy()

# make sure both have same columns (add missing cols as NaN)
for c in set(X.columns) - set(X_test.columns):
    X_test[c] = np.nan
for c in set(X_test.columns) - set(X.columns):
    X[c] = np.nan
X = X[X_test.columns]  # align order

# categorical feature indices (CatBoost needs indices, not names)
cat_cols = X.select_dtypes(include=["object", "category", "bool"]).columns.tolist()
for c in cat_cols:
    X[c] = X[c].astype(str).fillna("NA")
    X_test[c] = X_test[c].astype(str).fillna("NA")
cat_features = [X.columns.get_loc(c) for c in cat_cols]

# -----------------------------
# Train (CV bagging) + predict proba
# -----------------------------
N_FOLDS = 5
skf = StratifiedKFold(n_splits=N_FOLDS, shuffle=True, random_state=42)

test_proba = np.zeros(len(X_test), dtype=float)

for fold, (tr_idx, va_idx) in enumerate(skf.split(X, y), 1):
    model = CatBoostClassifier(
        loss_function="Logloss",
        eval_metric="F1",
        iterations=4000,
        learning_rate=0.03,
        depth=8,
        l2_leaf_reg=4.0,
        random_seed=42 + fold,
        verbose=200,
        auto_class_weights="Balanced"
    )
    model.fit(
        X.iloc[tr_idx], y.iloc[tr_idx],
        cat_features=cat_features,
        eval_set=(X.iloc[va_idx], y.iloc[va_idx]),
        use_best_model=True
    )
    test_proba += model.predict_proba(X_test)[:, 1] / N_FOLDS

# default threshold 0.5 (you can tune if you want)
y_pred = (test_proba >= 0.5).astype(int)
pred_yesno = pd.Series(y_pred).map({1: "Yes", 0: "No"})

# -----------------------------
# Build output.csv in required format
# -----------------------------
sub1 = pd.DataFrame({"subtaskID": [1], "datapointID": [0], "answer": [f"{ratio:.6f}"]})
sub2 = pd.DataFrame({"subtaskID": 2, "datapointID": test_ids.values, "answer": pred_yesno.values})
output = pd.concat([sub1, sub2], ignore_index=True)
output.to_csv("output.csv", index=False)

print("Saved: output.csv")
print(f"Subtask1 ratio = {ratio:.6f}")
print(f"Predictions: Yes={int((y_pred==1).sum())}, No={int((y_pred==0).sum())}")


# -- Code Cell --
# %%
import pandas as pd
import numpy as np
from sklearn.model_selection import StratifiedKFold, train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from catboost import CatBoostClassifier
from xgboost import XGBClassifier
from lightgbm import LGBMClassifier
from sklearn.metrics import f1_score
from scipy.optimize import minimize
from sklearn.impute import SimpleImputer
import warnings
warnings.filterwarnings('ignore')

# %%
# Load data
train = pd.read_csv('train.csv')
test = pd.read_csv('test.csv')

# %%
# Calculate ratio for Subtask 1
ratio = train['HadHeartAttack'].value_counts().min() / train['HadHeartAttack'].value_counts().max()
print(f"Ratio for Subtask 1: {ratio:.6f}")

# %%
# Identify categorical and numerical columns
cat = train.drop(columns=["HadHeartAttack"]).select_dtypes(include="object").columns
num = train.drop(columns=["HadHeartAttack"]).select_dtypes(exclude="object").columns

# Impute missing values
num_imputer = SimpleImputer(strategy="median")  # Changed to median for better outlier handling
train[num] = num_imputer.fit_transform(train[num])
cat_imputer = SimpleImputer(strategy="most_frequent")
train[cat] = cat_imputer.fit_transform(train[cat])

# %%
# Define mapping dictionaries
map_ch = {
    "Within past year (anytime less than 12 months ago)": 1,
    "Within past 2 years (1 year but less than 2 years ago)": 2,
    "Within past 5 years (2 years but less than 5 years ago)": 5,
    "5 or more years ago": 10
}
map_health = {'Poor': 0, 'Fair': 1, 'Good': 3, 'Very good': 4, 'Excellent': 5}
map_sex = {'Female': 0, 'Male': 1}
map_diab = {
    'No': 0,
    'No, pre-diabetes or borderline diabetes': 5,
    'Yes, but only during pregnancy (female)': 2,
    'Yes': 10
}
map_YES_NO = {'Yes': 1, 'No': 0}
map_covid = {
    'No': 0,
    'Tested positive using home test without a health professional': 1,
    'Yes': 2
}
map_teeth = {'None of them': 0, '1 to 5': 2, "6 or more, but not all": 3, "All": 10}
map_TETANUS = {
    'No, did not receive any tetanus shot in the past 10 years': 0,
    'Yes, received tetanus shot but not sure what type': 1,
    'Yes, received tetanus shot, but not Tdap': 2,
    'Yes, received Tdap': 3
}
map_SMOKER = {
    'Never smoked': 10,
    'Former smoker': 6,
    'Current smoker - now smokes some days': 3,
    'Current smoker - now smokes every day': 1
}
map_VAPE = {
    'Never used e-cigarettes in my entire life': 10,
    'Not at all (right now)': 6,
    'Use them some days': 3,
    'Use them every day': 1
}
map_age = {
    "Age 18 to 24": 0, "Age 25 to 29": 1, "Age 30 to 34": 2, "Age 35 to 39": 3,
    "Age 40 to 44": 4, "Age 45 to 49": 5, "Age 50 to 54": 6, "Age 55 to 59": 7,
    "Age 60 to 64": 8, "Age 65 to 69": 9, "Age 70 to 74": 10, "Age 75 to 79": 11,
    "Age 80 or older": 12
}

# %%
# Apply mappings to train data
train = pd.get_dummies(train, columns=["RaceEthnicityCategory"], drop_first=True)
train['LastCheckupTime'] = train['LastCheckupTime'].astype(str).map(map_ch)
train['Sex'] = train['Sex'].astype(str).map(map_sex)
train['GeneralHealth'] = train['GeneralHealth'].astype(str).map(map_health)
train['RemovedTeeth'] = train['RemovedTeeth'].astype(str).map(map_teeth)

yes_no_cols = [
    "HadHeartAttack", "HadAngina", "HadStroke", "HadAsthma",
    "HadSkinCancer", "HadCOPD", "HadDepressiveDisorder",
    "HadKidneyDisease", "HadArthritis",
    "DeafOrHardOfHearing", "BlindOrVisionDifficulty",
    "DifficultyConcentrating", "DifficultyWalking", 'PhysicalActivities',
    "DifficultyDressingBathing", "DifficultyErrands", "ChestScan",
    "AlcoholDrinkers", "HIVTesting", "FluVaxLast12", "PneumoVaxEver", "HighRiskLastYear"
]

for col in yes_no_cols:
    if col in train.columns:
        train[col] = train[col].astype(str).map(map_YES_NO)

train['SmokerStatus'] = train['SmokerStatus'].astype(str).map(map_SMOKER)
train['ECigaretteUsage'] = train['ECigaretteUsage'].astype(str).map(map_VAPE)
train['AgeCategory'] = train['AgeCategory'].astype(str).map(map_age)
train['TetanusLast10Tdap'] = train['TetanusLast10Tdap'].astype(str).map(map_TETANUS)
train['HadDiabetes'] = train['HadDiabetes'].astype(str).map(map_diab)
train['CovidPos'] = train['CovidPos'].astype(str).map(map_covid)

# %%
# Drop unnecessary columns and handle missing values
train = train.drop(columns=['ID', 'State'])
train = train.dropna()
train = train.astype(float)

print(f"Training samples after cleaning: {len(train)}")

# %%
# Prepare features and target
X = train.drop(columns='HadHeartAttack').copy()
y = train['HadHeartAttack'].copy()
feature_cols = X.columns.tolist()

print(f"Number of features: {len(feature_cols)}")
print(f"Class distribution: {y.value_counts().to_dict()}")

# %%
# OOF Ensemble with multiple models
N_FOLDS = 7  # Increased folds for better generalization
skf = StratifiedKFold(n_splits=N_FOLDS, shuffle=True, random_state=42)

# Initialize OOF arrays for 5 models
oof_lr = np.zeros(len(X))
oof_rf = np.zeros(len(X))
oof_cb = np.zeros(len(X))
oof_xgb = np.zeros(len(X))
oof_lgbm = np.zeros(len(X))

# Store models
models_lr = []
models_rf = []
models_cb = []
models_xgb = []
models_lgbm = []
scalers = []

print("=" * 60)
print("OOF Training with 5 Models")
print("=" * 60)

for fold, (train_idx, val_idx) in enumerate(skf.split(X, y)):
    print(f"\nFold {fold + 1}/{N_FOLDS}")
    
    X_train_fold = X.iloc[train_idx].values
    y_train_fold = y.iloc[train_idx].values
    X_val_fold = X.iloc[val_idx].values
    y_val_fold = y.iloc[val_idx].values
    
    # Scale features
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train_fold)
    X_val_scaled = scaler.transform(X_val_fold)
    scalers.append(scaler)
    
    # Calculate class weights for imbalanced data
    class_weight_dict = {0: 1, 1: len(y_train_fold[y_train_fold == 0]) / len(y_train_fold[y_train_fold == 1])}
    
    # Model 1: Logistic Regression
    lr = LogisticRegression(
        random_state=42,
        max_iter=2000,
        class_weight='balanced',
        C=0.1,
        solver='saga'
    )
    lr.fit(X_train_scaled, y_train_fold)
    oof_lr[val_idx] = lr.predict_proba(X_val_scaled)[:, 1]
    models_lr.append(lr)
    
    # Model 2: Random Forest
    rf = RandomForestClassifier(
        n_estimators=500,
        max_depth=15,
        min_samples_split=20,
        min_samples_leaf=10,
        class_weight='balanced',
        random_state=42,
        n_jobs=-1
    )
    rf.fit(X_train_scaled, y_train_fold)
    oof_rf[val_idx] = rf.predict_proba(X_val_scaled)[:, 1]
    models_rf.append(rf)
    
    # Model 3: CatBoost
    cb = CatBoostClassifier(
        iterations=1000,
        learning_rate=0.03,
        depth=6,
        l2_leaf_reg=3,
        random_state=42,
        verbose=0,
        scale_pos_weight=class_weight_dict[1]
    )
    cb.fit(X_train_scaled, y_train_fold)
    oof_cb[val_idx] = cb.predict_proba(X_val_scaled)[:, 1]
    models_cb.append(cb)
    
    # Model 4: XGBoost
    xgb = XGBClassifier(
        n_estimators=1000,
        learning_rate=0.03,
        max_depth=6,
        subsample=0.8,
        colsample_bytree=0.8,
        scale_pos_weight=class_weight_dict[1],
        random_state=42,
        n_jobs=-1,
        eval_metric='logloss'
    )
    xgb.fit(X_train_scaled, y_train_fold)
    oof_xgb[val_idx] = xgb.predict_proba(X_val_scaled)[:, 1]
    models_xgb.append(xgb)
    
    # Model 5: LightGBM
    lgbm = LGBMClassifier(
        n_estimators=1000,
        learning_rate=0.03,
        max_depth=6,
        num_leaves=31,
        subsample=0.8,
        colsample_bytree=0.8,
        scale_pos_weight=class_weight_dict[1],
        random_state=42,
        n_jobs=-1,
        verbose=-1
    )
    lgbm.fit(X_train_scaled, y_train_fold)
    oof_lgbm[val_idx] = lgbm.predict_proba(X_val_scaled)[:, 1]
    models_lgbm.append(lgbm)
    
    # Print fold scores
    print(f"  LR   F1: {f1_score(y_val_fold, (oof_lr[val_idx] >= 0.5).astype(int)):.4f}")
    print(f"  RF   F1: {f1_score(y_val_fold, (oof_rf[val_idx] >= 0.5).astype(int)):.4f}")
    print(f"  CB   F1: {f1_score(y_val_fold, (oof_cb[val_idx] >= 0.5).astype(int)):.4f}")
    print(f"  XGB  F1: {f1_score(y_val_fold, (oof_xgb[val_idx] >= 0.5).astype(int)):.4f}")
    print(f"  LGBM F1: {f1_score(y_val_fold, (oof_lgbm[val_idx] >= 0.5).astype(int)):.4f}")

print("\n" + "=" * 60)
print("Individual Model OOF F1 Scores")
print("=" * 60)
print(f"Logistic Regression: {f1_score(y, (oof_lr >= 0.5).astype(int)):.4f}")
print(f"Random Forest:       {f1_score(y, (oof_rf >= 0.5).astype(int)):.4f}")
print(f"CatBoost:            {f1_score(y, (oof_cb >= 0.5).astype(int)):.4f}")
print(f"XGBoost:             {f1_score(y, (oof_xgb >= 0.5).astype(int)):.4f}")
print(f"LightGBM:            {f1_score(y, (oof_lgbm >= 0.5).astype(int)):.4f}")

# %%
# Optimize ensemble weights and threshold
def objective(params):
    weights = params[:5]
    threshold = params[5]
    weights = weights / weights.sum()  # Normalize weights
    
    ensemble_pred = (weights[0] * oof_lr + 
                     weights[1] * oof_rf + 
                     weights[2] * oof_cb +
                     weights[3] * oof_xgb +
                     weights[4] * oof_lgbm)
    
    y_pred = (ensemble_pred >= threshold).astype(int)
    return -f1_score(y, y_pred)  # Negative because we minimize

# Initial guess: equal weights and threshold 0.5
initial_params = np.array([0.2, 0.2, 0.2, 0.2, 0.2, 0.5])

# Bounds: weights [0,1], threshold [0.3, 0.7]
bounds = [(0, 1)] * 5 + [(0.3, 0.7)]

# Optimize
result = minimize(objective, initial_params, method='Nelder-Mead', bounds=bounds)
optimal_weights = result.x[:5]
optimal_weights = optimal_weights / optimal_weights.sum()
optimal_threshold = result.x[5]

print("\n" + "=" * 60)
print("Optimized Ensemble")
print("=" * 60)
print(f"Optimal weights: LR={optimal_weights[0]:.4f}, RF={optimal_weights[1]:.4f}, "
      f"CB={optimal_weights[2]:.4f}, XGB={optimal_weights[3]:.4f}, LGBM={optimal_weights[4]:.4f}")
print(f"Optimal threshold: {optimal_threshold:.4f}")

ensemble_oof = (optimal_weights[0] * oof_lr + 
                optimal_weights[1] * oof_rf + 
                optimal_weights[2] * oof_cb +
                optimal_weights[3] * oof_xgb +
                optimal_weights[4] * oof_lgbm)

y_pred_oof = (ensemble_oof >= optimal_threshold).astype(int)
oof_f1 = f1_score(y, y_pred_oof)
print(f"Optimized Ensemble OOF F1: {oof_f1:.4f}")

# %%
# Process test data
test = pd.read_csv('test.csv')
test_ids = test['ID'].copy()

test[num] = num_imputer.transform(test[num])
test[cat] = cat_imputer.transform(test[cat])

test = pd.get_dummies(test, columns=["RaceEthnicityCategory"], drop_first=True)
test["LastCheckupTime"] = test["LastCheckupTime"].astype(str).map(map_ch)
test["Sex"] = test["Sex"].astype(str).map(map_sex)
test["GeneralHealth"] = test["GeneralHealth"].astype(str).map(map_health)
test["RemovedTeeth"] = test["RemovedTeeth"].astype(str).map(map_teeth)

yes_no_cols_test = [c for c in yes_no_cols if c in test.columns and c != "HadHeartAttack"]
for col in yes_no_cols_test:
    test[col] = test[col].astype(str).map(map_YES_NO)

test["SmokerStatus"] = test["SmokerStatus"].astype(str).map(map_SMOKER)
test["ECigaretteUsage"] = test["ECigaretteUsage"].astype(str).map(map_VAPE)
test["AgeCategory"] = test["AgeCategory"].astype(str).map(map_age)
test["TetanusLast10Tdap"] = test["TetanusLast10Tdap"].astype(str).map(map_TETANUS)
test["HadDiabetes"] = test["HadDiabetes"].astype(str).map(map_diab)
test["CovidPos"] = test["CovidPos"].astype(str).map(map_covid)

test = test.drop(columns=["ID", "State"], errors="ignore")

# Align columns with training data
for col in feature_cols:
    if col not in test.columns:
        test[col] = 0
test = test[feature_cols]

# %%
# Generate predictions on test set
test_preds_lr = np.zeros(len(test))
test_preds_rf = np.zeros(len(test))
test_preds_cb = np.zeros(len(test))
test_preds_xgb = np.zeros(len(test))
test_preds_lgbm = np.zeros(len(test))

for fold in range(N_FOLDS):
    X_test_scaled = scalers[fold].transform(test.values)
    test_preds_lr += models_lr[fold].predict_proba(X_test_scaled)[:, 1] / N_FOLDS
    test_preds_rf += models_rf[fold].predict_proba(X_test_scaled)[:, 1] / N_FOLDS
    test_preds_cb += models_cb[fold].predict_proba(X_test_scaled)[:, 1] / N_FOLDS
    test_preds_xgb += models_xgb[fold].predict_proba(X_test_scaled)[:, 1] / N_FOLDS
    test_preds_lgbm += models_lgbm[fold].predict_proba(X_test_scaled)[:, 1] / N_FOLDS

# Ensemble with optimal weights
test_ensemble = (optimal_weights[0] * test_preds_lr + 
                 optimal_weights[1] * test_preds_rf + 
                 optimal_weights[2] * test_preds_cb +
                 optimal_weights[3] * test_preds_xgb +
                 optimal_weights[4] * test_preds_lgbm)

y_pred_final = (test_ensemble >= optimal_threshold).astype(int)

print(f"\nTest Predictions: Yes={y_pred_final.sum()}, No={len(y_pred_final) - y_pred_final.sum()}")

# %%
# Create submission file
sub1 = pd.DataFrame({
    "subtaskID": [1],
    "datapointID": [0],
    "answer": [f"{ratio:.6f}"]
})

pred_yesno = pd.Series(y_pred_final).map({1: "Yes", 0: "No"})

sub2 = pd.DataFrame({
    "subtaskID": 2,
    "datapointID": test_ids.values,
    "answer": pred_yesno.values
})

output = pd.concat([sub1, sub2], ignore_index=True)
output.to_csv("output.csv", index=False)

print("\n" + "=" * 60)
print("Submission file created: output.csv")
print("=" * 60)