"""Source-stamped bilateral-feedback contracts used by formal experiments.

The current simulator historically delayed a residual and a wrench in separate
queues.  That is sufficient for a visual demonstration but cannot guarantee a
matched H3 comparison.  This module makes the comparison unit explicit: one
return packet carries the reconstructed wrench and ``J_s.T @ wrench`` together,
and both mapping conditions consume that immutable packet.

All spatial vectors use the repository convention ``[linear; angular]`` for
twists and ``[force; moment]`` for wrenches.
"""

from __future__ import annotations

from collections import OrderedDict
from dataclasses import dataclass
from enum import Enum, IntEnum
import hashlib
from typing import Iterable

import numpy as np


def _readonly_vector(value, size: int, name: str) -> np.ndarray:
    array = np.asarray(value, dtype=float).reshape(-1).copy()
    if array.shape != (size,):
        raise ValueError(f"{name} must have shape ({size},), got {array.shape}")
    if not np.all(np.isfinite(array)):
        raise ValueError(f"{name} must contain only finite values")
    array.setflags(write=False)
    return array


def _readonly_matrix(value, shape: tuple[int, int], name: str) -> np.ndarray:
    array = np.asarray(value, dtype=float).copy()
    if array.shape != shape:
        raise ValueError(f"{name} must have shape {shape}, got {array.shape}")
    if not np.all(np.isfinite(array)):
        raise ValueError(f"{name} must contain only finite values")
    array.setflags(write=False)
    return array


def _payload_hash(parts: Iterable[np.ndarray | int | float]) -> str:
    digest = hashlib.sha256()
    for part in parts:
        if isinstance(part, np.ndarray):
            contiguous = np.ascontiguousarray(part, dtype=np.float64)
            digest.update(str(contiguous.shape).encode("ascii"))
            digest.update(contiguous.tobytes())
        elif isinstance(part, (int, np.integer)):
            digest.update(f"i:{int(part)}".encode("ascii"))
        else:
            digest.update(f"f:{float(part):.17g}".encode("ascii"))
    return digest.hexdigest()


class MappingKind(str, Enum):
    """Feedback mappings required by the H3/H4 comparison."""

    DIFFERENTIAL_RESIDUAL = "differential_residual"
    MATCHED_DIFFERENTIAL_WRENCH = "matched_differential_wrench"
    DIRECT_MASTER_JACOBIAN = "direct_master_jacobian"


class MapPolicy(str, Enum):
    """Which differential is used for a delayed return packet."""

    CURRENT = "current"
    SOURCE_STAMPED = "source_stamped"


class PacketState(IntEnum):
    EMPTY = 0
    ACTIVE = 1
    HELD = 2
    TIMED_OUT = 3
    RECOVERING = 4


class PacketRejectReason(IntEnum):
    ACCEPTED = 0
    INVALID = 1
    CORRUPT = 2
    DUPLICATE_OR_STALE = 3


@dataclass(frozen=True)
class MapSnapshot:
    """One accepted 50-Hz retargeting differential."""

    map_id: int
    source_index: int
    source_time: float
    differential: np.ndarray
    valid: bool = True
    reason_code: int = 0

    def __post_init__(self) -> None:
        if self.map_id < 0 or self.source_index < 0:
            raise ValueError("map_id and source_index must be non-negative")
        if not np.isfinite(self.source_time):
            raise ValueError("source_time must be finite")
        matrix = np.asarray(self.differential, dtype=float)
        if matrix.ndim != 2 or matrix.shape[0] == 0 or matrix.shape[1] == 0:
            raise ValueError("differential must be a non-empty matrix")
        object.__setattr__(
            self,
            "differential",
            _readonly_matrix(matrix, matrix.shape, "differential"),
        )

    @property
    def digest(self) -> str:
        return _payload_hash(
            (
                self.map_id,
                self.source_index,
                self.source_time,
                self.differential,
            )
        )


class MapRegistry:
    """Bounded map-id registry for source-stamped delayed feedback."""

    def __init__(self, capacity: int = 512):
        if capacity <= 0:
            raise ValueError("capacity must be positive")
        self.capacity = int(capacity)
        self._maps: OrderedDict[int, MapSnapshot] = OrderedDict()

    def add(self, snapshot: MapSnapshot) -> None:
        existing = self._maps.get(snapshot.map_id)
        if existing is not None and existing.digest != snapshot.digest:
            raise ValueError(
                f"map_id {snapshot.map_id} cannot be reused for different data"
            )
        self._maps[snapshot.map_id] = snapshot
        self._maps.move_to_end(snapshot.map_id)
        while len(self._maps) > self.capacity:
            self._maps.popitem(last=False)

    def get(self, map_id: int) -> MapSnapshot:
        try:
            return self._maps[int(map_id)]
        except KeyError as exc:
            raise KeyError(f"map_id {map_id} is not available") from exc

    @property
    def latest(self) -> MapSnapshot:
        if not self._maps:
            raise KeyError("map registry is empty")
        return next(reversed(self._maps.values()))

    def __len__(self) -> int:
        return len(self._maps)


@dataclass(frozen=True)
class ForwardPacket:
    """Master-to-slave held reference packet."""

    seq: int
    source_index: int
    source_time: float
    map_id: int
    q_slave_ref: np.ndarray
    qd_slave_ctrl: np.ndarray
    valid: bool = True

    def __post_init__(self) -> None:
        if self.seq < 0 or self.source_index < 0 or self.map_id < 0:
            raise ValueError("packet identifiers must be non-negative")
        if not np.isfinite(self.source_time):
            raise ValueError("source_time must be finite")
        q = np.asarray(self.q_slave_ref, dtype=float).reshape(-1)
        qd = np.asarray(self.qd_slave_ctrl, dtype=float).reshape(-1)
        if q.size == 0 or q.shape != qd.shape:
            raise ValueError("q_slave_ref and qd_slave_ctrl must have equal shapes")
        object.__setattr__(
            self, "q_slave_ref", _readonly_vector(q, q.size, "q_slave_ref")
        )
        object.__setattr__(
            self, "qd_slave_ctrl", _readonly_vector(qd, qd.size, "qd_slave_ctrl")
        )


@dataclass(frozen=True)
class ReturnPacket:
    """Slave-to-master packet supporting every declared feedback condition."""

    seq: int
    source_index: int
    source_time: float
    echoed_map_id: int
    residual: np.ndarray
    wrench: np.ndarray
    js_t_wrench: np.ndarray
    qd_slave_actual: np.ndarray
    valid: bool = True

    def __post_init__(self) -> None:
        if self.seq < 0 or self.source_index < 0 or self.echoed_map_id < 0:
            raise ValueError("packet identifiers must be non-negative")
        if not np.isfinite(self.source_time):
            raise ValueError("source_time must be finite")
        residual = np.asarray(self.residual, dtype=float).reshape(-1)
        js_t = np.asarray(self.js_t_wrench, dtype=float).reshape(-1)
        qd = np.asarray(self.qd_slave_actual, dtype=float).reshape(-1)
        if residual.size == 0 or residual.shape != js_t.shape or residual.shape != qd.shape:
            raise ValueError(
                "residual, js_t_wrench and qd_slave_actual must have equal shapes"
            )
        object.__setattr__(
            self, "residual", _readonly_vector(residual, residual.size, "residual")
        )
        object.__setattr__(
            self, "wrench", _readonly_vector(self.wrench, 6, "wrench")
        )
        object.__setattr__(
            self,
            "js_t_wrench",
            _readonly_vector(js_t, js_t.size, "js_t_wrench"),
        )
        object.__setattr__(
            self,
            "qd_slave_actual",
            _readonly_vector(qd, qd.size, "qd_slave_actual"),
        )

    @property
    def matched_input_hash(self) -> str:
        """Hash proving that H3 matched conditions consumed one packet."""
        return _payload_hash(
            (
                self.seq,
                self.source_index,
                self.source_time,
                self.echoed_map_id,
                self.wrench,
                self.js_t_wrench,
            )
        )


@dataclass(frozen=True)
class FeedbackResult:
    tau_master_raw: np.ndarray
    tau_slave_source: np.ndarray
    selected_map_id: int | None
    matched_input_hash: str
    valid: bool
    reason: str


def map_return_feedback(
    *,
    kind: MappingKind,
    packet: ReturnPacket,
    master_jacobian: np.ndarray,
    maps: MapRegistry,
    map_policy: MapPolicy = MapPolicy.CURRENT,
) -> FeedbackResult:
    """Map one immutable return packet to a raw master generalized torque.

    ``MATCHED_DIFFERENTIAL_WRENCH`` and ``DIRECT_MASTER_JACOBIAN`` both use the
    packet's one reconstructed wrench.  The former consumes its stored
    ``J_s.T @ wrench`` value; the latter consumes ``wrench`` directly.
    """
    Jm = np.asarray(master_jacobian, dtype=float)
    if Jm.ndim != 2 or Jm.shape[0] != 6:
        raise ValueError("master_jacobian must have shape (6, nv_master)")
    if not np.all(np.isfinite(Jm)):
        raise ValueError("master_jacobian must contain only finite values")

    matched_hash = packet.matched_input_hash
    if not packet.valid:
        return FeedbackResult(
            tau_master_raw=np.zeros(Jm.shape[1], dtype=float),
            tau_slave_source=np.zeros_like(packet.residual),
            selected_map_id=None,
            matched_input_hash=matched_hash,
            valid=False,
            reason="invalid_return_packet",
        )

    if kind is MappingKind.DIRECT_MASTER_JACOBIAN:
        return FeedbackResult(
            tau_master_raw=np.asarray(Jm.T @ packet.wrench, dtype=float),
            tau_slave_source=packet.js_t_wrench.copy(),
            selected_map_id=None,
            matched_input_hash=matched_hash,
            valid=True,
            reason="ok",
        )

    if map_policy is MapPolicy.SOURCE_STAMPED:
        selected = maps.get(packet.echoed_map_id)
    else:
        selected = maps.latest
    A = selected.differential
    if A.shape[0] != packet.residual.size or A.shape[1] != Jm.shape[1]:
        raise ValueError(
            "differential shape is incompatible with slave/master dimensions"
        )
    if not selected.valid:
        return FeedbackResult(
            tau_master_raw=np.zeros(Jm.shape[1], dtype=float),
            tau_slave_source=np.zeros(A.shape[0], dtype=float),
            selected_map_id=selected.map_id,
            matched_input_hash=matched_hash,
            valid=False,
            reason="invalid_differential",
        )

    source = (
        packet.residual
        if kind is MappingKind.DIFFERENTIAL_RESIDUAL
        else packet.js_t_wrench
    )
    return FeedbackResult(
        tau_master_raw=np.asarray(A.T @ source, dtype=float),
        tau_slave_source=source.copy(),
        selected_map_id=selected.map_id,
        matched_input_hash=matched_hash,
        valid=True,
        reason="ok",
    )


def normalized_actual_power_mismatch(
    tau_master_raw: np.ndarray,
    qd_master: np.ndarray,
    tau_slave_source: np.ndarray,
    qd_slave_actual: np.ndarray,
    dt: float | np.ndarray,
    *,
    sigma_feedback: float = 1.0,
    epsilon_energy: float = 1e-12,
) -> float:
    """Compute the trajectory-level H3 endpoint from source-aligned samples."""
    tm = np.asarray(tau_master_raw, dtype=float)
    qm = np.asarray(qd_master, dtype=float)
    ts = np.asarray(tau_slave_source, dtype=float)
    qs = np.asarray(qd_slave_actual, dtype=float)
    if tm.ndim != 2 or tm.shape != qm.shape or ts.ndim != 2 or ts.shape != qs.shape:
        raise ValueError("torque and velocity arrays must be paired 2-D arrays")
    if tm.shape[0] != ts.shape[0]:
        raise ValueError("master and slave arrays must have the same sample count")
    delta = np.broadcast_to(np.asarray(dt, dtype=float), (tm.shape[0],))
    if np.any(delta <= 0.0) or not np.all(np.isfinite(delta)):
        raise ValueError("dt must contain finite positive values")
    if epsilon_energy <= 0.0:
        raise ValueError("epsilon_energy must be positive")

    pm = np.einsum("ij,ij->i", tm, qm)
    ps = float(sigma_feedback) * np.einsum("ij,ij->i", ts, qs)
    numerator = float(np.sum(np.abs(pm - ps) * delta))
    denominator = float(
        0.5 * np.sum((np.abs(pm) + np.abs(ps)) * delta) + epsilon_energy
    )
    return numerator / denominator