"""Inference path for the encoder demo.

Loads V3 (Option A: nocompress + mean-pool + attn-LoRA) and exposes a
single `run_inference(token_ids)` function returning multi-head
predictions in a UI-friendly shape.

Single model, no side-by-side — different from the parent's
pretrained-vs-random comparison. The encoder demo's narrative is about
the architecture itself, not the value of pretraining.

Runs on CPU (fp32) for HF Spaces basic tier and local Mac smoke tests.
On the production H100 (192.222.55.165 / demos.liquid.ai surface), the
caller can pass `dtype=torch.bfloat16, device="cuda"` for the bf16 fast
path.

Performance note: even on CPU, ~64-tx inference is dominated by the
backbone's 16-layer forward pass over 960 pseudo-tokens. Expect
~3-8 seconds per inference on a Mac M-series CPU. For an interactive
demo we cache the model in memory and accept the latency; the user
selects a customer, clicks Run, waits ~5 seconds, sees outputs. GPU
deployment cuts this to <100ms.
"""

from __future__ import annotations

from pathlib import Path
from typing import Any

import numpy as np
import torch
import torch.nn.functional as F
import yaml

from src.data.schema import SchemaConfig, load_schema
from src.training.trainer_utils import load_checkpoint

from encoder.src.model.transaction_fm import build_transaction_fm


class EncoderDemoModel:
    """Wraps the V3 stack with cached inference for the demo.

    Construction:
        model = EncoderDemoModel(
            model_config_path="encoder/configs/model_nocompress.yaml",
            schema_path="data/schema.yaml",
            checkpoint_path="encoder/experiments/.../step_004999.pt",
            dtype=torch.float32,
            device="cpu",
        )

    Inference:
        results = model.run_inference(token_ids)  # (64, 15) int64
        # results: dict[str, np.ndarray]
        #   fraud:         shape (1,)  — sigmoid probability in [0, 1]
        #   next_merchant: shape (10003,) — softmax over merchant_id vocab
        #   amount_range:  shape (16,)    — softmax over 16 amount buckets
        #   mcc:           shape (103,)   — softmax over MCC vocab
    """

    def __init__(
        self,
        model_config_path: str | Path,
        schema_path: str | Path,
        checkpoint_path: str | Path | None = None,
        dtype: torch.dtype = torch.float32,
        device: str = "cpu",
    ) -> None:
        self.device = torch.device(device)
        self.dtype = dtype

        with open(model_config_path) as f:
            mcfg = yaml.safe_load(f)
        self.schema: SchemaConfig = load_schema(schema_path)
        self.head_configs = mcfg["heads"]

        # device_map=None + manual .to(device) below. We deliberately don't
        # use device_map="auto" because the encoder/projector/heads are
        # constructed separately and would otherwise stay on CPU.
        self.model = build_transaction_fm(
            schema=self.schema,
            head_configs=self.head_configs,
            model_path=mcfg["backbone"]["hf_path"],
            architecture_cfg=mcfg.get("architecture"),
            encoder_cfg=mcfg.get("encoder"),
            projector_cfg=mcfg.get("projector"),
            lora_cfg=mcfg["backbone"]["lora"],
            dtype=dtype,
            device_map=None,
        ).to(self.device)
        self.model.eval()

        self.checkpoint_status = "no checkpoint (random-init)"
        if checkpoint_path is not None:
            ckpt_path = Path(checkpoint_path)
            if not ckpt_path.exists():
                raise FileNotFoundError(
                    f"Checkpoint not found: {ckpt_path}",
                )
            # Peek at the checkpoint to detect slim variant (LFM base
            # stripped). Slim checkpoints load with strict=False because
            # the freshly-loaded LFM base keys are missing from the
            # state_dict — but the runtime values are identical to what
            # was originally in the checkpoint, so the math is unchanged.
            ckpt_peek = torch.load(ckpt_path, map_location="cpu", weights_only=False)
            is_slim = ckpt_peek.get("model_state_dict_slim", False)
            del ckpt_peek  # release memory before load_checkpoint re-reads

            ckpt = load_checkpoint(ckpt_path, self.model, strict=not is_slim)
            step = ckpt.get("step", "?")
            slim_note = " (slim — LFM base from HF cache)" if is_slim else ""
            self.checkpoint_status = f"step {step}{slim_note}"

    @torch.no_grad()
    def run_inference(self, token_ids: np.ndarray) -> dict[str, np.ndarray]:
        """Run all heads on one customer sequence.

        Args:
            token_ids: (64, 15) int64 numpy array — one customer's
                64-transaction history.

        Returns:
            dict mapping head name to its prediction array. Fraud is a
            sigmoid probability; the rest are softmax-normalized class
            distributions.
        """
        tensor = torch.from_numpy(token_ids).unsqueeze(0).long().to(self.device)
        # tensor: (1, 64, 15)
        predictions = self.model(tensor)
        # predictions: dict[head_name, logits]
        #   fraud:         (1, 1)
        #   next_merchant: (1, 10003)
        #   amount_range:  (1, 16)
        #   mcc:           (1, 103)

        results: dict[str, np.ndarray] = {}
        for name, logits in predictions.items():
            if name == "fraud":
                prob = torch.sigmoid(logits).squeeze().cpu().numpy()
                # Cast to scalar then back to 1-element array for uniform downstream handling
                results[name] = np.array([float(prob)])
            else:
                probs = F.softmax(logits, dim=-1).squeeze(0).cpu().numpy()
                results[name] = probs

        return results

    def num_params(self) -> dict[str, int]:
        """Param breakdown for the architecture-display card in the UI."""
        total = sum(p.numel() for p in self.model.parameters())
        trainable = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
        return {"total": total, "trainable": trainable}