"""CPU functional implementation tests; run on authorized remote server only."""
import json
from pathlib import Path
import struct
import tempfile
import unittest
import numpy as np
from functional.model import StreamingLM,causal_conv,silu,attention,norm
from functional.run import probability_metrics,error_rows


class FunctionalTests(unittest.TestCase):
    def test_streamed_projection_embedding_and_output_head(self):
        rng=np.random.default_rng(7)
        arrays={'model.language_model.embed_tokens.weight':rng.normal(size=(19,4)).astype(np.float32),
                'model.language_model.test.weight':rng.normal(size=(9,4)).astype(np.float32)}
        header={};offset=0;body=b''
        for name,arr in arrays.items():
            raw=arr.tobytes();header[name]={'shape':list(arr.shape),'dtype':'F32','data_offsets':[offset,offset+len(raw)]}
            body+=raw;offset+=len(raw)
        encoded=json.dumps(header).encode()
        with tempfile.TemporaryDirectory() as d:
            Path(d,'model.safetensors').write_bytes(struct.pack('<Q',len(encoded))+encoded+body)
            lm=StreamingLM(d,{'rms_norm_eps':1e-6},enforce_memory=False)
            ids=np.array([[2,4,1,18],[7,3,12,5]])
            x=lm.embedding(ids)
            np.testing.assert_array_equal(x,arrays['model.language_model.embed_tokens.weight'][ids])
            actual=lm.linear(x,'test',rows=2)
            np.testing.assert_allclose(actual,x@arrays['model.language_model.test.weight'].T,atol=1e-6)
            logits=x[:,:-1]@arrays['model.language_model.embed_tokens.weight'].T
            m=logits.max(-1)
            expected=m+np.log(np.exp(logits-m[...,None]).sum(-1))-np.take_along_axis(logits,ids[:,1:,None],axis=-1)[...,0]
            np.testing.assert_allclose(lm.nll(x,ids,rows=3),expected,atol=2e-6)

    def test_causality_and_identical_function_metrics(self):
        rng=np.random.default_rng(8)
        q=rng.normal(size=(1,5,2,4)).astype(np.float32)
        k=rng.normal(size=(1,5,1,4)).astype(np.float32)
        v=rng.normal(size=(1,5,1,4)).astype(np.float32)
        y,p=attention(q,k,v)
        v2=v.copy();v2[:,3:]+=100
        y2,_=attention(q,k,v2)
        np.testing.assert_allclose(y[:,:3],y2[:,:3])
        self.assertTrue(np.all(p[...,np.triu_indices(5,1)[0],np.triu_indices(5,1)[1]]==0))
        m=probability_metrics(p,p)
        self.assertAlmostEqual(m['js_mean'],0)
        self.assertAlmostEqual(m['top5_overlap'],1)
        rows=error_rows(y,y,[{'id':'one','split':'heldout'}],{})
        self.assertEqual(rows[0]['relative_l2'],0)

    def test_causal_convolution_and_norm_boundary(self):
        x=np.arange(1,7,dtype=np.float32).reshape(1,3,2)
        kernel=np.array([[1,2],[3,4]],dtype=np.float32)
        expected=x*kernel[:,-1]
        expected[:,1:]+=x[:,:-1]*kernel[:,0]
        np.testing.assert_allclose(causal_conv(x,kernel),silu(expected))
        weight=np.zeros(2,dtype=np.float32)
        np.testing.assert_allclose(norm(x,weight),x/np.sqrt(np.mean(x*x,axis=-1,keepdims=True)+1e-6))
        np.testing.assert_array_equal(norm(x,weight,centered=False),np.zeros_like(x))


    def test_nonrotary_gain_identity_and_rotary_negative_control(self):
        from functional.model import rope
        from functional.gain_factorization import consolidate
        rng=np.random.default_rng(19)
        q=rng.standard_normal((2,7,8,256),dtype=np.float32)
        k=rng.standard_normal((2,7,2,256),dtype=np.float32)
        a=rng.uniform(.5,1.5,256).astype(np.float32)
        b=rng.uniform(.5,1.5,256).astype(np.float32)
        def scores(q,k):return np.einsum('bthd,bshd->bhts',rope(q),np.repeat(rope(k),4,axis=2))/16
        reference=scores(norm(q,a,centered=False),norm(k,b,centered=False))
        for direction in ('q_to_k','k_to_q'):
            q1,k1=consolidate(q,k,a,b,64,direction)
            np.testing.assert_allclose(scores(q1,k1),reference,atol=3e-6,rtol=3e-5)
        # Moving gains through rotated coordinates is not a valid general identity.
        q1,k1=consolidate(q,k,a,b,0,'q_to_k')
        self.assertGreater(float(np.max(np.abs(scores(q1,k1)-reference))),1e-3)


    def test_neuron_function_metrics_and_targeted_replacement(self):
        from functional.neuron_reuse import pair_metrics,NeuronReuseLM
        x=np.linspace(-2,2,128,dtype=np.float64)
        m=pair_metrics(x,x*1.05)
        self.assertAlmostEqual(m['relative_l2_a_to_b'],.05)
        self.assertAlmostEqual(m['pearson'],1.)
        self.assertGreater(m['effective_energy_tokens_a'],16)
        dead=pair_metrics(np.zeros(128),np.zeros(128))
        self.assertEqual(dead['rms_a'],0.)
        class Tiny(NeuronReuseLM):
            def __init__(self):
                self.arrays={'gate_proj':np.array([[1,0],[0,2],[1,1]],dtype=np.float32),
                             'up_proj':np.array([[2,0],[0,1],[1,-1]],dtype=np.float32),
                             'down_proj':np.array([[1,2,3],[-1,1,0]],dtype=np.float32)}
            def linear(self,h,name,rows=256):return h@self.arrays[name.rsplit('.',1)[-1]].T
        lm=Tiny();lm.configure(5,0,1)
        h=np.array([[[1,2],[-1,3]]],dtype=np.float32)
        z=silu(h@lm.arrays['gate_proj'].T)*(h@lm.arrays['up_proj'].T)
        z[...,0]=z[...,1]
        np.testing.assert_allclose(lm.mlp(h,5),z@lm.arrays['down_proj'].T)
        original=silu(h@lm.arrays['gate_proj'].T)*(h@lm.arrays['up_proj'].T)
        np.testing.assert_allclose(lm.mlp(h,6),original@lm.arrays['down_proj'].T)


    def test_uniform_rmsnorm_function_bound_is_tight(self):
        from functional.norm_repetition import uniform_bounds
        a=np.array([1.2,.8,1.5,1.1],dtype=np.float64)
        b=np.array([1.22,.78,1.45,1.13],dtype=np.float64)
        ab,ba=uniform_bounds(a,b)
        basis=np.eye(4)
        # Common RMS denominator cancels; basis inputs attain the exact maximum.
        ya,yb=basis*a,basis*b
        ea=np.linalg.norm(ya-yb,axis=-1)/np.linalg.norm(ya,axis=-1)
        eb=np.linalg.norm(ya-yb,axis=-1)/np.linalg.norm(yb,axis=-1)
        self.assertAlmostEqual(float(ea.max()),ab)
        self.assertAlmostEqual(float(eb.max()),ba)
        self.assertIsNone(uniform_bounds(np.array([0.,1.]),np.ones(2)))


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