import json
import subprocess
import sys
from pathlib import Path

import numpy as np
import pytest

from SR_analysis import stage_1_infer as stage1
from SR_analysis.utils.harmonics import gen_target_states_at, harmonics_document


def illusion_cfg():
    return {
        "scene_id": "illusion",
        "sample_interval": 600,
        "u0": 0.01,
    }


def test_target_extraction_uses_force_first_and_all_sensors():
    canonical, forces, sensors = stage1.extract_illusion_target(np.arange(12, dtype=np.float32))
    np.testing.assert_array_equal(canonical, np.arange(8, dtype=np.float32))
    np.testing.assert_array_equal(forces, [0, 1])
    np.testing.assert_array_equal(sensors, [2, 3, 4, 5, 6, 7])


def test_target_extraction_rejects_short_observation():
    with pytest.raises(ValueError, match="8 channels"):
        stage1.extract_illusion_target(np.arange(7))


def test_phase_index_matches_legacy_first_and_next_observations():
    assert stage1.target_phase_index(0) == 0
    assert stage1.target_phase_index(0, after_action=True) == 1
    assert [stage1.target_phase_index(i) for i in range(4)] == [0, 1, 2, 3]
    with pytest.raises(ValueError):
        stage1.target_phase_index(-1)


def test_harmonics_loader_accepts_legacy_list_and_v2_document():
    legacy = [{"dc": 1.0, "amps": [2.0], "freqs": [0.25], "phases": [0.0]}]
    document = harmonics_document(legacy, ["fx"], schema_version="test-v2")
    assert document["channel_names"] == ["fx"]
    np.testing.assert_allclose(gen_target_states_at(0, legacy), gen_target_states_at(0, document))
    assert gen_target_states_at(np.array([0, 1]), document).shape == (2, 1)


def test_gpu_identity_uses_explicit_visibility_mapping(monkeypatch):
    calls = []

    def fake_run(command, **kwargs):
        calls.append((command, kwargs))
        return subprocess.CompletedProcess(command, 0, stdout="GPU-uuid, Test GPU\n", stderr="")

    monkeypatch.setattr(stage1.subprocess, "run", fake_run)
    identity = stage1.gpu_identity(1, "GPU-first,GPU-second")

    assert identity == {"uuid": "GPU-uuid", "name": "Test GPU"}
    assert "--id=GPU-second" in calls[0][0]
    assert "--id=1" not in calls[0][0]


def test_gpu_identity_never_treats_logical_index_as_physical(monkeypatch):
    calls = []

    def fake_run(command, **kwargs):
        calls.append(command)
        return subprocess.CompletedProcess(
            command, 0, stdout="GPU-a, First\nGPU-b, Second\n", stderr="",
        )

    monkeypatch.setattr(stage1.subprocess, "run", fake_run)

    assert stage1.gpu_identity(1, None) == {"uuid": None, "name": None}
    assert all(not argument.startswith("--id=") for argument in calls[0])


def test_device_manifest_records_environment_without_gpu_hard_failure(monkeypatch):
    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "GPU-only")

    def unavailable(*args, **kwargs):
        raise FileNotFoundError("nvidia-smi")

    monkeypatch.setattr(stage1.subprocess, "run", unavailable)

    assert stage1.cfd_device_manifest(0) == {
        "cfd_device": 0,
        "cfd_device_logical": 0,
        "cuda_visible_devices": "GPU-only",
        "cfd_gpu_uuid": None,
        "cfd_gpu_name": None,
    }


def test_paths_are_repo_rooted_and_do_not_depend_on_cwd(tmp_path, monkeypatch):
    monkeypatch.chdir(tmp_path)
    assert stage1.resolve_output_root(None) == stage1.REPO_ROOT / "src/SR_analysis/data/runs"
    relative = stage1.resolve_output_root("scratch/stage1")
    assert relative == stage1.REPO_ROOT / "scratch/stage1"
    out = stage1.scene_output_dir(relative, "round-one", "illusion_1L", illusion_cfg())
    assert out == stage1.REPO_ROOT / "scratch/stage1/round-one/illusion/illusion_1L"


def test_output_refuses_overwrite(tmp_path):
    out = tmp_path / "run"; out.mkdir(); (out / "result.json").write_text("{}")
    with pytest.raises(FileExistsError, match="Refusing"):
        stage1.prepare_output_dir(out, overwrite=False)
    stage1.prepare_output_dir(out, overwrite=True)
    assert (out / "result.json").exists()


def test_npz_v2_schema_and_legacy_aliases(tmp_path):
    cfg = illusion_cfg()
    controlled = tmp_path / "controlled.npz"
    target = tmp_path / "target.npz"
    sensors = np.ones((2, 6), dtype=np.float32)
    forces = np.full((2, 6), 2, dtype=np.float32)
    actions = np.full((2, 3), 0.5, dtype=np.float32)
    commanded = np.full((2, 3), 0.03, dtype=np.float32)
    target_forces = np.array([[1, 2], [3, 4]], dtype=np.float32)
    stage1.save_controlled_npz(
        controlled, "illusion_1L", cfg, sensors=sensors, forces=forces,
        actions_norm=actions, actions_commanded=commanded,
        target_forces=target_forces, phase_indices=np.array([0, 1]),
    )
    stage1.save_target_npz(target, "illusion_1L", cfg, np.arange(16).reshape(2, 8), illusion=True)

    with np.load(controlled) as data:
        assert data["schema_version"].item() == stage1.SCHEMA_VERSION
        np.testing.assert_array_equal(data["actions"], data["actions_norm"])
        np.testing.assert_array_equal(data["sensors"], data["sensors_post"])
        np.testing.assert_array_equal(data["forces"], data["forces_post"])
        np.testing.assert_array_equal(data["target_forces"], data["target_forces_for_action"])
        np.testing.assert_array_equal(data["target_phase_index"], [0, 1])
        np.testing.assert_allclose(data["actions_omega_commanded"], commanded)
        np.testing.assert_allclose(data["actions_alpha_commanded"], commanded / cfg["u0"])
        assert data["dt_c"].item() == pytest.approx(0.3)
    with np.load(target) as data:
        assert data["target_states"].shape == (2, 8)
        assert data["target_forces"].shape == (2, 2)
        assert data["target_sensors"].shape == (2, 6)
        assert len(data["layout"]) == 8


def test_norm_validation_rejects_bad_shapes_and_zeros():
    with pytest.raises(ValueError):
        stage1.validate_norm({"force_norm_fact": 1, "sens_deviation": [0] * 5, "sens_norm_fact": [1] * 6})
    with pytest.raises(ValueError):
        stage1.validate_norm({"force_norm_fact": 1, "sens_deviation": [0] * 6, "sens_norm_fact": [0] * 6})


def test_cli_rejects_mutually_exclusive_scene_and_group():
    with pytest.raises(SystemExit) as error:
        stage1.build_parser().parse_args(["--scene", "illusion_1L", "--group", "illusion_trained"])
    assert error.value.code == 2


def test_dry_run_does_not_import_legacy_or_create_output(tmp_path):
    script = stage1.Path(stage1.__file__)
    output_root = tmp_path / "runs"
    code = (
        "import runpy,sys; "
        f"sys.argv={[str(script), '--scene', 'illusion_1L', '--run-id', 'dry', '--output-root', str(output_root), '--dry-run']!r}; "
        f"runpy.run_path({str(script)!r}, run_name='__main__')"
    )
    completed = subprocess.run([sys.executable, "-c", code], cwd=tmp_path, text=True, capture_output=True)
    assert completed.returncode == 0, completed.stderr
    assert "LegacyCelerisLab" not in completed.stderr
    assert not output_root.exists()


def test_harmonics_document_is_json_serializable():
    document = harmonics_document([{"dc": 0.0, "amps": [], "freqs": [], "phases": []}], ["fx"], schema_version="v2")
    assert json.loads(json.dumps(document))["schema_version"] == "v2"
