#!/bin/bash
# Pull raw records, logs and chain messages from the cluster nodes and merge the E1_core Llama shard into runs/raw/E1_core.jsonl
# (idempotent: shard records are appended only if their attempt_id is not already present).
set -e
cd "$(dirname "$0")/.."
mkdir -p runs/raw/remote runs/logs/remote runs/raw/chains
rsync -aq M3U96a:~/wp30-worker/wp30/runs/raw/E1_core.jsonl runs/raw/remote/E1_core_llama_m3u96a.jsonl
rsync -aq M3U96a:~/wp30-worker/wp30/runs/logs/ runs/logs/remote/m3u96a/
rsync -aq M3U96b:~/wp30-worker/wp30/runs/raw/ED_hosted_core.jsonl runs/raw/ 2>/dev/null || true
rsync -aq M3U96b:~/wp30-worker/wp30/runs/raw/E4_core.jsonl runs/raw/remote/E4_core_shard0_m3u96b.jsonl 2>/dev/null || true
rsync -aq M3U96b:~/wp30-worker/wp30/runs/raw/chains/ runs/raw/chains/ 2>/dev/null || true
rsync -aq M3U96b:~/wp30-worker/wp30/runs/logs/ runs/logs/remote/m3u96b/
python3 - <<'PY'
import json
seen = set()
for l in open("runs/raw/E1_core.jsonl"):
    seen.add(json.loads(l)["attempt_id"])
n = 0
with open("runs/raw/E1_core.jsonl", "a") as out:
    for l in open("runs/raw/remote/E1_core_llama_m3u96a.jsonl"):
        r = json.loads(l)
        if r["document_id"].startswith("lb_"):
            continue   # public families of the Llama pair were run on the laptop (D-019); the 25 remote public records are dropped
        if r["attempt_id"] not in seen:
            r["hardware_id"] = "m3u96a-m3ultra-96g"; out.write(json.dumps(r, ensure_ascii=False) + "\n"); seen.add(r["attempt_id"]); n += 1
print("merged", n, "new Llama-shard records; total", len(seen))
# merge the E4_core chain shard from M3U96b (even documents) into the local file (odd documents)
try:
    seen4 = {json.loads(l)["attempt_id"] for l in open("runs/raw/E4_core.jsonl")} if __import__("os").path.exists("runs/raw/E4_core.jsonl") else set()
    n4 = 0
    with open("runs/raw/E4_core.jsonl", "a") as out:
        for l in open("runs/raw/remote/E4_core_shard0_m3u96b.jsonl"):
            r = json.loads(l)
            if r["attempt_id"] not in seen4:
                r["hardware_id"] = "m3u96b-m3ultra-96g"; out.write(json.dumps(r, ensure_ascii=False) + "\n"); seen4.add(r["attempt_id"]); n4 += 1
    print("merged", n4, "E4 shard0 records; total", len(seen4))
except FileNotFoundError:
    pass
# (E5/E3 run on the laptop since D-019; the E4 shard0 records are labelled at merge time above)
PY