from __future__ import annotations

import numpy as np
from validation.artifacts import ValidationArtifact, load_artifact, save_artifact
from validation.compare import compare_artifacts


def _artifact(value: float = 1.0) -> ValidationArtifact:
    return ValidationArtifact(
        sample_id=np.asarray(["file.root:0", "file.root:1"]),
        labels=np.asarray([0, 1]),
        folds=np.asarray([0, 1]),
        weights=np.asarray([1.0, -2.0], dtype=np.float32),
        globals=np.asarray([[value, 2.0], [3.0, 4.0]], dtype=np.float32),
        node_features_flat=np.asarray(
            [[1.0, 0.0], [2.0, 1.0], [3.0, 2.0]], dtype=np.float32
        ),
        node_offsets=np.asarray([0, 1, 3]),
        edge_src_flat=np.asarray([0, 1]),
        edge_dst_flat=np.asarray([0, 0]),
        edge_features_flat=np.asarray([[0.1], [0.2]], dtype=np.float32),
        edge_offsets=np.asarray([0, 0, 2]),
        logits=np.asarray([[0.1], [0.2]], dtype=np.float32),
    )


def test_artifact_round_trip_preserves_variable_size_arrays(tmp_path):
    save_artifact(_artifact(), tmp_path / "artifact")
    restored = load_artifact(tmp_path / "artifact")
    assert restored.event_count == 2
    np.testing.assert_array_equal(restored.sample_id, ["file.root:0", "file.root:1"])
    np.testing.assert_array_equal(restored.event_nodes(1), [[2.0, 1.0], [3.0, 2.0]])
    np.testing.assert_array_equal(restored.event_edges(1)[0], [0, 1])


def test_comparator_reports_tolerance_and_exact_mismatches():
    same = compare_artifacts(_artifact(), _artifact(1.0000001))
    assert same["overall_status"] == "PASS"
    changed = _artifact()
    changed.labels[1] = 0
    report = compare_artifacts(_artifact(), changed)
    assert report["overall_status"] == "FAIL"
    assert report["stages"]["labels"]["status"] == "FAIL"


def test_comparator_rejects_duplicate_sample_ids():
    duplicate = _artifact()
    duplicate.sample_id[1] = duplicate.sample_id[0]
    try:
        compare_artifacts(_artifact(), duplicate)
    except ValueError as error:
        assert "duplicates" in str(error)
    else:
        raise AssertionError("duplicate IDs must fail loudly")