import importlib.util
import json
import tempfile
import unittest
from pathlib import Path

import numpy as np


ROOT = Path(__file__).parents[2]
SPEC = importlib.util.spec_from_file_location("cad_graphs", ROOT / "전처리" / "cad_graphs.py")
cg = importlib.util.module_from_spec(SPEC); SPEC.loader.exec_module(cg)


class GraphGeometryTests(unittest.TestCase):
    def test_topological_edges_are_bidirectional_with_endpoint_aligned_features(self):
        # Two coplanar triangles: the diagonal must have dihedral==0 in both
        # directions and attributes must match their reverse except direction.
        v = np.array([[0, 0, 0], [1, 0, 0], [1, 1, 0], [0, 1, 0]], float)
        f = np.array([[0, 1, 2], [0, 2, 3]], np.int32)
        topo = cg._topology(v, f, np.tile([0., 0., 1.], (4, 1)))
        ei, ea = topo["edge_index"], topo["edge_attr"]
        self.assertEqual(ei.shape[1], 10)  # 5 undirected mesh edges, not KNN links
        lookup = {(int(a), int(b)): i for i, (a, b) in enumerate(ei.T)}
        for (a, b), i in lookup.items():
            j = lookup[(b, a)]
            self.assertAlmostEqual(float(ea[i, 0]), float(ea[j, 0]))
            self.assertTrue(np.allclose(ea[i, 1:4], -ea[j, 1:4]))
            self.assertTrue(np.allclose(ea[i, 4:], ea[j, 4:]))
        diagonal = lookup[(0, 2)]
        self.assertAlmostEqual(float(ea[diagonal, 5]), 0.0, places=6)
        self.assertEqual(float(ea[diagonal, 6]), 0.0)

    def test_exact_welding_and_npz_correspondence_on_real_mesh_geometry(self):
        import open3d as o3d
        with tempfile.TemporaryDirectory() as td:
            td = Path(td); stl, out = td / "box.stl", td / "graphs"
            box = o3d.geometry.TriangleMesh.create_box(2, 3, 4)
            box.compute_vertex_normals()
            # STL round trip deliberately creates repeated triangle-corner vertices.
            self.assertTrue(o3d.io.write_triangle_mesh(str(stl), box))
            raw = o3d.io.read_triangle_mesh(str(stl)); welded = cg._load_mesh(stl)
            self.assertGreater(len(raw.vertices), len(welded.vertices))
            cg.prepare(stl, None, out, (20,), td / "figures")
            a = np.load(out / "graph_20.npz", allow_pickle=False)
            self.assertTrue(np.isfinite(a["x"]).all())
            self.assertTrue(np.isfinite(a["edge_attr"]).all())
            self.assertTrue((a["face"] >= 0).all() and (a["face"] < len(a["pos_mm"])).all())
            self.assertEqual(a["original_vertex_index"].shape[0], len(a["pos_mm"]))
            # The correspondence is measured against original (welded) mesh points.
            source = np.asarray(welded.vertices)
            self.assertTrue(np.allclose(a["pos_mm"], source[a["original_vertex_index"]], atol=np.max(a["original_vertex_distance_mm"]) + 1e-5))

    def test_generated_brep_masks_and_adjacency_when_available(self):
        p = ROOT / "data_processed" / "brep" / "faces_uv.npz"
        if not p.exists(): self.skipTest("full B-Rep preparation has not run")
        a = np.load(p, allow_pickle=False)
        self.assertTrue(np.isfinite(a["pos_mm"]).all() and np.isfinite(a["curvature"]).all())
        self.assertEqual(a["trim_valid"].dtype, np.dtype(bool))
        edge = {tuple(x) for x in a["edge_index"].T}
        self.assertGreater(len(edge), 0)
        self.assertTrue(all((b, a) in edge for a, b in edge))
        self.assertFalse(a["training_valid"].any())

    def test_viewer_triangles_match_npz_face_transpose(self):
        """Viewer JSON must preserve [F,3], while training NPZ is [3,F]."""
        viewer = ROOT / "results" / "graphs" / "viewer_data.js"
        if not viewer.exists(): self.skipTest("graph viewer has not been generated")
        text = viewer.read_text(encoding="utf8")
        self.assertTrue(text.startswith("window.CAD_GRAPH_DATA=") and text.endswith(";"))
        payload = json.loads(text[len("window.CAD_GRAPH_DATA="):-1])
        for graph in payload["graphs"]:
            a = np.load(ROOT / "data_processed" / "graphs" / f"graph_{graph['budget']}.npz", allow_pickle=False)
            self.assertEqual(graph["face"], a["face"].T.tolist())

    def test_brep_box_has_actual_shared_edge_adjacency(self):
        try:
            import importlib.util
            from OCP.BRepPrimAPI import BRepPrimAPI_MakeBox
            spec = importlib.util.spec_from_file_location("brep", ROOT / "전처리" / "brep.py")
            bp = importlib.util.module_from_spec(spec); spec.loader.exec_module(bp)
        except ImportError as exc:
            self.skipTest(str(exc))
        shape = BRepPrimAPI_MakeBox(10., 20., 30.).Shape()
        edge_index, edge_attr, boundary, nonmanifold = bp._face_adjacency(shape, 6, bp._ocp())
        links = {tuple(e) for e in edge_index.T}
        self.assertGreater(len(links), 0)
        self.assertTrue(all((b, a) in links for a, b in links))
        self.assertTrue(np.isfinite(edge_attr).all())
        self.assertEqual(boundary, 0)
        self.assertEqual(nonmanifold, 0)


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