use std::rc::Rc;
use std::{
    cell::RefCell,
    collections::{BTreeMap, BTreeSet},
};

use dimacs::Clause;
use itertools::Itertools;
use log::{debug, info, trace};
use priority_queue::PriorityQueue;
use rand::Rng;

use crate::mus::MusEngine;
use crate::proof::{CompleteArgs, Proof};
use crate::propositional::{frontier, resolve, to_i64, ClauseKey};

#[derive(Debug, Clone)]
pub struct Subproblem {
    /// Input formula
    formula: Rc<BTreeSet<ClauseKey>>,
    /// Proof prefix with all used axioms downloaded
    pub(crate) has_all_axioms: bool,
    /// Clauses in the most recent *complete* layer of the proof DAG
    pub(crate) current_layer: Vec<Clause>,
    /// Clauses in the layer currently under construction
    pub(crate) next_layer: Vec<Clause>,
    /// A set of clauses that could be used in the constructed proof layer
    /// together with the premise clauses needed to derive them
    action_set: Vec<(Clause, Option<(Clause, Clause)>)>,
    /// Proof prefix corresponding to the *complete* part of the subproblem
    ///
    /// This contains the same steps as `running_proof_prefix`
    /// but without the current layer.
    current_proof_prefix: Proof,
    /// Proof prefix corresponding to the subproblem
    running_proof_prefix: Proof,
    /// Core clauses of the current proof prefix
    pub(crate) frontier: BTreeSet<ClauseKey>,
    /// Core clauses of the running proof prefix
    pub(crate) running_frontier: BTreeSet<ClauseKey>,
    /// Clauses that are not used in any descendant of this subproblem
    pub(crate) forgotten: BTreeSet<ClauseKey>,
    /// Unique numeric identifier of the subproblem
    pub(crate) timestamp: usize,
    /// Lower bound for the proof length or None if that has not been computed yet
    bound_value: Option<usize>,
    /// Number of completed layers of derived clauses
    pub(crate) num_derived_layers: usize,
    /// Set to true if the subproblem is a part of an exhaustive search
    /// (without any branching factor caps)
    is_exhaustive: bool,
}

impl std::hash::Hash for Subproblem {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.timestamp.hash(state);
    }
}

#[derive(Default, Copy, Clone)]
pub enum PartitionOrder {
    #[default]
    LengthDscFreqDsc,
    LengthDscFreqAsc,
    LengthAscFreqDsc,
    LengthAscFreqAsc,
}

impl PartialEq<Self> for Subproblem {
    fn eq(&self, other: &Self) -> bool {
        self.frontier == other.frontier
            && self.current_layer == other.current_layer
            && self.next_layer == other.next_layer
            && self.action_set == other.action_set
            && self.forgotten == other.forgotten
    }
}

impl Eq for Subproblem {}

const MAX_SUBSUMPTION_FILTER: usize = 2_048;

impl Subproblem {
    pub fn running_proof_len(&self) -> usize {
        self.running_proof_prefix.len()
    }

    pub fn action_set(&self) -> Vec<Clause> {
        self.action_set
            .iter()
            .map(|(c, _)| c)
            .cloned()
            .collect_vec()
    }

    pub fn refresh_action_set(&mut self, ord: PartitionOrder, max_clauses: Option<usize>) {
        self.action_set.clear();
        let mut frequencies = BTreeMap::<i64, i64>::default();
        for clause in &self.running_frontier {
            for lit in clause.lits() {
                let rec = frequencies.entry(to_i64(lit)).or_default();
                *rec += 1;
            }
        }
        let mut derivable = BTreeSet::default();
        let mut top_clauses = PriorityQueue::new();
        if !self.has_all_axioms {
            // Collect all axioms that have not been added yet
            for key in self.formula.iter() {
                if self.frontier.contains(key) {
                    continue;
                }
                derivable.insert(key.clone());
                let clause = key.0.clone();
                self.action_set.push((clause.clone(), None));
                // Update the priority queue and dispose of the lowest-priority clauses
                if let Some(max_clauses) = max_clauses {
                    let clause_len = clause.lits().len() as i64;
                    let total_freq = clause
                        .lits()
                        .iter()
                        .map(|lit| frequencies.get(&to_i64(lit)).cloned().unwrap_or(0))
                        .sum::<i64>();
                    let priority = match ord {
                        PartitionOrder::LengthAscFreqAsc => (clause_len, total_freq),
                        PartitionOrder::LengthAscFreqDsc => (clause_len, -total_freq),
                        PartitionOrder::LengthDscFreqAsc => (-clause_len, total_freq),
                        PartitionOrder::LengthDscFreqDsc => (-clause_len, -total_freq),
                    };
                    top_clauses.push(ClauseKey(clause.clone()), priority);
                    if self.action_set.len() >= max_clauses {
                        if let Some((disposed_clause, _)) = top_clauses.pop() {
                            self.action_set.retain(|(c, _)| c != &*disposed_clause);
                        }
                    }
                }
            }
        } else {
            // Collect all valid resolvents between available clauses and last layer clauses
            // that are not in the formula yet
            let mut n_candidates = 0;
            for clause_lhs in &self.frontier {
                if let Some(max_clauses) = max_clauses {
                    if n_candidates >= 100 * max_clauses {
                        break;
                    }
                }
                for clause_rhs in &self.current_layer {
                    if let Some(new_clause) = resolve(clause_lhs, clause_rhs) {
                        // Disregard clauses already in the frontier set
                        let key = ClauseKey(new_clause.clone());
                        if self.frontier.contains(&key) {
                            continue;
                        }
                        let is_in_action_set = derivable.contains(&key);
                        let is_forgotten = self.forgotten.contains(&key);
                        // Disregard clauses already in the action/forgotten sets
                        if is_in_action_set || is_forgotten {
                            continue;
                        }
                        // Store the clause and its derivation
                        derivable.insert(key);
                        self.action_set.push((
                            new_clause.clone(),
                            Some((clause_lhs.clone().0, clause_rhs.clone())),
                        ));
                        // Update the priority queue and dispose of the lowest-priority clauses
                        if let Some(max_clauses) = max_clauses {
                            let clause_len = new_clause.lits().len() as i64;
                            let total_freq = new_clause
                                .lits()
                                .iter()
                                .map(|lit| frequencies.get(&to_i64(lit)).cloned().unwrap_or(0))
                                .sum::<i64>();
                            let priority = match ord {
                                PartitionOrder::LengthAscFreqAsc => (clause_len, total_freq),
                                PartitionOrder::LengthAscFreqDsc => (clause_len, -total_freq),
                                PartitionOrder::LengthDscFreqAsc => (-clause_len, total_freq),
                                PartitionOrder::LengthDscFreqDsc => (-clause_len, -total_freq),
                            };
                            top_clauses.push(ClauseKey(new_clause.clone()), priority);
                            n_candidates += 1;
                            if self.action_set.len() >= max_clauses {
                                if let Some((disposed_clause, _)) = top_clauses.pop() {
                                    self.action_set.retain(|(c, _)| c != &*disposed_clause);
                                }
                            }
                            if n_candidates >= 100 * max_clauses {
                                break;
                            }
                        }
                    }
                }
            }
        }
        debug!(
            "Filled action set with {} clauses pre-subsumption",
            self.action_set.len()
        );
        // Eliminate dominated clauses
        if self.action_set.len() <= MAX_SUBSUMPTION_FILTER {
            let clauses = frontier(
                &self
                    .action_set
                    .iter()
                    .map(|(c, _)| c)
                    .cloned()
                    .collect_vec(),
            );
            self.action_set.retain(|(c, _)| clauses.contains(c));
        };
        self.filter_action_set();
        self.action_set.sort_by_key(|(c, _)| {
            let clause_len = c.lits().len() as i64;
            let total_freq = c
                .lits()
                .iter()
                .map(|lit| frequencies.get(&to_i64(lit)).cloned().unwrap_or(0))
                .sum::<i64>();
            match ord {
                PartitionOrder::LengthAscFreqAsc => (clause_len, total_freq),
                PartitionOrder::LengthAscFreqDsc => (clause_len, -total_freq),
                PartitionOrder::LengthDscFreqAsc => (-clause_len, total_freq),
                PartitionOrder::LengthDscFreqDsc => (-clause_len, -total_freq),
            }
        });
    }

    pub fn filter_action_set(&mut self) {
        // Discard all clauses that have been explicitly forgotten
        self.action_set
            .retain(|(clause, _)| !self.forgotten.contains(&ClauseKey(clause.clone())));
    }

    /// Computes the lower bound on the proof length
    ///
    /// The arguments correspond to:
    /// - an external API for discovering MUS lower bounds,
    /// - a halting policy that limits the resources on a single bounding run,
    /// - and the best known proof length (smallest lower bound that triggers pruning).
    pub fn bound(&mut self, bounder: &mut Box<dyn MusEngine>) -> usize {
        if let Some(bound_value) = self.bound_value {
            // Return the value computed earlier
            bound_value
        } else {
            // Find the lower bound on the MUS size and bound the remaining proof length
            // by subtracting 1 from it (any proof has at least #MUS - 1 new clauses)
            let fixed = if self.is_exhaustive {
                self.running_proof_prefix
                    .unused_clauses()
                    .cloned()
                    .collect_vec()
            } else {
                vec![]
            };
            let mus = if self.has_all_axioms {
                bounder.optimize(
                    self.running_frontier
                        .iter()
                        .map(|x| x.0.clone())
                        .collect_vec()
                        .as_slice(),
                    &fixed,
                )
            } else {
                bounder.optimize(
                    &self
                        .running_frontier
                        .iter()
                        .map(|x| &x.0)
                        .chain(self.action_set.iter().map(|(x, _)| x))
                        .cloned()
                        .collect_vec(),
                    &fixed,
                )
            };
            let mut bound_value = match (self.is_exhaustive, self.has_all_axioms) {
                (false, _) => {
                    self.running_proof_len() - self.running_proof_prefix.num_unused_clauses()
                }
                (true, true) => self.running_proof_len(),
                (true, false) => mus.lower_bound,
            };
            bound_value += mus.lower_bound;
            bound_value = bound_value.saturating_sub(1);
            self.bound_value = Some(bound_value);
            bound_value
        }
    }

    /// Returns the bound value if it has already been computed before
    pub fn bound_cache(&self) -> Option<usize> {
        self.bound_value
    }

    /// Builds the root subproblem
    pub fn from_axioms(axioms: &[Clause], ord: PartitionOrder, max_actions: Option<usize>) -> Self {
        let mut proof_prefix = Proof::from_axioms(axioms);
        let axioms_vec = axioms.to_vec();
        let is_exhaustive = max_actions.is_none();
        if !is_exhaustive {
            for axiom in &axioms_vec {
                proof_prefix.axiom_step(axiom.clone());
            }
        }
        let frontier = proof_prefix
            .proof_frontier()
            .into_iter()
            .map(ClauseKey)
            .collect::<BTreeSet<_>>();
        if is_exhaustive {
            assert!(frontier.is_empty());
        }
        let mut subproblem = Self {
            formula: Rc::new(axioms_vec.iter().map(|x| ClauseKey(x.clone())).collect()),
            has_all_axioms: !is_exhaustive,
            action_set: vec![],
            running_frontier: frontier.clone(),
            frontier,
            current_proof_prefix: proof_prefix.clone(),
            running_proof_prefix: proof_prefix,
            timestamp: 0,
            forgotten: BTreeSet::default(),
            current_layer: if is_exhaustive { vec![] } else { axioms_vec },
            next_layer: vec![],
            bound_value: None,
            num_derived_layers: 0,
            is_exhaustive,
        };
        debug!("Collected the root subproblem data");
        subproblem.refresh_action_set(ord, max_actions);
        debug!(
            "Starting with an action set of cardinality {}",
            subproblem.action_set.len()
        );
        subproblem
    }

    /// Returns the trimmed proof from this subproblem,
    /// completing it with a CaDiCaL / DPLL run if needed
    pub fn complete(&self, args: &CompleteArgs, rng: &mut impl Rng) -> Option<Proof> {
        let mut res_proof = self.running_proof_prefix.clone();
        if !res_proof.is_complete() && !res_proof.conclude(args, rng) {
            return None;
        }
        trace!(
            "The proof for subproblem #{} has been completed with length {}",
            self.timestamp,
            res_proof.len()
        );
        res_proof.trim();
        Some(res_proof)
    }

    /// Partitions the subproblem on the next proof layer
    pub fn expand(
        &mut self,
        ord: PartitionOrder,
        max_actions: Option<usize>,
        args: &CompleteArgs,
        rng: Rc<RefCell<impl Rng>>,
    ) -> impl Iterator<Item = Subproblem> + '_ {
        debug!("Entered `expand` for subproblem #{}", self.timestamp);
        if !self.has_all_axioms {
            let n_original = self.action_set.len();
            self.fix_critical_axioms(args, &mut *rng.borrow_mut());
            let n_fixed = n_original - self.action_set.len();
            if n_fixed == n_original && self.formula.len() == n_fixed {
                info!("Received a MUS formula with {n_fixed} clauses");
            }
            debug!("Fixed {n_fixed} clauses out of {n_original}");
        }
        self.expand_fixed(ord, max_actions)
    }

    fn expand_fixed(
        &self,
        ord: PartitionOrder,
        max_actions: Option<usize>,
    ) -> impl Iterator<Item = Subproblem> + '_ {
        // Create a new subproblem for each clause, forgetting all of the previous ones,
        // and updating the next layer with the new clause
        self.action_set
            .iter()
            .enumerate()
            .filter_map(|(clause_ix, (clause, premises))| {
                let mut new_proof = self.running_proof_prefix.clone();
                let step_result = match premises {
                    None => new_proof.axiom_step(clause.clone()),
                    Some((lhs, rhs)) => new_proof.resolution_step(lhs.clone(), rhs.clone()),
                };
                if step_result.is_some() {
                    let mut new_subproblem = Subproblem {
                        formula: self.formula.clone(),
                        has_all_axioms: self.has_all_axioms,
                        action_set: self.action_set.clone(),
                        frontier: self.frontier.clone(),
                        running_frontier: new_proof
                            .proof_frontier()
                            .into_iter()
                            .map(ClauseKey)
                            .collect(),
                        current_proof_prefix: self.current_proof_prefix.clone(),
                        running_proof_prefix: new_proof,
                        timestamp: 0,
                        forgotten: self.forgotten.clone(),
                        current_layer: self.current_layer.clone(),
                        next_layer: self.next_layer.clone(),
                        bound_value: None,
                        num_derived_layers: self.num_derived_layers,
                        is_exhaustive: self.is_exhaustive,
                    };
                    new_subproblem.next_layer.push(clause.clone());
                    for (forgotten, _) in &self.action_set[..=clause_ix] {
                        new_subproblem
                            .forgotten
                            .insert(ClauseKey(forgotten.clone()));
                    }
                    new_subproblem.filter_action_set();
                    trace!(
                        "Subproblem #{}, expanded {} / {} sub-subproblems",
                        self.timestamp,
                        clause_ix + 1,
                        self.action_set.len()
                    );
                    Some(new_subproblem)
                } else {
                    None
                }
            })
            // Add an extra subproblem that forgets everything and fixes the current layer
            .chain(
                std::iter::once_with(move || {
                    let mut next_layer_subproblem = Subproblem {
                        formula: self.formula.clone(),
                        has_all_axioms: true,
                        current_layer: self.next_layer.clone(),
                        next_layer: vec![],
                        action_set: vec![],
                        current_proof_prefix: self.running_proof_prefix.clone(),
                        running_proof_prefix: self.running_proof_prefix.clone(),
                        frontier: self
                            .running_proof_prefix
                            .proof_frontier()
                            .into_iter()
                            .map(ClauseKey)
                            .collect(),
                        running_frontier: self
                            .running_proof_prefix
                            .proof_frontier()
                            .into_iter()
                            .map(ClauseKey)
                            .collect(),
                        forgotten: self.forgotten.clone(),
                        timestamp: 0,
                        bound_value: self.bound_value,
                        num_derived_layers: self.num_derived_layers + 1,
                        is_exhaustive: self.is_exhaustive,
                    };
                    for (clause, _) in &self.action_set {
                        next_layer_subproblem
                            .forgotten
                            .insert(ClauseKey(clause.clone()));
                    }
                    next_layer_subproblem.refresh_action_set(ord, max_actions);
                    next_layer_subproblem
                })
                .filter(|next| !next.current_layer.is_empty()),
            )
            .filter(move |subproblem| {
                if max_actions.is_none() {
                    // Complete search only:
                    // filter out subproblems with clauses that are
                    // (a) unused and (b) not in frontier
                    !subproblem.has_all_axioms
                        || subproblem
                            .running_proof_prefix
                            .unused_clauses()
                            .all(|c| subproblem.running_frontier.contains(&ClauseKey(c.clone())))
                } else {
                    // Otherwise, do the same for all clauses but axioms
                    subproblem.running_proof_prefix.unused_clauses().all(|c| {
                        let key = ClauseKey(c.clone());
                        subproblem.formula.contains(&key)
                            || subproblem.running_frontier.contains(&key)
                    })
                }
            })
    }

    fn fix_critical_axioms(&mut self, args: &CompleteArgs, rng: &mut impl Rng) {
        assert!(!self.has_all_axioms);
        // Mark all clauses such that removing them from the union of the already added
        // axioms and the unresolved ones (in the action set) yields a SAT formula
        let mut fix_indices = (0..self.action_set.len())
            .filter(|&forbidden_index| {
                let mut proof_prefix = self.running_proof_prefix.clone();
                assert_eq!(proof_prefix.len(), self.next_layer.len());
                for (axiom_index, (axiom, _)) in self.action_set.iter().enumerate() {
                    if axiom_index == forbidden_index {
                        continue;
                    }
                    proof_prefix.axiom_step(axiom.clone());
                }
                !proof_prefix.conclude(args, rng)
            })
            .collect_vec();
        // Move all marked clauses from the action set to the next layer
        for &index in &fix_indices {
            let axiom = self.action_set[index].0.clone();
            self.next_layer.push(axiom.clone());
            self.running_proof_prefix.axiom_step(axiom);
        }
        // Discard clauses in the decreasing order of their indices in order to
        // not invalidate the later indices
        fix_indices.sort_by(|a, b| a.cmp(b).reverse());
        for index in fix_indices {
            self.action_set.remove(index);
        }
    }
}

impl std::fmt::Display for Subproblem {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
        writeln!(f, "Subproblem #{}", self.timestamp)?;
        writeln!(
            f,
            "Axiom layer is {}",
            if self.has_all_axioms {
                "complete"
            } else {
                "incomplete"
            }
        )?;
        let mut curr = self
            .current_layer
            .iter()
            .map(|c| c.lits().iter().map(to_i64).collect_vec())
            .collect_vec();
        curr.sort();
        writeln!(f, "Current layer: {curr:?}")?;
        let mut next = self
            .next_layer
            .iter()
            .map(|c| c.lits().iter().map(to_i64).collect_vec())
            .collect_vec();
        next.sort();
        writeln!(f, "Next layer: {next:?}")?;
        let mut actions = self
            .action_set
            .iter()
            .map(|(c, _)| c.lits().iter().map(to_i64).collect_vec())
            .collect_vec();
        actions.sort();
        writeln!(f, "Action set: {actions:?}")?;
        let mut curr_frontier = self
            .frontier
            .iter()
            .map(|c| c.lits().iter().map(to_i64).collect_vec())
            .collect_vec();
        curr_frontier.sort();
        writeln!(f, "Frontier: {curr_frontier:?}")?;
        let mut forgotten = self
            .forgotten
            .iter()
            .map(|c| c.lits().iter().map(to_i64).collect_vec())
            .collect_vec();
        forgotten.sort();
        write!(f, "Forgotten: {forgotten:?}")?;
        Ok(())
    }
}

#[cfg(test)]
mod subproblem_tests {
    use super::*;
    use dimacs::{parse_dimacs, Instance, Lit};

    #[test]
    fn test_clause_dedup() {
        let clauses =
            if let Instance::Cnf { clauses, .. } = parse_dimacs("p cnf 2 2\n1 0\n1 0\n").unwrap() {
                clauses
            } else {
                unreachable!()
            };
        let subproblem = Subproblem::from_axioms(&clauses, PartitionOrder::default(), None);
        assert_eq!(
            subproblem.frontier,
            vec![Clause::from_vec(vec![Lit::from_i64(1)])]
        );
    }

    #[test]
    fn test_clause_removal() {
        let clauses = if let Instance::Cnf { clauses, .. } =
            parse_dimacs("p cnf 2 2\n1 2 0\n1 0\n").unwrap()
        {
            clauses
        } else {
            unreachable!()
        };
        let subproblem = Subproblem::from_axioms(&clauses, PartitionOrder::default(), None);
        assert_eq!(
            subproblem.frontier,
            vec![Clause::from_vec(vec![Lit::from_i64(1)])]
        );
    }

    #[test]
    fn test_clause_sort() {
        let clauses = if let Instance::Cnf { clauses, .. } =
            parse_dimacs("p cnf 4 2\n1 2 0\n1 0\n2 0\n-1 -2 0\n").unwrap()
        {
            clauses
        } else {
            unreachable!()
        };
        let subproblem = Subproblem::from_axioms(&clauses, PartitionOrder::default(), None);
        assert_eq!(
            subproblem.frontier,
            vec![
                Clause::from_vec(vec![Lit::from_i64(-1), Lit::from_i64(-2)]),
                Clause::from_vec(vec![Lit::from_i64(1)]),
                Clause::from_vec(vec![Lit::from_i64(2)]),
            ]
        );
    }
}
