"""Model factory — build neural nets in one line."""

from __future__ import annotations

import torch
import torch.nn as nn


class _Classifier(nn.Module):
    def __init__(
        self,
        input_size: int,
        hidden: list[int],
        output_size: int,
        activation: str,
        dropout: float,
        batch_norm: bool,
        name: str,
    ):
        super().__init__()
        self.name = name
        self._activation_name = activation
        self._dropout = dropout
        self._batch_norm = batch_norm
        self._act_cls = {"relu": nn.ReLU, "tanh": nn.Tanh, "sigmoid": nn.Sigmoid, "gelu": nn.GELU}[activation]

        layers: list[nn.Module] = []
        prev = input_size
        for h in hidden:
            layers.append(nn.Linear(prev, h))
            if batch_norm:
                layers.append(nn.BatchNorm1d(h))
            layers.append(self._act_cls())
            if dropout > 0:
                layers.append(nn.Dropout(dropout))
            prev = h
        layers.append(nn.Linear(prev, output_size))
        self.net = nn.Sequential(*layers)
        self._task = "classification" if output_size > 1 else "regression"
        self._output_size = output_size
        self._hidden = list(hidden)
        self._input_size = input_size

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)

    def grow(self, new_size: int | None = None) -> int:
        """Insert a new hidden layer before the output layer.

        The new layer uses the same activation/dropout/batchnorm settings
        as the rest of the network. Existing weights are preserved.

        Args:
            new_size: Width of the new layer. Defaults to the last hidden
                      layer's size (or half the input size if no hidden layers).

        Returns:
            The size of the layer that was added.
        """
        # Figure out the size of the layer feeding into the output
        old_output = self.net[-1]
        assert isinstance(old_output, nn.Linear)
        in_features = old_output.in_features

        if new_size is None:
            new_size = in_features  # match the previous layer width

        # Build the new block: Linear → [BN] → Activation → [Dropout]
        new_layers: list[nn.Module] = [nn.Linear(in_features, new_size)]
        if self._batch_norm:
            new_layers.append(nn.BatchNorm1d(new_size))
        new_layers.append(self._act_cls())
        if self._dropout > 0:
            new_layers.append(nn.Dropout(self._dropout))

        # New output layer connects from the grown layer
        new_output = nn.Linear(new_size, self._output_size)

        # Rebuild the sequential: everything except old output + new block + new output
        all_layers = list(self.net.children())[:-1] + new_layers + [new_output]
        self.net = nn.Sequential(*all_layers)
        self._hidden.append(new_size)

        return new_size

    @property
    def num_layers(self) -> int:
        return len(self._hidden)

    def __repr__(self) -> str:
        params = sum(p.numel() for p in self.parameters())
        return f"{self.name} (Classifier, {len(self._hidden)} hidden layers, {params:,} params, task={self._task})"


class _VisionModel(nn.Module):
    def __init__(
        self,
        num_classes: int,
        backbone: str,
        pretrained: bool,
        freeze_backbone: bool,
        dropout: float,
        name: str,
    ):
        super().__init__()
        self.name = name
        import torchvision.models as models

        weights = "DEFAULT" if pretrained else None
        self.backbone = getattr(models, backbone)(weights=weights)

        if freeze_backbone:
            for param in self.backbone.parameters():
                param.requires_grad = False

        # Replace the final classification head
        if hasattr(self.backbone, "fc"):
            in_features = self.backbone.fc.in_features
            head = [nn.Dropout(dropout), nn.Linear(in_features, num_classes)] if dropout > 0 else [nn.Linear(in_features, num_classes)]
            self.backbone.fc = nn.Sequential(*head)
        elif hasattr(self.backbone, "classifier"):
            if isinstance(self.backbone.classifier, nn.Sequential):
                in_features = self.backbone.classifier[-1].in_features
                self.backbone.classifier[-1] = nn.Linear(in_features, num_classes)
            else:
                in_features = self.backbone.classifier.in_features
                self.backbone.classifier = nn.Linear(in_features, num_classes)

        self._task = "classification"
        self._output_size = num_classes

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.backbone(x)

    def __repr__(self) -> str:
        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
        total = sum(p.numel() for p in self.parameters())
        return f"{self.name} (Vision, {trainable:,}/{total:,} trainable params)"


class _TextModel(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        embed_dim: int,
        hidden_dim: int,
        num_classes: int,
        num_layers: int,
        dropout: float,
        name: str,
    ):
        super().__init__()
        self.name = name
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.lstm = nn.LSTM(
            embed_dim,
            hidden_dim,
            num_layers=num_layers,
            batch_first=True,
            bidirectional=True,
            dropout=dropout if num_layers > 1 else 0.0,
        )
        self.drop = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
        self.fc = nn.Linear(hidden_dim * 2, num_classes)
        self._task = "classification"
        self._output_size = num_classes

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        embedded = self.embedding(x)
        _, (hidden, _) = self.lstm(embedded)
        # Concatenate forward and backward hidden states
        combined = torch.cat((hidden[-2], hidden[-1]), dim=1)
        return self.fc(self.drop(combined))

    def __repr__(self) -> str:
        params = sum(p.numel() for p in self.parameters())
        return f"{self.name} (Text/LSTM, {params:,} params)"


class Model:
    """Factory for building neural networks in one line.

    Examples::

        model = Model.classifier(784, [128, 64], 10)
        model = Model.classifier(784, [256, 128, 64], 10, dropout=0.3, batch_norm=True, name="my_net")
        model = Model.vision(num_classes=10, freeze_backbone=True)
        model = Model.text(vocab_size=10000, num_classes=5, num_layers=2)
    """

    @staticmethod
    def classifier(
        input_size: int,
        hidden: list[int],
        output_size: int,
        *,
        activation: str = "relu",
        dropout: float = 0.0,
        batch_norm: bool = False,
        name: str = "neurologic-classifier",
    ) -> _Classifier:
        """Feedforward neural network for tabular data.

        Args:
            input_size: Number of input features.
            hidden: List of hidden layer sizes, e.g. [128, 64].
            output_size: Number of output classes (or 1 for regression).
            activation: One of "relu", "tanh", "sigmoid", "gelu".
            dropout: Dropout rate between layers (0 = off).
            batch_norm: Add batch normalization between layers.
            name: Friendly name for this model.
        """
        return _Classifier(input_size, hidden, output_size, activation, dropout, batch_norm, name)

    @staticmethod
    def vision(
        num_classes: int,
        *,
        backbone: str = "resnet18",
        pretrained: bool = True,
        freeze_backbone: bool = False,
        dropout: float = 0.0,
        name: str = "neurologic-vision",
    ) -> nn.Module:
        """Image classifier with a pretrained backbone.

        Args:
            num_classes: Number of output classes.
            backbone: Torchvision model name (resnet18, resnet50, mobilenet_v2, etc.).
            pretrained: Use pretrained weights.
            freeze_backbone: Freeze all backbone layers (only train the head).
            dropout: Dropout before the final layer.
            name: Friendly name for this model.
        """
        return _VisionModel(num_classes, backbone, pretrained, freeze_backbone, dropout, name)

    @staticmethod
    def text(
        vocab_size: int,
        num_classes: int,
        *,
        embed_dim: int = 128,
        hidden_dim: int = 256,
        num_layers: int = 1,
        dropout: float = 0.0,
        name: str = "neurologic-text",
    ) -> nn.Module:
        """Bidirectional LSTM text classifier.

        Args:
            vocab_size: Size of the vocabulary.
            num_classes: Number of output classes.
            embed_dim: Embedding dimension.
            hidden_dim: LSTM hidden dimension.
            num_layers: Number of stacked LSTM layers.
            dropout: Dropout rate.
            name: Friendly name for this model.
        """
        return _TextModel(vocab_size, embed_dim, hidden_dim, num_classes, num_layers, dropout, name)
