"""A shared Bank-gated neural recurrence, executed from the public Formula IR.

Adapted from ARTI 3.3.0 examples/formula_recurrence.py (MIT).
Source: github.com/ragnarok-io/ARTI at dcb321d6e108d0167c5fbca0a77501f34b154af2.
This is a forward/gradient/state trace, not a task-quality benchmark.
Run: python neural-recurrence.py --output recurrence-record.json [--jax]
"""
import argparse
import json
import os
from pathlib import Path
import sys
import hashlib
sys.dont_write_bytecode=True
import torch
import arti
from arti import mechanisms as m


def build(dim=3):
    tensor=m.TensorType(('B','D'),('B',dim),dtype='float32')
    x,h,total=(m.InputBinding(name,tensor) for name in ('x','h','total'))
    banks={name:m.BankBinding(name,'arti/recurrent-example@1',name,
            m.TensorType(('D',),(dim,),dtype='float32'))
            for name in ('gain','retention','unit','negative_unit','zero')}
    gate=m.scalar_map(m.add(x,m.scale(h,banks['retention'])),mode='sigmoid')
    target=m.scalar_map(m.add(m.scale(x,banks['gain']),h),mode='tanh')
    unit=m.broadcast(banks['unit'],output_axes=tensor.axis_names,output_sizes=tensor.sizes)
    complement=m.add(unit,m.scale(gate,banks['negative_unit']))
    next_h=m.add(m.scale(gate,h),m.scale(complement,target))
    next_total=m.add(total,next_h)
    zero=m.broadcast(banks['zero'],output_axes=tensor.axis_names,output_sizes=tensor.sizes)
    positive=m.compare(next_h,zero,mode='gt')
    body=m.FormulaProgram.build(outputs=(next_h,next_total,positive))
    scan=m.FormulaScan(body,axis='T',sequence_types={
        'x':m.TensorType(('B','T','D'),('B','T',dim),dtype='float32')},
        carry_outputs={'h':body.outputs[0],'total':body.outputs[1]},
        emissions={'hidden':body.outputs[0],'positive':body.outputs[2]},max_length=8)
    return scan,banks,body


def main():
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--output',type=Path,default=Path('recurrence-record.json'))
    parser.add_argument('--jax',action='store_true')
    args=parser.parse_args()
    if arti.__version__ != '3.3.0':
        raise RuntimeError('This recorded example targets arti-fit==3.3.0')
    scan,bindings,body=build()
    weights={'gain':torch.nn.Parameter(torch.tensor([.7,1.1,.9])),
             'retention':torch.nn.Parameter(torch.tensor([.2,.4,.6])),
             'unit':torch.ones(3),'negative_unit':-torch.ones(3),'zero':torch.zeros(3)}
    sequence=torch.tensor([[[.3,-.2,.6],[.1,.4,-.3],[-.4,.5,.2],[.7,-.1,.1],[.2,.3,-.5]]])
    inputs={'x':sequence,'h':torch.zeros(1,3),'total':torch.zeros(1,3)}
    originals={name:value.clone() for name,value in inputs.items()}
    limits=m.FormulaLimits(max_instructions=2048,max_slots=4096,max_steps=1024)
    outputs=scan(inputs=inputs,banks={name:bindings[name].bind(value) for name,value in weights.items()},limits=limits)
    objective=outputs['emit.hidden'].square().mean()
    objective.backward()
    grads={name:value.grad.detach().tolist() for name,value in weights.items() if value.requires_grad}
    assert all(torch.isfinite(weights[name].grad).all() and weights[name].grad.norm()>0 for name in grads)
    assert all(torch.equal(value,originals[name]) for name,value in inputs.items())
    record={'purpose':'execution / gradient / functional-state trace, not task effectiveness',
            'package':'arti-fit==3.3.0','source_revision':'dcb321d6e108d0167c5fbca0a77501f34b154af2',
            'steps':5,'max_length':8,'instruction_count':len(body.instructions),'trainable_values':6,
            'inputs':{name:value.tolist() for name,value in inputs.items()},
            'outputs':{name:value.detach().tolist() for name,value in outputs.items()},
            'bank_gradients':grads,'input_state_unchanged':True,'optimizer_updates':0,
            'source_sha256':hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
            'jax':{'executed':False}}
    if args.jax:
        os.environ.setdefault('JAX_PLATFORMS','cpu')
        import jax
        import jax.numpy as jnp
        import numpy as np
        import arti.jax as aj
        native=aj.bind_formula_scan(scan)
        params=native.init(jax.random.key(0),bank_initializers={
            name:jnp.asarray(value.detach().numpy()) for name,value in weights.items()})
        native_inputs={name:jnp.asarray(value.numpy()) for name,value in inputs.items()}
        actual=jax.jit(native.apply)(params,native_inputs)
        difference=max(float(np.max(np.abs(np.asarray(actual[name]).astype(float)-value.detach().numpy().astype(float)))) for name,value in outputs.items())
        def loss(parameters):return jnp.mean(jnp.square(native.apply(parameters,native_inputs)['emit.hidden']))
        native_grads=jax.grad(loss)(params)
        gradient_difference=max(float(np.max(np.abs(np.asarray(native_grads[name])-weights[name].grad.numpy()))) for name in grads)
        assert difference<2e-6 and gradient_difference<2e-6
        record['jax']={'executed':True,'platform':'cpu',
                       'jit_output_max_abs':difference,'gradient_max_abs':gradient_difference,
                       'scope':'this bounded shared FormulaScan only'}
    args.output.parent.mkdir(parents=True,exist_ok=True)
    args.output.write_text(json.dumps(record,indent=2)+'\n',encoding='utf8')
    print(json.dumps({'steps':record['steps'],'instructions':record['instruction_count'],
                      'input_state_unchanged':True,'jax':record['jax']}))


if __name__=='__main__':main()