#!/usr/bin/env python3
"""Full-vocabulary distances at fixed prefixes; no target-answer likelihood score.

Requires numpy only. These metrics remain prefix dependent, not task accuracy.
All computations use float64; raw logits are centered before RMSE comparison.
"""
import argparse
import json
import shlex
import subprocess
from pathlib import Path
import numpy as np


def candidate_chunks(path, count, remote):
    expected = count*157184*4
    if not remote:
        if path.stat().st_size != expected: raise ValueError('candidate shape mismatch')
        data = np.memmap(path, dtype='<f4', mode='r', shape=(count,157184))
        for start in range(0,count,4): yield data[start:start+4]
        return
    quoted = shlex.quote(str(path))
    # Only read the remote artifact. Stream it without a second local disk copy.
    command = 'stat -c %s -- '+quoted+' && cat -- '+quoted
    with subprocess.Popen(['ssh','-o','BatchMode=yes',remote,command],stdout=subprocess.PIPE) as proc:
        try:
            if int(proc.stdout.readline()) != expected: raise ValueError('remote candidate shape mismatch')
            for start in range(0,count,4):
                rows = min(4,count-start)
                raw = proc.stdout.read(rows*157184*4)
                if len(raw) != rows*157184*4: raise ValueError('truncated remote logits')
                yield np.frombuffer(raw,dtype='<f4').reshape(rows,157184)
            if proc.stdout.read(1) or proc.wait() != 0: raise ValueError('remote logit read failed')
        finally:
            if proc.poll() is None: proc.kill()


def metrics(reference, candidate):
    r, c = np.asarray(reference, dtype=np.float64), np.asarray(candidate, dtype=np.float64)
    if not np.isfinite(r).all() or not np.isfinite(c).all(): raise ValueError('nonfinite logits')
    def distribution(x):
        shifted = x-x.max(axis=-1, keepdims=True)
        lp = shifted-np.log(np.exp(shifted).sum(axis=-1, keepdims=True))
        return lp, np.exp(lp)
    lp, p = distribution(r)
    lq, q = distribution(c)
    lm = np.logaddexp(lp, lq)-np.log(2.)
    diff = c-r
    diff -= diff.mean(axis=-1, keepdims=True)
    rtop, ctop = r.argmax(axis=-1), c.argmax(axis=-1)
    top2 = np.partition(r, -2, axis=-1)[:, -2:]
    gap = top2.max(axis=-1)-top2.min(axis=-1)
    return dict(kl=(p*(lp-lq)).sum(axis=-1),
                js=(p*(lp-lm)+q*(lq-lm)).sum(axis=-1)/2,
                tv=np.abs(p-q).sum(axis=-1)/2,
                centered_logit_rmse=np.sqrt(np.square(diff).mean(axis=-1)),
                top1_equal=(rtop == ctop).astype(np.float64),
                reference_top_gap=gap,
                reference_preference_regret=lp[np.arange(len(r)), rtop]-lp[np.arange(len(r)), ctop])


def summarize(data):
    if not data or not len(data['kl']): return None
    out = dict(tokens=len(data['kl']))
    for key, value in data.items():
        x = np.asarray(value)
        out[key] = dict(mean=float(x.mean()), p50=float(np.median(x)),
                        p95=float(np.quantile(x, .95)), max=float(x.max()))
    confident = np.asarray(data['reference_top_gap']) >= 1.
    out['reference_gap_ge_1_nat'] = dict(tokens=int(confident.sum()),
        top1_agreement=float(np.asarray(data['top1_equal'])[confident].mean()) if confident.any() else None)
    return out


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument('--suite', type=Path, required=True)
    p.add_argument('--reference', type=Path, required=True)
    p.add_argument('--candidate', type=Path, required=True)
    p.add_argument('--candidate-ssh', help='Read remote logits over SSH; compute on this host')
    p.add_argument('--suffix', default='rknn')
    p.add_argument('--output', type=Path, required=True)
    a = p.parse_args()
    suite = json.loads(a.suite.read_text())
    pooled = {'think':{}, 'nothink':{}}
    result = dict(reference=str(a.reference), candidate=str(a.candidate), candidate_ssh=a.candidate_ssh, suffix=a.suffix,
                  scope='Fixed teacher prefixes, full vocabulary, natural logarithms; not free generation or accuracy', cases=[])
    for case in suite['cases']:
        count = len(case['teacher_ids'])
        file = a.reference/(case['id']+'.bf16.f32')
        if file.stat().st_size != count*157184*4: raise ValueError('shape mismatch: '+str(file))
        reference = np.memmap(file, dtype='<f4', mode='r', shape=(count,157184))
        candidate_path = a.candidate/(case['id']+'.'+a.suffix+'.f32')
        values = {}
        for i, candidate in enumerate(candidate_chunks(candidate_path, count, a.candidate_ssh)):
            start = i*4
            for key, x in metrics(reference[start:start+4], candidate).items():
                values.setdefault(key, []).extend(x.tolist())
        mode = 'think' if case['thinking'] else 'nothink'
        for key, x in values.items(): pooled[mode].setdefault(key, []).extend(x)
        def part(lo, hi): return summarize({key:x[lo:hi] for key,x in values.items()})
        row = dict(id=case['id'], mode=mode, whole=summarize(values),
                   reasoning=part(0, case['answer_start']), answer=part(case['answer_start'], count),
                   windows={f'{i}-{min(i+256,count)}':part(i,min(i+256,count)) for i in range(0,count,256)})
        result['cases'].append(row)
        print(json.dumps(dict(case=case['id'], tv=row['whole']['tv'], kl=row['whole']['kl'])), flush=True)
    result['aggregate'] = {mode:summarize(values) for mode,values in pooled.items()}
    a.output.parent.mkdir(parents=True, exist_ok=True)
    a.output.write_text(json.dumps(result, ensure_ascii=False, indent=2)+'\n')
    print(json.dumps(result['aggregate']), flush=True)


if __name__ == '__main__': main()