from __future__ import annotations

import os

import numpy as np
import pytest

from SR_analysis.checks.order_contract import expected_ledger, impulse_commands
from SR_analysis.configs import (
    ILLUSION_TARGET_FORCE_SLICE,
    ILLUSION_TARGET_SENSOR_SLICE,
    LEGACY_CFG_DIR,
    MODEL_DIR,
    PROJECT_ROOT,
    control_dt,
    get_scene,
)
from SR_analysis.utils.cfd_interface import (
    LEGACY_DTW_VERSION,
    calc_dtw_sim,
    compute_similarity,
    compute_similarity_illusion_legacy_v1,
    compute_similarity_karman_legacy_v1,
)
from SR_analysis.utils.feature_builder import apply_G_alpha as feature_apply_G_alpha
from SR_analysis.utils.feature_builder import compute_dimensionless
from SR_analysis.utils.g_operator import (
    apply_G_alpha,
    apply_G_raw12,
    apply_G_target_force,
)


def test_project_paths_and_karman_contract() -> None:
    assert PROJECT_ROOT == os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
    assert MODEL_DIR == os.path.join(PROJECT_ROOT, "models")
    assert LEGACY_CFG_DIR == os.path.join(PROJECT_ROOT, "configs", "legacy_configs")

    cfg = get_scene("karman_re100")
    assert cfg["re_d"] == 50.0
    assert cfg["control_time_scale_steps"] == 2000
    assert cfg["control_dt"] == control_dt(800) == 0.4
    assert cfg["fifo_len"] == 150
    assert cfg["conv_len"] == 30
    assert cfg["body_layout"] == ("front", "upper", "lower")
    assert cfg["force_layout"] == (
        "front_fx", "front_fy", "upper_fx", "upper_fy", "lower_fx", "lower_fy"
    )
    assert cfg["fifo_init_action"] == cfg["action_bias"]
    assert cfg["policy_init_action"] == cfg["action_bias"]


def test_illusion_target_contract_and_slices() -> None:
    cfg = get_scene("illusion_1L")
    target = np.arange(16).reshape(2, 8)

    assert cfg["conv_len"] == 36
    assert cfg["target_center_x"] == 20.0
    assert cfg["target_radius"] == cfg["target_diameter"] == 1.0
    assert cfg["target_radius_lattice"] == 20.0
    assert cfg["fifo_init_action"] == (0.0, -1.0, 1.0)
    assert cfg["policy_init_action"] == (0.0, -1.0, 1.0)
    assert cfg["target_force_slice"] == ILLUSION_TARGET_FORCE_SLICE == (0, 2)
    assert cfg["target_sensor_slice"] == ILLUSION_TARGET_SENSOR_SLICE == (2, 8)
    assert np.array_equal(target[:, slice(*cfg["target_force_slice"])], target[:, :2])
    assert np.array_equal(target[:, slice(*cfg["target_sensor_slice"])], target[:, 2:8])


@pytest.mark.parametrize(
    ("operator", "values"),
    [
        (apply_G_raw12, np.arange(24, dtype=np.float32).reshape(2, 12)),
        (apply_G_alpha, np.arange(6, dtype=np.float32).reshape(2, 3)),
        (apply_G_target_force, np.arange(4, dtype=np.float32).reshape(2, 2)),
    ],
)
def test_g_is_an_involution(operator, values) -> None:
    transformed = operator(values)
    assert transformed.dtype == values.dtype
    np.testing.assert_array_equal(operator(transformed), values)


def test_g_uses_native_front_upper_lower_order() -> None:
    raw = np.arange(12, dtype=np.float64)
    np.testing.assert_array_equal(
        apply_G_raw12(raw),
        [4, -5, 2, -3, 0, -1, 6, -7, 10, -11, 8, -9],
    )
    np.testing.assert_array_equal(apply_G_alpha(np.array([1.0, 2.0, 3.0])), [-1.0, -3.0, -2.0])
    np.testing.assert_array_equal(apply_G_target_force(np.array([4.0, 5.0])), [4.0, -5.0])


def test_feature_builder_reexports_canonical_batch_action_g() -> None:
    actions = np.arange(6, dtype=np.float32).reshape(2, 3)
    transformed = feature_apply_G_alpha(actions)
    np.testing.assert_array_equal(transformed, apply_G_alpha(actions))
    assert transformed.dtype == actions.dtype


def test_order_ledger_and_impulse_slots_follow_front_upper_lower() -> None:
    karman = expected_ledger("karman_re100")
    illusion = expected_ledger("illusion_1L")
    assert karman["body_object_ids"] == {"front": 4, "upper": 5, "lower": 6}
    assert illusion["body_object_ids"] == {"front": 3, "upper": 4, "lower": 5}
    commands = impulse_commands(7, 0.125)
    np.testing.assert_array_equal(np.flatnonzero(commands["front"]), [4])
    np.testing.assert_array_equal(np.flatnonzero(commands["upper"]), [5])
    np.testing.assert_array_equal(np.flatnonzero(commands["lower"]), [6])


def test_dimensionless_force_columns_follow_front_upper_lower_order() -> None:
    sensors = np.zeros((1, 6), dtype=np.float64)
    forces = np.arange(1, 7, dtype=np.float64).reshape(1, 6)
    dim = compute_dimensionless(sensors, forces, u0=1.0, d=2.0, rho=1.0)

    assert dim["Cd_F"][0] == 1.0
    assert dim["Cl_F"][0] == 2.0
    assert dim["Cd_T"][0] == 3.0
    assert dim["Cl_T"][0] == 4.0
    assert dim["Cd_B"][0] == 5.0
    assert dim["Cl_B"][0] == 6.0


def test_legacy_dtw_fixture_and_scene_wrappers() -> None:
    assert LEGACY_DTW_VERSION == "legacy_dtw_v1_abs_n_unclipped"
    assert calc_dtw_sim(np.zeros(2), np.full(2, 3.0)) == -2.0

    conv_len = 2
    sensors = np.arange(24, dtype=np.float64).reshape(4, 6) / 10.0
    state = sensors[2:4].copy()
    illusion_target = np.column_stack((np.full(4, 99.0), np.full(4, -99.0), sensors))

    assert compute_similarity_karman_legacy_v1(sensors, state, conv_len) == pytest.approx(1.0)
    assert compute_similarity(sensors, state, conv_len) == pytest.approx(1.0)
    assert compute_similarity_illusion_legacy_v1(illusion_target, state, conv_len) == pytest.approx(1.0)
