"""APE model and compact cache used by the isolated Keyframe baseline."""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import torch
from torch import nn


class ActionCopycat(nn.Module):
    """Official one-hidden-layer action-only copycat architecture."""

    def __init__(self, action_dim: int = 7, history: int = 9, hidden: int = 300):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(action_dim * history, hidden),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(hidden, action_dim),
        )

    def forward(self, action_history: torch.Tensor) -> torch.Tensor:
        return self.net(action_history)


def q99_normalize(
    actions: np.ndarray,
    statistics: dict,
    *,
    binary_action_indices: tuple[int, ...] = (),
) -> np.ndarray:
    """Apply q01/q99 normalization with an explicit binary-action contract.

    A zero percentile span is not sufficient evidence that an action dimension is
    binary.  This matters for RoboTwin absolute qpos, where a joint can be constant
    in a dataset shard while remaining a continuous joint.  Such continuous,
    zero-span dimensions are mapped to zero; only explicitly declared indices use
    the binary threshold.
    """

    actions = np.asarray(actions, dtype=np.float32)
    q01 = np.asarray(statistics["q01"], dtype=np.float32)
    q99 = np.asarray(statistics["q99"], dtype=np.float32)
    if actions.shape[-1] != len(q01) or q01.shape != q99.shape:
        raise ValueError(
            f"action/statistics shape mismatch: actions={actions.shape}, "
            f"q01={q01.shape}, q99={q99.shape}"
        )
    binary = np.zeros(len(q01), dtype=bool)
    for index in binary_action_indices:
        if index < 0 or index >= len(q01):
            raise ValueError(f"binary action index {index} is out of range")
        binary[index] = True
    continuous = ~binary
    scaled = continuous & (q01 != q99)
    constant = continuous & (q01 == q99)
    normalized = np.zeros_like(actions, dtype=np.float32)
    normalized[..., scaled] = (
        2.0
        * (actions[..., scaled] - q01[scaled])
        / (q99[scaled] - q01[scaled])
        - 1.0
    )
    normalized[..., scaled] = np.clip(
        normalized[..., scaled], -2.2, 2.2
    )
    normalized[..., constant] = 0.0
    normalized[..., binary] = (
        actions[..., binary] > 0.5
    ).astype(np.float32)
    return normalized


def history_target(actions: np.ndarray, anchor: int, history: int = 9):
    """Repeat the earliest physical action at a trajectory boundary."""

    indices = np.arange(anchor - history, anchor, dtype=np.int64)
    indices = np.clip(indices, 0, len(actions) - 1)
    return actions[indices].reshape(-1), actions[anchor]


@dataclass(frozen=True)
class CacheRecord:
    ape: float
    weight: float


class KeyframeWeightCache:
    """Read a compact NPZ cache without altering sampling."""

    def __init__(self, path: str | Path):
        self.path = Path(path)
        payload = np.load(self.path, allow_pickle=False)
        self.metadata = json.loads(str(payload["metadata_json"].item()))
        if self.metadata.get("method") != "Keyframe-Focused-adapted":
            raise RuntimeError("unexpected keyframe cache method")
        self._datasets = {}
        for index, name in enumerate(self.metadata["dataset_names"]):
            trajectory_ids = payload[f"d{index}_trajectory_ids"].astype(np.int64)
            offsets = payload[f"d{index}_offsets"].astype(np.int64)
            anchors = payload[f"d{index}_anchors"].astype(np.int64)
            ape = payload[f"d{index}_ape"].astype(np.float32)
            weights = payload[f"d{index}_weights"].astype(np.float32)
            if not (len(anchors) == len(ape) == len(weights)):
                raise RuntimeError(f"cache array lengths differ for {name}")
            if len(offsets) != len(trajectory_ids) + 1 or int(offsets[-1]) != len(anchors):
                raise RuntimeError(f"cache offsets are invalid for {name}")
            if np.any(~np.isfinite(ape)) or np.any(~np.isfinite(weights)):
                raise RuntimeError(f"cache contains non-finite values for {name}")
            if np.any(weights <= 0):
                raise RuntimeError(f"cache contains non-positive weights for {name}")
            self._datasets[name] = {
                "trajectory_to_row": {
                    int(trajectory_id): row
                    for row, trajectory_id in enumerate(trajectory_ids)
                },
                "offsets": offsets,
                "anchors": anchors,
                "ape": ape,
                "weights": weights,
            }

    def lookup(self, dataset_name: str, trajectory_id: int, base_index: int) -> CacheRecord:
        try:
            dataset = self._datasets[dataset_name]
            row = dataset["trajectory_to_row"][int(trajectory_id)]
        except KeyError as exc:
            raise KeyError(
                f"keyframe cache has no trajectory {(dataset_name, trajectory_id)}"
            ) from exc
        start, end = dataset["offsets"][row : row + 2]
        local = dataset["anchors"][start:end]
        position = int(np.searchsorted(local, int(base_index)))
        if position >= len(local) or int(local[position]) != int(base_index):
            raise KeyError(
                f"keyframe cache has no anchor {(dataset_name, trajectory_id, base_index)}"
            )
        absolute = int(start + position)
        return CacheRecord(
            ape=float(dataset["ape"][absolute]),
            weight=float(dataset["weights"][absolute]),
        )