# SPDX-License-Identifier: AGPL-3.0-or-later | Commercial license available
# © Concepts 1996–2026 Miroslav Šotek. All rights reserved.
# © Code 2020–2026 Miroslav Šotek. All rights reserved.
# ORCID: 0009-0009-3560-0851
# Contact: www.anulum.li | protoscience@anulum.li
# SC-NeuroCore — Training and evaluation loops for SNN models

"""Training and evaluation loops for SNN models."""

from __future__ import annotations

from typing import Callable, Tuple

import torch
from torch.utils.data import DataLoader

from .losses import spike_count_loss


def auto_device() -> torch.device:
    """Select best available device: CUDA > MPS > CPU."""
    if torch.cuda.is_available():
        return torch.device("cuda")
    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def train_epoch(
    model: torch.nn.Module,
    loader: DataLoader,
    optimizer: torch.optim.Optimizer,
    n_timesteps: int,
    loss_fn: Callable = spike_count_loss,
    device: str | torch.device = "cpu",
    max_grad_norm: float | None = None,
) -> Tuple[float, float]:
    """One training epoch. Returns (avg_loss, accuracy)."""
    model.train()
    total_loss = 0.0
    correct = 0
    total = 0

    for data, targets in loader:
        data, targets = data.to(device), targets.to(device)
        data = data.view(data.shape[0], -1)
        # (batch, features) -> (T, batch, features)
        data = data.unsqueeze(0).expand(n_timesteps, -1, -1)

        spike_counts, _ = model(data)
        loss = loss_fn(spike_counts, targets)

        optimizer.zero_grad()
        loss.backward()
        if max_grad_norm is not None:
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
        optimizer.step()

        total_loss += loss.item() * targets.shape[0]
        correct += (spike_counts.argmax(dim=1) == targets).sum().item()
        total += targets.shape[0]

    return total_loss / total, correct / total


@torch.no_grad()
def evaluate(
    model: torch.nn.Module,
    loader: DataLoader,
    n_timesteps: int,
    loss_fn: Callable = spike_count_loss,
    device: str = "cpu",
) -> Tuple[float, float]:
    """Evaluate model. Returns (avg_loss, accuracy)."""
    model.eval()
    total_loss = 0.0
    correct = 0
    total = 0

    for data, targets in loader:
        data, targets = data.to(device), targets.to(device)
        data = data.view(data.shape[0], -1)
        data = data.unsqueeze(0).expand(n_timesteps, -1, -1)

        spike_counts, _ = model(data)
        loss = loss_fn(spike_counts, targets)

        total_loss += loss.item() * targets.shape[0]
        correct += (spike_counts.argmax(dim=1) == targets).sum().item()
        total += targets.shape[0]

    return total_loss / total, correct / total
