#!/usr/bin/env python3
"""Audit the photon-number distribution of every Belenos export.

For each *.result.json we decode the BSCount and tally how many recorded
events had 0, 1, 2, or >=2 detected photons across the eight ports. This
quantifies the directly-observed multi-photon detection rate and lets us
bound the multi-photon contamination of the single-photon postselected set.

Two scopes are reported:
  * the manuscript six-experiment job set, taken from the unique
    ``raw_result_file`` entries of ``data/manifest/job_manifest.csv``;
  * the full archive of all ``*.result.json`` exports (which includes
    additional completed campaign jobs beyond the manuscript tables).
"""

from __future__ import annotations

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

ROOT = Path(__file__).resolve().parents[1]
RAW_DIR = ROOT / "data" / "raw" / "quandela_exports"
MANIFEST = ROOT / "data" / "manifest" / "job_manifest.csv"


def decode_pcvl_zip(value: str) -> str:
    value = value.lstrip(":")
    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)
    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(p) for p in state_text.strip("|<>").split(","))
        counts[state] = int(count_text)
    return counts


def tally(path: Path) -> tuple[int, int, int, int]:
    """Return (n0, n1, n2, n3plus) detected-photon-number counts for one export."""
    data = json.loads(path.read_text())
    counts = parse_bscount(decode_pcvl_zip(data["results"]))
    n0 = n1 = n2 = n3p = 0
    for state, c in counts.items():
        s = sum(state)
        if s == 0:
            n0 += c
        elif s == 1:
            n1 += c
        elif s == 2:
            n2 += c
        else:
            n3p += c
    return n0, n1, n2, n3p


def manuscript_files() -> list[Path]:
    """Unique raw result files backing the six manuscript experiments."""
    seen: dict[str, Path] = {}
    with MANIFEST.open() as fh:
        for row in csv.DictReader(fh):
            fname = row["raw_result_file"].strip()
            if fname and fname not in seen:
                seen[fname] = RAW_DIR / fname
    return list(seen.values())


def report(label: str, files: list[Path]) -> None:
    tot_n0 = tot_n1 = tot_n2 = tot_n3p = 0
    worst = 0.0
    for path in files:
        n0, n1, n2, n3p = tally(path)
        tot_n0 += n0
        tot_n1 += n1
        tot_n2 += n2
        tot_n3p += n3p
        all_f = n0 + n1 + n2 + n3p
        if all_f:
            worst = max(worst, (n2 + n3p) / all_f)
    tot_all = tot_n0 + tot_n1 + tot_n2 + tot_n3p
    multi = tot_n2 + tot_n3p
    print(f"\n=== {label} ({len(files)} exports) ===")
    print(f"all detected records : {tot_all}")
    print(f"n=1 (single-photon)  : {tot_n1}")
    print(f"n=2 (two-photon)     : {tot_n2}")
    print(f"n>=3                 : {tot_n3p}")
    print(f"multi (n>=2)         : {multi}  "
          f"({100 * multi / tot_all:.3f}% of all records)")
    print(f"worst single config  : {100 * worst:.3f}% multi-photon")
    # Lost-twin contamination: a 2-photon emission detected as 1 photon.
    # ratio (1 detected)/(2 detected) = 2 eta (1-eta)/eta^2 = 2(1-eta)/eta.
    eta = 0.0484
    factor = 2 * (1 - eta) / eta
    masq = multi * factor
    print(f"lost-twin factor 2(1-eta)/eta = {factor:.1f} (eta={eta})")
    print(f"=> masquerading 2-as-1 events ~ {masq:.0f} "
          f"({100 * masq / tot_n1:.2f}% of single-photon counts)")


def main() -> None:
    report("MANUSCRIPT six-experiment job set", manuscript_files())
    report("FULL archive", sorted(RAW_DIR.glob("*.result.json")))


if __name__ == "__main__":
    main()
