#!/usr/bin/env python3
"""Small schema/plot controls; these are not stellar-equilibrium validation."""
import csv
from pathlib import Path
import tempfile
import unittest

import matplotlib.pyplot as plt
import numpy as np

import plot_model_slices as slices


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


def fixture(directory):
    # Two disconnected material triangles leave a genuine empty gap. Their
    # enthalpy is independent of density and one density sample is negative.
    coordinates = [(-2, 0), (-1, 0), (-1.5, 1), (1, 0), (2, 0), (1.5, 1)]
    rho, h = [-.4, .2, 1, .5, 1, 2], [1, 2, 3, 4, 5, 6]
    vertices = [[i, i // 3, u, .2, v, rho[i], h[i], 0] for i, (u, v) in enumerate(coordinates)]
    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"], [[0, 0, 0, 1, 2], [1, 1, 3, 4, 5]])
    segments = []
    for element in range(2):
        for i, j in ((0, 1), (1, 2), (2, 0)):
            a, b = coordinates[3 * element + i], coordinates[3 * element + j]
            segments.append([a[0], .2, a[1], b[0], .2, b[1], 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=.2,
                    root_tolerance=1e-11, deformation_scale=1, vertex_count=6, triangle_count=2,
                    surface_segments=6, reference_surface_segments=6, surface_coplanar_triangles=0,
                    reference_surface_coplanar_triangles=0, subdivisions=6, u_min=-2, u_max=2, v_min=0, v_max=1)
    write_csv(directory / "metadata.csv", ["metric", "value"], metadata.items())


class ModelSlicePlotTests(unittest.TestCase):
    def setUp(self):
        self.temporary = tempfile.TemporaryDirectory()
        self.directory = Path(self.temporary.name)
        fixture(self.directory)

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

    def test_signed_independent_fields_are_preserved(self):
        data = slices.load_slice(self.directory)
        self.assertEqual(data.vertices[0, 5], -.4)
        self.assertEqual(data.vertices[0, 6], 1.)
        norm, _, _ = slices.color_scale(data.vertices[:, 5])
        self.assertEqual(norm(0.), .5)
        self.assertEqual(norm(-.4), 0.)

    def test_no_triangulation_across_vacuum(self):
        data = slices.load_slice(self.directory)
        np.testing.assert_array_equal(data.triangles, [[0, 1, 2], [3, 4, 5]])
        for triangle in data.triangles:
            horizontal = data.vertices[triangle, 2]
            self.assertTrue(np.all(horizontal < 0) or np.all(horizontal > 0))
        raster = slices.rasterize_fields(data, (-2, 2, 0, 1), 64)
        self.assertTrue(raster.mask[:, :, 24:40].all())
        figure = slices.render(data, reference_outline=True, title="Analytic geometry/schema control")
        self.assertEqual(figure.axes[0].get_aspect(), 1.)
        self.assertEqual(figure.axes[1].get_aspect(), 1.)
        for suffix in ("png", "pdf"):
            path = self.directory / ("control." + suffix)
            figure.savefig(path)
            self.assertGreater(path.stat().st_size, 1000)

    def test_cross_element_triangle_is_rejected(self):
        write_csv(self.directory / "triangles.csv", ["triangle", "element", "v0", "v1", "v2"], [[0, 0, 0, 1, 4], [1, 1, 3, 4, 5]])
        with self.assertRaisesRegex(ValueError, "different material elements"):
            slices.load_slice(self.directory)

    def test_raster_overlap_uses_first_triangle_without_trace_averaging(self):
        data = slices.load_slice(self.directory)
        original = slices.rasterize_fields(data, (-2, 2, 0, 1), 64)
        extra = data.vertices[:3].copy()
        extra[:, 0] += len(data.vertices)
        extra[:, 1] = 2
        extra[:, 5] = 100
        extra[:, 6] = -100
        data.vertices = np.concatenate((data.vertices, extra))
        data.triangles = np.concatenate((data.triangles, [[6, 7, 8]]))
        overlapped = slices.rasterize_fields(data, (-2, 2, 0, 1), 64)
        np.testing.assert_array_equal(original.mask, overlapped.mask)
        np.testing.assert_array_equal(original.data, overlapped.data)

    def test_wrong_physical_plane_is_rejected(self):
        path = self.directory / "vertices.csv"
        text = path.read_text().replace("0,0,-2,0.2,0,", "0,0,-2,0.25,0,")
        path.write_text(text)
        with self.assertRaisesRegex(ValueError, "physical plane"):
            slices.load_slice(self.directory)

    def test_wrong_reference_plane_is_rejected(self):
        path = self.directory / "reference_surface.csv"
        values = slices.table(path, ["x0", "y0", "z0", "x1", "y1", "z1", "maximum_plane_error"])
        values[0, 1] = 0.
        write_csv(path, ["x0", "y0", "z0", "x1", "y1", "z1", "maximum_plane_error"], values)
        with self.assertRaisesRegex(ValueError, "outline"):
            slices.load_slice(self.directory)

    def test_shape_exaggeration_is_rejected(self):
        path = self.directory / "metadata.csv"
        path.write_text(path.read_text().replace("deformation_scale,1", "deformation_scale,10"))
        with self.assertRaisesRegex(ValueError, "scale exactly one"):
            slices.load_slice(self.directory)


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