"""Populate data/raw_port_counts.csv from the experimental dataset.

Reads the manifest + decoded per-port counts in `../dataset/` and fills the
n0..n7 columns of `data/raw_port_counts.csv` with the measured single-photon
port distributions, resolving each manuscript row to its job through the
manifest.  Re-run this whenever the dataset is updated.

Exp 5b (selectivity) is filled too: its reported quantity is the within-syndrome
ratio ports[tp]/(ports 0-3), which `data_io.py` computes via the "selectivity"
metric (the full single-photon total is preserved for event counting).  The
underlying jobs (RS_845_test4/5/6) are also checked by verify_package.py.

Usage:
    python build_raw_counts.py [path-to-dataset]
"""

from __future__ import annotations

import csv
import os
import sys

HERE = os.path.dirname(__file__)
PKG = sys.argv[1] if len(sys.argv) > 1 else os.path.join(HERE, "..", "dataset")
PER_PORT = os.path.join(PKG, "data", "processed", "per_port_counts.csv")
MANIFEST = os.path.join(PKG, "data", "manifest", "job_manifest.csv")
RAW_CSV = os.path.join(HERE, "data", "raw_port_counts.csv")

# (pipeline experiment, pipeline label) -> (manifest experiment, manifest configuration)
MAP = {
    ("exp1", "input mode 0"): ("1", "mode 0"),
    ("exp1", "input mode 1"): ("1", "mode 1"),
    ("exp1", "input mode 2"): ("1", "mode 2"),
    ("exp1", "input mode 3"): ("1", "mode 3"),
    ("exp1", "input mode 4"): ("1", "mode 4"),
    ("exp1", "input mode 5"): ("1", "mode 5"),
    ("exp1", "input mode 6"): ("1", "mode 6"),
    ("exp1", "input mode 7"): ("1", "mode 7"),
    ("exp2", "Control (non-neutral)"): ("2", "Control non-neutral"),
    ("exp2", "|0> - |1>"): ("2", "|0>-|1> neutral"),
    ("exp2", "|0> - |4>"): ("2", "|0>-|4> neutral"),
    ("exp2", "Balanced 4+4-"): ("2", "Balanced 4+4-"),
    ("exp2", "Uniform (pure DC)"): ("2", "Uniform pure DC"),
    ("exp3", "c=0.00"): ("3", "c=0.00"),
    ("exp3", "c=0.05"): ("3", "c=0.05"),
    ("exp3", "c=0.10"): ("3", "c=0.10"),
    ("exp3", "c=0.20"): ("3", "c=0.20"),
    ("exp3", "c=0.40"): ("3", "c=0.40"),
    ("exp3", "c=0.80"): ("3", "c=0.80"),
    ("exp3", "c=1.60"): ("3", "c=1.60"),
    ("exp4", "Neutral, 0x core"): ("4", "Neutral no gate cycle"),
    ("exp4", "Neutral, 1x core"): ("4", "Neutral 1x gate cycle"),
    ("exp4", "Neutral, 2x core"): ("4", "Neutral 2x gate cycle"),
    ("exp4", "Neutral, 3x core"): ("4", "Neutral 3x gate cycle"),
    ("exp4", "Control, 0x core"): ("4", "Control non-neutral no gate cycle"),
    ("exp4", "Control, 1x core"): ("4", "Control non-neutral 1x gate cycle"),
    ("exp5a", "b0+b1"): ("5", "Code b0 xor b1"),
    ("exp5a", "b0+b2"): ("5", "Code b0 xor b2"),
    ("exp5a", "b1+b2"): ("5", "Code b1 xor b2"),
    ("exp5a", "b0+b1+b2"): ("5", "Code b0 xor b1 xor b2"),
    ("exp5b", "|0>-|1> (Sx)"): ("5", "|0>-|1> Sx violation"),
    ("exp5b", "|0>-|2> (Sy)"): ("5", "|0>-|2> Sy violation"),
    ("exp5b", "|0>-|4> (Sz)"): ("5", "|0>-|4> Sz violation"),
    ("exp5c", "Non-neutral control (syndrome)"): ("5", "Mode |0> control"),
    ("exp6", "Calibrated (Apr 4)"): ("6", "Calibrated Apr 4"),
    ("exp6", "Degraded (Apr 5-6)"): ("6", "Degraded Apr 5-6"),
    ("exp6", "Restored (Apr 6)"): ("6", "Restored Apr 6"),
}


def load_per_port() -> dict[str, list[int]]:
    jobs: dict[str, list[int]] = {}
    with open(PER_PORT, newline="") as fh:
        for r in csv.DictReader(fh):
            jobs.setdefault(r["job_name"], [0] * 8)[int(r["port"])] = int(r["count"])
    return jobs


def load_manifest(jobs: dict[str, list[int]]) -> dict[tuple[str, str], str]:
    out: dict[tuple[str, str], str] = {}
    with open(MANIFEST, newline="") as fh:
        for r in csv.DictReader(fh):
            for cand in (r["raw_result_file"].removesuffix(".result.json"),
                         r["expected_job_name"]):
                if cand in jobs:
                    out[(r["experiment"], r["configuration"])] = cand
                    break
    return out


def main() -> None:
    jobs = load_per_port()
    man = load_manifest(jobs)

    with open(RAW_CSV, newline="") as fh:
        reader = csv.DictReader(fh)
        fields = reader.fieldnames
        rows = list(reader)

    filled = skipped = 0
    for row in rows:
        key = (row["experiment"], row["label"])
        if key not in MAP:
            skipped += 1
            continue
        man_key = MAP[key]
        if man_key not in man:
            raise SystemExit(f"No manifest job for {key} -> {man_key}")
        ports = jobs[man[man_key]]
        for i in range(8):
            row[f"n{i}"] = ports[i]
        filled += 1

    with open(RAW_CSV, "w", newline="") as fh:
        w = csv.DictWriter(fh, fieldnames=fields)
        w.writeheader()
        w.writerows(rows)

    print(f"Filled {filled} rows from the dataset"
          + (f"; left {skipped} rows on the summary path." if skipped else "."))
    print(f"Wrote {RAW_CSV}")


if __name__ == "__main__":
    main()
