"""CPU request descriptors and pinned static transform, never a CUDA test."""
import copy
import unittest
import numpy as np
import normal_inputs as I
from normal_io import load_source, LAB
from normal_runner import compose_requests, derive_spec

class InputsTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.inputs=load_source('_normal_test_CI', LAB/'conditioning_inputs.py',I.PINS['conditioning_inputs.py'])

    def task(self,arm='paired',cs=True):
        return dict(task_id='block/normal/'+arm+'/000',kind='train',arm=arm,steps=1200,mode='physiological',
            recipe=dict(cs=None if not cs else dict(cue='A',seed=2**62+37,start_tick=0,stop_tick_exclusive=1000,source_stream_id='s'),
                us=dict(nominal_slot_tick=900,rng_seed=None,rng_draws=0,requested=True,tick=900),
                D_command=0. if arm=='d_zero' else 1.,update_enabled=arm!='frozen'))

    def test_full63_seed_and_aliases(self):
        a=self.task();b=self.task('d_zero')
        self.assertEqual(I.descriptor(a),I.descriptor(b))
        self.assertEqual(I.descriptor(a)['seed'],2**62+37)

    def test_empty_CS_is_not_resampled(self):
        self.assertIsNone(I.descriptor(self.task('dan_only',False)))

    def test_all_DAN_ids_and_counterfactual_US(self):
        dan=np.array([10,20],np.int32);empty=dict(tick=np.empty(0,np.int32),neuron_idx=np.empty(0,np.int32))
        for arm in ('paired','d_zero','frozen'):
            composed,us=compose_requests(empty,dan,self.task(arm),self.inputs)
            np.testing.assert_array_equal(composed['requests']['neuron_idx'],dan)
            self.assertEqual(us['tick'].tolist(),[900,900])

    def test_no_US_and_zero_command(self):
        task=self.task('cs_only');task['recipe']['us'].update(requested=False,tick=None);task['recipe']['D_command']=0.
        cs=dict(tick=np.array([1],np.int32),neuron_idx=np.array([5],np.int32))
        combined,us=compose_requests(cs,np.array([10,20],np.int32),task,self.inputs)
        self.assertEqual(len(us['tick']),0);self.assertEqual(combined['requests']['neuron_idx'].tolist(),[5])

    def test_conflict_and_command_fail_closed(self):
        cs=dict(tick=np.array([900],np.int32),neuron_idx=np.array([10],np.int32))
        with self.assertRaises(ValueError):compose_requests(cs,np.array([10,20],np.int32),self.task(),self.inputs)
        task=self.task('d_zero');task['recipe']['D_command']=1.
        with self.assertRaises(ValueError):compose_requests(cs,np.array([10,20],np.int32),task,self.inputs)

    def test_probe_no_US_no_update(self):
        p=dict(kind='probe',mode='matched_energy',steps=1200,feedback_to_training_or_idle=False,
            recipe=dict(cue='B',input=dict(seed=2**61,stream_id='eval'),cs_start_tick=0,cs_stop_tick_exclusive=1000,
                US_requested=False,D_command=0.,weight_updates_enabled=False))
        d=I.descriptor(p);self.assertEqual(d['table_id'],'matched_energy--B')
        empty=dict(tick=np.empty(0,np.int32),neuron_idx=np.empty(0,np.int32))
        combined,us=compose_requests(empty,np.array([10,20],np.int32),p,self.inputs)
        self.assertEqual(len(combined['requests']['tick']),0)
        p['feedback_to_training_or_idle']=True
        with self.assertRaises(ValueError):compose_requests(empty,np.array([10,20],np.int32),p,self.inputs)

    def test_clock_rejected(self):
        task=self.task();task['steps']=1800
        with self.assertRaises(ValueError):I.descriptor(task)

    def test_case_math_config_retained(self):
        original=dict(kind='fullgraph',dt_ms=.1,apl_delay_steps=18,precision='float64',apl_precision='float32',
            record_history=False,seed=1073741825,model=dict(tau=20),apl=dict(gain=.01),graph_dir='/frozen/graph',
            rate_sum_hz=1,mode='physiological')
        task=self.task();task['mode']='matched_energy';before=copy.deepcopy(original)
        spec=derive_spec(original,task,budget_root='/new/exp',output_cap=1234)
        self.assertEqual(original,before)
        for key in ('dt_ms','apl_delay_steps','precision','apl_precision','seed','model','apl','graph_dir'):
            self.assertEqual(spec[key],original[key])
        self.assertEqual(spec['steps'],1200);self.assertFalse(spec['within_trial_weight_updates'])

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