use std::fmt::Display;
use std::fmt::Formatter;
use std::fmt::Result as FmtResult;

use itertools::Itertools;
use ndarray::ArrayView1;
use twenty_first::math::traits::FiniteField;
use twenty_first::prelude::*;

use crate::table::challenges::Challenges;
use crate::table::master_table::MasterExtTable;

/// The implementations of these functions are automatically generated using the
/// command `cargo run --bin constraint-evaluation-generator` and live in
/// `constraints.rs`.
pub trait Evaluable<FF: FiniteField> {
    /// The code for this method must be generated by running
    /// `cargo run --bin constraint-evaluation-generator`
    fn evaluate_initial_constraints(
        base_row: ArrayView1<FF>,
        ext_row: ArrayView1<XFieldElement>,
        challenges: &Challenges,
    ) -> Vec<XFieldElement>;

    /// The code for this method must be generated by running
    /// `cargo run --bin constraint-evaluation-generator`
    fn evaluate_consistency_constraints(
        base_row: ArrayView1<FF>,
        ext_row: ArrayView1<XFieldElement>,
        challenges: &Challenges,
    ) -> Vec<XFieldElement>;

    /// The code for this method must be generated by running
    /// `cargo run --bin constraint-evaluation-generator`
    fn evaluate_transition_constraints(
        current_base_row: ArrayView1<FF>,
        current_ext_row: ArrayView1<XFieldElement>,
        next_base_row: ArrayView1<FF>,
        next_ext_row: ArrayView1<XFieldElement>,
        challenges: &Challenges,
    ) -> Vec<XFieldElement>;

    /// The code for this method must be generated by running
    /// `cargo run --bin constraint-evaluation-generator`
    fn evaluate_terminal_constraints(
        base_row: ArrayView1<FF>,
        ext_row: ArrayView1<XFieldElement>,
        challenges: &Challenges,
    ) -> Vec<XFieldElement>;
}

pub trait Quotientable: Evaluable<BFieldElement> {
    const NUM_INITIAL_CONSTRAINTS: usize;
    const NUM_CONSISTENCY_CONSTRAINTS: usize;
    const NUM_TRANSITION_CONSTRAINTS: usize;
    const NUM_TERMINAL_CONSTRAINTS: usize;

    /// The total number of constraints. The number of quotients is identical.
    const NUM_CONSTRAINTS: usize = Self::NUM_INITIAL_CONSTRAINTS
        + Self::NUM_CONSISTENCY_CONSTRAINTS
        + Self::NUM_TRANSITION_CONSTRAINTS
        + Self::NUM_TERMINAL_CONSTRAINTS;

    fn initial_quotient_degree_bounds(interpolant_degree: isize) -> Vec<isize>;

    fn consistency_quotient_degree_bounds(
        interpolant_degree: isize,
        padded_height: usize,
    ) -> Vec<isize>;

    fn transition_quotient_degree_bounds(
        interpolant_degree: isize,
        padded_height: usize,
    ) -> Vec<isize>;

    fn terminal_quotient_degree_bounds(interpolant_degree: isize) -> Vec<isize>;
}

/// The type of constraint. Can be used to determine the degree bounds for the quotient
/// polynomials. Concretely, the degree of the zerofier polynomials differs between the
/// constraint types.
#[derive(Debug, Copy, Clone, Eq, PartialEq, Ord, PartialOrd)]
pub(crate) enum ConstraintType {
    Initial,
    Consistency,
    Transition,
    Terminal,
}

impl Display for ConstraintType {
    fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
        match self {
            ConstraintType::Initial => write!(f, "initial"),
            ConstraintType::Consistency => write!(f, "consistency"),
            ConstraintType::Transition => write!(f, "transition"),
            ConstraintType::Terminal => write!(f, "terminal"),
        }
    }
}

/// Helps debugging and benchmarking. The maximal degree achieved in any table dictates the length
/// of the FRI domain, which in turn is responsible for the main performance bottleneck.
#[derive(Debug, Clone, Eq, PartialEq, Ord, PartialOrd)]
pub(crate) struct DegreeWithOrigin {
    pub degree: isize,
    pub interpolant_degree: isize,
    pub zerofier_degree: isize,
    pub origin_index: usize,
    pub origin_table_height: usize,
    pub origin_constraint_type: ConstraintType,
}

impl Display for DegreeWithOrigin {
    fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
        assert!(self.degree > 0);
        let zerofier_corrected_degree = self.degree + self.zerofier_degree;
        let degree = zerofier_corrected_degree / self.interpolant_degree;
        let idx = self.origin_index;
        let constraint_type = self.origin_constraint_type;
        write!(
            f,
            "Degree of polynomial {idx:02} of type “{constraint_type}” is {degree}."
        )
    }
}

/// Compute the degrees of the quotients from all AIR constraints that apply to the table.
pub(crate) fn all_degrees_with_origin(
    interpolant_degree: isize,
    padded_height: usize,
) -> Vec<DegreeWithOrigin> {
    let initial_degrees_with_origin =
        MasterExtTable::initial_quotient_degree_bounds(interpolant_degree)
            .into_iter()
            .enumerate()
            .map(|(origin_index, degree)| DegreeWithOrigin {
                degree,
                interpolant_degree,
                zerofier_degree: 1,
                origin_index,
                origin_table_height: padded_height,
                origin_constraint_type: ConstraintType::Initial,
            })
            .collect_vec();

    let consistency_degrees_with_origin =
        MasterExtTable::consistency_quotient_degree_bounds(interpolant_degree, padded_height)
            .into_iter()
            .enumerate()
            .map(|(origin_index, degree)| DegreeWithOrigin {
                degree,
                interpolant_degree,
                zerofier_degree: padded_height as isize,
                origin_index,
                origin_table_height: padded_height,
                origin_constraint_type: ConstraintType::Consistency,
            })
            .collect();

    let transition_degrees_with_origin =
        MasterExtTable::transition_quotient_degree_bounds(interpolant_degree, padded_height)
            .into_iter()
            .enumerate()
            .map(|(origin_index, degree)| DegreeWithOrigin {
                degree,
                interpolant_degree,
                zerofier_degree: padded_height as isize - 1,
                origin_index,
                origin_table_height: padded_height,
                origin_constraint_type: ConstraintType::Transition,
            })
            .collect();

    let terminal_degrees_with_origin =
        MasterExtTable::terminal_quotient_degree_bounds(interpolant_degree)
            .into_iter()
            .enumerate()
            .map(|(origin_index, degree)| DegreeWithOrigin {
                degree,
                interpolant_degree,
                zerofier_degree: 1,
                origin_index,
                origin_table_height: padded_height,
                origin_constraint_type: ConstraintType::Terminal,
            })
            .collect();

    [
        initial_degrees_with_origin,
        consistency_degrees_with_origin,
        transition_degrees_with_origin,
        terminal_degrees_with_origin,
    ]
    .concat()
}