"""Regression checks for the capture-only released 1.2 pressure export."""
import hashlib
import json
import sys
import tempfile
import unittest
from pathlib import Path

import numpy as np
import torch

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "models/trackformer_1_2_field"))
from predict import export_pressure_fields
from plot_pressure import render_pressure_map


class PressureFieldExportTest(unittest.TestCase):
    def setUp(self):
        self.contract = {"normalization": {"std": [40], "mean": [1000]}}
        lat, lon = torch.meshgrid(torch.linspace(60, 0, 25), torch.linspace(100, 180, 33), indexing="ij")
        rlat, rlon = torch.meshgrid(torch.linspace(30, 10, 121), torch.linspace(120, 140, 121), indexing="ij")
        static = lambda y, x: torch.stack((y/90, x/180-1, torch.zeros_like(y), torch.zeros_like(y)))[None]
        self.inputs = {"global_static": static(lat, lon), "regional_static": static(rlat, rlon),
                       "center": torch.tensor([[20., 130.]])}
        self.outputs = []
        for i in range(20):
            y, x = torch.meshgrid(torch.linspace(25, 15, 65), torch.linspace(125+i*.5, 135+i*.5, 65), indexing="ij")
            core = -torch.exp(-((x-(130+i*.5))**2+(y-20)**2)/3)[None, None]
            mask = torch.ones_like(core, dtype=torch.bool)
            mask[:, :, :, 0] = False
            self.outputs.append({"core": core, "core_lat": y[None], "core_lon": x[None], "core_valid": mask,
                "regional_valid": torch.ones((1,1,121,121),dtype=torch.bool), "track_valid": torch.tensor([True]),
                "pressure": torch.tensor([800.])})
        self.issue = 1536624000 * 10**9

    def test_core_units_geography_and_every_valid_time(self):
        result = export_pressure_fields(self.outputs, self.inputs, self.contract, self.issue)
        self.assertEqual(result["core_mslp_hpa"].shape, (20,65,65))
        self.assertEqual(result["core_latitude_deg"].shape, (20,65,65))
        self.assertEqual(result["core_longitude_deg"].shape, (20,65,65))
        self.assertEqual(result["core_valid"].dtype, np.dtype(bool))
        self.assertEqual(result["regional_valid"].shape, (20,121,121))
        np.testing.assert_array_equal(result["valid_time_ns"], self.issue+np.arange(6,121,6,dtype=np.int64)*3600*10**9)
        self.assertEqual(int(result["member_count"]), 1)
        self.assertEqual(float(result["core_information_spacing_km"]), 20)
        self.assertAlmostEqual(float(result["core_mslp_hpa"][0,32,32]),960)
        np.testing.assert_array_equal(result["core_longitude_deg"][19], self.outputs[19]["core_lon"][0].numpy())
        self.assertFalse(result["core_valid"][:,:,0].any())

    def test_capture_preserves_tensors_and_never_inserts_scalar_pressure(self):
        before = [{k:v.clone() for k,v in o.items()} for o in self.outputs]
        result = export_pressure_fields(self.outputs, self.inputs, self.contract, self.issue)
        self.assertGreater(float(result["core_mslp_hpa"].min()), float(self.outputs[0]["pressure"][0]))
        for old, new in zip(before, self.outputs):
            for key in old:
                self.assertTrue(torch.equal(old[key], new[key]), key)

    def test_a_moving_core_can_leave_the_original_fixed_patch(self):
        result = export_pressure_fields(self.outputs, self.inputs, self.contract, self.issue)
        fixed_east = float(result["regional_longitude_deg"].max())
        self.assertGreater(float(result["core_longitude_deg"][-1].max()), fixed_east)
        np.testing.assert_array_equal(result["regional_longitude_deg"], (self.inputs["regional_static"][0,1].numpy()+1)*180)

    def test_incomplete_rollout_is_rejected(self):
        with self.assertRaisesRegex(ValueError, "twenty"):
            export_pressure_fields(self.outputs[:-1], self.inputs, self.contract, self.issue)

    def plot_data(self):
        result = export_pressure_fields(self.outputs, self.inputs, self.contract, self.issue)
        result.update(lead_hours=np.arange(6,121,6), basin_mslp_hpa=np.stack([
            1005+self.inputs["global_static"][0,0].numpy()*10+i*.1 for i in range(20)]),
            central_pressure_hpa=np.full(20,960),
            track_lat_lon=np.column_stack((np.full(20,20),130+np.arange(20)*.5)))
        return result

    def test_renderer_accepts_real_geography_and_distinct_leads(self):
        with tempfile.TemporaryDirectory(prefix="pressure-export-test-") as folder:
            first, last = Path(folder)/"first.png", Path(folder)/"last.png"
            data = self.plot_data()
            render_pressure_map(data, first, 6, 4)
            render_pressure_map(data, last, 120, 4)
            self.assertGreater(first.stat().st_size, 10000)
            self.assertNotEqual(hashlib.sha256(first.read_bytes()).digest(), hashlib.sha256(last.read_bytes()).digest())

    def test_missing_core_is_masked_not_filled_from_scalar(self):
        data = self.plot_data()
        data["core_valid"][:] = False
        original = data["core_mslp_hpa"].copy()
        with tempfile.TemporaryDirectory(prefix="pressure-export-test-") as folder:
            render_pressure_map(data, Path(folder)/"masked.png", 120)
        np.testing.assert_array_equal(data["core_mslp_hpa"], original)

    def test_invalid_contour_interval_and_unregistered_ensemble_are_rejected(self):
        for interval in (0,-1,np.nan):
            with self.assertRaisesRegex(ValueError, "interval"):
                render_pressure_map(self.plot_data(), "unused.png", interval=interval)
        data = self.plot_data()
        data["member_count"] = np.array(50)
        with self.assertRaisesRegex(ValueError, "one-member"):
            render_pressure_map(data, "unused.png")

    def test_released_neural_source_hashes_are_unchanged(self):
        directory = ROOT/"models/trackformer_1_2_field"
        manifest = json.loads((directory/"manifest.json").read_text())
        for name, expected in manifest["source_module_sha256"].items():
            self.assertEqual(hashlib.sha256((directory/name).read_bytes()).hexdigest(), expected, name)


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