from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Mapping, Optional, Tuple, Union, Literal

import numpy as np


ArrayLike = Union[np.ndarray]
PadMode = Literal["none", "zero", "one", "repeat"]


def _as_5d(frames: np.ndarray) -> np.ndarray:
    if frames.ndim == 4:
        return frames[None, ...]
    if frames.ndim == 5:
        return frames
    raise ValueError(f"Expected frames with rank 4 or 5, got shape={frames.shape}")


def _normalize01(x: np.ndarray) -> np.ndarray:
    x = x.astype(np.float32, copy=False)
    if x.max() > 1.5:
        x = x / 255.0
    return np.clip(x, 0.0, 1.0)


def _pad_last_frame(frames: np.ndarray, pad_mode: PadMode, target_len: int) -> np.ndarray:
    """
    frames: [B,T,H,W,C] float in 0..1
    Returns [B,target_len,H,W,C].
    """
    if pad_mode == "none":
        return frames
    b, t, h, w, c = frames.shape
    if t > target_len:
        return frames[:, :target_len]
    if t == target_len:
        return frames
    if t < target_len - 1:
        raise ValueError(f"pad_last_frame supports T=={target_len-1} or T=={target_len}, got T={t}")

    if pad_mode == "repeat":
        last = frames[:, -1:, :, :, :]
    elif pad_mode == "zero":
        last = np.zeros((b, 1, h, w, c), dtype=frames.dtype)
    elif pad_mode == "one":
        last = np.ones((b, 1, h, w, c), dtype=frames.dtype)
    else:
        raise ValueError(f"Unknown pad_mode={pad_mode}")
    return np.concatenate([frames, last], axis=1)


def _pick_signature(module) -> Tuple[str, Any]:
    sigs = getattr(module, "signatures", None)
    if not isinstance(sigs, Mapping) or not sigs:
        raise RuntimeError("SavedModel has no callable signatures.")
    if "serving_default" in sigs:
        return "serving_default", sigs["serving_default"]
    name = next(iter(sigs.keys()))
    return name, sigs[name]


def _signature_io(fn) -> Tuple[Dict[str, Any], Dict[str, Any]]:
    # structured_input_signature = (args, kwargs)
    args, kwargs = fn.structured_input_signature
    if args:
        raise RuntimeError("This SavedModel signature expects positional args; unsupported.")
    return dict(kwargs), dict(fn.structured_outputs)


def _infer_hw_from_inputs(input_specs: Mapping[str, Any]) -> Optional[Tuple[int, int]]:
    # Try to find a spec with shape [B,T,H,W,C] or [B,H,W,C]
    for spec in input_specs.values():
        shape = getattr(spec, "shape", None)
        if shape is None:
            continue
        dims = list(shape)
        if len(dims) == 5 and dims[2] is not None and dims[3] is not None:
            return int(dims[2]), int(dims[3])
        if len(dims) == 4 and dims[1] is not None and dims[2] is not None:
            return int(dims[1]), int(dims[2])
    return None


def _select_prediction_tensor(outputs: Mapping[str, Any]):
    # Heuristic: prefer key contains 'pred', else first rank-5 tensor, else first tensor-like output
    items = list(outputs.items())
    for k, v in items:
        if "pred" in k.lower():
            return k, v
    for k, v in items:
        shape = getattr(v, "shape", None)
        if shape is not None and len(shape) == 5:
            return k, v
    return items[0]


@dataclass
class Predictor:
    model_dir: Union[str, Path]
    signature_name: Optional[str] = None
    input_key: Optional[str] = None
    resize_hw: Optional[Tuple[int, int]] = None

    def __post_init__(self) -> None:
        import tensorflow as tf

        self._tf = tf

        # Prefer GPU if available; otherwise CPU is used automatically.
        try:
            gpus = tf.config.list_physical_devices("GPU")
            for gpu in gpus:
                # Avoid TF grabbing all GPU memory up front.
                tf.config.experimental.set_memory_growth(gpu, True)
            self.device: str = "GPU" if gpus else "CPU"
        except Exception:
            self.device = "CPU"

        self.model_dir = Path(self.model_dir)
        self._module = tf.saved_model.load(str(self.model_dir))
        name, fn = _pick_signature(self._module)
        if self.signature_name is None:
            self.signature_name = name
        self._fn = self._module.signatures[self.signature_name]

        in_specs, out_specs = _signature_io(self._fn)
        self._in_specs = in_specs
        self._out_specs = out_specs

        if self.input_key is None:
            # Prefer single input; otherwise pick first
            if len(in_specs) == 1:
                self.input_key = next(iter(in_specs.keys()))
            else:
                self.input_key = next(iter(in_specs.keys()))

        if self.resize_hw is None:
            self.resize_hw = _infer_hw_from_inputs(in_specs) or (128, 128)

        # Infer expected sequence length if fixed (e.g. 4)
        spec = self._in_specs.get(self.input_key) if self.input_key else None
        shape = getattr(spec, "shape", None)
        self.sequence_len: Optional[int] = None
        if shape is not None and len(shape) == 5 and shape[1] is not None:
            self.sequence_len = int(shape[1])

    def _prep(self, frames: np.ndarray, pad_last_frame: PadMode = "none") -> np.ndarray:
        frames = _as_5d(frames)
        frames = _normalize01(frames)

        # Resize with TF to match model expectation (keeps it black-box)
        b, t, h, w, c = frames.shape
        rh, rw = self.resize_hw or (h, w)
        if (h, w) != (rh, rw):
            tf = self._tf
            x = tf.convert_to_tensor(frames, dtype=tf.float32)
            x = tf.reshape(x, (b * t, h, w, c))
            x = tf.image.resize(x, (rh, rw), method="bilinear", antialias=True)
            x = tf.reshape(x, (b, t, rh, rw, c))
            frames = x.numpy()

        if self.sequence_len is not None:
            frames = _pad_last_frame(frames, pad_last_frame, target_len=self.sequence_len)
        return frames

    def predict_sequence(self, frames: np.ndarray, pad_last_frame: PadMode = "none") -> np.ndarray:
        tf = self._tf
        x = self._prep(frames, pad_last_frame=pad_last_frame)
        inp = {self.input_key: tf.convert_to_tensor(x, dtype=tf.float32)}
        outs = self._fn(**inp)
        pred = outs.get("predictions")
        if pred is None:
            _, pred = _select_prediction_tensor(outs)
        pred_np = pred.numpy()
        return np.clip(pred_np, 0.0, 1.0)

    def predict_outputs(self, frames: np.ndarray, pad_last_frame: PadMode = "none") -> Dict[str, np.ndarray]:
        tf = self._tf
        x = self._prep(frames, pad_last_frame=pad_last_frame)
        inp = {self.input_key: tf.convert_to_tensor(x, dtype=tf.float32)}
        outs = self._fn(**inp)
        out_np: Dict[str, np.ndarray] = {}
        for k, v in outs.items():
            try:
                out_np[k] = v.numpy()
            except Exception:
                pass
        if "predictions" in out_np:
            out_np["predictions"] = np.clip(out_np["predictions"], 0.0, 1.0)
        return out_np

    def predict_last_frame(self, frames: np.ndarray, pad_last_frame: PadMode = "none") -> np.ndarray:
        seq = self.predict_sequence(frames, pad_last_frame=pad_last_frame)
        if seq.ndim != 5:
            raise RuntimeError(f"Expected prediction rank 5, got shape={seq.shape}")
        return seq[:, -1]

    def predict_next_from_3(self, frames3: np.ndarray, pad_last_frame: PadMode = "zero") -> np.ndarray:
        """
        Predict next frame from 3 context frames by padding the 4th input slot.
        Returns [B,H,W,C].
        """
        if self.sequence_len is not None and self.sequence_len != 4:
            raise RuntimeError(f"predict_next_from_3 expects sequence_len=4, got {self.sequence_len}")
        last = self.predict_last_frame(frames3, pad_last_frame=pad_last_frame)
        return last