#!/usr/bin/env python3
"""
correlate_doe.py <raw_log> <out_csv>

Reconstructs which DoE trial (run_id, id_idx, id_label, rep) was active
for every heartbeat-tick CSV row in a std79-doe raw log, despite the
async HB-ON tick printer splicing CSV rows mid-token into the trial
loop's own console output on the shared serial line -- including mid a
DOE-RUN,... marker itself.

Method:
  1. Strip ANSI codes and the per-line "[VM-TAG] " console prefix (an
     artifact of the log capture layer), concatenate every stripped
     line with no separator. Find every [HADES][DOE ] header/row
     occurrence and physically remove exactly those spans, recording
     each removed data row's insertion offset in the CLEANED stream's
     own coordinates.
  2. Parse only "DOE-RUN,<run_id> ," (the run_id digits) from the clean
     stream -- reliable even when a CSV row spliced into the id_idx/
     id_label/rep text right after it (observed live: a row's last
     digit field can land with zero separator against the immediately
     following trial digit, e.g. "...11642" + "6" -> "116426", making
     id_idx unrecoverable from THIS log alone for the affected trial).
     id_idx/id_label/rep are looked up from RUN_ID_MAP below instead of
     re-parsed -- deterministic given the campaign's fixed seed (12345),
     independently verified identical across 5+ prior clean runs
     (results-20260911-donor-floor-fix/amd64-doe-raw.log and others).
  3. Assign each CSV row to whichever run_id marker's start offset is
     the largest one at or before that row's insertion offset.
"""
import re
import sys
import csv
import bisect

# From results-20260911-donor-floor-fix/amd64-doe-raw.log (seed 12345),
# independently re-verified identical (md5) across every std79-doe run
# since e2abc56. run_id -> (id_idx, id_label, rep).
RUN_ID_MAP = {
    0: (6, "04", 2), 1: (3, "01", 1), 2: (0, "zuse", 1), 3: (1, "rajames", 0),
    4: (5, "03", 0), 5: (0, "zuse", 2), 6: (2, "00", 2), 7: (8, "06", 2),
    8: (0, "zuse", 0), 9: (6, "04", 1), 10: (6, "04", 0), 11: (4, "02", 2),
    12: (1, "rajames", 1), 13: (7, "05", 1), 14: (7, "05", 2), 15: (4, "02", 0),
    16: (2, "00", 0), 17: (8, "06", 0), 18: (7, "05", 0), 19: (2, "00", 1),
    20: (8, "06", 1), 21: (5, "03", 2), 22: (1, "rajames", 2), 23: (3, "01", 2),
    24: (5, "03", 1), 25: (3, "01", 0), 26: (4, "02", 1),
}

def build_stream(raw_text):
    parts = []
    for line in raw_text.split('\n'):
        m = re.match(r'^\[[^\]]+\]\s?(.*)$', line)
        parts.append(m.group(1) if m else line)
    return ''.join(parts)

ROW_RE = re.compile(
    r'\[HADES\]\[DOE \] '
    r'(?:(?P<header>tick_number,elapsed_ns,[^\[]*?)(?=\[HADES\]\[DOE \]|DOE-RUN|$)'
    r'|(?P<row>[0-9]+(?:,[0-9]+)+))'
)
RUN_ID_RE = re.compile(r'DOE-RUN,\s*([0-9]+)\s*,')

def main():
    if len(sys.argv) != 3:
        print("usage: correlate_doe.py <raw_log> <out_csv>", file=sys.stderr)
        sys.exit(1)
    log_path, out_path = sys.argv[1], sys.argv[2]

    with open(log_path, "r", errors="replace") as f:
        raw = f.read()
    raw = re.sub(r'\x1b\[[0-9;]*[a-zA-Z]', '', raw)
    stream = build_stream(raw)

    cleaned_parts = []
    cleaned_offset = 0
    last_end = 0
    csv_rows = []  # (cleaned_insertion_offset, fields)
    header_fields = None
    for m in ROW_RE.finditer(stream):
        before = stream[last_end:m.start()]
        cleaned_parts.append(before)
        cleaned_offset += len(before)
        if m.group('header') is not None:
            header_fields = [x.strip() for x in m.group('header').split(',')]
        else:
            csv_rows.append((cleaned_offset, m.group('row').split(',')))
        last_end = m.end()
    cleaned_parts.append(stream[last_end:])
    cleaned_stream = ''.join(cleaned_parts)

    if header_fields is None:
        print("FATAL: never found the DOE CSV header row", file=sys.stderr)
        sys.exit(1)

    markers = [(m.start(), int(m.group(1))) for m in RUN_ID_RE.finditer(cleaned_stream)]
    # Sanity: run_ids should appear in increasing offset order with values
    # 0..26 each exactly once. Warn, don't crash, on anything unexpected --
    # a genuinely dropped/garbled run_id digit is possible in principle.
    seen = [rid for _, rid in markers]
    if sorted(set(seen)) != list(range(27)) or len(seen) != 27:
        print(f"WARNING: expected exactly run_ids 0..26 once each, got {sorted(seen)}", file=sys.stderr)

    marker_offsets = [off for off, _ in markers]

    out_rows = []
    unlabeled = 0
    for offset, fields in csv_rows:
        i = bisect.bisect_right(marker_offsets, offset) - 1
        if i >= 0:
            run_id = markers[i][1]
            id_idx, id_label, rep = RUN_ID_MAP.get(run_id, (None, None, None))
        elif 0 not in seen:
            # No marker precedes this row because run_id 0's own marker
            # (always the very first trial, loop index 0) got garbled
            # beyond recovery in this particular log -- safe to assume
            # every row before the first RECOVERED marker still belongs
            # to trial 0, not truly unlabeled.
            run_id = 0
            id_idx, id_label, rep = RUN_ID_MAP[0]
        else:
            run_id, id_idx, id_label, rep = None, None, None, None
            unlabeled += 1
        out_rows.append(((run_id, id_idx, id_label, rep), fields))

    with open(out_path, "w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["run_id", "id_idx", "id_label", "rep"] + header_fields)
        for (run_id, id_idx, id_label, rep), fields in out_rows:
            w.writerow([run_id, id_idx, id_label, rep] + fields)

    print(f"trial markers found: {len(markers)}")
    print(f"data rows: {len(out_rows)}, unlabeled (before first marker): {unlabeled}")
    labeled_run_ids = sorted(set(t[0] for (t, _) in out_rows if t[0] is not None))
    print(f"distinct run_ids covered: {len(labeled_run_ids)} -> {labeled_run_ids}")

if __name__ == "__main__":
    main()