use alloc::vec;
use alloc::vec::Vec;
use core::fmt::Debug;

use p3_matrix::dense::RowMajorMatrix;
use p3_matrix::{Dimensions, Matrix};
use serde::de::DeserializeOwned;
use serde::Serialize;

/// A "Mixed Matrix Commitment Scheme" (MMCS) is a generalization of a vector commitment scheme; it
/// supports committing to matrices and then opening rows. It is also batch-oriented; one can commit
/// to a batch of matrices at once even if their widths and heights differ.
///
/// When a particular row index is opened, it is interpreted directly as a row index for matrices
/// with the largest height. For matrices with smaller heights, some bits of the row index are
/// removed (from the least-significant side) to get the effective row index. These semantics are
/// useful in the FRI protocol. See the documentation for `open_batch` for more details.
pub trait Mmcs<T: Send + Sync>: Clone {
    type ProverData<M>;
    type Commitment: Clone + Serialize + DeserializeOwned + Send + Sync;
    type Proof: Clone + Serialize + DeserializeOwned + Send + Sync;
    type Error: Debug;

    fn commit<M: Matrix<T>>(&self, inputs: Vec<M>) -> (Self::Commitment, Self::ProverData<M>);

    fn commit_matrix<M: Matrix<T>>(&self, input: M) -> (Self::Commitment, Self::ProverData<M>) {
        self.commit(vec![input])
    }

    fn commit_vec(&self, input: Vec<T>) -> (Self::Commitment, Self::ProverData<RowMajorMatrix<T>>)
    where
        T: Clone + Send + Sync,
    {
        self.commit_matrix(RowMajorMatrix::new_col(input))
    }

    /// Opens a batch of rows from committed matrices
    /// returns `(openings, proof)`
    /// where `openings` is a vector whose `i`th element is the `j`th row of the ith matrix `M[i]`,
    /// and `j = index >> (log2_ceil(max_height) - log2_ceil(M[i].height))`.
    fn open_batch<M: Matrix<T>>(
        &self,
        index: usize,
        prover_data: &Self::ProverData<M>,
    ) -> (Vec<Vec<T>>, Self::Proof);

    /// Get the matrices that were committed to.
    fn get_matrices<'a, M: Matrix<T>>(&self, prover_data: &'a Self::ProverData<M>) -> Vec<&'a M>;

    fn get_matrix_heights<M: Matrix<T>>(&self, prover_data: &Self::ProverData<M>) -> Vec<usize> {
        self.get_matrices(prover_data)
            .iter()
            .map(|matrix| matrix.height())
            .collect()
    }

    /// Get the largest height of any committed matrix.
    fn get_max_height<M: Matrix<T>>(&self, prover_data: &Self::ProverData<M>) -> usize {
        self.get_matrix_heights(prover_data)
            .into_iter()
            .max()
            .unwrap_or_else(|| panic!("No committed matrices?"))
    }

    /// Verify a batch opening.
    /// `index` is the row index we're opening for each matrix, following the same
    /// semantics as `open_batch`.
    /// `dimensions` is a slice whose ith element is the dimensions of the matrix being opened
    /// in the ith opening
    fn verify_batch(
        &self,
        commit: &Self::Commitment,
        dimensions: &[Dimensions],
        index: usize,
        opened_values: &[Vec<T>],
        proof: &Self::Proof,
    ) -> Result<(), Self::Error>;

    // ============================================================
    // Path-pruning extension (Phase 2.1 / B-Stage 1')
    // ============================================================

    /// Pruned proof for a batch of `N` row indices opened together against the
    /// same committed data. Compared to `Vec<Self::Proof>`, this representation
    /// shares Merkle siblings between paths and can save 50-75% of the wire
    /// bytes when `N` is large relative to `log_max_height`.
    ///
    /// Implementations that do not optimize this path may set
    /// `type PrunedProof = Vec<Self::Proof>` and have `open_batch_pruned`
    /// just loop over `open_batch`.
    type PrunedProof: Clone + Serialize + DeserializeOwned + Send + Sync;

    /// Open a batch of rows with shared sibling deduplication.
    ///
    /// `indices` is a slice of `N` row indices (interpreted at the largest
    /// matrix height; smaller matrices follow the same `index >> shift`
    /// semantics as `open_batch`). Indices need not be sorted; implementations
    /// must internally sort + dedup and the returned values follow that order.
    ///
    /// Returns:
    ///   * `openings_per_query` (len = N_unique), each entry being one
    ///     per-query opening in the same shape as `open_batch`, and
    ///   * a single pruned proof that covers all N_unique paths jointly.
    ///     Implementations are expected to embed the sorted index list inside
    ///     the proof so verifiers can recover query order without an extra arg.
    fn open_batch_pruned<M: Matrix<T>>(
        &self,
        indices: &[usize],
        prover_data: &Self::ProverData<M>,
    ) -> (Vec<Vec<Vec<T>>>, Self::PrunedProof);

    /// Verify a pruned batch opening produced by `open_batch_pruned`.
    ///
    /// The query indices are recovered from the proof itself (they were
    /// embedded by `open_batch_pruned`). Callers that need to cross-check
    /// against externally-supplied indices should do so at their own layer.
    ///
    /// `opened_values_per_query[i]` must match the i-th opened query in the
    /// proof's internal sorted order, in the same shape returned by `open_batch`.
    fn verify_batch_pruned(
        &self,
        commit: &Self::Commitment,
        dimensions: &[Dimensions],
        opened_values_per_query: &[Vec<Vec<T>>],
        proof: &Self::PrunedProof,
    ) -> Result<(), Self::Error>;

    /// Fixed-frontier proof for a batch of `N` row indices opened together
    /// against the same committed data.
    ///
    /// Unlike [`Self::PrunedProof`], this proof does not do exact query-path
    /// deduplication. Instead, every query carries a fixed-length path up to a
    /// frontier layer, and the whole frontier is opened to the root. This keeps
    /// verifier control flow fixed for a given `(N, log_height, top_layers)`.
    type FixedTopProof: Clone + Serialize + DeserializeOwned + Send + Sync;

    /// Open a batch of rows with fixed-frontier path sharing.
    ///
    /// `top_layers` is the requested number of tree layers above the frontier.
    /// Implementations may reduce it when the committed matrix shape requires
    /// matrix-injection below the frontier.
    fn open_batch_fixed_top<M: Matrix<T>>(
        &self,
        indices: &[usize],
        top_layers: usize,
        prover_data: &Self::ProverData<M>,
    ) -> (Vec<Vec<Vec<T>>>, Self::FixedTopProof);

    /// Verify a fixed-frontier batch opening.
    ///
    /// Unlike [`Self::verify_batch_pruned`], the verifier supplies the expected
    /// query indices explicitly and implementations must cross-check them
    /// against the proof.
    fn verify_batch_fixed_top(
        &self,
        commit: &Self::Commitment,
        dimensions: &[Dimensions],
        indices: &[usize],
        opened_values_per_query: &[Vec<Vec<T>>],
        proof: &Self::FixedTopProof,
    ) -> Result<(), Self::Error>;
}
