from __future__ import annotations

import sys
from pathlib import Path

import numpy as np

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

from env_karman import compute_similarity
from normalization import (
    CALIBRATION_SCHEMA_NATIVE_V2,
    LEGACY_SENSOR_FACTOR,
    POLICY_SENSOR_UNITS_LEGACY,
    SENSOR_UNITS_NATIVE,
    convert_legacy_calibration,
    convert_legacy_harmonics,
    convert_legacy_target,
    dtw_sensor_factor,
    policy_sensor_scale,
)


def _legacy_calibration():
    return {"SENSOR_CC": 78.0, "SENS_SCALE": 0.78, "dtw_norm_scale": 0.234}


def test_calibration_conversion_and_policy_schema_behavior():
    legacy = _legacy_calibration()
    native = convert_legacy_calibration(legacy)
    assert native["schema_version"] == CALIBRATION_SCHEMA_NATIVE_V2
    assert native["sensor_units"] == SENSOR_UNITS_NATIVE
    assert "SENSOR_CC" not in native
    assert np.isclose(native["SENS_SCALE"], 0.01)
    assert np.isclose(native["dtw_norm_scale"], 0.003)
    assert dtw_sensor_factor(legacy) == LEGACY_SENSOR_FACTOR
    assert dtw_sensor_factor(native) == 1.0
    assert np.isclose(policy_sensor_scale(native), 0.01)
    assert np.isclose(policy_sensor_scale(native, POLICY_SENSOR_UNITS_LEGACY), 0.78)
    assert np.isclose(policy_sensor_scale(legacy), 0.78)


def test_target_and_harmonics_conversion_only_scale_sensor_channels():
    target = np.arange(24, dtype=np.float32).reshape(3, 8)
    converted = convert_legacy_target(target)
    np.testing.assert_allclose(converted[:, :6], target[:, :6] / 78.0)
    np.testing.assert_allclose(converted[:, 6:], target[:, 6:])

    harmonics = [
        {"dc": float(i + 1), "amps": [float(i + 2)], "freqs": [0.1], "phases": [0.2]}
        for i in range(8)
    ]
    converted_h = convert_legacy_harmonics(harmonics)
    for i in range(6):
        assert np.isclose(converted_h[i]["dc"], harmonics[i]["dc"] / 78.0)
        assert np.isclose(converted_h[i]["amps"][0], harmonics[i]["amps"][0] / 78.0)
    assert converted_h[6:] == harmonics[6:]


def test_dtw_similarity_is_invariant_under_unit_conversion():
    rng = np.random.default_rng(7)
    target = rng.normal(size=(150, 6))
    state = rng.normal(size=(90, 6))
    legacy_scale = 0.25
    legacy_sim = compute_similarity(target, state, 30, legacy_scale)
    native_sim = compute_similarity(
        target / LEGACY_SENSOR_FACTOR,
        state / LEGACY_SENSOR_FACTOR,
        30,
        legacy_scale / LEGACY_SENSOR_FACTOR,
    )
    assert np.isclose(native_sim, legacy_sim, atol=1e-12)
