#!/usr/bin/env python3
"""Sequential before/after KMeans measurement; retain all samples and evidence.

Build tests/kmeans_benchmark.cpp against preserved baseline ml.cpp with
ANER_KMEANS_BASELINE, and against the corrected library without that macro.
Finish builds/tests before invoking. No speed or RSS claim is pre-populated.
"""
import argparse
from datetime import datetime
from zoneinfo import ZoneInfo
import hashlib
import json
import os
from pathlib import Path
import platform
import statistics
import subprocess
import sys
import time

CASES = [
    ('tiny','Tiny',32,1,4,2,20,'balanced'),
    ('batch','Batch',4096,32768,8,4,20,'balanced'),
    ('wide','Wide',2048,4096,32,8,20,'balanced'),
    ('larger','Larger',8192,32768,16,16,20,'balanced'),
    ('many','Many clusters',4096,8192,8,32,20,'balanced'),
    ('everyrow','Every row a center',128,4096,4,128,10,'balanced'),
    ('batch-overlap','Overlapping batch',4096,32768,8,4,20,'overlap'),
    ('batch-skew','Skewed batch',4096,32768,8,4,20,'skew'),
    ('batch-cap1','One-update batch',4096,32768,8,4,1,'balanced'),
]
SOURCES=['src/ml.cpp','src/ml_exact.hpp','src/parallel.cpp','include/aner/ml.hpp',
         'include/aner/parallel.hpp','include/aner/random.hpp','src/nn.cpp','src/tensor.cpp',
         'tests/kmeans_benchmark.cpp','tools/kmeans_benchmark.py','CMakeLists.txt']

def sha(path): return hashlib.sha256(Path(path).read_bytes()).hexdigest()
def main():
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--baseline',required=True,type=Path)
    parser.add_argument('--corrected',required=True,type=Path)
    parser.add_argument('--source-root',required=True,type=Path)
    parser.add_argument('--baseline-source',required=True,type=Path)
    parser.add_argument('--output',required=True,type=Path)
    parser.add_argument('--repeats',default=7,type=int)
    parser.add_argument('--rss-repeats',default=3,type=int)
    parser.add_argument('--oracle-report',required=True,type=Path)
    parser.add_argument('--integration-report',required=True,type=Path)
    parser.add_argument('--native-checks',required=True,type=int)
    args=parser.parse_args()
    if args.output.exists(): parser.error('output exists: preserve raw evidence')
    if args.repeats<5 or args.rss_repeats<1: parser.error('need >=5 timing and >=1 RSS repetitions')
    cpu=platform.processor()
    if platform.system()=='Darwin': cpu=subprocess.check_output(['/usr/sbin/sysctl','-n','machdep.cpu.brand_string'],text=True).strip()
    compiler=subprocess.check_output(['/usr/bin/c++','--version'],text=True).strip()
    oracle=json.loads(args.oracle_report.read_text())
    integration=json.loads(args.integration_report.read_text())
    if not oracle.get('passed') or not integration.get('passed'): raise RuntimeError('independent numerical validation must pass')
    def portable(value):
        if isinstance(value,dict): return {key:portable(item) for key,item in value.items()}
        if isinstance(value,list): return [portable(item) for item in value]
        if isinstance(value,str): return value.replace(str(args.source_root.resolve().parent.parent),'$STUDY_ROOT')
        return value
    evidence={'schema' :'aner.kmeans-evidence/1','status':'verified','recorded_date':datetime.now(ZoneInfo('America/Denver')).date().isoformat(),
      'timestamp_utc':time.strftime('%Y-%m-%dT%H:%M:%SZ',time.gmtime()),
      'machine':{'system':platform.system(),'release':platform.release(),'platform':platform.platform(),
                 'machine':platform.machine(),'cpu_model':cpu,'logical_cpu_count':os.cpu_count()},
      'build':{'cpu_profile':'apple-m1','configuration':'Release','compiler':compiler,'numerical_flags':['-fno-fast-math','-ffp-contract=off'],
               'optimization_flags':['-O3','-DNDEBUG','-mcpu=apple-m1'],'lto':'ThinLTO',
               'baseline_source_revision':'32e34b92ef103156c9919212999b97b020e65aba',
               'baseline_scope':'Preserved baseline ml.cpp with unchanged Tensor/NN dependencies and the same current executor library; legacy KMeans does not invoke that executor.',
               'baseline_ml_sha256':sha(args.baseline_source),
               'sources':{p:sha(args.source_root/p) for p in SOURCES},
               'binaries':{'legacy':sha(args.baseline),'corrected':sha(args.corrected)}},
      'validation':{'numerical_oracle':{'passed':True,'cases':oracle['case_count'],
                       'method':'Independent Python Fraction from binary64 values; exact sum/division and exact-square accumulation; compare final Float64 bits.',
                       'report':portable(oracle)},'integration_oracle':{'passed':True,'report':portable(integration)},
                    'determinism':{'passed':True,'checks':args.native_checks,
                       'scope':'Corrected serial versus explicit 1/2/4 workers: all center/label/inertia bits and iterations/convergence, prediction bits. Same binary/platform FP contract.'},
                    'native_tests':{'passed':True,'scope':'Release numerical/parallel/language/session and existing ML/KNN/decision-tree/executor regressions; sanitizer checks recorded separately.'}},
      'methodology':[
        'All cases and configurations run sequentially after builds/tests finish. Seven raw timed repetitions after two warmups.',
        'Fit timer includes seeded initialization, all updates, final labels/inertia and workspace allocation. Prediction timer excludes fitting.',
        'Inputs are already materialized CPU Float64. Parsing, fixture creation and correctness checks are outside timed regions.',
        'Each result is bitwise checked after its timer; corrected execution modes use the corrected serial model as reference.',
        'Legacy mathematics is intentionally a distinct numerical baseline; compare completed iterations, not just elapsed time.',
        'Cold pool timers include pool creation/destruction. Reused pools are created before timing. Serial uses no pool.',
        'One-update batch fixes update count to separate numerical changes in stopping from arithmetic/scheduling cost.',
        'Workers are native pool concurrency limits, not measured CPU utilization or dedicated cores.',
        'No benchmark command runs publishing, remote SSH, cache cleanup or package distribution.',
      ],
      'workspace':{'records':[],'default_cap_bytes':64*1024*1024,'measurement_kind':'minimum accepted max_workspace_bytes confirmed with exact and minus-one budgets',
        'included':'Fit new model centers/labels, live vector capacities, one fixed exact accumulator per allocated lane, 272 B active mean finalization temporary per mean lane, bounded owner/control objects. Prediction output/lane controls; existing model excluded.',
        'excluded':'Input/host-conversion caches, executor internals and thread stacks, observer/context snapshot payloads, opaque shared ownership and allocator overhead, ordinary call-stack frames, process RSS.',
        'legacy_logical_fit_payload_formula_64bit':'24*N + 40*K*D + 8*K bytes, excluding vector/owner controls, allocator overhead and input.',
        'cap_is_allocation':False},
      'rss':{'records':[],'measurement_kind':'fresh-process getrusage lifetime peak',
        'aggregation':'maximum of raw independent fresh-process lifetime peaks',
        'notes':['macOS getrusage reports bytes; Linux KiB is converted to bytes.',
                 'Full process lifetime includes runtime, input/query tensors, executor threads and the operation. Prediction peak includes its prerequisite serial fit.',
                 'RSS-only processes avoid oracle model copies/timing repetitions; no subtraction produces a per-call allocation estimate.']},
      'artifacts':[], 'workloads':[]}
    def invoke(binary,case,operation,mode,workers,workspace=False,rss=False):
        ident,name,n,q,d,k,cap,distribution=case
        command=[str(binary.resolve()),'--n',str(n),'--q',str(q),'--d',str(d),'--k',str(k),'--cap',str(cap),
            '--distribution',distribution,'--seed','42','--operation',operation,'--mode',mode,'--workers',str(workers),
            '--repeats',str(args.repeats),'--warmups','2','--workspace',str(int(workspace)),'--rss-only',str(int(rss))]
        result=subprocess.run(command,capture_output=True,text=True,check=True,timeout=300)
        record=json.loads(result.stdout)
        if result.stderr: raise RuntimeError(result.stderr)
        return record
    for case in CASES:
        ident,name,n,q,d,k,cap,distribution=case
        workload={'id':ident,'name':name,'n':n,'q':q,'d':d,'k':k,'max_iterations':cap,'seed':42,'distribution':distribution}
        for operation in ('fit','predict'):
            old=invoke(args.baseline,case,operation,'serial',1)
            new=invoke(args.corrected,case,operation,'serial',1,workspace=ident in ('batch','wide'))
            study={'legacy':old,'serial':new,'pools':[],'coldPools':[]}
            for workers in (1,2,4):
                study['pools'].append(invoke(args.corrected,case,operation,'context',workers,workspace=ident in ('batch','wide')))
                study['coldPools'].append(invoke(args.corrected,case,operation,'cold',workers))
            for row in [new]+study['pools']:
                if row['charged_workspace_bytes']:
                    evidence['workspace']['records'].append({'workload':ident,'operation':operation,**{key:row[key] for key in ('workers','charged_workspace_bytes','active_lanes','exact_budget_accepted','minus_one_rejected')}})
            for row in [old,new]+study['pools']+study['coldPools']:
                row['speedup_vs_corrected_serial']=new['median_seconds']/row['median_seconds']
                row['speedup_vs_legacy_serial']=old['median_seconds']/row['median_seconds']
                row['speedup_vs_pool1']=study['pools'][0]['median_seconds']/row['median_seconds']
            workload[operation]=study
            if ident in ('batch','wide'):
                for flavor,binary,mode,workers in [('legacy',args.baseline,'serial',1),('corrected',args.corrected,'serial',1)]+[('corrected',args.corrected,'context',w) for w in (1,2,4)]:
                    samples=[invoke(binary,case,operation,mode,workers,rss=True)['peak_rss_bytes'] for _ in range(args.rss_repeats)]
                    evidence['rss']['records'].append({'workload':ident,'operation':operation,'flavor':flavor,'mode':mode,
                       'workers':0 if mode=='serial' else workers,'fresh_process':True,'raw_peak_rss_bytes':samples,
                       'peak_rss_bytes':max(samples),'measurement_kind':'fresh-process getrusage lifetime peak'})
            print(f'{name:20} {operation:7} old {old["median_seconds"]*1000:.4f} ms new {new["median_seconds"]*1000:.4f} ms pool4 {study["pools"][-1]["median_seconds"]*1000:.4f} ms ({new["median_seconds"]/study["pools"][-1]["median_seconds"]:.3f}x vs corrected serial), I={old["iterations"]}/{new["iterations"]}',file=sys.stderr,flush=True)
        evidence['workloads'].append(workload)
    # Fingerprint only public, reviewable source/sample evidence. No native
    # packages, private connection settings, binaries or runtime data are copied.
    public_files=['tests/kmeans_benchmark.cpp','tools/kmeans_benchmark.py','tests/kmeans_oracle_probe.cpp',
                  'tools/verify_kmeans_numeric.py','tools/verify_kmeans_exact.py','tests/ml_exact_tests.cpp','src/ml_exact.hpp']
    for rel in public_files:evidence['artifacts'].append({'path':'/reference/'+rel,'sha256':sha(args.source_root/rel)})
    evidence['artifacts'].append({'path':'/reference/docs/benchmarks/kmeans_legacy_ml_32e.cpp','sha256':sha(args.baseline_source)})
    if evidence['build']['sources']!={p:sha(args.source_root/p) for p in SOURCES}: raise RuntimeError('benchmark sources changed during capture')
    args.output.parent.mkdir(parents=True,exist_ok=True)
    with args.output.open('x') as output:json.dump(evidence,output,indent=2);output.write('\n')
    print(args.output.resolve())
if __name__=='__main__':main()