from __future__ import annotations

import importlib.util
import json
import tempfile
import unittest
from pathlib import Path

import numpy as np


REPO_ROOT = Path(__file__).resolve().parents[1]


def load_module(name: str, relative_path: str):
    spec = importlib.util.spec_from_file_location(name, REPO_ROOT / relative_path)
    assert spec is not None and spec.loader is not None
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


evaluate = load_module("evaluate_voc2012_segmentation", "scripts/stages/evaluate_voc2012_segmentation.py")
validate = load_module("validate_voc2012_segmentation_quality", "scripts/stages/validate_voc2012_segmentation_quality.py")


class Voc2012QualityTests(unittest.TestCase):
    def test_official_top_left_padding_and_ignore_label(self) -> None:
        image = np.arange(2 * 3 * 3, dtype=np.uint8).reshape(2, 3, 3)
        label = np.asarray([[0, 1, 2], [3, 4, 5]], dtype=np.uint8)
        padded_image, padded_label = evaluate.pad_common_input(image, label, 5, 6, 128, 255)
        self.assertEqual(padded_image.shape, (1, 5, 6, 3))
        self.assertEqual(padded_label.shape, (5, 6))
        np.testing.assert_array_equal(padded_image[0, :2, :3], image)
        np.testing.assert_array_equal(padded_label[:2, :3], label)
        self.assertTrue(np.all(padded_image[0, 2:, :] == 128))
        self.assertTrue(np.all(padded_image[0, :, 3:] == 128))
        self.assertTrue(np.all(padded_label[2:, :] == 255))
        self.assertTrue(np.all(padded_label[:, 3:] == 255))

    def test_confusion_excludes_ignore_255(self) -> None:
        label = np.asarray([[0, 1, 255], [1, 2, 255]], dtype=np.uint8)
        prediction = np.asarray([[[0, 2, 1], [1, 2, 0]]], dtype=np.int64)
        histogram, valid_pixels = evaluate.confusion_matrix(label, prediction, 3, 255)
        self.assertEqual(valid_pixels, 4)
        expected = np.asarray([[1, 0, 0], [0, 1, 1], [0, 0, 1]], dtype=np.int64)
        np.testing.assert_array_equal(histogram, expected)

    def test_metric_matches_independent_recompute(self) -> None:
        histogram = np.asarray([[4, 1, 0], [1, 3, 0], [0, 1, 2]], dtype=np.int64)
        first = evaluate.metrics_from_confusion(histogram)
        second = validate.recompute_metrics(histogram)
        self.assertAlmostEqual(first["miou"], second["miou"], places=15)
        self.assertAlmostEqual(first["pixel_accuracy"], second["pixel_accuracy"], places=15)
        self.assertEqual(first["valid_pixels"], second["valid_pixels"])
        self.assertEqual(first["valid_class_count"], second["valid_class_count"])

    def test_checkpoint_rejects_changed_run_fingerprint(self) -> None:
        with tempfile.TemporaryDirectory() as temporary:
            path = Path(temporary) / "checkpoint.jsonl"
            path.write_text(json.dumps({"run_fingerprint": "old", "image_id": "a"}) + "\n")
            with self.assertRaisesRegex(ValueError, "fingerprint mismatch"):
                evaluate.load_checkpoint(path, "new")

    def test_common_report_contract_is_literal_in_evaluator(self) -> None:
        source = (REPO_ROOT / "scripts/stages/evaluate_voc2012_segmentation.py").read_text()
        for token in [
            '"acceptance_status": "MEASURED_NO_ACCEPTANCE_THRESHOLD"',
            '"metric_name": "mean_iou"',
            '"threshold": None',
            '"fp32"',
            '"public_int8"',
            '"sample_count"',
        ]:
            self.assertIn(token, source)

    def test_metadata_finalizer_preserves_metrics_and_adds_formats(self) -> None:
        finalizer = load_module(
            "finalize_voc2012_quality_summary",
            "scripts/stages/finalize_voc2012_quality_summary.py",
        )
        with tempfile.TemporaryDirectory() as temporary:
            path = Path(temporary) / "quality_summary.json"
            path.write_text(
                json.dumps(
                    {
                        "model_id": "SGXX",
                        "quality": {
                            "fp32": {"miou": 0.75, "sample_count": 1449},
                            "public_int8": {"miou": 0.74, "sample_count": 1449},
                        },
                    }
                )
            )
            original_argv = __import__("sys").argv
            try:
                __import__("sys").argv = ["finalize", "--summary", str(path)]
                self.assertEqual(finalizer.main(), 0)
            finally:
                __import__("sys").argv = original_argv
            value = json.loads(path.read_text())
            self.assertEqual(value["quality"]["fp32"]["format"], "tensorflow_graphdef")
            self.assertEqual(value["quality"]["public_int8"]["format"], "tflite")
            self.assertTrue(value["pair_comparison"]["same_semantic_input"])
            self.assertEqual(value["quality"]["fp32"]["miou"], 0.75)


if __name__ == "__main__":
    unittest.main()