symbolica 0.9.0

use std::fmt::{self, Display, Write};

use colored::Colorize;

use crate::{
    atom::{
        representation::FunView, AddView, Atom, AtomView, MulView, NumView, PowView, Symbol,
        VarView,
    },
    coefficient::CoefficientView,
    domains::{
        factorized_rational_polynomial::FactorizedRationalPolynomial,
        finite_field::FiniteFieldCore, rational_polynomial::RationalPolynomial, Ring, RingPrinter,
    },
    poly::{polynomial::MultivariatePolynomial, Exponent, MonomialOrder},
    state::State,
    tensors::matrix::Matrix,
};

#[derive(Debug, Copy, Clone)]
pub struct PrintOptions {
    pub terms_on_new_line: bool,
    pub color_top_level_sum: bool,
    pub color_builtin_symbols: bool,
    pub print_finite_field: bool,
    pub symmetric_representation_for_finite_field: bool,
    pub explicit_rational_polynomial: bool,
    pub number_thousands_separator: Option<char>,
    pub multiplication_operator: char,
    pub double_star_for_exponentiation: bool,
    pub square_brackets_for_function: bool,
    pub num_exp_as_superscript: bool,
    pub latex: bool,
}

impl PrintOptions {
    /// Print the output in a Mathematica-readable format.
    pub const fn mathematica() -> PrintOptions {
        Self {
            terms_on_new_line: false,
            color_top_level_sum: false,
            color_builtin_symbols: false,
            print_finite_field: true,
            symmetric_representation_for_finite_field: false,
            explicit_rational_polynomial: false,
            number_thousands_separator: None,
            multiplication_operator: ' ',
            double_star_for_exponentiation: false,
            square_brackets_for_function: true,
            num_exp_as_superscript: false,
            latex: false,
        }
    }

    /// Print the output in a Latex input format.
    pub const fn latex() -> PrintOptions {
        Self {
            terms_on_new_line: false,
            color_top_level_sum: false,
            color_builtin_symbols: false,
            print_finite_field: true,
            symmetric_representation_for_finite_field: false,
            explicit_rational_polynomial: false,
            number_thousands_separator: None,
            multiplication_operator: ' ',
            double_star_for_exponentiation: false,
            square_brackets_for_function: false,
            num_exp_as_superscript: false,
            latex: true,
        }
    }

    /// Print the output suitable for a file.
    pub const fn file() -> PrintOptions {
        Self {
            terms_on_new_line: false,
            color_top_level_sum: false,
            color_builtin_symbols: false,
            print_finite_field: false,
            symmetric_representation_for_finite_field: false,
            explicit_rational_polynomial: false,
            number_thousands_separator: None,
            multiplication_operator: '*',
            double_star_for_exponentiation: false,
            square_brackets_for_function: false,
            num_exp_as_superscript: false,
            latex: false,
        }
    }

    /// Print the output in a sympy input format.
    pub const fn sympy() -> PrintOptions {
        Self {
            double_star_for_exponentiation: true,
            ..Self::file()
        }
    }
}

impl Default for PrintOptions {
    fn default() -> Self {
        Self {
            terms_on_new_line: false,
            color_top_level_sum: true,
            color_builtin_symbols: true,
            print_finite_field: true,
            symmetric_representation_for_finite_field: false,
            explicit_rational_polynomial: false,
            number_thousands_separator: None,
            multiplication_operator: '*',
            double_star_for_exponentiation: false,
            square_brackets_for_function: false,
            num_exp_as_superscript: false,
            latex: false,
        }
    }
}

#[derive(Debug, Copy, Clone)]
pub struct PrintState {
    pub level: usize,
    pub top_level_add_child: bool,
    pub explicit_sign: bool,
    pub superscript: bool,
}

macro_rules! define_formatters {
    ($($a:ident),*) => {
        $(
        trait $a {
            fn fmt_debug(
                &self,
                f: &mut fmt::Formatter,
            ) -> fmt::Result;

            fn fmt_output(
                &self,
                f: &mut fmt::Formatter,
                print_opts: &PrintOptions,
                print_state: PrintState,
            ) -> fmt::Result;
        })+
    };
}

define_formatters!(
    FormattedPrintVar,
    FormattedPrintNum,
    FormattedPrintFn,
    FormattedPrintPow,
    FormattedPrintMul,
    FormattedPrintAdd
);

pub struct AtomPrinter<'a> {
    pub atom: AtomView<'a>,
    pub print_opts: PrintOptions,
}

impl<'a> AtomPrinter<'a> {
    /// Create a new atom printer with default printing options.
    pub fn new(atom: AtomView<'a>) -> AtomPrinter<'a> {
        AtomPrinter {
            atom,
            print_opts: PrintOptions::default(),
        }
    }

    pub fn new_with_options(atom: AtomView<'a>, print_opts: PrintOptions) -> AtomPrinter<'a> {
        AtomPrinter { atom, print_opts }
    }
}

impl<'a> fmt::Display for AtomPrinter<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let print_state = PrintState {
            level: 0,
            top_level_add_child: false,
            explicit_sign: false,
            superscript: false,
        };
        self.atom.fmt_output(f, &self.print_opts, print_state)
    }
}

impl std::fmt::Display for Symbol {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(State::get_name(*self))
    }
}

impl Atom {
    /// Construct a printer for the atom with special options.
    pub fn printer<'a>(&'a self, opts: PrintOptions) -> AtomPrinter<'a> {
        AtomPrinter::new_with_options(self.as_view(), opts)
    }

    /// Print the atom in a form that is unique and independent of any implementation details.
    ///     
    /// Anti-symmetric functions are not supported.
    pub fn to_canonical_string(&self) -> String {
        self.as_view().to_canonical_string()
    }
}

impl<'a> AtomView<'a> {
    fn fmt_debug(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        match self {
            AtomView::Num(n) => n.fmt_debug(fmt),
            AtomView::Var(v) => v.fmt_debug(fmt),
            AtomView::Fun(f) => f.fmt_debug(fmt),
            AtomView::Pow(p) => p.fmt_debug(fmt),
            AtomView::Mul(m) => m.fmt_debug(fmt),
            AtomView::Add(a) => a.fmt_debug(fmt),
        }
    }

    fn fmt_output(
        &self,
        fmt: &mut fmt::Formatter,
        opts: &PrintOptions,
        print_state: PrintState,
    ) -> fmt::Result {
        match self {
            AtomView::Num(n) => n.fmt_output(fmt, opts, print_state),
            AtomView::Var(v) => v.fmt_output(fmt, opts, print_state),
            AtomView::Fun(f) => f.fmt_output(fmt, opts, print_state),
            AtomView::Pow(p) => p.fmt_output(fmt, opts, print_state),
            AtomView::Mul(t) => t.fmt_output(fmt, opts, print_state),
            AtomView::Add(e) => e.fmt_output(fmt, opts, print_state),
        }
    }

    /// Construct a printer for the atom with special options.
    pub fn printer(&self, opts: PrintOptions) -> AtomPrinter {
        AtomPrinter::new_with_options(*self, opts)
    }

    /// Print the atom in a form that is unique and independent of any implementation details.
    ///
    /// Anti-symmetric functions are not supported.
    pub fn to_canonical_string(&self) -> String {
        let mut s = String::new();
        self.to_canonical_view_impl(&mut s);
        s
    }

    fn to_canonical_view_impl(&self, out: &mut String) {
        fn add_paren(cur: AtomView, s: AtomView) -> bool {
            if let AtomView::Pow(_) = cur {
                matches!(s, AtomView::Add(_) | AtomView::Mul(_))
            } else if let AtomView::Mul(_) = cur {
                matches!(s, AtomView::Add(_))
            } else {
                false
            }
        }

        match self {
            AtomView::Num(_) => write!(out, "{}", self).unwrap(),
            AtomView::Var(v) => write!(out, "{}", v.get_symbol()).unwrap(),
            AtomView::Fun(f) => {
                write!(out, "{}(", f.get_symbol()).unwrap();

                let mut args = vec![];

                for x in f.iter() {
                    let mut arg = String::new();
                    x.to_canonical_view_impl(&mut arg);
                    args.push(arg);
                }

                // TODO: anti-symmetric may generate minus sign...
                if f.is_symmetric() {
                    args.sort();
                }

                if f.is_antisymmetric() {
                    unimplemented!(
                        "Antisymmetric functions are not supported yet for canonical view"
                    );
                }

                for (i, arg) in args.iter().enumerate() {
                    if i > 0 {
                        write!(out, ",").unwrap();
                    }
                    write!(out, "{}", arg).unwrap();
                }

                write!(out, ")").unwrap();
            }
            AtomView::Pow(p) => {
                let (b, e) = p.get_base_exp();

                if add_paren(*self, b) {
                    write!(out, "(").unwrap();
                    b.to_canonical_view_impl(out);
                    write!(out, ")").unwrap();
                } else {
                    b.to_canonical_view_impl(out);
                }

                if add_paren(*self, e) {
                    write!(out, "^(").unwrap();
                    e.to_canonical_view_impl(out);
                    write!(out, ")").unwrap();
                } else {
                    write!(out, "^").unwrap();
                    e.to_canonical_view_impl(out);
                }
            }
            AtomView::Mul(m) => {
                let mut terms = vec![];

                for x in m.iter() {
                    let mut term = if add_paren(*self, x) {
                        "(".to_string()
                    } else {
                        String::new()
                    };

                    x.to_canonical_view_impl(&mut term);

                    if add_paren(*self, x) {
                        term.push(')');
                    }

                    terms.push(term);
                }

                terms.sort();

                for (i, term) in terms.iter().enumerate() {
                    if i > 0 {
                        write!(out, "*").unwrap();
                    }

                    write!(out, "{}", term).unwrap();
                }
            }
            AtomView::Add(a) => {
                let mut terms = vec![];

                for x in a.iter() {
                    let mut term = if add_paren(*self, x) {
                        "(".to_string()
                    } else {
                        String::new()
                    };

                    x.to_canonical_view_impl(&mut term);

                    if add_paren(*self, x) {
                        term.push(')');
                    }

                    terms.push(term);
                }

                terms.sort();

                for (i, term) in terms.iter().enumerate() {
                    if i > 0 {
                        write!(out, "+").unwrap();
                    }
                    write!(out, "{}", term).unwrap();
                }
            }
        }
    }
}

impl<'a> fmt::Debug for AtomView<'a> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        self.fmt_debug(fmt)
    }
}

impl<'a> FormattedPrintVar for VarView<'a> {
    fn fmt_output(
        &self,
        f: &mut fmt::Formatter,
        opts: &PrintOptions,
        print_state: PrintState,
    ) -> fmt::Result {
        if print_state.explicit_sign {
            if print_state.top_level_add_child && opts.color_top_level_sum {
                f.write_fmt(format_args!("{}", "+".yellow()))?;
            } else {
                f.write_char('+')?;
            }
        }

        let id = self.get_symbol();
        let name = State::get_name(id);

        if opts.latex {
            match id {
                State::E => f.write_char('e'),
                State::PI => f.write_str("\\pi"),
                State::I => f.write_char('i'),
                _ => f.write_str(name),
            }
        } else if opts.color_builtin_symbols && name.ends_with('_') {
            f.write_fmt(format_args!("{}", name.cyan().italic()))
        } else if opts.color_builtin_symbols && State::is_builtin(id) {
            f.write_fmt(format_args!("{}", name.purple()))
        } else {
            f.write_str(name)
        }
    }

    fn fmt_debug(&self, f: &mut fmt::Formatter) -> fmt::Result {
        <Self as std::fmt::Debug>::fmt(self, f)
    }
}

impl<'a> FormattedPrintNum for NumView<'a> {
    fn fmt_debug(&self, f: &mut fmt::Formatter) -> fmt::Result {
        <Self as std::fmt::Debug>::fmt(self, f)
    }

    fn fmt_output(
        &self,
        f: &mut fmt::Formatter,
        opts: &PrintOptions,
        print_state: PrintState,
    ) -> fmt::Result {
        /// Input must be digits only.
        fn format_num(
            mut s: String,
            opts: &PrintOptions,
            print_state: &PrintState,
            f: &mut fmt::Formatter,
        ) -> fmt::Result {
            if print_state.superscript {
                let map = ['⁰', '¹', '²', '³', '⁴', '⁵', '⁶', '⁷', '⁸', '⁹'];
                s = s
                    .as_bytes()
                    .iter()
                    .map(|x| map[(x - b'0') as usize])
                    .collect();

                return f.write_str(&s);
            }

            if let Some(c) = opts.number_thousands_separator {
                let mut first = true;
                for triplet in s.as_bytes().chunks(3) {
                    if !first {
                        f.write_char(c)?;
                    }
                    f.write_str(std::str::from_utf8(triplet).unwrap())?;
                    first = false;
                }

                Ok(())
            } else {
                f.write_str(&s)
            }
        }

        let d = self.get_coeff_view();

        let is_negative = match d {
            CoefficientView::Natural(n, _) => n < 0,
            CoefficientView::Large(r) => r.is_negative(),
            _ => false,
        };

        if is_negative {
            if print_state.top_level_add_child && opts.color_top_level_sum {
                f.write_fmt(format_args!("{}", "-".yellow()))?;
            } else if print_state.superscript {
                f.write_char('⁻')?;
            } else {
                f.write_char('-')?;
            }
        } else if print_state.explicit_sign {
            if print_state.top_level_add_child && opts.color_top_level_sum {
                f.write_fmt(format_args!("{}", "+".yellow()))?;
            } else {
                f.write_char('+')?;
            }
        }

        match d {
            CoefficientView::Natural(num, den) => {
                if !opts.latex
                    && (opts.number_thousands_separator.is_some() || print_state.superscript)
                {
                    format_num(num.unsigned_abs().to_string(), opts, &print_state, f)?;
                    if den != 1 {
                        f.write_char('/')?;
                        format_num(den.to_string(), opts, &print_state, f)?;
                    }
                    Ok(())
                } else if den != 1 {
                    if opts.latex {
                        f.write_fmt(format_args!("\\frac{{{}}}{{{}}}", num.unsigned_abs(), den))
                    } else {
                        f.write_fmt(format_args!("{}/{}", num.unsigned_abs(), den))
                    }
                } else {
                    f.write_fmt(format_args!("{}", num.unsigned_abs()))
                }
            }
            CoefficientView::Float(fl) => {
                let float = fl.to_float();
                f.write_fmt(format_args!("{}", float))
            }
            CoefficientView::Large(r) => {
                let rat = r.to_rat().abs();
                if !opts.latex
                    && (opts.number_thousands_separator.is_some() || print_state.superscript)
                {
                    format_num(rat.numerator().to_string(), opts, &print_state, f)?;
                    if !rat.is_integer() {
                        f.write_char('/')?;
                        format_num(rat.denominator().to_string(), opts, &print_state, f)?;
                    }
                    Ok(())
                } else if !rat.is_integer() {
                    if opts.latex {
                        f.write_fmt(format_args!(
                            "\\frac{{{}}}{{{}}}",
                            rat.numerator(),
                            rat.denominator(),
                        ))
                    } else {
                        f.write_fmt(format_args!(
                            "{}/{}",
                            rat.numerator_ref(),
                            rat.denominator_ref()
                        ))
                    }
                } else {
                    f.write_fmt(format_args!("{}", rat.numerator_ref()))
                }
            }
            CoefficientView::FiniteField(num, fi) => {
                let ff = State::get_finite_field(fi);
                f.write_fmt(format_args!(
                    "[{}%{}]",
                    ff.from_element(&num),
                    ff.get_prime()
                ))
            }
            CoefficientView::RationalPolynomial(p) => f.write_fmt(format_args!(
                "[{}]",
                RationalPolynomialPrinter {
                    poly: &p.deserialize(),
                    opts: *opts,
                    add_parentheses: false,
                }
            )),
        }
    }
}

impl<'a> FormattedPrintMul for MulView<'a> {
    fn fmt_debug(&self, f: &mut fmt::Formatter) -> fmt::Result {
        <Self as std::fmt::Debug>::fmt(self, f)
    }

    fn fmt_output(
        &self,
        f: &mut fmt::Formatter,
        opts: &PrintOptions,
        mut print_state: PrintState,
    ) -> fmt::Result {
        // write the coefficient first
        let mut first = true;
        let mut skip_num = false;
        if let Some(AtomView::Num(n)) = self.iter().last() {
            // write -1*x as -x
            if n.get_coeff_view() == CoefficientView::Natural(-1, 1) {
                if print_state.top_level_add_child && opts.color_top_level_sum {
                    f.write_fmt(format_args!("{}", "-".yellow()))?;
                } else {
                    f.write_char('-')?;
                }

                first = true;
            } else {
                n.fmt_output(f, opts, print_state)?;
                first = false;
            }

            skip_num = true;
        } else if print_state.explicit_sign {
            if print_state.top_level_add_child && opts.color_top_level_sum {
                f.write_fmt(format_args!("{}", "+".yellow()))?;
            } else {
                f.write_char('+')?;
            }
        }

        print_state.top_level_add_child = false;
        print_state.level += 1;
        print_state.explicit_sign = false;
        for x in self.iter().take(if skip_num {
            self.get_nargs() - 1
        } else {
            self.get_nargs()
        }) {
            if !first {
                if opts.latex {
                    f.write_char(' ')?;
                } else {
                    f.write_char(opts.multiplication_operator)?;
                }
            }
            first = false;

            if let AtomView::Add(_) = x {
                if opts.latex {
                    f.write_str("\\left(")?;
                } else {
                    f.write_char('(')?;
                }
                x.fmt_output(f, opts, print_state)?;
                if opts.latex {
                    f.write_str("\\right)")?;
                } else {
                    f.write_char(')')?;
                }
            } else {
                x.fmt_output(f, opts, print_state)?;
            }
        }
        Ok(())
    }
}

impl<'a> FormattedPrintFn for FunView<'a> {
    fn fmt_output(
        &self,
        f: &mut fmt::Formatter,
        opts: &PrintOptions,
        mut print_state: PrintState,
    ) -> fmt::Result {
        if print_state.explicit_sign {
            if print_state.top_level_add_child && opts.color_top_level_sum {
                f.write_fmt(format_args!("{}", "+".yellow()))?;
            } else {
                f.write_char('+')?;
            }
        }

        let id = self.get_symbol();
        let name = State::get_name(id);

        if opts.latex {
            if name == "cos" || name == "sin" || name == "exp" || name == "log" {
                f.write_fmt(format_args!("\\{}\\!\\left(", name))?;
            } else {
                f.write_fmt(format_args!("{}\\!\\left(", name))?;
            }
        } else {
            if opts.color_builtin_symbols && name.ends_with('_') {
                f.write_fmt(format_args!("{}", name.cyan().italic()))?;
            } else if opts.color_builtin_symbols && State::is_builtin(id) {
                f.write_fmt(format_args!("{}", name.purple()))?;
            } else {
                f.write_str(name)?;
            }

            if opts.square_brackets_for_function {
                f.write_char('[')?;
            } else {
                f.write_char('(')?;
            }
        }

        print_state.top_level_add_child = false;
        print_state.level += 1;
        print_state.explicit_sign = false;
        let mut first = true;
        for x in self.iter() {
            if !first {
                f.write_char(',')?;
            }
            first = false;

            x.fmt_output(f, opts, print_state)?;
        }

        if opts.latex {
            f.write_str("\\right)")
        } else if opts.square_brackets_for_function {
            f.write_char(']')
        } else {
            f.write_char(')')
        }
    }

    fn fmt_debug(&self, f: &mut fmt::Formatter) -> fmt::Result {
        <Self as std::fmt::Debug>::fmt(self, f)
    }
}

impl<'a> FormattedPrintPow for PowView<'a> {
    fn fmt_output(
        &self,
        f: &mut fmt::Formatter,
        opts: &PrintOptions,
        mut print_state: PrintState,
    ) -> fmt::Result {
        if print_state.explicit_sign {
            if print_state.top_level_add_child && opts.color_top_level_sum {
                f.write_fmt(format_args!("{}", "+".yellow()))?;
            } else {
                f.write_char('+')?;
            }
        }

        let b = self.get_base();
        let e = self.get_exp();

        print_state.top_level_add_child = false;
        print_state.level += 1;
        print_state.explicit_sign = false;

        let mut superscript_exponent = false;
        if opts.latex {
            if let AtomView::Num(n) = e {
                if n.get_coeff_view() == CoefficientView::Natural(-1, 1) {
                    // TODO: construct the numerator
                    f.write_str("\\frac{1}{")?;
                    b.fmt_output(f, opts, print_state)?;
                    return f.write_char('}');
                }
            }
        } else if opts.num_exp_as_superscript {
            if let AtomView::Num(n) = e {
                superscript_exponent = n.get_coeff_view().is_integer()
            }
        }

        let base_needs_parentheses =
            matches!(b, AtomView::Add(_) | AtomView::Mul(_) | AtomView::Pow(_))
                || if let AtomView::Num(n) = b {
                    match n.get_coeff_view() {
                        CoefficientView::Natural(n, d) => n < 0 || d != 1,
                        CoefficientView::Float(_) => true, // TODO
                        CoefficientView::Large(r) => r.is_negative() || !r.to_rat().is_integer(),
                        CoefficientView::FiniteField(n, i) => {
                            opts.symmetric_representation_for_finite_field
                                && n.0 * 2 > State::get_finite_field(i).get_prime()
                        }
                        CoefficientView::RationalPolynomial(_) => true,
                    }
                } else {
                    false
                };

        if base_needs_parentheses {
            if opts.latex {
                f.write_str("\\left(")?;
            } else {
                f.write_char('(')?;
            }
            b.fmt_output(f, opts, print_state)?;
            if opts.latex {
                f.write_str("\\right)")?;
            } else {
                f.write_char(')')?;
            }
        } else {
            b.fmt_output(f, opts, print_state)?;
        }

        if !superscript_exponent {
            if !opts.latex && opts.double_star_for_exponentiation {
                f.write_str("**")?;
            } else {
                f.write_char('^')?;
            }
        }

        if opts.latex {
            f.write_char('{')?;
            e.fmt_output(f, opts, print_state)?;
            f.write_char('}')
        } else {
            let exp_needs_parentheses = matches!(e, AtomView::Add(_) | AtomView::Mul(_))
                || if let AtomView::Num(n) = e {
                    !n.get_coeff_view().is_integer()
                } else {
                    false
                };

            if exp_needs_parentheses {
                f.write_char('(')?;
                e.fmt_output(f, opts, print_state)?;
                f.write_char(')')
            } else {
                print_state.superscript = superscript_exponent;
                e.fmt_output(f, opts, print_state)
            }
        }
    }

    fn fmt_debug(&self, f: &mut fmt::Formatter) -> fmt::Result {
        <Self as std::fmt::Debug>::fmt(self, f)
    }
}

impl<'a> FormattedPrintAdd for AddView<'a> {
    fn fmt_output(
        &self,
        f: &mut fmt::Formatter,
        opts: &PrintOptions,
        mut print_state: PrintState,
    ) -> fmt::Result {
        let mut first = true;
        print_state.top_level_add_child = print_state.level == 0;
        print_state.level += 1;

        for x in self.iter() {
            if !first && print_state.top_level_add_child && opts.terms_on_new_line {
                f.write_char('\n')?;
                f.write_char('\t')?;
            }
            print_state.explicit_sign = !first;
            first = false;

            x.fmt_output(f, opts, print_state)?;
        }
        Ok(())
    }

    fn fmt_debug(&self, f: &mut fmt::Formatter) -> fmt::Result {
        <Self as std::fmt::Debug>::fmt(self, f)
    }
}

pub struct FactorizedRationalPolynomialPrinter<'a, R: Ring, E: Exponent> {
    pub poly: &'a FactorizedRationalPolynomial<R, E>,
    pub opts: PrintOptions,
    pub add_parentheses: bool,
}

impl<'a, R: Ring, E: Exponent> FactorizedRationalPolynomialPrinter<'a, R, E> {
    pub fn new(
        poly: &'a FactorizedRationalPolynomial<R, E>,
    ) -> FactorizedRationalPolynomialPrinter<'a, R, E> {
        FactorizedRationalPolynomialPrinter {
            poly,
            opts: PrintOptions::default(),
            add_parentheses: false,
        }
    }

    pub fn new_with_options(
        poly: &'a FactorizedRationalPolynomial<R, E>,
        opts: PrintOptions,
    ) -> FactorizedRationalPolynomialPrinter<'a, R, E> {
        FactorizedRationalPolynomialPrinter {
            poly,
            opts,
            add_parentheses: false,
        }
    }
}

impl<'a, R: Ring, E: Exponent> Display for FactorizedRationalPolynomialPrinter<'a, R, E> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if self.opts.explicit_rational_polynomial {
            if !R::is_zero(&self.poly.numer_coeff)
                && !self.poly.numerator.ring.is_one(&self.poly.numer_coeff)
            {
                f.write_fmt(format_args!(
                    "[{}]*",
                    RingPrinter {
                        ring: &self.poly.numerator.ring,
                        element: &self.poly.numer_coeff,
                        opts: self.opts,
                        in_product: false
                    }
                ))?;
            }

            if self.poly.denominators.is_empty()
                && self.poly.numerator.ring.is_one(&self.poly.denom_coeff)
            {
                if self.poly.numerator.is_zero() {
                    f.write_char('0')?;
                } else {
                    f.write_fmt(format_args!(
                        "[{}]",
                        PolynomialPrinter {
                            poly: &self.poly.numerator,
                            opts: self.opts,
                        }
                    ))?;
                }
            } else {
                f.write_fmt(format_args!(
                    "[{}",
                    PolynomialPrinter {
                        poly: &self.poly.numerator,

                        opts: self.opts,
                    },
                ))?;

                if !self.poly.numerator.ring.is_one(&self.poly.denom_coeff) {
                    f.write_fmt(format_args!(
                        ",{},1",
                        RingPrinter {
                            ring: &self.poly.numerator.ring,
                            element: &self.poly.denom_coeff,
                            opts: self.opts,
                            in_product: false
                        },
                    ))?;
                }

                for (d, p) in &self.poly.denominators {
                    f.write_fmt(format_args!(
                        ",{}",
                        PolynomialPrinter {
                            poly: d,

                            opts: self.opts,
                        }
                    ))?;
                    f.write_fmt(format_args!(",{}", p))?;
                }

                f.write_char(']')?;
            }

            return Ok(());
        }

        if R::is_zero(&self.poly.numer_coeff) {
            return f.write_char('0');
        }

        if self.poly.denominators.is_empty()
            && self.poly.numerator.ring.is_one(&self.poly.denom_coeff)
        {
            if !self.poly.numerator.ring.is_one(&self.poly.numer_coeff) {
                f.write_fmt(format_args!(
                    "{}",
                    RingPrinter {
                        ring: &self.poly.numerator.ring,
                        element: &self.poly.numer_coeff,
                        opts: self.opts,
                        in_product: false
                    }
                ))?;
            }

            if (self.poly.numerator.ring.is_one(&self.poly.numer_coeff) && !self.add_parentheses)
                || self.poly.numerator.nterms() < 2
            {
                if !self.poly.numerator.ring.is_one(&self.poly.numer_coeff) {
                    if self.poly.numerator.is_one() {
                        return Ok(());
                    }

                    f.write_char('*')?;
                }

                f.write_fmt(format_args!(
                    "{}",
                    PolynomialPrinter {
                        poly: &self.poly.numerator,

                        opts: self.opts,
                    }
                ))
            } else {
                if !self.poly.numerator.ring.is_one(&self.poly.numer_coeff) {
                    if self.poly.numerator.is_one() {
                        return Ok(());
                    }

                    f.write_char('*')?;
                }

                f.write_fmt(format_args!(
                    "({})",
                    PolynomialPrinter {
                        poly: &self.poly.numerator,

                        opts: self.opts,
                    }
                ))
            }
        } else {
            if self.opts.latex {
                if !self.poly.numerator.ring.is_one(&self.poly.numer_coeff) {
                    f.write_fmt(format_args!(
                        "{} ",
                        RingPrinter {
                            ring: &self.poly.numerator.ring,
                            element: &self.poly.numer_coeff,
                            opts: self.opts,
                            in_product: false
                        }
                    ))?;
                }

                f.write_fmt(format_args!(
                    "\\frac{{{}}}{{",
                    PolynomialPrinter {
                        poly: &self.poly.numerator,

                        opts: self.opts,
                    },
                ))?;

                if !self.poly.numerator.ring.is_one(&self.poly.denom_coeff) {
                    f.write_fmt(format_args!(
                        "{}",
                        RingPrinter {
                            ring: &self.poly.numerator.ring,
                            element: &self.poly.denom_coeff,
                            opts: self.opts,
                            in_product: false
                        }
                    ))?;
                }

                for (d, p) in &self.poly.denominators {
                    if *p == 1 {
                        f.write_fmt(format_args!(
                            "({})",
                            PolynomialPrinter {
                                poly: d,
                                opts: self.opts,
                            }
                        ))?;
                    } else {
                        f.write_fmt(format_args!(
                            "({})^{}",
                            PolynomialPrinter {
                                poly: d,
                                opts: self.opts,
                            },
                            p
                        ))?;
                    }
                }

                return f.write_str("}}");
            }

            if !self.poly.numerator.ring.is_one(&self.poly.numer_coeff) {
                f.write_fmt(format_args!(
                    "{}*",
                    RingPrinter {
                        ring: &self.poly.numerator.ring,
                        element: &self.poly.numer_coeff,
                        opts: self.opts,
                        in_product: false
                    }
                ))?;
            }

            if self.poly.numerator.nterms() < 2 {
                f.write_fmt(format_args!(
                    "{}",
                    PolynomialPrinter {
                        poly: &self.poly.numerator,
                        opts: self.opts,
                    }
                ))?;
            } else {
                f.write_fmt(format_args!(
                    "({})",
                    PolynomialPrinter {
                        poly: &self.poly.numerator,
                        opts: self.opts,
                    }
                ))?;
            }

            f.write_char('/')?;

            if self.poly.denominators.is_empty() {
                return f.write_fmt(format_args!(
                    "{}",
                    RingPrinter {
                        ring: &self.poly.numerator.ring,
                        element: &self.poly.denom_coeff,
                        opts: self.opts,
                        in_product: true
                    }
                ));
            }

            if self.poly.numerator.ring.is_one(&self.poly.denom_coeff)
                && self.poly.denominators.len() == 1
                && self.poly.denominators[0].0.nterms() == 1
                && self.poly.denominators[0].1 == 1
            {
                let (d, _) = &self.poly.denominators[0];
                let var_count = d.exponents.iter().filter(|x| !x.is_zero()).count();

                if var_count == 0 || d.ring.is_one(&d.coefficients[0]) && var_count == 1 {
                    return f.write_fmt(format_args!(
                        "{}",
                        PolynomialPrinter {
                            poly: d,
                            opts: self.opts,
                        }
                    ));
                }
            }

            f.write_char('(')?; // TODO: add special cases for 1 argument

            if !self.poly.numerator.ring.is_one(&self.poly.denom_coeff) {
                f.write_fmt(format_args!(
                    "{}",
                    RingPrinter {
                        ring: &self.poly.numerator.ring,
                        element: &self.poly.denom_coeff,
                        opts: self.opts,
                        in_product: true
                    }
                ))?;
            }

            for (d, p) in &self.poly.denominators {
                if *p == 1 {
                    f.write_fmt(format_args!(
                        "({})",
                        PolynomialPrinter {
                            poly: d,
                            opts: self.opts,
                        }
                    ))?;
                } else {
                    f.write_fmt(format_args!(
                        "({}){}{}",
                        PolynomialPrinter {
                            poly: d,
                            opts: self.opts,
                        },
                        if self.opts.double_star_for_exponentiation {
                            "**"
                        } else {
                            "^"
                        },
                        p
                    ))?;
                }
            }

            f.write_char(')')
        }
    }
}

pub struct RationalPolynomialPrinter<'a, R: Ring, E: Exponent> {
    pub poly: &'a RationalPolynomial<R, E>,
    pub opts: PrintOptions,
    pub add_parentheses: bool,
}

impl<'a, R: Ring, E: Exponent> RationalPolynomialPrinter<'a, R, E> {
    pub fn new(poly: &'a RationalPolynomial<R, E>) -> RationalPolynomialPrinter<'a, R, E> {
        RationalPolynomialPrinter {
            poly,
            opts: PrintOptions::default(),
            add_parentheses: false,
        }
    }

    pub fn new_with_options(
        poly: &'a RationalPolynomial<R, E>,
        opts: PrintOptions,
    ) -> RationalPolynomialPrinter<'a, R, E> {
        RationalPolynomialPrinter {
            poly,
            opts,
            add_parentheses: false,
        }
    }
}

impl<'a, R: Ring, E: Exponent> Display for RationalPolynomialPrinter<'a, R, E> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if self.opts.explicit_rational_polynomial {
            if self.poly.denominator.is_one() {
                if self.poly.numerator.is_zero() {
                    f.write_char('0')?;
                } else {
                    f.write_fmt(format_args!(
                        "[{}]",
                        PolynomialPrinter {
                            poly: &self.poly.numerator,
                            opts: self.opts,
                        }
                    ))?;
                }
            } else {
                f.write_fmt(format_args!(
                    "[{},{}]",
                    PolynomialPrinter {
                        poly: &self.poly.numerator,
                        opts: self.opts,
                    },
                    PolynomialPrinter {
                        poly: &self.poly.denominator,
                        opts: self.opts,
                    }
                ))?;
            }

            return Ok(());
        }

        if self.poly.denominator.is_one() {
            if !self.add_parentheses || self.poly.numerator.nterms() < 2 {
                f.write_fmt(format_args!(
                    "{}",
                    PolynomialPrinter {
                        poly: &self.poly.numerator,
                        opts: self.opts,
                    }
                ))
            } else {
                f.write_fmt(format_args!(
                    "({})",
                    PolynomialPrinter {
                        poly: &self.poly.numerator,
                        opts: self.opts,
                    }
                ))
            }
        } else {
            if self.opts.latex {
                return f.write_fmt(format_args!(
                    "\\frac{{{}}}{{{}}}",
                    PolynomialPrinter {
                        poly: &self.poly.numerator,
                        opts: self.opts,
                    },
                    PolynomialPrinter {
                        poly: &self.poly.denominator,
                        opts: self.opts,
                    }
                ));
            }

            if self.poly.numerator.nterms() < 2 {
                f.write_fmt(format_args!(
                    "{}",
                    PolynomialPrinter {
                        poly: &self.poly.numerator,
                        opts: self.opts,
                    }
                ))?;
            } else {
                f.write_fmt(format_args!(
                    "({})",
                    PolynomialPrinter {
                        poly: &self.poly.numerator,
                        opts: self.opts,
                    }
                ))?;
            }

            if self.poly.denominator.nterms() == 1 {
                let var_count = self
                    .poly
                    .denominator
                    .exponents
                    .iter()
                    .filter(|x| !x.is_zero())
                    .count();

                if var_count == 0
                    || self
                        .poly
                        .denominator
                        .ring
                        .is_one(&self.poly.denominator.coefficients[0])
                        && var_count == 1
                {
                    return f.write_fmt(format_args!(
                        "/{}",
                        PolynomialPrinter {
                            poly: &self.poly.denominator,
                            opts: self.opts,
                        }
                    ));
                }
            }

            f.write_fmt(format_args!(
                "/({})",
                PolynomialPrinter {
                    poly: &self.poly.denominator,

                    opts: self.opts,
                }
            ))
        }
    }
}
pub struct PolynomialPrinter<'a, F: Ring + Display, E: Exponent, O: MonomialOrder> {
    pub poly: &'a MultivariatePolynomial<F, E, O>,
    pub opts: PrintOptions,
}

impl<'a, R: Ring + Display, E: Exponent, O: MonomialOrder> PolynomialPrinter<'a, R, E, O> {
    pub fn new(poly: &'a MultivariatePolynomial<R, E, O>) -> PolynomialPrinter<'a, R, E, O> {
        PolynomialPrinter {
            poly,
            opts: PrintOptions::default(),
        }
    }

    pub fn new_with_options(
        poly: &'a MultivariatePolynomial<R, E, O>,
        opts: PrintOptions,
    ) -> PolynomialPrinter<'a, R, E, O> {
        PolynomialPrinter { poly, opts }
    }
}

impl<'a, F: Ring + Display, E: Exponent, O: MonomialOrder> Display
    for PolynomialPrinter<'a, F, E, O>
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if f.sign_plus() {
            f.write_char('+')?;
        }

        let var_map: Vec<String> = self
            .poly
            .variables
            .as_ref()
            .iter()
            .map(|v| v.to_string())
            .collect();

        let mut is_first_term = true;
        for monomial in self.poly {
            let mut is_first_factor = true;
            if self.poly.ring.is_one(monomial.coefficient) {
                if !is_first_term {
                    write!(f, "+")?;
                }
            } else if monomial
                .coefficient
                .eq(&self.poly.ring.neg(&self.poly.ring.one()))
            {
                write!(f, "-")?;
            } else {
                if is_first_term {
                    self.poly
                        .ring
                        .fmt_display(monomial.coefficient, &self.opts, true, f)?;
                } else {
                    write!(
                        f,
                        "{:+}",
                        RingPrinter {
                            ring: &self.poly.ring,
                            element: monomial.coefficient,
                            opts: self.opts,
                            in_product: true
                        }
                    )?;
                }
                is_first_factor = false;
            }
            is_first_term = false;
            for (var_id, e) in var_map.iter().zip(monomial.exponents) {
                if e.is_zero() {
                    continue;
                }
                if is_first_factor {
                    is_first_factor = false;
                } else if !self.opts.latex {
                    write!(f, "*")?;
                }

                f.write_str(var_id)?;

                if e.to_u32() != 1 {
                    if self.opts.latex {
                        write!(f, "^{{{}}}", e)?;
                    } else if self.opts.double_star_for_exponentiation {
                        write!(f, "**{}", e)?;
                    } else {
                        write!(f, "^{}", e)?;
                    }
                }
            }
            if is_first_factor {
                write!(f, "1")?;
            }
        }
        if is_first_term {
            write!(f, "0")?;
        }

        if self.opts.print_finite_field {
            Display::fmt(&self.poly.ring, f)?;
        }

        Ok(())
    }
}

pub struct MatrixPrinter<'a, F: Ring + Display> {
    pub matrix: &'a Matrix<F>,
    pub opts: PrintOptions,
}

impl<'a, F: Ring + Display> MatrixPrinter<'a, F> {
    pub fn new(matrix: &'a Matrix<F>) -> MatrixPrinter<'a, F> {
        MatrixPrinter {
            matrix,

            opts: PrintOptions::default(),
        }
    }

    pub fn new_with_options(matrix: &'a Matrix<F>, opts: PrintOptions) -> MatrixPrinter<'a, F> {
        MatrixPrinter { matrix, opts }
    }
}

impl<'a, F: Ring + Display> Display for MatrixPrinter<'a, F> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if self.opts.latex {
            f.write_str("\\begin{pmatrix}")?;

            for (ri, r) in self.matrix.row_iter().enumerate() {
                for (ci, c) in r.iter().enumerate() {
                    f.write_fmt(format_args!(
                        "{}",
                        RingPrinter {
                            ring: &self.matrix.field,
                            element: c,
                            opts: self.opts,
                            in_product: false,
                        },
                    ))?;

                    if ci + 1 < self.matrix.ncols as usize {
                        f.write_str(" & ")?;
                    }
                }
                if ri + 1 < self.matrix.nrows as usize {
                    f.write_str(r" \\ ")?;
                }
            }

            f.write_str("\\end{pmatrix}")
        } else {
            f.write_char('{')?;
            for (ri, r) in self.matrix.row_iter().enumerate() {
                f.write_char('{')?;
                for (ci, c) in r.iter().enumerate() {
                    f.write_fmt(format_args!(
                        "{}",
                        RingPrinter {
                            ring: &self.matrix.field,
                            element: c,
                            opts: self.opts,
                            in_product: false,
                        },
                    ))?;

                    if ci + 1 < self.matrix.ncols as usize {
                        f.write_char(',')?;
                    }
                }
                f.write_char('}')?;
                if ri + 1 < self.matrix.nrows as usize {
                    f.write_char(',')?;
                }
            }
            f.write_char('}')
        }
    }
}

#[cfg(test)]
mod test {
    use colored::control::ShouldColorize;

    use crate::{
        atom::Atom,
        domains::{finite_field::Zp, integer::Z},
        printer::{AtomPrinter, PolynomialPrinter, PrintOptions},
        state::{FunctionAttribute, State},
    };

    #[test]
    fn atoms() {
        let a = Atom::parse("f(x,y^2)^(x+z)/5+3").unwrap();

        if ShouldColorize::from_env().should_colorize() {
            assert_eq!(format!("{}", a), "1/5*f(x,y^2)^(x+z)\u{1b}[33m+\u{1b}[0m3");
        } else {
            assert_eq!(format!("{}", a), "1/5*f(x,y^2)^(x+z)+3");
        }

        assert_eq!(
            format!(
                "{}",
                AtomPrinter::new_with_options(a.as_view(), PrintOptions::latex())
            ),
            "\\frac{1}{5} f\\!\\left(x,y^{2}\\right)^{x+z}+3"
        );

        assert_eq!(
            format!(
                "{}",
                AtomPrinter::new_with_options(a.as_view(), PrintOptions::mathematica())
            ),
            "1/5 f[x,y^2]^(x+z)+3"
        );

        let a = Atom::parse("8127389217 x^2").unwrap();
        assert_eq!(
            format!(
                "{}",
                AtomPrinter::new_with_options(
                    a.as_view(),
                    PrintOptions {
                        number_thousands_separator: Some('_'),
                        multiplication_operator: ' ',
                        num_exp_as_superscript: true,
                        ..PrintOptions::file()
                    }
                )
            ),
            "812_738_921_7 x²"
        );
    }

    #[test]
    fn polynomials() {
        let a = Atom::parse("15 x^2")
            .unwrap()
            .to_polynomial::<_, u8>(&Zp::new(17), None);
        assert_eq!(
            format!(
                "{}",
                PolynomialPrinter::new_with_options(
                    &a,
                    PrintOptions {
                        print_finite_field: true,
                        symmetric_representation_for_finite_field: true,
                        ..PrintOptions::file()
                    }
                )
            ),
            "-2*x^2 % 17"
        );
    }

    #[test]
    fn rational_polynomials() {
        let a = Atom::parse("15 x^2 / (1+x)")
            .unwrap()
            .to_rational_polynomial::<_, _, u8>(&Z, &Z, None);
        assert_eq!(format!("{}", a), "15*x^2/(1+x)");

        let a = Atom::parse("(15 x^2 + 6) / (1+x)")
            .unwrap()
            .to_rational_polynomial::<_, _, u8>(&Z, &Z, None);
        assert_eq!(format!("{}", a), "(6+15*x^2)/(1+x)");
    }

    #[test]
    fn factorized_rational_polynomials() {
        let a = Atom::parse("15 x^2 / ((1+x)(x+2))")
            .unwrap()
            .to_factorized_rational_polynomial::<_, _, u8>(&Z, &Z, None);
        assert!(
            format!("{}", a) == "15*x^2/((1+x)(2+x))" || format!("{}", a) == "15*x^2/((2+x)(1+x))"
        );

        let a = Atom::parse("(15 x^2 + 6) / ((1+x)(x+2))")
            .unwrap()
            .to_factorized_rational_polynomial::<_, _, u8>(&Z, &Z, None);
        assert!(
            format!("{}", a) == "3*(2+5*x^2)/((1+x)(2+x))"
                || format!("{}", a) == "3*(2+5*x^2)/((2+x)(1+x))"
        );
    }

    #[test]
    fn base_parentheses() {
        let a = Atom::parse("(-1)^(x+1)-(1/2)^x").unwrap();
        assert_eq!(
            format!(
                "{}",
                AtomPrinter::new_with_options(a.as_view(), PrintOptions::file())
            ),
            "(-1)^(x+1)-(1/2)^x"
        )
    }

    #[test]
    fn canon() {
        let _ =
            State::get_symbol_with_attributes("canon_f", &[FunctionAttribute::Symmetric]).unwrap();
        let _ = State::get_symbol("canon_y");
        let _ = State::get_symbol("canon_x");

        let a = Atom::parse("canon_x^2 + 2*canon_x*canon_y + canon_y^2*(canon_x+canon_y) + canon_f(canon_x,canon_y)").unwrap();
        assert_eq!(
            a.to_canonical_string(),
            "(canon_x+canon_y)*canon_y^2+2*canon_x*canon_y+canon_f(canon_x,canon_y)+canon_x^2"
        );
    }
}