#!/usr/bin/env python3
"""Locate representable-range clipping in exported KDA W4 input matrices."""
import argparse
import json
from pathlib import Path
import numpy as np
import torch
from audit_absolute_weights import Package, TensorSource, source_linear_names, family, blank, add, measure, finish


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument('--source', type=Path, required=True)
    p.add_argument('--original', type=Path, required=True)
    p.add_argument('--calibrated', type=Path, required=True)
    p.add_argument('--output', type=Path, required=True)
    a = p.parse_args()
    torch.set_num_threads(2)
    src = TensorSource(a.source)
    packs = {v:Package(getattr(a,v)) for v in ('original','calibrated')}
    results = []
    for name,e in packs['original'].entries.items():
        if e[4] != 3: continue
        base = name.removesuffix('.weight')
        if family(base) != 'kda_in': continue
        offset = 0
        for key in source_linear_names(base):
            t = src.tensor(key)
            acc = {v:blank() for v in packs}
            extra = {v:dict(out_of_range=0, clipped_error_sum_sq=0., worst=None) for v in packs}
            for lo in range(0,len(t),256):
                hi = min(lo+256,len(t))
                ref = t[lo:hi].float().numpy()
                for v,pack in packs.items():
                    q = pack.rows(base, offset+lo, offset+hi)
                    s = pack.array(base+'.scales','<f4')[offset+lo:offset+hi,None]
                    clipped = (ref > 7*s) | (ref < -8*s)
                    d = q.astype(np.float64)-ref.astype(np.float64)
                    m = measure(ref,q)
                    add(acc[v],m)
                    extra[v]['out_of_range'] += int(clipped.sum())
                    extra[v]['clipped_error_sum_sq'] += float(np.square(d[clipped]).sum())
                    i,j = np.unravel_index(np.abs(d).argmax(),d.shape)
                    worst = dict(row=int(lo+i), column=int(j), original=float(ref[i,j]),
                                 reconstructed=float(q[i,j]), scale=float(s[i,0]),
                                 abs_error=float(abs(d[i,j])), outside_range=bool(clipped[i,j]))
                    if extra[v]['worst'] is None or worst['abs_error'] > extra[v]['worst']['abs_error']:
                        extra[v]['worst'] = worst
            results.append(dict(name=key, matrix=base, metrics={v:dict(**finish(acc[v]),**extra[v]) for v in packs}))
            offset += len(t)
    a.output.parent.mkdir(parents=True,exist_ok=True)
    a.output.write_text(json.dumps(dict(scope='KDA input only; outside [-8s,7s] inclusive FP32 boundaries', tensors=results),indent=2)+'\n')
    for v in packs:
        count = sum(r['metrics'][v]['count'] for r in results)
        clipped = sum(r['metrics'][v]['out_of_range'] for r in results)
        print(json.dumps(dict(variant=v,count=count,clipped=clipped,clipped_fraction=clipped/count,
            worst=max(results,key=lambda r:r['metrics'][v]['worst']['abs_error']))),flush=True)


if __name__ == '__main__': main()