qec-code 0.3.0

use crate::error::{QecError, Result};
use crate::finite_group::{
    FiniteGroupSpec, GroupAlgebraElement, left_regular_lift, right_regular_lift,
};

/// Maximum number of group-algebra cells in either ring-level lifted-product
/// check matrix. This bounds CLI input amplification before dense ring rows are
/// materialized.
pub const MAX_LIFTED_PRODUCT_RING_CELLS: usize = 1_000_000;

/// Maximum dense binary cells in either lifted CSS check matrix.
///
/// The common CSS construction contract verifies orthogonality and ranks after
/// construction. Those paths use pairwise sparse row intersections and dense
/// rank matrices, so the lifted-product constructor must bound the post-lift
/// matrices before materializing them.
pub const MAX_LIFTED_PRODUCT_BINARY_CELLS: usize = 10_000_000;

/// Maximum row pairs considered by the common CSS orthogonality check.
pub const MAX_LIFTED_PRODUCT_ORTHOGONALITY_ROW_PAIRS: usize = 10_000_000;

/// Ring-level lifted-product matrix shape before regular binary lifting.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct LiftedProductRingShape {
    pub h_x_rows: usize,
    pub h_z_rows: usize,
    pub num_cols: usize,
}

/// Ring-level lifted-product checks over the group algebra.
///
/// `h_x` has `shape.h_x_rows` rows and `shape.num_cols` columns; `h_z` has
/// `shape.h_z_rows` rows and the same columns. Transposed protograph blocks
/// use group inversion on every support element.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LiftedProductRingChecks {
    pub shape: LiftedProductRingShape,
    pub h_x: Vec<Vec<GroupAlgebraElement>>,
    pub h_z: Vec<Vec<GroupAlgebraElement>>,
}

/// Binary CSS checks after applying regular group lifts.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LiftedProductBinaryChecks {
    pub num_cols: usize,
    pub h_x: Vec<Vec<usize>>,
    pub h_z: Vec<Vec<usize>>,
}

/// Return the ring-level shape, rejecting overflow and oversized dense ring
/// materializations.
pub fn checked_lifted_product_ring_shape(
    left_rows: usize,
    left_cols: usize,
    right_rows: usize,
    right_cols: usize,
) -> Result<LiftedProductRingShape> {
    let h_x_rows = left_rows.checked_mul(right_cols).ok_or(ring_overflow())?;
    let h_z_rows = left_cols.checked_mul(right_rows).ok_or(ring_overflow())?;
    let left_block_cols = left_cols.checked_mul(right_cols).ok_or(ring_overflow())?;
    let right_block_cols = left_rows.checked_mul(right_rows).ok_or(ring_overflow())?;
    let num_cols = left_block_cols
        .checked_add(right_block_cols)
        .ok_or(ring_overflow())?;
    check_ring_cell_limit(h_x_rows, num_cols)?;
    check_ring_cell_limit(h_z_rows, num_cols)?;
    Ok(LiftedProductRingShape {
        h_x_rows,
        h_z_rows,
        num_cols,
    })
}

/// Return the ring-level shape and also preflight the post-lift binary shape.
pub fn checked_lifted_product_binary_shape(
    group: &FiniteGroupSpec,
    left_rows: usize,
    left_cols: usize,
    right_rows: usize,
    right_cols: usize,
) -> Result<LiftedProductRingShape> {
    let shape = checked_lifted_product_ring_shape(left_rows, left_cols, right_rows, right_cols)?;
    let binary_h_x_rows = checked_binary_dimension(shape.h_x_rows, group.order())?;
    let binary_h_z_rows = checked_binary_dimension(shape.h_z_rows, group.order())?;
    let binary_cols = checked_binary_dimension(shape.num_cols, group.order())?;
    check_binary_cell_limit("H_X", binary_h_x_rows, binary_cols)?;
    check_binary_cell_limit("H_Z", binary_h_z_rows, binary_cols)?;
    check_binary_row_pair_limit(binary_h_x_rows, binary_h_z_rows)?;
    Ok(shape)
}

/// Build ring-level lifted-product checks.
pub fn lifted_product_ring_checks(
    group: &FiniteGroupSpec,
    left: &[Vec<GroupAlgebraElement>],
    right: &[Vec<GroupAlgebraElement>],
) -> Result<LiftedProductRingChecks> {
    let (left_rows, left_cols) = group_algebra_matrix_shape(left)?;
    let (right_rows, right_cols) = group_algebra_matrix_shape(right)?;
    validate_group_orders(group, left)?;
    validate_group_orders(group, right)?;
    let shape = checked_lifted_product_ring_shape(left_rows, left_cols, right_rows, right_cols)?;

    let h_x = hconcat(
        &matrix_kron_identity(group, left, right_cols)?,
        &identity_kron_matrix(group, left_rows, &inverse_transpose(group, right)?)?,
    )?;
    let h_z = hconcat(
        &identity_kron_matrix(group, left_cols, right)?,
        &matrix_kron_identity(group, &inverse_transpose(group, left)?, right_rows)?,
    )?;
    debug_assert_eq!(h_x.len(), shape.h_x_rows);
    debug_assert_eq!(h_z.len(), shape.h_z_rows);
    debug_assert!(h_x.iter().all(|row| row.len() == shape.num_cols));
    debug_assert!(h_z.iter().all(|row| row.len() == shape.num_cols));
    Ok(LiftedProductRingChecks { shape, h_x, h_z })
}

/// Build binary CSS checks using commuting left/right regular actions.
///
/// Ring-level rows are ordered as the Kronecker constructors emit them: for
/// `A kron I_q`, rows are `(a_row, q_index)` and columns are
/// `(a_col, q_index)` with the left index major; for `I_p kron B`, rows are
/// `(p_index, b_row)` and columns are `(p_index, b_col)`. The binary lift then
/// expands every ring column into a contiguous `|G|` basis block.
///
/// For support element `g` and basis row `x`, the left action places a 1 at
/// `g^-1 x`, while the right action places a 1 at `x g`. The binary CSS check
/// order is:
///
/// - `H_X = [left_lift(A kron I), right_lift(I kron B^T)]`
/// - `H_Z = [right_lift(I kron inverse_entries(B)), left_lift(inverse_transpose(A) kron I)]`
///
/// This fixes the left qubit block before the right qubit block and makes the
/// extra involutions explicit for callers comparing ring and binary outputs.
pub fn lifted_product_binary_checks(
    group: &FiniteGroupSpec,
    left: &[Vec<GroupAlgebraElement>],
    right: &[Vec<GroupAlgebraElement>],
) -> Result<LiftedProductBinaryChecks> {
    let (left_rows, left_cols) = group_algebra_matrix_shape(left)?;
    let (right_rows, right_cols) = group_algebra_matrix_shape(right)?;
    validate_group_orders(group, left)?;
    validate_group_orders(group, right)?;
    let shape =
        checked_lifted_product_binary_shape(group, left_rows, left_cols, right_rows, right_cols)?;

    let h_x_left = matrix_kron_identity(group, left, right_cols)?;
    let h_x_right = identity_kron_matrix(group, left_rows, &transpose_without_inversion(right)?)?;
    let h_z_left = identity_kron_matrix(group, left_cols, &invert_entries(group, right)?)?;
    let h_z_right = matrix_kron_identity(group, &inverse_transpose(group, left)?, right_rows)?;

    let h_x =
        left_regular_lift(group, &h_x_left)?.hconcat(&right_regular_lift(group, &h_x_right)?)?;
    let h_z =
        right_regular_lift(group, &h_z_left)?.hconcat(&left_regular_lift(group, &h_z_right)?)?;
    debug_assert_eq!(h_x.num_cols(), h_z.num_cols());
    debug_assert_eq!(h_x.num_rows(), shape.h_x_rows * group.order());
    debug_assert_eq!(h_z.num_rows(), shape.h_z_rows * group.order());
    Ok(LiftedProductBinaryChecks {
        num_cols: h_x.num_cols(),
        h_x: h_x.rows().to_vec(),
        h_z: h_z.rows().to_vec(),
    })
}

fn ring_overflow() -> QecError {
    QecError::GroupAlgebraDimensionOverflow {
        operation: "lifted product ring shape",
    }
}

fn binary_overflow() -> QecError {
    QecError::GroupAlgebraDimensionOverflow {
        operation: "lifted product binary shape",
    }
}

fn checked_binary_dimension(value: usize, group_order: usize) -> Result<usize> {
    value.checked_mul(group_order).ok_or(binary_overflow())
}

fn check_ring_cell_limit(rows: usize, num_cols: usize) -> Result<()> {
    let cell_count = rows.checked_mul(num_cols).ok_or(ring_overflow())?;
    if cell_count > MAX_LIFTED_PRODUCT_RING_CELLS {
        return Err(QecError::InvalidCssConstruction {
            construction: "lifted_product".to_owned(),
            reason: format!(
                "ring cell count {cell_count} exceeds maximum supported {MAX_LIFTED_PRODUCT_RING_CELLS}"
            ),
        });
    }
    Ok(())
}

fn check_binary_cell_limit(matrix_name: &str, rows: usize, num_cols: usize) -> Result<()> {
    let cell_count = rows.checked_mul(num_cols).ok_or(binary_overflow())?;
    if cell_count > MAX_LIFTED_PRODUCT_BINARY_CELLS {
        return Err(QecError::InvalidCssConstruction {
            construction: "lifted_product".to_owned(),
            reason: format!(
                "binary {matrix_name} cell count {cell_count} exceeds maximum supported {MAX_LIFTED_PRODUCT_BINARY_CELLS}"
            ),
        });
    }
    Ok(())
}

fn check_binary_row_pair_limit(h_x_rows: usize, h_z_rows: usize) -> Result<()> {
    let row_pairs = h_x_rows.checked_mul(h_z_rows).ok_or(binary_overflow())?;
    if row_pairs > MAX_LIFTED_PRODUCT_ORTHOGONALITY_ROW_PAIRS {
        return Err(QecError::InvalidCssConstruction {
            construction: "lifted_product".to_owned(),
            reason: format!(
                "binary orthogonality row-pair count {row_pairs} exceeds maximum supported {MAX_LIFTED_PRODUCT_ORTHOGONALITY_ROW_PAIRS}"
            ),
        });
    }
    Ok(())
}

fn group_algebra_matrix_shape(matrix: &[Vec<GroupAlgebraElement>]) -> Result<(usize, usize)> {
    let Some(first_row) = matrix.first() else {
        return Err(invalid_protograph("must contain at least one row"));
    };
    if first_row.is_empty() {
        return Err(invalid_protograph("must contain at least one column"));
    }
    for row in matrix {
        if row.len() != first_row.len() {
            return Err(QecError::GroupAlgebraMatrixRowWidthMismatch {
                expected: first_row.len(),
                actual: row.len(),
            });
        }
    }
    Ok((matrix.len(), first_row.len()))
}

fn validate_group_orders(
    group: &FiniteGroupSpec,
    matrix: &[Vec<GroupAlgebraElement>],
) -> Result<()> {
    for row in matrix {
        for element in row {
            if element.group_order() != group.order() {
                return Err(QecError::GroupAlgebraOrderMismatch {
                    expected: group.order(),
                    actual: element.group_order(),
                });
            }
        }
    }
    Ok(())
}

fn invalid_protograph(reason: &str) -> QecError {
    QecError::InvalidCssConstruction {
        construction: "lifted_product".to_owned(),
        reason: reason.to_owned(),
    }
}

fn inverse_transpose(
    group: &FiniteGroupSpec,
    matrix: &[Vec<GroupAlgebraElement>],
) -> Result<Vec<Vec<GroupAlgebraElement>>> {
    let rows = matrix.len();
    let cols = matrix[0].len();
    let mut output = Vec::with_capacity(cols);
    for col in 0..cols {
        let mut row = Vec::with_capacity(rows);
        for source_row in matrix {
            let support = source_row[col]
                .support()
                .iter()
                .map(|&element| group.inverse(element))
                .collect::<Result<Vec<_>>>()?;
            row.push(GroupAlgebraElement::new(group, support)?);
        }
        output.push(row);
    }
    Ok(output)
}

fn transpose_without_inversion(
    matrix: &[Vec<GroupAlgebraElement>],
) -> Result<Vec<Vec<GroupAlgebraElement>>> {
    let rows = matrix.len();
    let cols = matrix[0].len();
    let mut output = Vec::with_capacity(cols);
    for col in 0..cols {
        let mut row = Vec::with_capacity(rows);
        for source_row in matrix {
            row.push(source_row[col].clone());
        }
        output.push(row);
    }
    Ok(output)
}

fn invert_entries(
    group: &FiniteGroupSpec,
    matrix: &[Vec<GroupAlgebraElement>],
) -> Result<Vec<Vec<GroupAlgebraElement>>> {
    matrix
        .iter()
        .map(|row| {
            row.iter()
                .map(|entry| {
                    let support = entry
                        .support()
                        .iter()
                        .map(|&element| group.inverse(element))
                        .collect::<Result<Vec<_>>>()?;
                    GroupAlgebraElement::new(group, support)
                })
                .collect()
        })
        .collect()
}

fn matrix_kron_identity(
    group: &FiniteGroupSpec,
    matrix: &[Vec<GroupAlgebraElement>],
    identity_size: usize,
) -> Result<Vec<Vec<GroupAlgebraElement>>> {
    let zero = GroupAlgebraElement::new(group, Vec::new())?;
    let mut output = Vec::with_capacity(matrix.len() * identity_size);
    for source_row in matrix {
        for diagonal in 0..identity_size {
            let mut row = Vec::with_capacity(source_row.len() * identity_size);
            for element in source_row {
                for column in 0..identity_size {
                    row.push(if column == diagonal {
                        element.clone()
                    } else {
                        zero.clone()
                    });
                }
            }
            output.push(row);
        }
    }
    Ok(output)
}

fn identity_kron_matrix(
    group: &FiniteGroupSpec,
    identity_size: usize,
    matrix: &[Vec<GroupAlgebraElement>],
) -> Result<Vec<Vec<GroupAlgebraElement>>> {
    let zero = GroupAlgebraElement::new(group, Vec::new())?;
    let matrix_cols = matrix[0].len();
    let mut output = Vec::with_capacity(identity_size * matrix.len());
    for diagonal in 0..identity_size {
        for source_row in matrix {
            let mut row = Vec::with_capacity(identity_size * matrix_cols);
            for block in 0..identity_size {
                if block == diagonal {
                    row.extend(source_row.iter().cloned());
                } else {
                    row.extend(std::iter::repeat_n(zero.clone(), matrix_cols));
                }
            }
            output.push(row);
        }
    }
    Ok(output)
}

fn hconcat(
    left: &[Vec<GroupAlgebraElement>],
    right: &[Vec<GroupAlgebraElement>],
) -> Result<Vec<Vec<GroupAlgebraElement>>> {
    if left.len() != right.len() {
        return Err(invalid_protograph("internal lifted-product row mismatch"));
    }
    Ok(left
        .iter()
        .zip(right)
        .map(|(left_row, right_row)| {
            let mut row = left_row.clone();
            row.extend(right_row.iter().cloned());
            row
        })
        .collect())
}