use std::{
    cmp::Reverse,
    hash::{Hash, Hasher},
};

use ahash::{AHasher, HashMap, HashSet, HashSetExt};
use rand::{thread_rng, Rng};

use crate::rings::{
    float::NumericalFloatLike,
    rational::{Rational, RationalField},
    EuclideanDomain,
};
use crate::{rings::Ring, state::State};

use super::{polynomial::MultivariatePolynomial, Exponent};

/// A borrowed version of a Horner node, suitable as a key in a
/// hashmap. It uses precomputed hashes for the complete node
/// `var^pow*content+rest` and for its children `var^pow*content` and `var^pow`.
#[derive(Debug, Clone, Copy)]
pub struct BorrowedHornerNode<'a, R: Ring>
where
    R::Element: Hash + Eq,
{
    pub var: usize,
    pub pow: usize,
    pub gcd: R::Element, // only used for counting number of operations
    pub content: Option<&'a HornerScheme<R>>,
    pub rest: Option<&'a HornerScheme<R>>,
    pub hash: (u64, u64, u64),
}

impl<'a, R: Ring> PartialEq for BorrowedHornerNode<'a, R>
where
    R::Element: Hash + Eq,
{
    fn eq(&self, other: &Self) -> bool {
        // hash and gcd is skipped
        self.var == other.var
            && self.pow == other.pow
            && self.content == other.content
            && self.rest == other.rest
    }
}

impl<'a, R: Ring> Eq for BorrowedHornerNode<'a, R> where R::Element: Hash + Eq {}

impl<'a> Hash for BorrowedHornerNode<'a, RationalField> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        let hash = if self.content.is_some() {
            if self.rest.is_some() {
                self.hash.2
            } else {
                self.hash.1
            }
        } else {
            self.hash.0
        };

        state.write_u64(hash);
    }
}

impl<'a, R: Ring> From<&'a HornerNode<R>> for BorrowedHornerNode<'a, R>
where
    R::Element: Hash + Eq,
{
    fn from(n: &'a HornerNode<R>) -> Self {
        BorrowedHornerNode {
            var: n.var,
            pow: n.pow,
            gcd: n.gcd.clone(),
            content: n.content_rest.0.as_ref(),
            rest: n.content_rest.1.as_ref(),
            hash: n.hash,
        }
    }
}

#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub enum BorrowedHornerScheme<'a, R: Ring>
where
    R::Element: Hash + Eq,
{
    Node(BorrowedHornerNode<'a, R>),
    Leaf(&'a R::Element),
}

impl<'a, R: Ring> From<&'a HornerScheme<R>> for BorrowedHornerScheme<'a, R>
where
    R::Element: Hash + Eq,
{
    fn from(value: &'a HornerScheme<R>) -> Self {
        match value {
            HornerScheme::Node(n) => BorrowedHornerScheme::Node(BorrowedHornerNode::from(n)),
            HornerScheme::Leaf(_, l) => BorrowedHornerScheme::Leaf(l),
        }
    }
}

impl<'a> BorrowedHornerScheme<'a, RationalField> {
    /// Determine the number of operations required to evaluate the Horner scheme.
    /// Common subexpressions are only counted once.
    pub fn op_count_cse(&self) -> usize {
        let mut h = HashSet::default();
        self.op_count_cse_impl(&mut h)
    }

    fn op_count_cse_impl(&self, set: &mut HashSet<BorrowedHornerNode<'a, RationalField>>) -> usize {
        match self {
            BorrowedHornerScheme::Node(n) => {
                let gcd_op = if n.gcd.abs() != Rational::one() { 1 } else { 0 };

                // check if n = var^pow*a+b is seen before
                if set.contains(n) {
                    return gcd_op;
                }

                // check if var^pow*a is seen before
                let mut b = n.clone();
                b.rest = None;

                if set.contains(&b) {
                    set.insert(n.clone());
                    gcd_op
                        + 1
                        + n.rest
                            .map(|x| BorrowedHornerScheme::from(x).op_count_cse_impl(set))
                            .unwrap() // can fail now because of gcd?
                } else {
                    // add var^pow to seen list
                    let instr = if n.pow > 1 {
                        let mut c = b.clone();
                        c.content = None;

                        if set.contains(&c) {
                            0
                        } else {
                            set.insert(c.clone());
                            n.pow - 1
                        }
                    } else {
                        0
                    };

                    set.insert(b.clone());
                    set.insert(n.clone());

                    gcd_op
                        + instr
                        + n.content
                            .map(|x| 1 + BorrowedHornerScheme::from(x).op_count_cse_impl(set))
                            .unwrap_or(0)
                        + n.rest
                            .map(|x| 1 + BorrowedHornerScheme::from(x).op_count_cse_impl(set))
                            .unwrap_or(0)
                }
            }
            BorrowedHornerScheme::Leaf(_) => 0,
        }
    }
}

#[derive(Debug, Clone)]
pub struct HornerNode<R: Ring>
where
    R::Element: Hash + Eq,
{
    pub var: usize,
    pub pow: usize,
    pub gcd: R::Element,
    pub content_rest: Box<(Option<HornerScheme<R>>, Option<HornerScheme<R>>)>,
    pub hash: (u64, u64, u64),
}

impl<R: Ring> PartialEq for HornerNode<R>
where
    R::Element: Hash + Eq,
{
    fn eq(&self, other: &Self) -> bool {
        // hash is skipped, but the gcd is included
        self.var == other.var
            && self.pow == other.pow
            && self.gcd == other.gcd
            && self.content_rest == other.content_rest
    }
}

/// A polynomial written in a Horner scheme, i.e.:
/// `x^2*(x*(y*(z + 1) + y) + 5) + z^2`.
#[derive(Debug, Clone)]
pub enum HornerScheme<R: Ring>
where
    R::Element: Hash + Eq,
{
    Node(HornerNode<R>),
    Leaf(u64, R::Element), // hash and number
}

impl<R: Ring> PartialEq for HornerScheme<R>
where
    R::Element: Hash + Eq,
{
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (Self::Node(l0), Self::Node(r0)) => l0 == r0,
            (Self::Leaf(_, l1), Self::Leaf(_, r1)) => l1 == r1,
            _ => false,
        }
    }
}

impl HornerScheme<RationalField> {
    /// Evaluate a polynomial written in a Horner scheme. For faster
    /// evaluation, convert the Horner scheme into an `InstructionList`.
    pub fn evaluate(&self, samples: &[Rational]) -> Rational {
        let field = RationalField::new();
        match self {
            HornerScheme::Node(n) => {
                let e = match &n.content_rest.0 {
                    Some(s) => match &n.content_rest.1 {
                        Some(s1) => field.add(
                            &field.mul(
                                &field.pow(&samples[n.var], n.pow as u64),
                                &s.evaluate(samples),
                            ),
                            &s1.evaluate(samples),
                        ),
                        None => field.mul(
                            &field.pow(&samples[n.var], n.pow as u64),
                            &s.evaluate(samples),
                        ),
                    },
                    None => match &n.content_rest.1 {
                        Some(s1) => field.add(
                            &field.pow(&samples[n.var], n.pow as u64),
                            &s1.evaluate(samples),
                        ),
                        None => field.pow(&samples[n.var], n.pow as u64),
                    },
                };
                &e * &n.gcd
            }
            HornerScheme::Leaf(_, l) => l.clone(),
        }
    }

    /// Clean up the Horner scheme and save the allocations in `boxes`.
    fn cleanup(
        &mut self,
        boxes: &mut Vec<
            Box<(
                Option<HornerScheme<RationalField>>,
                Option<HornerScheme<RationalField>>,
            )>,
        >,
    ) {
        let private = std::mem::replace(self, HornerScheme::Leaf(0, Rational::zero()));
        match private {
            HornerScheme::Node(mut n) => {
                if let Some(c) = &mut n.content_rest.0 {
                    c.cleanup(boxes);
                }
                if let Some(c) = &mut n.content_rest.1 {
                    c.cleanup(boxes);
                }

                *n.content_rest = (None, None);

                boxes.push(n.content_rest);
            }
            HornerScheme::Leaf(_, _) => {}
        }
    }
}

impl<R: Ring> std::fmt::Display for HornerScheme<R>
where
    R::Element: Hash + Eq,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            HornerScheme::Leaf(_, l) => f.write_fmt(format_args!("{:?}", l)),
            HornerScheme::Node(n) => {
                f.write_fmt(format_args!("+{:?}*(", n.gcd))?;

                if n.pow == 1 {
                    f.write_fmt(format_args!("x{}", n.var))?;
                } else {
                    f.write_fmt(format_args!("x{}^{}", n.var, n.pow))?;
                }
                if let Some(s) = &n.content_rest.0 {
                    f.write_fmt(format_args!("*("))?;
                    s.fmt(f)?;
                    f.write_fmt(format_args!(")"))?;
                }

                if let Some(s) = &n.content_rest.1 {
                    if let HornerScheme::Leaf(_, _) = &s {
                        f.write_str("+")?;
                        s.fmt(f)?;
                    } else {
                        s.fmt(f)?;
                    }
                }

                f.write_str(")")?;

                Ok(())
            }
        }
    }
}

impl<E: Exponent> MultivariatePolynomial<RationalField, E> {
    /// Write the polynomial in a Horner scheme with the variable ordering
    /// defined in `order`.
    pub fn to_horner_scheme(&self, order: &[usize]) -> HornerScheme<RationalField> {
        let mut indices: Vec<_> = (0..self.nterms).collect();
        let mut power_sub = vec![E::zero(); self.nvars];
        let mut horner_boxes = vec![];

        self.to_horner_scheme_impl(order, &mut indices, 0, &mut power_sub, &mut horner_boxes)
    }

    /// Create a Horner scheme using the variable order specified in `order`,
    /// using the terms from `self` indexed by `indices[index_start..]`.
    /// The exponents of each term should be reduced by `power_sub`
    /// and allocations from previous Horner scheme constructions can
    /// be provided using `boxes`.
    ///
    /// This function is optimized, as it is called many times during
    /// the Horner scheme optimization.
    fn to_horner_scheme_impl(
        &self,
        order: &[usize],
        indices: &mut Vec<usize>,
        index_start: usize,
        power_sub: &mut [E],
        boxes: &mut Vec<
            Box<(
                Option<HornerScheme<RationalField>>,
                Option<HornerScheme<RationalField>>,
            )>,
        >,
    ) -> HornerScheme<RationalField> {
        if order.is_empty() {
            debug_assert!(indices.len() <= index_start + 1);

            let num = if indices.len() == index_start {
                self.field.zero()
            } else {
                self.coefficients[indices[index_start]].clone()
            };

            let mut h = AHasher::default();
            h.write_u8(1);
            num.hash(&mut h);
            return HornerScheme::Leaf(h.finish(), num);
        }

        // find the lowest non-zero power of the next variable in the ordering
        let var = order[0];
        let mut min_pow = E::zero();

        for i in &indices[index_start..] {
            let e = self.exponents(*i)[var];
            if e > power_sub[var] {
                let p = e - power_sub[var];
                if min_pow == E::zero() || p < min_pow {
                    min_pow = p;
                }
            }
        }

        if min_pow == E::zero() {
            return self.to_horner_scheme_impl(&order[1..], indices, index_start, power_sub, boxes);
        }

        let new_index_start = indices.len();

        // create the branch for terms that do not contain var^min_pow
        for t in index_start..new_index_start {
            if self.exponents(indices[t])[var] <= power_sub[var] {
                indices.push(indices[t]);
            }
        }

        let mut rest =
            self.to_horner_scheme_impl(&order[1..], indices, new_index_start, power_sub, boxes);

        // create the branch for terms that do contain var^min_pow and lower the power
        indices.truncate(new_index_start);
        for t in index_start..new_index_start {
            if self.exponents(indices[t])[var] > power_sub[var] {
                indices.push(indices[t]);
            }
        }

        power_sub[var] += min_pow;

        let mut content =
            self.to_horner_scheme_impl(order, indices, new_index_start, power_sub, boxes);

        power_sub[var] = power_sub[var] - min_pow;

        indices.truncate(new_index_start);

        // compute the gcd of both branches
        // normalize such that the first branch is positive
        let mut gcd = match &content {
            HornerScheme::Node(n) => n.gcd.clone(),
            HornerScheme::Leaf(_, l) => l.clone(),
        };

        let gcd_norm = match &content {
            HornerScheme::Node(n) => {
                if n.gcd.is_negative() {
                    Rational::Natural(-1, 1)
                } else {
                    Rational::Natural(1, 1)
                }
            }
            HornerScheme::Leaf(_, l) => {
                if l.is_negative() {
                    Rational::Natural(-1, 1)
                } else {
                    Rational::Natural(1, 1)
                }
            }
        };

        gcd = RationalField::new().gcd(
            &gcd,
            match &rest {
                HornerScheme::Node(n) => &n.gcd,
                HornerScheme::Leaf(_, l) => l,
            },
        );

        gcd *= &gcd_norm;

        if !gcd.is_one() {
            for s in [&mut content, &mut rest] {
                match s {
                    HornerScheme::Node(n) => n.gcd = &n.gcd / &gcd,
                    HornerScheme::Leaf(n, l) => {
                        *l = &*l / &gcd;

                        // overwrite the hash of the number
                        let mut h = AHasher::default();
                        h.write_u8(1);
                        l.hash(&mut h);
                        *n = h.finish();
                    }
                };
            }
        }

        // compute the hash of the node and its components
        let mut h = AHasher::default();
        h.write_u8(0);
        var.hash(&mut h);
        (min_pow.to_u32() as usize).hash(&mut h);

        let pow_hash = h.finish(); // hash var^pow

        match &content {
            HornerScheme::Node(n) => h.write_u64(n.hash.2),
            HornerScheme::Leaf(n, _) => h.write_u64(*n),
        }

        let pow_content_hash = h.finish(); // hash var^pow*content

        match &rest {
            HornerScheme::Node(n) => h.write_u64(n.hash.2),
            HornerScheme::Leaf(n, _) => h.write_u64(*n),
        }

        let full_hash = h.finish(); // hash var^pow*content+rest

        let children = (
            if let HornerScheme::Leaf(_, Rational::Natural(1, 1)) = content {
                None
            } else {
                Some(content)
            },
            if let HornerScheme::Leaf(_, Rational::Natural(0, 1)) = rest {
                None
            } else {
                Some(rest)
            },
        );

        // recycle the box if it is available
        let boxed_children = if let Some(mut b) = boxes.pop() {
            *b = children;
            b
        } else {
            Box::new(children)
        };

        HornerScheme::Node(HornerNode {
            var,
            pow: min_pow.to_u32() as usize,
            gcd,
            hash: (pow_hash, pow_content_hash, full_hash),
            content_rest: boxed_children,
        })
    }

    /// Find the Horner scheme that minimizes the number of operations,
    /// taking common subexpressions into account.
    pub fn optimize_horner_scheme(
        &self,
        num_tries: usize,
    ) -> (HornerScheme<RationalField>, usize, Vec<usize>) {
        let (mut hs, op_count, scheme) =
            HornerScheme::optimize_multiple(std::slice::from_ref(&self), num_tries);
        (hs.pop().unwrap(), op_count, scheme)
    }

    /// Optimize an expression for evaluation, given `num_iter` tries.
    pub fn optimize(&self, num_iter: usize) -> InstructionListOutput<Rational> {
        let (h, _ops, _scheme) = self.optimize_horner_scheme(num_iter);
        let mut i = h.to_instr(self.nvars);
        i.fuse_operations();
        while i.common_pair_elimination() {
            i.fuse_operations();
        }

        i.to_output(self.var_map.as_ref().unwrap().to_vec(), true)
    }
}

impl HornerScheme<RationalField> {
    pub fn optimize_multiple<E: Exponent>(
        polys: &[&MultivariatePolynomial<RationalField, E>],
        num_tries: usize,
    ) -> (Vec<HornerScheme<RationalField>>, usize, Vec<usize>) {
        if polys.is_empty() {
            return (vec![], 0, vec![]);
        }

        assert!(
            polys
                .windows(2)
                .all(|r| r[0].var_map == r[1].var_map && r[0].nvars == r[1].nvars),
            "Variable maps of all polynomials must be the same"
        );

        // the starting scheme is the descending order of occurrence of variables
        let mut occurrence: Vec<_> = (0..polys[0].nvars).map(|x| (x, 0)).collect();
        for es in polys[0].exponents.chunks(polys[0].nvars) {
            for ((_, o), e) in occurrence.iter_mut().zip(es) {
                if *e > E::zero() {
                    *o += 1;
                }
            }
        }
        occurrence.sort_by_key(|e| Reverse(e.1));

        let mut scheme: Vec<_> = occurrence.into_iter().map(|(v, _)| v).collect();

        let mut indices: Vec<_> = vec![];
        let mut power_sub = vec![E::zero(); polys[0].nvars];

        let mut horner_boxes = vec![];

        let mut best = Vec::with_capacity(polys.len());
        let mut best_score = 0;
        for x in polys {
            indices.clear();
            indices.extend(0..x.nterms);

            let h = x.to_horner_scheme_impl(
                &scheme,
                &mut indices,
                0,
                &mut power_sub,
                &mut horner_boxes,
            );
            best_score += BorrowedHornerScheme::from(&h).op_count_cse();
            best.push(h);
        }

        let mut best_scheme = scheme.clone();
        let mut rng = thread_rng();

        let mut new_best = Vec::with_capacity(polys.len());

        // TODO: for few variables, test all permutations
        for i in 0..num_tries {
            let a = rng.gen_range(0..polys[0].nvars);
            let b = rng.gen_range(0..polys[0].nvars);
            scheme.swap(a, b);

            let mut new_oc = 0;

            // use the same hash set for all polynomials
            let mut hash_set = HashSet::with_capacity(best_score * 2);
            for x in polys {
                indices.clear();
                indices.extend(0..x.nterms);

                let h = x.to_horner_scheme_impl(
                    &scheme,
                    &mut indices,
                    0,
                    &mut power_sub,
                    &mut horner_boxes,
                );
                new_best.push(h);
            }

            for x in &new_best {
                new_oc += BorrowedHornerScheme::from(x).op_count_cse_impl(&mut hash_set);
            }

            if new_oc <= best_score {
                // accept move
                for mut x in best.drain(..) {
                    x.cleanup(&mut horner_boxes);
                }

                for x in new_best.drain(..) {
                    best.push(x);
                }

                best_score = new_oc;
                best_scheme.copy_from_slice(&scheme);
            } else {
                for mut x in new_best.drain(..) {
                    x.cleanup(&mut horner_boxes);
                }
            }

            // see if we reject the change
            if new_oc as f64 / best_score as f64 > 1. + 0.5 * (1. - i as f64 / num_tries as f64) {
                //if new_oc as f64 > best_score as f64 {
                scheme.swap(a, b);
            }
        }

        (best, best_score, best_scheme)
    }

    fn get_constants(&self, constants: &mut HashMap<Rational, usize>, shift: usize) {
        match self {
            HornerScheme::Node(n) => {
                if n.gcd != Rational::one() && !constants.contains_key(&n.gcd) {
                    constants.insert(n.gcd.clone(), constants.len() + shift);
                }
                if let Some(content) = n.content_rest.0.as_ref() {
                    content.get_constants(constants, shift);
                }
                if let Some(content) = n.content_rest.1.as_ref() {
                    content.get_constants(constants, shift);
                }
            }
            HornerScheme::Leaf(_, l) => {
                if !constants.contains_key(l) {
                    constants.insert(l.clone(), constants.len() + shift);
                }
            }
        }
    }

    /// Convert the Horner schemes to a list of instructions, suitable for numerical evaluation.
    pub fn to_instr_multiple(
        schemes: &[HornerScheme<RationalField>],
        nvars: usize,
    ) -> InstructionList {
        let mut seen: HashMap<BorrowedHornerNode<'_, RationalField>, usize> = HashMap::default();

        // collect all unique constants
        let mut constant_map: HashMap<Rational, usize> = HashMap::default();

        for s in schemes {
            s.get_constants(&mut constant_map, nvars);
        }

        let mut constants: Vec<_> = constant_map.iter().collect();
        constants.sort_by_key(|(_, c)| *c);

        let mut instr: Vec<_> = (0..nvars)
            .map(|i| Instruction::Init(Variable::Var(i)))
            .collect();

        for x in constants {
            instr.push(Instruction::Init(Variable::Constant(x.0.clone())));
        }

        for s in schemes {
            let v = s.to_instr_rec(&mut instr, &mut seen, &mut constant_map);
            instr.push(Instruction::Yield(v));
        }

        InstructionList { instr }
    }

    /// Convert the Horner scheme to a list of instructions, suitable for numerical evaluation.
    pub fn to_instr(&self, nvars: usize) -> InstructionList {
        Self::to_instr_multiple(std::slice::from_ref(self), nvars)
    }

    fn to_instr_rec<'a>(
        &'a self,
        instr: &mut Vec<Instruction<Rational>>,
        seen: &mut HashMap<BorrowedHornerNode<'a, RationalField>, usize>,
        constant_seen: &mut HashMap<Rational, usize>,
    ) -> usize {
        match self {
            HornerScheme::Node(n) => {
                let nb = BorrowedHornerNode::from(n);

                // check if n = var^pow*content+rest is seen before
                if let Some(v) = seen.get(&nb) {
                    return if n.gcd != Rational::one() {
                        let gcd = *constant_seen.get(&n.gcd).unwrap();
                        instr.push(Instruction::Mul(vec![gcd, *v]));
                        instr.len() - 1
                    } else {
                        *v
                    };
                }

                // check if var^pow*content is seen before
                let mut b = BorrowedHornerNode::from(n);
                b.rest = None;

                let v1 = if let Some(v) = seen.get(&b) {
                    *v
                } else {
                    // check if var^pow is seen before
                    let vp = if n.pow > 1 {
                        let mut c = b.clone();
                        c.content = None;
                        if let Some(v) = seen.get(&c) {
                            *v
                        } else {
                            /// Binary exponentiation
                            fn bin_exp(
                                var: usize,
                                p: usize,
                                instr: &mut Vec<Instruction<Rational>>,
                                seen: &mut HashMap<BorrowedHornerNode<'_, RationalField>, usize>,
                            ) -> usize {
                                if p == 1 {
                                    return var;
                                }

                                // create a borrowed node with the proper power
                                let mut h = AHasher::default();
                                h.write_u8(0);
                                var.hash(&mut h);
                                p.hash(&mut h);
                                let hash = h.finish();
                                let a = BorrowedHornerNode {
                                    var,
                                    pow: p,
                                    gcd: Rational::one(),
                                    content: None,
                                    rest: None,
                                    hash: (hash, hash, hash),
                                };

                                if let Some(a) = seen.get(&a) {
                                    return *a;
                                }

                                if p % 2 == 0 {
                                    let p_half = bin_exp(var, p / 2, instr, seen);
                                    instr.push(Instruction::Mul(vec![p_half, p_half]));
                                } else {
                                    let p_minone = bin_exp(var, p - 1, instr, seen);
                                    instr.push(Instruction::Mul(vec![var, p_minone]));
                                }

                                seen.insert(a, instr.len() - 1);

                                instr.len() - 1
                            }

                            bin_exp(n.var, n.pow, instr, seen)
                        }
                    } else {
                        n.var
                    };

                    if let Some(s) = &n.content_rest.0 {
                        let v1 = s.to_instr_rec(instr, seen, constant_seen);

                        if v1 < vp {
                            instr.push(Instruction::Mul(vec![v1, vp]));
                        } else {
                            instr.push(Instruction::Mul(vec![vp, v1]));
                        }
                        seen.insert(b, instr.len() - 1);
                        instr.len() - 1
                    } else {
                        vp
                    }
                };

                let vr = if let Some(s) = &n.content_rest.1 {
                    let v2 = s.to_instr_rec(instr, seen, constant_seen);

                    if v1 < v2 {
                        instr.push(Instruction::Add(vec![v1, v2]));
                    } else {
                        instr.push(Instruction::Add(vec![v2, v1]));
                    }

                    seen.insert(nb, instr.len() - 1);
                    instr.len() - 1
                } else {
                    v1
                };

                if n.gcd != Rational::one() {
                    let gcd = *constant_seen.get(&n.gcd).unwrap();
                    instr.push(Instruction::Mul(vec![gcd, vr]));
                    instr.len() - 1
                } else {
                    vr
                }
            }
            HornerScheme::Leaf(_, l) => *constant_seen.get(l).unwrap(),
        }
    }
}

// An arithmetical instruction that is part of an `InstructionList`.
#[derive(Debug, Clone)]
pub enum Instruction<N: NumericalFloatLike> {
    Init(Variable<N>),
    Add(Vec<usize>),
    Mul(Vec<usize>),
    Yield(usize),
    Empty,
}

// An variable that is part of an `InstructionList`,
// which may refer to another instruction in the instruction list.
#[derive(Debug, Clone, PartialEq)]
pub enum Variable<N: NumericalFloatLike> {
    Var(usize),
    Constant(N),
}

impl Variable<Rational> {
    fn to_pretty_string(
        &self,
        var_map: &[super::Variable],
        state: &State,
        mode: InstructionSetMode,
    ) -> String {
        match self {
            Variable::Var(v) => var_map[*v].to_string(state),
            Variable::Constant(c) => match mode {
                InstructionSetMode::Plain => format!("{}", c),
                InstructionSetMode::CPP(_) => {
                    if c.is_integer() {
                        format!("T({})", c.numerator())
                    } else {
                        format!("T({})/T({})", c.numerator(), c.denominator())
                    }
                }
            },
        }
    }
}

impl<N: NumericalFloatLike> std::fmt::Display for Variable<N> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Variable::Var(v) => f.write_fmt(format_args!("x{}", v)),
            Variable::Constant(c) => <N as std::fmt::Display>::fmt(c, f),
        }
    }
}

impl<N: NumericalFloatLike> Variable<N> {
    pub fn convert<'a, NO: NumericalFloatLike + From<&'a N>>(&'a self) -> Variable<NO> {
        match self {
            Variable::Var(v) => Variable::Var(*v),
            Variable::Constant(c) => Variable::Constant(NO::from(c)),
        }
    }
}

/// A list of instructions suitable for fast numerical evaluation.
pub struct InstructionList {
    instr: Vec<Instruction<Rational>>,
}

impl InstructionList {
    /// Evaluate the instructions and yield the result.
    /// For a more efficient evaluation, call `to_output()` first.
    pub fn evaluate(&self, samples: &[Rational]) -> Rational {
        let mut eval: Vec<Rational> = vec![Rational::zero(); self.instr.len()];

        for (reg, x) in self.instr.iter().enumerate() {
            match x {
                Instruction::Add(a) => {
                    let mut r = Rational::zero();
                    for x in a {
                        r += &eval[*x];
                    }
                    eval[reg] = r;
                }
                Instruction::Mul(m) => {
                    let mut r = Rational::one();
                    for x in m {
                        r *= &eval[*x];
                    }
                    eval[reg] = r;
                }
                Instruction::Init(i) => match i {
                    Variable::Var(v) => eval[reg] = samples[*v].clone(),
                    Variable::Constant(c) => eval[reg] = c.clone(),
                },
                Instruction::Yield(y) => return eval[*y].clone(),
                Instruction::Empty => {}
            }
        }
        unreachable!()
    }

    /// Return the number of arithmetical operations required for evaluation.
    pub fn op_count(&self) -> usize {
        let mut sum = 0;
        // TODO: detect -1 in init and discount it to stay in line with the FORM counting
        for x in &self.instr {
            sum += match x {
                Instruction::Add(a) => a.len() - 1,
                Instruction::Mul(m) => m.len() - 1,
                Instruction::Yield(_) => 0,
                Instruction::Empty => 0,
                Instruction::Init(_) => 0,
            };
        }
        sum
    }

    /// Fuse `Z1=a+b`, `Z2=Z1+c` to `Z1=a+b+c` if `Z1` is only used in `Z2`
    pub fn fuse_operations(&mut self) {
        let mut use_count: Vec<usize> = vec![0; self.instr.len()];

        for (i, x) in self.instr.iter().enumerate() {
            match x {
                Instruction::Add(a) => {
                    for &v in a {
                        if let Instruction::Mul(_) = self.instr[v] {
                            use_count[v] = 2; // different type, so disable
                        } else {
                            use_count[v] += 1;
                        }
                    }
                }
                Instruction::Mul(m) => {
                    for &v in m {
                        if let Instruction::Add(_) = self.instr[v] {
                            use_count[v] = 2;
                        } else {
                            use_count[v] += 1;
                        }
                    }
                }
                Instruction::Yield(v) => {
                    use_count[*v] = 2; // always different type
                }
                Instruction::Empty => {}
                Instruction::Init(_) => {
                    use_count[i] = 2;
                }
            };
        }

        for i in 0..self.instr.len() {
            // we could be in chain of single use -> single use -> etc so work from the start
            if let Instruction::Add(a) | Instruction::Mul(a) = &self.instr[i] {
                if a.iter().any(|v| use_count[*v] == 1) {
                    let mut instr = std::mem::replace(&mut self.instr[i], Instruction::Empty);

                    if let Instruction::Add(a) | Instruction::Mul(a) = &mut instr {
                        let mut new_a = Vec::with_capacity(a.len());
                        for v in a.drain(..) {
                            if use_count[v] == 1 {
                                if let Instruction::Add(aa) | Instruction::Mul(aa) = &self.instr[v]
                                {
                                    for x in aa {
                                        new_a.push(*x);
                                    }
                                    self.instr[v] = Instruction::Empty;
                                } else {
                                    unreachable!()
                                }
                            } else {
                                new_a.push(v);
                            }
                        }
                        new_a.sort();
                        *a = new_a;
                    }

                    self.instr[i] = instr;
                }
            }
        }

        self.remove_empty_ops();
    }

    /// Remove empty instructions from the list.
    fn remove_empty_ops(&mut self) {
        // now remove the old labels and renumber all
        let mut num_removed_entries = 0;
        let mut new_instr = Vec::with_capacity(self.instr.len());
        let mut cum_step = vec![0; self.instr.len()];
        for (i, mut x) in std::mem::take(&mut self.instr).into_iter().enumerate() {
            cum_step[i] = num_removed_entries;
            if let Instruction::Empty = x {
                num_removed_entries += 1;
                continue;
            }

            match &mut x {
                Instruction::Add(a) | Instruction::Mul(a) => {
                    for v in a {
                        *v -= cum_step[*v];
                    }
                }
                Instruction::Yield(v) => {
                    *v -= cum_step[*v];
                }
                Instruction::Empty => {}
                Instruction::Init(_) => {}
            };
            new_instr.push(x);
        }

        self.instr = new_instr;
    }

    /// Find and extract pairs of variables that appear in more than one instruction.
    /// This reduces the number of operations. Returns `true` iff an extraction could be performed.
    ///
    /// This function can be called multiple times such that common subexpressions that
    /// are larger than pairs can also be extracted.
    pub fn common_pair_elimination(&mut self) -> bool {
        let mut pair_count = HashMap::default();
        let mut last_init = 0;

        for (i, x) in self.instr.iter().enumerate() {
            match x {
                Instruction::Add(m) | Instruction::Mul(m) => {
                    for i in 0..m.len() - 1 {
                        for j in &m[i + 1..] {
                            *pair_count
                                .entry((matches!(x, Instruction::Add(_)), m[i], *j))
                                .or_insert(0) += 1;
                        }
                    }
                }
                Instruction::Yield(_) | Instruction::Empty => {}
                Instruction::Init(_) => {
                    last_init = i + 1;
                }
            };
        }

        let mut v: Vec<_> = pair_count.into_iter().collect();
        v.sort_by_key(|k| Reverse(k.1));

        if v.is_empty() || v[0].1 < 2 {
            return false;
        }

        // TODO: instead of only doing the first, do all non-overlapping replacements

        let (is_add, idx1, idx2) = v[0].0;

        let insert_index = last_init.max(idx1 + 1).max(idx2 + 1);
        self.instr.insert(
            insert_index,
            if is_add {
                Instruction::Add(vec![idx1, idx2])
            } else {
                Instruction::Mul(vec![idx1, idx2])
            },
        );

        // substitute the pair in all instructions
        for x in self.instr.iter_mut().skip(insert_index + 1) {
            let same_op = is_add == matches!(x, Instruction::Add(_));

            match x {
                Instruction::Add(a) | Instruction::Mul(a) => {
                    for v in &mut *a {
                        if *v >= insert_index {
                            *v += 1;
                        }
                    }

                    if same_op {
                        let mut idx1_count = 0;
                        let mut idx2_count = 0;
                        for v in &*a {
                            if *v == idx1 {
                                idx1_count += 1;
                            }
                            if *v == idx2 {
                                idx2_count += 1;
                            }
                        }

                        let pair_count = idx1_count.min(idx2_count);

                        if pair_count > 0 {
                            a.retain(|x| x != &idx1 && x != &idx2);

                            // add back removed indices in cases such as idx1*idx2*idx2
                            if idx1_count > pair_count {
                                a.extend(std::iter::repeat(idx1).take(idx1_count - pair_count));
                            }
                            if idx2_count > pair_count {
                                a.extend(std::iter::repeat(idx2).take(idx2_count - pair_count));
                            }

                            // TODO: Z2=Z1 can be detected here with a.is_empty() && pair_count == 1
                            a.extend(std::iter::repeat(insert_index).take(pair_count));
                            a.sort();
                        }
                    }
                }
                Instruction::Yield(v) => {
                    if *v >= insert_index {
                        *v += 1;
                    }
                }
                Instruction::Empty | Instruction::Init(_) => {}
            };
        }

        // remove trivial relabeling instructions, e.g. Z2=Z1, that could have been created
        let mut map: Vec<_> = (0..self.instr.len()).collect();
        for (i, x) in self.instr.iter_mut().enumerate() {
            match x {
                Instruction::Add(a) | Instruction::Mul(a) => {
                    if a.len() == 1 {
                        map[i] = a[0];
                        *x = Instruction::Empty;
                        continue;
                    }

                    for v in a {
                        *v = map[*v];
                    }
                }
                Instruction::Yield(v) => {
                    *v = map[*v];
                }
                Instruction::Empty => {}
                Instruction::Init(_) => {}
            };
        }

        self.remove_empty_ops();

        true
    }

    /// Convert the instruction list into its output format, where
    /// intermediate registers can be recycled. This vastly improves
    /// the computational memory needed for evaluations.
    pub fn to_output(
        self,
        input_map: Vec<super::Variable>,
        recycle_registers: bool,
    ) -> InstructionListOutput<Rational> {
        if !recycle_registers {
            return InstructionListOutput {
                instr: self.instr.into_iter().enumerate().collect(),
                input_map,
            };
        }

        let mut last_use: Vec<usize> = (0..self.instr.len()).collect();

        for (i, x) in self.instr.iter().enumerate() {
            match x {
                Instruction::Add(a) | Instruction::Mul(a) => {
                    for v in a {
                        last_use[*v] = i;
                    }
                }
                Instruction::Yield(v) => {
                    last_use[*v] = i;
                }
                Instruction::Empty | Instruction::Init(_) => {}
            };
        }

        // prevent init slots from being overwritten
        for (i, x) in self.instr.iter().enumerate() {
            if let Instruction::Init(_) = x {
                last_use[i] = self.instr.len();
            }
        }

        let mut rename_map: Vec<_> = (0..self.instr.len()).collect();

        let mut output = Vec::with_capacity(self.instr.len());

        for (i, mut x) in self.instr.into_iter().enumerate() {
            let cur_last_use = last_use[i];
            // find first free variable
            let reg = if let Some((new_v, lu)) =
                last_use[..i].iter_mut().enumerate().find(|(_, r)| **r <= i)
            {
                *lu = cur_last_use; // set the last use to the current variable last use
                last_use[i] = 0; // make the current index available
                rename_map[i] = new_v; // set the rename map so that every occurrence on the rhs is replaced
                new_v
            } else {
                i
            };

            match &mut x {
                Instruction::Add(a) | Instruction::Mul(a) => {
                    for v in a {
                        *v = rename_map[*v];
                    }
                }
                Instruction::Yield(v) => {
                    *v = rename_map[*v];
                }
                Instruction::Empty | Instruction::Init(_) => {}
            };

            output.push((reg, x));
        }

        InstructionListOutput {
            instr: output,
            input_map,
        }
    }
}

/// A list of instructions suitable for fast numerical evaluation.
pub struct InstructionListOutput<N: NumericalFloatLike> {
    instr: Vec<(usize, Instruction<N>)>,
    input_map: Vec<super::Variable>,
}

/// An efficient structure that performs a range of operations.
/// `Add(reg,index,len)` means: `eval[reg] = eval[indices[index]] +...+ eval[indices[index_pos]]`.
#[derive(Clone, Copy)]
enum InstructionRange {
    Add(usize, usize, usize), // reg, index, len
    Mul(usize, usize, usize),
    Out(usize),
}
/// A fast polynomial evaluator that evaluates polynomials written
/// in the form:
/// ```
/// Z0 = x
/// Z1 = y
/// Z2 = 2.
/// Z3 = 5.
/// Z4 = Z0*Z2
/// Z5 = Z3+Z1+Z0
/// Z4 = Z4*Z3
/// ```
/// where `Z0,Z1` is the range of sample points that get overwritten
/// at every call, `Z2,Z3` is a range of numerical constants
/// that are set only once, and `Z4,Z5` are instructions that
/// get overwritten at every call. These instructions only use
/// indices in the `Z` array and their evaluation can therefore
/// be done efficiently.
#[derive(Clone)]
pub struct InstructionEvaluator<N: NumericalFloatLike> {
    instr: Vec<InstructionRange>,
    indices: Vec<usize>,
    eval: Vec<N>, // evaluation buffer
    out: Vec<N>,  // output buffer
}

impl<N: NumericalFloatLike> InstructionEvaluator<N> {
    /// Evaluate the converted polynomials at a given sample point and
    /// write the values in `out`.
    ///
    /// The user must ensure that `samples` has the
    /// same length as the number of variables in the
    /// polynomials (including non-occurring ones).
    pub fn evaluate(&mut self, samples: &[N]) -> &[N] {
        macro_rules! get_eval {
            ($i:expr) => {
                unsafe { self.eval.get_unchecked(*self.indices.get_unchecked($i)) }
            };
        }

        // write the sample point into the evaluation buffer
        // all constant numbers are still in the evaluation buffer
        self.eval[..samples.len()].clone_from_slice(samples);

        let mut out_counter = 0;

        for x in &self.instr {
            match x {
                InstructionRange::Add(reg, pos, len) => {
                    // unroll the loop for additional performance
                    *unsafe { self.eval.get_unchecked_mut(*reg) } = match len {
                        2 => get_eval!(*pos).clone() + get_eval!(*pos + 1),
                        3 => get_eval!(*pos).clone() + get_eval!(*pos + 1) + get_eval!(*pos + 2),
                        4 => {
                            get_eval!(*pos).clone()
                                + get_eval!(*pos + 1)
                                + get_eval!(*pos + 2)
                                + get_eval!(*pos + 3)
                        }
                        _ => {
                            let mut tmp = N::zero();
                            for aa in unsafe { self.indices.get_unchecked(*pos..(pos + len)) } {
                                tmp += unsafe { self.eval.get_unchecked(*aa) };
                            }
                            tmp
                        }
                    };
                }
                InstructionRange::Mul(reg, pos, len) => {
                    *unsafe { self.eval.get_unchecked_mut(*reg) } = match len {
                        2 => get_eval!(*pos).clone() * get_eval!(*pos + 1),
                        3 => get_eval!(*pos).clone() * get_eval!(*pos + 1) * get_eval!(*pos + 2),
                        4 => {
                            get_eval!(*pos).clone()
                                * get_eval!(*pos + 1)
                                * get_eval!(*pos + 2)
                                * get_eval!(*pos + 3)
                        }
                        _ => {
                            let mut tmp = N::one();
                            for aa in unsafe { self.indices.get_unchecked(*pos..(pos + len)) } {
                                tmp *= unsafe { self.eval.get_unchecked(*aa) };
                            }
                            tmp
                        }
                    };
                }
                InstructionRange::Out(pos) => {
                    unsafe {
                        *self.out.get_unchecked_mut(out_counter) =
                            self.eval.get_unchecked(*pos).clone()
                    };
                    out_counter += 1;
                }
            }
        }

        &self.out
    }
}

impl<N: NumericalFloatLike> InstructionListOutput<N> {
    /// Convert all numbers in the instruction list from the field `N` to the field `NO`.
    pub fn convert<'a, NO: NumericalFloatLike + From<&'a N>>(
        &'a self,
    ) -> InstructionListOutput<NO> {
        let mut instr = Vec::with_capacity(self.instr.len());

        for (reg, inst) in &self.instr {
            let new_instr = match inst {
                Instruction::Add(a) => Instruction::Add(a.clone()),
                Instruction::Mul(a) => Instruction::Mul(a.clone()),
                Instruction::Yield(y) => Instruction::Yield(*y),
                Instruction::Empty => unreachable!("No empty slots allowed in output"),
                Instruction::Init(v) => Instruction::Init(v.convert()),
            };
            instr.push((*reg, new_instr));
        }
        InstructionListOutput {
            instr,
            input_map: self.input_map.clone(),
        }
    }

    /// Create a fast numerical evaluator.
    pub fn evaluator(&self) -> InstructionEvaluator<N> {
        let mut eval = vec![N::zero(); self.instr.len()];

        let mut out_counter = 0;
        let mut simple_instr = vec![];
        let mut indices: Vec<usize> = vec![];
        for (reg, ins) in &self.instr {
            match ins {
                Instruction::Init(x) => {
                    // fill in all constants into the buffer
                    // these entries will never be overwritten
                    if let Variable::Constant(c) = x {
                        eval[*reg] = c.clone();
                    }
                }
                Instruction::Add(a) => {
                    let len = indices.len();
                    indices.extend(a);
                    simple_instr.push(InstructionRange::Add(*reg, len, indices.len() - len));
                }
                Instruction::Mul(a) => {
                    let len = indices.len();
                    indices.extend(a);
                    simple_instr.push(InstructionRange::Mul(*reg, len, indices.len() - len));
                }
                Instruction::Yield(i) => {
                    simple_instr.push(InstructionRange::Out(*i));
                    out_counter += 1;
                }
                Instruction::Empty => {}
            }
        }

        InstructionEvaluator {
            instr: simple_instr,
            indices,
            eval,
            out: vec![N::zero(); out_counter],
        }
    }
}

impl std::fmt::Display for InstructionList {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut out_counter = 0;
        for (reg, x) in self.instr.iter().enumerate() {
            match x {
                Instruction::Add(a) => f.write_fmt(format_args!(
                    "Z{} = {};\n",
                    reg,
                    a.iter()
                        .map(|x| format!("Z{}", x))
                        .collect::<Vec<_>>()
                        .join("+")
                ))?,
                Instruction::Mul(m) => f.write_fmt(format_args!(
                    "Z{} = {};\n",
                    reg,
                    m.iter()
                        .map(|x| format!("Z{}", x))
                        .collect::<Vec<_>>()
                        .join("*")
                ))?,
                Instruction::Yield(y) => {
                    f.write_fmt(format_args!("OUT{} = Z{};\n", out_counter, y))?;
                    out_counter += 1;
                }
                Instruction::Empty => f.write_fmt(format_args!("Z{} = NOP;\n", reg))?,
                Instruction::Init(i) => f.write_fmt(format_args!("Z{} = {};\n", reg, i))?,
            }
        }

        Ok(())
    }
}

pub struct CPPPrinter {}

impl CPPPrinter {
    pub fn format_number(num: &Rational) -> String {
        if num.is_integer() {
            format!("T({})", num.numerator())
        } else {
            format!("T({})/T({})", num.numerator(), num.denominator())
        }
    }
}

#[derive(Clone, Copy)]
pub struct InstructionSetModeCPPSettings {
    pub write_header_and_test: bool,
    pub always_pass_output_array: bool,
}

#[derive(Clone, Copy)]
pub enum InstructionSetMode {
    Plain,
    CPP(InstructionSetModeCPPSettings),
}

pub struct InstructionSetPrinter<'a> {
    pub instr: &'a InstructionListOutput<Rational>,
    pub state: &'a State,
    pub mode: InstructionSetMode,
}

impl<'a> std::fmt::Display for InstructionSetPrinter<'a> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let has_only_one_return_value = self
            .instr
            .instr
            .iter()
            .filter(|x| matches!(x.1, Instruction::Yield(_)))
            .count()
            == 1;
        let use_return_value = has_only_one_return_value
            && match self.mode {
                InstructionSetMode::Plain => false,
                InstructionSetMode::CPP(s) => !s.always_pass_output_array,
            };

        if let InstructionSetMode::CPP(s) = self.mode {
            if s.write_header_and_test {
                f.write_str("#include <iostream>\n")?;
                f.write_str("\n")?;
            }

            f.write_str("template<typename T>\n")?;

            f.write_fmt(format_args!(
                "{} evaluate({}{}) {{\n",
                if use_return_value { "T" } else { "void" },
                self.instr
                    .input_map
                    .iter()
                    .map(|x| format!("T {}", x.to_string(self.state)))
                    .collect::<Vec<_>>()
                    .join(","),
                if use_return_value { "" } else { ", T* out" }
            ))?;

            let max_register = self.instr.instr.iter().map(|r| r.0).max().unwrap_or(0);
            f.write_fmt(format_args!(
                "\tT {};\n",
                (0..=max_register)
                    .map(|x| format!("Z{}", x))
                    .collect::<Vec<_>>()
                    .join(","),
            ))?;
        }

        let mut out_counter = 0;

        for (reg, x) in &self.instr.instr {
            match x {
                Instruction::Add(a) => f.write_fmt(format_args!(
                    "\tZ{} = {};\n",
                    reg,
                    a.iter()
                        .map(|x| format!("Z{}", x))
                        .collect::<Vec<_>>()
                        .join("+")
                ))?,
                Instruction::Mul(m) => f.write_fmt(format_args!(
                    "\tZ{} = {};\n",
                    reg,
                    m.iter()
                        .map(|x| format!("Z{}", x))
                        .collect::<Vec<_>>()
                        .join("*")
                ))?,
                Instruction::Yield(y) => {
                    match self.mode {
                        InstructionSetMode::Plain => {
                            f.write_fmt(format_args!("\tOUT{} = Z{};\n", out_counter, y))?
                        }
                        InstructionSetMode::CPP(_) => {
                            if use_return_value {
                                f.write_fmt(format_args!("\treturn Z{};\n", y))?
                            } else {
                                f.write_fmt(format_args!("\tout[{}] = Z{};\n", out_counter, y))?
                            }
                        }
                    }
                    out_counter += 1;
                }
                Instruction::Empty => f.write_fmt(format_args!("\tZ{} = NOP;\n", reg))?,
                Instruction::Init(x) => f.write_fmt(format_args!(
                    "\tZ{} = {};\n",
                    reg,
                    x.to_pretty_string(&self.instr.input_map, self.state, self.mode)
                ))?,
            }
        }

        if let InstructionSetMode::CPP(s) = self.mode {
            f.write_str("}\n")?;

            if s.write_header_and_test {
                let points: Vec<_> = (0..self.instr.input_map.len())
                    .map(|i| ((i + 1) as f64 / (self.instr.input_map.len() + 2) as f64).to_string())
                    .collect();

                if use_return_value {
                    f.write_fmt(format_args!(
                        "\nint main() {{\n\tstd::cout << evaluate<double>({}) << std::endl;\n}}",
                        points.join(",")
                    ))?;
                } else {
                    f.write_fmt(format_args!(
                        "\nint main() {{\n\tdouble out[{}];\n\tevaluate({}, out);\n\tstd::cout << {} << std::endl;\n}}",
                        out_counter,
                        points.join(","),
                        (0..out_counter)
                            .map(|i| format!("out[{}]", i))
                            .collect::<Vec<_>>()
                            .join(" << \", \" << ")
                    ))?;
                }
            }
        }

        Ok(())
    }
}
