#!/usr/bin/env python3
"""Decode Quandela/Perceval job exports into per-port count tables.

Input files are the JSON result exports downloaded from Quandela Cloud.
The script expects filenames ending in ``.result.json`` and writes:

  * data/processed/per_port_counts.csv
  * data/processed/job_summary.csv

It decodes the ``:PCVL:zip:...`` BSCount payload, keeps both all detected
events and the single-photon postselected subset, and reports Wilson 95%
confidence intervals for port-0 dump probabilities.
"""

from __future__ import annotations

import base64
import csv
import json
import math
import re
import zlib
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
RAW_DIR = ROOT / "data" / "raw" / "quandela_exports"
OUT_DIR = ROOT / "data" / "processed"


def wilson_ci(successes: int, total: int, z: float = 1.96) -> tuple[float, float]:
    if total == 0:
        return (float("nan"), float("nan"))
    phat = successes / total
    denom = 1 + z * z / total
    center = (phat + z * z / (2 * total)) / denom
    half = z * math.sqrt(phat * (1 - phat) / total + z * z / (4 * total * total)) / denom
    return center - half, center + half


def decode_pcvl_zip(value: str) -> str:
    value = value.lstrip(":")
    if not value.startswith("PCVL:zip:"):
        raise ValueError(f"Unsupported Perceval serialization prefix: {value[:32]}")
    payload = value[len("PCVL:zip:") :]
    return zlib.decompress(base64.b64decode(payload)).decode()


def parse_bscount(text: str) -> dict[tuple[int, ...], int]:
    match = re.search(r"BSCount:\{(.*)\}$", text)
    if not match:
        raise ValueError("Decoded payload does not contain a BSCount object")

    counts: dict[tuple[int, ...], int] = {}
    for entry in match.group(1).split(";"):
        if not entry.strip():
            continue
        state_text, count_text = entry.split("=")
        state = tuple(int(part) for part in state_text.strip("|<>").split(","))
        counts[state] = int(count_text)
    return counts


def load_counts(path: Path) -> dict[tuple[int, ...], int]:
    data = json.loads(path.read_text())
    return parse_bscount(decode_pcvl_zip(data["results"]))


def summarize_file(path: Path) -> tuple[list[dict[str, object]], dict[str, object]]:
    job_name = path.name.removesuffix(".result.json")
    counts = load_counts(path)

    all_total = sum(counts.values())
    single_total = 0
    port_counts = {port: 0 for port in range(8)}

    for state, count in counts.items():
        if sum(state) != 1:
            continue
        single_total += count
        # final_mode_number is 8 in these experiments; ignore any unused modes.
        port = state.index(1)
        if port < 8:
            port_counts[port] += count

    rows = []
    for port, count in port_counts.items():
        rows.append(
            {
                "job_name": job_name,
                "port": port,
                "count": count,
                "single_photon_total": single_total,
                "probability": count / single_total if single_total else "",
            }
        )

    dump = port_counts[0]
    lo, hi = wilson_ci(dump, single_total)
    summary = {
        "job_name": job_name,
        "all_detected_events": all_total,
        "single_photon_total": single_total,
        "dump_count": dump,
        "dump_probability": dump / single_total if single_total else "",
        "dump_wilson95_low": lo,
        "dump_wilson95_high": hi,
    }
    return rows, summary


def main() -> None:
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    result_files = sorted(RAW_DIR.glob("*.result.json"))
    if not result_files:
        raise SystemExit(f"No *.result.json files found in {RAW_DIR}")

    port_rows: list[dict[str, object]] = []
    summary_rows: list[dict[str, object]] = []
    for path in result_files:
        rows, summary = summarize_file(path)
        port_rows.extend(rows)
        summary_rows.append(summary)

    with (OUT_DIR / "per_port_counts.csv").open("w", newline="") as fh:
        writer = csv.DictWriter(
            fh,
            fieldnames=["job_name", "port", "count", "single_photon_total", "probability"],
        )
        writer.writeheader()
        writer.writerows(port_rows)

    with (OUT_DIR / "job_summary.csv").open("w", newline="") as fh:
        writer = csv.DictWriter(
            fh,
            fieldnames=[
                "job_name",
                "all_detected_events",
                "single_photon_total",
                "dump_count",
                "dump_probability",
                "dump_wilson95_low",
                "dump_wilson95_high",
            ],
        )
        writer.writeheader()
        writer.writerows(summary_rows)

    print(f"Decoded {len(result_files)} result export(s)")
    print(f"Wrote {OUT_DIR / 'per_port_counts.csv'}")
    print(f"Wrote {OUT_DIR / 'job_summary.csv'}")


if __name__ == "__main__":
    main()
