"""One allocation-reduction attempt, with a similarly optimized baseline.

Does not change weights or run the network again. Checks bitwise equality to the
network-tested factorization on every captured real Q/K input before timing.
"""
import argparse
import json
from pathlib import Path
import resource
import time
import numpy as np
from .model import StreamingLM
from .gain_factorization import consolidate,normalized,LAYERS
from .run import save_csv,write_json


def baseline_inplace(q,k,a,b):
    q=normalized(q);q*=a
    k=normalized(k);k*=b
    return q,k


def factored_inplace(q,k,a,b,r,combined,move='q_to_k'):
    q,k=normalized(q),normalized(k)
    if move=='q_to_k':
        q[...,:r]*=a[:r];k*=combined
    else:
        q*=combined;k[...,:r]*=b[:r]
    return q,k


def main():
    p=argparse.ArgumentParser();p.add_argument('--root',required=True,type=Path);a=p.parse_args();root=a.root
    ledger_path=root/'compute-budget-ledger.json';ledger=json.loads(ledger_path.read_text());remaining=int(7200-ledger['spent_seconds']-20)
    if remaining<60:raise TimeoutError('Cumulative budget insufficient')
    start=time.monotonic();result={'state':'running'}
    try:
        cfg=json.loads((root/'model_config.json').read_text());lm=StreamingLM(root/'weights',cfg,seconds=min(remaining,300))
        r=64;data=np.load(root/'baseline-activations.npz');checks=[]
        for layer in LAYERS:
            lm.guard();pfx=f'layers.{layer}.self_attn.'
            ag,bg=1+lm.small(pfx+'q_norm.weight'),1+lm.small(pfx+'k_norm.weight')
            q,k=data[f'qraw_{layer}'],data[f'kraw_{layer}']
            for move in ('q_to_k','k_to_q'):
                combined=np.concatenate([(bg if move=='q_to_k' else ag)[:r],ag[r:]*bg[r:]])
                expected=consolidate(q,k,ag,bg,r,move)
                actual=factored_inplace(q,k,ag,bg,r,combined,move)
                for x,y in zip(expected,actual):np.testing.assert_array_equal(x,y)
                checks.append({'layer':layer,'direction':move,'query_elements':int(q.size),'key_elements':int(k.size),'bitwise_equal':True})
        q,k=data['qraw_3'][:1],data['kraw_3'][:1];pfx='layers.3.self_attn.'
        ag,bg=1+lm.small(pfx+'q_norm.weight'),1+lm.small(pfx+'k_norm.weight')
        combined=np.concatenate([bg[:r],ag[r:]*bg[r:]])
        funcs={'baseline_legacy':lambda:(normalized(q)*ag,normalized(k)*bg),
               'baseline_inplace':lambda:baseline_inplace(q,k,ag,bg),
               'factored_inplace':lambda:factored_inplace(q,k,ag,bg,r,combined)}
        for _ in range(3):
            for fn in funcs.values():fn()
        timing=[];names=list(funcs)
        for repeat in range(30):
            order=names[repeat%3:]+names[:repeat%3]
            for name in order:
                lm.guard();t=time.perf_counter()
                for _ in range(50):funcs[name]()
                timing.append({'repeat':repeat,'implementation':name,'seconds_per_call':(time.perf_counter()-t)/50,'calls_per_block':50})
        save_csv(root/'gain-factorization-optimized-timing.csv',timing)
        values={name:np.array([row['seconds_per_call'] for row in timing if row['implementation']==name]) for name in names}
        ratio=values['baseline_inplace']/values['factored_inplace']
        medians={name:float(np.median(v)) for name,v in values.items()}
        result.update(state='completed',bitwise_checks=checks,medians_seconds=medians,
                      paired_baseline_over_factored_median=float(np.median(ratio)),
                      paired_ratio_p10=float(np.quantile(ratio,.1)),paired_ratio_p90=float(np.quantile(ratio,.9)),
                      fraction_blocks_factored_faster=float(np.mean(ratio>1)),
                      descriptive_local_gain=bool(np.median(ratio)>1.05 and np.mean(ratio>1)>=.8 and np.quantile(ratio,.1)>1),
                      claim_limit='Hot-array CPU normalization microbenchmark against equally allocation-optimized baseline; no full-model speedup or independent statistical significance claim.')
    except Exception as exc:
        result.update(state='stopped',error_type=type(exc).__name__,error=str(exc));raise
    finally:
        elapsed=time.monotonic()-start;result.update(elapsed_seconds=elapsed,peak_rss_kib=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss)
        ledger=json.loads(ledger_path.read_text());ledger['spent_seconds']+=elapsed;ledger['calls'].append({'module':'functional.benchmark_factorization','seconds':elapsed});write_json(ledger_path,ledger)
        write_json(root/'gain-factorization-optimized-summary.json',result)


if __name__=='__main__':main()