"""
Core model architecture for WastePredictor.

These classes are **identical** to the originals in local-module/train.py so that
existing .pkl files deserialise correctly via the CustomUnpickler in loader.py.

DO NOT rename classes or move them to sub-modules — pickle stores fully-qualified
class paths; the loader remaps `train.*` → `model_lib._core.architecture.*`.
"""
from __future__ import annotations

import os
import warnings

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.ensemble import (
    GradientBoostingRegressor,
    RandomForestRegressor,
    StackingRegressor,
)
from sklearn.linear_model import Ridge
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from sklearn.model_selection import KFold, train_test_split
from sklearn.preprocessing import RobustScaler

warnings.filterwarnings("ignore")

os.environ.setdefault("LOKY_MAX_CPU_COUNT", "1")

try:
    import xgboost as xgb        # type: ignore[import]
    HAS_XGB = True
except ImportError:
    HAS_XGB = False

try:
    import lightgbm as lgb       # type: ignore[import]
    HAS_LGB = True
except ImportError:
    HAS_LGB = False


# ── Feature Engineering ──────────────────────────────────────────────────────

class AdvancedFeatureEngineer:
    """Domain-driven feature engineering for waste prediction."""

    OUTPUT_COLS = [
        "Total_Waste_kg",
        "Solid_Waste_Limestone_kg",
        "Solid_Waste_Gypsum_kg",
        "Solid_Waste_Industrial_Salt_kg",
        "Liquid_Waste_Bittern_Liters",
        "Potential_Epsom_Salt_kg",
        "Potential_Potash_kg",
        "Potential_Magnesium_Oil_Liters",
    ]

    def __init__(self) -> None:
        self.feature_names: list[str] = []
        self.prod_mean: float | None = None

    def fit_transform(self, df: pd.DataFrame) -> pd.DataFrame:
        self.prod_mean = df["production_volume"].mean()
        return self._create_features(df)

    def transform(self, df: pd.DataFrame) -> pd.DataFrame:
        return self._create_features(df)

    def _create_features(self, df: pd.DataFrame) -> pd.DataFrame:
        f = pd.DataFrame()

        f["production_volume"] = df["production_volume"]
        f["log_production"] = np.log1p(df["production_volume"])
        f["production_sq"] = df["production_volume"] ** 2 / 1e10
        f["production_sqrt"] = np.sqrt(df["production_volume"])

        f["rain_sum"] = df["rain_sum"]
        f["log_rain"] = np.log1p(df["rain_sum"])
        f["temperature_mean"] = df["temperature_mean"]
        f["humidity_mean"] = df["humidity_mean"]
        f["wind_speed_mean"] = df["wind_speed_mean"]

        evap_power = (df["temperature_mean"] * df["wind_speed_mean"]) / (
            df["humidity_mean"] + 1
        )
        brine_density = 20 + 15 * (evap_power / (evap_power.max() + 1e-6))
        dilution_penalty = np.exp(-df["rain_sum"] / 100)
        f["brine_density_be"] = brine_density * dilution_penalty

        f["epsom_salt_precip"] = (f["brine_density_be"] > 30).astype(float)
        f["potash_precip"] = (f["brine_density_be"] > 32).astype(float)
        f["density_rain_dilution"] = f["brine_density_be"] * np.exp(-df["rain_sum"] / 100)

        if "Month" in df.columns:
            f["month_sin"] = np.sin(2 * np.pi * df["Month"] / 12)
            f["month_cos"] = np.cos(2 * np.pi * df["Month"] / 12)
            f["is_dry_season"] = (
                (df["Month"] >= 3) & (df["Month"] <= 8)
            ).astype(float)
            f["is_peak_prod"] = df["Month"].isin([3, 4, 7, 8]).astype(float)

        f["wet_index"] = (df["rain_sum"] / 400) * (df["humidity_mean"] / 100)
        f["dry_index"] = (
            (1 - df["rain_sum"] / 400).clip(0, 1)
            * (1 - df["humidity_mean"] / 100).clip(0, 1)
        )
        f["evap_index"] = evap_power

        f["prod_x_wet"] = df["production_volume"] * f["wet_index"] / 1e5
        f["prod_x_dry"] = df["production_volume"] * f["dry_index"] / 1e5
        f["prod_x_evap"] = df["production_volume"] * f["evap_index"] / 1e6
        f["prod_x_rain"] = df["production_volume"] * df["rain_sum"] / 1e7
        f["prod_x_humidity"] = df["production_volume"] * df["humidity_mean"] / 1e6
        f["prod_x_temp"] = df["production_volume"] * df["temperature_mean"] / 1e6

        f["expected_waste_base"] = df["production_volume"] * 2.1
        f["waste_multiplier"] = 1 + (f["wet_index"] - 0.5) * 0.3

        f["rain_x_humidity"] = df["rain_sum"] * df["humidity_mean"] / 1e4
        f["temp_x_wind"] = df["temperature_mean"] * df["wind_speed_mean"] / 100
        f["rain_sq"] = df["rain_sum"] ** 2 / 1e5

        if self.prod_mean:
            f["prod_relative"] = df["production_volume"] / self.prod_mean
        else:
            f["prod_relative"] = df["production_volume"] / df["production_volume"].mean()

        self.feature_names = list(f.columns)
        return f


# ── Gradient Boosting ─────────────────────────────────────────────────────────

class GradientBoostingWasteModel:
    """Per-target XGBoost/GBR ensemble."""

    def __init__(
        self,
        n_estimators: int = 500,
        learning_rate: float = 0.05,
        max_depth: int = 5,
    ) -> None:
        self.n_estimators = n_estimators
        self.learning_rate = learning_rate
        self.max_depth = max_depth
        self.models: dict = {}
        self.scalers: dict = {}
        self.feature_engineer = AdvancedFeatureEngineer()
        self.output_cols = AdvancedFeatureEngineer.OUTPUT_COLS

    def fit(self, df: pd.DataFrame, verbose: bool = True) -> "GradientBoostingWasteModel":
        X = self.feature_engineer.fit_transform(df)
        self.scaler_X = RobustScaler()
        X_scaled = self.scaler_X.fit_transform(X)

        for col in self.output_cols:
            if verbose:
                print(f"  Training GB for {col}…")
            y_log = np.log1p(df[col].values)
            model = (
                xgb.XGBRegressor(
                    n_estimators=self.n_estimators,
                    learning_rate=self.learning_rate,
                    max_depth=self.max_depth,
                    subsample=0.8,
                    colsample_bytree=0.8,
                    reg_alpha=0.1,
                    reg_lambda=1.0,
                    random_state=42,
                    n_jobs=1,
                    verbosity=0,
                )
                if HAS_XGB
                else GradientBoostingRegressor(
                    n_estimators=self.n_estimators,
                    learning_rate=self.learning_rate,
                    max_depth=self.max_depth,
                    subsample=0.8,
                    random_state=42,
                )
            )
            model.fit(X_scaled, y_log)
            self.models[col] = model
        return self

    def predict(self, df: pd.DataFrame) -> pd.DataFrame:
        X = self.feature_engineer.transform(df)
        if hasattr(self.scaler_X, 'feature_names_in_'):
            X = X[list(self.scaler_X.feature_names_in_)]
        X_scaled = self.scaler_X.transform(X)
        preds = {
            col: np.expm1(self.models[col].predict(X_scaled)).clip(0)
            for col in self.output_cols
        }
        return self._enforce_consistency(pd.DataFrame(preds))

    def _enforce_consistency(self, df: pd.DataFrame) -> pd.DataFrame:
        df = df.copy()
        solids = [
            "Solid_Waste_Limestone_kg",
            "Solid_Waste_Gypsum_kg",
            "Solid_Waste_Industrial_Salt_kg",
        ]
        solid_sum = df[solids].sum(axis=1)
        scale = np.where(solid_sum > 0, df["Total_Waste_kg"] / solid_sum, 1.0)
        for c in solids:
            df[c] = df[c] * scale
        return df

    def evaluate(self, df: pd.DataFrame) -> dict:
        y_true = df[self.output_cols].values
        y_pred = self.predict(df).values
        r2_per = r2_score(y_true, y_pred, multioutput="raw_values")
        return {
            "r2": r2_score(y_true, y_pred),
            "mae": mean_absolute_error(y_true, y_pred),
            "rmse": float(np.sqrt(mean_squared_error(y_true, y_pred))),
            "r2_per_target": dict(zip(self.output_cols, r2_per)),
        }


# ── Stacked Ensemble ──────────────────────────────────────────────────────────

class StackedEnsembleModel:
    """Level-0: XGB + LGB + RF + GBR. Level-1: Ridge meta-learner."""

    def __init__(self) -> None:
        self.feature_engineer = AdvancedFeatureEngineer()
        self.output_cols = AdvancedFeatureEngineer.OUTPUT_COLS
        self.models: dict = {}
        self.scaler_X: RobustScaler | None = None

    def _base_models(self) -> list:
        models: list = []
        if HAS_XGB:
            models.append(
                (
                    "xgb",
                    xgb.XGBRegressor(
                        n_estimators=300,
                        learning_rate=0.05,
                        max_depth=6,
                        subsample=0.8,
                        colsample_bytree=0.8,
                        reg_alpha=0.1,
                        random_state=42,
                        verbosity=0,
                    ),
                )
            )
        if HAS_LGB:
            models.append(
                (
                    "lgb",
                    lgb.LGBMRegressor(
                        n_estimators=300,
                        learning_rate=0.05,
                        max_depth=6,
                        subsample=0.8,
                        colsample_bytree=0.8,
                        reg_alpha=0.1,
                        random_state=43,
                        verbose=-1,
                    ),
                )
            )
        models.append(
            (
                "rf",
                RandomForestRegressor(
                    n_estimators=200,
                    max_depth=10,
                    min_samples_split=5,
                    random_state=44,
                    n_jobs=1,
                ),
            )
        )
        models.append(
            (
                "gbr",
                GradientBoostingRegressor(
                    n_estimators=200,
                    learning_rate=0.05,
                    max_depth=5,
                    subsample=0.8,
                    random_state=45,
                ),
            )
        )
        return models

    def fit(self, df: pd.DataFrame, verbose: bool = True) -> "StackedEnsembleModel":
        X = self.feature_engineer.fit_transform(df)
        self.scaler_X = RobustScaler()
        X_scaled = self.scaler_X.fit_transform(X)

        for col in self.output_cols:
            if verbose:
                print(f"  Training stacked ensemble for {col}…")
            y = np.log1p(df[col].values)
            stacked = StackingRegressor(
                estimators=self._base_models(),
                final_estimator=Ridge(alpha=1.0),
                cv=5,
                n_jobs=1,
                passthrough=True,
            )
            stacked.fit(X_scaled, y)
            self.models[col] = stacked
        return self

    def predict(self, df: pd.DataFrame) -> pd.DataFrame:
        X = self.feature_engineer.transform(df)
        if hasattr(self.scaler_X, 'feature_names_in_'):
            X = X[list(self.scaler_X.feature_names_in_)]
        X_scaled = self.scaler_X.transform(X)  # type: ignore[union-attr]
        preds = {
            col: np.expm1(self.models[col].predict(X_scaled)).clip(0)
            for col in self.output_cols
        }
        return self._enforce_consistency(pd.DataFrame(preds))

    def _enforce_consistency(self, df: pd.DataFrame) -> pd.DataFrame:
        df = df.copy()
        solids = [
            "Solid_Waste_Limestone_kg",
            "Solid_Waste_Gypsum_kg",
            "Solid_Waste_Industrial_Salt_kg",
        ]
        solid_sum = df[solids].sum(axis=1)
        scale = np.where(solid_sum > 0, df["Total_Waste_kg"] / solid_sum, 1.0)
        for c in solids:
            df[c] = df[c] * scale
        return df

    def evaluate(self, df: pd.DataFrame) -> dict:
        y_true = df[self.output_cols].values
        y_pred = self.predict(df).values
        r2_per = r2_score(y_true, y_pred, multioutput="raw_values")
        return {
            "r2": r2_score(y_true, y_pred),
            "mae": mean_absolute_error(y_true, y_pred),
            "rmse": float(np.sqrt(mean_squared_error(y_true, y_pred))),
            "r2_per_target": dict(zip(self.output_cols, r2_per)),
        }


# ── Neural Network ────────────────────────────────────────────────────────────

class DeepNeuralNetworkModel(nn.Module):
    """Deep NN with skip connections and GELU activations."""

    def __init__(
        self,
        input_dim: int,
        output_dim: int = 8,
        hidden_dims: list[int] | None = None,
    ) -> None:
        super().__init__()
        if hidden_dims is None:
            hidden_dims = [256, 512, 256, 128]

        self.input_bn = nn.BatchNorm1d(input_dim)

        layers: list[nn.Module] = []
        prev_dim = input_dim
        for h_dim in hidden_dims:
            layers += [
                nn.Linear(prev_dim, h_dim),
                nn.BatchNorm1d(h_dim),
                nn.GELU(),
                nn.Dropout(0.2),
            ]
            prev_dim = h_dim

        self.hidden = nn.Sequential(*layers)
        self.skip = nn.Linear(input_dim, 64)
        self.output = nn.Linear(hidden_dims[-1] + 64, output_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x_norm = self.input_bn(x)
        hidden = self.hidden(x_norm)
        skip = torch.relu(self.skip(x_norm))
        return self.output(torch.cat([hidden, skip], dim=1))


class NeuralNetworkTrainer:
    """Train DeepNeuralNetworkModel with AdamW + cosine annealing + early stopping."""

    def __init__(
        self,
        epochs: int = 1000,
        lr: float = 0.001,
        patience: int = 50,
    ) -> None:
        self.epochs = epochs
        self.lr = lr
        self.patience = patience
        self.feature_engineer = AdvancedFeatureEngineer()
        self.output_cols = AdvancedFeatureEngineer.OUTPUT_COLS
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    def fit(self, df: pd.DataFrame, verbose: bool = True) -> "NeuralNetworkTrainer":
        X = self.feature_engineer.fit_transform(df)
        y = df[self.output_cols].values

        self.scaler_X = RobustScaler()
        self.scaler_y = RobustScaler()
        X_scaled = self.scaler_X.fit_transform(X)
        y_scaled = self.scaler_y.fit_transform(np.log1p(y))

        X_tr, X_val, y_tr, y_val = train_test_split(
            X_scaled, y_scaled, test_size=0.15, random_state=42
        )
        to_t = lambda a: torch.tensor(a, dtype=torch.float32).to(self.device)
        X_tr_t, y_tr_t = to_t(X_tr), to_t(y_tr)
        X_val_t, y_val_t = to_t(X_val), to_t(y_val)

        self.model = DeepNeuralNetworkModel(X_tr.shape[1], len(self.output_cols)).to(
            self.device
        )
        optimizer = optim.AdamW(
            self.model.parameters(), lr=self.lr, weight_decay=1e-4
        )
        scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(
            optimizer, T_0=50, T_mult=2
        )
        criterion = nn.HuberLoss()

        best_val, best_state, patience_ctr = float("inf"), None, 0

        for epoch in range(self.epochs):
            self.model.train()
            optimizer.zero_grad()
            loss = criterion(self.model(X_tr_t), y_tr_t)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
            optimizer.step()
            scheduler.step()

            self.model.eval()
            with torch.no_grad():
                val_loss = criterion(self.model(X_val_t), y_val_t).item()

            if val_loss < best_val - 1e-5:
                best_val = val_loss
                best_state = {k: v.cpu().clone() for k, v in self.model.state_dict().items()}
                patience_ctr = 0
            else:
                patience_ctr += 1

            if patience_ctr >= self.patience:
                if verbose:
                    print(f"  Early stopping at epoch {epoch + 1}")
                break

            if verbose and (epoch + 1) % 100 == 0:
                print(f"  Epoch {epoch + 1}: train={loss.item():.6f} val={val_loss:.6f}")

        if best_state:
            self.model.load_state_dict(best_state)
        return self

    def predict(self, df: pd.DataFrame) -> pd.DataFrame:
        X = self.feature_engineer.transform(df)
        if hasattr(self.scaler_X, 'feature_names_in_'):
            X = X[list(self.scaler_X.feature_names_in_)]
        X_t = torch.tensor(
            self.scaler_X.transform(X), dtype=torch.float32
        ).to(self.device)
        self.model.eval()
        with torch.no_grad():
            pred_scaled = self.model(X_t).cpu().numpy()
        pred = np.expm1(self.scaler_y.inverse_transform(pred_scaled)).clip(0)
        return self._enforce_consistency(
            pd.DataFrame(pred, columns=self.output_cols)
        )

    def _enforce_consistency(self, df: pd.DataFrame) -> pd.DataFrame:
        df = df.copy()
        solids = [
            "Solid_Waste_Limestone_kg",
            "Solid_Waste_Gypsum_kg",
            "Solid_Waste_Industrial_Salt_kg",
        ]
        solid_sum = df[solids].sum(axis=1)
        scale = np.where(solid_sum > 0, df["Total_Waste_kg"] / solid_sum, 1.0)
        for c in solids:
            df[c] = df[c] * scale
        return df

    def evaluate(self, df: pd.DataFrame) -> dict:
        y_true = df[self.output_cols].values
        y_pred = self.predict(df).values
        r2_per = r2_score(y_true, y_pred, multioutput="raw_values")
        return {
            "r2": r2_score(y_true, y_pred),
            "mae": mean_absolute_error(y_true, y_pred),
            "rmse": float(np.sqrt(mean_squared_error(y_true, y_pred))),
            "r2_per_target": dict(zip(self.output_cols, r2_per)),
        }


# ── Production Ensemble ───────────────────────────────────────────────────────

class ProductionWastePredictor:
    """
    Weighted ensemble of GradientBoosting + StackedEnsemble + NeuralNetwork.
    This is the object serialised inside waste_predictor_v4.pkl.
    """

    OUTPUT_COLS = AdvancedFeatureEngineer.OUTPUT_COLS

    def __init__(self) -> None:
        self.models: dict = {}
        self.weights: dict[str, float] = {}
        self.feature_engineer = AdvancedFeatureEngineer()
        self.output_cols = self.OUTPUT_COLS
        self.test_metrics: dict = {}

    # ------------------------------------------------------------------
    # Training
    # ------------------------------------------------------------------

    def fit(self, df: pd.DataFrame, verbose: bool = True) -> "ProductionWastePredictor":
        train_df, test_df = train_test_split(df, test_size=0.15, random_state=42)
        results: dict = {}

        print("─── Gradient Boosting ───")
        gb = GradientBoostingWasteModel(n_estimators=500, learning_rate=0.03, max_depth=6)
        gb.fit(train_df, verbose=False)
        results["gradient_boosting"] = {"model": gb, "metrics": gb.evaluate(test_df)}

        if HAS_XGB or HAS_LGB:
            print("─── Stacked Ensemble ───")
            se = StackedEnsembleModel()
            se.fit(train_df, verbose=False)
            results["stacked"] = {"model": se, "metrics": se.evaluate(test_df)}

        print("─── Neural Network ───")
        nn_trainer = NeuralNetworkTrainer(epochs=800, lr=0.001, patience=60)
        nn_trainer.fit(train_df, verbose=False)
        results["neural_network"] = {"model": nn_trainer, "metrics": nn_trainer.evaluate(test_df)}

        total_r2 = sum(v["metrics"]["r2"] for v in results.values())
        for name, data in results.items():
            self.models[name] = data["model"]
            self.weights[name] = data["metrics"]["r2"] / total_r2

        pred = self._ensemble_predict(test_df)
        y_true = test_df[self.output_cols].values
        r2_per = r2_score(y_true, pred, multioutput="raw_values")
        self.test_metrics = {
            "r2": r2_score(y_true, pred),
            "mae": mean_absolute_error(y_true, pred),
            "rmse": float(np.sqrt(mean_squared_error(y_true, pred))),
            "r2_per_target": dict(zip(self.output_cols, r2_per)),
        }

        self._cross_validate(df)
        return self

    def _cross_validate(self, df: pd.DataFrame) -> None:
        kf = KFold(n_splits=5, shuffle=True, random_state=42)
        cv_scores = []
        for fold, (tr_idx, val_idx) in enumerate(kf.split(df)):
            gb = GradientBoostingWasteModel(n_estimators=300, learning_rate=0.05, max_depth=5)
            gb.fit(df.iloc[tr_idx], verbose=False)
            cv_scores.append(gb.evaluate(df.iloc[val_idx])["r2"])
            print(f"  Fold {fold + 1}: R²={cv_scores[-1]:.4f}")
        print(f"  CV mean R²: {np.mean(cv_scores):.4f} ± {np.std(cv_scores):.4f}")

    # ------------------------------------------------------------------
    # Inference
    # ------------------------------------------------------------------

    def _ensemble_predict(self, df: pd.DataFrame) -> np.ndarray:
        preds = [m.predict(df).values for m in self.models.values()]
        weights = np.array(list(self.weights.values()))
        return np.average(preds, axis=0, weights=weights)

    def predict(self, df: pd.DataFrame) -> pd.DataFrame:
        return pd.DataFrame(self._ensemble_predict(df), columns=self.output_cols)

    # ------------------------------------------------------------------
    # Serialisation
    # ------------------------------------------------------------------

    def save(self, path: str) -> None:
        import json
        import pickle

        with open(path, "wb") as f:
            pickle.dump(
                {
                    "models": self.models,
                    "weights": self.weights,
                    "feature_names": self.feature_engineer.feature_names,
                    "output_cols": self.output_cols,
                    "metrics": self.test_metrics,
                },
                f,
            )
        meta_path = path.replace(".pkl", "_metadata.json")
        with open(meta_path, "w") as f:
            json.dump(
                {
                    "output_columns": self.output_cols,
                    "feature_names": self.feature_engineer.feature_names,
                    "model_weights": self.weights,
                    "test_metrics": {k: float(v) for k, v in self.test_metrics.items() if isinstance(v, float)},
                    "per_target_r2": {k: float(v) for k, v in self.test_metrics.get("r2_per_target", {}).items()},
                },
                f,
                indent=2,
            )

    @classmethod
    def load(cls, path: str) -> "ProductionWastePredictor":
        """Load via the custom unpickler so any origin repo works."""
        from model_lib._core.loader import load_model_file

        return load_model_file(path)
