from __future__ import annotations

import importlib.util
import json
from pathlib import Path

import numpy as np

MODULE_PATH = Path(__file__).resolve().parents[1] / "tools/export_flow_comparison.py"
spec = importlib.util.spec_from_file_location("export_flow_comparison", MODULE_PATH)
module = importlib.util.module_from_spec(spec)
assert spec.loader
spec.loader.exec_module(module)


def test_wrapped_phase_difference_crosses_branch_cut():
    delta = module.wrapped_phase_difference(-np.pi + 0.02, np.pi - 0.03)
    assert np.isclose(delta, 0.05)


def test_joint_phase_selection_is_deterministic_and_nearby():
    n = 20
    phase = np.linspace(0, 4 * np.pi, n, endpoint=False)
    def trace(offset: float) -> np.ndarray:
        out = np.zeros((n, 6))
        out[:, 2] = 2.0 + 3.0 * np.cos(phase + offset)
        out[:, 3] = -1.0 + 0.5 * np.sin(phase + offset)
        return out
    first = module.select_joint_phase_match(trace(0), trace(0.05), trace(-0.04), 0, n, temporal_weight=0.001)
    second = module.select_joint_phase_match(trace(0), trace(0.05), trace(-0.04), 0, n, temporal_weight=0.001)
    assert first == second
    assert len(set(first["indices"].values())) == 1
    assert np.isclose(first["wrapped_angle_errors_rad"]["PPO"], 0.05)
    assert np.isclose(first["wrapped_angle_errors_rad"]["SR"], 0.04)
    assert first["objective_terms"]["temporal_distance_samples"] == 0


def test_exact_candidate_selection_uses_trace_index_mapping():
    candidates = np.empty((4, 2, 3), dtype=np.float32)
    for candidate_index, trace_index in enumerate(range(96, 100)):
        candidates[candidate_index].fill(trace_index)
    field, candidate_index = module.select_exact_candidate_field(
        candidates, np.arange(96, 100), 98, expected_start=96, expected_stop=100
    )
    assert candidate_index == 2
    np.testing.assert_array_equal(field, np.full((2, 3), 98, dtype=np.float32))


def test_exact_candidate_selection_rejects_mapping_mismatch():
    candidates = np.zeros((4, 2, 3), dtype=np.float32)
    mismatched = np.array([96, 97, 99, 100])
    with np.testing.assert_raises_regex(ValueError, "does not exactly cover"):
        module.select_exact_candidate_field(
            candidates, mismatched, 99, expected_start=96, expected_stop=100
        )


def test_exact_candidate_selection_requires_float32():
    candidates = np.zeros((4, 2, 3), dtype=np.float64)
    with np.testing.assert_raises_regex(ValueError, "float32"):
        module.select_exact_candidate_field(
            candidates, np.arange(96, 100), 98, expected_start=96, expected_stop=100
        )


def test_orientation_transposes_xy_to_image_yx():
    omega_xy = np.arange(12).reshape(4, 3)
    oriented = module.orient_vorticity_xy_to_yx(omega_xy, (4, 3))
    assert oriented.shape == (3, 4)
    np.testing.assert_array_equal(oriented[:, 2], omega_xy[2, :])


def test_crop_uses_physical_x_and_centered_y_coordinates():
    field = np.arange(7 * 10).reshape(7, 10)
    cropped, meta = module.crop_yx(field, (1.0, 3.0), (-1.0, 1.0), d_lattice=2.0, center_y_lattice=3.0)
    np.testing.assert_array_equal(cropped, field[1:6, 2:7])
    assert meta["x_slice"] == [2, 7]
    assert meta["y_slice"] == [1, 6]
    assert meta["extent_xD_yD"] == [1.0, 3.0, -1.0, 1.0]


def test_manifest_includes_npz(tmp_path: Path):
    repo = tmp_path
    package = repo / "package"
    package.mkdir()
    (package / "data.npz").write_bytes(b"npz")
    (package / "figure.png").write_bytes(b"png")
    (package / "ignored.txt").write_text("ignored", encoding="utf-8")
    (package / "manifest.json").write_text(json.dumps({"source_policy": "test", "summary": {}}), encoding="utf-8")
    module.update_package_manifest(package, repo)
    manifest = json.loads((package / "manifest.json").read_text(encoding="utf-8"))
    paths = {entry["path"] for entry in manifest["artifacts"]}
    assert "package/data.npz" in paths
    assert "package/figure.png" in paths
    assert "package/ignored.txt" not in paths
    assert manifest["summary"]["flow_field_npz"] == 1
