"""Your first ARTI rule: edit my_rule, then run, save and reload.

Python 3.10+, arti-fit==3.3.0, and PyTorch on CPU.
Four fixed 1024-dimensional neural-layer inputs; no base model or training.
The declaration is reloaded with the same caller-supplied values in this run.
Those values are not saved as a complete model asset by this exercise.
"""
import json
from pathlib import Path

import torch
from arti import mechanisms as api


def my_rule(hidden):
    # Add your operation here, before return.
    return hidden


def with_my_rule(program):
    names = {binding.name: binding for binding in program.bindings}
    hidden = api.contract(names['x'], names['lora.A'],
                          reduce_axes=(('Din', 'Din'),), output_axes=('B', 'S', 'R'))
    hidden = my_rule(hidden)
    update = api.contract(hidden, names['lora.B'],
                          reduce_axes=(('R', 'R'),), output_axes=('B', 'S', 'Dout'))
    return api.FormulaProgram.build(outputs=(
        api.add(names['base'], api.scale(update, names['lora.gain'],
                                        accumulation_dtype='float32'),
                accumulation_dtype='float32'),))


def main():
    original = api.build_lora_program(input_dim=1024, output_dim=1024, rank=32,
        source_ref='arti/formula-operand-bank@1', dtype='float32',
        pointwise_accumulation_dtype='float32')
    changed = with_my_rule(original)
    generator = torch.Generator().manual_seed(7)
    values = {name: torch.randn(32, 1024, generator=generator) * .02
              for name in ['lora.A', 'lora.B']}
    values['lora.B'] = values['lora.B'].t().contiguous()
    bindings = {binding.name: binding for binding in original.bindings
                if isinstance(binding, api.BankBinding)}
    banks = {name: bindings[name].bind(value) for name, value in values.items()}
    inputs = {'x': torch.randn(1, 4, 1024, generator=generator),
              'base': torch.randn(1, 4, 1024, generator=generator),
              'lora.gain': torch.tensor(1., dtype=torch.float32)}
    execute = lambda program: api.FormulaFabricV2(program)(
        inputs=inputs, banks=banks).values[0]
    before, after = execute(original), execute(changed)
    destination = Path('program-change')
    destination.mkdir(exist_ok=True)
    file = destination / 'program.json'
    file.write_text(json.dumps(changed.to_dict(), indent=2)+'\n', encoding='utf8')
    restored = api.FormulaProgram.from_dict(json.loads(file.read_text(encoding='utf8')))
    reloaded = execute(restored)
    record = {
        'scope': 'Fixed neural-layer computation edit; zero training updates',
        'input_shape': list(inputs['x'].shape), 'rank': 32,
        'bank_values': sum(value.numel() for value in values.values()),
        'before_instructions': len(original.instructions),
        'after_instructions': len(changed.instructions),
        'same_bindings': original.bindings == changed.bindings,
        'output_max_abs_change': (after-before).abs().max().item(),
        'reload_max_abs_difference': (after-reloaded).abs().max().item(),
        'output_sample': {'position': [0, 0, 0],
                          'before': before[0, 0, 0].item(),
                          'after': after[0, 0, 0].item()},
    }
    assert record['same_bindings'], 'The exercise expects the original bindings.'
    assert torch.equal(after, reloaded), 'The reloaded declaration must agree.'
    (destination/'result.json').write_text(json.dumps(record, indent=2)+'\n', encoding='utf8')
    print(json.dumps(record, indent=2))
    print('Saved program-change/program.json and program-change/result.json')
    if record['output_max_abs_change'] == 0:
        print('The rule is unchanged. Add SiLU in my_rule, then run this command again.')
    else:
        print('Your rule changed the output. Reload agrees with the changed program.')


if __name__ == '__main__':
    main()