//! The two published record sets, read as they unpack.

use std::path::{Path, PathBuf};
use wai_quantum::quantum_frame::ErrorModel;
use wai_quantum::quantum_memory::{Shots, fit_priors, read_01, read_b8};
use wai_quantum::quantum_qec_text::parse_dem;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Set {
    /// doi:10.5281/zenodo.6804040: distance 3 (four patches) and 5, rounds
    /// 1 to 25, 50,000 shots an experiment.
    Y2023,
    /// doi:10.5281/zenodo.13273331, the 105-qubit set: distance 3 (nine
    /// patches), 5 (four) and 7 (one), rounds 1 to 250.
    Y2024,
}

impl Set {
    pub fn parse(s: &str) -> Option<Set> {
        match s {
            "2023" => Some(Set::Y2023),
            "2024" => Some(Set::Y2024),
            _ => None,
        }
    }
    pub fn name(self) -> &'static str {
        match self {
            Set::Y2023 => "2023",
            Set::Y2024 => "2024",
        }
    }
    /// The experiment, as cited.
    pub fn citation(self) -> &'static str {
        match self {
            Set::Y2023 => "Google Quantum AI, Suppressing quantum errors by scaling a surface code logical qubit, Nature 614, 676 (2023), doi:10.1038/s41586-022-05434-1 (allow-name: the citation and CC BY 4.0 credit the record's licence requires)",
            Set::Y2024 => "Google Quantum AI and Collaborators, Quantum error correction below the surface code threshold, Nature 638, 920 (2025), doi:10.1038/s41586-024-08449-y (allow-name: the citation and CC BY 4.0 credit the record's licence requires)",
        }
    }
    /// The record, pinned.
    pub fn source(self) -> serde_json::Value {
        match self {
            Set::Y2023 => serde_json::json!({
                "doi": "10.5281/zenodo.6804040",
                "url": "https://zenodo.org/api/records/6804040/files/google_qec3v5_experiment_data.zip/content", // allow-name: the record's own file name and its CC BY 4.0 credit
                "file": "google_qec3v5_experiment_data.zip", "bytes": 315_490_804u64, // allow-name: the record's own file name and its CC BY 4.0 credit
                "md5": "a7fd8b481c3087090093106382dc217d",
                "sha256": "ef569742126b305e163cc7fccbf6275e03bdd54d090dffb1aef12a78af0a74c4",
                "licence": "CC BY 4.0", "by": "Google Quantum AI", // allow-name: the record's own file name and its CC BY 4.0 credit
            }),
            Set::Y2024 => serde_json::json!({
                "doi": "10.5281/zenodo.13273331",
                "url": "https://zenodo.org/api/records/13273331/files/google_105Q_surface_code_d3_d5_d7.zip/content", // allow-name: the record's own file name and its CC BY 4.0 credit
                "file": "google_105Q_surface_code_d3_d5_d7.zip", "bytes": 5_716_907_033u64, // allow-name: the record's own file name and its CC BY 4.0 credit
                "md5": "21fa6ad35b395d838ebcdbc92e364a12",
                "sha256": "1e8e3b4f5f35ba4fd9b4a8448473f37a7090b404fd6eff0c00188092876070dd",
                "licence": "CC BY 4.0", "by": "Google Quantum AI", // allow-name: the record's own file name and its CC BY 4.0 credit
            }),
        }
    }
    /// The fit, as the receipt states it.
    pub fn fit_rule(self) -> &'static str {
        match self {
            Set::Y2023 => "per patch and basis, a weighted least-squares line through ln(1 − 2p) over rounds 3 to 25 (binomial weights; σ scaled by √(χ²/dof) when above one); ε averaged over patches and bases; Λ = ε₃/ε₅",
            Set::Y2024 => "per patch and basis, a weighted least-squares line through ln(1 − 2p) over rounds 10 to 250 less the 13-round calibration experiments (binomial weights; σ scaled by √(χ²/dof) when above one); ε averaged over patches and bases; Λ from a weighted line through ln ε_d against (d + 1)/2",
        }
    }
    /// Whether a publication's fits use an experiment of `rounds` rounds:
    /// from round 3 (2023); from round 10 (2024), less the 13-round
    /// experiments, which calibrated the optimized prior and carry no
    /// ensembled-synthesis predictions.
    pub fn fitted(self, rounds: u32) -> bool {
        match self {
            Set::Y2023 => rounds >= 3,
            Set::Y2024 => rounds >= 10 && rounds != 13,
        }
    }
}

/// One experiment of a set.
#[derive(Clone, Debug)]
pub struct Experiment {
    pub name: String,
    pub dir: PathBuf,
    pub distance: u32,
    pub basis: String,
    pub rounds: u32,
    pub patch: String,
    pub shots: usize,
    pub detectors: usize,
}

fn kv(text: &str, key: &str) -> Option<String> {
    text.lines().find_map(|l| l.strip_prefix(key).and_then(|r| r.strip_prefix(':')).map(|v| v.trim().to_string()))
}

/// Every surface-code experiment of the set under `dir`, in name order.
pub fn list(set: Set, dir: &Path) -> Vec<Experiment> {
    let mut out = Vec::new();
    match set {
        Set::Y2023 => {
            let mut names: Vec<String> = std::fs::read_dir(dir).unwrap().map(|e| e.unwrap().file_name().into_string().unwrap()).filter(|n| n.starts_with("surface_code_")).collect();
            names.sort();
            for n in names {
                let d = dir.join(&n);
                let p = std::fs::read_to_string(d.join("properties.yml")).unwrap();
                let num = |k: &str| kv(&p, k).unwrap_or_else(|| panic!("{n}: no {k}")).parse::<usize>().unwrap();
                let (row, col) = (num("center_data_qubit_row"), num("center_data_qubit_col"));
                out.push(Experiment {
                    name: n.clone(),
                    dir: d.clone(),
                    distance: num("distance") as u32,
                    basis: kv(&p, "basis").unwrap(),
                    rounds: num("rounds") as u32,
                    patch: format!("{row}_{col}"),
                    shots: num("shots"),
                    detectors: num("circuit_detectors"),
                });
            }
        }
        Set::Y2024 => {
            let root = dir.join("google_105Q_surface_code_d3_d5_d7"); // allow-name: the record's own file name and its CC BY 4.0 credit
            let mut patches: Vec<String> = std::fs::read_dir(&root).unwrap().map(|e| e.unwrap().file_name().into_string().unwrap()).filter(|n| n.starts_with('d') && n.contains("_at_")).collect();
            patches.sort();
            for patch in patches {
                for basis in ["X", "Z"] {
                    let bdir = root.join(&patch).join(basis);
                    if !bdir.is_dir() {
                        continue;
                    }
                    let mut rs: Vec<String> = std::fs::read_dir(&bdir).unwrap().map(|e| e.unwrap().file_name().into_string().unwrap()).filter(|n| n.starts_with('r')).collect();
                    rs.sort();
                    for r in rs {
                        let d = bdir.join(&r);
                        let meta: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(d.join("metadata.json")).unwrap()).unwrap();
                        let get = |k: &str| meta.get(k).and_then(serde_json::Value::as_u64).unwrap_or_else(|| panic!("{}: metadata has no {k}: {meta}", d.display())) as usize;
                        // The circuit is written out round by round, one
                        // DETECTOR line a detector.
                        let circuit = std::fs::read_to_string(d.join("circuit_ideal.stim")).unwrap();
                        assert!(!circuit.contains("REPEAT"), "{}: a circuit with blocks", d.display());
                        let detectors = circuit.lines().filter(|l| l.starts_with("DETECTOR")).count();
                        out.push(Experiment {
                            name: format!("{patch}/{basis}/{r}"),
                            dir: d.clone(),
                            distance: patch[1..].split('_').next().unwrap().parse().unwrap(),
                            basis: basis.to_string(),
                            rounds: r[1..].parse().unwrap(),
                            patch: patch.split("_at_").nth(1).unwrap().to_string(),
                            shots: get("shots"),
                            detectors,
                        });
                    }
                }
            }
        }
    }
    out
}

/// The experiment's detection events and observable flips.
pub fn shots(set: Set, e: &Experiment) -> Shots {
    let det = std::fs::read(e.dir.join("detection_events.b8")).unwrap();
    let obs = match set {
        Set::Y2023 => read_01(&std::fs::read_to_string(e.dir.join("obs_flips_actual.01")).unwrap(), 1, e.shots).unwrap(),
        Set::Y2024 => read_b8(&std::fs::read(e.dir.join("obs_flips_actual.b8")).unwrap(), 1, e.shots).unwrap(),
    };
    Shots::from_b8(det, e.detectors, obs, 1).unwrap_or_else(|err| panic!("{}: {err}", e.name))
}

/// The published decoders' predictions for the experiment's shots `skip` to
/// `skip + n`: each decoder's name and its predictions on those shots, shot
/// `skip` first.
pub fn published(set: Set, e: &Experiment, skip: usize, n: usize) -> Vec<(String, Vec<u64>)> {
    let mut out = Vec::new();
    let cut = |pred: Vec<u64>| -> Vec<u64> { pred[skip..skip + n].to_vec() };
    match set {
        Set::Y2023 => {
            for name in ["pymatching", "correlated_matching", "belief_matching", "tensor_network_contraction"] {
                let f = e.dir.join(format!("obs_flips_predicted_by_{name}.01"));
                if let Ok(t) = std::fs::read_to_string(&f) {
                    out.push((name.to_string(), cut(read_01(&t, 1, e.shots).unwrap())));
                }
            }
        }
        Set::Y2024 => {
            let dr = e.dir.join("decoding_results");
            if let Ok(rd) = std::fs::read_dir(&dr) {
                let mut names: Vec<String> = rd.map(|x| x.unwrap().file_name().into_string().unwrap()).collect();
                names.sort();
                for name in names {
                    let f = dr.join(&name).join("obs_flips_predicted.b8");
                    if let Ok(b) = std::fs::read(&f) {
                        out.push((name, cut(read_b8(&b, 1, e.shots).unwrap())));
                    }
                }
            }
        }
    }
    out
}

fn read_model(path: &Path) -> ErrorModel {
    let t = std::fs::read_to_string(path).unwrap_or_else(|e| panic!("{}: {e}", path.display()));
    parse_dem(&t).unwrap_or_else(|e| panic!("{}: {e}", path.display())).to_model().unwrap_or_else(|e| panic!("{}: {e}", path.display()))
}

/// The models a prior names, one per shot residue: shot `s` is decoded with
/// model `s mod len`.
pub fn models(set: Set, e: &Experiment, prior: &str, shots: &Shots) -> Vec<ErrorModel> {
    // This crate's fit, over a published model's structure: even shots
    // decoded with the model fitted to the odd ones, and odd with even.
    let fitted = |structure: ErrorModel| -> Vec<ErrorModel> {
        let (for_even, a) = fit_priors(&structure, &shots.select(|s| s % 2 == 1), true, 200);
        let (for_odd, b) = fit_priors(&structure, &shots.select(|s| s % 2 == 0), true, 200);
        eprintln!("{}: fitted {} symptom sets to {} parities ({} and {} sweeps, residual {:.3e} → {:.3e})", e.name, a.unknowns, a.equations, a.sweeps, b.sweeps, a.residual.0, a.residual.1);
        vec![for_even, for_odd]
    };
    match (set, prior) {
        (Set::Y2023, "fit") => fitted(read_model(&e.dir.join("circuit_detector_error_model.dem"))),
        (Set::Y2023, "fitpij") => fitted(read_model(&e.dir.join("pij_from_even_for_odd.dem"))),
        (Set::Y2024, "fit") => fitted(read_model(&e.dir.join("decoding_results/correlated_matching_decoder_with_si1000_prior/error_model.dem"))),
        (Set::Y2024, "fitrl") => fitted(read_model(&e.dir.join("decoding_results/correlated_matching_decoder_with_rl_optimized_prior/error_model.dem"))),
        // Even shots with the model fitted to the odd ones, and odd shots
        // with the model fitted to the even ones.
        (Set::Y2023, "pij") => vec![read_model(&e.dir.join("pij_from_odd_for_even.dem")), read_model(&e.dir.join("pij_from_even_for_odd.dem"))],
        (Set::Y2023, "circuit") => vec![read_model(&e.dir.join("circuit_detector_error_model.dem"))],
        (Set::Y2024, "si1000") => vec![read_model(&e.dir.join("decoding_results/correlated_matching_decoder_with_si1000_prior/error_model.dem"))],
        (Set::Y2024, "rl") => vec![read_model(&e.dir.join("decoding_results/correlated_matching_decoder_with_rl_optimized_prior/error_model.dem"))],
        _ => panic!("no prior {prior} for set {}", set.name()),
    }
}