import importlib.util
from pathlib import Path
import unittest

import numpy as np

MODULE = Path(__file__).parents[1] / "representations.py"
spec = importlib.util.spec_from_file_location("representations", MODULE)
rep = importlib.util.module_from_spec(spec); spec.loader.exec_module(rep)


class RepresentationTests(unittest.TestCase):
    def test_z_buffer_and_barycentric_backprojection(self):
        # Same projected triangle twice: the +z triangle must be the front surface.
        p = np.array([[-1,-1,0],[1,-1,0],[0,1,0],[-1,-1,1],[1,-1,1],[0,1,1]],float)
        f = np.array([[0,1,2],[3,4,5]])
        out = rep.rasterize_view(p,f,np.tile([0,0,1.],(6,1)),np.zeros(6),np.zeros(6),np.ones(6,bool),[0,0,1],33)
        y,x=np.argwhere(out["visibility"])[len(np.argwhere(out["visibility"]))//2]
        self.assertEqual(out["source_triangle"][y,x], 1)
        b=out["barycentric"][y,x]; self.assertTrue(np.isclose(b.sum(),1))
        self.assertTrue(np.allclose(out["xyz_mm"][y,x], b @ p[f[1]], atol=1e-5))

    def test_sampling_is_deterministic_and_missing_masks_are_explicit(self):
        cad={"pos_mm":np.arange(90,dtype=float).reshape(30,3),"normal":np.tile([0,0,1.],(30,1)),"k1_per_mm":np.ones(30),"k2_per_mm":np.ones(30),"curvature_valid":np.ones(30,bool),"normal_turn_deg":np.zeros(30)}
        cad["curvature_valid"][0] = False
        a,b=rep.cad_points(cad,10),rep.cad_points(cad,10)
        self.assertTrue(np.array_equal(a["source_triangle"],b["source_triangle"]))
        # Explicitly exercise an invalid curvature sample, regardless of RNG draw.
        one=rep.cad_points({k:(v[:1] if isinstance(v,np.ndarray) and v.ndim else v) for k,v in cad.items()}, 1)
        self.assertFalse(one["feature_valid"][0,7:9].any())
        self.assertTrue(one["feature_valid"][0,9])
        self.assertEqual(float(one["x"][0,7]), 0.0)
        self.assertFalse(a["feature_valid"][:,10:].any())  # thickness / thermal / stress
        self.assertFalse(a["training_valid"].any())

    def test_surface_voxel_aggregation_is_surface_only_keeps_pitch_and_ignores_invalid(self):
        # The first triangle has orientation-sign-flipped principal curvatures;
        # its invariant magnitude/Gaussian must not cancel.  The second is invalid.
        cad={"pos_mm":np.array([[0,0,0],[.1,0,0],[0,.1,0],[.1,.1,0],[.2,.1,0],[.1,.2,0]],float),"normal":np.tile([0,0,1.],(6,1)),"k1_per_mm":np.array([1.,-1.,1.,100.,100.,100.]),"k2_per_mm":np.array([1.,-1.,1.,100.,100.,100.]),"curvature_valid":np.array([1,1,1,0,0,0],bool),"normal_turn_deg":np.zeros(6),"face":np.array([[0,3],[1,4],[2,5]])}
        v=rep.surface_voxels(cad,1)
        self.assertEqual(tuple(v["grid_dims"]), (1,1,1))
        self.assertTrue(v["surface_only"] and not v["solid_occupancy"])
        self.assertTrue(v["x"][:,0].min() == 1 and not v["feature_valid"][:,9:].any())
        self.assertAlmostEqual(float(v["x"][0,6]), np.sqrt(2), places=5)
        self.assertAlmostEqual(float(v["x"][0,7]), 1.0, places=5)
        self.assertTrue(v["feature_valid"][0,6] and v["feature_valid"][0,7] and v["feature_valid"][0,8])
        self.assertAlmostEqual(float(v["voxel_pitch_mm"]), .2, places=6)

    def test_image_contract_masks_and_physical_screen_scale(self):
        p=np.array([[0,0,0],[10,0,0],[0,1,0]],float)
        cad={"pos_mm":p,"normal":np.tile([0,0,1.],(3,1)),"face":np.array([[0],[1],[2]]),"k1_per_mm":np.array([1.,1.,99.]),"k2_per_mm":np.array([1.,1.,99.]),"curvature_valid":np.array([1,1,0],bool),"normal_turn_deg":np.zeros(3)}
        raw=rep.rasterize_view(p,cad["face"],cad["normal"],np.ones(3),np.ones(3),cad["curvature_valid"],[0,0,1],64)
        b=raw["orthographic_bounds_uw_mm"]
        self.assertAlmostEqual(float(b[1]-b[0]), float(b[3]-b[2]), places=5)
        mv=rep.make_multiview(cad,32)
        self.assertEqual(mv["x"].shape, (6,len(rep.IMAGE_FEATURE_NAMES),32,32))
        self.assertEqual(mv["feature_valid"].shape, mv["x"].shape)
        self.assertFalse(mv["feature_valid"][:,14:].any())  # thickness/thermal/stress
        self.assertFalse(mv["feature_valid"][:,7:9].any())  # triangle curvature invalid
        self.assertTrue(np.isfinite(mv["x"]).all())

    def test_thermal_requires_real_history_and_has_stated_rate_semantics(self):
        x=rep.thermal_history_features(np.array([0.,1.,2.]),np.array([[10.],[8.],[5.]]))
        self.assertTrue(x["max"][0] == 10 and x["cooling_rate_per_s"][0] < 0)
        with self.assertRaises(ValueError): rep.thermal_history_features(np.array([0.,np.nan]),np.array([[1.],[2.]]))

    def test_degenerate_scan_neighborhood_is_masked(self):
        x, valid=rep.estimate_scan_features(np.array([[1.,2.,3.]]))
        self.assertTrue(valid[0,:3].all())
        self.assertFalse(valid[0,3:].any())
        self.assertTrue(np.isfinite(x).all())


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