"""LightGBM probability models for HK weather prediction targets.

Each model uses a 3-layer calibration stack:
  Layer 1: LightGBM binary classifier → raw log-odds
  Layer 2: Platt scaling (logistic regression on validation logits)
  Layer 3: Isotonic regression fallback (non-linear calibration)

Calibration parameters are saved/loaded with each model.
"""

import os
import json
import pickle
from pathlib import Path
from typing import Dict, Optional, Tuple, List

import numpy as np
import pandas as pd

try:
    import lightgbm as lgb
except ImportError:
    lgb = None

try:
    from sklearn.isotonic import IsotonicRegression
    from sklearn.linear_model import LogisticRegression
except ImportError:
    IsotonicRegression = None
    LogisticRegression = None

from config import DATA_DIR, PROJECT_ROOT


MODEL_DIR = Path(DATA_DIR) / "models"


TARGET_DEFINITIONS = {
    "rain_gt_0mm_24h": {
        "variable": "precipitation_sum", "threshold": 0.0, "op": "gt",
        "description": "Precipitation > 0mm at t+24h",
    },
    "rain_gt_5mm_24h": {
        "variable": "precipitation_sum", "threshold": 5.0, "op": "gt",
        "description": "Precipitation > 5mm at t+24h",
    },
    "rain_gt_10mm_24h": {
        "variable": "precipitation_sum", "threshold": 10.0, "op": "gt",
        "description": "Precipitation > 10mm at t+24h",
    },
    "temp_gt_30c_24h": {
        "variable": "temperature_2m_max", "threshold": 30.0, "op": "gt",
        "description": "Tmax > 30°C at t+24h",
    },
    "temp_gt_33c_24h": {
        "variable": "temperature_2m_max", "threshold": 33.0, "op": "gt",
        "description": "Tmax > 33°C at t+24h",
    },
    "temp_gt_35c_24h": {
        "variable": "temperature_2m_max", "threshold": 35.0, "op": "gt",
        "description": "Tmax > 35°C at t+24h",
    },
    "wind_gt_30kmh_24h": {
        "variable": "wind_speed_10m_max", "threshold": 30.0, "op": "gt",
        "description": "Wind gust > 30 km/h at t+24h",
    },
}

LGBM_PARAMS = {
    "objective": "binary",
    "metric": "binary_logloss",
    "boosting_type": "gbdt",
    "num_leaves": 15,
    "learning_rate": 0.03,
    "feature_fraction": 0.7,
    "bagging_fraction": 0.7,
    "bagging_freq": 5,
    "min_data_in_leaf": 50,
    "min_gain_to_split": 0.05,
    "lambda_l1": 0.5,
    "lambda_l2": 1.0,
    "max_depth": 4,
    "verbose": -1,
    "random_state": 42,
}


class ProbabilityCalibrator:
    """
    Post-hoc probability calibration using Platt scaling + isotonic regression.

    Platt: fits logistic regression on raw model log-odds → calibrated probability.
           Works well when raw scores follow a sigmoidal miscalibration pattern.
    Isotonic: non-parametric, fits step-wise monotonic function.
              Better for non-sigmoidal patterns but needs more data.

    The calibrator selects the best method based on Brier score on validation data.
    """

    def __init__(self, min_obs_isotonic: int = 100):
        self.min_obs_isotonic = min_obs_isotonic
        self.platt_model: Optional[LogisticRegression] = None
        self.iso_model: Optional[IsotonicRegression] = None
        self.method: Optional[str] = None      # "platt", "isotonic", or "none"
        self.fitted: bool = False

    def fit(self, raw_scores: np.ndarray, y_true: np.ndarray):
        """
        Fit calibration on validation data.

        Parameters
        ----------
        raw_scores : np.ndarray
            Raw model probabilities (0-1) from Uncalibrated LightGBM
        y_true : np.ndarray
            Binary ground truth labels
        """
        if len(raw_scores) < 10:
            self.method = "none"
            self.fitted = True
            return

        raw_scores = np.clip(raw_scores, 0.001, 0.999).reshape(-1, 1)
        y_true = np.asarray(y_true).ravel()

        from sklearn.metrics import brier_score_loss

        # Platt scaling (logistic regression on raw scores)
        self.platt_model = LogisticRegression(C=1.0, solver="lbfgs")
        self.platt_model.fit(raw_scores, y_true)
        platt_proba = self.platt_model.predict_proba(raw_scores)[:, 1]
        platt_brier = brier_score_loss(y_true, platt_proba)

        # Isotonic regression
        iso_brier = float("inf")
        if len(y_true) >= self.min_obs_isotonic and IsotonicRegression is not None:
            try:
                self.iso_model = IsotonicRegression(
                    y_min=0.001, y_max=0.999, out_of_bounds="clip"
                )
                self.iso_model.fit(raw_scores.ravel(), y_true)
                iso_proba = self.iso_model.predict(raw_scores.ravel())
                iso_brier = brier_score_loss(y_true, iso_proba)
            except Exception:
                self.iso_model = None

        # Select best method (prefer Platt for smooth calibration)
        # Isotonic can produce step functions with few unique points
        base_brier = brier_score_loss(y_true, raw_scores.ravel())
        scores = {"platt": platt_brier, "base": base_brier}

        # Only consider isotonic if it's significantly better and has enough unique outputs
        if self.iso_model is not None and iso_brier < platt_brier * 0.95:
            scores["isotonic"] = iso_brier
        else:
            scores["isotonic"] = float("inf")

        best = min(scores, key=scores.get)

        if best == "isotonic" and self.iso_model is not None:
            self.method = "isotonic"
        elif best == "platt" and self.platt_model is not None:
            self.method = "platt"
        else:
            self.method = "none"  # Raw scores are already best

        self.fitted = True
        print(f"    Calibration: {self.method} (platt_brier={platt_brier:.4f}, "
              f"iso_brier={iso_brier:.4f}, raw_brier={base_brier:.4f})")

    def calibrate(self, raw_scores: np.ndarray) -> np.ndarray:
        """Apply fitted calibration to raw scores (0-1)."""
        if not self.fitted or self.method == "none":
            raw = np.clip(raw_scores, 0.01, 0.99)
            return np.clip(raw, 0.01, 0.99)

        raw = np.atleast_1d(raw_scores)
        raw_clipped = np.clip(raw, 0.001, 0.999)

        if self.method == "platt" and self.platt_model is not None:
            cal = self.platt_model.predict_proba(raw_clipped.reshape(-1, 1))[:, 1]
        elif self.method == "isotonic" and self.iso_model is not None:
            cal = self.iso_model.predict(raw_clipped.ravel())
        else:
            cal = raw_clipped.ravel()

        # Gentle blending toward 0.5 for extreme probabilities
        # Only blend when raw is very extreme (>0.95 or <0.05)
        extremes = np.abs(raw_clipped.ravel() - 0.5)
        blend = np.clip((extremes - 0.4) / 0.1, 0, 0.3)
        cal_smoothed = cal * (1 - blend) + 0.5 * blend

        return np.clip(cal_smoothed, 0.01, 0.99)

    def save(self, path: str):
        """Save calibration params."""
        data = {
            "method": self.method,
            "platt": pickle.dumps(self.platt_model) if self.platt_model else None,
            "iso": pickle.dumps(self.iso_model) if self.iso_model else None,
        }
        with open(path, "wb") as f:
            pickle.dump(data, f)

    def load(self, path: str):
        """Load calibration params."""
        with open(path, "rb") as f:
            data = pickle.load(f)
        self.method = data.get("method", "none")
        if data.get("platt"):
            self.platt_model = pickle.loads(data["platt"])
        if data.get("iso"):
            self.iso_model = pickle.loads(data["iso"])
        self.fitted = True


class WeatherModel:
    """Probability model for a single weather target.

    Two model modes:
      - 'lgb': LightGBM gradient boosting (for real ERA5 data)
      - 'lr': Logistic regression (for synthetic/bootstrap data, prevents overfitting)
    """

    def __init__(self, target_name: str, mode: str = "lgb"):
        if target_name not in TARGET_DEFINITIONS:
            raise ValueError(f"Unknown target: {target_name}")
        self.target_name = target_name
        self.target_def = TARGET_DEFINITIONS[target_name]
        self.model: Optional[lgb.Booster] = None
        self.lr_model = None  # LogisticRegression for 'lr' mode
        self.mode = mode
        self.feature_importance: Dict[str, float] = {}
        self.calibrator = ProbabilityCalibrator()
        self._trained = False

    def train(
        self,
        X_train: np.ndarray,
        y_train: np.ndarray,
        X_val: Optional[np.ndarray] = None,
        y_val: Optional[np.ndarray] = None,
        params: Optional[Dict] = None,
        early_stopping_rounds: int = 50,
        verbose: bool = True,
    ):
        """Train model + calibrate."""
        if self.mode == "lr":
            self._train_lr(X_train, y_train, X_val, y_val)
        else:
            self._train_lgb(X_train, y_train, X_val, y_val, params, early_stopping_rounds, verbose)

        self._trained = True
        if X_val is not None and y_val is not None:
            self.calibrator.fit(self.predict_raw(X_val), y_val)
        elif X_train is not None and y_train is not None:
            self.calibrator.fit(self.predict_raw(X_train), y_train)

    def _train_lr(self, X_train, y_train, X_val, y_val):
        """Train logistic regression model."""
        if LogisticRegression is None:
            raise ImportError("scikit-learn not installed")
        from sklearn.preprocessing import StandardScaler
        self.scaler = StandardScaler()
        X_train_scaled = self.scaler.fit_transform(X_train)
        self.lr_model = LogisticRegression(
            C=0.1,                # Strong L2 regularization
            solver="lbfgs",
            max_iter=2000,
            class_weight="balanced",
        )
        self.lr_model.fit(X_train_scaled, y_train)
        self._trained = True

    def _train_lgb(self, X_train, y_train, X_val, y_val, params, early_stopping_rounds, verbose):
        """Train LightGBM model."""
        if lgb is None:
            raise ImportError("lightgbm not installed")
        train_params = {**LGBM_PARAMS, **(params or {})}
        dtrain = lgb.Dataset(X_train, label=y_train)
        if X_val is not None and y_val is not None:
            dval = lgb.Dataset(X_val, label=y_val, reference=dtrain)
            valid_sets, valid_names = [dtrain, dval], ["train", "valid"]
        else:
            valid_sets, valid_names = None, None
        self.model = lgb.train(
            train_params, dtrain, num_boost_round=500,
            valid_sets=valid_sets, valid_names=valid_names,
            callbacks=[lgb.early_stopping(early_stopping_rounds),
                       lgb.log_evaluation(period=50 if verbose else 0)]
            if X_val is not None else None,
        )
        self._compute_feature_importance()

    def predict_raw(self, X: np.ndarray) -> np.ndarray:
        """Raw probability (0-1) before calibration."""
        if not self._trained:
            raise RuntimeError("Model not trained")
        if self.mode == "lr" and self.lr_model is not None:
            X_scaled = self.scaler.transform(X)
            return self.lr_model.predict_proba(X_scaled)[:, 1]
        elif self.model is not None:
            return self.model.predict(X)
        else:
            return np.full(len(X), 0.5)

    def predict_proba(self, X: np.ndarray) -> np.ndarray:
        """Calibrated probability (0-100)."""
        raw = self.predict_raw(X)
        cal = self.calibrator.calibrate(raw)
        return cal * 100.0

    def predict(self, X: np.ndarray, threshold: float = 50.0) -> np.ndarray:
        return (self.predict_proba(X) >= threshold).astype(int)

    def evaluate(self, X: np.ndarray, y: np.ndarray) -> Dict[str, float]:
        proba = self.predict_proba(X) / 100.0
        pred = (proba >= 0.5).astype(int)
        from sklearn.metrics import accuracy_score, brier_score_loss, roc_auc_score, log_loss
        return {
            "accuracy": float(accuracy_score(y, pred)),
            "brier_score": float(brier_score_loss(y, proba)),
            "roc_auc": float(roc_auc_score(y, proba)) if len(np.unique(y)) > 1 else 0.5,
            "log_loss": float(log_loss(y, proba)),
            "n_samples": len(y),
            "p_yes_actual": float(y.mean() * 100),
            "p_yes_predicted": float(proba.mean() * 100),
            "calibration_method": self.calibrator.method,
        }

    def _compute_feature_importance(self):
        if self.model is None:
            return
        gain = self.model.feature_importance(importance_type="gain")
        names = self.model.feature_name()
        self.feature_importance = dict(sorted(zip(names, gain), key=lambda x: x[1], reverse=True))

    def top_features(self, n: int = 15) -> Dict[str, float]:
        items = sorted(self.feature_importance.items(), key=lambda x: x[1], reverse=True)
        return dict(items[:n])

    def save(self, path: Optional[str] = None):
        MODEL_DIR.mkdir(parents=True, exist_ok=True)
        p = path or (MODEL_DIR / f"{self.target_name}.lgb")
        if self.model:
            self.model.save_model(str(p))
        elif not os.path.exists(p):
            # Create marker for LR models
            with open(p, "w") as f:
                f.write("lr")

        meta = {
            "target_name": self.target_name,
            "target_definition": self.target_def,
            "mode": self.mode,
            "feature_importance": self.feature_importance,
            "trained": self._trained,
            "calibration_method": self.calibrator.method,
        }
        with open(str(p).replace(".lgb", "_meta.json"), "w") as f:
            json.dump(meta, f, indent=2, default=str)
        self.calibrator.save(str(p).replace(".lgb", "_cal.pkl"))
        if self.lr_model is not None:
            import pickle
            with open(str(p).replace(".lgb", "_lr.pkl"), "wb") as f:
                pickle.dump({"model": self.lr_model, "scaler": self.scaler}, f)

    def load(self, path: Optional[str] = None):
        p = path or (MODEL_DIR / f"{self.target_name}.lgb")
        if not os.path.exists(p):
            raise FileNotFoundError(f"Model not found: {p}")
        meta_path = str(p).replace(".lgb", "_meta.json")
        if os.path.exists(meta_path):
            with open(meta_path) as f:
                meta = json.load(f)
                self.mode = meta.get("mode", "lgb")
                self.feature_importance = meta.get("feature_importance", {})

        if self.mode == "lr":
            import pickle
            lr_path = str(p).replace(".lgb", "_lr.pkl")
            if os.path.exists(lr_path):
                with open(lr_path, "rb") as f:
                    data = pickle.load(f)
                    self.lr_model = data["model"]
                    self.scaler = data["scaler"]
        else:
            if lgb is None:
                raise ImportError("lightgbm not installed")
            self.model = lgb.Booster(model_file=str(p))
        self._trained = True
        cal_path = str(p).replace(".lgb", "_cal.pkl")
        if os.path.exists(cal_path):
            self.calibrator.load(cal_path)

    @staticmethod
    def build_target(df: pd.DataFrame, variable: str, threshold: float, op: str = "gt") -> np.ndarray:
        values = df[variable].values
        if op == "gt":
            return (values > threshold).astype(int)
        elif op == "ge":
            return (values >= threshold).astype(int)
        elif op == "lt":
            return (values < threshold).astype(int)
        elif op == "le":
            return (values <= threshold).astype(int)
        else:
            raise ValueError(f"Unknown operator: {op}")


class ModelEnsemble:
    """Manage multiple WeatherModel instances."""

    def __init__(self):
        self.models: Dict[str, WeatherModel] = {}

    def load_all(self):
        MODEL_DIR.mkdir(parents=True, exist_ok=True)
        for target in TARGET_DEFINITIONS:
            model_path = MODEL_DIR / f"{target}.lgb"
            if model_path.exists():
                model = WeatherModel(target)
                model.load(str(model_path))
                self.models[target] = model
        return self.models

    def predict_all(self, X: np.ndarray) -> Dict[str, float]:
        if X.ndim == 1:
            X = X.reshape(1, -1)
        return {name: float(model.predict_proba(X)[0]) for name, model in self.models.items()}

    def predict(self, target: str, X: np.ndarray) -> float:
        if target not in self.models:
            raise KeyError(f"Model '{target}' not loaded.")
        return float(self.models[target].predict_proba(X)[0])

    def has(self, target: str) -> bool:
        return target in self.models

    @property
    def available_targets(self) -> List[str]:
        return list(self.models.keys())