use crate::core::{
    actually_used_field::ActuallyUsedField,
    bounds::{FieldBounds, IsBounds},
    compile_passes::optimizer::Optimizer,
    expressions::{
        bit_expr::BitExpr,
        conversion_expr::{ConversionExpr, ConversionValue},
        expr::{expr_false, expr_true, Expr},
    },
    ir_builder::ExprStore,
};
use std::cmp::Ordering;

impl Optimizer {
    fn unfold_bit_expr(&self, bit_expr_id: usize) -> UnfoldedBitExpr {
        use BitExpr::*;
        match self.expr_store.get_expr(bit_expr_id) {
            Expr::Bit(Val(false)) => UnfoldedBitExpr::False,
            Expr::Bit(Val(true)) => UnfoldedBitExpr::True,
            Expr::Bit(Not(e)) => {
                if let Expr::Bit(Binop(e1, e2, truth_table)) = self.expr_store.get_expr(*e) {
                    debug_assert!(!truth_table[0], "this should always be false");
                    UnfoldedBitExpr::BinOp(Bin(*e, *e1, *e2, neg_truth_table(*truth_table)))
                } else {
                    debug_assert!(matches!(
                        self.unfold_bit_expr(*e),
                        UnfoldedBitExpr::Input(..)
                    ));
                    UnfoldedBitExpr::NotInput(*e)
                }
            }
            Expr::Bit(Binop(e1, e2, truth_table)) => {
                debug_assert!(!truth_table[0], "this should always be false again");
                UnfoldedBitExpr::BinOp(Bin(bit_expr_id, *e1, *e2, *truth_table))
            }
            _ => UnfoldedBitExpr::Input(bit_expr_id),
        }
    }

    fn optimize_not(&mut self, e: usize) -> Expr<usize> {
        use BitExpr::*;
        match self.unfold_bit_expr(e) {
            UnfoldedBitExpr::False => expr_true(),
            UnfoldedBitExpr::True => expr_false(),
            UnfoldedBitExpr::Input(e_input) => Expr::Bit(Not(e_input)),
            UnfoldedBitExpr::BinOp(Bin(_, e1, e2, truth_table)) => {
                if truth_table[0] {
                    debug_assert!(matches!(self.expr_store.get_expr(e), Expr::Bit(Not(_))));
                    Expr::Bit(Binop(e1, e2, neg_truth_table(truth_table)))
                } else {
                    let base = self
                        .expr_store
                        .new_expr(Expr::Bit(Binop(e1, e2, truth_table)));
                    Expr::Bit(Not(base))
                }
            }
            UnfoldedBitExpr::NotInput(e) => self.expr_store.get_expr(e).clone(),
        }
    }
    // We assume e1 < e2.
    fn binop(&mut self, e1: usize, e2: usize, truth_table: [bool; 4]) -> Expr<usize> {
        use BitExpr::*;
        if truth_table.iter().all(|x| *x) {
            // truth_table is all true
            Expr::Bit(Val(true))
        } else if truth_table.iter().all(|x| !x) {
            // truth_table is all false
            Expr::Bit(Val(false))
        } else if truth_table[0] == truth_table[1] && truth_table[2] == truth_table[3] {
            // truth_table doesn't depend on e2
            if truth_table[0] {
                self.optimize_not(e1)
            } else {
                self.expr_store.get_expr(e1).clone()
            }
        } else if truth_table[0] == truth_table[2] && truth_table[1] == truth_table[3] {
            // truth_table doesn't depend on e1
            if truth_table[0] {
                self.optimize_not(e2)
            } else {
                self.expr_store.get_expr(e2).clone()
            }
        } else if truth_table[0] {
            let base =
                self.expr_store
                    .new_expr(Expr::Bit(Binop(e1, e2, neg_truth_table(truth_table))));
            Expr::Bit(Not(base))
        } else {
            Expr::Bit(Binop(e1, e2, truth_table))
        }
    }
    fn make_binop(
        &mut self,
        e1: usize,
        e2: usize,
        func: impl Fn(bool, bool) -> bool,
    ) -> Expr<usize> {
        match e1.cmp(&e2) {
            Ordering::Less => self.binop(e1, e2, make_truth_table(func)),
            Ordering::Greater => self.binop(e2, e1, make_rev_truth_table(func)),
            Ordering::Equal => {
                let res_false = func(false, false);
                let res_true = func(true, true);
                if res_false == res_true {
                    Expr::Bit(BitExpr::Val(res_false))
                } else if res_false {
                    self.optimize_not(e1)
                } else {
                    self.expr_store.get_expr(e1).clone()
                }
            }
        }
    }
    fn merge_unary_with_unary(
        &mut self,
        e_1: usize,
        f_1: fn(bool) -> bool,
        e_2: usize,
        f_2: fn(bool) -> bool,
        func: impl Fn(bool, bool) -> bool,
    ) -> Expr<usize> {
        self.make_binop(e_1, e_2, |b_1, b_2| func(f_1(b_1), f_2(b_2)))
    }
    fn merge_binop_with_unary(
        &mut self,
        Bin(e_binop, e1, e2, truth_table): Bin,
        e_unary: usize,
        func: impl Fn(bool, bool) -> bool,
    ) -> Expr<usize> {
        if e_unary == e1 {
            self.make_binop(e1, e2, |x1, x2| {
                func(BitExpr::eval_binop(x1, x2, truth_table), x1)
            })
        } else if e_unary == e2 {
            self.make_binop(e1, e2, |x1, x2| {
                func(BitExpr::eval_binop(x1, x2, truth_table), x2)
            })
        } else {
            let binop_offset = truth_table[0];
            self.make_binop(e_binop, e_unary, |x1, x2| func(x1 ^ binop_offset, x2))
        }
    }
    fn merge_binop_with_binop(
        &mut self,
        Bin(a, a_1, a_2, a_truth_table): Bin,
        Bin(b, b_1, b_2, b_truth_table): Bin,
        func: impl Fn(bool, bool) -> bool,
    ) -> Expr<usize> {
        debug_assert!(a_1 < a_2);
        debug_assert!(b_1 < b_2);
        let a_offset = a_truth_table[0];
        let b_offset = b_truth_table[0];
        if a_1 == b_1 && a_2 == b_2 {
            self.make_binop(a_1, a_2, |x1, x2| {
                func(
                    BitExpr::eval_binop(x1, x2, a_truth_table),
                    BitExpr::eval_binop(x1, x2, b_truth_table),
                )
            })
        } else if a == b_1 {
            self.make_binop(b_1, b_2, |x1, x2| {
                func(x1 ^ a_offset, BitExpr::eval_binop(x1, x2, b_truth_table))
            })
        } else if a == b_2 {
            self.make_binop(b_1, b_2, |x1, x2| {
                func(x2 ^ a_offset, BitExpr::eval_binop(x1, x2, b_truth_table))
            })
        } else if b == a_1 {
            self.make_binop(a_1, a_2, |x1, x2| {
                func(BitExpr::eval_binop(x1, x2, a_truth_table), x1 ^ b_offset)
            })
        } else if b == a_2 {
            self.make_binop(a_1, a_2, |x1, x2| {
                func(BitExpr::eval_binop(x1, x2, a_truth_table), x2 ^ b_offset)
            })
        } else {
            self.make_binop(a, b, |xa, xb| func(xa ^ a_offset, xb ^ b_offset))
        }
    }
    fn optimize_binop(
        &mut self,
        e1: usize,
        e2: usize,
        op: impl Fn(bool, bool) -> bool,
    ) -> Expr<usize> {
        match (self.unfold_bit_expr(e1), self.unfold_bit_expr(e2)) {
            (UnfoldedBitExpr::BinOp(b1), UnfoldedBitExpr::BinOp(b2)) => {
                self.merge_binop_with_binop(b1, b2, op)
            }
            (UnfoldedBitExpr::BinOp(bin), b) => {
                let (f, e_b) = b.to_unary(e2);
                self.merge_binop_with_unary(bin, e_b, |x, y| op(x, f(y)))
            }
            (b, UnfoldedBitExpr::BinOp(bin)) => {
                let (f, e_b) = b.to_unary(e1);
                self.merge_binop_with_unary(bin, e_b, |x, y| op(f(y), x))
            }
            (a, b) => {
                let (f_a, e_a) = a.to_unary(e1);
                let (f_b, e_b) = b.to_unary(e2);
                self.merge_unary_with_unary(e_a, f_a, e_b, f_b, op)
            }
        }
    }

    pub fn optimize_bit_expr(&mut self, expr: BitExpr<usize>) -> Expr<usize> {
        use BitExpr::*;

        match expr {
            Not(e) => self.optimize_not(e),
            And(e1, e2) => self.optimize_binop(e1, e2, |a, b| a & b),
            Xor(e1, e2) => self.optimize_binop(e1, e2, |a, b| a ^ b),
            Reveal(e) => self.expr_store.get_expr(e).clone(),
            Val(_) => Expr::Bit(expr),
            Binop(e1, e2, truth_table) => {
                self.optimize_binop(e1, e2, |a, b| BitExpr::eval_binop(a, b, truth_table))
            }
            Random(_) => Expr::Bit(expr),
            KeccakF1600(..) => Expr::Bit(expr),
            Input(..) => Expr::Bit(expr),
        }
    }

    pub fn optimize_conversion_expr<F: ActuallyUsedField>(
        &mut self,
        expr: ConversionExpr<F, usize>,
    ) -> Expr<usize> {
        use ConversionExpr::*;
        let expr_store = &mut self.expr_store;
        match expr {
            BitNumToBit(e, i, mut signed) => {
                let bounds: FieldBounds<F> = expr_store.bounds(e);
                let is_nonnegative = !bounds.has_negatives();
                if is_nonnegative {
                    signed = false;
                }
                let (i, with_not) = {
                    if let Some(val) = bounds.as_constant() {
                        let ConversionValue::Bit(v) = BitNumToBit(val, i, signed).eval().unwrap()
                        else {
                            unreachable!("A BitNumToBit should return a bit");
                        };
                        return expr_bit(v);
                    } else if i > bounds.lowest_included_unique_power_of_2(signed) {
                        let ref_i = bounds.lowest_included_unique_power_of_2(signed);
                        let bits = bounds.bits_in_pos(i, signed);
                        let ref_bits = bounds.bits_in_pos(ref_i, signed);
                        if ref_bits.0 == ref_bits.1 {
                            unreachable!(
                                "this is impossible to happen: bits number {ref_i} (starting at 0 from lsb) of min and max in {bounds:?} should be different, but they are {ref_bits:?}"
                            )
                        }
                        match bits {
                            (true, true) => return expr_true(),
                            (false, false) => return expr_false(),
                            _ => {
                                if bits.0 == ref_bits.0 {
                                    (ref_i, false)
                                } else {
                                    (ref_i, true)
                                }
                            }
                        }
                    } else {
                        (i, false)
                    }
                };
                let new_expr = if let Some(BitToBitNum(v, signed2)) =
                    F::expr_to_conversion_expr(expr_store.get_expr(e).clone())
                {
                    if (signed == signed2 && v.len() <= F::CAPACITY as usize)
                        || (signed && !signed2 && v.len() < F::CAPACITY as usize)
                        || is_nonnegative
                    {
                        if i < v.len() {
                            expr_store.get_expr(v[i]).clone()
                        } else if !signed2 || v.is_empty() || is_nonnegative {
                            expr_false()
                        } else {
                            expr_store.get_expr(*v.last().unwrap()).clone()
                        }
                    } else {
                        F::conversion_expr_to_expr(BitNumToBit(e, i, signed))
                    }
                } else {
                    F::conversion_expr_to_expr(BitNumToBit(e, i, signed))
                };
                if with_not {
                    let n = expr_store.new_expr(new_expr);
                    self.optimize_not(n)
                } else {
                    new_expr
                }
            }
            BitToBitNum(ref v, false) => {
                if v.len() == 1 {
                    if let Some(BitNumToBit(e, 0, _)) =
                        F::expr_to_conversion_expr(self.expr_store.get_expr(v[0]).clone())
                    {
                        if F::bounds_to_field_bounds(*self.expr_store.get_bounds(e))
                            .is_arithmetic_boolean()
                        {
                            return self.expr_store.get_expr(e).clone();
                        }
                    }
                }
                F::conversion_expr_to_expr(expr)
            }
            ScalarFromPlaintextBit(e) => {
                if let Some(BitNumToBit(e_2, 0, _)) =
                    F::expr_to_conversion_expr(self.expr_store.get_expr(e).clone())
                {
                    if F::bounds_to_field_bounds(*self.expr_store.get_bounds(e_2))
                        .is_arithmetic_boolean()
                    {
                        return self.expr_store.get_expr(e_2).clone();
                    }
                }
                F::conversion_expr_to_expr(expr)
            }
            expr => F::conversion_expr_to_expr(expr),
        }
    }
}
fn expr_bit(v: bool) -> Expr<usize> {
    Expr::Bit(BitExpr::Val(v))
}

fn make_truth_table(func: impl Fn(bool, bool) -> bool) -> [bool; 4] {
    [
        func(false, false),
        func(false, true),
        func(true, false),
        func(true, true),
    ]
}
fn make_rev_truth_table(func: impl Fn(bool, bool) -> bool) -> [bool; 4] {
    [
        func(false, false),
        func(true, false),
        func(false, true),
        func(true, true),
    ]
}

struct Bin(usize, usize, usize, [bool; 4]);

enum UnfoldedBitExpr {
    False,
    True,
    Input(usize),
    NotInput(usize),
    BinOp(Bin),
}
impl UnfoldedBitExpr {
    fn to_unary(&self, e: usize) -> (fn(bool) -> bool, usize) {
        match self {
            UnfoldedBitExpr::False => (|_| false, e),
            UnfoldedBitExpr::True => (|_| true, e),
            UnfoldedBitExpr::Input(e) => (|x| x, *e),
            UnfoldedBitExpr::NotInput(e) => (|x| !x, *e),
            UnfoldedBitExpr::BinOp(Bin(..)) => {
                panic!("not unary")
            }
        }
    }
}

fn neg_truth_table(truth_table: [bool; 4]) -> [bool; 4] {
    map_truth_table(truth_table, |b| !b)
}

fn map_truth_table(truth_table: [bool; 4], func: fn(bool) -> bool) -> [bool; 4] {
    [
        func(truth_table[0]),
        func(truth_table[1]),
        func(truth_table[2]),
        func(truth_table[3]),
    ]
}
