"""Prepare one V4.4 ablation arm's frozen 20-epoch LM+chat mixture.

Derives the arm request from the frozen V4.1 mixture exactly as V4.2 did (same
joined banks, exclusion masks and crop minimum; conversion pool removed; 20
epochs), then overrides only the arm's seed, chat fraction or crop probability.
Request and manifest are written with LF newlines so the manifest SHA is the
same on Windows and Linux; a configured arm must reproduce its pinned SHA.

The 'control' arm rebuilds V4.2 itself: epochs 0-15 must equal frozen V4.1, and
all 20 must equal the V4.2 mixture when that directory is present.
"""
import argparse
import json
from pathlib import Path
import shutil
import sys
import tomllib

ROOT=Path(__file__).resolve().parents[2]
sys.path.insert(0,str(ROOT/'src'))
from vimeml.training.data import file_sha
from vimeml.training.joint_v44 import ARMS,MIXTURE_EPOCHS
from vimeml.training.v41_data import prepare_maps

OLD=ROOT/'artifacts/training-data/corpus-v41-mixed-c128'
V42=ROOT/'artifacts/training-data/corpus-v42-mixed-c128'


def directory(arm):
    return ROOT/f'artifacts/training-data/corpus-v44-{arm}-c128'


def request_for(arm):
    request=json.loads((OLD/'request.json').read_text('utf-8'))
    if request['epochs']!=16:
        raise ValueError('Unexpected V4.1 mixture shape')
    seed,chat,crop=ARMS[arm]
    for key in ('conversion_dir','conversion_manifest_sha256','conversion_keep',
                'conversion_keep_sha256','conversion_examples','numpy_version'):
        request.pop(key)
    request.update(arm=arm,epochs=MIXTURE_EPOCHS,seed=seed,chat_fraction=chat,crop_probability=crop,
        old_keep=directory(arm).relative_to(ROOT).as_posix()+'/old.keep.npy',
        derived_from='artifacts/training-data/corpus-v41-mixed-c128',
        derivation=('V4.2 request (V4.1 with epochs 16->20 and the conversion pool removed) '
                    f'with seed={seed}, chat_fraction={chat}, crop_probability={crop}'))
    return request


def write_lf(path,text):
    temporary=path.with_suffix(path.suffix+'.tmp')
    temporary.write_bytes(text.encode('utf-8'))
    temporary.replace(path)


def main():
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--arm',choices=sorted(ARMS),required=True)
    args=parser.parse_args()
    target=directory(args.arm)
    request=request_for(args.arm)
    if (target/'request.json').exists():
        if json.loads((target/'request.json').read_text('utf-8'))!=request:
            raise ValueError(f'{target} holds a different request')
    else:
        target.mkdir(parents=True,exist_ok=True)
        shutil.copyfile(OLD/'old.keep.npy',target/'old.keep.npy')
        write_lf(target/'request.json',json.dumps(request,ensure_ascii=False,indent=1)+'\n')
    result=prepare_maps(target,root=ROOT)
    manifest=target/'manifest.json'
    raw=manifest.read_bytes()
    if b'\r\n' in raw:
        write_lf(manifest,raw.decode('utf-8').replace('\r\n','\n'))
    if len(result['epochs'])!=MIXTURE_EPOCHS:
        raise ValueError('prepare_maps did not produce the 20-epoch plan')
    report={'arm':args.arm,'directory':target.relative_to(ROOT).as_posix(),'epochs':MIXTURE_EPOCHS,
        'dataset_sha256':file_sha(manifest),
        'chat_fraction':[round(r['chat_fraction'],6) for r in result['epochs'][:2]],
        'chat_effective_passes':round(result['epochs'][0]['chat_effective_passes'],3),
        'windows_epoch1':result['epochs'][0]['windows']}
    if args.arm=='control':
        old=json.loads((OLD/'manifest.json').read_text('utf-8'))['epochs']
        if result['epochs'][:16]!=old:
            raise ValueError('Control maps differ from frozen V4.1 epochs 0-15')
        report['identical_to_v41']='epochs 0-15'
        if (V42/'manifest.json').exists():
            if result['epochs']!=json.loads((V42/'manifest.json').read_text('utf-8'))['epochs']:
                raise ValueError('Control maps differ from frozen V4.2')
            report['identical_to_v42']='epochs 0-19'
    else:
        config=tomllib.loads((ROOT/f'configs/train-v44-{args.arm}.toml').read_text('utf-8'))
        if config['data_mixture']['directory']!=report['directory']:
            raise ValueError('Config points at another mixture directory')
        if config['dataset_sha256'] and config['dataset_sha256']!=report['dataset_sha256']:
            raise ValueError(f"{args.arm} mixture SHA {report['dataset_sha256']} differs from the pinned "
                             f"{config['dataset_sha256']}; do not train, report this mismatch")
        report['matches_pinned_sha']=bool(config['dataset_sha256'])
    print(json.dumps(report),flush=True)


if __name__=='__main__':
    main()