"""Add SiLU to a LoRA Formula and reload its declaration with the same bindings.

Fixed 1024-dimensional neural-layer inputs, rank 32, no training or model download.
This exercise checks computation editing and reload, not reading-task quality.
Requires arti-fit==3.3.0 and PyTorch, as installed by the CPU starter guide.
"""
import json
from pathlib import Path
import torch
from arti import mechanisms as api


def with_nonlinearity(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 = api.scalar_map(hidden, mode='silu')  # Try changing 'silu' to 'tanh'.
    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_nonlinearity(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()}
    assert record['same_bindings'] and record['output_max_abs_change'] > 0
    assert torch.equal(after, reloaded)
    (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 __name__ == '__main__':
    main()