"""
automl_lib.model_zoo
~~~~~~~~~~~~~~~~~~~~
Registry that maps (data_type, task) → candidate models and provides
factory functions to instantiate them.
"""

from __future__ import annotations

from typing import Any

import torch
import torch.nn as nn

from .config import AutoMLConfig
from .utils import get_logger

logger = get_logger(__name__)


# ═══════════════════════════════════════════════════════════════════════════
# Built-in PyTorch models
# ═══════════════════════════════════════════════════════════════════════════


class TabularNet(nn.Module):
    """Fully-connected network for tabular data."""

    def __init__(self, input_dim: int, num_classes: int,
                 layers: list[int] | None = None,
                 activation: str = "relu", dropout: float = 0.3):
        super().__init__()
        layers = layers or [256, 128]
        act_fn = _get_activation(activation)

        modules: list[nn.Module] = []
        in_dim = input_dim
        for hidden in layers:
            modules += [nn.Linear(in_dim, hidden), act_fn(), nn.Dropout(dropout)]
            in_dim = hidden
        modules.append(nn.Linear(in_dim, num_classes))
        self.net = nn.Sequential(*modules)

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


class LSTMModel(nn.Module):
    """LSTM for time-series forecasting."""

    def __init__(self, input_dim: int, hidden_dim: int = 128,
                 num_layers: int = 2, output_dim: int = 1,
                 dropout: float = 0.2):
        super().__init__()
        self.lstm = nn.LSTM(
            input_dim, hidden_dim, num_layers=num_layers,
            batch_first=True, dropout=dropout if num_layers > 1 else 0,
        )
        self.fc = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        # x: (batch, seq_len, features)
        out, _ = self.lstm(x)
        return self.fc(out[:, -1, :])


class TextCNN(nn.Module):
    """Simple 1-D CNN text classifier over token embeddings."""

    def __init__(self, vocab_size: int = 30522, embed_dim: int = 128,
                 num_classes: int = 2, num_filters: int = 128,
                 filter_sizes: tuple = (3, 4, 5), dropout: float = 0.3):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.convs = nn.ModuleList([
            nn.Conv1d(embed_dim, num_filters, k) for k in filter_sizes
        ])
        self.dropout = nn.Dropout(dropout)
        self.fc = nn.Linear(num_filters * len(filter_sizes), num_classes)

    def forward(self, input_ids, attention_mask=None, **kwargs):
        x = self.embedding(input_ids)            # (B, L, E)
        x = x.transpose(1, 2)                    # (B, E, L)
        pooled = [torch.relu(c(x)).max(dim=2).values for c in self.convs]
        x = torch.cat(pooled, dim=1)
        x = self.dropout(x)
        return self.fc(x)


# ═══════════════════════════════════════════════════════════════════════════
# Model Registry
# ═══════════════════════════════════════════════════════════════════════════

_REGISTRY: dict[tuple[str, str], list[str]] = {
    # (data_type, task) → list of model keys
    ("tabular", "classification"): ["tabular_net", "sklearn_rf", "sklearn_xgb"],
    ("tabular", "regression"):     ["tabular_net", "sklearn_rf", "sklearn_xgb"],
    ("image", "classification"):   ["resnet18", "resnet50", "efficientnet_b0"],
    ("text", "classification"):    ["text_cnn", "distilbert"],
    ("timeseries", "forecasting"): ["lstm", "transformer_ts"],
    ("timeseries", "regression"):  ["lstm"],
    ("audio", "classification"):   ["resnet18"],  # spectrogram approach
}


def list_models(data_type: str, task: str) -> list[str]:
    """Return available model keys for a given (data_type, task)."""
    return _REGISTRY.get((data_type, task), ["tabular_net"])


def build_model(
    model_name: str,
    *,
    input_shape: tuple,
    num_classes: int,
    config: AutoMLConfig,
) -> nn.Module:
    """Instantiate and return a model by key name."""
    logger.info(f"Building model: {model_name}")

    if model_name == "tabular_net":
        return TabularNet(
            input_dim=input_shape[0],
            num_classes=num_classes,
            layers=config.layers,
            activation=config.activations,
            dropout=config.dropout,
        )

    if model_name in ("resnet18", "resnet50", "efficientnet_b0"):
        return _build_torchvision(model_name, num_classes, config.pretrained)

    if model_name == "text_cnn":
        return TextCNN(num_classes=num_classes, dropout=config.dropout)

    if model_name == "distilbert":
        return _build_distilbert(num_classes)

    if model_name == "lstm":
        input_dim = input_shape[-1] if len(input_shape) > 1 else 1
        return LSTMModel(input_dim=input_dim, output_dim=num_classes)

    if model_name == "transformer_ts":
        input_dim = input_shape[-1] if len(input_shape) > 1 else 1
        return _build_transformer_ts(input_dim, num_classes)

    # Scikit-learn models (returned as-is; trainer handles differently)
    if model_name.startswith("sklearn_"):
        return _build_sklearn(model_name, num_classes, config)

    raise ValueError(f"Unknown model: {model_name}")


# ── Factory helpers ───────────────────────────────────────────────────────

def _build_torchvision(name: str, num_classes: int, pretrained: bool) -> nn.Module:
    import torchvision.models as tvm

    weights = "IMAGENET1K_V1" if pretrained else None
    if name == "resnet18":
        model = tvm.resnet18(weights=weights)
        model.fc = nn.Linear(model.fc.in_features, num_classes)
    elif name == "resnet50":
        model = tvm.resnet50(weights=weights)
        model.fc = nn.Linear(model.fc.in_features, num_classes)
    elif name == "efficientnet_b0":
        model = tvm.efficientnet_b0(weights=weights)
        model.classifier[1] = nn.Linear(
            model.classifier[1].in_features, num_classes
        )
    else:
        raise ValueError(f"Unsupported torchvision model: {name}")
    return model


def _build_distilbert(num_classes: int) -> nn.Module:
    from transformers import AutoModelForSequenceClassification
    model = AutoModelForSequenceClassification.from_pretrained(
        "distilbert-base-uncased", num_labels=num_classes
    )
    return model


def _build_transformer_ts(input_dim: int, output_dim: int) -> nn.Module:
    """Lightweight Transformer encoder for time-series."""

    class TransformerTS(nn.Module):
        def __init__(self):
            super().__init__()
            d_model = 64
            self.input_proj = nn.Linear(input_dim, d_model)
            encoder_layer = nn.TransformerEncoderLayer(
                d_model=d_model, nhead=4, dim_feedforward=128,
                dropout=0.1, batch_first=True,
            )
            self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=2)
            self.fc = nn.Linear(d_model, output_dim)

        def forward(self, x):
            x = self.input_proj(x)
            x = self.encoder(x)
            return self.fc(x[:, -1, :])

    return TransformerTS()


def _build_sklearn(name: str, num_classes: int, config: AutoMLConfig):
    """Return a scikit-learn estimator wrapped for our API."""

    class SklearnWrapper(nn.Module):
        """Thin wrapper so sklearn models pass through our pipeline."""

        def __init__(self, estimator):
            super().__init__()
            self.estimator = estimator
            self._is_sklearn = True

        def forward(self, x):
            raise NotImplementedError("Use estimator.predict() for sklearn")

        def fit(self, X, y):
            self.estimator.fit(X, y)

        def predict(self, X):
            return self.estimator.predict(X)

        def predict_proba(self, X):
            if hasattr(self.estimator, "predict_proba"):
                return self.estimator.predict_proba(X)
            return None

    if name == "sklearn_rf":
        from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
        task = config.task or "classification"
        est = (RandomForestClassifier(n_estimators=200, n_jobs=-1, random_state=config.seed)
               if task == "classification"
               else RandomForestRegressor(n_estimators=200, n_jobs=-1, random_state=config.seed))
    elif name == "sklearn_xgb":
        try:
            from xgboost import XGBClassifier, XGBRegressor
            task = config.task or "classification"
            est = (XGBClassifier(n_estimators=200, use_label_encoder=False,
                                 eval_metric="logloss", random_state=config.seed)
                   if task == "classification"
                   else XGBRegressor(n_estimators=200, random_state=config.seed))
        except ImportError:
            from sklearn.ensemble import GradientBoostingClassifier, GradientBoostingRegressor
            task = config.task or "classification"
            est = (GradientBoostingClassifier(n_estimators=200, random_state=config.seed)
                   if task == "classification"
                   else GradientBoostingRegressor(n_estimators=200, random_state=config.seed))
    else:
        raise ValueError(f"Unknown sklearn model: {name}")

    return SklearnWrapper(est)


# ── Utilities ─────────────────────────────────────────────────────────────

def _get_activation(name: str):
    """Return an activation class by name."""
    mapping = {
        "relu": nn.ReLU,
        "gelu": nn.GELU,
        "silu": nn.SiLU,
        "tanh": nn.Tanh,
        "leaky_relu": nn.LeakyReLU,
        "elu": nn.ELU,
    }
    if name.lower() not in mapping:
        raise ValueError(
            f"Unknown activation: {name}. Choose from {list(mapping)}"
        )
    return mapping[name.lower()]
