"""

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 trainiqConfig

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("trainiq_output/checkpoints/best_model.pt",

    ...                                       model, config)

    >>> predictions = predictor.predict(new_data)
    """


    def __init__(self, model: nn.Module, config: trainiqConfig):

        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: trainiqConfig,

    ) -> "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()

