#!/usr/bin/env python3
"""
build_baseline.py — Generate markov_baseline.bin for the T1 RID sensor.

Simulates 100,000 benign RID events using the same cadence priors as
markov_rid.py (steady heartbeats, periodic ID_BROADCAST, rare GAP), counts
transitions, Laplace-smooths, computes -log probabilities, quantises to
int16 Q1.15 scaled to NLL_SCALE=8.0 nats, and writes the binary file that
markov_init() in markov_scorer.cpp expects.

Output file layout:
   magic (4 bytes "MKV1", LE)
   version (uint16 LE, = 1)
   n_states (uint16 LE, = 35)
   transition matrix [35 * 35] int16 LE, row-major (from-state rows)

File size = 8 + 2 * 35 * 35 = 2458 bytes

Usage:
   python build_baseline.py [--out ../../data/models/markov_baseline.bin]
"""

import argparse
import math
import pathlib
import random
import struct
import sys

# ---------------------------------------------------------------------------
# Mirror the C++ enum and IAT bucket definitions
# ---------------------------------------------------------------------------
MK_BEACON       = 0
MK_HEARTBEAT    = 1
MK_TELEMETRY    = 2
MK_ID_BROADCAST = 3
MK_HANDOFF      = 4
MK_LOSS         = 5
MK_GAP          = 6
MK_EVENT_COUNT  = 7

MK_IAT_BUCKETS = 5
MK_N_STATES    = MK_EVENT_COUNT * MK_IAT_BUCKETS  # 35

# IAT bucket edges in seconds (mirrors C++ IAT_EDGES_MS / 1000)
IAT_EDGES_S = (0.1, 0.5, 2.0, 10.0)

MAGIC   = b'MKV1'
VERSION = 1

# NLL scale used in C++ q15_to_nll() — must match exactly
NLL_SCALE = 8.0


def iat_bucket(dt_s: float) -> int:
    for i, e in enumerate(IAT_EDGES_S):
        if dt_s < e:
            return i
    return MK_IAT_BUCKETS - 1


def state_index(ev: int, iat: int) -> int:
    return ev * MK_IAT_BUCKETS + iat


# ---------------------------------------------------------------------------
# Benign event generator
# Cadence priors (matches markov_rid.py _demo / gen_benign_track):
#   HEARTBEAT    : every ~1 s  (± 50 ms jitter)
#   TELEMETRY    : every ~3 s  (± 150 ms jitter)
#   ID_BROADCAST : every ~10 s (± 400 ms jitter)
#   HANDOFF      : very rare, ~every 120 s
#   GAP          : rare, ~1 in 500 heartbeat intervals
#   LOSS         : very rare, ~1 in 2000 events
#   BEACON       : at start of each track only
# ---------------------------------------------------------------------------
def gen_benign_events(n_events: int, rng: random.Random):
    """Yield (event_type, timestamp_s) pairs for n_events total."""
    t = 0.0
    hb_next = t + 1.0
    tm_next = t + 3.0
    id_next = t + 10.0
    ho_next = t + 120.0
    count   = 0

    # Start each simulated track with a BEACON
    yield MK_BEACON, t
    count += 1

    while count < n_events:
        candidates = [
            (hb_next, MK_HEARTBEAT),
            (tm_next, MK_TELEMETRY),
            (id_next, MK_ID_BROADCAST),
            (ho_next, MK_HANDOFF),
        ]
        candidates.sort()
        t, ev = candidates[0]

        # Occasional GAP injection (low probability)
        if ev == MK_HEARTBEAT and rng.random() < 0.002:
            ev = MK_GAP
            t += rng.uniform(0.5, 3.0)

        # Very occasional LOSS
        if rng.random() < 0.0005:
            yield MK_LOSS, t
            count += 1
            t += rng.uniform(0.1, 0.5)

        yield ev, t
        count += 1

        if ev == MK_HEARTBEAT:
            hb_next = t + 1.0 + rng.gauss(0, 0.05)
        elif ev == MK_TELEMETRY:
            tm_next = t + 3.0 + rng.gauss(0, 0.15)
        elif ev == MK_ID_BROADCAST:
            id_next = t + 10.0 + rng.gauss(0, 0.4)
        elif ev == MK_HANDOFF:
            ho_next = t + 120.0 + rng.gauss(0, 5.0)
            # Restart a new "track" with a beacon
            yield MK_BEACON, t + rng.uniform(0.05, 0.3)
            count += 1


def build_transition_counts(n_events: int = 100_000, seed: int = 42):
    rng = random.Random(seed)

    # counts[from_state][to_state]
    counts = [[0] * MK_N_STATES for _ in range(MK_N_STATES)]

    prev_state = None
    prev_ts    = None

    gen = gen_benign_events(n_events, rng)
    for ev, ts in gen:
        dt    = (ts - prev_ts) if prev_ts is not None else 0.0
        iat   = iat_bucket(dt)
        state = state_index(ev, iat)

        if prev_state is not None:
            counts[prev_state][state] += 1

        prev_state = state
        prev_ts    = ts

    return counts


def laplace_smooth_log_probs(counts, smoothing: float = 0.5):
    """Return log-prob matrix from raw counts with Laplace smoothing."""
    V = MK_N_STATES
    log_probs = [[0.0] * V for _ in range(V)]
    for row in range(V):
        total = sum(counts[row]) + smoothing * V
        for col in range(V):
            p = (counts[row][col] + smoothing) / total
            log_probs[row][col] = math.log(p)
    return log_probs


def nll_to_q15(nll: float) -> int:
    """Quantise NLL (nats) to int16 Q1.15 with NLL_SCALE=8.0."""
    if nll <= 0.0:
        return 0
    v = (nll / NLL_SCALE) * 32767.0
    v = min(v, 32767.0)
    return int(v + 0.5)


def write_bin(log_probs, out_path: pathlib.Path):
    n = MK_N_STATES
    with open(out_path, 'wb') as fh:
        # Header
        fh.write(MAGIC)                             # 4 bytes
        fh.write(struct.pack('<H', VERSION))        # 2 bytes
        fh.write(struct.pack('<H', n))              # 2 bytes
        # Transition matrix: row-major, int16 LE
        for row in range(n):
            for col in range(n):
                nll  = -log_probs[row][col]         # convert log-prob → NLL
                q15  = nll_to_q15(nll)
                fh.write(struct.pack('<h', q15))    # signed int16 LE

    size = out_path.stat().st_size
    expected = 8 + 2 * n * n
    assert size == expected, f"Size mismatch: {size} != {expected}"
    return size


def verify_bin(out_path: pathlib.Path):
    """Sanity-check the written file by re-reading and checking a few entries."""
    with open(out_path, 'rb') as fh:
        magic = fh.read(4)
        assert magic == MAGIC, f"Bad magic: {magic!r}"
        version, n_states = struct.unpack('<HH', fh.read(4))
        assert version == VERSION
        assert n_states == MK_N_STATES

        # Read first row (BEACON|iat0 → anything)
        row0 = struct.unpack(f'<{n_states}h', fh.read(n_states * 2))
        # All values should be in [0, 32767]
        assert all(0 <= v <= 32767 for v in row0), "Q15 range violation"
        # The most common follow-on from BEACON|fast should have a lower NLL
        # than an OOV sentinel (32767). Just assert not all are max.
        assert not all(v == 32767 for v in row0), "All OOV — something wrong"


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        '--out',
        default=str(pathlib.Path(__file__).parent.parent.parent /
                    'data' / 'models' / 'markov_baseline.bin'),
        help='Output path for markov_baseline.bin'
    )
    parser.add_argument('--n-events', type=int, default=100_000)
    parser.add_argument('--seed',     type=int, default=42)
    parser.add_argument('--smoothing', type=float, default=0.1)
    args = parser.parse_args()

    out_path = pathlib.Path(args.out)
    out_path.parent.mkdir(parents=True, exist_ok=True)

    print(f"Simulating {args.n_events:,} benign RID events (seed={args.seed})…")
    counts = build_transition_counts(args.n_events, args.seed)

    observed = sum(counts[r][c] for r in range(MK_N_STATES)
                                for c in range(MK_N_STATES))
    print(f"  Transitions observed : {observed:,}")
    print(f"  States with outgoing : "
          f"{sum(1 for r in range(MK_N_STATES) if any(counts[r]))}")

    print("Computing Laplace-smoothed log-probabilities…")
    log_probs = laplace_smooth_log_probs(counts, args.smoothing)

    print(f"Writing binary to: {out_path}")
    size = write_bin(log_probs, out_path)

    verify_bin(out_path)

    expected = 8 + 2 * MK_N_STATES * MK_N_STATES
    print(f"  File size  : {size} bytes  (expected {expected})")
    print(f"  NLL_SCALE  : {NLL_SCALE}")
    print(f"  Smoothing  : {args.smoothing}")

    # Spot-check: HEARTBEAT|iat2 → HEARTBEAT|iat2 should be a low-NLL path
    hb_s2 = state_index(MK_HEARTBEAT, 2)  # HB arriving 500ms-2s after prev HB
    self_nll = -log_probs[hb_s2][hb_s2]
    print(f"  HEARTBEAT|iat2 -> HEARTBEAT|iat2 NLL : {self_nll:.4f} nats")

    # Spot-check: LOSS|iat4 -> GAP|iat4 should have a high NLL (rare pair)
    loss_s4 = state_index(MK_LOSS, 4)
    gap_s4  = state_index(MK_GAP, 4)
    rare_nll = -log_probs[loss_s4][gap_s4]
    print(f"  LOSS|iat4  -> GAP|iat4  NLL          : {rare_nll:.4f} nats")

    assert self_nll < rare_nll, (
        f"Expected common path NLL ({self_nll:.2f}) < rare path NLL "
        f"({rare_nll:.2f})"
    )
    print("BASELINE_BUILD_OK")
    return 0


if __name__ == '__main__':
    sys.exit(main())
