"""
bishop.neural_networks – PRML Chapter 5
========================================
Uses PyTorch for all heavy lifting: autograd, optimisers, layers.
Provides a high-level ``MLP`` builder that mirrors the PRML notation
(single hidden layer, configurable activation, MSE or cross-entropy loss).

    from ml_ai_library.bishop.neural_networks import MLP, train, evaluate

Example
-------
>>> import torch
>>> model = MLP(in_features=2, hidden=[64, 64], out_features=1,
...             activation="tanh", task="regression")
>>> train(model, X_tensor, y_tensor, epochs=200, lr=1e-3)
>>> preds = evaluate(model, X_test_tensor)
"""

from __future__ import annotations
from typing import List, Optional
import torch
import torch.nn as nn
from torch.utils.data import TensorDataset, DataLoader


# ---------------------------------------------------------------------------
# MLP builder (PRML §5.1)
# ---------------------------------------------------------------------------

_ACTIVATIONS = {
    "tanh":    nn.Tanh,
    "relu":    nn.ReLU,
    "sigmoid": nn.Sigmoid,
    "gelu":    nn.GELU,
    "linear":  nn.Identity,
}


class MLP(nn.Module):
    """
    Feed-forward multilayer perceptron (PRML §5.1).

    Parameters
    ----------
    in_features : int
        Input dimensionality D.
    hidden : list[int]
        Number of units in each hidden layer.
    out_features : int
        Output dimensionality K.
    activation : str
        Hidden layer activation: 'tanh' | 'relu' | 'sigmoid' | 'gelu'.
    task : str
        'regression'   → linear output, MSELoss
        'binary'       → sigmoid output, BCELoss
        'multiclass'   → softmax output (use CrossEntropyLoss externally)
    dropout : float
        Dropout probability (0 = disabled).
    """

    def __init__(
        self,
        in_features: int,
        hidden: List[int],
        out_features: int,
        activation: str = "tanh",
        task: str = "regression",
        dropout: float = 0.0,
    ) -> None:
        super().__init__()
        self.task = task
        act_cls = _ACTIVATIONS[activation]
        layers: List[nn.Module] = []
        prev = in_features
        for h in hidden:
            layers += [nn.Linear(prev, h), act_cls()]
            if dropout > 0:
                layers.append(nn.Dropout(dropout))
            prev = h
        layers.append(nn.Linear(prev, out_features))
        if task == "binary":
            layers.append(nn.Sigmoid())
        self.net = nn.Sequential(*layers)

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


# ---------------------------------------------------------------------------
# Convenience training loop
# ---------------------------------------------------------------------------

def train(
    model: nn.Module,
    X: torch.Tensor,
    y: torch.Tensor,
    epochs: int = 100,
    lr: float = 1e-3,
    batch_size: int = 32,
    weight_decay: float = 0.0,
    verbose: bool = False,
) -> List[float]:
    """
    Mini-batch SGD training loop using Adam optimiser.

    Returns list of per-epoch training losses.
    """
    task = getattr(model, "task", "regression")
    if task == "multiclass":
        criterion = nn.CrossEntropyLoss()
    elif task == "binary":
        criterion = nn.BCELoss()
    else:
        criterion = nn.MSELoss()

    optimiser = torch.optim.Adam(model.parameters(), lr=lr,
                                 weight_decay=weight_decay)
    loader = DataLoader(TensorDataset(X, y), batch_size=batch_size, shuffle=True)
    history: List[float] = []
    model.train()
    for epoch in range(epochs):
        total = 0.0
        for xb, yb in loader:
            optimiser.zero_grad()
            loss = criterion(model(xb), yb)
            loss.backward()
            optimiser.step()
            total += loss.item() * len(xb)
        avg = total / len(X)
        history.append(avg)
        if verbose and (epoch + 1) % max(1, epochs // 10) == 0:
            print(f"  epoch {epoch+1}/{epochs}  loss={avg:.6f}")
    return history


def evaluate(model: nn.Module, X: torch.Tensor) -> torch.Tensor:
    """Run model in eval mode and return predictions."""
    model.eval()
    with torch.no_grad():
        return model(X)
