import tempfile
import unittest
from pathlib import Path
import numpy as np
import torch
from sklearn.preprocessing import StandardScaler
from cfref import CFREF
from cfref._core import train_cfref


class TestCFREF(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        torch.set_num_threads(1)
        rng = np.random.RandomState(42)
        cls.X = rng.normal(size=(48, 6))
        cls.y = np.tile([0]*16+[1]*8, 2)
        cls.c = np.repeat(['source_A','source_B'], 24)
        cls.names = ['bin_'+str(i) for i in range(6)]
        cls.R = rng.normal(size=(5,6)); cls.Q = rng.normal(size=(9,6))
        cls.kw = dict(hidden_dim=8, embedding_dim=4, episodes=12,
                      reference_range=(3,5), threshold_references=3, threshold_draws=4)
        cls.model = CFREF(**cls.kw).fit(cls.X,cls.y,cls.c,feature_names=cls.names)

    def test_original_core_parity(self):
        Z = StandardScaler().fit_transform(self.X).astype(np.float32)
        net = train_cfref(Z,self.y,self.c,n_ref_range=(3,5),episodes=12,d_h=8,d_e=4)
        for k,v in net.state_dict().items():
            np.testing.assert_array_equal(v.detach().numpy(),self.model.net_.state_dict()[k].numpy())
        rng=np.random.RandomState(0); out=[]; Xt=torch.tensor(Z)
        with torch.no_grad():
            for c in np.unique(self.c):
                ctrl=np.where((self.c==c)&(self.y==0))[0]
                for _ in range(4):
                    ref=rng.choice(ctrl,3,replace=False)
                    out.append(net(Xt[np.setdiff1d(ctrl,ref)],Xt[ref]).numpy())
        self.assertEqual(self.model.threshold_,float(np.quantile(np.concatenate(out),.95)))

    def test_save_load(self):
        with tempfile.TemporaryDirectory() as d:
            p=Path(d)/'model.cfref';self.model.save(p); m=CFREF.load(p)
            np.testing.assert_array_equal(m.decision_function(self.Q,self.R,feature_names=self.names),self.model.decision_function(self.Q,self.R,feature_names=self.names))
            self.assertEqual(m.threshold_,self.model.threshold_)

    def test_permutation_invariance(self):
        np.testing.assert_allclose(self.model.decision_function(self.Q,self.R,feature_names=self.names),self.model.decision_function(self.Q,self.R[::-1],feature_names=self.names),rtol=1e-5,atol=1e-6)

    def test_feature_mismatch(self):
        with self.assertRaises(ValueError):self.model.predict(self.Q,self.R,feature_names=self.names[::-1])

    def test_small_reference(self):
        with self.assertRaises(ValueError):self.model.predict(self.Q,self.R[:1],feature_names=self.names)

    def test_missing_threshold_contributors(self):
        with self.assertRaisesRegex(ValueError,'No source cohort'):
            CFREF().fit(self.X,self.y,self.c,feature_names=self.names)

    def test_invalid_values(self):
        X=self.X.copy();X[0,0]=np.nan
        with self.assertRaises(ValueError):CFREF(**self.kw).fit(X,self.y,self.c,feature_names=self.names)

    def test_scaling_is_fixed_at_inference(self):
        before=self.model.mean_.copy()
        self.model.predict(self.Q+10,self.R+10,feature_names=self.names)
        np.testing.assert_array_equal(before,self.model.mean_)
        np.testing.assert_allclose(before,self.X.mean(axis=0))

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