"""Abstract base model that all SoftASP models must implement."""

from __future__ import annotations

from abc import ABC, abstractmethod
from typing import Any

import torch
import torch.nn as nn


class BaseSoftASPModel(nn.Module, ABC):
    """Base class for all SoftASP models.

    Every model receives a batch dict and returns an output dict containing
    at minimum ``loss`` and ``logits`` (or ``predictions``).
    """

    @abstractmethod
    def forward(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
        """Forward pass.

        Args:
            batch: Dict from DataLoader with keys like ``images``, ``labels``,
                and optional masks.

        Returns:
            Dict with at least ``loss`` and ``logits`` keys.
        """

    def count_parameters(self) -> dict[str, int]:
        """Return total and trainable parameter counts."""
        total = sum(p.numel() for p in self.parameters())
        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
        return {"total_params": total, "trainable_params": trainable}

    def get_backbone(self) -> nn.Module | None:
        """Optional component controlled by the trainer's freeze/ramp schedule."""
        return None