import unittest
import json
import numpy as np
import torch
from cfref import nested_lopo
from cfref.validation import training_mask

class TestNestedValidation(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        torch.set_num_threads(1)
        rng=np.random.RandomState(21)
        cls.S=rng.poisson(100,size=(60,5)).astype(float)
        cls.L=rng.poisson(100,size=(60,5)).astype(float)
        cls.y=np.tile([0]*12+[1]*8,3);cls.c=np.repeat(['A','B','C'],20)
        cls.ids=np.array([str(i) for i in range(60)])
        cls.S[20:,-1]=0;cls.L[20:,-1]=0
        cls.kw=dict(outer_cohorts=['A'],seeds=(0,),capacities=((8,4),(12,6)),
                    episodes=3,reference_sizes=(3,),draws=3,threshold_draws=2,
                    inner_references=3,outer_threshold_references=3)
        cls.out=nested_lopo(cls.S,cls.L,cls.y,cls.c,cls.ids,**cls.kw)

    def test_reference_and_query_disjoint(self):
        for row in self.out['assignments'].itertuples():
            self.assertFalse(set(json.loads(row.reference_ids))&set(json.loads(row.query_ids)))

    def test_auc_identical_when_only_threshold_changes(self):
        p=self.out['draw_metrics'].pivot(index=['held_out','seed','n_ref','draw'],columns='method',values='auc')
        np.testing.assert_array_equal(p.logreg_locked,p.logreg_recalibrated)

    def test_pairing(self):
        m=self.out['draw_metrics']
        self.assertTrue((m.groupby(['held_out','seed','n_ref','draw']).query_sha256.nunique()==1).all())
        self.assertEqual(len(self.out['paired_differences']),6)

    def test_outer_feature_exclusion_and_selection_isolation(self):
        a=self.S.copy();a[:20]=a[:20]*1000
        b=nested_lopo(a,self.L,self.y,self.c,self.ids,**self.kw)
        self.assertNotIn(4,self.out['fold_masks'][0]['kept_indices'])
        self.assertEqual(self.out['fold_masks'],b['fold_masks'])
        self.assertTrue(self.out['capacity_picks'].equals(b['capacity_picks']))
        self.assertTrue(self.out['inner_selection'].equals(b['inner_selection']))

    def test_inner_mask_does_not_use_inner_target(self):
        a=self.S.copy();train=self.c=='B';a[~train]=1e9
        np.testing.assert_array_equal(training_mask(self.S,self.L,train),training_mask(a,self.L,train))

    def test_duplicate_ids_rejected(self):
        ids=self.ids.copy();ids[1]=ids[0]
        with self.assertRaisesRegex(ValueError,'unique'):
            nested_lopo(self.S,self.L,self.y,self.c,ids,**self.kw)

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