# %% 
import pandas as pd
import numpy as np
from sklearn.compose import ColumnTransformer
from sklearn.impute import SimpleImputer
from sklearn.model_selection import StratifiedKFold
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from catboost import CatBoostClassifier
from sklearn.metrics import f1_score
from scipy.optimize import minimize

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

# %% Imputation
cat = train.drop(columns=["HadHeartAttack"]).select_dtypes(include="object").columns
num = train.drop(columns=["HadHeartAttack"]).select_dtypes(exclude="object").columns

num_imputer = SimpleImputer(strategy="mean")
train[num] = num_imputer.fit_transform(train[num])

cat_imputer = SimpleImputer(strategy="most_frequent")
train[cat] = cat_imputer.fit_transform(train[cat])

# %% Calculate ratio for subtask 1
ratio = train['HadHeartAttack'].value_counts().iloc[train['HadHeartAttack'].value_counts().argmin()] / \
        train['HadHeartAttack'].value_counts().iloc[train['HadHeartAttack'].value_counts().argmax()]
print(f"Ratio: {ratio:.6f}")

# %% Mappings
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, 'Excellent': 2, 'Good': 3, 'Very good': 4}
map_sex = {'Female': 0, 'Male': 1}
map_diab = {'Yes': 10, 'No': 0, 'No, pre-diabetes or borderline diabetes': 5, 
            'Yes, but only during pregnancy (female)': 2}
map_YES_NO = {'Yes': 1, 'No': 0}
map_covid = {'Yes': 2, 'No': 0, 'Tested positive using home test without a health professional': 1}
map_theeth = {'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 Tdap': 2,
    'Yes, received tetanus shot, but not Tdap': 3
}
map_SMOKER = {
    'Never smoked': 10, 'Former smoker': 6,
    'Current smoker - now smokes every day': 2,
    'Current smoker - now smokes some days': 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
}

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"
]

# %% Apply mappings to train
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_theeth)

for col in yes_no_cols:
    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)

train = train.drop(columns=['ID', 'State'])
train = train.dropna()
train = train.astype(int)

# %% Prepare X, y
X = train.drop(columns='HadHeartAttack').copy()
y = train['HadHeartAttack'].copy()

feature_cols = X.columns.tolist()

print(f"Training samples: {len(X)}")
print(f"Features: {len(feature_cols)}")

# %% OOF Predictions with StratifiedKFold
N_FOLDS = 5
skf = StratifiedKFold(n_splits=N_FOLDS, shuffle=True, random_state=42)

# Initialize OOF arrays
oof_lr = np.zeros(len(X))
oof_rf = np.zeros(len(X))
oof_cb = np.zeros(len(X))

# Store models for final prediction
models_lr = []
models_rf = []
models_cb = []
scalers = []

print("\n" + "="*60)
print("OOF Training")
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
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train_fold)
    X_val_scaled = scaler.transform(X_val_fold)
    scalers.append(scaler)
    
    # Logistic Regression
    lr = LogisticRegression(random_state=42, max_iter=1000)
    lr.fit(X_train_scaled, y_train_fold)
    oof_lr[val_idx] = lr.predict_proba(X_val_scaled)[:, 1]
    models_lr.append(lr)
    
    # Random Forest
    rf = RandomForestClassifier(n_estimators=300, 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)
    
    # CatBoost
    cb = CatBoostClassifier(random_state=42, verbose=0)
    cb.fit(X_train_scaled, y_train_fold)
    oof_cb[val_idx] = cb.predict_proba(X_val_scaled)[:, 1]
    models_cb.append(cb)
    
    # Print fold scores
    for name, oof_pred in [("LR", oof_lr), ("RF", oof_rf), ("CB", oof_cb)]:
        fold_f1 = f1_score(y_val_fold, (oof_pred[val_idx] >= 0.5).astype(int))
        print(f"  {name} F1: {fold_f1:.4f}")

# %% Print individual model OOF scores
print("\n" + "="*60)
print("Individual Model OOF F1 Scores (threshold=0.5)")
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}")

# %% Optimize weights using scipy.minimize
print("\n" + "="*60)
print("Optimizing Ensemble Weights with scipy.minimize")
print("="*60)

def neg_f1_score(weights, oof_preds, y_true, threshold):
    """Negative F1 score for minimization"""
    w = weights / weights.sum()  # Normalize weights
    ensemble_prob = w[0]*oof_preds[0] + w[1]*oof_preds[1] + w[2]*oof_preds[2]
    preds = (ensemble_prob >= threshold).astype(int)
    return -f1_score(y_true, preds)

def optimize_weights_and_threshold(oof_preds, y_true):
    """Find optimal weights and threshold"""
    best_f1 = 0
    best_weights = None
    best_threshold = None
    
    # Grid search over thresholds
    thresholds = np.linspace(0.1, 0.6, 51)
    
    for thresh in thresholds:
        # Optimize weights for this threshold
        result = minimize(
            neg_f1_score,
            x0=np.array([1/3, 1/3, 1/3]),  # Initial equal weights
            args=(oof_preds, y_true, thresh),
            method='Nelder-Mead',
            options={'maxiter': 500}
        )
        
        weights = result.x / result.x.sum()
        f1 = -result.fun
        
        if f1 > best_f1:
            best_f1 = f1
            best_weights = weights
            best_threshold = thresh
    
    return best_weights, best_threshold, best_f1

oof_preds = [oof_lr, oof_rf, oof_cb]
optimal_weights, optimal_threshold, optimal_f1 = optimize_weights_and_threshold(oof_preds, y.values)

print(f"\nOptimal Weights:")
print(f"  Logistic Regression: {optimal_weights[0]:.4f}")
print(f"  Random Forest:       {optimal_weights[1]:.4f}")
print(f"  CatBoost:            {optimal_weights[2]:.4f}")
print(f"\nOptimal Threshold: {optimal_threshold:.4f}")
print(f"Optimal OOF F1 Score: {optimal_f1:.4f}")

# %% Verify OOF ensemble score
oof_ensemble = optimal_weights[0]*oof_lr + optimal_weights[1]*oof_rf + optimal_weights[2]*oof_cb
oof_preds_final = (oof_ensemble >= optimal_threshold).astype(int)
print(f"\nVerification - OOF Ensemble F1: {f1_score(y, oof_preds_final):.4f}")

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

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_theeth)

yes_no_cols_test = [c for c in yes_no_cols if c in test.columns]
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]

# %% Final predictions using all fold models (averaging)
print("\n" + "="*60)
print("Making Final Predictions")
print("="*60)

test_preds_lr = np.zeros(len(test))
test_preds_rf = np.zeros(len(test))
test_preds_cb = 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

# Ensemble with optimal weights
test_ensemble = optimal_weights[0]*test_preds_lr + optimal_weights[1]*test_preds_rf + optimal_weights[2]*test_preds_cb
y_pred_final = (test_ensemble >= optimal_threshold).astype(int)

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

# %% Create submission
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("DONE!")
print("="*60)
print(f"Output saved to output.csv")
print(f"\nSummary:")
print(f"  - Ratio (subtask 1): {ratio:.6f}")
print(f"  - Optimal weights: LR={optimal_weights[0]:.3f}, RF={optimal_weights[1]:.3f}, CB={optimal_weights[2]:.3f}")
print(f"  - Optimal threshold: {optimal_threshold:.4f}")
print(f"  - OOF F1 Score: {optimal_f1:.4f}")
# %%
