from __future__ import annotations

from dataclasses import dataclass

import numpy as np
import torch
from sklearn.ensemble import GradientBoostingRegressor, RandomForestRegressor
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from sklearn.multioutput import MultiOutputRegressor

try:
    from xgboost import XGBRegressor
except Exception:  # pragma: no cover - optional dependency
    XGBRegressor = None

from gfis.data import theoretical_methane_upper_bound


class MethaneLSTM(torch.nn.Module):
    def __init__(self, input_size: int, hidden_size: int = 32, num_layers: int = 1):
        super().__init__()
        self.lstm = torch.nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
        self.head = torch.nn.Sequential(
            torch.nn.Linear(hidden_size, 24),
            torch.nn.ReLU(),
            torch.nn.Linear(24, 1),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        output, _ = self.lstm(x)
        return self.head(output[:, -1, :]).squeeze(-1)


def make_tabular_regressor(random_state: int = 42):
    if XGBRegressor is not None:
        return XGBRegressor(
            n_estimators=220,
            max_depth=4,
            learning_rate=0.045,
            subsample=0.9,
            colsample_bytree=0.9,
            objective="reg:squarederror",
            random_state=random_state,
        )
    return GradientBoostingRegressor(random_state=random_state)


def make_soft_sensor(random_state: int = 42):
    return RandomForestRegressor(n_estimators=160, random_state=random_state, min_samples_leaf=3)


def train_lstm(
    train_windows: np.ndarray,
    train_targets: np.ndarray,
    input_size: int,
    epochs: int = 8,
    lr: float = 0.004,
) -> MethaneLSTM:
    torch.manual_seed(42)
    model = MethaneLSTM(input_size=input_size)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    loss_fn = torch.nn.MSELoss()

    x_tensor = torch.tensor(train_windows, dtype=torch.float32)
    y_tensor = torch.tensor(train_targets, dtype=torch.float32)

    model.train()
    for _ in range(epochs):
        optimizer.zero_grad()
        pred = model(x_tensor)
        loss = loss_fn(pred, y_tensor)
        loss.backward()
        optimizer.step()
    return model


def predict_lstm(model: MethaneLSTM, windows: np.ndarray) -> np.ndarray:
    model.eval()
    with torch.no_grad():
        return model(torch.tensor(windows, dtype=torch.float32)).numpy()


def ensemble_predictions(lstm_pred: np.ndarray, tabular_pred: np.ndarray, lstm_weight: float = 0.45) -> np.ndarray:
    n = min(len(lstm_pred), len(tabular_pred))
    return lstm_weight * lstm_pred[-n:] + (1.0 - lstm_weight) * tabular_pred[-n:]


def apply_physics_constraint(predictions: np.ndarray, volatile_solids: np.ndarray) -> tuple[np.ndarray, int]:
    upper = theoretical_methane_upper_bound(volatile_solids[-len(predictions) :])
    violations = int(np.sum(predictions > upper))
    corrected = np.minimum(predictions, upper)
    return corrected, violations


def regression_metrics(y_true: np.ndarray, y_pred: np.ndarray) -> dict[str, float]:
    n = min(len(y_true), len(y_pred))
    yt = y_true[-n:]
    yp = y_pred[-n:]
    return {
        "r2": float(r2_score(yt, yp)),
        "rmse": float(np.sqrt(mean_squared_error(yt, yp))),
        "mae": float(mean_absolute_error(yt, yp)),
    }


@dataclass
class PredictionBundle:
    lstm: np.ndarray
    tabular: np.ndarray
    ensemble: np.ndarray
    physics_guided: np.ndarray
    physics_violations: int