from __future__ import annotations

import unittest

import numpy as np

from utils.section_generator import SectionGenerationOptions, format_station, generate_sections
from utils.surface_elevation_sampler import CallableSurfaceSampler, DtmGridSampler


class SectionGeneratorTests(unittest.TestCase):
    def setUp(self) -> None:
        self.sampler = CallableSurfaceSampler(lambda xy: 100.0 + 0.1 * xy[:, 0] + 0.2 * xy[:, 1])

    def test_stations_are_twenty_meter_interval_plus_endpoint_without_cp(self) -> None:
        route = [[0, 0, 0], [35, 0, 0], [55, 0, 0]]
        result = generate_sections(route, self.sampler)
        self.assertEqual(
            [station["chainage_m"] for station in result["longitudinal"]["stations"]],
            [0.0, 20.0, 40.0, 55.0],
        )
        self.assertEqual([station["kind"] for station in result["longitudinal"]["stations"]], ["bp", "regular", "regular", "ep"])

    def test_cross_axis_is_perpendicular_and_center_matches_station(self) -> None:
        result = generate_sections(
            [[0, 0, 0], [100, 0, 0]],
            self.sampler,
            SectionGenerationOptions(cross_half_width_m=10, cross_sample_interval_m=5),
        )
        station = result["cross_sections"][1]
        tangent = np.asarray(station["frame"]["tangent_xy"])
        left = np.asarray(station["frame"]["left_xy"])
        self.assertAlmostEqual(float(np.dot(tangent, left)), 0.0, places=8)
        center = next(sample for sample in station["samples"] if sample["offset_m"] == 0.0)
        self.assertAlmostEqual(center["x"], station["frame"]["origin"]["x"])
        self.assertAlmostEqual(center["y"], station["frame"]["origin"]["y"])
        self.assertAlmostEqual(center["z"], station["frame"]["origin"]["z"])
        self.assertAlmostEqual(center["elevation_m"], station["center_z"])
        self.assertAlmostEqual(station["center_x"], station["frame"]["origin"]["x"])

    def test_output_contains_stable_webcad_exchange_coordinates(self) -> None:
        result = generate_sections(
            [[1000, 2000, 0], [1010, 2010, 0]],
            self.sampler,
            crs="EPSG:5186",
            source_snapshot={"source_filter": "csf", "surface_method": "dtm", "smooth": True},
        )
        self.assertEqual(result["schema_version"], 1)
        self.assertEqual(result["coordinate_reference"]["crs"], "EPSG:5186")
        self.assertEqual(result["cad_exchange"]["cross_local_axes"], {"x": "offset_m", "y": "elevation_m"})
        self.assertIn("tangent_xy", result["cross_sections"][0]["frame"])
        self.assertEqual(format_station(1020.25), "STA.1+020.250")

    def test_dtm_sampler_marks_queries_crossing_invalid_cells_as_no_data(self) -> None:
        x = np.array([0.0, 1.0, 2.0])
        y = np.array([0.0, 1.0, 2.0])
        xx, yy = np.meshgrid(x, y)
        z = xx + yy
        mask = np.ones_like(z, dtype=bool)
        mask[1, 1] = False
        sampler = DtmGridSampler(x, y, z, mask)
        values, valid = sampler.sample_xy(np.array([[0.25, 0.25], [1.75, 1.75], [3.0, 3.0]]))
        self.assertEqual(valid.tolist(), [False, False, False])
        self.assertTrue(np.isnan(values).all())


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