"""Inference — load a saved model and run predictions without training code."""

from __future__ import annotations

from pathlib import Path

import torch

from neurologic.trainer import _best_device, _to_tensor


class Inference:
    """Load a saved NeuroLogic checkpoint and run predictions.

    Handles MPS/CUDA/CPU auto-detection, proper checkpoint unwrapping,
    and automatic task detection.

    Examples::

        engine = Inference("model.pt")
        prediction = engine.run([1.0, 2.0, 3.0])
        batch = engine.run([[1, 2, 3], [4, 5, 6]])
    """

    def __init__(self, model_path: str | Path, *, device: str | None = None):
        self.device = torch.device(device) if device else _best_device()
        checkpoint = torch.load(model_path, map_location=self.device, weights_only=False)

        # Unwrap from Trainer's save format
        if isinstance(checkpoint, dict) and "model" in checkpoint:
            self.model = checkpoint["model"]
        else:
            self.model = checkpoint

        self.model.to(self.device)
        self.model.eval()

        # Auto-detect task from saved model
        self._task = getattr(self.model, "_task", "regression")

    def run(self, data: object) -> list[float] | int:
        """Run inference on input data.

        Args:
            data: Input data — list, nested list, numpy array, pandas object, or tensor.

        Returns:
            For classification: predicted class index (single) or list of indices (batch).
            For regression: list of predicted values.
        """
        tensor = _to_tensor(data)
        if tensor.dim() == 1:
            tensor = tensor.unsqueeze(0)
        tensor = tensor.to(self.device, dtype=torch.float32)

        with torch.no_grad():
            raw_output = self.model(tensor)

            if self._task == "classification":
                preds = torch.argmax(raw_output, dim=1)
                result = preds.cpu().tolist()
                return result[0] if len(result) == 1 else result

            result = raw_output.cpu().view(-1).tolist()
            return result[0] if len(result) == 1 else result
