#!/usr/bin/env python3
"""Affine same-state controls for physical-slice refinement summaries."""
import csv
from pathlib import Path
import tempfile
import unittest

import numpy as np

import compare_model_slices as compare
from plot_model_slices import load_slice, plt, rasterize_fields


def write_csv(path, header, rows):
    with path.open("w", newline="") as stream:
        writer = csv.writer(stream)
        writer.writerow(header)
        writer.writerows(rows)


def square(directory, subdivisions, offset=.2):
    directory.mkdir()
    vertices, triangles = [], []
    n = subdivisions
    for j in range(n):
        for i in range(n):
            corners = [(i / n, j / n), ((i + 1) / n, j / n), (i / n, (j + 1) / n), ((i + 1) / n, (j + 1) / n)]
            for indices in ((0, 1, 3), (0, 3, 2)):
                start = len(vertices)
                triangles.append([len(triangles), 0, start, start + 1, start + 2])
                for corner in indices:
                    x, z = corners[corner]
                    vertices.append([len(vertices), 0, x, offset, z, x - .5, 1 + x + 2 * z, 0])
    write_csv(directory / "vertices.csv", ["vertex", "element", "x", "y", "z", "rho", "h", "plane_error"], vertices)
    write_csv(directory / "triangles.csv", ["triangle", "element", "v0", "v1", "v2"], triangles)
    segments = []
    for i in range(n):
        a, b = i / n, (i + 1) / n
        for x0, z0, x1, z1 in ((a, 0, b, 0), (1, a, 1, b), (b, 1, a, 1), (0, b, 0, a)):
            segments.append([x0, offset, z0, x1, offset, z1, 0])
    for name in ("surface.csv", "reference_surface.csv"):
        write_csv(directory / name, ["x0", "y0", "z0", "x1", "y1", "z1", "maximum_plane_error"], segments)
    metadata = dict(plane="xz", horizontal_axis=0, vertical_axis=2, normal_axis=1, offset=offset,
                    root_tolerance=1e-11, deformation_scale=1, vertex_count=len(vertices), triangle_count=len(triangles),
                    surface_segments=len(segments), reference_surface_segments=len(segments), surface_coplanar_triangles=0,
                    reference_surface_coplanar_triangles=0, subdivisions=n, u_min=0, u_max=1, v_min=0, v_max=1)
    write_csv(directory / "metadata.csv", ["metric", "value"], metadata.items())


class ComparisonTests(unittest.TestCase):
    def setUp(self):
        self.temporary = tempfile.TemporaryDirectory()
        self.root = Path(self.temporary.name)
        square(self.root / "coarse", 2)
        square(self.root / "fine", 4)
        self.coarse = load_slice(self.root / "coarse")
        self.fine = load_slice(self.root / "fine")
        self.source = self.root / "synthetic_state.txt"
        self.source.write_text("rho=x-0.5; h=1+x+2z; unit square; y=0.2\n")

    def tearDown(self):
        plt.close("all")
        self.temporary.cleanup()

    def test_affine_integrals_and_perimeter_do_not_change_under_refinement(self):
        report = compare.comparison(self.coarse, self.fine, self.source)
        for key, expected in (("triangle_area_sum", 1.), ("piecewise_linear_density_integral", 0.),
                              ("piecewise_linear_enthalpy_integral", 2.5)):
            self.assertAlmostEqual(report["changes"][key]["coarse"], expected, places=14)
            self.assertAlmostEqual(report["changes"][key]["fine"], expected, places=14)
            self.assertLess(report["changes"][key]["absolute_change"], 1e-14)
        self.assertEqual(report["changes"]["surface"]["segment_length_sum"]["fine"], 4.)
        self.assertEqual(report["changes"]["surface"]["u_min"]["fine"], 0.)
        self.assertEqual(report["changes"]["surface"]["v_max"]["fine"], 1.)
        self.assertFalse(report["state_provenance"]["export_binding_verified"])
        self.assertEqual(len(report["state_provenance"]["source_sha256"]), 64)

    def test_overlap_is_counted_not_hidden_as_union_area(self):
        self.fine.triangles = np.concatenate((self.fine.triangles, self.fine.triangles))
        metrics = compare.slice_metrics(self.fine)
        self.assertEqual(metrics["triangle_area_sum"], 2.)
        self.assertEqual(metrics["piecewise_linear_enthalpy_integral"], 5.)

    def test_shared_axes_and_signed_color_ranges(self):
        figure = compare.render_comparison(self.coarse, self.fine, reference_outline=True)
        axes = figure.axes[:4]
        for ax in axes:
            self.assertEqual(ax.get_xlim(), axes[0].get_xlim())
            self.assertEqual(ax.get_ylim(), axes[0].get_ylim())
            self.assertEqual(ax.get_aspect(), 1.)
        self.assertIs(axes[0].images[0].norm, axes[1].images[0].norm)
        self.assertIs(axes[2].images[0].norm, axes[3].images[0].norm)
        self.assertEqual(axes[0].images[0].norm.vmin, -.5)
        self.assertEqual(axes[0].images[0].norm.vmax, .5)
        for extension in ("png", "pdf"):
            path = self.root / ("comparison." + extension)
            figure.savefig(path)
            self.assertGreater(path.stat().st_size, 1000)

    def test_affine_scalar_raster_is_refinement_invariant(self):
        a = rasterize_fields(self.coarse, (0, 1, 0, 1), 65)
        b = rasterize_fields(self.fine, (0, 1, 0, 1), 65)
        self.assertFalse(a.mask.any())
        self.assertFalse(b.mask.any())
        np.testing.assert_allclose(a.data, b.data, rtol=0, atol=1e-14)

    def test_different_plane_offset_is_rejected(self):
        square(self.root / "wrong_offset", 4, offset=.3)
        with self.assertRaisesRegex(ValueError, "offsets"):
            compare.comparison(self.coarse, load_slice(self.root / "wrong_offset"), self.source)

    def test_reversed_refinement_is_rejected(self):
        with self.assertRaisesRegex(ValueError, "exceed"):
            compare.comparison(self.fine, self.coarse, self.source)

    def test_relative_zero_denominator_is_explicit(self):
        result = compare.differences({"zero": 0., "negative": -2.}, {"zero": 1., "negative": -3.})
        self.assertIsNone(result["zero"]["relative_change"])
        self.assertEqual(result["negative"]["relative_change"], -.5)
        self.assertEqual(result["negative"]["relative_absolute_change"], .5)


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