# =======================
# Nested Cross-Validation with Phi Methods
# =======================

# What the Code Will Do (Summarized):

# Outer 5-Fold (Model Evaluation):
# - Splits the data into 5 folds.
# - Each fold is held out once as a test set (20%).
# - The remaining 4 folds (80%) are used for training & hyperparameter tuning.

# Inner 5-Fold (Model Selection / Tuning):
# - On the 80% training data of each outer fold, a GridSearchCV is run.
# - This internally performs 5-fold CV to select the best hyperparameters.

# Then:
# - The best model (from inner 5-fold CV) is retrained on the full 80% outer training data.
# - It is then evaluated on the outer test fold (the 20% never seen before).
# - Evaluation metrics include: MAE, RMSE, φ-weighted MAE (MAE_phi), and φ-weighted RMSE (RMSE_phi).

# =======================
# ALGORITHM STRUCTURE
# =======================

# 1. Loop over each φ method (based on available housing_phi_<method>.csv files).
# 2. Load and engineer features from the base housing dataset.
# 3. For each outer fold (5 total):
#     - Split the data into training (80%) and test (20%).
#     - On the 80% training:
#         - Run inner 5-fold CV via GridSearchCV with specified hyperparameter grid.
#         - Identify best hyperparameters for this training fold.
#     - Retrain the best model on the full 80% training split.
#     - Evaluate it on the 20% outer test split.
#     - Store:
#         - MAE, RMSE, MAE_phi, RMSE_phi for this fold
#         - Test indices + predictions
#         - Best hyperparameters found
# 4. Repeat for all models (LinearRegression, DecisionTree, RandomForest, ANN).
# 5. Combine all 5-fold outer scores for each model and φ-method:
#     - Compute average and standard deviation across folds.
#     - Report most common/best hyperparameters.
# 6. Save all results to:
#     - `phi_nested_summary_results.csv` (consolidated table)
#     - One prediction file per model × φ combination with true/predicted values for each fold

# =======================
# Final Output Specification (Per φ Method)
# =======================

# For each φ method (e.g. default, gmm, kde, hist, spec):
# - Perform full 5x5 Nested CV (5 outer × 5 inner)
# - In each outer fold:
#     - Run inner CV to find best model config (via φ-weighted error)
#     - Retrain on full outer training data
#     - Evaluate on outer test set
#     - Save fold-specific results and test predictions
# - At the end, report:
#     - Mean + Std Dev for MAE, RMSE, MAE_phi, RMSE_phi
#     - List of best hyperparameters used per fold

import pandas as pd
import numpy as np
import os
from math import sqrt
import warnings

from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.linear_model import LinearRegression
from sklearn.tree import DecisionTreeRegressor
from sklearn.ensemble import RandomForestRegressor
from xgboost import XGBRegressor
from sklearn.metrics import mean_absolute_error, mean_squared_error
from sklearn.model_selection import KFold, GridSearchCV
from sklearn.pipeline import Pipeline

from scikeras.wrappers import KerasRegressor
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Dropout
from tensorflow.keras.optimizers import Adam

warnings.filterwarnings("ignore")
SEED = 42
np.random.seed(SEED)

# --- 1) Load insurance data and feature engineer --- #
df = pd.read_csv("data/insurance.csv")

def engineer_insurance(df):
    df = df.copy()
    # interaction and derived features
    df['age_bmi'] = df['age'] * df['bmi']
    df['bmi_over_30'] = (df['bmi'] > 30).astype(int)
    df['children_per_age'] = df['children'] / (df['age'] + 1)
    return df

df = engineer_insurance(df)

# --- 2) φ-method files (adjust paths to your uploaded phi CSVs) --- #
phi_methods = {
    "default": "CSV_output/insurance_phi_default.csv",
    "gmm":     "CSV_output/insurance_phi_gmm.csv",
    "kde":     "CSV_output/insurance_phi_kde.csv",
    "hist":    "CSV_output/insurance_phi_hist.csv",
    "spec":    "CSV_output/insurance_phi_spectral.csv"
}

# --- 3) Prepare frozen category lists --- #
base_X = df.drop(columns=["charges"])
cat_cols = base_X.select_dtypes(include="object").columns.tolist()
num_cols = base_X.select_dtypes(exclude="object").columns.tolist()
categories = [ base_X[col].unique().tolist() for col in cat_cols ]

def get_preprocessor(num_cols, cat_cols, categories):
    return ColumnTransformer([
        ("num", StandardScaler(), num_cols),
        ("cat", OneHotEncoder(
            categories=categories,
            handle_unknown="ignore",
            sparse_output=False
        ), cat_cols)
    ])

# --- 4) Compute input_dim once --- #
pre_global = get_preprocessor(num_cols, cat_cols, categories)
pre_global.fit(base_X)
input_dim = pre_global.transform(base_X).shape[1]

# --- 5) ANN builder function --- #
def build_ann(input_shape, lr=0.001, dropout=0.2, units=64):
    m = Sequential([
        Dense(units, input_shape=input_shape, activation="relu"),
        Dropout(dropout),
        Dense(units//2, activation="relu"),
        Dense(1)
    ])
    m.compile(optimizer=Adam(learning_rate=lr), loss="mse")
    return m

# --- 6) Metrics including φ-weighting --- #
def compute_metrics(y_true, y_pred, phi):
    mae     = mean_absolute_error(y_true, y_pred)
    rmse    = sqrt(mean_squared_error(y_true, y_pred))
    mae_phi = np.mean(phi * np.abs(y_true - y_pred))
    rmse_phi= np.sqrt(np.mean(phi * (y_true - y_pred)**2))
    return mae, rmse, mae_phi, rmse_phi

# --- 7) Nested CV over φ-methods and models --- #
results = []
all_best_params = []

y = df['charges'].values

for phi_name, phi_path in phi_methods.items():
    print(f"\nProcessing φ method: {phi_name}")
    d = df.copy()
    d['phi'] = pd.read_csv(phi_path)['phi']
    phi = d['phi'].values
    X = d.drop(columns=['charges','phi'])

    outer_kf = KFold(n_splits=5, shuffle=True, random_state=SEED)

    for model_name in ["LinearRegression","DecisionTree","RandomForest","XGBoost","ANN"]:
        print(f" Model: {model_name}")
        outer_metrics = []
        preds_list    = []

        for fold,(train_idx,test_idx) in enumerate(outer_kf.split(X), start=1):
            X_tr, X_te = X.iloc[train_idx], X.iloc[test_idx]
            y_tr, y_te = y[train_idx], y[test_idx]
            phi_tr,phi_te = phi[train_idx], phi[test_idx]

            pre = get_preprocessor(num_cols, cat_cols, categories)

            if model_name=="LinearRegression":
                estimator, grid = LinearRegression(), {}
            elif model_name=="DecisionTree":
                estimator = DecisionTreeRegressor(random_state=SEED)
                grid = {"regressor__max_depth":[4,6,8],
                        "regressor__min_samples_split":[2,5]}
            elif model_name=="RandomForest":
                estimator = RandomForestRegressor(random_state=SEED)
                grid = {"regressor__n_estimators":[50,100],
                        "regressor__max_depth":[6,8],
                        "regressor__min_samples_split":[2,5]}
            elif model_name == "XGBoost":
                estimator = XGBRegressor(
                    random_state=SEED,
                    objective='reg:squarederror',
                    verbosity=0
                )
                grid = {
                    "regressor__n_estimators": [100, 200],
                    "regressor__max_depth": [3, 6],
                    "regressor__learning_rate": [0.1, 0.01]
                }
            else:
                estimator = KerasRegressor(
                    model=build_ann,
                    verbose=0,
                    model__input_shape=(input_dim,)
                )
                grid = {"regressor__model__lr":[0.001,0.0005],
                        "regressor__model__dropout":[0.2,0.3],
                        "regressor__model__units":[32,64],
                        "regressor__epochs":[100],
                        "regressor__batch_size":[32]}

            pipe = Pipeline([("preprocessor",pre),("regressor",estimator)])

            if grid:
                search = GridSearchCV(pipe, grid, cv=5,
                                      scoring="neg_mean_squared_error",
                                      n_jobs=-1)
                search.fit(X_tr, y_tr)
                best = search.best_estimator_
                bp   = search.best_params_
            else:
                pipe.fit(X_tr, y_tr)
                best = pipe; bp = {}

            # collect best params
            all_best_params.append({**bp,
                                     "phi_method":phi_name,
                                     "model":model_name,
                                     "fold":fold})

            # predict & metrics
            y_pred = best.predict(X_te)
            m1,m2,m3,m4 = compute_metrics(y_te, y_pred, phi_te)
            outer_metrics.append([m1,m2,m3,m4])

            preds_list.append(pd.DataFrame({
                "phi_method":phi_name,
                "model":model_name,
                "fold":fold,
                "y_true":y_te,
                "y_pred":y_pred,
                "phi":phi_te
            }))

        # summary
        arr = np.array(outer_metrics)
        results.append({
            "phi_method":phi_name,
            "model":model_name,
            "MAE_mean":arr[:,0].mean(),
            "RMSE_mean":arr[:,1].mean(),
            "MAE_phi_mean":arr[:,2].mean(),
            "RMSE_phi_mean":arr[:,3].mean(),
            "MAE_std":arr[:,0].std(ddof=1),
            "RMSE_std":arr[:,1].std(ddof=1),
            "MAE_phi_std":arr[:,2].std(ddof=1),
            "RMSE_phi_std":arr[:,3].std(ddof=1)
        })

        # save fold preds
        pd.concat(preds_list).to_csv(
            f"predictions/Insurance_preds_{phi_name}_{model_name}.csv",
            index=False
        )


# final writes
pd.DataFrame(results).to_csv(
    "Dataset wise nested cv results/nested_summary_insurance_phi.csv",
    index=False
)
pd.DataFrame(all_best_params).to_csv(
    "ParameterPerFold/insurance_best_params.csv",
    index=False
)

print("Done. See summary_insurance_phi/, preds_insurance_phi/, params_insurance_phi/")
