"""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))


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