"""

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 trainiqConfig

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: trainiqConfig,

) -> 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: trainiqConfig):

    """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()]

