from __future__ import annotations

from typing import Optional, Dict, Any, List
import math
import numpy as np
import pandas as pd
import tensorflow as tf
from tensorflow.keras import Sequential
from tensorflow.keras.layers import Dense, LSTM, Dropout, Bidirectional, BatchNormalization, Flatten
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.regularizers import l2
from tensorflow.keras.callbacks import EarlyStopping
from tensorflow.keras.losses import MeanAbsoluteError, MeanSquaredError
from sklearn.model_selection import TimeSeriesSplit
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
from keras_tuner import BayesianOptimization, HyperModel

from .layers import CustomAttentionLayer
from .utils import normalize_xy, create_rnn_dataset


def _set_global_seed(seed: int):
    np.random.seed(seed)
    tf.random.set_seed(seed)


def _resolve_loss(loss: str):
    """
    Map 'mae'/'mse' -> (loss_obj, metrics_list).
    We keep a complementary metric for visibility.
    """
    loss = (loss or "mae").lower()
    if loss == "mae":
        return MeanAbsoluteError(), ["mse"]
    if loss == "mse":
        return MeanSquaredError(), ["mae"]
    raise ValueError(f"loss must be 'mae' or 'mse', got: {loss!r}")


class _CVHyperModel(HyperModel):
    def __init__(self, lookback: int, n_features: int, n_outputs: int = 1, loss: str = "mae"):
        self.input_shape = (lookback, n_features)
        self.n_outputs = n_outputs
        self.loss_obj, self.metrics_list = _resolve_loss(loss)

    def build(self, hp):
        m = Sequential()
        # 3 BiLSTM blocks
        for i, key in enumerate(("1", "2", "3")):
            m.add(Bidirectional(LSTM(
                units=hp.Int(f"units_{key}", 32, 512, step=32),
                activation="tanh", return_sequences=True,
                kernel_regularizer=l2(hp.Float(f"l2_{key}", 1e-6, 1e-3, sampling="log")),
                input_shape=self.input_shape if i == 0 else None,
            )))
            m.add(BatchNormalization())
            m.add(Dropout(hp.Float(f"dropout_{key}", 0.1, 0.5, step=0.1)))

        m.add(CustomAttentionLayer(hp.Float("emphasis_factor", 1.0, 2.0, step=0.1)))
        m.add(Flatten())
        m.add(Dense(self.n_outputs))

        m.compile(
            loss=self.loss_obj,
            optimizer=Adam(hp.Float("learning_rate", 1e-4, 1e-2, sampling="log")),
            metrics=self.metrics_list,
        )
        return m


def TSCV_model(
    csv_path: str,
    lookback: int = 24,
    n_splits: int = 10,
    epochs: int = 300,
    batch_size: int = 128,
    validation_split: float = 0.2,
    patience: int = 5,
    max_trials: int = 10,
    loss: str = "mae",
    seed: int = 42,
    save_model_path: Optional[str] = None,
    return_model: bool = False,
) -> Dict[str, Any]:
    """
    Time-series cross-validation pipeline with hyperparameter tuning,
    followed by a final fit on all training data and evaluation on a holdout.

    Args mirror TT_model with added:
        n_splits: number of folds for TimeSeriesSplit.

    Returns:
        dict with:
          - cv_folds (list of per-fold evals)
          - best_hyperparams
          - holdout_metrics (mse, rmse, mae, r2) on test set
          - y_test_pred (inverse-scaled predictions)
          - (optional) model
    """
    _set_global_seed(seed)

    data = pd.read_csv(csv_path, header=0, parse_dates=[0])
    dt_cols = [c for c in ("Date", "Time (GMT)") if c in data.columns]
    core = data.drop(columns=dt_cols) if dt_cols else data
    X = core.iloc[:, :-1].to_numpy()
    y = core.iloc[:, -1].to_numpy()

    Xn, yn, x_scaler, y_scaler = normalize_xy(X, y)

    # Holdout 20% for final evaluation; CV within the 80% training portion
    train_size = int(len(Xn) * 0.8)
    X_train, X_test = Xn[:train_size], Xn[train_size:]
    y_train, y_test = yn[:train_size], yn[train_size:]

    X_train_rnn, y_train_rnn = create_rnn_dataset(X_train, y_train, lookback)
    X_test_rnn,  y_test_rnn  = create_rnn_dataset(X_test,  y_test,  lookback)

    n_features = Xn.shape[1]
    tuner = BayesianOptimization(
        _CVHyperModel(lookback=lookback, n_features=n_features, loss=loss),
        objective="val_loss",
        max_trials=max_trials,
        seed=seed,
        project_name="Trials",
    )
    es = EarlyStopping(monitor="val_loss", patience=patience, mode="min", restore_best_weights=True)

    # Initial tuning on the train_rnn set
    tuner.search(
        X_train_rnn, y_train_rnn,
        epochs=epochs,
        batch_size=batch_size,
        validation_split=0.3,
        callbacks=[es],
        verbose=0,
    )
    best_hps = tuner.get_best_hyperparameters(1)[0]

    # TimeSeries CV across the *pre-windowed* training sequence (non-RNN form),
    # building windowed folds inside the loop.
    tscv = TimeSeriesSplit(n_splits=n_splits)
    fold_scores: List[Dict[str, float]] = []
    for tr_idx, va_idx in tscv.split(X_train):
        X_tr, X_va = X_train[tr_idx], X_train[va_idx]
        y_tr, y_va = y_train[tr_idx], y_train[va_idx]
        X_tr_rnn, y_tr_rnn = create_rnn_dataset(X_tr, y_tr, lookback)
        X_va_rnn, y_va_rnn = create_rnn_dataset(X_va, y_va, lookback)

        m = tuner.hypermodel.build(best_hps)
        m.fit(
            X_tr_rnn, y_tr_rnn,
            epochs=epochs,
            batch_size=batch_size,
            validation_data=(X_va_rnn, y_va_rnn),
            callbacks=[es],
            verbose=0,
        )
        eval_vals = m.evaluate(X_va_rnn, y_va_rnn, verbose=0)
        fold_scores.append({"val_loss": float(eval_vals[0]), "val_metric": float(eval_vals[1])})

    # Final model on full training windows
    model = tuner.hypermodel.build(best_hps)
    model.fit(
        X_train_rnn, y_train_rnn,
        epochs=epochs,
        batch_size=batch_size,
        validation_split=validation_split,
        callbacks=[es],
        verbose=0,
    )

    # Holdout evaluation on test windows
    test_pred = model.predict(X_test_rnn, verbose=0)
    test_pred_orig = y_scaler.inverse_transform(test_pred)
    y_all_orig = y_scaler.inverse_transform(yn)

    start = train_size + lookback
    mse = mean_squared_error(y_all_orig[start:], test_pred_orig)
    mae = mean_absolute_error(y_all_orig[start:], test_pred_orig)
    rmse = math.sqrt(mse)
    r2 = r2_score(y_all_orig[start:], test_pred_orig)

    if save_model_path:
        model.save(save_model_path)

    result: Dict[str, Any] = {
        "cv_folds": fold_scores,
        "best_hyperparams": {
            "units_1": best_hps.get("units_1"),
            "units_2": best_hps.get("units_2"),
            "units_3": best_hps.get("units_3"),
            "dropout_1": best_hps.get("dropout_1"),
            "dropout_2": best_hps.get("dropout_2"),
            "dropout_3": best_hps.get("dropout_3"),
            "learning_rate": best_hps.get("learning_rate"),
            "emphasis_factor": best_hps.get("emphasis_factor"),
        },
        "holdout_metrics": {"mse": float(mse), "rmse": float(rmse), "mae": float(mae), "r2": float(r2)},
        "y_test_pred": test_pred_orig.flatten().tolist(),
    }
    if return_model:
        result["model"] = model
    return result
