import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, Optional, Sequence

import numpy as np
from sklearn.metrics import mean_absolute_error, r2_score
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

try:
    import torch
    import torch.nn as nn
    import torch.optim as optim
    _TORCH_IMPORT_ERROR = None
except Exception as exc:
    torch = None
    nn = None
    optim = None
    _TORCH_IMPORT_ERROR = exc


class TorchScaledPredictWrapper:
    def __init__(self, torch_model: Any, x_scaler: StandardScaler, y_scaler: StandardScaler, device: str):
        self.torch_model = torch_model
        self.x_scaler = x_scaler
        self.y_scaler = y_scaler
        self.device = str(device)

    def predict(self, x_in, **kwargs):
        x_in = np.asarray(x_in, dtype=np.float64)
        if x_in.ndim == 3 and x_in.shape[-1] == 1:
            x_in = x_in.squeeze(axis=-1)
        if x_in.ndim > 2:
            x_in = x_in.reshape(len(x_in), -1)
        x_scaled = self.x_scaler.transform(x_in).astype(np.float32)
        xt = torch.as_tensor(x_scaled, dtype=torch.float32, device=self.device)
        if xt.ndim == 2:
            pass
        else:
            xt = xt.reshape(len(xt), -1)

        self.torch_model.eval()
        with torch.no_grad():
            if hasattr(self.torch_model, "expects_channel") and self.torch_model.expects_channel:
                # Conv1d on this runtime can hit intermittent cuDNN mapping errors;
                # keep GPU execution but bypass cuDNN for CNN forward passes.
                with torch.backends.cudnn.flags(enabled=False):
                    y_scaled = self.torch_model(xt.unsqueeze(1)).detach().cpu().numpy().reshape(-1, 1)
            else:
                y_scaled = self.torch_model(xt).detach().cpu().numpy().reshape(-1, 1)

        y_raw = self.y_scaler.inverse_transform(y_scaled)
        return y_raw.reshape(-1, 1)


class TorchMLPRegressor(nn.Module):
    expects_channel = False

    def __init__(self, input_dims: int, hidden_units: Sequence[int], dropout: float):
        super().__init__()
        layers = []
        prev = int(input_dims)
        for w in hidden_units:
            layers.append(nn.Linear(prev, int(w)))
            layers.append(nn.ELU())
            if dropout > 1e-8:
                layers.append(nn.Dropout(float(dropout)))
            prev = int(w)
        layers.append(nn.Linear(prev, 1))
        self.net = nn.Sequential(*layers)

    def forward(self, x):
        return self.net(x)


class TorchCNNRegressor(nn.Module):
    expects_channel = True

    def __init__(self, input_dims: int, dropout: float):
        super().__init__()
        flat_dim = 16 * max(1, int(input_dims) - 2)
        self.net = nn.Sequential(
            nn.Conv1d(1, 128, kernel_size=3, padding=1),
            nn.ELU(),
            nn.MaxPool1d(kernel_size=2, stride=1),
            nn.Dropout(float(dropout)),
            nn.Conv1d(128, 64, kernel_size=3, padding=1),
            nn.ELU(),
            nn.MaxPool1d(kernel_size=2, stride=1),
            nn.Dropout(float(dropout)),
            nn.Conv1d(64, 32, kernel_size=3, padding=1),
            nn.ELU(),
            nn.Conv1d(32, 16, kernel_size=3, padding=1),
            nn.ELU(),
            nn.Flatten(),
            nn.Linear(flat_dim, 64),
            nn.ELU(),
            nn.Linear(64, 1),
        )

    def forward(self, x):
        return self.net(x)


@dataclass
class FlowControlSurrogateV6:
    input_dims: int
    learning_rate: float = 1e-3
    batch_size: int = 64
    epochs: int = 500
    test_size: float = 0.2
    train_test_split_random_state: int = 42
    patience: int = 30
    check_point_path: Path = Path("surrogate_v6.pt")
    hidden_units: Sequence[int] = field(default_factory=lambda: (128, 64, 32))
    architecture: str = "mlp"
    dropout: float = 0.10
    weight_decay: float = 1e-6
    tf_device_id: Optional[int] = None
    require_gpu: bool = True

    def __post_init__(self):
        self.model: Optional[Any] = None
        self.x_scaler = StandardScaler()
        self.y_scaler = StandardScaler()
        self.last_fit_metrics: Dict[str, float] = {}

    @staticmethod
    def _forward_with_runtime_guard(model: Any, x: Any):
        if hasattr(model, "expects_channel") and model.expects_channel:
            with torch.backends.cudnn.flags(enabled=False):
                return model(x.unsqueeze(1))
        return model(x)

    @staticmethod
    def _ensure_torch_ready() -> None:
        if torch is None:
            raise ImportError(
                f"PyTorch is required for FlowControlSurrogateV6 but is not available: {_TORCH_IMPORT_ERROR}"
            )

    def _pick_device(self) -> str:
        self._ensure_torch_ready()
        if not torch.cuda.is_available():
            if self.require_gpu:
                raise RuntimeError("FlowControlSurrogateV6 requires GPU but torch.cuda is not available")
            return "cpu"

        if self.tf_device_id is None:
            return "cuda:0"
        idx = int(max(0, min(torch.cuda.device_count() - 1, int(self.tf_device_id))))
        return f"cuda:{idx}"

    def _build_model(self):
        arch = str(self.architecture).lower()
        if arch == "cnn":
            return TorchCNNRegressor(input_dims=self.input_dims, dropout=float(self.dropout))
        return TorchMLPRegressor(
            input_dims=self.input_dims,
            hidden_units=self.hidden_units,
            dropout=float(self.dropout),
        )

    def __call__(self, x, y, verbose: int = 0):
        self._ensure_torch_ready()
        x = np.asarray(x, dtype=np.float64)
        y = np.asarray(y, dtype=np.float64).reshape(-1)

        x_train, x_val, y_train, y_val = train_test_split(
            x,
            y,
            test_size=self.test_size,
            random_state=self.train_test_split_random_state,
            shuffle=True,
        )

        x_train_s = self.x_scaler.fit_transform(x_train).astype(np.float32)
        x_val_s = self.x_scaler.transform(x_val).astype(np.float32)
        y_train_s = self.y_scaler.fit_transform(y_train.reshape(-1, 1)).reshape(-1).astype(np.float32)

        device = self._pick_device()
        model = self._build_model().to(device)
        optimizer = optim.Adam(model.parameters(), lr=float(self.learning_rate), weight_decay=float(self.weight_decay))
        criterion = nn.MSELoss()

        xt = torch.as_tensor(x_train_s, dtype=torch.float32, device=device)
        yt = torch.as_tensor(y_train_s, dtype=torch.float32, device=device).reshape(-1, 1)

        xv = torch.as_tensor(x_val_s, dtype=torch.float32, device=device)

        batch = int(max(8, min(int(self.batch_size), len(x_train_s))))
        best_loss = np.inf
        best_state = None
        bad_epochs = 0
        epochs_ran = 0

        model.train()
        for ep in range(int(self.epochs)):
            perm = torch.randperm(xt.shape[0], device=device)
            epoch_loss = 0.0
            n_batches = 0

            for i in range(0, xt.shape[0], batch):
                idx = perm[i : i + batch]
                xb = xt[idx]
                yb = yt[idx]
                optimizer.zero_grad(set_to_none=True)
                pred = self._forward_with_runtime_guard(model, xb)
                loss = criterion(pred, yb)
                loss.backward()
                optimizer.step()
                epoch_loss += float(loss.item())
                n_batches += 1

            train_loss = epoch_loss / max(1, n_batches)
            epochs_ran = ep + 1

            if train_loss + 1e-10 < best_loss:
                best_loss = train_loss
                best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
                bad_epochs = 0
            else:
                bad_epochs += 1

            if bad_epochs >= int(self.patience):
                break

        if best_state is not None:
            model.load_state_dict(best_state)

        model.eval()
        with torch.no_grad():
            y_val_pred_s = self._forward_with_runtime_guard(model, xv).detach().cpu().numpy().reshape(-1, 1)

        y_val_pred = self.y_scaler.inverse_transform(y_val_pred_s).reshape(-1)

        arch = str(self.architecture).lower()
        gpu_idx = int(str(device).split(":")[1]) if str(device).startswith("cuda") else -1
        self.last_fit_metrics = {
            "val_r2": float(r2_score(y_val, y_val_pred)),
            "val_mae": float(mean_absolute_error(y_val, y_val_pred)),
            "val_count": int(len(y_val)),
            "best_val_loss": float(best_loss),
            "epochs_ran": int(epochs_ran),
            "architecture": arch,
            "device": f"GPU:{gpu_idx}" if gpu_idx >= 0 else "CPU",
            "gpu_count": int(torch.cuda.device_count() if torch.cuda.is_available() else 0),
            "selected_gpu": int(gpu_idx),
        }

        self.model = model
        return TorchScaledPredictWrapper(model, self.x_scaler, self.y_scaler, device=device)
