"""Bounded inline temporal state for the warm perception worker.

This module intentionally depends only on the Python standard library and
NumPy so the exact implementation can be included in the minimal, hashed
worker package.  It has no command or navigation authority.
"""

from __future__ import annotations

import hashlib
import json
import math
import time
from collections import deque
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np

PROFILE_SCHEMA = "missioncore.e23-inline-temporal-profile/v1"
PIPELINE_ID = "warm-worker-inline-bounded-temporal-2d-3d-semantic/v1"
TARGET_CLASS_COUNT = 16


@dataclass(slots=True)
class _TrackState:
    canonical_id: int
    source_ids: set[int]
    label: str
    group: str
    center: np.ndarray
    size: np.ndarray
    raw_center: np.ndarray
    velocity: np.ndarray
    last_frame: int
    last_seconds: float
    last_observed_frame: int
    score: float
    template: dict[str, Any]
    cuboid_center: np.ndarray | None = None
    cuboid_size: np.ndarray | None = None
    cuboid_yaw: float | None = None
    cuboid_velocity: np.ndarray | None = None
    cuboid_seconds: float | None = None
    distance_m: float | None = None


class TemporalStabilizer:
    """One-pass bounded state machine for 2D tracks and map-frame cuboids."""

    def __init__(self, profile: dict[str, Any]) -> None:
        self.profile = profile
        self.states: dict[int, _TrackState] = {}
        self.aliases: dict[int, int] = {}
        self.next_id = 1
        self.peak_states = 0
        self.stitched_tracks = 0
        self.held_2d = 0
        self.held_3d = 0
        self.reset_2d = 0
        self.reset_3d = 0

    def update(
        self,
        *,
        frame_index: int,
        session_seconds: float,
        objects: list[dict[str, Any]],
    ) -> list[dict[str, Any]]:
        self._prune(session_seconds)
        assigned: set[int] = set()
        result: list[dict[str, Any]] = []
        for source in sorted(objects, key=lambda item: int(item["track_id"])):
            state, stitched = self._resolve(source, frame_index, assigned)
            assigned.add(state.canonical_id)
            result.append(
                self._observe(
                    state,
                    source,
                    frame_index=frame_index,
                    session_seconds=session_seconds,
                    stitched=stitched,
                )
            )

        hold_frames = int(self.profile["tracking_2d"]["hold_frames"])
        for state in sorted(self.states.values(), key=lambda value: value.canonical_id):
            if state.canonical_id in assigned:
                continue
            gap = frame_index - state.last_observed_frame
            if not 1 <= gap <= hold_frames:
                continue
            held = json.loads(json.dumps(state.template))
            state.center = state.center + state.velocity
            state.last_frame = frame_index
            held["track_id"] = state.canonical_id
            held["bbox_xyxy"] = _box(state.center, state.size)
            held["score"] = max(0.0, state.score * (0.82**gap))
            held["temporal_2d_status"] = f"held-{gap}-frame"
            held["temporal_source_track_id"] = min(state.source_ids)
            self.held_2d += 1
            if not self._hold_cuboid(held, state, session_seconds):
                _clear_cuboid(held, "rejected-temporal-hold-expired-e23")
            result.append(held)

        self.peak_states = max(self.peak_states, len(self.states))
        return sorted(result, key=lambda item: int(item["track_id"]))

    def snapshot(self) -> dict[str, int]:
        return {
            "active_track_states": len(self.states),
            "peak_track_states": self.peak_states,
            "alias_count": len(self.aliases),
            "stitched_tracks": self.stitched_tracks,
            "held_2d": self.held_2d,
            "held_3d": self.held_3d,
            "reset_2d": self.reset_2d,
            "reset_3d": self.reset_3d,
        }

    def _resolve(
        self,
        source: dict[str, Any],
        frame_index: int,
        assigned: set[int],
    ) -> tuple[_TrackState, bool]:
        source_id = int(source["track_id"])
        canonical = self.aliases.get(source_id)
        if canonical is not None and canonical in self.states:
            return self.states[canonical], False

        bbox = np.asarray(source["bbox_xyxy"], dtype=np.float64)
        center, size = _center_size(bbox)
        config = self.profile["tracking_2d"]
        best: tuple[float, _TrackState] | None = None
        for state in self.states.values():
            gap = frame_index - state.last_observed_frame
            if (
                state.canonical_id in assigned
                or state.label != str(source["label"])
                or not 1 <= gap <= int(config["stitch_gap_frames"])
            ):
                continue
            predicted = state.center + state.velocity * gap
            predicted_box = np.asarray(_box(predicted, state.size), dtype=np.float64)
            iou = _iou(predicted_box, bbox)
            scale = max(1.0, math.sqrt(float(np.prod(np.maximum(state.size, 1.0)))))
            normalized_distance = float(np.linalg.norm(center - predicted) / scale)
            if iou < float(config["stitch_minimum_iou"]) and normalized_distance > float(
                config["stitch_maximum_normalized_center_distance"]
            ):
                continue
            score = iou - 0.25 * normalized_distance
            if best is None or score > best[0]:
                best = (score, state)
        if best is not None:
            state = best[1]
            state.source_ids.add(source_id)
            self.aliases[source_id] = state.canonical_id
            self.stitched_tracks += 1
            return state, True

        canonical_id = source_id
        if canonical_id in self.states:
            canonical_id = max(self.next_id, max(self.states, default=0) + 1)
        self.next_id = max(self.next_id, canonical_id + 1)
        state = _TrackState(
            canonical_id=canonical_id,
            source_ids={source_id},
            label=str(source["label"]),
            group=str(source.get("association_group", source["label"])),
            center=center.copy(),
            size=size.copy(),
            raw_center=center.copy(),
            velocity=np.zeros(2, dtype=np.float64),
            last_frame=frame_index,
            last_seconds=0.0,
            last_observed_frame=frame_index,
            score=float(source["score"]),
            template=json.loads(json.dumps(source)),
        )
        self.states[canonical_id] = state
        self.aliases[source_id] = canonical_id
        return state, False

    def _observe(
        self,
        state: _TrackState,
        source: dict[str, Any],
        *,
        frame_index: int,
        session_seconds: float,
        stitched: bool,
    ) -> dict[str, Any]:
        bbox = np.asarray(source["bbox_xyxy"], dtype=np.float64)
        observed_center, observed_size = _center_size(bbox)
        gap = max(1, frame_index - state.last_observed_frame)
        config = self.profile["tracking_2d"]
        predicted = state.center + state.velocity * gap
        scale = max(1.0, math.sqrt(float(np.prod(np.maximum(state.size, 1.0)))))
        innovation = float(np.linalg.norm(observed_center - predicted) / scale)
        if innovation > float(config["maximum_normalized_innovation"]):
            state.center = observed_center
            state.size = observed_size
            state.velocity.fill(0.0)
            status = "reset-large-innovation"
            self.reset_2d += 1
        else:
            blend = _adaptive(
                innovation,
                float(config["adaptive_innovation_low"]),
                float(config["adaptive_innovation_high"]),
            )
            center_alpha = _lerp(
                float(config["center_alpha_low"]),
                float(config["center_alpha_high"]),
                blend,
            )
            size_alpha = _lerp(
                float(config["size_alpha_low"]),
                float(config["size_alpha_high"]),
                blend,
            )
            state.center = predicted + center_alpha * (observed_center - predicted)
            state.size = state.size + size_alpha * (observed_size - state.size)
            observed_velocity = (observed_center - state.raw_center) / gap
            velocity_alpha = float(config["velocity_alpha"])
            state.velocity = (
                1.0 - velocity_alpha
            ) * state.velocity + velocity_alpha * observed_velocity
            status = "stitched-observed" if stitched else "observed"
        state.raw_center = observed_center
        state.last_frame = frame_index
        state.last_observed_frame = frame_index
        state.last_seconds = session_seconds
        state.score = float(source["score"])
        state.template = json.loads(json.dumps(source))

        normalized = json.loads(json.dumps(source))
        source_id = int(source["track_id"])
        normalized["track_id"] = state.canonical_id
        normalized["temporal_source_track_id"] = source_id
        normalized["bbox_xyxy"] = _box(state.center, state.size)
        normalized["temporal_2d_status"] = status
        self._observe_or_hold_cuboid(normalized, state, session_seconds)
        state.template = json.loads(json.dumps(normalized))
        return normalized

    def _observe_or_hold_cuboid(
        self,
        item: dict[str, Any],
        state: _TrackState,
        session_seconds: float,
    ) -> None:
        if str(item.get("cuboid_status", "")).startswith("accepted-"):
            self._observe_cuboid(item, state, session_seconds)
        else:
            self._hold_cuboid(item, state, session_seconds)

    def _observe_cuboid(
        self,
        item: dict[str, Any],
        state: _TrackState,
        session_seconds: float,
    ) -> None:
        center = np.asarray(item["cuboid_center_map"], dtype=np.float64)
        size = np.asarray(item["cuboid_half_size"], dtype=np.float64)
        yaw = _yaw(np.asarray(item["cuboid_quaternion_xyzw"], dtype=np.float64))
        config = self.profile["cuboids_3d"]
        status = "observed"
        if (
            state.cuboid_center is not None
            and state.cuboid_size is not None
            and state.cuboid_yaw is not None
            and state.cuboid_seconds is not None
        ):
            dt = max(1e-3, session_seconds - state.cuboid_seconds)
            velocity = (
                np.zeros(3, dtype=np.float64)
                if state.cuboid_velocity is None
                else state.cuboid_velocity
            )
            predicted = state.cuboid_center + velocity * dt
            innovation = float(np.linalg.norm(center - predicted))
            yaw_delta = _yaw_delta(yaw, state.cuboid_yaw)
            if innovation > float(config["maximum_center_innovation_m"]) or abs(
                math.degrees(yaw_delta)
            ) > float(config["maximum_yaw_innovation_degrees"]):
                status = "reset-large-innovation"
                self.reset_3d += 1
                velocity = np.zeros(3, dtype=np.float64)
            else:
                raw_velocity = (center - state.cuboid_center) / dt
                velocity_alpha = float(config["velocity_alpha"])
                velocity = (1.0 - velocity_alpha) * velocity + velocity_alpha * raw_velocity
                center = predicted + float(config["center_alpha"]) * (center - predicted)
                size = state.cuboid_size + float(config["size_alpha"]) * (size - state.cuboid_size)
                yaw = state.cuboid_yaw + float(config["yaw_alpha"]) * yaw_delta
            state.cuboid_velocity = velocity
        else:
            state.cuboid_velocity = np.zeros(3, dtype=np.float64)
        state.cuboid_center = center
        state.cuboid_size = size
        state.cuboid_yaw = yaw
        state.cuboid_seconds = session_seconds
        distance = item.get("distance_smoothed_m")
        if isinstance(distance, int | float) and not isinstance(distance, bool):
            state.distance_m = float(distance)
        item["cuboid_center_map"] = center.tolist()
        item["cuboid_half_size"] = size.tolist()
        item["cuboid_quaternion_xyzw"] = _quaternion(yaw)
        item["cuboid_status"] = "accepted-temporally-stabilized-e23-v1"
        item["temporal_status"] = f"e23-{status}"

    def _hold_cuboid(
        self,
        item: dict[str, Any],
        state: _TrackState,
        session_seconds: float,
    ) -> bool:
        if (
            state.cuboid_center is None
            or state.cuboid_size is None
            or state.cuboid_yaw is None
            or state.cuboid_seconds is None
            or session_seconds - state.cuboid_seconds
            > float(self.profile["cuboids_3d"]["hold_seconds"])
        ):
            return False
        age_ms = max(0.0, (session_seconds - state.cuboid_seconds) * 1000.0)
        item["cuboid_center_map"] = state.cuboid_center.tolist()
        item["cuboid_half_size"] = state.cuboid_size.tolist()
        item["cuboid_quaternion_xyzw"] = _quaternion(state.cuboid_yaw)
        item["cuboid_status"] = "accepted-temporal-hold-e23-v1"
        item["geometry"] = "temporally-held-last-supported-cuboid"
        item["distance_smoothed_m"] = state.distance_m
        item["temporal_status"] = f"e23-held-{age_ms:.0f}ms"
        item["observed_cuboid_center_map"] = None
        item["observed_cuboid_half_size"] = None
        item["observed_cuboid_quaternion_xyzw"] = None
        self.held_3d += 1
        return True

    def _prune(self, now: float) -> None:
        maximum_idle = float(self.profile["bounds"]["maximum_track_idle_seconds"])
        stale = [
            key
            for key, state in self.states.items()
            if state.last_seconds > 0 and now - state.last_seconds > maximum_idle
        ]
        for key in stale:
            state = self.states.pop(key)
            for source_id in state.source_ids:
                self.aliases.pop(source_id, None)
        maximum = int(self.profile["bounds"]["maximum_track_states"])
        if len(self.states) <= maximum:
            return
        for state in sorted(self.states.values(), key=lambda value: value.last_seconds)[
            : len(self.states) - maximum
        ]:
            self.states.pop(state.canonical_id, None)
            for source_id in state.source_ids:
                self.aliases.pop(source_id, None)


class StreamingSemanticStabilizer:
    """One-mask semantic hysteresis with bounded diagnostic samples."""

    def __init__(self, profile: dict[str, Any]) -> None:
        self.minimum_same_label_neighbors = int(profile["semantic"]["minimum_same_label_neighbors"])
        self.previous_raw: np.ndarray | None = None
        self.previous_stabilized: np.ndarray | None = None
        self.baseline_unsupported = deque(maxlen=4096)
        self.stabilized_unsupported = deque(maxlen=4096)
        self.processing_ms = deque(maxlen=4096)
        self.frames = 0

    def update(self, mask: np.ndarray) -> np.ndarray:
        if mask.dtype != np.uint8 or mask.shape != (600, 800):
            raise RuntimeError("LAB E23 semantic mask is invalid")
        started = time.perf_counter()
        if self.previous_raw is None or self.previous_stabilized is None:
            stabilized = mask.copy()
        else:
            support = _same_label_neighbor_count(mask)
            baseline_change = mask != self.previous_raw
            unsupported_baseline = baseline_change & (support < self.minimum_same_label_neighbors)
            stabilized = mask.copy()
            hold = (mask != self.previous_stabilized) & (
                support < self.minimum_same_label_neighbors
            )
            stabilized[hold] = self.previous_stabilized[hold]
            stabilized_support = _same_label_neighbor_count(stabilized)
            unsupported_stabilized = (stabilized != self.previous_stabilized) & (
                stabilized_support < self.minimum_same_label_neighbors
            )
            self.baseline_unsupported.append(float(np.mean(unsupported_baseline)))
            self.stabilized_unsupported.append(float(np.mean(unsupported_stabilized)))
        self.previous_raw = mask.copy()
        self.previous_stabilized = stabilized
        self.frames += 1
        self.processing_ms.append((time.perf_counter() - started) * 1000.0)
        return stabilized

    def snapshot(self) -> dict[str, Any]:
        baseline = _percentiles(self.baseline_unsupported)
        stabilized = _percentiles(self.stabilized_unsupported)
        baseline_mean = float(baseline["mean"])
        stabilized_mean = float(stabilized["mean"])
        reduction = (
            0.0
            if baseline_mean <= 0
            else max(0.0, (baseline_mean - stabilized_mean) / baseline_mean)
        )
        return {
            "frames": self.frames,
            "history_masks": int(self.previous_stabilized is not None),
            "diagnostic_sample_capacity": self.processing_ms.maxlen,
            "baseline_unsupported_change_fraction": baseline,
            "stabilized_unsupported_change_fraction": stabilized,
            "unsupported_change_reduction_fraction": reduction,
            "processing_ms": _percentiles(self.processing_ms),
        }


def read_inline_profile(path: Path) -> tuple[dict[str, Any], str]:
    resolved = path.expanduser().resolve(strict=True)
    profile = json.loads(resolved.read_text(encoding="utf-8-sig"))
    if not isinstance(profile, dict):
        raise RuntimeError("LAB E23 inline temporal profile is invalid")
    source = profile.get("source")
    tracking = profile.get("tracking_2d")
    cuboids = profile.get("cuboids_3d")
    semantic = profile.get("semantic")
    bounds = profile.get("bounds")
    acceptance = profile.get("acceptance")
    authority = profile.get("authority")
    if (
        profile.get("schema_version") != PROFILE_SCHEMA
        or profile.get("mode") != "inline-shadow-qualification"
        or profile.get("stage") != "warm-worker-after-fusion-before-result-publication"
        or not all(
            isinstance(value, dict)
            for value in (
                source,
                tracking,
                cuboids,
                semantic,
                bounds,
                acceptance,
            )
        )
        or source.get("source_id") != "sensor.camera.right"
        or source.get("resolution") != [800, 600]
        or source.get("calibration_slot") != "camera_1"
        or semantic.get("mode") != "spatially-supported-streaming-hysteresis-v1"
        or int(semantic.get("class_count", 0)) != TARGET_CLASS_COUNT
        or authority != {"commands_enabled": False, "navigation_or_safety_accepted": False}
        or not 1 <= int(bounds.get("maximum_track_states", 0)) <= 512
        or int(bounds.get("semantic_history_masks", 0)) != 1
    ):
        raise RuntimeError("LAB E23 inline temporal profile is invalid")
    for owner, keys in (
        (
            tracking,
            (
                "center_alpha_low",
                "center_alpha_high",
                "size_alpha_low",
                "size_alpha_high",
                "velocity_alpha",
            ),
        ),
        (cuboids, ("center_alpha", "size_alpha", "yaw_alpha", "velocity_alpha")),
    ):
        if any(not 0 < float(owner.get(key, 0)) <= 1 for key in keys):
            raise RuntimeError("LAB E23 smoothing coefficient is invalid")
    if any(
        float(acceptance.get(key, 0)) <= 0
        for key in (
            "maximum_camera_frame_processing_p95_ms",
            "maximum_semantic_frame_processing_p95_ms",
            "maximum_rss_growth_mib",
        )
    ):
        raise RuntimeError("LAB E23 acceptance contract is invalid")
    return profile, _sha256(resolved)


def stabilize_world_state(
    source: dict[str, Any],
    fusion_objects: list[dict[str, Any]],
    memory: dict[int, dict[str, Any]],
) -> dict[str, Any]:
    world = json.loads(json.dumps(source))
    source_objects = {
        int(item["track_id"]): item
        for item in source.get("objects", [])
        if isinstance(item, dict) and isinstance(item.get("track_id"), int)
    }
    objects: list[dict[str, Any]] = []
    active: set[int] = set()
    for fusion in fusion_objects:
        if not str(fusion.get("cuboid_status", "")).startswith("accepted-"):
            continue
        canonical = int(fusion["track_id"])
        source_id = int(fusion.get("temporal_source_track_id", canonical))
        template = source_objects.get(source_id) or memory.get(canonical) or {}
        item = json.loads(json.dumps(template))
        item.update(
            {
                "track_id": canonical,
                "class": str(fusion.get("association_group", "object")),
                "detector_label": str(fusion.get("label", "object")),
                "confidence": float(fusion.get("score", 0.0)),
                "position_map_m": fusion["cuboid_center_map"],
                "orientation_map_xyzw": fusion["cuboid_quaternion_xyzw"],
                "size_m": [2.0 * float(value) for value in fusion["cuboid_half_size"]],
                "range_m": fusion.get("distance_smoothed_m"),
                "support_points": int(fusion.get("clustered_points", 0)),
                "geometry": fusion.get("geometry"),
                "temporal_status": fusion.get("temporal_status"),
            }
        )
        if "held" in str(fusion.get("temporal_status", "")):
            item["velocity_status"] = "temporally-held-diagnostic"
        memory[canonical] = json.loads(json.dumps(item))
        active.add(canonical)
        objects.append(item)
    for canonical in set(memory) - active:
        memory.pop(canonical, None)
    world["objects"] = objects
    world["object_count"] = len(objects)
    world.setdefault("delivery", {})["temporal_stability"] = "e23-inline-bounded"
    return world


def _same_label_neighbor_count(mask: np.ndarray) -> np.ndarray:
    padded = np.pad(mask, 1, mode="edge")
    count = np.zeros(mask.shape, dtype=np.uint8)
    for y in range(3):
        for x in range(3):
            if y == 1 and x == 1:
                continue
            count += padded[y : y + mask.shape[0], x : x + mask.shape[1]] == mask
    return count


def _center_size(box: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    return (box[:2] + box[2:]) * 0.5, np.maximum(box[2:] - box[:2], 1.0)


def _box(center: np.ndarray, size: np.ndarray) -> list[float]:
    half = np.maximum(size, 1.0) * 0.5
    values = np.concatenate((center - half, center + half))
    values[[0, 2]] = np.clip(values[[0, 2]], 0.0, 799.0)
    values[[1, 3]] = np.clip(values[[1, 3]], 0.0, 599.0)
    return [round(float(value), 6) for value in values]


def _iou(left: np.ndarray, right: np.ndarray) -> float:
    intersection_min = np.maximum(left[:2], right[:2])
    intersection_max = np.minimum(left[2:], right[2:])
    intersection_size = np.maximum(0.0, intersection_max - intersection_min)
    intersection = float(np.prod(intersection_size))
    left_area = float(np.prod(np.maximum(0.0, left[2:] - left[:2])))
    right_area = float(np.prod(np.maximum(0.0, right[2:] - right[:2])))
    union = left_area + right_area - intersection
    return 0.0 if union <= 0 else intersection / union


def _adaptive(value: float, lower: float, upper: float) -> float:
    if upper <= lower:
        return 1.0
    return min(1.0, max(0.0, (value - lower) / (upper - lower)))


def _lerp(lower: float, upper: float, fraction: float) -> float:
    return lower + (upper - lower) * fraction


def _yaw(quaternion: np.ndarray) -> float:
    x, y, z, w = quaternion
    return math.atan2(2.0 * (w * z + x * y), 1.0 - 2.0 * (y * y + z * z))


def _yaw_delta(value: float, reference: float) -> float:
    return math.atan2(math.sin(value - reference), math.cos(value - reference))


def _quaternion(yaw: float) -> list[float]:
    return [0.0, 0.0, math.sin(yaw * 0.5), math.cos(yaw * 0.5)]


def _clear_cuboid(item: dict[str, Any], status: str) -> None:
    for key in (
        "cuboid_center_map",
        "cuboid_half_size",
        "cuboid_quaternion_xyzw",
        "observed_cuboid_center_map",
        "observed_cuboid_half_size",
        "observed_cuboid_quaternion_xyzw",
    ):
        item[key] = None
    item["cuboid_status"] = status
    item["temporal_status"] = status


def _percentiles(values: Any) -> dict[str, float | int]:
    samples = np.asarray(list(values), dtype=np.float64)
    if samples.size == 0:
        return {"count": 0, "mean": 0.0, "p50": 0.0, "p95": 0.0, "max": 0.0}
    return {
        "count": int(samples.size),
        "mean": round(float(np.mean(samples)), 6),
        "p50": round(float(np.percentile(samples, 50)), 6),
        "p95": round(float(np.percentile(samples, 95)), 6),
        "max": round(float(np.max(samples)), 6),
    }


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        while chunk := stream.read(1024 * 1024):
            digest.update(chunk)
    return digest.hexdigest()
