"""Remote-only numerical comparison with pinned official Torch definitions.

Reads only a known official Transformers source file; extracts named mathematical
operators/classes via AST, removing optional kernel decorators. Never imports the
complete model or executes checkpoint repository code. No compilation/build.
"""
import argparse
import ast
import hashlib
import json
from pathlib import Path
from types import SimpleNamespace
import numpy as np
from .model import StreamingLM, delta_rule, rope, norm

EXPECTED_COMMIT = 'fd3cf0140de03f7619ac95b5cfed05b711ff9e22'
NAMES = {'apply_mask_to_padding_states', 'l2norm', 'torch_chunk_gated_delta_rule', 'torch_recurrent_gated_delta_rule',
         'causal_conv1d_fn', 'Qwen3_5RMSNormGated', 'Qwen3_5GatedDeltaNet',
         'rotate_half', 'apply_rotary_pos_emb', 'repeat_kv', 'eager_attention_forward',
         'Qwen3_5Attention', 'Qwen3_5MLP', 'Qwen3_5RMSNorm', 'Qwen3_5DecoderLayer',
         'Qwen3_5TextRotaryEmbedding'}


def official_namespace(source):
    import torch
    import torch.nn.functional as F
    tree = ast.parse(source)
    nodes = []
    for node in tree.body:
        if isinstance(node, (ast.FunctionDef, ast.ClassDef)) and node.name in NAMES:
            for child in ast.walk(node):
                if isinstance(child, (ast.FunctionDef, ast.ClassDef)):
                    child.decorator_list = [d for d in child.decorator_list if isinstance(d, ast.Name) and d.id in ('staticmethod','classmethod')]
            nodes.append(node)
    assert {x.name for x in nodes} == NAMES
    tree = ast.Module(body=[ast.ImportFrom(module='__future__',names=[ast.alias(name='annotations')],level=0),*nodes],type_ignores=[])
    ns = {'torch':torch, 'nn':torch.nn, 'F':F, 'ACT2FN':{'silu':F.silu},
          'GradientCheckpointingLayer':torch.nn.Module,
          'is_torchdynamo_exporting':lambda:False,
          'ALL_ATTENTION_FUNCTIONS':SimpleNamespace(get_interface=lambda name, fallback:fallback)}
    exec(compile(ast.fix_missing_locations(tree),'pinned_transformers_reference','exec'),ns)
    return ns


class ArrayLM(StreamingLM):
    def __init__(self, arrays, config):
        self.arrays, self.config, self.eps = arrays,config,config['rms_norm_eps']
    def guard(self):
        pass
    def small(self,name):
        return self.arrays[name]
    def linear(self,x,name,rows=256):
        return x@self.arrays[name+'.weight'].T
    def embedding(self,ids):
        return self.arrays['embed_tokens.weight'][ids]


def main():
    p=argparse.ArgumentParser();p.add_argument('--source',type=Path,required=True);p.add_argument('--output',type=Path,required=True)
    a=p.parse_args()
    import torch
    torch.set_num_threads(1)
    torch.manual_seed(42)
    provenance=json.loads(Path(__file__).with_name('reference.json').read_text())
    assert hashlib.sha256(a.source.read_bytes()).hexdigest()==provenance['sha256'], 'Reference source hash mismatch'
    ns=official_namespace(a.source.read_text())
    cfg={'hidden_size':16,'num_attention_heads':2,'num_key_value_heads':1,'head_dim':8,
         'intermediate_size':24,'hidden_act':'silu','attention_bias':False,'attention_dropout':0.,
         'linear_num_value_heads':2,'linear_num_key_heads':2,'linear_key_head_dim':4,
         'linear_value_head_dim':4,'linear_conv_kernel_dim':4,'rms_norm_eps':1e-6,
         'layer_types':['linear_attention','full_attention'],'max_position_embeddings':4096,
         'rope_parameters':{'rope_type':'default','rope_theta':10000000.,'partial_rotary_factor':.5,'mrope_section':[1,0,0]},
         '_attn_implementation':'eager'}
    config=SimpleNamespace(**cfg)
    layers=[ns['Qwen3_5DecoderLayer'](config,i).eval() for i in range(2)]
    for layer in layers:
        for name,t in layer.named_parameters():
            with torch.no_grad():
                t.uniform_(-.2,.2)
                if name=='linear_attn.norm.weight':t.add_(1.)
    arrays={f'layers.{i}.'+k:v.detach().numpy().copy() for i,l in enumerate(layers) for k,v in l.state_dict().items()}
    lm=ArrayLM(arrays,cfg)
    rng=np.random.default_rng(123)
    x=rng.standard_normal((2,7,16),dtype=np.float32)
    tx=torch.from_numpy(x)
    positions=torch.arange(7)[None,:]
    rotary=ns['Qwen3_5TextRotaryEmbedding'](config)
    cos,sin=rotary(tx,positions)
    mask=torch.triu(torch.full((7,7),-float('inf')),diagonal=1)[None,None]
    evidence=[]
    def check(name,actual,expected,atol=2e-5,rtol=2e-5):
        actual=np.asarray(actual);expected=expected.detach().numpy() if hasattr(expected,'detach') else np.asarray(expected)
        np.testing.assert_allclose(actual,expected,atol=atol,rtol=rtol)
        evidence.append({'test':name,'shape':list(actual.shape),'max_abs_error':float(np.max(np.abs(actual-expected))),'atol':atol,'rtol':rtol,'passed':True})
    with torch.no_grad():
        for i,layer in enumerate(layers):
            n=lm.normalization(x,f'layers.{i}.input_layernorm')
            tn=layer.input_layernorm(torch.from_numpy(x))
            check(f'layer{i}_rmsnorm',n,tn)
            if i==0:
                out=lm.linear_attention(n,i)
                tout=layer.linear_attn(tn)
            else:
                out,pr=lm.full_attention(n,i)
                tout,tpr=layer.self_attn(tn,(cos,sin),mask)
                check('attention_probabilities',pr,tpr)
            check(f'layer{i}_mixer',out,tout)
            h=x+out
            n=lm.normalization(h,f'layers.{i}.post_attention_layernorm')
            check(f'layer{i}_mlp',lm.mlp(n,i),layer.mlp(layer.post_attention_layernorm(torch.from_numpy(h))))
            actual=h+lm.mlp(n,i)
            expected=layer(torch.from_numpy(x),(cos,sin),attention_mask=mask if i else None)
            check(f'layer{i}_complete',actual,expected)
        # Independent recurrence/chunk equality, including nonmultiple sequence length.
        for length in (1,7,33):
            q=rng.standard_normal((2,length,2,4),dtype=np.float32)
            k=rng.standard_normal(q.shape,dtype=np.float32)
            v=rng.standard_normal((2,length,2,5),dtype=np.float32)
            g=-np.abs(rng.standard_normal((2,length,2),dtype=np.float32))
            beta=rng.uniform(size=g.shape).astype(np.float32)
            out,state=delta_rule(q,k,v,g,beta)
            args=[torch.from_numpy(z) for z in (q,k,v,g,beta)]
            for kind in ('torch_recurrent_gated_delta_rule','torch_chunk_gated_delta_rule'):
                ref,refstate=ns[kind](*args,output_final_state=True,use_qk_l2norm_in_kernel=True)
                check(f'{kind}_length{length}_output',out,ref)
                check(f'{kind}_length{length}_state',state,refstate)
        # Full two-layer residual forward, final norm and embedding lookup.
        arrays['embed_tokens.weight']=rng.standard_normal((31,16),dtype=np.float32)
        arrays['norm.weight']=rng.standard_normal((16,),dtype=np.float32)*.1
        ids=rng.integers(0,31,size=(2,7))
        h=torch.from_numpy(arrays['embed_tokens.weight'][ids])
        for i,l in enumerate(layers):h=l(h,(cos,sin),attention_mask=mask if i else None)
        expected=h*torch.rsqrt(h.square().mean(-1,keepdim=True)+1e-6)*(1+torch.from_numpy(arrays['norm.weight']))
        actual,_=lm.forward(ids)
        check('complete_two_layer_text_forward',actual,expected)
    result={'passed':True,'reference_commit':EXPECTED_COMMIT,'source_sha256':hashlib.sha256(a.source.read_bytes()).hexdigest(),
            'torch':torch.__version__,'numpy':np.__version__,'implementation_sha256':hashlib.sha256(Path(__file__).with_name('model.py').read_bytes()).hexdigest(),'tests':evidence,'limitations':'FP32 CPU semantics on tiny synthetic configurations; not full-checkpoint BF16 kernel equivalence'}
    a.output.write_text(json.dumps(result,indent=2))
    print(json.dumps({'passed':True,'checks':len(evidence),'maximum_absolute_discrepancy':max(x['max_abs_error'] for x in evidence)}))


if __name__=='__main__':main()