"""CSV-backed data layer for the hardware counts.

This module is a drop-in replacement for importing `hardware_data` directly:

    import data_io as hw      # instead of: import hardware_data as hw

It resolves the per-configuration counts from three sources, in priority order:

  1. `data/raw_port_counts.csv` -- the measured per-output-port single-photon
     count distributions (built from the dataset by `build_raw_counts.py`).  For
     any row whose eight port columns are filled in, the relevant event count
     `x` and the postselected total `n` are computed directly from the ports
     (n = sum of ports; x depends on the experiment's metric).  Primary source.
  2. `data/hardware_counts.csv` -- a summary (x, n) table used as a fallback.
  3. the embedded constants in `hardware_data.py` -- ultimate fallback.

Per-experiment metric (how `x` is read from a port distribution):
  dump          -> x = port 0                      (Exp. 1-4, 6 dump port)
  syndrome_sum  -> x = ports 0+1+2+3               (Exp. 5A, 5C syndrome leakage)
  selectivity   -> x = ports[target_port],         (Exp. 5B selectivity)
                   n = ports 0+1+2+3 (within-syndrome denominator)

For `syndrome_sum` rows that also carry a `target_port` (Exp. 5A), the code-side
fidelity is computed as ports[target_port] / n.  For `selectivity` rows the
Wilson interval uses the within-syndrome denominator (ports 0-3) while the full
postselected single-photon total is preserved in `extra["events"]` so the
event-count totals stay correct.
"""

from __future__ import annotations

import csv
import os
from dataclasses import replace

import hardware_data as _hd
from hardware_data import Row

DATA_DIR = os.path.join(os.path.dirname(__file__), "data")
SUMMARY_CSV = os.path.join(DATA_DIR, "hardware_counts.csv")
RAW_CSV = os.path.join(DATA_DIR, "raw_port_counts.csv")

# Metric used to read x from a port distribution, per experiment group.
METRIC = {
    "exp1": "dump",
    "exp2": "dump",
    "exp3": "dump",
    "exp4": "dump",
    "exp5a": "syndrome_sum",
    "exp5b": "selectivity",
    "exp5c": "syndrome_sum",
    "exp6": "dump",
}

SUMMARY_FIELDS = [
    "experiment", "label", "kind", "x", "n", "target_port", "expected",
    "expected_port", "fidelity", "depth", "hom_visibility", "hom_from_report",
    "c", "events",
]
PORT_FIELDS = [f"n{i}" for i in range(8)]
RAW_FIELDS = [
    "experiment", "label", "kind", "metric", "target_port", "expected_port",
    "expected", "depth", "hom_visibility", "hom_from_report", "c",
] + PORT_FIELDS


# --------------------------------------------------------------------------
# Embedded canonical structure (experiment key -> list of Rows)
# --------------------------------------------------------------------------
def _embedded() -> dict[str, list[Row]]:
    exp3 = [replace(r, extra={**r.extra, "c": c})
            for r, c in zip(_hd.EXP3, _hd.EXP3_C)]
    return {
        "exp2": list(_hd.EXP2),
        "exp3": exp3,
        "exp4": list(_hd.EXP4),
        "exp5a": list(_hd.EXP5A),
        "exp5b": list(_hd.EXP5B),
        "exp5c": [_hd.EXP5C_CONTROL],
        "exp6": list(_hd.EXP6),
        "exp1": list(_hd.EXP1_PER_MODE),
    }


# --------------------------------------------------------------------------
# Parsing helpers
# --------------------------------------------------------------------------
def _i(s):
    s = (s or "").strip()
    return int(float(s)) if s else None


def _f(s):
    s = (s or "").strip()
    return float(s) if s else None


def _b(s):
    return (s or "").strip().lower() in ("true", "1", "yes", "y")


def _extra_from_fields(expected_port, fidelity, depth, hom_vis, hom_rep, c):
    extra: dict = {}
    if expected_port is not None:
        extra["expected_port"] = expected_port
    if fidelity is not None:
        extra["fidelity"] = fidelity
    if depth is not None:
        extra["depth"] = depth
    if hom_vis is not None:
        extra["hom_visibility"] = hom_vis
    if hom_rep:
        extra["hom_calibration_report"] = True
        extra["hom_from_calibration_report"] = True
    if c is not None:
        extra["c"] = c
    return extra


# --------------------------------------------------------------------------
# Readers
# --------------------------------------------------------------------------
def _read_summary(path: str) -> dict[tuple[str, str], Row]:
    out: dict[tuple[str, str], Row] = {}
    with open(path, newline="") as fh:
        for r in csv.DictReader(fh):
            if not r.get("experiment") or r["experiment"].startswith("#"):
                continue
            extra = _extra_from_fields(
                _i(r.get("expected_port")), _f(r.get("fidelity")),
                _i(r.get("depth")), _f(r.get("hom_visibility")),
                _b(r.get("hom_from_report")), _f(r.get("c")))
            ev = _i(r.get("events"))
            if ev is not None:
                extra["events"] = ev
            row = Row(label=r["label"], x=_i(r["x"]), n=_i(r["n"]),
                      expected=_f(r.get("expected")), kind=r.get("kind", ""),
                      port=_i(r.get("target_port")), extra=extra)
            out[(r["experiment"], r["label"])] = row
    return out


def _read_raw_ports(path: str) -> dict[tuple[str, str], Row]:
    """Compute (x, n) from filled per-port rows; skip rows with blank ports."""
    out: dict[tuple[str, str], Row] = {}
    with open(path, newline="") as fh:
        for r in csv.DictReader(fh):
            if not r.get("experiment") or r["experiment"].startswith("#"):
                continue
            ports = [_i(r.get(f"n{i}")) for i in range(8)]
            if any(p is None for p in ports):
                continue  # row has no per-port data; fall back to summary
            n = sum(ports)
            if n == 0:
                continue
            metric = (r.get("metric") or METRIC.get(r["experiment"], "dump")).strip()
            tp = _i(r.get("target_port"))
            events = n  # total postselected single-photon events for this config
            if metric == "dump":
                x = ports[0]
            elif metric == "syndrome_sum":
                x = ports[0] + ports[1] + ports[2] + ports[3]
            elif metric == "port":
                if tp is None:
                    continue
                x = ports[tp]
            elif metric == "selectivity":
                if tp is None:
                    continue
                syn = ports[0] + ports[1] + ports[2] + ports[3]
                if syn == 0:
                    continue
                x = ports[tp]
                n = syn  # within-syndrome denominator for the Wilson interval
            else:
                continue
            extra = _extra_from_fields(
                _i(r.get("expected_port")), None, _i(r.get("depth")),
                _f(r.get("hom_visibility")), _b(r.get("hom_from_report")),
                _f(r.get("c")))
            if metric == "syndrome_sum" and tp is not None:
                extra["fidelity"] = ports[tp] / n
            if metric == "selectivity":
                extra["events"] = events
            row = Row(label=r["label"], x=x, n=n, expected=_f(r.get("expected")),
                      kind=r.get("kind", ""), port=tp, extra=extra)
            out[(r["experiment"], r["label"])] = row
    return out


# --------------------------------------------------------------------------
# Resolve final data set
# --------------------------------------------------------------------------
def load() -> dict[str, list[Row]]:
    data = _embedded()
    summary = _read_summary(SUMMARY_CSV) if os.path.exists(SUMMARY_CSV) else {}
    raw = _read_raw_ports(RAW_CSV) if os.path.exists(RAW_CSV) else {}

    overrides: dict[tuple[str, str], Row] = {}
    overrides.update(summary)
    overrides.update(raw)  # raw per-port wins over summary

    if overrides:
        # apply overrides by (experiment, label); also collect new exp1 rows
        for key in list(data):
            rows = data[key]
            new_rows = []
            for row in rows:
                new_rows.append(overrides.get((key, row.label), row))
            data[key] = new_rows
        # exp1 per-mode rows may exist only in the CSVs
        exp1_keys = sorted(k for k in overrides if k[0] == "exp1")
        if exp1_keys and not data.get("exp1"):
            data["exp1"] = [overrides[k] for k in exp1_keys]
    return data


_DATA = load()

# Public names matching hardware_data, so callers can `import data_io as hw`.
EXP1_PER_MODE = _DATA.get("exp1", [])
EXP1_SUMMARY = _hd.EXP1_SUMMARY
EXP2 = _DATA["exp2"]
EXP3 = _DATA["exp3"]
EXP3_C = [r.extra["c"] for r in EXP3]
EXP4 = _DATA["exp4"]
EXP5A = _DATA["exp5a"]
EXP5B = _DATA["exp5b"]
EXP5C_CONTROL = _DATA["exp5c"][0]
EXP6 = _DATA["exp6"]
DEVICE_METADATA = _hd.DEVICE_METADATA


def total_events() -> int:
    n = EXP1_SUMMARY["n_per_input"] * EXP1_SUMMARY["n_inputs"]
    for g in (EXP2, EXP3, EXP4, EXP5A, EXP5B, [EXP5C_CONTROL], EXP6):
        n += sum(r.extra.get("events", r.n) for r in g)
    return n


# --------------------------------------------------------------------------
# Exporters (used to (re)generate the shipped CSVs from the embedded data)
# --------------------------------------------------------------------------
def export_summary_csv(path: str = SUMMARY_CSV) -> None:
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, "w", newline="") as fh:
        w = csv.DictWriter(fh, fieldnames=SUMMARY_FIELDS)
        w.writeheader()
        for exp, rows in _embedded().items():
            if exp == "exp1":
                continue
            for r in rows:
                e = r.extra
                w.writerow({
                    "experiment": exp, "label": r.label, "kind": r.kind,
                    "x": r.x, "n": r.n,
                    "target_port": "" if r.port is None else r.port,
                    "expected": "" if r.expected is None else r.expected,
                    "expected_port": e.get("expected_port", ""),
                    "fidelity": e.get("fidelity", ""),
                    "depth": e.get("depth", ""),
                    "hom_visibility": e.get("hom_visibility", ""),
                    "hom_from_report": "true" if e.get(
                        "hom_from_calibration_report") else "",
                    "c": e.get("c", ""),
                    "events": e.get("events", ""),
                })


def export_raw_template(path: str = RAW_CSV) -> None:
    """Write a per-port template: structure filled in, port columns blank.

    Each blank n0..n7 cell holds the exported per-output-port single-photon
    counts for that configuration; `build_raw_counts.py` fills them from the
    dataset.  Filled rows override the summary counts automatically.
    """
    os.makedirs(os.path.dirname(path), exist_ok=True)
    # eight Experiment-1 per-mode rows in addition to the rest.
    exp1_rows = [Row(label=f"input mode {k}", x=0, n=0, expected=0.125,
                     kind="control", port=0) for k in range(8)]
    groups = dict(_embedded())
    groups["exp1"] = exp1_rows
    order = ["exp1", "exp2", "exp3", "exp4", "exp5a", "exp5b", "exp5c", "exp6"]
    with open(path, "w", newline="") as fh:
        w = csv.DictWriter(fh, fieldnames=RAW_FIELDS)
        w.writeheader()
        for exp in order:
            for r in groups[exp]:
                e = r.extra
                row = {
                    "experiment": exp, "label": r.label, "kind": r.kind,
                    "metric": METRIC[exp],
                    "target_port": "" if r.port is None else r.port,
                    "expected_port": e.get("expected_port", ""),
                    "expected": "" if r.expected is None else r.expected,
                    "depth": e.get("depth", ""),
                    "hom_visibility": e.get("hom_visibility", ""),
                    "hom_from_report": "true" if e.get(
                        "hom_from_calibration_report") else "",
                    "c": e.get("c", ""),
                }
                for i in range(8):
                    row[f"n{i}"] = ""   # filled from the dataset
                w.writerow(row)


if __name__ == "__main__":  # pragma: no cover
    export_summary_csv()
    export_raw_template()
    print(f"wrote {SUMMARY_CSV}")
    print(f"wrote {RAW_CSV}")
    print("loaded total events:", f"{total_events():,}")
