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 bike-sharing day data and feature engineer --- #
df = pd.read_csv("data/day.csv")
# parse date and drop unused
df["dteday"] = pd.to_datetime(df["dteday"])
# engineer date-based features
df["year"]       = df["dteday"].dt.year
df["month"]      = df["dteday"].dt.month
df["weekday"]    = df["dteday"].dt.weekday
df["is_weekend"] = df["weekday"].isin([5,6]).astype(int)
# interaction feature
if "temp" in df.columns and "hum" in df.columns:
    df["temp_hum"] = df["temp"] * df["hum"]
# drop unneeded columns
drop_cols = ["instant","dteday","casual","registered"]
for c in drop_cols:
    if c in df.columns:
        df.drop(columns=c, inplace=True)

# target
y = df["cnt"].values
# we'll attach phi later, drop cnt for X
X_base = df.drop(columns=["cnt"])  

# --- 2) φ-method files for day.csv --- #
phi_methods = {
    "default": "CSV_output/day_phi_default.csv",
    "gmm":     "CSV_output/day_phi_gmm.csv",
    "kde":     "CSV_output/day_phi_kde.csv",
    "hist":    "CSV_output/day_phi_hist.csv",
    "spec":    "CSV_output/day_phi_spectral.csv"
}

# --- 3) Prepare frozen category lists --- #
# treat low-cardinality numeric codes as categories
cat_cols = [col for col in X_base.columns if X_base[col].nunique() < 12 and X_base[col].dtype != float]
num_cols = [col for col in X_base.columns if col not in cat_cols]
# cast cats to str to ensure encoder consistency
X_base[cat_cols] = X_base[cat_cols].astype(str)
# build categories
categories = [ X_base[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(X_base)
input_dim = pre_global.transform(X_base).shape[1]

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

# --- 6) Metrics with φ 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 = []

for phi_name, phi_file in phi_methods.items():
    print(f"\nProcessing φ method: {phi_name}")
    # attach phi
    d = X_base.copy()
    d['phi'] = pd.read_csv(phi_file)['phi']
    phi = d['phi'].values
    X = d.drop(columns=['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 = {}

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

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

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

        pd.concat(preds_list).to_csv(
            f"predictions/day_preds_{phi_name}_{model_name}.csv", index=False
        )

# write summary and best params
pd.DataFrame(results).to_csv(
    "Dataset wise nested cv results/nested_summary_day_phi.csv", index=False
)
pd.DataFrame(all_best_params).to_csv(
    "ParameterPerFold/day_best_params.csv", index=False
)

print("Done: summaries in summary_day_phi/, predictions in preds_day_phi/")
