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

ROOT=Path(__file__).resolve().parents[1]
sys.path[:0]=[str(ROOT),str(ROOT/"forma3d"),str(ROOT/"vendor/forma/src")]
from common.runtime import assert_disjoint,fleet_score,immutable,read,paired_ci,stable_object_splits

class SelectionTests(unittest.TestCase):
    def scores(self,value=.2):
        return {l:{"error_rate":value} for l in ("en","de","fr","es","it","nl","pl","ko")}

    def test_korean_regression_rejected_despite_mean_gain(self):
        base=self.scores(); candidate=self.scores(.1); candidate["ko"]["error_rate"]=.201
        eligible,score=fleet_score(base,candidate)
        self.assertFalse(eligible); self.assertGreater(score,0)

    def test_non_korean_tolerance(self):
        base=self.scores(); candidate=self.scores(.19)
        candidate["nl"]["error_rate"]=.206
        self.assertFalse(fleet_score(base,candidate)[0])
        candidate["nl"]["error_rate"]=.204
        self.assertTrue(fleet_score(base,candidate)[0])

    def test_split_leakage(self):
        with self.assertRaises(ValueError):
            assert_disjoint({"train":[{"sha":"abc"}],"test":[{"sha":"abc"}]},("sha",))

    def test_object_split_preservation_and_conflicts(self):
        def row(uid): return {"metadata":{"asset_id":uid}}
        groups={"oldgroup":[row("a"),row("new-duplicate")], **{str(i):[row(str(i))] for i in range(20)}}
        assignment=stable_object_splits(groups,{"a":"test"},confirmation=5)
        self.assertEqual(assignment["oldgroup"],"test")
        self.assertEqual(list(assignment.values()).count("confirmation"),5)
        groups["oldgroup"].append(row("b"))
        with self.assertRaises(ValueError): stable_object_splits(groups,{"a":"test","b":"train"},confirmation=5)

    def test_resume_settings(self):
        with tempfile.TemporaryDirectory() as d:
            path=Path(d)/"settings.json"
            immutable(path,{"seed":42}); immutable(path,{"seed":42})
            with self.assertRaises(RuntimeError): immutable(path,{"seed":43})
            self.assertEqual(read(path),{"seed":42})

    def test_cluster_bootstrap_and_ratio(self):
        a=[{"id":str(i),"group":str(i//2),"edits":2,"units":10,"metric":.3} for i in range(8)]
        b=[{**r,"edits":1,"metric":.4} for r in a]
        result=paired_ci(a,b,"edits",ratio=True,repeats=50)
        self.assertEqual(result["clusters"],4)
        self.assertAlmostEqual(result["delta_candidate_minus_baseline"],-.1)
        self.assertAlmostEqual(paired_ci(a,b,"metric",repeats=50)["ci95"][0],.1)
        with self.assertRaises(ValueError): paired_ci(a,b[:-1],"metric")

class GeometryTests(unittest.TestCase):
    def test_rotation_and_scale(self):
        import numpy as np
        from geometry import rotation,rotate_splats
        from image3d_student.splats import Splats,rotation_matrices
        q=np.array([[1,0,0,0]],np.float32)
        a=Splats(np.array([[.1,.2,.3]]),np.ones((1,3)),np.array([[.01,.02,.03]]),q,np.ones((1,1)))
        m=rotation(90,20); b=rotate_splats(a,m,.8)
        np.testing.assert_allclose(b.xyz,a.xyz@m.T*.8,atol=1e-6)
        np.testing.assert_allclose(rotation_matrices(b.quat)[0],m,atol=1e-6)
        np.testing.assert_allclose(b.scale,a.scale*.8)

    def test_silhouette_black_objects_and_empty_outputs(self):
        import numpy as np
        from geometry import splat_mask,mesh_mask,rotation,iou,chamfer
        from image3d_student.splats import Splats
        a=Splats(np.zeros((1,3)),np.zeros((1,3)),np.full((1,3),.1),np.array([[1,0,0,0]]),np.ones((1,1)))
        self.assertGreater(splat_mask(a,rotation(0)).sum(),0)
        self.assertEqual(splat_mask(None,rotation(0)).sum(),0)
        mask=mesh_mask(np.array([[-.2,-.2,0],[.2,-.2,0],[0,.2,0]]),np.array([[0,1,2]]),rotation(0))
        self.assertEqual(iou(mask,mask),1.)
        self.assertEqual(chamfer(np.empty((0,3)),np.zeros((1,3))),2.)
        self.assertEqual(chamfer(np.zeros((1,3)),np.zeros((1,3))),0.)

    def test_ply_round_trip(self):
        import numpy as np
        from image3d_student.splats import Splats,write_ply,read_ply
        a=Splats(np.array([[.1,0,0]]),np.full((1,3),.4),np.full((1,3),.1),np.array([[1,0,0,0]]),np.full((1,1),.8))
        with tempfile.TemporaryDirectory() as d:
            path=Path(d)/"test.ply"; write_ply(path,a); b=read_ply(path)
            for key in a.__dataclass_fields__:
                np.testing.assert_allclose(getattr(a,key),getattr(b,key),atol=1e-6)

    def test_torch_resume_preserves_optimizer_and_rng(self):
        import random
        import torch
        from image3d_student.model import Student
        spec=importlib.util.spec_from_file_location("forma_run",ROOT/"forma3d/run.py")
        mod=importlib.util.module_from_spec(spec); spec.loader.exec_module(mod)
        torch.manual_seed(44); model=Student(16,64)
        optimizer=torch.optim.AdamW(model.parameters(),lr=5e-5)
        x=torch.rand(1,3,128,128)
        model(x)["occ"].square().mean().backward(); optimizer.step()
        rng=random.Random(44)
        with tempfile.TemporaryDirectory() as d:
            path=Path(d)/"last.pt"; mod.save(path,model,optimizer,{"step":1},rng)
            ck=torch.load(path,map_location="cpu",weights_only=True)
            resumed=Student(16,64); resumed.load_state_dict(ck["model"])
            opt2=torch.optim.AdamW(resumed.parameters(),lr=5e-5); opt2.load_state_dict(ck["optimizer"])
            for m,o in ((model,optimizer),(resumed,opt2)):
                o.zero_grad(); m(x)["occ"].square().mean().backward(); o.step()
            for a,b in zip(model.parameters(),resumed.parameters()): torch.testing.assert_close(a,b,rtol=0,atol=0)
            other=random.Random(); other.setstate(ck["python_rng"])
            self.assertEqual(rng.random(),other.random())

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