#!/usr/bin/env python3
"""
expand-fseq.py — Gothic Folly rose-window FSEQ pixel-expansion (issue #70)

The rose window is DESIGNED in xLights at the cell level: 14 tracery cells x 16
petals = 224 logical pixels (Rose Window A = petals 1-8 in u1, B = petals 9-16 in
u2; 42 channels/petal). But each petal's physical strip has ~517 individually
addressed WS2815 LEDs. This script expands the cell-level colors to the physical
LED counts and APPENDS the expanded channels at a high universe block (u76+),
leaving every original channel untouched (so xLights' controller config for all
other zones stays valid). Run it once per finalized sequence, then play the
expanded .fseq on FPP.

  Pipeline:  xLights .fseq  ->  expand-fseq.py  ->  expanded .fseq  ->  FPP -> F48V5

Cell -> physical mapping comes from pixel-map/petal-config-v2.json (per-petal,
calibrated). Passage LEDs (skip_before) stay dark. Color values are copied RGB as-is;
the physical strips' color order (RGB — Falcon-confirmed 2026-07-15) is applied by the F48V5 port config,
not here.

Usage:
    python3 pixel-map/expand-fseq.py SOURCE.fseq [OUT.fseq]        # uncompressed output
    python3 pixel-map/expand-fseq.py --zstd SOURCE.fseq [OUT.fseq] # zstd-compressed output
    python3 pixel-map/expand-fseq.py --map        # print petal -> expanded universe map
    python3 pixel-map/expand-fseq.py --test       # self-test the expansion logic

Requires the `zstd` CLI on PATH for reading (and, with --zstd, writing) xLights'
zstd-compressed .fseq (brew install zstd). Uncompressed source .fseq needs no
external tool. Default output is uncompressed dense (~10 GB/hr at 40 fps); the
expanded frames are mostly zeros, so `--zstd` shrinks that dramatically and FPP
reads compressed FSEQ v2 natively.
"""
import sys, os, json, struct, subprocess, time

HERE = os.path.dirname(os.path.abspath(__file__))
PETAL_CONFIG = os.path.join(HERE, 'petal-config-v2.json')

# ---- Source (cell-level) layout — authoritative, matches generate-rose-window-submodels.py + the sim ----
CELL_ORDER = ['1b', '2a', '2c', '3b', '4a', '4c', '5a', '5b', '5c', '6a', '6c', '7a', '7b', '7c']
CELL_INDEX = {c: i for i, c in enumerate(CELL_ORDER)}
CELLS_PER_PETAL = 14
CH_PER_PIXEL = 3                     # RGB (source) / WS2815 (physical)
PETAL_CH = CELLS_PER_PETAL * CH_PER_PIXEL   # 42

def source_petal_base(P):
    """0-based absolute channel of petal P's 42-channel cell block in the source."""
    if 1 <= P <= 8:
        return (P - 1) * PETAL_CH               # Rose Window A, universe 1 (abs 0)
    return 510 + (P - 9) * PETAL_CH             # Rose Window B, universe 2 (abs 510)

# ---- Expanded (physical) layout — this script DEFINES it; the F48V5 runbook (#21) follows it ----
PHYS_PER_PETAL   = 517                 # physical LEDs per petal (calibrated max; strips left uncut)
EXP_BASE_UNIV    = 76                  # first expanded universe (u76+; towers end u75, arches u48)
UNIV_PER_PETAL   = 4                   # 517px*3 = 1551 ch = 3.04 universes -> 4, universe-aligned per petal
CH_PER_UNIVERSE  = 510

def exp_petal_universe(P):
    return EXP_BASE_UNIV + (P - 1) * UNIV_PER_PETAL
def exp_petal_base(P):
    """0-based absolute channel where petal P's physical pixels start."""
    return (exp_petal_universe(P) - 1) * CH_PER_UNIVERSE
EXP_MAX_CHANNEL = exp_petal_base(16) + PHYS_PER_PETAL * CH_PER_PIXEL   # end of last petal

# ---- petal segment map (which physical LEDs belong to which cell, per petal) ----
def load_petal_segments():
    cfg = json.load(open(PETAL_CONFIG))
    default = cfg['default_segments']
    overrides = cfg.get('petal_overrides', {})
    segs = {}
    for P in range(1, 17):
        key = f'P{P:02d}'
        ov = overrides.get(key)
        if isinstance(ov, dict):
            segs[P] = ov.get('segments') or ov.get('default_segments') or default
        else:
            segs[P] = default          # P02/P15/P16 (uncalibrated) use the P01 baseline
    return segs

# ============================== FSEQ v2 reader ==============================

def _zstd_decompress(blk):
    r = subprocess.run(['zstd', '-d', '--stdout', '-'], input=blk, capture_output=True)
    if r.returncode != 0:
        raise RuntimeError('zstd decompress failed: ' + r.stderr.decode()[:200])
    return r.stdout

ZSTD_LEVEL = 19          # offline one-time step; mostly-zero frames compress fast even at max

def _zstd_compress(data, level=ZSTD_LEVEL):
    r = subprocess.run(['zstd', f'-{level}', '-c', '-'], input=data, capture_output=True)
    if r.returncode != 0:
        raise RuntimeError('zstd compress failed: ' + r.stderr.decode()[:200])
    return r.stdout

def read_fseq(path):
    d = open(path, 'rb').read()
    if d[0:4] != b'PSEQ':
        raise ValueError('not a v2 (.fseq/PSEQ) file: ' + repr(d[0:4]))
    chan_data_offset = struct.unpack('<H', d[4:6])[0]
    chan_count = struct.unpack('<I', d[10:14])[0]     # stored (sparse-packed) channels/frame
    num_frames = struct.unpack('<I', d[14:18])[0]
    step_time = d[18]
    comp_type = d[20] & 0x0f                            # 0=none 1=zstd 2=zlib
    num_blocks = ((d[20] & 0xf0) << 4) | d[21]
    num_sparse = d[22]
    off = 32
    blocks = []
    for _ in range(num_blocks):
        fn, ln = struct.unpack('<II', d[off:off+8]); off += 8
        blocks.append((fn, ln))
    sparse = []
    for _ in range(num_sparse):
        start = d[off] | (d[off+1] << 8) | (d[off+2] << 16)
        length = d[off+3] | (d[off+4] << 8) | (d[off+5] << 16)
        off += 6; sparse.append((start, length))
    if not sparse:
        sparse = [(0, chan_count)]                     # dense file
    # decompress the block stream into one packed buffer
    payload = d[chan_data_offset:]
    if comp_type == 0:
        raw = payload
    elif comp_type in (1, 2):
        raw = bytearray(); pos = 0
        for fn, ln in blocks:
            if ln == 0: continue
            blk = payload[pos:pos+ln]; pos += ln
            raw += _zstd_decompress(blk) if comp_type == 1 else __import__('zlib').decompress(blk)
    else:
        raise ValueError('unknown compression type %d' % comp_type)
    max_chan = max(s + l for s, l in sparse)
    frames = []
    for fi in range(num_frames):
        packed = raw[fi*chan_count:(fi+1)*chan_count]
        dense = bytearray(max_chan)
        o = 0
        for start, length in sparse:
            dense[start:start+length] = packed[o:o+length]; o += length
        frames.append(dense)
    return {'num_frames': num_frames, 'step_time': step_time, 'max_chan': max_chan}, frames

# ============================== expansion ==============================

def expand_frames(frames, petal_segs):
    out_max = max(EXP_MAX_CHANNEL, frames[0].__len__() if frames else 0)
    out_frames = []
    for dense in frames:
        out = bytearray(out_max)
        out[0:len(dense)] = dense                      # pass ALL original channels through untouched
        for P in range(1, 17):
            sbase = source_petal_base(P)
            ebase = exp_petal_base(P)
            idx = 0
            for seg in petal_segs[P]:
                idx += seg.get('skip_before', 0)       # passage LEDs stay black
                ci = CELL_INDEX[seg['cell']]
                src = sbase + ci * CH_PER_PIXEL
                r, g, b = dense[src], dense[src+1], dense[src+2]
                for k in range(seg['count']):
                    o = ebase + (idx + k) * CH_PER_PIXEL
                    out[o] = r; out[o+1] = g; out[o+2] = b
                idx += seg['count']
        out_frames.append(out)
    return out_frames, out_max

# ============================== FSEQ v2 writer (uncompressed, dense) ==============================

def write_fseq(path, frames, chan_count, step_time):
    h = bytearray(32)
    h[0:4] = b'PSEQ'
    struct.pack_into('<H', h, 4, 32)                   # channel data offset
    h[6] = 0; h[7] = 2                                 # v2.0
    struct.pack_into('<H', h, 8, 32)                   # variable-header offset (= data start; none)
    struct.pack_into('<I', h, 10, chan_count)
    struct.pack_into('<I', h, 14, len(frames))
    h[18] = step_time
    h[19] = 0                                          # flags
    h[20] = 0                                          # compType=0, numBlocks hi=0
    h[21] = 0                                          # numBlocks lo=0
    h[22] = 0                                          # numSparse=0 (dense)
    h[23] = 0
    struct.pack_into('<Q', h, 24, int(time.time() * 1_000_000) & 0xFFFFFFFFFFFFFFFF)
    with open(path, 'wb') as f:
        f.write(h)
        for fr in frames:
            f.write(fr if len(fr) == chan_count else (bytes(fr) + b'\x00' * (chan_count - len(fr))))

# ============================== FSEQ v2 writer (zstd-compressed, dense) ==============================

def write_fseq_zstd(path, frames, chan_count, step_time, level=ZSTD_LEVEL):
    """Write a zstd-compressed FSEQ v2. Frames are grouped into blocks (~2 MB
    uncompressed each, block count capped at the format's 4095 limit); each block
    is zstd-compressed. Dense (no sparse ranges) — the block compression is what
    crushes the mostly-zero expanded frames. Round-trips through read_fseq()."""
    import math
    num_frames = len(frames)
    def norm(fr):
        return bytes(fr) if len(fr) == chan_count else (bytes(fr) + b'\x00' * (chan_count - len(fr)))
    fpb = max(1, (2 * 1024 * 1024) // max(1, chan_count))          # ~2 MB uncompressed per block
    if num_frames and math.ceil(num_frames / fpb) > 4095:
        fpb = math.ceil(num_frames / 4095)                         # stay within 12-bit block count
    comp_blocks = []
    for first in range(0, num_frames, fpb):
        chunk = b''.join(norm(frames[i]) for i in range(first, min(first + fpb, num_frames)))
        comp_blocks.append((first, _zstd_compress(chunk, level)))
    num_blocks = len(comp_blocks)
    header_len = 32 + num_blocks * 8                               # numSparse=0, no variable headers
    h = bytearray(header_len)
    h[0:4] = b'PSEQ'
    struct.pack_into('<H', h, 4, header_len)                       # channel data offset
    h[6] = 0; h[7] = 2                                             # v2.0
    struct.pack_into('<H', h, 8, header_len)                       # variable-header offset (= data start)
    struct.pack_into('<I', h, 10, chan_count)
    struct.pack_into('<I', h, 14, num_frames)
    h[18] = step_time
    h[19] = 0                                                      # flags
    h[20] = (1 & 0x0f) | (((num_blocks >> 8) & 0x0f) << 4)         # compType=1 (zstd) + block-count hi nibble
    h[21] = num_blocks & 0xFF                                      # block-count lo byte
    h[22] = 0                                                      # numSparse=0 (dense)
    h[23] = 0
    struct.pack_into('<Q', h, 24, int(time.time() * 1_000_000) & 0xFFFFFFFFFFFFFFFF)
    off = 32
    for first, cblk in comp_blocks:                               # block index: (firstFrame, compLen)
        struct.pack_into('<II', h, off, first, len(cblk)); off += 8
    with open(path, 'wb') as f:
        f.write(h)
        for _, cblk in comp_blocks:
            f.write(cblk)

# ============================== map / cli ==============================

def print_map():
    print('Expanded rose-window universe map (this script defines it; F48V5 ports 33-48 follow it):')
    print(f'{"Petal":6} {"Port":5} {"StartUniv":9} {"StartCh":8} {"Count":6} {"Type":7} {"Color":6} {"Bright"}')
    for P in range(1, 17):
        port = 32 + P
        print(f'{P:<6} {port:<5} u{exp_petal_universe(P):<8} {1:<8} {PHYS_PER_PETAL:<6} {"WS2815":7} {"RGB":6} 50%')
    print(f'\nExpanded block: u{EXP_BASE_UNIV}-u{exp_petal_universe(16)+UNIV_PER_PETAL-1}  '
          f'({16*UNIV_PER_PETAL} universes, {UNIV_PER_PETAL}/petal, universe-aligned). '
          f'Max channel {EXP_MAX_CHANNEL} = u{(EXP_MAX_CHANNEL-1)//510+1}. F48V5 ceiling 192.')

def self_test():
    segs = load_petal_segments()
    # synthetic 1-frame source: give every cell of every petal a distinct-ish color
    frame = bytearray(1020)
    for P in range(1, 17):
        base = source_petal_base(P)
        for cell, ci in CELL_INDEX.items():
            o = base + ci*3
            frame[o] = (P*10) & 0xFF; frame[o+1] = (ci*15) & 0xFF; frame[o+2] = 77
    out, omax = expand_frames([frame], segs)
    of = out[0]
    ok = True
    # verify: for petal 3 (has override), cell 3b -> its physical LEDs carry petal 3 / cell-3b color
    for P in (1, 3, 16):
        base = source_petal_base(P); ebase = exp_petal_base(P)
        idx = 0
        for seg in segs[P]:
            idx += seg.get('skip_before', 0)
            ci = CELL_INDEX[seg['cell']]
            exp_rgb = (frame[base+ci*3], frame[base+ci*3+1], frame[base+ci*3+2])
            # check first + last LED of this segment
            for k in (0, seg['count']-1):
                o = ebase + (idx+k)*3
                got = (of[o], of[o+1], of[o+2])
                if got != exp_rgb:
                    ok = False; print(f'  MISMATCH P{P} cell {seg["cell"]} led {idx+k}: got {got} exp {exp_rgb}')
            idx += seg['count']
        # passage LEDs (a skip_before>0 segment) should be black
    total_lit = sum(1 for i in range(exp_petal_base(1), omax, 3) if of[i] or of[i+1] or of[i+2])
    print(f'self-test: expansion {"OK" if ok else "FAILED"}; out max channel {omax}; lit physical pixels {total_lit}')
    print(f'  expected ~ sum(cell_leds per petal) ; P01 cell_leds=509, 16 petals ~ {509*16} (minus uncalibrated variance)')
    return ok

def main():
    args = sys.argv[1:]
    if not args or args[0] in ('-h', '--help'):
        print(__doc__); return
    if args[0] == '--map':
        print_map(); return
    if args[0] == '--test':
        sys.exit(0 if self_test() else 1)
    use_zstd = False
    if '--zstd' in args:
        use_zstd = True; args = [a for a in args if a != '--zstd']
    src = args[0]
    out = args[1] if len(args) > 1 else os.path.splitext(src)[0] + '-expanded.fseq'
    print(f'Reading {src} ...')
    meta, frames = read_fseq(src)
    print(f'  {meta["num_frames"]} frames, {meta["step_time"]}ms/frame ({1000//meta["step_time"]}fps), '
          f'source max channel {meta["max_chan"]}')
    segs = load_petal_segments()
    print('Expanding rose window (16 petals -> physical LEDs, appended at u76+) ...')
    out_frames, out_max = expand_frames(frames, segs)
    print(f'  expanded max channel {out_max} = u{(out_max-1)//510+1}')
    if use_zstd:
        print(f'Writing zstd-compressed (level {ZSTD_LEVEL}) ...')
        write_fseq_zstd(out, out_frames, out_max, meta['step_time'])
        sz = os.path.getsize(out)
        raw = out_max * len(out_frames)
        print(f'Wrote {out}  ({sz/1e6:.1f} MB zstd, vs {raw/1e6:.1f} MB dense — {raw/max(sz,1):.0f}x smaller)')
    else:
        write_fseq(out, out_frames, out_max, meta['step_time'])
        sz = os.path.getsize(out)
        print(f'Wrote {out}  ({sz/1e6:.1f} MB, uncompressed — use --zstd to shrink)')
    print('  -> copy to the FPP Pi; configure F48V5 rose ports 33-48 per `expand-fseq.py --map`.')

if __name__ == '__main__':
    main()