#!/usr/bin/env python3
"""Generate a TINY, self-contained `wai.neural.int_hyper` conformance vector.

The pretrained int_hyper vector needs a ~23 MB PTQ model, too large to commit —
so it is regenerated per run (tools/wai_sota_vectors.sh, wai-sota-vectors.yml)
and checked on one machine at a time. This emits a SMALL synthetic conditional
codec (deterministic random integer weights) that drives the SAME decode path
end to end — factorized z -> integer hyper-decoder h_s -> integer-sigma entropy
bucketing -> integer rANS -> integer synthesis g_s -> RGB — with a committed
expected FNV-1a. A few KB, so it commits
and the int_hyper decode is CI-verifiable with no external artifacts.

  python tools/wai_int_hyper_tiny_vector.py <out-dir>

The numbers are not a "real" codec (random weights); the point is a deterministic
exercise of every integer op (convT / conv / igdn / leaky / factorized-z /
conditional-y / rANS), byte-reproducible on every machine.
"""
import sys, os, json
import numpy as np
from compressai._CXX import pmf_to_quantized_cdf
from compressai.ans import RansEncoder

# Canonical integer-net runner + helpers (the same the pretrained-model export uses).
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from wai_sota_int_synth_export import run_int_net, det_build_index_vec, fnv1a, model_to_bin

AFRAC = 12; GF = 16; SHIFT = 24
SCALE = 1 << SHIFT


def quant_cdf(pmf):
    """CompressAI quantized CDF for a pmf over n symbols (+ 1 tail) -> len n+2."""
    tail = max(1e-6, 1.0 - sum(pmf))
    return list(pmf_to_quantized_cdf([float(p) for p in pmf] + [tail], 16))


def gauss_pmf(sigma, n):
    c = n / 2.0
    p = np.exp(-((np.arange(n) - c) ** 2) / (2.0 * sigma * sigma))
    return (p / p.sum()).tolist()


def convT(rng, cin, cout, k, stride, pad, opad):
    w = rng.integers(-160, 160, (cin, cout, k, k), dtype=np.int64)
    mult = rng.integers(SCALE >> 7, SCALE >> 6, cout, dtype=np.int64)  # scale ~1/128..1/64
    bias = rng.integers(-30, 30, cout, dtype=np.int64)
    return {'type': 'convT', 'in_ch': cin, 'out_ch': cout, 'k': k, 'stride': stride, 'pad': pad,
            'opad': opad, 'weight': w.reshape(-1).tolist(), 'mult': mult.tolist(), 'bias': bias.tolist()}


def conv(rng, cin, cout, k, stride, pad):
    w = rng.integers(-160, 160, (cout, cin, k, k), dtype=np.int64)
    mult = rng.integers(SCALE >> 7, SCALE >> 6, cout, dtype=np.int64)
    bias = rng.integers(-30, 30, cout, dtype=np.int64)
    return {'type': 'conv', 'in_ch': cin, 'out_ch': cout, 'k': k, 'stride': stride, 'pad': pad,
            'weight': w.reshape(-1).tolist(), 'mult': mult.tolist(), 'bias': bias.tolist()}


def igdn(rng, c):
    gamma = rng.integers(0, max(2, (1 << GF) // (4 * c)), (c, c), dtype=np.int64)
    beta = rng.integers(1 << (2 * AFRAC - 1), 1 << (2 * AFRAC), c, dtype=np.int64)  # positive
    return {'type': 'igdn', 'c': c, 'gamma': gamma.reshape(-1).tolist(), 'beta': beta.tolist()}


def main(outdir):
    os.makedirs(outdir, exist_ok=True)
    rng = np.random.default_rng(20260625)  # fixed seed -> deterministic vector
    M, N = 4, 2                      # tiny y / z channel counts
    zh = zw = 2                      # z spatial; y spatial = 4; image = 16
    # tiny synthesis g_s: convT -> igdn -> convT (-> 3 RGB)
    g_s = {'afrac': AFRAC, 'gf': GF, 'shift': SHIFT, 'layers': [
        convT(rng, M, M, 3, 2, 1, 1), igdn(rng, M), convT(rng, M, 3, 3, 2, 1, 1)]}
    # tiny hyper-decoder h_s: convT -> leaky -> conv (-> 2M = sigma||mu)
    h_s = {'afrac': AFRAC, 'gf': GF, 'shift': SHIFT, 'layers': [
        convT(rng, N, M, 3, 2, 1, 1), {'type': 'leaky', 'slope_q': 655, 'slope_frac': 16},
        conv(rng, M, 2 * M, 3, 1, 1)]}

    # entropy tables — y Gaussian-conditional buckets + z factorized (one CDF/channel)
    scale_table = [1.0, 2.0, 4.0, 8.0]
    edges = np.floor(np.array(scale_table) * (1 << AFRAC)).astype(np.int64)
    NY = 16
    cdfs = [quant_cdf(gauss_pmf(s, NY)) for s in scale_table]
    cdf_lengths = [len(c) for c in cdfs]
    offsets = [-NY // 2] * len(scale_table)
    NZ = 16
    z_cdf = quant_cdf(gauss_pmf(2.0, NZ))
    z_cdfs = [z_cdf for _ in range(N)]
    z_cdf_lengths = [len(z_cdf)] * N
    z_offsets = [-NZ // 2] * N
    z_median_q = [0] * N

    # bitstream symbols (deterministic), kept in range of the tiny CDFs
    z_sym = rng.integers(-3, 4, (N, zh, zw), dtype=np.int64)      # factorized z symbols
    z_q = (z_sym << AFRAC)                                        # + median_q (0)
    sm = run_int_net(h_s, z_q)                                    # [2M, 4, 4]
    sigma_int, mu_int = sm[:M], sm[M:]
    sh, sw = sigma_int.shape[1], sigma_int.shape[2]
    resid = rng.integers(-3, 4, (M, sh, sw), dtype=np.int64)      # y residuals
    indexes = det_build_index_vec(sigma_int, edges)
    y_hat = (resid << AFRAC) + mu_int
    rgb_chw = run_int_net(g_s, y_hat)                             # [3, 16, 16] @ 2^afrac
    one = 1 << AFRAC
    out_c = np.clip(rgb_chw, 0, one)
    pix = (((out_c * 255 + (1 << (AFRAC - 1))) >> AFRAC)).clip(0, 255).astype(np.uint8)
    rgb = np.ascontiguousarray(pix.transpose(1, 2, 0))           # HWC u8
    rgb_hash = format(fnv1a(rgb.tobytes()), '016x')

    enc = RansEncoder()
    zc = np.repeat(np.arange(N), zh * zw)
    z_bits = enc.encode_with_indexes(z_sym.reshape(-1).tolist(), zc.tolist(), z_cdfs, z_cdf_lengths, z_offsets)
    enc2 = RansEncoder()
    y_bits = enc2.encode_with_indexes(resid.reshape(-1).tolist(), indexes.reshape(-1).tolist(),
                                      cdfs, cdf_lengths, offsets)

    json.dump(g_s, open(f'{outdir}/tiny_g_s.json', 'w'))
    json.dump(h_s, open(f'{outdir}/tiny_h_s.json', 'w'))
    open(f'{outdir}/tiny_g_s.bin', 'wb').write(model_to_bin(g_s))
    open(f'{outdir}/tiny_h_s.bin', 'wb').write(model_to_bin(h_s))
    json.dump({'afrac': AFRAC, 'frac': AFRAC, 'scale_edges': edges.tolist(),
               'cdfs': cdfs, 'cdf_lengths': cdf_lengths, 'offsets': offsets,
               'z_cdfs': z_cdfs, 'z_cdf_lengths': z_cdf_lengths, 'z_offsets': z_offsets,
               'z_median_q': z_median_q}, open(f'{outdir}/tiny_entropy.json', 'w'))
    json.dump({'m': M, 'sh': int(sh), 'sw': int(sw), 'n_z': N, 'zh': zh, 'zw': zw,
               'z_rans': list(z_bits), 'rans': list(y_bits)}, open(f'{outdir}/tiny_bitstream.json', 'w'))
    json.dump({'rgb_hash': rgb_hash, 'out_c': 3, 'out_h': int(pix.shape[1]), 'out_w': int(pix.shape[2])},
              open(f'{outdir}/tiny.json', 'w'))
    sz = sum(os.path.getsize(f'{outdir}/{f}') for f in os.listdir(outdir))
    print(f'tiny int_hyper vector: image {pix.shape[1]}x{pix.shape[2]}, '
          f'{len(z_bits)} z + {len(y_bits)} y rANS bytes, rgb_hash {rgb_hash}, total {sz} B')
    return rgb, int(pix.shape[1]), int(pix.shape[2])  # the decoded I-frame, for the video keyframe encoder


if __name__ == '__main__':
    main(sys.argv[1])