use std::{cell::RefCell, collections::BTreeMap, time::Instant};

use dimacs::Clause;
use itertools::Itertools;
use log::{error, info, trace, warn};
use rustsat::encodings::CollectClauses;
use rustsat::lit;
use rustsat::solvers::SolveIncremental;
use rustsat::types::Clause as RustsatClause;
use rustsat_cadical::CaDiCaL;

use crate::{propositional::to_i64, subproblem};

use super::{MusEngine, MusStatus};

#[derive(Debug, Clone)]
enum SatResult {
    Unsat,
    Sat,
}

#[derive(Default)]
struct SatSolver<'sat> {
    sat: RefCell<CaDiCaL<'sat, 'sat, 'sat>>,
    mandatory_clauses: Vec<Clause>,
    clause_vars: Vec<(Clause, usize)>,
}

impl<'sat> SatSolver<'sat> {
    fn from_bounding_problem(
        formula: &[Clause],
        mandatory: &[Clause],
    ) -> Result<Self, rustsat::OutOfMemory> {
        let mut solver = CaDiCaL::default();
        let mut next_clause_variable = formula
            .iter()
            .flat_map(|clause| clause.lits().iter().map(|lit| lit.var().0 as i32))
            .max()
            .unwrap_or(0)
            + 1;
        let mut mandatory_clauses = vec![];
        let mut clause_vars = vec![];
        for clause in formula {
            let mut lits = clause
                .lits()
                .iter()
                .map(to_i64)
                .map(|x| x as i32)
                .collect_vec();
            if mandatory.contains(clause) {
                mandatory_clauses.push(clause.clone());
            } else {
                lits.push(-next_clause_variable);
                clause_vars.push((clause.clone(), next_clause_variable as usize));
                next_clause_variable += 1;
            }
            solver.add_clause(RustsatClause::from(
                lits.iter()
                    .map(|l| {
                        if *l > 0 {
                            lit![*l as u32 - 1]
                        } else {
                            !lit![-*l as u32 - 1]
                        }
                    })
                    .collect_vec()
                    .as_slice(),
            ))?;
        }
        Ok(Self {
            sat: RefCell::new(solver),
            mandatory_clauses,
            clause_vars,
        })
    }

    fn solve(&self, formula: &[Clause]) -> Result<SatResult, rustsat::OutOfMemory> {
        let mut assumptions = vec![];
        for formula_clause in formula {
            if self.mandatory_clauses.contains(formula_clause) {
                continue;
            } else if let Some(&clause_var) =
                self.clause_vars.iter().find_map(|(clause, clause_var)| {
                    if clause == formula_clause {
                        Some(clause_var)
                    } else {
                        None
                    }
                })
            {
                assumptions.push(clause_var as i32);
            } else {
                error!("Encountered an unknown clause {:?}", formula_clause);
                panic!();
            }
        }
        let solver_res = self.sat.borrow_mut().solve_assumps(
            assumptions
                .iter()
                .map(|l| {
                    if *l > 0 {
                        lit![*l as u32 - 1]
                    } else {
                        !lit![-*l as u32 - 1]
                    }
                })
                .collect_vec()
                .as_slice(),
        );
        let verdict = match solver_res {
            Ok(res) => match res {
                rustsat::solvers::SolverResult::Sat => {
                    // SAT solver has discovered a model,
                    // return SAT and proceed without changes
                    SatResult::Sat
                }
                rustsat::solvers::SolverResult::Unsat => {
                    // SAT solver has proven unsatisfiability,
                    // extract the core from the assumptions,
                    // push it into the solver, and return
                    let clause = self.sat.borrow_mut().core().unwrap().as_slice().into();
                    self.sat.borrow_mut().add_clause(clause)?;
                    SatResult::Unsat
                }
                rustsat::solvers::SolverResult::Interrupted => {
                    // SAT solver has been interrupted, so we assume
                    // that the instance is unsatisfiable
                    // (we need a lower-bound, and this way
                    // we relax the underlying MUS problem)
                    SatResult::Unsat
                }
            },
            Err(e) => {
                // SAT solver has crashed, so we assume
                // that the instance is unsatisfiable
                // (we need a lower-bound, and this way
                // we relax the underlying MUS problem)
                warn!("CaDiCaL has crashed with message: {e}");
                SatResult::Unsat
            }
        };
        Ok(verdict)
    }
}

#[derive(Debug, Clone)]
struct BnbSubproblem {
    formula: Vec<Clause>,
    mandatory: Vec<Clause>,
}

impl BnbSubproblem {
    fn lower_bound(&self, bounder: &mut impl MusEngine) -> usize {
        bounder.optimize(&self.formula, &self.mandatory).lower_bound
    }

    fn fix_undecided(&mut self, solver: &mut SatSolver) -> Result<(), rustsat::OutOfMemory> {
        // Collect all clauses `omega` such that formula \ {omega} is SAT
        // and push them to the mandatory set
        for clause in &self.formula {
            if self.mandatory.contains(clause) {
                continue;
            }
            let formula = self
                .formula
                .iter()
                .filter(|formula_clause| formula_clause != &clause)
                .cloned()
                .collect_vec();
            match solver.solve(&formula)? {
                SatResult::Sat => {
                    self.mandatory.push(clause.clone());
                }
                SatResult::Unsat => {}
            }
        }
        Ok(())
    }
}

#[derive(Default)]
pub struct BnbMusEngine<'sat, SubMus: MusEngine + Default> {
    queue: Vec<BnbSubproblem>,
    formula_size: BTreeMap<usize, usize>,
    mandatory_size: BTreeMap<usize, usize>,
    solver: Option<SatSolver<'sat>>,
    pub sub_bounder: SubMus,
}

impl<'sat, SubMus: MusEngine + Default> BnbMusEngine<'sat, SubMus> {
    pub fn seed_solver(&mut self, formula: &[dimacs::Clause]) {
        self.solver = if let Ok(solver) = SatSolver::from_bounding_problem(formula, &[]) {
            Some(solver)
        } else {
            error!("SAT solver has ran out of memory while creating a bounding SAT formula");
            None
        }
    }
}

impl<'sat, SubMus: MusEngine + Default> From<SubMus> for BnbMusEngine<'sat, SubMus> {
    fn from(sub_bounder: SubMus) -> Self {
        Self {
            sub_bounder,
            ..Default::default()
        }
    }
}

impl<'sat, SubMus: MusEngine + Default> MusEngine for BnbMusEngine<'sat, SubMus> {
    fn optimize_with_deadline(
        &mut self,
        formula: &[Clause],
        mandatory: &[Clause],
        deadline: Option<Instant>,
    ) -> super::MusStatus {
        trace!("Running the BnB bounder");
        trace!("Formula:");
        for clause in formula {
            trace!("{:?}", clause.lits().iter().map(to_i64).collect_vec());
        }
        trace!("Mandatory:");
        for clause in mandatory {
            trace!("{:?}", clause.lits().iter().map(to_i64).collect_vec());
        }
        self.mandatory_size
            .entry(mandatory.len())
            .and_modify(|c| *c += 1)
            .or_insert(1);
        self.formula_size
            .entry(formula.len())
            .and_modify(|c| *c += 1)
            .or_insert(1);
        let mut root_subproblem = BnbSubproblem {
            formula: formula.iter().cloned().collect_vec(),
            mandatory: mandatory.iter().cloned().collect_vec(),
        };
        let is_temporary_solver = self.solver.is_none();
        if is_temporary_solver {
            self.solver = if let Ok(solver) = SatSolver::from_bounding_problem(formula, mandatory) {
                Some(solver)
            } else {
                error!("SAT solver has ran out of memory while creating a bounding SAT formula");
                return super::MusStatus {
                    smallest_subset: formula.to_vec(),
                    lower_bound: mandatory.len(),
                };
            }
        }
        if root_subproblem
            .fix_undecided(
                self.solver
                    .as_mut()
                    .expect("no SAT solver is configured in the bounder"),
            )
            .is_err()
        {
            error!("SAT solver has ran out of memory while fixing undecided clauses");
            return super::MusStatus {
                smallest_subset: formula.to_vec(),
                lower_bound: mandatory.len(),
            };
        }
        let mut status = MusStatus {
            smallest_subset: formula.iter().cloned().collect_vec(),
            lower_bound: root_subproblem.lower_bound(&mut self.sub_bounder),
        };
        self.queue.clear();
        self.queue.push(root_subproblem);
        trace!(
            "Starting the search with {} clauses, {} mandatory clauses, and a lower bound of {}",
            formula.len(),
            mandatory.len(),
            status.lower_bound
        );
        loop {
            if deadline.map(|d| Instant::now() >= d).unwrap_or(false) {
                trace!("Timed out the bound computation");
                break;
            }
            if let Some(subproblem) = self.queue.pop() {
                if subproblem.formula.len() < status.smallest_subset.len() {
                    status.smallest_subset = subproblem.formula.clone();
                }
                if status.lower_bound >= status.upper_bound() {
                    status.lower_bound = status.upper_bound();
                    trace!("Bounds have collapsed, halting the search");
                    break;
                }
                if subproblem.lower_bound(&mut self.sub_bounder) > status.lower_bound {
                    status.lower_bound = subproblem.lower_bound(&mut self.sub_bounder);
                    trace!("New lower bound is {}", status.lower_bound);
                }
                if let Some(clause) = subproblem
                    .formula
                    .iter()
                    .filter(|clause| !subproblem.mandatory.contains(clause))
                    .max_by_key(|clause| clause.len())
                {
                    let mut fix = subproblem.clone();
                    fix.mandatory.push(clause.clone());
                    assert_eq!(fix.mandatory.len(), subproblem.mandatory.len() + 1);
                    let mut remove = subproblem.clone();
                    remove.formula.retain(|c| c != clause);
                    assert_eq!(remove.formula.len() + 1, subproblem.formula.len());
                    if remove
                        .fix_undecided(
                            self.solver
                                .as_mut()
                                .expect("no SAT solver is configured in the bounder"),
                        )
                        .is_err()
                    {
                        error!("SAT solver has ran out of memory while fixing undecided clauses");
                        break;
                    }
                    self.queue.push(fix);
                    self.queue.push(remove);
                    self.queue.sort_by_key(|subproblem| {
                        -(subproblem.lower_bound(&mut self.sub_bounder) as i32)
                    });
                } else {
                    assert_eq!(subproblem.mandatory.len(), subproblem.formula.len());
                }
            } else {
                break;
            }
        }
        trace!(
            "Returning the lower bound of {} and an UNSAT set of cardinality {}",
            status.lower_bound,
            status.upper_bound()
        );
        if is_temporary_solver {
            self.solver = None;
        }
        self.queue.clear();
        status
    }

    fn report(&self) {
        info!("Formula size distribution: {:?}", self.formula_size);
        info!("Mandatory size distribution: {:?}", self.mandatory_size);
    }
}
