"""Run on the authorized CPU server: python -m unittest model_shell.test_analysis."""
import unittest
import numpy as np
from model_shell.analyze import sample,geometry,analyze,seed
from finetune_compare.lora import EffectiveWeight

class ArrayTensor:
    def __init__(self,x):self.x=x;self.shape=x.shape;self.numel=x.size;self.dtype='F32'
    def read(self,s=slice(None)):return self.x.reshape(-1)[s].copy()
    def evict(self):pass

class Checks(unittest.TestCase):
    def test_sampling_and_lora(self):
        rng=np.random.default_rng(17);w=rng.normal(size=(80,64)).astype('f');a=rng.normal(size=(2,64)).astype('f');b=rng.normal(size=(80,2)).astype('f')
        effective=EffectiveWeight(ArrayTensor(w),a,b,2);explicit=ArrayTensor(w+2*b@a)
        for axis,n,d,start,stop in [('rows',80,64,0,80),('columns',64,80,0,80),('rows',32,64,16,48)]:
            g=dict(id=axis+str(start),orientation=axis,population=n,dimension=d,start=start,stop=stop)
            x,h=sample(effective,g);y,k=sample(explicit,g);np.testing.assert_allclose(x,y,atol=2e-6,rtol=1e-5);self.assertEqual(h,k)
            self.assertEqual(x.shape,(n,d))
            rr=np.random.default_rng(seed(g['id']+'/indices'));indices=np.sort(rr.choice(n,n,replace=False));order=rr.permutation(n)
            cloud=explicit.x[start:stop] if axis=='rows' else explicit.x.T
            np.testing.assert_allclose(y,cloud[indices][order],atol=0,rtol=0)
    def test_center_scale_and_rank(self):
        rng=np.random.default_rng(5);x=rng.normal(size=(128,96)).astype('f')
        for center in ('origin','mean'):
            p=geometry(x,64,center);q=geometry(x*3,64,center)
            np.testing.assert_allclose(p['cdf_rms'],q['cdf_rms']);self.assertLessEqual(p['covariance_participation'],p['rank_ceiling']+1e-6)
        p=geometry(x,64,'mean');q=geometry(x+4,64,'mean');np.testing.assert_allclose(p['median_rms'],q['median_rms'],rtol=1e-6)
        self.assertFalse(geometry(np.zeros((32,64)),16,'mean')['valid'])
    def test_ball_fit_and_sparse_counterexample(self):
        rng=np.random.default_rng(33);d=32;ncal=256
        # Analytic expectation over independent calibration maxima.
        c=rng.random((20000,ncal))**(1/d);v=rng.random(20000)**(1/d)
        observed=np.mean(v<=.99*c.max(1));expected=.99**d*ncal/(ncal+1)
        self.assertLess(abs(observed-expected),.015)
        radii=rng.random(512)**(1/d);x=np.zeros((512,d),dtype='f');x[:,0]=radii
        p=geometry(x,256,'origin');self.assertGreater(p['median_max'],.9)
        # Sparse support alone retains the full-dimensional radial law.
        self.assertLess(p['cdf_max'][2],.08)
        x[:,0]=rng.random(512)**.5
        self.assertGreater(geometry(x,256,'origin')['cdf_max'][2],.65)
        sphere=rng.normal(size=(512,d));sphere/=np.linalg.norm(sphere,axis=1)[:,None]
        self.assertEqual(geometry(sphere,256,'origin')['cdf_max'][2],0.)
    def test_reproducible_controls(self):
        x=np.random.default_rng(7).normal(size=(32,40)).astype('f');a=analyze(x,'test');b=analyze(x,'test')
        self.assertEqual(a,b)
        self.assertTrue(a['mean']['valid'])

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