# SPDX-License-Identifier: Apache-2.0
"""Host pre-processing (no device, no weights): the node's normalization rule and the encoder's in-graph
pre-processing as host features, against the verified research reference (``dp_common.py``) and against the rules
of ``preprocessing_utils.cpp:34-84`` / SPEC 3.8 on hand-made cases. (The features are proven against the exported
graph itself, tensor by tensor, in ``test_reference_cpu.py``.)

    TT_VISIBLE_DEVICES=none python -m pytest -q code/tt_diffusion_planner/tests/test_host_preprocess.py
"""
from __future__ import annotations

import json
import math

import numpy as np
import pytest

from tt_diffusion_planner.host.features import atan2_onnx, decoder_masks, encoder_features, line_features
from tt_diffusion_planner.host.normalize import FLT_EPSILON, normalize_array, normalize_inputs, speed_masks
from tt_diffusion_planner.host.pipeline import prepare
from tt_diffusion_planner.reference import config as C
from tt_diffusion_planner.reference.weights import load_param_json
from tt_diffusion_planner.tests import _research as R

PARAM_JSON = None
for _cand in (R.RESEARCH.parents[1] / "assets" / "diffusion-planner" / "hf_diffusion_planner" / C.PARAM_JSON,):
    if _cand.is_file():
        PARAM_JSON = _cand


@pytest.fixture(scope="module")
def normalization():
    if PARAM_JSON is None:
        pytest.skip("diffusion_planner.param.json not found")
    return load_param_json(PARAM_JSON)


def _zero_raw():
    raw = {k: np.zeros(s, np.float32) for k, s in C.INPUT_SHAPES.items()}
    raw["ego_current_state"][0, :4] = (0, 0, 1, 0)
    raw["ego_agent_past"][0, :, 2] = 1.0
    raw["turn_indicators"][:] = 1.0
    raw["ego_shape"][0] = (2.79, 4.89, 1.896)
    return raw


# ------------------------------------------------------------------------------------------- normalization

def test_normalize_rule_zero_rows_and_columns():
    mean = np.array([10.0, 0.0, 0.0, 0.0], np.float32)
    std = np.array([20.0, 20.0, 1.0, 1.0], np.float32)
    v = np.array([[0, 0, 0, 0],                                   # padding: untouched
                  [1e-8, -1e-8, 0, 1e-9],                          # every |v| < FLT_EPSILON: untouched (kept as is)
                  [30.0, 2.0, 1.0, 0.0],                           # a real row
                  [0, 0, 1.0, 0]], np.float32)                     # cos = 1 makes the row real: x -> -0.5
    out = normalize_array(v, mean, std)
    np.testing.assert_array_equal(out[0], v[0])
    np.testing.assert_array_equal(out[1], v[1])
    np.testing.assert_array_equal(out[2], ((v[2] - mean) / std).astype(np.float32))
    np.testing.assert_array_equal(out[3], np.array([-0.5, 0, 1, 0], np.float32))
    one = normalize_array(np.array([[0.0], [10.0], [-FLT_EPSILON]], np.float32), np.array([0.0]), np.array([20.0]))
    np.testing.assert_array_equal(one.ravel(), np.array([0.0, 0.5, -FLT_EPSILON / 20], np.float32))
    with pytest.raises(ValueError, match="Standard deviation is zero"):
        normalize_array(v, mean, np.array([20, 20, 0, 1], np.float32))
    with pytest.raises(KeyError, match="Missing key lanes"):
        normalize_inputs({"lanes": np.zeros(C.INPUT_SHAPES["lanes"], np.float32)}, {})


def test_normalize_skips_raw_keys(normalization):
    raw = _zero_raw()
    raw["sampled_trajectories"][0, 0, 5] = (1, 2, 3, 4)
    norm = normalize_inputs(raw, normalization.observation)
    for k in C.SKIP_NORMALIZATION:
        np.testing.assert_array_equal(norm[k], raw[k])
    assert norm["ego_agent_past"][0, 0].tolist() == [-0.5, 0.0, 1.0, 0.0]


@pytest.mark.skipif(R.load_script("dp_common") is None, reason="research scripts not present")
@pytest.mark.parametrize("scene", R.research_scenes() or ["none"])
def test_normalization_and_masks_equal_research_reference(normalization, scene):
    dp = R.load_script("dp_common")
    raw = R.research_scene_raw(scene)
    if raw is None:
        pytest.skip("no research scene")
    _, obs, _ = dp.load_param_json(str(PARAM_JSON))
    want = dp.normalize_inputs(raw, obs)
    got = normalize_inputs(raw, normalization.observation)
    for k in want:
        np.testing.assert_array_equal(got[k], want[k], err_msg=k)
    for k, v in dp.speed_masks(want).items():
        np.testing.assert_array_equal(speed_masks(got)[k], v, err_msg=k)
    with np.load(R.ORT_GOLDENS / f"golden_{scene}.npz") as z:  # what the ORT goldens were fed
        for k in C.ENCODER_INPUTS:
            src = got[k] if k in got else speed_masks(got)[k]
            np.testing.assert_array_equal(src, z["in/" + k], err_msg=k)


# ------------------------------------------------------------------------------------------------- features

def test_atan2_onnx_quadrants_and_quirks():
    rng = np.random.default_rng(0)
    y = rng.normal(size=1000).astype(np.float32)
    x = rng.normal(size=1000).astype(np.float32)
    np.testing.assert_allclose(atan2_onnx(y, x), np.arctan2(y, x), rtol=0, atol=2e-6)
    # x = 0: dy / 0 = +-inf -> +-pi/2 like atan2
    np.testing.assert_allclose(atan2_onnx(np.float32([2.0, -3.0]), np.float32([0.0, 0.0])),
                               [math.pi / 2, -math.pi / 2], atol=1e-7)
    # the exported decomposition maps y = 0, x < 0 to -pi (IEEE atan2 gives +pi); 0 / 0 is NaN (masked later)
    assert atan2_onnx(np.float32(0.0), np.float32(-1.0)) == np.float32(-3.1415927)
    assert np.isnan(atan2_onnx(np.float32(0.0), np.float32(0.0)))


def test_line_features_append_point_deltas():
    x = np.zeros((2, 4, 3), np.float32)
    x[0, :, 0] = [0, 1, 3, 6]
    x[0, :, 1] = [0, -1, -1, 2]
    x[0, :, 2] = 1.0
    f = line_features(x)
    assert f.shape == (2, 4, 5)
    np.testing.assert_array_equal(f[0, :, 3], [1, 2, 3, 0])
    np.testing.assert_array_equal(f[0, :, 4], [-1, 0, 3, 0])
    np.testing.assert_array_equal(f[1], 0)


def test_encoder_features_rules(normalization):
    raw = _zero_raw()
    raw["ego_agent_past"][0, :, 0] = np.arange(31) - 30.0                    # x: -30 .. 0
    nb = raw["neighbor_agents_past"][0]
    nb[0, :, :4] = (12.0, 3.0, 1.0, 0.0)                                       # agent 0: valid everywhere
    nb[0, :, 4:8] = (5.0, 1.0, 1.8, 4.5)
    nb[0, :, 8] = 1.0
    nb[1, 10, 0] = 50.0                                                        # agent 1: only an old (dropped) row
    nb[2, 30, 4] = 2.0                                                         # agent 2: only a velocity at t = 30
    nb[3, 27, :4] = (20.0, 1.0, 0.0, 1.0)                                      # agent 3: a pose at t = 27 only
    raw["lanes"][0, 0, :, 0] = np.linspace(0, 19, 20)                          # lane 0 along +x
    raw["lanes"][0, 0, :19, 2] = 1.0
    raw["lanes"][0, 0, 0, 12] = 1.0                                            # no traffic light (dim 12)
    raw["lanes_speed_limit"][0, 0, 0] = 8.0
    raw["polygons"][0, 0, :, 0] = np.linspace(0, 39, 40)
    raw["polygons"][0, 0, :, 2] = 1.0
    raw["line_strings"][0, 1, :, 1] = np.linspace(5, 24, 20)
    raw["line_strings"][0, 1, :, 3] = 1.0                                      # road border
    norm = normalize_inputs(raw, normalization.observation)
    norm.update(speed_masks(norm))
    f = encoder_features(norm, norm)
    # ego: the 6 OLDEST rows kept, the rest zeroed; the token is valid with an all-zero position
    np.testing.assert_array_equal(f.ego[:6], norm["ego_agent_past"][0, :6])
    assert not f.ego[6:].any() and f.valid["ego"].all() and not f.pos[0, :4].any() and f.pos[0, 4] == 1
    # neighbours: rows 0..24 dropped; agent 1 (old row only) invalid; agent 2's t = 30 row is non-zero (a velocity),
    # so the normalisation moved its x to (0 - 10) / 20 = -0.5; the velocity is zeroed after the validity test and
    # the valid-step flag set
    assert f.valid["neighbor"][:4].tolist() == [True, False, True, True]
    assert not f.neighbor[:, :25].any()
    np.testing.assert_array_equal(f.neighbor[0, 25:, 4:6], 0)
    np.testing.assert_array_equal(f.neighbor[0, 25:, 8], 1)
    np.testing.assert_array_equal(f.neighbor[2, 30], [-0.5, 0, 0, 0, 0, 0, 0, 0, 1])
    np.testing.assert_array_equal(f.neighbor[3, :, 8], (np.arange(31) == 27).astype(np.float32))
    np.testing.assert_array_equal(f.neighbor_type[0], [1, 0, 0])
    sl = C.TOKEN_SLICES
    np.testing.assert_allclose(f.pos[sl["neighbor"].start, :4], [0.1, 0.15, 1.0, 0.0], atol=1e-7)
    # lanes: heading of point 10 from (dx, dy); speed and its TRT mask; attributes from point 0
    assert f.valid["lane"][0] and not f.valid["lane"][1:].any()
    np.testing.assert_allclose(f.pos[sl["lane"].start, :4], [norm["lanes"][0, 0, 10, 0], 0, 1, 0], atol=1e-7)
    assert f.lane_has_speed[0, 0] and not f.lane_has_speed[1:].any() and f.lane_speed[0, 0] == np.float32(0.4)
    assert f.lane_attr[0, 12 - C.LANE_FEATURE_DIM] == 1
    # polygons: the pseudo-heading atan2(dx, is_intersection_area) of point 20 (dx = 1/20 normalised)
    h = math.atan2(1.0 / 20.0, 1.0)
    np.testing.assert_allclose(f.pos[sl["polygon"].start, 2:4], [math.cos(h), math.sin(h)], atol=1e-6)
    # line strings: atan2(is_road_border, is_stop_line) = pi/2 for a road border, whatever its geometry
    np.testing.assert_allclose(f.pos[sl["line_string"].start + 1, 2:4], [0, 1], atol=1e-6)
    assert not f.valid["line_string"][0] and f.valid["line_string"][1]
    # static objects are zeros in Autoware: never valid; goal / ego shape / turn always valid
    assert not f.valid["static"].any() and f.valid["goal"].all() and f.valid["turn"].all()
    np.testing.assert_array_equal(f.pos[sl["turn"].start, :4 + 9], [0, 0, 1, 0] + [0] * 8 + [1])
    np.testing.assert_array_equal(f.turn, raw["turn_indicators"][0, :30])
    # masks: invalid tokens have no position feature; the ego key is always valid
    assert not f.pos[~f.token_valid].any()
    assert f.key_valid[0] and (f.key_valid[1:] == f.token_valid[1:]).all()
    assert f.token_valid.sum() == 1 + 3 + 1 + 1 + 1 + 3
    # the decoder keys on the CURRENT (normalised) pose only (dit.py:155-156): agent 3, whose only pose is at t = 27,
    # is an encoder entity but not a decoder key; agent 2's normalised x = -0.5 makes it a key
    d = decoder_masks(norm)
    assert d.agent_valid[:5].tolist() == [True, True, False, True, False]
    np.testing.assert_array_equal(d.current_states[0], [-0.5, 0, 1, 0])
    np.testing.assert_array_equal(d.current_states[1], norm["neighbor_agents_past"][0, 0, 30, :4])


def test_prepare_on_the_shipped_samples(normalization):
    for stem, n_agents in (("kashiwanoha_dense", 88), ("straight_road", 12)):
        raw = R.sample_raw(stem)
        p = prepare(raw, normalization.observation)
        assert p.neighbor_rows.tolist() == list(range(n_agents))
        assert p.enable_force_stop and p.prev_report == 1
        assert p.features.valid["neighbor"].sum() == n_agents == p.decoder.agent_valid.sum() - 1
        np.testing.assert_array_equal(p.x_T, raw["sampled_trajectories"][0])
        assert set(p.norm) >= set(C.INPUT_NAMES) | {"lanes_has_speed_limit", "route_lanes_has_speed_limit"}


def test_prepare_refuses_bad_inputs(normalization):
    from tt_diffusion_planner.ttaw.io import InputError

    raw = _zero_raw()
    del raw["delay"]
    with pytest.raises(InputError, match="missing"):
        prepare(raw, normalization.observation)
    raw = _zero_raw()
    raw["lanes"] = np.zeros((1, 70, 20, 33), np.float32)
    with pytest.raises(InputError, match="shape"):
        prepare(raw, normalization.observation)


def test_param_json_contract(normalization):
    assert normalization.major_version == 5
    mean, std = normalization.state()
    assert mean.shape == (C.MAX_NUM_AGENTS, 1, 4)
    np.testing.assert_array_equal(np.unique(mean.reshape(-1, 4), axis=0), [[10, 0, 0, 0]])
    np.testing.assert_array_equal(np.unique(std.reshape(-1, 4), axis=0), [[20, 20, 1, 1]])
    assert set(normalization.observation) == set(C.INPUT_NAMES) - set(C.SKIP_NORMALIZATION)
    args = normalization.args
    assert (args["agent_num"], args["time_len"], args["future_len"], args["hidden_dim"]) == (320, 31, 80, 256)
    assert (args["encoder_mixer_depth"], args["encoder_fusion_depth"], args["decoder_depth"]) == (6, 6, 3)
    assert json.loads(json.dumps(args))["diffusion_model_type"] == "x_start"