from pathlib import Path
import hashlib,json,math,sys,unittest
import numpy as np
import pandas as pd
sys.path.insert(0,str(Path(__file__).parent))
from checkpoint_analysis import COLS,STEPS,corner,cosine,lowrank_inner
from analyze_tables import main,segmented
ROOT=Path(__file__).resolve().parents[1]

class SourceDataTests(unittest.TestCase):
 @classmethod
 def setUpClass(cls):
  cls.g=np.load(ROOT/'data/gradient.npy');cls.e=np.load(ROOT/'data/eval.npy');cls.p=np.load(ROOT/'data/pca.npy');cls.s=main()
 def test_01_artifact_digests(self):
  v=json.loads((ROOT/'provenance/source_verification.json').read_text())
  for f,h in v['distributed_derived_artifact_sha256'].items():self.assertEqual(hashlib.sha256((ROOT/f).read_bytes()).hexdigest(),h)
 def test_02_counts(self):self.assertEqual(self.g.shape,(792,2));self.assertEqual(self.e.shape,(79,2))
 def test_03_step_grid(self):np.testing.assert_array_equal(self.g[:,0],np.arange(1,793));np.testing.assert_array_equal(self.e[:,0],np.arange(10,791,10))
 def test_04_finite(self):self.assertTrue(np.isfinite(self.g).all() and np.isfinite(self.e).all() and np.isfinite(self.p).all())
 def test_05_gradient_peak(self):self.assertEqual(self.s['raw_gradient_peak']['step'],169);self.assertEqual(self.s['raw_gradient_peak']['value'],40.9891242980957)
 def test_06_prefix_digest(self):self.assertEqual(self.s['original_prefix_numeric_digests']['gradient'],'93ef3714f2ac07d5eccbb661b271763765644035dd92402be298d760a4cbe66b');self.assertEqual(self.s['original_prefix_numeric_digests']['eval'],'4c61d98d4151417004068d424fd41be2599cb43c24d60fac2d4238a0272a409b')
 def test_07_eval_interval(self):self.assertEqual(self.s['largest_eval_decrease']['to_step'],170)
 def test_08_late_reversal(self):self.assertEqual([r['step'] for r in self.s['nondecreasing_eval_intervals']],[770])
 def test_09_two_line_prefix(self):self.assertEqual(self.s['prefix_two_line_best']['split_after_step'],181)
 def test_10_pca_grid(self):np.testing.assert_array_equal(self.p[:,0],np.arange(5,791,5));self.assertEqual(self.p.shape,(158,3))
 def test_11_pca_center(self):np.testing.assert_allclose(self.p[:,1:].mean(axis=0),0,atol=1e-12)
 def test_12_selected_norms(self):
  d=pd.read_csv(ROOT/'data/selected_checkpoint_metrics.csv');self.assertTrue(np.all(np.diff(d.iloc[:,1])>0));self.assertAlmostEqual(d.iloc[-1,1],5.324990583460912)
 def test_13_filter_peaks(self):
  d=pd.read_csv(ROOT/'data/corner_sensitivity.csv');self.assertTrue((d[d.amplitude_threshold==.002].peak_step_B_and_effective==190).all());self.assertEqual(len(d),18)
 def test_14_schedule(self):self.assertEqual(len(STEPS),167);self.assertEqual(STEPS[-1],792)

class AlgebraTests(unittest.TestCase):
 def setUp(self):
  r=np.random.default_rng(314159);self.a=r.normal(size=(3,11));self.b=r.normal(size=(3,7));self.U=np.stack([64*np.outer(b,a) for a,b in zip(self.a,self.b)])
 def test_15_gram_equals_dense(self):
  G=4096*(self.a@self.a.T)*(self.b@self.b.T);np.testing.assert_allclose(G,self.U.reshape(3,-1)@self.U.reshape(3,-1).T,rtol=1e-13)
 def test_16_sign_invariance(self):np.testing.assert_array_equal(np.outer(-self.b[0],-self.a[0]),np.outer(self.b[0],self.a[0]))
 def test_17_scale_gauge(self):np.testing.assert_allclose(np.outer(8*self.b[0],self.a[0]/8),np.outer(self.b[0],self.a[0]),rtol=1e-14)
 def test_18_corner_equals_dense(self):
  U=self.U.reshape(3,-1);G=U@U.T;u,v=U[0]-U[1],U[2]-U[1];self.assertAlmostEqual(corner(G,0,1,2),float(u@v/np.sqrt((u@u)*(v@v))),places=13)
 def test_19_zero_cosine_undefined(self):self.assertTrue(math.isnan(cosine(np.array([[0.,0.],[0.,1.]]),0,1)))
 def test_20_straight_corner(self):
  u=np.array([[0.,0.],[1.,0.],[2.,0.]]);self.assertEqual(corner(u@u.T,0,1,2),-1.)
 def test_21_stable_difference_expansion(self):
  def d(j,i):return [(self.b[j]-self.b[i],self.a[j]),(self.b[i],self.a[j]-self.a[i])]
  actual=np.sum((self.U[0]-self.U[1])*(self.U[2]-self.U[1]))/4096
  self.assertAlmostEqual(lowrank_inner(d(0,1),d(2,1)),float(actual),places=11)
 def test_22_origin_inequality_not_error(self):
  origin,current,estimate=2.,1.,1.;self.assertNotEqual(estimate,origin);self.assertEqual(estimate,current)
if __name__=='__main__':unittest.main(verbosity=2)