"""The sar3d processor on the real CUDA backend: notebook reference and contracts.

Every case runs build/native_tests/sar3d_fixture (SAR3D_FIXTURE overrides it).
Without a usable GPU the processor publishes sar3d_status 3: then only that
warning is checked and the remaining tests are skipped. The notebook is needed
only to regenerate sar3d_reference.npz with generate_sar3d_reference.py.
"""
from __future__ import annotations

import copy
import json
import os
from pathlib import Path
import subprocess
import unittest

import numpy as np

ROOT = Path(__file__).resolve().parents[2]
FIXTURE = Path(os.environ.get("SAR3D_FIXTURE", ROOT / "build/native_tests/sar3d_fixture"))
REFERENCE = Path(__file__).with_name("sar3d_reference.npz")
VIEWS = ("sar3d_front", "sar3d_side", "sar3d_top", "sar3d_oblique")


def run(request: dict) -> dict[str, dict]:
    output = subprocess.run([str(FIXTURE)], input=json.dumps(request), text=True,
                            capture_output=True, timeout=30, check=True)
    return {item["name"]: item for item in json.loads(output.stdout)["payloads"]}


def status(payloads: dict[str, dict]) -> tuple[int, str]:
    table = payloads["sar3d_status"]["table"]
    return int(table[0]), bytes(map(int, table[1:])).decode("utf-8")


def reference_cases(archive) -> list[dict]:
    """Requests of the reference archive, one per notebook parameter case."""
    cases = []
    for case in json.loads(archive["cases"].item()):
        request = json.loads(archive["request"].item())
        request.update(case["settings"])
        for member, windows in zip([*request["previous"], request["latest"]], case["windows"], strict=True):
            for trace, (start, end) in zip(member["traces"], windows, strict=True):
                trace.update(start=start, end=end)
        cases.append(request)
    return cases


class Sar3dTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls) -> None:
        with np.load(REFERENCE, allow_pickle=False) as archive:
            cls.reference = {key: archive[key] for key in archive.files}
        cls.cases = reference_cases(cls.reference)
        cls.scene_payloads = run(cls.cases[0])
        code, message = status(cls.scene_payloads)
        cls.unavailable = message if code == 3 else None

    def scene(self, **settings) -> dict:
        if self.unavailable:
            self.skipTest(self.unavailable)
        request = copy.deepcopy(self.cases[0])
        request.update(settings)
        return request

    def test_unavailable_gpu_is_a_visible_warning(self) -> None:
        code, message = status(self.scene_payloads)
        if self.unavailable:
            self.assertIn("CUDA", message)
            self.assertFalse(any(item["image"] for item in self.scene_payloads.values()))
        else:
            self.assertEqual(code, 0, message)

    def test_cuda_matches_notebook(self) -> None:
        if self.unavailable:
            self.skipTest(self.unavailable)
        for index, request in enumerate(self.cases):
            with self.subTest(case=index):
                payloads = run(request)
                self.assertEqual(status(payloads)[0], 0, status(payloads))
                for view, horizontal, vertical in (("front", "x", "y"), ("side", "z", "y"), ("top", "x", "z")):
                    payload = payloads[f"sar3d_{view}"]
                    expected = self.reference[f"case_{index}_{view}"]
                    np.testing.assert_allclose(payload["x"], self.reference[f"case_{index}_{horizontal}"],
                                               rtol=2e-6, atol=1e-6)
                    np.testing.assert_allclose(payload["y"], self.reference[f"case_{index}_{vertical}"],
                                               rtol=2e-6, atol=1e-6)
                    np.testing.assert_allclose(np.asarray(payload["image"]).reshape(expected.shape),
                                               expected, rtol=2e-3, atol=3e-4)
                peak = payloads["sar3d_peak"]["table"]
                expected_peak = self.reference[f"case_{index}_peak"]
                np.testing.assert_allclose(peak[1:], expected_peak[1:], rtol=0, atol=2e-6)
                np.testing.assert_allclose(peak[0], expected_peak[0], rtol=2e-3, atol=0)

    def test_invalid_input_is_an_error_with_a_reason(self) -> None:
        def no_traces(request): request["latest"]["traces"] = []
        def empty_window(request): request["latest"]["traces"][0]["end"] = request["latest"]["traces"][0]["start"]
        def falling_frequency(request): request["latest"]["traces"][0]["frequency"].reverse()
        for name, settings, damage in (
            ("no traces", {}, no_traces),
            ("empty capture window", {}, empty_window),
            ("falling frequency", {}, falling_frequency),
            ("grid beyond 32-bit indexing", {"voxel": 1e-8}, None),
            ("vertical look angle", {"pitch": 90}, None),
            ("result larger than the ring slot", {"slot": 1024}, None),
        ):
            with self.subTest(name):
                request = self.scene(**settings)
                if damage:
                    damage(request)
                code, message = status(run(request))
                self.assertEqual(code, 2, message)
                self.assertTrue(message)

    def test_empty_signal_is_reported_without_a_peak(self) -> None:
        request = self.scene(previous=[], radius=0)
        for trace in request["latest"]["traces"]:
            trace["s21"] = [[0, 0]] * len(trace["frequency"])
        payloads = run(request)
        self.assertEqual(status(payloads)[0], 1)
        self.assertNotIn("sar3d_peak", payloads)
        for name in VIEWS:
            self.assertFalse(np.any(payloads[name]["image"]))

    def test_weak_signal_stays_finite(self) -> None:
        request = self.scene(previous=[], radius=0, normalize=False)
        for trace in request["latest"]["traces"]:
            trace["s21"] = (np.asarray(trace["s21"]) * 1e-38).tolist()
        payloads = run(request)
        self.assertEqual(status(payloads)[0], 0, status(payloads))
        for name in VIEWS:
            values = np.asarray(payloads[name]["image"])
            self.assertTrue(np.isfinite(values).all())
            self.assertGreater(values.max(), 0)
            self.assertLess(values.max(), 1e-10)


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