from __future__ import annotations

from dataclasses import dataclass
from typing import Optional

import numpy as np
from sklearn.base import BaseEstimator, RegressorMixin


class TorchCNN1DRegressor(BaseEstimator, RegressorMixin):
    """A small optional 1-D CNN regressor over the feature vector.

    This is intentionally lightweight and disabled by default because it requires
    PyTorch, which is often not installed on HPC systems. It treats the feature
    vector as a short 1-D sequence. This is not a spatial CNN over the lat-lon
    field; it is mainly a neural-network benchmark between MLP and full spatial
    CNN. Use only after checking that torch is available.
    """

    def __init__(
        self,
        n_epochs: int = 120,
        batch_size: int = 512,
        learning_rate: float = 1e-3,
        weight_decay: float = 1e-4,
        channels: int = 16,
        patience: int = 20,
        validation_fraction: float = 0.15,
        random_seed: int = 42,
        verbose: bool = False,
    ):
        self.n_epochs = n_epochs
        self.batch_size = batch_size
        self.learning_rate = learning_rate
        self.weight_decay = weight_decay
        self.channels = channels
        self.patience = patience
        self.validation_fraction = validation_fraction
        self.random_seed = random_seed
        self.verbose = verbose
        self.model_ = None

    def _build_model(self, n_features: int):
        import torch
        import torch.nn as nn

        class Net(nn.Module):
            def __init__(self, n_features: int, channels: int):
                super().__init__()
                self.net = nn.Sequential(
                    nn.Conv1d(1, channels, kernel_size=3, padding=1),
                    nn.ReLU(),
                    nn.Conv1d(channels, channels, kernel_size=3, padding=1),
                    nn.ReLU(),
                    nn.AdaptiveAvgPool1d(1),
                    nn.Flatten(),
                    nn.Linear(channels, 1),
                )
            def forward(self, x):
                return self.net(x).squeeze(-1)

        return Net(n_features, self.channels)

    def fit(self, X, y):
        try:
            import torch
            from torch.utils.data import DataLoader, TensorDataset
        except Exception as exc:  # pragma: no cover
            raise ImportError("cnn1d_small requires PyTorch. Install torch or remove cnn1d_small from methods.") from exc

        rng = np.random.default_rng(self.random_seed)
        X = np.asarray(X, dtype=np.float32)
        y = np.asarray(y, dtype=np.float32)
        n = X.shape[0]
        if n < 50:
            raise ValueError("Not enough samples for CNN training")
        idx = np.arange(n)
        rng.shuffle(idx)
        n_val = max(1, int(round(n * self.validation_fraction)))
        val_idx = idx[:n_val]
        tr_idx = idx[n_val:]

        torch.manual_seed(self.random_seed)
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model = self._build_model(X.shape[1]).to(device)
        opt = torch.optim.AdamW(model.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay)
        loss_fn = torch.nn.MSELoss()

        Xtr = torch.from_numpy(X[tr_idx]).unsqueeze(1)
        ytr = torch.from_numpy(y[tr_idx])
        Xva = torch.from_numpy(X[val_idx]).unsqueeze(1).to(device)
        yva = torch.from_numpy(y[val_idx]).to(device)
        loader = DataLoader(TensorDataset(Xtr, ytr), batch_size=self.batch_size, shuffle=True)

        best_loss = np.inf
        best_state = None
        bad = 0
        for epoch in range(int(self.n_epochs)):
            model.train()
            for xb, yb in loader:
                xb = xb.to(device)
                yb = yb.to(device)
                opt.zero_grad(set_to_none=True)
                loss = loss_fn(model(xb), yb)
                loss.backward()
                opt.step()
            model.eval()
            with torch.no_grad():
                val_loss = float(loss_fn(model(Xva), yva).detach().cpu().numpy())
            if val_loss < best_loss - 1e-7:
                best_loss = val_loss
                best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
                bad = 0
            else:
                bad += 1
            if self.verbose and epoch % 20 == 0:
                print(f"CNN epoch={epoch} val_loss={val_loss:.6g}")
            if bad >= int(self.patience):
                break
        if best_state is not None:
            model.load_state_dict(best_state)
        self.model_ = model.to("cpu")
        self.n_features_in_ = X.shape[1]
        return self

    def predict(self, X):
        if self.model_ is None:
            raise RuntimeError("Model is not fitted")
        import torch
        X = np.asarray(X, dtype=np.float32)
        self.model_.eval()
        out = []
        bs = int(self.batch_size)
        with torch.no_grad():
            for i in range(0, X.shape[0], bs):
                xb = torch.from_numpy(X[i:i+bs]).unsqueeze(1)
                yb = self.model_(xb).detach().cpu().numpy()
                out.append(yb)
        return np.concatenate(out).astype(np.float32)