"""Test: does the second '20 NN' block in SS0 continue Tran samples?"""
import sys
sys.path.insert(0, ".")
from analysis.load_bundle import _parse_txt
from minimateplus.waveform_codec import walk_body, find_data_start


def s4(n):
    return n if n < 8 else n - 16


def i8(b):
    return b if b < 128 else b - 256


def main():
    stem = "M529LL1A.SS0"
    path = f"tests/fixtures/5-11-26/{stem}"
    with open(path, "rb") as f:
        body = f.read()[43:-26]
    _, samples = _parse_txt(path + ".TXT")
    truth_T_16 = [round(v * 200) for v in samples["Tran"]]

    # Preamble
    T0 = int.from_bytes(body[3:5], "big", signed=True)
    T1 = int.from_bytes(body[5:7], "big", signed=True)

    # Walk blocks
    start = find_data_start(body)
    blocks = walk_body(body, start)

    print(f"=== {stem} ===  T[0]={T0} T[1]={T1}")

    # Hypothesis: Tran continues through ALL 10 NN and 20 NN blocks
    # in order, until the next 40 02 segment header (which resets).
    T = [T0, T1]
    cur = T1
    decoded_count = 2  # T[0], T[1] from preamble
    for bi, blk in enumerate(blocks):
        if blk.tag_hi == 0x10:
            for byte in blk.data:
                for nib in ((byte >> 4) & 0xF, byte & 0xF):
                    cur += s4(nib)
                    T.append(cur)
                    decoded_count += 1
        elif blk.tag_hi == 0x20:
            for byte in blk.data:
                cur += i8(byte)
                T.append(cur)
                decoded_count += 1
        elif blk.tag_hi == 0x40:
            # Segment header — stop here for this test
            break
        # 00 and 30 NN don't contribute to Tran (in this hypothesis)

    # Compare to truth
    print(f"  Decoded {len(T)} T samples up to first 40 02")
    matches = sum(1 for i in range(min(len(T), len(truth_T_16))) if T[i] == truth_T_16[i])
    print(f"  Matches in first {min(len(T), len(truth_T_16))}: {matches}")
    # Print first divergence
    for i in range(min(len(T), len(truth_T_16))):
        if T[i] != truth_T_16[i]:
            print(f"  First divergence: sample {i}: pred={T[i]}, truth={truth_T_16[i]}")
            # Show context
            print(f"    pred  [{i-3}:{i+5}]: {T[max(0,i-3):i+5]}")
            print(f"    truth [{i-3}:{i+5}]: {truth_T_16[max(0,i-3):i+5]}")
            break


if __name__ == "__main__":
    main()