#!/usr/bin/env python3
"""Pure numerical contracts for matched H3 feedback and packet semantics."""

from pathlib import Path
import sys
import unittest

import numpy as np


CODE_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(CODE_ROOT))

from core.feedback_protocol import (  # noqa: E402
    ForwardPacket,
    MapPolicy,
    MapRegistry,
    MapSnapshot,
    MappingKind,
    PacketRejectReason,
    PacketState,
    ReturnPacket,
    map_return_feedback,
    normalized_actual_power_mismatch,
)
from core.network_emulator import (  # noqa: E402
    DeterministicChannel,
    NetworkTraceEntry,
    PacketReceiver,
)


class FeedbackProtocolTest(unittest.TestCase):
    def make_packet(self, seq=3, map_id=1):
        wrench = np.array([1.0, -2.0, 0.5, 0.2, -0.3, 0.1])
        Js = np.arange(42, dtype=float).reshape(6, 7) / 50.0
        return ReturnPacket(
            seq=seq,
            source_index=10 + seq,
            source_time=0.02 * seq,
            echoed_map_id=map_id,
            residual=np.linspace(-1.0, 1.0, 7),
            wrench=wrench,
            js_t_wrench=Js.T @ wrench,
            qd_slave_actual=np.linspace(0.1, 0.7, 7),
        )

    def test_matched_conditions_share_one_input_hash(self):
        registry = MapRegistry()
        A = np.eye(7)
        registry.add(MapSnapshot(1, 2, 0.04, A))
        packet = self.make_packet()
        Jm = np.arange(42, dtype=float).reshape(6, 7) / 70.0

        differential = map_return_feedback(
            kind=MappingKind.MATCHED_DIFFERENTIAL_WRENCH,
            packet=packet,
            master_jacobian=Jm,
            maps=registry,
        )
        direct = map_return_feedback(
            kind=MappingKind.DIRECT_MASTER_JACOBIAN,
            packet=packet,
            master_jacobian=Jm,
            maps=registry,
        )

        self.assertEqual(
            differential.matched_input_hash, direct.matched_input_hash
        )
        np.testing.assert_allclose(
            differential.tau_master_raw, A.T @ packet.js_t_wrench
        )
        np.testing.assert_allclose(
            direct.tau_master_raw, Jm.T @ packet.wrench
        )

    def test_source_stamped_policy_uses_echoed_map(self):
        registry = MapRegistry()
        registry.add(MapSnapshot(1, 0, 0.0, np.eye(7)))
        registry.add(MapSnapshot(2, 1, 0.02, 2.0 * np.eye(7)))
        packet = self.make_packet(map_id=1)
        Jm = np.zeros((6, 7))
        source = map_return_feedback(
            kind=MappingKind.MATCHED_DIFFERENTIAL_WRENCH,
            packet=packet,
            master_jacobian=Jm,
            maps=registry,
            map_policy=MapPolicy.SOURCE_STAMPED,
        )
        current = map_return_feedback(
            kind=MappingKind.MATCHED_DIFFERENTIAL_WRENCH,
            packet=packet,
            master_jacobian=Jm,
            maps=registry,
            map_policy=MapPolicy.CURRENT,
        )
        self.assertEqual(source.selected_map_id, 1)
        self.assertEqual(current.selected_map_id, 2)
        np.testing.assert_allclose(
            current.tau_master_raw, 2.0 * source.tau_master_raw
        )

    def test_power_endpoint_is_zero_for_exact_source_alignment(self):
        tau_s = np.array([[1.0, 2.0], [-1.0, 0.5]])
        qd_s = np.array([[0.3, -0.1], [0.2, 0.4]])
        tau_m = tau_s.copy()
        qd_m = qd_s.copy()
        mismatch = normalized_actual_power_mismatch(
            tau_m, qd_m, tau_s, qd_s, np.array([0.01, 0.02])
        )
        self.assertAlmostEqual(mismatch, 0.0)

    def test_duplicate_out_of_order_and_timeout_are_explicit(self):
        packets = [
            ForwardPacket(i, i, i * 0.01, i, np.zeros(2), np.zeros(2))
            for i in range(3)
        ]
        # seq 0 arrives after seq 1, and seq 1 is duplicated.
        channel = DeterministicChannel(
            [
                NetworkTraceEntry(0, 0.04),
                NetworkTraceEntry(1, 0.01, duplicate=True),
                NetworkTraceEntry(2, 0.02),
            ]
        )
        for packet in packets:
            channel.send(packet, packet.source_time)
        receiver = PacketReceiver(timeout_s=0.05)
        receptions = []
        for delivery in channel.poll(0.05):
            receptions.append(receiver.accept(delivery))

        self.assertTrue(receptions[0].accepted)
        self.assertIn(
            PacketRejectReason.DUPLICATE_OR_STALE,
            [item.reason for item in receptions if not item.accepted],
        )
        held = receiver.sample(0.051)
        self.assertIn(held.state, (PacketState.ACTIVE, PacketState.RECOVERING))
        self.assertEqual(receiver.sample(0.2).state, PacketState.TIMED_OUT)


if __name__ == "__main__":
    unittest.main()