import csv
from pathlib import Path

import numpy as np

from SR_analysis import stage_3_validate as stage3
from SR_analysis.tools.telemetry_to_csv import (
    canonical_arrays,
    export,
    legacy_similarity,
    time_coordinate,
)
from SR_analysis.utils.metrics import legacy_cycle_similarity


def test_stage1_prefers_post_state_and_physical_actions():
    data = {
        "sensors": np.zeros((2, 6)),
        "sensors_post": np.ones((2, 6)),
        "forces_post": np.full((2, 6), 2.0),
        "actions_alpha_commanded": np.full((2, 3), 3.0),
    }
    arrays = canonical_arrays(data, "ppo")
    assert np.all(arrays["sensors"] == 1.0)
    assert np.all(arrays["forces"] == 2.0)
    assert np.all(arrays["actions_alpha"] == 3.0)


def test_stage3_ppo_schema_preserves_physical_actions():
    data = {
        "sensors": np.ones((2, 6)),
        "forces": np.ones((2, 6)) * 2,
        "actions_alpha": np.ones((2, 3)) * 3,
        "actions_omega": np.ones((2, 3)) * 0.03,
    }
    arrays = canonical_arrays(data, "ppo")
    assert np.all(arrays["actions_alpha"] == 3.0)
    assert np.all(arrays["actions_omega"] == 0.03)


def test_time_coordinate_uses_d_over_u0():
    np.testing.assert_allclose(time_coordinate(3, 800), [0.0, 0.4, 0.8])


def test_csv_and_stage3_share_exact_legacy_metric():
    target = np.arange(48, dtype=np.float64).reshape(8, 6) / 10.0
    state = target[4:].copy()
    canonical = legacy_cycle_similarity(target, state, 4)
    stage3_metric = stage3.compute_versioned_metrics(target, state, 4)[
        "legacy_reference_cycle_vs_last_recorded_cycle"
    ]
    csv_metric = legacy_similarity(target, state, 4)

    assert stage3_metric == canonical
    assert csv_metric["similarity"] == canonical["similarity"]
    assert csv_metric["lag"] == canonical["lag"]
    assert csv_metric["per_channel"] == [
        canonical["per_channel"][str(index)] for index in range(6)
    ]


def test_export_writes_wide_and_long_csv(tmp_path: Path):
    trajectory = tmp_path / "trajectory.npz"
    np.savez(
        trajectory,
        sensors=np.ones((4, 6)),
        forces=np.ones((4, 6)) * 2,
        actions_alpha=np.ones((4, 3)) * 3,
        target_sensors=np.ones((60, 6)),
    )
    output = tmp_path / "csv"
    export(trajectory, output, "case", "sr", 800, 30)
    assert (output / "sr_wide.csv").is_file()
    assert (output / "sr_long.csv").is_file()
    with (output / "sr_wide.csv").open(newline="", encoding="utf-8") as handle:
        rows = list(csv.reader(handle))
    assert len(rows) == 5
    assert "sensors_upper_ux" in rows[0]
