#!/usr/bin/env python3
"""Check that the `wai.image.jpegai` exact decoder's gain scaler equals the reference
decoder's, although the two compute it in different arithmetic.

The per-channel gain scaler is

    scaler[c] = round( exp( (g[c] + beta) * log_k / 2^7 ) * 2^10 ) / 2^10

with `log_k = (ln 54.82 - ln 0.11) / 31`. The two decoders evaluate it differently.

  Exact decoder (`jpegai_pipeline::gain_unit_scaler_vec`): f64 throughout, with
  `GAIN_LOG_K` as its f64 `log_k`, `exp` from `wai-rs/src/float_repro.rs` (built from
  IEEE-754 basic operations, so its bits are the same on every target and within a
  few ULPs of the exact value), rounding half away from zero. It calls this with
  `beta = 0` only: `jpegai_decode::check_codestream` refuses any codestream whose
  `beta_displacement_log` is not 0, because the beta displacement is not implemented.

  Reference software: float32. The log-domain gain `g + beta` is an int32 tensor and
  `log_k` an f64 scalar, so the product and everything after it are float32:
      x32 = f32( f32(g + beta) * f32(log_k) ) / 2^7      (the divide is exact)
      t   = exp32(x32) * 2^10                            (the multiply is exact)
      scaler = round_half_even(t) / 2^10
  `exp32` is the float32 `exp` of its tensor library on the platform it runs on,
  which is not required to be correctly rounded.

The multiplies, divides and conversions are IEEE-754 basic operations, so both `x`
values are fixed. Only the two `exp`s vary, and each can move the scaler only if the
exact value `exp(x) * 2^10` lies close to a rounding boundary (k + 0.5). This script
computes those exact values with 60-digit decimal arithmetic and measures:

  * decoder margin: distance from the boundary in ULPs of the f64 value. The
    decoder's `exp` is within a few ULPs, so this must be large.
  * reference margin: distance from the boundary in float32 ULPs (of the larger of
    the value and the boundary). Any float32 `exp` whose error is below this many
    ULPs gives the scaler of the exact value. The script also reports what a
    correctly rounded float32 `exp` gives, since `t` can land exactly on k + 0.5.
  * that both exact values round to the same integer.

It fails unless `GAIN_LOG_K` is the f64 `log_k` the reference computes (f64
logarithms, correctly rounded, then the difference, then / 31), and unless, at
beta = 0 over every table entry, the decoder margin is at least MIN_MARGIN_ULPS,
the reference margin at least MIN_MARGIN_F32_ULPS, and both round to the same
integer.

    python wai_jpegai_gain_exp_margin.py <assets_dir> [--beta-sweep]

<assets_dir> is the jpegai-vm-assets root (the decoder assets exported from the
reference software; see wai_jpegai_build_e2e_inputs.py). Every
model/CCS_SGMM/tools_*/model_{y,uv} gain table present is checked; the decoder's
scope is tools_0. GAIN_LOG_K is read from the Rust source so the check follows the
code. Standard library only.

--beta-sweep also reports, without gating, the same measures over g + beta for every
beta the 12-bit header field can carry (-2048..2047): what a decoder that implements
the beta displacement would face.
"""
import glob
import math
import os
import re
import struct
import sys
from decimal import ROUND_FLOOR, ROUND_HALF_EVEN, Decimal, getcontext

getcontext().prec = 60
MIN_MARGIN_ULPS = 1e6  # f64 ULPs; far beyond float_repro::exp's error (≤ 2 ULPs, tested)
MIN_MARGIN_F32_ULPS = 8.0  # float32 ULPs: several times a float32 exp accurate to ~1 ULP
SCALER_PRECISION = 10
SIGMA_PRECISION = 7

HERE = os.path.dirname(os.path.abspath(__file__))
PIPELINE_RS = os.path.join(HERE, "..", "wai-rs", "src", "jpegai_pipeline.rs")


def gain_log_k():
    src = open(PIPELINE_RS).read()
    m = re.search(r"const GAIN_LOG_K: f64 = ([0-9_.]+);", src)
    if not m:
        sys.exit(f"GAIN_LOG_K not found in {PIPELINE_RS}")
    return float(m.group(1).replace("_", ""))


def reference_log_k():
    # (ln 54.82 - ln 0.11) / 31 as the reference computes it in f64: each logarithm
    # rounded to f64 (here correctly), then the difference, then the division.
    ln = lambda v: float(Decimal(v).ln())  # Decimal(v) is the exact f64 value of v
    return (ln(54.82) - ln(0.11)) / 31


def f32(v):
    return struct.unpack("<f", struct.pack("<f", v))[0]  # round to nearest even


def ulp32(v):
    return 2.0 ** (math.frexp(v)[1] - 24)  # float32 spacing at positive normal v


def boundary(v):
    """The rounding boundary nearest to v > 0 (a half-integer) and the distance."""
    b = Decimal(int(v)) + Decimal("0.5")
    return b, abs(v - b)


def decoder_side(n, k):
    # Same f64 operations, same order, as the Rust: (((g + beta) as f64) * K) / 128.0.
    x = (float(n) * k) / 2.0 ** SIGMA_PRECISION
    v = Decimal(x).exp() * (1 << SCALER_PRECISION)  # exact to 60 digits
    _, d = boundary(v)
    r = int((v + Decimal("0.5")).to_integral_value(rounding=ROUND_FLOOR))
    return r, float(d) / math.ulp(float(v))


def reference_side(n, k_ref):
    # The product of two float32 values is exact in f64, so one rounding gives f32(a*b).
    x = f32(f32(float(n)) * f32(k_ref)) / 2.0 ** SIGMA_PRECISION
    e = Decimal(x).exp()
    v = e * (1 << SCALER_PRECISION)
    b, d = boundary(v)
    r = int(v.to_integral_value(rounding=ROUND_HALF_EVEN))
    margin = float(d) / ulp32(max(float(v), float(b)))
    u = Decimal(ulp32(float(e)))
    t = (e / u).to_integral_value(rounding=ROUND_HALF_EVEN) * u * (1 << SCALER_PRECISION)
    r_cr = int(t.to_integral_value(rounding=ROUND_HALF_EVEN))  # correctly rounded exp32
    return r, margin, r_cr


def measure(xs, k, k_ref):
    out = {}
    for n in xs:
        rd, md = decoder_side(n, k)
        rr, mr, rcr = reference_side(n, k_ref)
        out[n] = (rd, md, rr, mr, rcr)
    return out


def summary(res):
    md = min(r[1] for r in res.values())
    mr_n = min(res, key=lambda n: res[n][3])
    return (md, res[mr_n][3], mr_n,
            sorted(n for n, r in res.items() if r[0] != r[2]),
            sorted(n for n, r in res.items() if r[3] < 1.0),
            sorted(n for n, r in res.items() if r[0] != r[4]))


BETA_RANGE = range(-2048, 2048)  # 12-bit beta_displacement_log, de-biased by 2048


def main():
    args = sys.argv[1:]
    sweep = "--beta-sweep" in args
    args = [a for a in args if a != "--beta-sweep"]
    if len(args) != 1:
        sys.exit(__doc__)
    k = gain_log_k()
    k_ref = reference_log_k()
    pattern = os.path.join(args[0], "model", "CCS_SGMM", "tools_*", "model_*",
                           "common_modules", "quantizer", "gain_unit_mlog.csv")
    tables = sorted(glob.glob(pattern))
    if not tables:
        sys.exit(f"no gain tables under {pattern}")
    print(f"GAIN_LOG_K = {k!r}; the reference's f64 log_k = {k_ref!r}; "
          f"as float32 {f32(k_ref)!r}")
    fail = []
    if k != k_ref:
        fail.append(f"GAIN_LOG_K differs from the reference's log_k by {k - k_ref:.3e} "
                    f"({(k - k_ref) / math.ulp(k_ref):.3e} f64 ULPs)")
    worst_d, worst_r = math.inf, math.inf
    swept = {}
    for path in tables:
        gains = [int(s) for s in open(path).read().split()]
        rel = os.path.relpath(os.path.dirname(path), args[0]).split(os.sep + "common_modules")[0]
        md, mr, mr_n, differ, under1, cr_differ = summary(measure(gains, k, k_ref))
        worst_d, worst_r = min(worst_d, md), min(worst_r, mr)
        print(f"{rel}: {len(gains)} entries, range [{min(gains)}, {max(gains)}]")
        print(f"  beta = 0: decoder margin {md:.3e} f64 ULPs; reference margin "
              f"{mr:.2f} float32 ULPs (at g = {mr_n}); exact values round differently "
              f"at {differ or 'none'}")
        if differ:
            fail.append(f"{rel}: decoder and reference scalers differ at g = {differ}")
        if sweep:
            xs = {g + b for g in gains for b in BETA_RANGE}
            new = [x for x in xs if x not in swept]
            swept.update(measure(new, k, k_ref))
            md, mr, mr_n, differ, under1, cr_differ = summary({x: swept[x] for x in xs})
            print(f"  beta in [{BETA_RANGE.start}, {BETA_RANGE.stop - 1}] (not gated): "
                  f"{len(xs)} distinct g + beta; decoder margin {md:.3e} f64 ULPs; "
                  f"reference margin {mr:.2f} float32 ULPs (at g + beta = {mr_n})")
            print(f"    reference margin under 1 float32 ULP at g + beta = "
                  f"{under1 or 'none'}; exact values round differently at "
                  f"{differ or 'none'}; a correctly rounded float32 exp gives a "
                  f"different scaler from the decoder's at {cr_differ or 'none'}")
    print(f"beta = 0, all tables: decoder margin {worst_d:.3e} f64 ULPs (threshold "
          f"{MIN_MARGIN_ULPS:.0e}); reference margin {worst_r:.2f} float32 ULPs "
          f"(threshold {MIN_MARGIN_F32_ULPS:g})")
    if sweep:
        md, mr, mr_n, differ, under1, cr_differ = summary(swept)
        print(f"beta sweep, all tables (not gated): {len(swept)} distinct g + beta; "
              f"decoder margin {md:.3e} f64 ULPs; reference margin {mr:.2f} float32 "
              f"ULPs (at {mr_n}); under 1 float32 ULP at {under1 or 'none'}; exact "
              f"values round differently at {differ or 'none'}; correctly rounded "
              f"float32 exp differs from the decoder at {cr_differ or 'none'}")
    if worst_d < MIN_MARGIN_ULPS:
        fail.append("an entry is close enough to a rounding boundary that a last-place "
                    "difference in the decoder's exp could change its scaler")
    if worst_r < MIN_MARGIN_F32_ULPS:
        fail.append("an entry is close enough to a rounding boundary that the "
                    "reference's float32 exp could change its scaler")
    if fail:
        sys.exit("FAIL: " + "; ".join(fail))
    print(f"OK: at beta = 0 the decoder's scaler equals the reference's for every "
          f"float32 exp within {math.floor(worst_r)} ULPs of the exact value")


if __name__ == "__main__":
    main()