"""
automl_lib.inference
~~~~~~~~~~~~~~~~~~~~
Unified prediction interface for any trained model.
"""

from __future__ import annotations

import os
from typing import Any

import numpy as np
import torch
import torch.nn as nn

from .config import AutoMLConfig
from .hardware import detect_device
from .utils import get_logger

logger = get_logger(__name__)


class Predictor:
    """Load a trained model and run inference on new data.

    Example
    -------
    >>> predictor = Predictor.from_checkpoint("automl_output/checkpoints/best_model.pt",
    ...                                       model, config)
    >>> predictions = predictor.predict(new_data)
    """

    def __init__(self, model: nn.Module, config: AutoMLConfig):
        self.config = config
        self.device = detect_device(config.device)
        self.model = model.to(self.device)
        self.model.eval()

    @classmethod
    def from_checkpoint(
        cls,
        checkpoint_path: str,
        model: nn.Module,
        config: AutoMLConfig,
    ) -> "Predictor":
        """Load model weights from a checkpoint file."""
        state = torch.load(checkpoint_path, map_location="cpu")
        model.load_state_dict(state)
        logger.info(f"Loaded checkpoint: {checkpoint_path}")
        return cls(model, config)

    @torch.no_grad()
    def predict(self, data: Any) -> np.ndarray:
        """Run inference on *data* and return predictions.

        *data* can be:
        - A ``numpy.ndarray`` or ``torch.Tensor`` (tabular / image)
        - A ``list[str]`` (text)
        - A ``dict`` with ``input_ids`` / ``attention_mask`` (tokenised text)
        """
        self.model.eval()

        # ── Text (list of strings) ────────────────────────────────────
        if isinstance(data, list) and all(isinstance(s, str) for s in data):
            return self._predict_text(data)

        # ── Dict batch (HuggingFace-style) ────────────────────────────
        if isinstance(data, dict):
            inputs = {k: v.to(self.device) for k, v in data.items()}
            out = self.model(**inputs)
            logits = out.logits if hasattr(out, "logits") else out
            return logits.argmax(dim=1).cpu().numpy()

        # ── Tensor / ndarray ──────────────────────────────────────────
        if isinstance(data, np.ndarray):
            data = torch.tensor(data, dtype=torch.float32)
        if isinstance(data, torch.Tensor):
            data = data.to(self.device)
            out = self.model(data)
            if out.dim() > 1 and out.size(1) > 1:
                return out.argmax(dim=1).cpu().numpy()
            return out.squeeze().cpu().numpy()

        raise TypeError(f"Unsupported input type: {type(data)}")

    def predict_proba(self, data: Any) -> np.ndarray:
        """Return class probabilities (softmax) for classification models."""
        self.model.eval()
        if isinstance(data, np.ndarray):
            data = torch.tensor(data, dtype=torch.float32)
        data = data.to(self.device)
        with torch.no_grad():
            out = self.model(data)
            probs = torch.softmax(out, dim=1)
        return probs.cpu().numpy()

    # ── Private ───────────────────────────────────────────────────────

    def _predict_text(self, texts: list[str]) -> np.ndarray:
        """Tokenize and predict for raw text strings."""
        try:
            from transformers import AutoTokenizer
        except ImportError:
            raise ImportError("transformers is required for text prediction")

        tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
        encoded = tokenizer(
            texts, padding=True, truncation=True,
            max_length=256, return_tensors="pt",
        )
        inputs = {k: v.to(self.device) for k, v in encoded.items()}
        out = self.model(**inputs)
        logits = out.logits if hasattr(out, "logits") else out
        return logits.argmax(dim=1).cpu().numpy()
