#!/usr/bin/env python3
"""Workload generator and independent golden model for cv_x1_pro_axi.

The golden model below re-implements the CV-X1 Pro arithmetic contract from
public/forge/source/SPEC.md sections 1-2 in plain Python integers (it does not
read or execute the RTL):

    deq_i  = sext4(kv_i) - sext4(zp)                 (lane i = 0..7)
    raw    = sum_i deq_i * sext8(q_i)
    scaled = raw * uint8(scale)
    base   = 0 if clr else previous output (0 after reset)
    sum    = base + scaled
    out    = clamp32(sum) if sat else wrap32(sum); sticky |= sat and overflow

Writes hex vector files for the testbench (one value per line):
    x1_<phase>_w.hex  64 x 32-bit weight words, index = {row[3:0], bank[1:0]}
    x1_<phase>_t.hex  tiles {tlast[84], tuser[83:64], tdata[63:0]}
    x1_<phase>_e.hex  expected results {bank[34:33], tlast[32], tdata[31:0]}
    x1_counts.hex     tile counts and expected sticky flags (see COUNT_ORDER)
Usage: gen_x1_vectors.py OUT_DIR [SEED]
"""
import os
import random
import sys

INT32_MIN, INT32_MAX = -(1 << 31), (1 << 31) - 1
COUNT_ORDER = ["A", "G", "R", "B", "C"]


def s4(v):
    v &= 0xF
    return v - 16 if v & 8 else v


def s8(v):
    v &= 0xFF
    return v - 256 if v & 0x80 else v


def wrap32(v):
    v &= 0xFFFFFFFF
    return v - (1 << 32) if v & 0x80000000 else v


class Golden:
    def __init__(self, weights):
        self.w = list(weights)
        self.acc = 0          # m_axis_tdata after reset
        self.sticky = False

    def tile(self, qdata, kv_addr, scale, zp, clr, sat):
        word = self.w[kv_addr]
        raw = 0
        for i in range(8):
            kv = s4(word >> (4 * i))
            q = s8(qdata >> (8 * i))
            raw += (kv - s4(zp)) * q
        scaled = raw * (scale & 0xFF)
        base = 0 if clr else self.acc
        total = base + scaled
        if sat:
            if total > INT32_MAX:
                out, self.sticky = INT32_MAX, True
            elif total < INT32_MIN:
                out, self.sticky = INT32_MIN, True
            else:
                out = total
        else:
            out = wrap32(total)
        self.acc = out
        return out


def pack_q(qs):
    v = 0
    for i, q in enumerate(qs):
        v |= (q & 0xFF) << (8 * i)
    return v


def pack_kv(kvs):
    v = 0
    for i, k in enumerate(kvs):
        v |= (k & 0xF) << (4 * i)
    return v


def tuser(kv_addr, scale, zp, clr, sat):
    return ((kv_addr & 0x3F) | ((scale & 0xFF) << 6) | ((zp & 0xF) << 14)
            | ((clr & 1) << 18) | ((sat & 1) << 19))


class Phase:
    def __init__(self, name, weights):
        self.name = name
        self.w = list(weights)
        self.g = Golden(weights)
        self.tiles = []
        self.exp = []

    def add(self, qs, kv_addr, scale, zp, clr, sat, last):
        if not self.tiles:
            clr = 1           # each phase starts a new accumulation (phases run back to back)
        qd = pack_q(qs)
        out = self.g.tile(qd, kv_addr, scale, zp, clr, sat)
        self.tiles.append((last << 84) | (tuser(kv_addr, scale, zp, clr, sat) << 64) | qd)
        self.exp.append(((kv_addr & 3) << 33) | (last << 32) | (out & 0xFFFFFFFF))


def rand_tile(rng, extreme=False):
    if extreme and rng.random() < 0.6:
        qs = [rng.choice([-128, 127, 0, -1, 1]) for _ in range(8)]
        scale = rng.choice([0, 1, 254, 255])
        zp = rng.choice([-8, 7, 0, -1])
    else:
        qs = [rng.randint(-128, 127) for _ in range(8)]
        scale = rng.randint(0, 255)
        zp = rng.randint(-8, 7)
    return (qs, rng.randint(0, 63), scale, zp,
            int(rng.random() < 0.25), int(rng.random() < 0.5), int(rng.random() < 0.2))


def special_weights(rng):
    w = [rng.getrandbits(32) for _ in range(64)]
    w[0] = pack_kv([-8] * 8)          # all -8
    w[1] = pack_kv([7] * 8)           # all +7
    w[2] = 0
    w[3] = pack_kv([-8, 7] * 4)
    return w


def sat_run(ph, n, kv_addr, zp, q, sat, scale=255):
    for k in range(n):
        ph.add([q] * 8, kv_addr, scale, zp, int(k == 0), sat, int(k == n - 1))


def main():
    out = sys.argv[1]
    seed = int(sys.argv[2]) if len(sys.argv) > 2 else 20261106
    rng = random.Random(seed)
    os.makedirs(out, exist_ok=True)
    phases = {}

    # Phase A: random + saturation / wrap runs (weights A)
    wa = special_weights(rng)
    a = Phase("A", wa)
    for _ in range(1500):
        a.add(*rand_tile(rng))
    sat_run(a, 600, 0, 7, -128, 1)      # +3,916,800 per tile -> clamps at INT32_MAX
    for _ in range(200):
        a.add(*rand_tile(rng, extreme=True))
    sat_run(a, 600, 1, -8, -128, 1)     # -3,916,800 per tile -> clamps at INT32_MIN
    sat_run(a, 600, 0, 7, -128, 0)      # same magnitude, wrap-around mode
    sat_run(a, 600, 1, -8, -128, 0)
    for _ in range(500):
        a.add(*rand_tile(rng, extreme=True))
    phases["A"] = a

    # Phase G: three independent tiles (clr=1) for the busy-guard test.
    # They read kv addresses 0..2 only; the test rewrites address 63.
    g = Phase("G", wa)
    for k in range(3):
        g.add([rng.randint(-128, 127) for _ in range(8)], k, rng.randint(0, 255),
              rng.randint(-8, 7), 1, 1, int(k == 2))
    phases["G"] = g

    # Phase R: traffic interrupted by reset (only a prefix is checked)
    r = Phase("R", wa)
    for _ in range(1500):
        r.add(*rand_tile(rng))
    phases["R"] = r

    # Phase B: after reset the weight banks are zero
    b = Phase("B", [0] * 64)
    for _ in range(300):
        b.add(*rand_tile(rng, extreme=True))
    phases["B"] = b

    # Phase C: new weights, extreme-heavy workload incl. saturation in both directions
    wc = special_weights(rng)
    c = Phase("C", wc)
    for _ in range(1200):
        c.add(*rand_tile(rng, extreme=True))
    sat_run(c, 560, 0, 7, -128, 1)
    sat_run(c, 560, 1, -8, -128, 1)
    for _ in range(1200):
        c.add(*rand_tile(rng))
    phases["C"] = c

    for name, ph in phases.items():
        with open(os.path.join(out, f"x1_{name}_w.hex"), "w") as f:
            f.writelines(f"{x:08x}\n" for x in ph.w)
        with open(os.path.join(out, f"x1_{name}_t.hex"), "w") as f:
            f.writelines(f"{x:022x}\n" for x in ph.tiles)
        with open(os.path.join(out, f"x1_{name}_e.hex"), "w") as f:
            f.writelines(f"{x:09x}\n" for x in ph.exp)
    counts = [len(phases[n].tiles) for n in COUNT_ORDER]
    counts += [int(phases[n].g.sticky) for n in COUNT_ORDER]
    counts += [0] * (16 - len(counts))
    with open(os.path.join(out, "x1_counts.hex"), "w") as f:
        f.writelines(f"{x:08x}\n" for x in counts)
    total = sum(len(p.tiles) for p in phases.values())
    sats = sum(1 for p in phases.values() for e in p.exp
               if (e & 0xFFFFFFFF) in (0x7FFFFFFF, 0x80000000))
    print(f"gen_x1_vectors seed={seed} tiles={total} "
          + " ".join(f"{n}={len(phases[n].tiles)}" for n in COUNT_ORDER)
          + f" clamp_valued_outputs={sats} "
          + " ".join(f"sticky_{n}={int(phases[n].g.sticky)}" for n in COUNT_ORDER))


if __name__ == "__main__":
    main()