qec-code 0.3.0

use std::fs;
use std::path::PathBuf;

use clap::{Args, Parser, Subcommand, ValueEnum};

use crate::QecError;
use crate::codes::built_in_css::{built_in_css_catalog, built_in_css_checks};
use crate::codes::quantum_tanner::{
    QuantumTannerSpec, quantum_tanner_css_checks, quantum_tanner_spec_from_json_str,
};
use crate::codes::steane::Steane;
use crate::css::{CssCode, SparseRowsMatrix, sparse_rows_matrix_from_json_str};
use crate::distance::{compute_distance, compute_distance_with_solver_options};
use crate::distance_bound::{
    RandomWindowUpperBoundOptions, RandomizedUpperBoundOptions, random_window_css_upper_bound,
    randomized_css_upper_bound,
};
use crate::distance_exact::{
    ExactCssDistanceBackend, ExactCssDistanceInput, ExactCssDistanceOptions,
    ExactCssDistanceResult, ExactCssDistanceSolverOptions,
};
use crate::error::CssMatrixReadSource;
use crate::family_contract::{
    CssConstructionSpec, CssFamilySpec, construct_css, parse_css_construction_json,
};
use crate::family_verifier::{FamilyVerificationReport, verify_checked_in_family_manifest};

#[derive(Debug, Parser)]
#[command(name = "qec-code")]
pub struct Cli {
    #[command(subcommand)]
    pub command: Commands,
}

#[derive(Debug, Subcommand)]
pub enum Commands {
    Code {
        #[command(subcommand)]
        command: CodeCommands,
    },
}

#[derive(Debug, Subcommand)]
pub enum CodeCommands {
    Steane {
        #[command(subcommand)]
        command: SteaneCommands,
    },
    Css(CssArgs),
    CssDistance {
        #[command(subcommand)]
        command: CssDistanceCommands,
    },
}

#[derive(Debug, Args)]
#[command(args_conflicts_with_subcommands = true)]
#[command(subcommand_negates_reqs = true)]
#[command(arg_required_else_help = true)]
pub struct CssArgs {
    #[command(subcommand)]
    command: Option<CssCommands>,
    #[arg(value_name = "CODE_ID", required = true)]
    code_id: Option<String>,
    #[arg(value_name = "MATRIX", required = true)]
    matrix: Option<CssMatrixKind>,
}

impl CssArgs {
    pub fn list() -> Self {
        Self {
            command: Some(CssCommands::List),
            code_id: None,
            matrix: None,
        }
    }

    pub fn export(code_id: String, matrix: CssMatrixKind) -> Self {
        Self {
            command: None,
            code_id: Some(code_id),
            matrix: Some(matrix),
        }
    }

    pub fn export_subcommand(code_id: String, matrix: CssMatrixKind) -> Self {
        Self {
            command: Some(CssCommands::Export { code_id, matrix }),
            code_id: None,
            matrix: None,
        }
    }
}

#[derive(Debug, Subcommand)]
pub enum CssCommands {
    List,
    VerifyFamilies,
    Export {
        code_id: String,
        matrix: CssMatrixKind,
    },
    Construct {
        #[arg(long)]
        spec: PathBuf,
        output: CssConstructionOutput,
    },
    QuantumTanner {
        #[arg(long)]
        spec: PathBuf,
        matrix: CssMatrixKind,
    },
}

#[derive(Debug, Clone, Copy, ValueEnum)]
pub enum CssMatrixKind {
    Hx,
    Hz,
}

#[derive(Debug, Clone, Copy, ValueEnum)]
pub enum CssConstructionOutput {
    Hx,
    Hz,
    Metadata,
}

#[derive(Debug, Subcommand)]
pub enum CssDistanceCommands {
    Exact(ExactCssDistanceCli),
    RandomizedUpperBound(RandomizedUpperBoundCli),
    RandomWindowUpperBound(RandomWindowUpperBoundCli),
}

#[derive(Debug, Args)]
pub struct ExactCssDistanceCli {
    #[arg(long)]
    code_id: Option<String>,
    #[arg(long)]
    hx: Option<PathBuf>,
    #[arg(long)]
    hz: Option<PathBuf>,
    #[arg(long)]
    quantum_tanner_spec: Option<PathBuf>,
    #[arg(long, value_enum, default_value_t = ExactCssDistanceBackend::Auto)]
    backend: ExactCssDistanceBackend,
    #[arg(long, allow_hyphen_values = true)]
    time_limit_seconds: Option<f64>,
    #[arg(long, allow_hyphen_values = true)]
    mip_gap: Option<f64>,
    #[arg(long)]
    threads: Option<u32>,
    #[arg(long)]
    verbose_solver: bool,
    #[arg(long)]
    json: bool,
}

#[derive(Debug, Args)]
pub struct RandomizedUpperBoundCli {
    #[arg(long)]
    code_id: Option<String>,
    #[arg(long)]
    hx: Option<PathBuf>,
    #[arg(long)]
    hz: Option<PathBuf>,
    #[arg(long)]
    quantum_tanner_spec: Option<PathBuf>,
    #[arg(long)]
    iterations: usize,
    #[arg(long, default_value_t = 1)]
    restarts: usize,
    #[arg(long)]
    seed: u64,
    #[arg(long)]
    target_weight: Option<usize>,
    #[arg(long)]
    json: bool,
}

#[derive(Debug, Args)]
pub struct RandomWindowUpperBoundCli {
    #[arg(long)]
    code_id: Option<String>,
    #[arg(long)]
    hx: Option<PathBuf>,
    #[arg(long)]
    hz: Option<PathBuf>,
    #[arg(long)]
    quantum_tanner_spec: Option<PathBuf>,
    #[arg(long)]
    iterations: usize,
    #[arg(long, default_value_t = 1)]
    restarts: usize,
    #[arg(long)]
    seed: u64,
    #[arg(long)]
    target_weight: Option<usize>,
    #[arg(long)]
    json: bool,
}

#[derive(Debug, Subcommand)]
pub enum SteaneCommands {
    Summary,
    Stabilizers,
    Logicals,
    Distance,
}

pub fn run(cli: Cli) -> Result<String, QecError> {
    match cli.command {
        Commands::Code { command } => run_code(command),
    }
}

fn run_code(command: CodeCommands) -> Result<String, QecError> {
    match command {
        CodeCommands::Steane { command } => run_steane(command),
        CodeCommands::Css(args) => run_css_args(args),
        CodeCommands::CssDistance { command } => run_css_distance(command),
    }
}

fn run_css_args(args: CssArgs) -> Result<String, QecError> {
    match args.command {
        Some(CssCommands::List) => Ok(run_css_list()),
        Some(CssCommands::VerifyFamilies) => run_css_verify_families(),
        Some(CssCommands::Export { code_id, matrix }) => run_css(&code_id, matrix),
        Some(CssCommands::Construct { spec, output }) => run_css_construction_spec(&spec, output),
        Some(CssCommands::QuantumTanner { spec, matrix }) => run_css_quantum_tanner(&spec, matrix),
        None => {
            let code_id = args
                .code_id
                .expect("clap requires CODE_ID when no css subcommand is used");
            let matrix = args
                .matrix
                .expect("clap requires MATRIX when no css subcommand is used");

            run_css(&code_id, matrix)
        }
    }
}

fn run_css_verify_families() -> Result<String, QecError> {
    family_report_to_cli_result(verify_checked_in_family_manifest()?)
}

fn family_report_to_cli_result(report: FamilyVerificationReport) -> Result<String, QecError> {
    if report.failed == 0 {
        Ok(report.output)
    } else {
        Err(QecError::FamilyVerificationFailed {
            report: report.output,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn family_report_to_cli_result_returns_error_for_failed_report() {
        let output = "FAIL generalized_bicycle expected rank_x=5 actual rank_x=4\nSUMMARY FAIL supported=12 deferred=2 failed=1".to_owned();
        let report = FamilyVerificationReport {
            output: output.clone(),
            failed: 1,
        };

        let error = family_report_to_cli_result(report).unwrap_err();

        assert_eq!(error, QecError::FamilyVerificationFailed { report: output });
    }
}

fn run_css_list() -> String {
    let catalog = built_in_css_catalog();
    let width = catalog
        .iter()
        .map(|entry| entry.spec.len())
        .max()
        .unwrap_or(0);
    let mut lines = Vec::with_capacity(catalog.len() + 1);

    lines.push("Built-in CSS codes:".to_owned());
    lines.extend(catalog.iter().map(|entry| {
        format!(
            "  {:width$}  {}",
            entry.spec,
            entry.description,
            width = width
        )
    }));

    lines.join("\n")
}

fn run_css(code_id: &str, matrix: CssMatrixKind) -> Result<String, QecError> {
    let spec = CssConstructionSpec::from_inline(code_id)?;
    let construction = construct_css(spec)?;
    let rows = match matrix {
        CssMatrixKind::Hx => construction.checks.h_x,
        CssMatrixKind::Hz => construction.checks.h_z,
    };

    let matrix = SparseRowsMatrix::new(construction.stats.n, rows)?;
    Ok(matrix.to_json_string())
}

fn run_css_construction_spec(
    path: &PathBuf,
    output: CssConstructionOutput,
) -> Result<String, QecError> {
    let input = read_css_spec_file(path)?;
    let spec = parse_css_construction_json(&input)?;
    export_css_construction_output(spec, output)
}

fn run_css_quantum_tanner(spec: &PathBuf, matrix: CssMatrixKind) -> Result<String, QecError> {
    let spec = read_quantum_tanner_spec(spec)?;
    export_css_construction(CssFamilySpec::QuantumTanner(spec).into(), matrix)
}

fn export_css_construction(
    spec: CssConstructionSpec,
    matrix: CssMatrixKind,
) -> Result<String, QecError> {
    export_css_construction_matrix(construct_css(spec)?, matrix)
}

fn export_css_construction_output(
    spec: CssConstructionSpec,
    output: CssConstructionOutput,
) -> Result<String, QecError> {
    let construction = construct_css(spec)?;
    match output {
        CssConstructionOutput::Hx => {
            export_css_construction_matrix(construction, CssMatrixKind::Hx)
        }
        CssConstructionOutput::Hz => {
            export_css_construction_matrix(construction, CssMatrixKind::Hz)
        }
        CssConstructionOutput::Metadata => Ok(serde_json::to_string(&construction)
            .expect("validated CSS construction result should always serialize")),
    }
}

fn export_css_construction_matrix(
    construction: crate::family_contract::CssConstructionResult,
    matrix: CssMatrixKind,
) -> Result<String, QecError> {
    let rows = match matrix {
        CssMatrixKind::Hx => construction.checks.h_x,
        CssMatrixKind::Hz => construction.checks.h_z,
    };

    let matrix = SparseRowsMatrix::new(construction.stats.n, rows)?;
    Ok(matrix.to_json_string())
}

fn read_quantum_tanner_spec(path: &PathBuf) -> Result<QuantumTannerSpec, QecError> {
    let input = read_css_spec_file(path)?;

    quantum_tanner_spec_from_json_str(&input)
}

fn read_css_spec_file(path: &PathBuf) -> Result<String, QecError> {
    fs::read_to_string(path).map_err(|err| QecError::CssMatrixReadFailed {
        path: path.display().to_string(),
        source: CssMatrixReadSource(err.to_string()),
    })
}

fn run_css_distance(command: CssDistanceCommands) -> Result<String, QecError> {
    match command {
        CssDistanceCommands::Exact(options) => run_css_exact_distance(options),
        CssDistanceCommands::RandomizedUpperBound(options) => {
            run_css_randomized_upper_bound(options)
        }
        CssDistanceCommands::RandomWindowUpperBound(options) => {
            run_css_random_window_upper_bound(options)
        }
    }
}

fn run_css_exact_distance(cli: ExactCssDistanceCli) -> Result<String, QecError> {
    const COMMAND: &str = "code css-distance exact";

    if !cli.json {
        return Err(QecError::JsonOutputRequired { command: COMMAND });
    }

    let (css, options) = css_code_and_exact_options_from_cli(&cli)?;
    let computation = compute_distance_with_solver_options(css.code(), options.solver)?;
    let result = ExactCssDistanceResult::completed_with_solver_report(
        computation.distance,
        options,
        computation.solver_report,
    );

    serde_json::to_string(&result).map_err(|err| QecError::InvalidCssDistanceInput(err.to_string()))
}

fn css_code_and_exact_options_from_cli(
    cli: &ExactCssDistanceCli,
) -> Result<(CssCode, ExactCssDistanceOptions), QecError> {
    let solver = validate_exact_css_solver_options(cli)?;
    match css_distance_input_selection(&cli.code_id, &cli.hx, &cli.hz, &cli.quantum_tanner_spec)? {
        CssDistanceInputSelection::CodeId(code_id) => Ok((
            css_code_from_built_in(code_id)?,
            ExactCssDistanceOptions {
                input: ExactCssDistanceInput::CodeId {
                    code_id: code_id.to_owned(),
                },
                solver,
            },
        )),
        CssDistanceInputSelection::Files { hx, hz } => Ok((
            css_code_from_files(hx, hz)?,
            ExactCssDistanceOptions {
                input: ExactCssDistanceInput::Files {
                    hx: hx.display().to_string(),
                    hz: hz.display().to_string(),
                },
                solver,
            },
        )),
        CssDistanceInputSelection::QuantumTannerSpec(spec) => Ok((
            css_code_from_quantum_tanner_spec(spec)?,
            ExactCssDistanceOptions {
                input: ExactCssDistanceInput::QuantumTannerSpec {
                    quantum_tanner_spec: spec.display().to_string(),
                },
                solver,
            },
        )),
    }
}

fn validate_exact_css_solver_options(
    cli: &ExactCssDistanceCli,
) -> Result<ExactCssDistanceSolverOptions, QecError> {
    if let Some(value) = cli.time_limit_seconds {
        if !value.is_finite() || value <= 0.0 {
            return Err(QecError::InvalidCssDistanceInput(
                "invalid exact CSS distance solver option time_limit_seconds".to_owned(),
            ));
        }
    }
    if let Some(value) = cli.mip_gap {
        if !value.is_finite() || value < 0.0 {
            return Err(QecError::InvalidCssDistanceInput(
                "invalid exact CSS distance solver option mip_gap".to_owned(),
            ));
        }
    }
    if let Some(value) = cli.threads {
        if value == 0 {
            return Err(QecError::InvalidCssDistanceInput(
                "invalid exact CSS distance solver option threads".to_owned(),
            ));
        }
    }

    Ok(ExactCssDistanceSolverOptions {
        backend: cli.backend,
        time_limit_seconds: cli.time_limit_seconds,
        mip_gap: cli.mip_gap,
        threads: cli.threads,
        verbose_solver: cli.verbose_solver,
    })
}

fn run_css_randomized_upper_bound(cli: RandomizedUpperBoundCli) -> Result<String, QecError> {
    const COMMAND: &str = "code css-distance randomized-upper-bound";

    if !cli.json {
        return Err(QecError::JsonOutputRequired { command: COMMAND });
    }

    let css = css_code_from_randomized_upper_bound_cli(&cli)?;
    let options = RandomizedUpperBoundOptions {
        iterations: cli.iterations,
        restarts: cli.restarts,
        seed: cli.seed,
        target_weight: cli.target_weight,
    };
    let result = randomized_css_upper_bound(&css, options)?;

    serde_json::to_string(&result).map_err(|err| QecError::InvalidCssDistanceInput(err.to_string()))
}

fn css_code_from_randomized_upper_bound_cli(
    cli: &RandomizedUpperBoundCli,
) -> Result<CssCode, QecError> {
    match css_distance_input_selection(&cli.code_id, &cli.hx, &cli.hz, &cli.quantum_tanner_spec)? {
        CssDistanceInputSelection::CodeId(code_id) => css_code_from_built_in(code_id),
        CssDistanceInputSelection::Files { hx, hz } => css_code_from_files(hx, hz),
        CssDistanceInputSelection::QuantumTannerSpec(spec) => {
            css_code_from_quantum_tanner_spec(spec)
        }
    }
}

fn run_css_random_window_upper_bound(cli: RandomWindowUpperBoundCli) -> Result<String, QecError> {
    const COMMAND: &str = "code css-distance random-window-upper-bound";

    if !cli.json {
        return Err(QecError::JsonOutputRequired { command: COMMAND });
    }

    let css = css_code_from_random_window_upper_bound_cli(&cli)?;
    let options = RandomWindowUpperBoundOptions {
        iterations: cli.iterations,
        restarts: cli.restarts,
        seed: cli.seed,
        target_weight: cli.target_weight,
    };
    let result = random_window_css_upper_bound(&css, options)?;

    serde_json::to_string(&result).map_err(|err| QecError::InvalidCssDistanceInput(err.to_string()))
}

fn css_code_from_random_window_upper_bound_cli(
    cli: &RandomWindowUpperBoundCli,
) -> Result<CssCode, QecError> {
    match css_distance_input_selection(&cli.code_id, &cli.hx, &cli.hz, &cli.quantum_tanner_spec)? {
        CssDistanceInputSelection::CodeId(code_id) => css_code_from_built_in(code_id),
        CssDistanceInputSelection::Files { hx, hz } => css_code_from_files(hx, hz),
        CssDistanceInputSelection::QuantumTannerSpec(spec) => {
            css_code_from_quantum_tanner_spec(spec)
        }
    }
}

enum CssDistanceInputSelection<'a> {
    CodeId(&'a str),
    Files { hx: &'a PathBuf, hz: &'a PathBuf },
    QuantumTannerSpec(&'a PathBuf),
}

fn css_distance_input_selection<'a>(
    code_id: &'a Option<String>,
    hx: &'a Option<PathBuf>,
    hz: &'a Option<PathBuf>,
    quantum_tanner_spec: &'a Option<PathBuf>,
) -> Result<CssDistanceInputSelection<'a>, QecError> {
    let source_count = usize::from(code_id.is_some())
        + usize::from(hx.is_some() || hz.is_some())
        + usize::from(quantum_tanner_spec.is_some());

    if source_count == 0 {
        return Err(QecError::InvalidCssDistanceInput(
            "provide --code-id, --quantum-tanner-spec, or both --hx and --hz".to_owned(),
        ));
    }

    if source_count > 1 {
        return Err(QecError::InvalidCssDistanceInput(
            "use only one input source: --code-id, --quantum-tanner-spec, or --hx/--hz".to_owned(),
        ));
    }

    match (
        code_id.as_deref(),
        hx.as_ref(),
        hz.as_ref(),
        quantum_tanner_spec.as_ref(),
    ) {
        (Some(code_id), None, None, None) => Ok(CssDistanceInputSelection::CodeId(code_id)),
        (None, Some(hx), Some(hz), None) => Ok(CssDistanceInputSelection::Files { hx, hz }),
        (None, None, None, Some(spec)) => Ok(CssDistanceInputSelection::QuantumTannerSpec(spec)),
        (None, Some(_), None, None) | (None, None, Some(_), None) => Err(
            QecError::InvalidCssDistanceInput("--hx and --hz must be provided together".to_owned()),
        ),
        _ => Err(QecError::InvalidCssDistanceInput(
            "use only one input source: --code-id, --quantum-tanner-spec, or --hx/--hz".to_owned(),
        )),
    }
}

fn css_code_from_built_in(code_id: &str) -> Result<CssCode, QecError> {
    let checks = built_in_css_checks(code_id)?;
    let hx = SparseRowsMatrix::new(checks.num_cols, checks.hx)?.to_dense_rows();
    let hz = SparseRowsMatrix::new(checks.num_cols, checks.hz)?.to_dense_rows();
    CssCode::from_hx_hz(hx, hz)
}

fn css_code_from_quantum_tanner_spec(path: &PathBuf) -> Result<CssCode, QecError> {
    let spec = read_quantum_tanner_spec(path)?;
    let checks = quantum_tanner_css_checks(&spec)?;
    let hx = SparseRowsMatrix::new(checks.num_cols, checks.hx)?.to_dense_rows();
    let hz = SparseRowsMatrix::new(checks.num_cols, checks.hz)?.to_dense_rows();

    CssCode::from_hx_hz(hx, hz)
}

fn css_code_from_files(hx_path: &PathBuf, hz_path: &PathBuf) -> Result<CssCode, QecError> {
    let hx = read_css_sparse_rows_matrix(hx_path)?;
    let hz = read_css_sparse_rows_matrix(hz_path)?;

    if hx.num_cols() != hz.num_cols() {
        return Err(QecError::InvalidCssDistanceInput(format!(
            "hx width {} does not match hz width {}",
            hx.num_cols(),
            hz.num_cols()
        )));
    }

    CssCode::from_hx_hz(hx.to_dense_rows(), hz.to_dense_rows())
}

fn read_css_sparse_rows_matrix(path: &PathBuf) -> Result<SparseRowsMatrix, QecError> {
    let input = fs::read_to_string(path).map_err(|err| QecError::CssMatrixReadFailed {
        path: path.display().to_string(),
        source: CssMatrixReadSource(err.to_string()),
    })?;

    sparse_rows_matrix_from_json_str(&input)
}

fn run_steane(command: SteaneCommands) -> Result<String, QecError> {
    let steane = Steane::new()?;
    let code = steane.code();

    match command {
        SteaneCommands::Summary => Ok(format!(
            "name: steane\nn: {}\nstabilizer_rank: {}\nk: {}",
            code.n(),
            code.stabilizer_rank(),
            code.num_logical_qubits()
        )),
        SteaneCommands::Stabilizers => {
            let lines = code
                .stabilizers()
                .iter()
                .enumerate()
                .map(|(index, stabilizer)| format!("g{}: {}", index + 1, format_pauli(stabilizer)))
                .collect::<Vec<_>>();
            Ok(lines.join("\n"))
        }
        SteaneCommands::Logicals => {
            let basis = code.logical_basis()?;
            Ok(format!(
                "k: {}\nlogical_x:\n{}\nlogical_z:\n{}",
                basis.k,
                format_pauli_list(&basis.logical_x),
                format_pauli_list(&basis.logical_z)
            ))
        }
        SteaneCommands::Distance => {
            let distance = compute_distance(code)?;
            Ok(format!(
                "distance: {}\nlogical_class: {:?}\nwitness: {}",
                distance.distance,
                distance.logical_class,
                format_pauli(&distance.witness)
            ))
        }
    }
}

fn format_pauli_list(paulis: &[crate::Pauli]) -> String {
    paulis
        .iter()
        .enumerate()
        .map(|(index, pauli)| format!("  {}: {}", index + 1, format_pauli(pauli)))
        .collect::<Vec<_>>()
        .join("\n")
}

fn format_pauli(pauli: &crate::Pauli) -> String {
    format!(
        "x={:?} z={:?} weight={}",
        pauli.x_bits(),
        pauli.z_bits(),
        pauli.weight()
    )
}