use std::{
    collections::{BTreeMap, BTreeSet},
    io::Write,
    sync::Mutex,
};

use anyhow::bail;
use dimacs::{Clause, Lit, Var};
use itertools::Itertools;
use log::warn;
use rand::{
    seq::{IteratorRandom, SliceRandom as _},
    Rng,
};
use rustsat::types::Clause as RustsatClause;
use rustsat::{lit, solvers::Solve};
use rustsat_cadical::Lrat;

use crate::propositional::{frontier, negate, resolve, to_i64, var_to_neg, var_to_pos, ClauseKey};

const MAX_PROOF_FRONTIER: usize = 2_000;

#[derive(Debug, Clone, PartialEq, Eq)]
/// An enumeration that annotates a proof clause sequence with
/// justifications for each of the steps
pub enum Derivation {
    /// The clause was downloaded from the formula
    Axiom,
    /// The clause is the resolvent of two other clauses
    Resolution(Clause, Clause),
}

#[derive(Debug, Clone, PartialEq, Eq)]
/// A structure storing a clause in a resolution proof with
/// its reason (either a part of the input or a resolvent)
pub struct ProofStep {
    clause: Clause,
    reason: Derivation,
}

#[derive(Debug, Clone)]
/// A structure for storing proof prefixes with auxiliary data
pub struct Proof {
    /// List of derivation steps
    ///
    /// Must terminate at an empty clause if the proof is complete.
    derivations: Vec<ProofStep>,
    /// A map of clauses derived at any point of a proof with their
    /// indices in `derivations` vector
    derived_clauses: BTreeMap<ClauseKey, usize>,
    /// Axioms from the input formula
    available_axioms: BTreeSet<ClauseKey>,
    /// Variables in the input formula
    variables: Vec<Var>,
    /// Clauses in the input formula that have not been used yet
    unused_clauses: BTreeSet<ClauseKey>,
}

#[derive(Debug, Clone, Default)]
/// Arguments for the SAT solver completion calls
pub struct CompleteArgs {
    pub randomize: bool,
    pub use_cadical: bool,
}

impl Proof {
    /// Produces a trivial proof prefix from the input formula
    pub fn from_axioms(axioms: &[Clause]) -> Self {
        let available_axioms = axioms.iter().map(|x| ClauseKey(x.clone())).collect();
        Proof {
            derivations: vec![],
            derived_clauses: Default::default(),
            unused_clauses: Default::default(),
            variables: axioms
                .iter()
                .flat_map(|axiom| axiom.lits().iter().map(|lit| lit.var()))
                .unique()
                .collect_vec(),
            available_axioms,
        }
    }

    pub fn vars(&self) -> &[Var] {
        &self.variables
    }

    pub fn unused_clauses(&self) -> impl Iterator<Item = &Clause> {
        self.unused_clauses.iter().map(|clause_key| &clause_key.0)
    }

    pub fn num_unused_clauses(&self) -> usize {
        self.unused_clauses.len()
    }

    pub fn has_derived(&self, clause: Clause) -> bool {
        self.derived_clauses.contains_key(&ClauseKey(clause))
    }

    pub fn derived_at(&self, clause: Clause) -> Option<usize> {
        self.derived_clauses.get(&ClauseKey(clause)).cloned()
    }

    fn push(&mut self, derivation: ProofStep) -> usize {
        let clause = derivation.clause.clone();
        if let Some(clause_index) = self.derived_at(clause.clone()) {
            return clause_index;
        }
        let new_index = self.derivations.len();
        self.derived_clauses
            .insert(ClauseKey(clause.clone()), new_index);
        // Mark the resolvent as unused
        self.unused_clauses.insert(ClauseKey(clause));
        match &derivation.reason {
            Derivation::Axiom => {
                // No clauses were used
            }
            Derivation::Resolution(lhs_clause, rhs_clause) => {
                // Mark premises as used
                self.unused_clauses.remove(&ClauseKey(lhs_clause.clone()));
                self.unused_clauses.remove(&ClauseKey(rhs_clause.clone()));
            }
        }
        self.derivations.push(derivation);
        new_index
    }

    /// Number of derivations steps in the proof prefix
    pub fn len(&self) -> usize {
        self.derivations.len()
    }

    /// Record the clause download
    pub fn axiom_step(&mut self, clause: Clause) -> Option<usize> {
        if self.available_axioms.contains(&ClauseKey(clause.clone())) {
            Some(self.push(ProofStep {
                clause,
                reason: Derivation::Axiom,
            }))
        } else {
            None
        }
    }

    /// Record the next proof step as the resolution of input clauses.
    ///
    /// Returns Some variant with the proof step index if the resolution clauses is well-defined
    /// and is not a tautology and None otherwise.
    pub fn resolution_step(&mut self, lhs_clause: Clause, rhs_clause: Clause) -> Option<usize> {
        if self.is_complete() {
            return None;
        }
        if let Some(clause) = resolve(&lhs_clause, &rhs_clause) {
            if !self.has_derived(clause.clone()) {
                Some(self.push(ProofStep {
                    clause: clause.clone(),
                    reason: Derivation::Resolution(lhs_clause, rhs_clause),
                }))
            } else {
                None
            }
        } else {
            None
        }
    }

    /// Return true if the last proof clause is empty and false otherwise
    pub fn is_complete(&self) -> bool {
        self.derivations.last().map(|step| step.clause.len()) == Some(0)
    }

    /// Complete the proof prefix with a proof of the available formula
    pub fn conclude(&mut self, args: &CompleteArgs, rng: &mut impl Rng) -> bool {
        if self.is_complete() {
            warn!("Tried to complete a proof that already has been completed");
            return true;
        }
        // Run the configured solver, and either write the proof or report
        // that the formula is SAT (which can happen until all axioms are downloaded)
        if args.use_cadical {
            match self.conclude_with_cadical(args, rng) {
                Ok(true) => {
                    // CaDiCaL has found the proof, proceed normally
                }
                Ok(false) => {
                    // CaDiCaL has shown that the formula is SAT
                    return false;
                }
                Err(e) => {
                    // CaDiCaL has failed to reach a conclusion, report and proceed normally
                    warn!("CaDiCaL error: {e}");
                }
            }
        }
        if self.is_complete() {
            true
        } else {
            match derive_empty(
                &self.proof_frontier(),
                self.vars(),
                &[],
                args.randomize,
                rng,
            ) {
                Some((last_clause, derivations)) => {
                    assert_eq!(last_clause.len(), 0);
                    for derivation in derivations {
                        self.push(derivation);
                    }
                    true
                }
                None => false,
            }
        }
    }

    fn conclude_with_cadical(
        &mut self,
        args: &CompleteArgs,
        rng: &mut impl Rng,
    ) -> anyhow::Result<bool> {
        let mut frontier = self.proof_frontier();
        let mut var_perm = self.variables.iter().map(|var| var.0 as u32).collect_vec();
        if args.randomize {
            var_perm.shuffle(rng);
            frontier.shuffle(rng);
        }
        let var_lookup = BTreeMap::from_iter(
            var_perm
                .iter()
                .enumerate()
                .map(|(ix, var)| (*var, ix as u32)),
        );
        let clause_map_mutex = Mutex::<LratClauseMap>::default();
        let mut solver = rustsat_cadical::CaDiCaL::default();
        if solver
            .set_configuration(rustsat_cadical::Config::UNSAT)
            .is_err()
        {
            warn!("Failed to configure CaDiCaL to work in UNSAT configuration");
        }
        solver.attach_lrat_callback(|index, _, deps| {
            let mut deps_vec = deps.iter().collect_vec();
            let mut clause_map = clause_map_mutex.lock().unwrap();
            let mut running_clause = clause_map.get(*deps_vec.pop().unwrap()).cloned().unwrap();
            while let Some(&next_index) = deps_vec.pop() {
                let premise = clause_map.get(next_index).unwrap();
                let new_clause = resolve(&running_clause, premise).unwrap();
                if !self.has_derived(new_clause.clone()) {
                    self.resolution_step(running_clause.clone(), premise.clone());
                }
                running_clause = new_clause;
            }
            clause_map.add_derived(running_clause, index);
        });
        if args.randomize {
            solver.set_option("seed", rng.next_u32() as i32).unwrap();
            solver.set_option("shuffle", 1).unwrap();
            solver.set_option("shufflerandom", 1).unwrap();
        }
        for clause in &frontier {
            {
                let mut clause_map = clause_map_mutex.lock().unwrap();
                clause_map.add_axiom(clause.clone());
            }
            if let Err(e) = solver.add_clause(RustsatClause::from(
                clause
                    .lits()
                    .iter()
                    .map(to_i64)
                    .map(|l| {
                        let var_id = *var_lookup
                            .get(&(l.unsigned_abs() as u32))
                            .expect("cannot find a variable index");
                        if l > 0 {
                            lit![var_id]
                        } else {
                            !lit![var_id]
                        }
                    })
                    .collect_vec()
                    .as_slice(),
            )) {
                warn!("CaDiCaL has failed to add a clause {clause:?}, error message: {e}");
                return Err(e);
            }
        }
        match solver.solve() {
            Ok(x) => match x {
                rustsat::solvers::SolverResult::Sat => Ok(false),
                rustsat::solvers::SolverResult::Unsat => Ok(true),
                rustsat::solvers::SolverResult::Interrupted => {
                    bail!("Interrupted CaDiCaL")
                }
            },
            Err(e) => {
                warn!("CaDiCaL has failed to solve the bounding formula, reverting to DPLL; error message: {e}");
                Err(e)
            }
        }
    }

    /// Remove the derivations that are not used to deduce the empty clause
    pub fn trim(&mut self) {
        if !self.is_complete() {
            warn!("Tried to trim an incomplete proof, abort the trimming procedure");
            return;
        }
        // Start with the final (empty) clause and proceed adding the dependencies of an
        // empty clause -- and only them! -- into the trimmed proof
        let mut visited = BTreeSet::from_iter([self.derivations.len() - 1]);
        for (index, step) in self.derivations.iter().enumerate().rev() {
            // If the current proof step has not appeared in `visited`
            // that means that it was not used on the trace from the empty clause
            if !visited.contains(&index) {
                continue;
            }
            // Otherwise, find the two clauses used in resolution or skip if an axiom
            match &step.reason {
                Derivation::Axiom => continue,
                Derivation::Resolution(lhs, rhs) => {
                    let lhs_index = self
                        .derived_at(lhs.clone())
                        .expect("unjustified proof step");
                    let rhs_index = self
                        .derived_at(rhs.clone())
                        .expect("unjustified proof step");
                    visited.insert(lhs_index);
                    visited.insert(rhs_index);
                }
            }
        }
        // Introduce the used steps in the reverse order
        self.derivations = visited
            .iter()
            .sorted()
            .map(|ix| self.derivations[*ix].clone())
            .collect_vec();
        self.derived_clauses = BTreeMap::from_iter(
            self.derivations
                .iter()
                .enumerate()
                .map(|(ix, step)| (ClauseKey(step.clause.clone()), ix)),
        );
    }

    /// Returns a subset of clauses in the proof sufficient to complete any optimal proof
    ///
    /// The current implementation excludes all subsumed clauses; in other words, if
    /// clauses A and B have been derived and B contains all the literals in A, then
    /// B is excluded from the formula.
    pub(crate) fn proof_frontier(&self) -> Vec<Clause> {
        let clauses = self
            .derivations
            .iter()
            .map(|step| step.clause.clone())
            .collect_vec();
        if clauses.len() < MAX_PROOF_FRONTIER {
            frontier(&clauses)
        } else {
            clauses
        }
    }

    /// Write the proof in TraceCheck format
    ///
    /// This implementation uses the following subset of TraceCheck for the validation:
    /// - axioms: `<index> <literal>* 0 0`,
    /// - resolutions: `<index> <literal>* 0 <lhs> <rhs> 0`, where `lhs` and `rhs` are
    /// the clause indices used in the given resolution step.
    pub fn write_tracecheck(&self, mut writer: impl Write) -> std::io::Result<()> {
        for (current_index, step) in self.derivations.iter().enumerate() {
            let clause_str = step
                .clause
                .lits()
                .iter()
                .map(|lit| format!("{}", to_i64(lit)))
                .join(" ");
            match &step.reason {
                Derivation::Axiom => {
                    writeln!(writer, "{} {clause_str} 0 0", current_index + 1)?;
                }
                Derivation::Resolution(lhs, rhs) => {
                    let lhs_index = 1 + self.derived_at(lhs.clone()).expect("unjustified clause");
                    let rhs_index = 1 + self.derived_at(rhs.clone()).expect("unjustified clause");
                    writeln!(
                        writer,
                        "{} {clause_str} 0 {lhs_index} {rhs_index} 0",
                        current_index + 1
                    )?;
                }
            }
        }
        Ok(())
    }
}

#[derive(Debug, Default)]
struct LratClauseMap {
    last_clause_index: u64,
    clause_map: BTreeMap<u64, Clause>,
}

impl LratClauseMap {
    fn add_axiom(&mut self, clause: Clause) {
        self.last_clause_index += 1;
        self.clause_map.insert(self.last_clause_index, clause);
    }

    fn add_derived(&mut self, clause: Clause, next_index: u64) {
        if next_index != self.last_clause_index + 1 {
            warn!("Inconsistent LRAT indexing; last visited index is {} but received a clause with index {}", self.last_clause_index, next_index);
        }
        self.last_clause_index = next_index;
        self.clause_map.insert(self.last_clause_index, clause);
    }

    fn get(&self, index: u64) -> Option<&Clause> {
        self.clause_map.get(&index)
    }
}

/// DPLL recursive routine; subproblems are identified by variable assignment lists,
/// and every call performs unit propagation, pick an unassigned variable.
///
/// Returns the clause derived by this subproblem with the proof steps leading to it;
/// the root invokation should return an empty clause, while the invokations without branching
/// should return one of the clauses in the formula.
fn derive_empty(
    formula: &[Clause],
    vars: &[Var],
    assignment: &[Lit],
    randomize: bool,
    rng: &mut impl Rng,
) -> Option<(Clause, Vec<ProofStep>)> {
    // If one of the clauses is falsified by the assignment,
    // return the clause without any further steps
    if let Some(falsified_clause) = formula.iter().find(|clause| {
        clause.lits().iter().all(|clause_lit| {
            let negated_lit = negate(clause_lit);
            assignment.contains(&negated_lit)
        })
    }) {
        return Some((
            falsified_clause.clone(),
            vec![ProofStep {
                clause: falsified_clause.clone(),
                reason: Derivation::Axiom,
            }],
        ));
    }
    // Pick an unassigned variable among non-pure literals,
    // with unit propagations being in the priority
    let mut free_vars = vars
        .iter()
        .filter(|&&var| {
            let pos_lit = var_to_pos(&var);
            let neg_lit = var_to_neg(&var);
            assignment
                .iter()
                .all(|assigned_lit| assigned_lit.var() != var)
                && formula
                    .iter()
                    .any(|clause| clause.lits().contains(&pos_lit))
                && formula
                    .iter()
                    .any(|clause| clause.lits().contains(&neg_lit))
        })
        .cloned()
        .collect_vec();
    let mut short_list_vars = free_vars
        .iter()
        .filter(|&&var| {
            formula.iter().any(|clause| {
                clause.lits().iter().all(|&lit| {
                    let neg_lit = Lit::from_i64(
                        (lit.var().to_u64() as i64)
                            * match lit.sign() {
                                dimacs::Sign::Pos => -1,
                                dimacs::Sign::Neg => 1,
                            },
                    );
                    lit.var() == var || assignment.contains(&neg_lit)
                })
            })
        })
        .cloned()
        .collect_vec();
    if short_list_vars.is_empty() {
        short_list_vars.append(&mut free_vars);
    }
    let free_var = if randomize {
        short_list_vars.iter().choose(rng)
    } else {
        short_list_vars.first()
    }
    .cloned();
    // Split on the chosen variable; if it does not exist, we have found a SAT assignment
    free_var.and_then(|free_var| {
        // Generate two new sub-proofs, one for each of the assignments
        let mut assignment_plus = assignment.iter().cloned().collect_vec();
        let lit_plus = var_to_pos(&free_var);
        assignment_plus.push(lit_plus);
        let mut assignment_minus = assignment.iter().cloned().collect_vec();
        let lit_minus = var_to_neg(&free_var);
        assignment_minus.push(lit_minus);
        // If either of the subproblems is SAT, so is the whole problem
        let lhs = derive_empty(formula, vars, &assignment_plus, randomize, rng);
        let rhs = derive_empty(formula, vars, &assignment_minus, randomize, rng);
        match (lhs, rhs) {
            (Some((clause_plus, proof_plus)), Some((clause_minus, proof_minus))) => Some((free_var, lit_plus, lit_minus, clause_plus, proof_plus, clause_minus, proof_minus)),
            (_, _) => None
        }
    }).and_then(|(free_var, lit_plus, lit_minus, clause_plus, proof_plus, clause_minus, proof_minus)| {
        // Check whether the sub-proofs actually use the branching variable
        let has_plus = clause_plus.lits().contains(&lit_minus);
        let has_minus = clause_minus.lits().contains(&lit_plus);
        let ret = match (has_plus, has_minus) {
            // If one of the branches has refuted the assignment without the new variable, return it as is
            (false, true) => (clause_plus, proof_plus),
            (true, false) => (clause_minus, proof_minus),
            (false, false) => {
                if proof_plus.len() < proof_minus.len() {
                    (clause_plus, proof_plus)
                } else {
                    (clause_minus, proof_minus)
                }
            }
            // Otherwise, merge these proofs and top off with the resolution on two output clauses
            (true, true) => {
                let clause_res = resolve(&clause_plus, &clause_minus).unwrap_or_else(|| {
                panic!(
                    "Cannot resolve clauses {:?} and {:?} from DPLL sub-trees after branching on variable {}",
                    clause_plus, clause_minus, free_var.to_u64()
                )
            });
                (
                    clause_res.clone(),
                    proof_plus
                        .into_iter()
                        .chain(proof_minus)
                        .chain(std::iter::once(ProofStep {
                            clause: clause_res,
                            reason: Derivation::Resolution(clause_plus, clause_minus),
                        }))
                        .collect_vec(),
                )
            }
        };
        Some(ret)
    })
}
