import copy
import json
from pathlib import Path

root = Path(__file__).resolve().parent
state = json.loads((root / 'bs2-mlp-control/state.json').read_text())
calibrated = json.loads((root / 'real-bs1/state.json').read_text())
# Reuse the BS2 geometry, applying the established calibrated precision policy.
# The compiler still compares the full graph/configuration context and validates
# all layouts, exchanges and physical placement when this checkpoint is loaded.
graph, _, mapping = state['context'].split('\n')
config = calibrated['context'].split('\n')[1]
config = config.replace('capacity_baseline: false', 'capacity_baseline: true')
config = config.replace('exchange_table_budget_bytes: 81920', 'exchange_table_budget_bytes: 114688')
state['context'] = '\n'.join([graph, config, mapping])
for op, plan in state['recipe']['plans'].items():
    if 'Gemm' not in plan['operator']:
        continue
    precision = copy.deepcopy(calibrated['recipe']['plans'][op]['operator']['Gemm']['multiply'])
    plan['operator']['Gemm']['multiply'] = precision
    for operand in plan['requirements']['inputs']:
        operand['format']['precision'] = copy.deepcopy(precision)
for value, fmt in state['inputs'].items():
    fmt['precision'] = copy.deepcopy(calibrated['inputs'][value]['precision'])
for op in ('12', '35'):
    state['recipe']['plans'][op]['dispatch']['Attention']['fp8_scales'] = copy.deepcopy(
        calibrated['recipe']['plans'][op]['dispatch']['Attention']['fp8_scales'])
state['recipe']['cast_before_copies'] = sorted(
    {tuple(site) for site in state['recipe']['cast_before_copies'] if site != [0, 0]}
    | {(12, 1), (12, 2)})
state['alternatives'] = {}
state['visited'] = []
state['attempts'] = 0
(root / 'real-bs2-seed.json').write_text(json.dumps(state))
out = root / 'real-bs2'
out.mkdir(exist_ok=True)
script = (root / 'real-bs2-baseline/build.sh').read_text()
script = script.replace('real-bs2-baseline/', 'real-bs2/')
script = script.replace('--save-search-state', '--load-search-state artifacts/full-qk-20260913/real-bs2-seed.json --save-search-state')
(out / 'build.sh').write_text(script)