use dimacs::Clause;
use itertools::Itertools;
use log::{debug, info, trace, warn};
use priority_queue::DoublePriorityQueue;

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

use crate::{
    mus::MusEngine,
    profiling::ObservationLog,
    proof::{CompleteArgs, Proof},
    propositional::ClauseKey,
    subproblem::{PartitionOrder, Subproblem},
};

#[derive(Clone, Debug, Default)]
/// Data structure storing the current-best solution and bound
pub struct ProofMinimizerState {
    pub proof: Option<Proof>,
    pub bound: usize,
}

impl ProofMinimizerState {
    fn proof_length(&self) -> Option<usize> {
        self.proof.as_ref().map(|proof| proof.len())
    }

    fn show_proof(&mut self, new_proof: Proof) {
        let update = match self.proof_length() {
            Some(inc_length) => {
                if new_proof.len() < inc_length {
                    if inc_length < self.bound {
                        warn!("Inconsistent bound computation! Discovered a proof of length {} for a subproblem with lower bound {}", inc_length, self.bound);
                        self.bound = inc_length;
                    }
                    true
                } else {
                    false
                }
            }
            None => true,
        };
        if update {
            self.proof = Some(new_proof);
            self.notify();
        }
    }

    fn show_bound(&mut self, new_bound: usize) {
        if new_bound <= self.bound {
            return;
        }
        self.bound = new_bound;
        self.notify();
        if let Some(inc_length) = self.proof_length() {
            if inc_length < self.bound {
                warn!("Inconsistent bound computation! Discovered a proof of length {} for a subproblem with lower bound {}", inc_length, self.bound);
                self.bound = inc_length;
            }
        }
    }

    fn collapse(&mut self) {
        match self.proof_length() {
            Some(inc_length) => {
                self.bound = inc_length;
            }
            None => {
                unreachable!("Trying to collapse a search state without a proof");
            }
        }
    }

    fn is_optimal(&self) -> bool {
        self.proof_length()
            .map(|proof_len| proof_len == self.bound)
            .unwrap_or(false)
    }

    fn notify(&self) {
        if let Some(inc_length) = self.proof_length() {
            info!("Best available proof has length {}", inc_length);
        }
        info!(
            "Best available lower bound on proof length is {}",
            self.bound
        );
    }
}

/// Structure for running the branch-and-bound proof minimizer
pub struct ProofMinimizer<Rng: rand::Rng> {
    /// Queue of unexplored subproblems.
    ///
    /// The priority key is lexicographic:
    /// - choose the subproblem with the lowest bound,
    /// - break ties in favor of the highest timestamps
    subproblems: DoublePriorityQueue<Subproblem, (usize, i64)>,
    /// Deadline for terminating the optimization loop if imposed by user
    deadline: Option<Instant>,
    /// MUS solver used in the bounding procedure calls
    mus_engine: Box<dyn MusEngine>,
    /// Completion call arguments
    complete_args: CompleteArgs,
    /// Exponential backoff base for the completion call
    backoff_base: f64,
    /// Bounding call statistics
    bounding_durations: ObservationLog<Duration, u32>,
    /// Completion call statistics
    completion_durations: RefCell<ObservationLog<Duration, u32>>,
    /// Clause sorting order during the subproblem partitioning
    partition_order: PartitionOrder,
    /// Maximum number of subproblems explored in each partitioning.
    ///
    /// None corresponds to the canonincal branch-and-bound procedure, whereas
    /// the Some(k) variant only considers the first k clauses returned by
    /// `Subproblem::expand`.
    max_branching: Option<usize>,
    /// Maximum number of subproblems in the queue
    max_queue_size: Option<usize>,
    /// Number of opened search tree nodes
    num_nodes: usize,
    /// Number of dominated search tree nodes
    num_dominated: usize,
    /// Number of bounded search tree nodes
    num_bounded: usize,
    /// Number of search tree nodes discarded due to exceeding the queue size
    num_forgotten: usize,
    /// Peak number of subproblems in the queue
    peak_queue: usize,
    /// Random number generator
    rng: Rc<RefCell<Rng>>,
}

impl<Rng: rand::Rng> ProofMinimizer<Rng> {
    /// Creates a new solver with a given bounding configuration
    pub fn new(
        mus_engine: Box<dyn MusEngine>,
        complete_args: CompleteArgs,
        partition_order: PartitionOrder,
        rng: Rng,
    ) -> Self {
        info!("Initializing SAT completion solver with arguments {complete_args:?}");
        ProofMinimizer {
            subproblems: DoublePriorityQueue::new(),
            deadline: None,
            max_branching: None,
            max_queue_size: None,
            mus_engine,
            complete_args,
            backoff_base: 1.0,
            partition_order,
            rng: RefCell::from(rng).into(),
            bounding_durations: ObservationLog::default(),
            completion_durations: RefCell::<_>::default(),
            num_nodes: 0,
            num_dominated: 0,
            num_bounded: 0,
            num_forgotten: 0,
            peak_queue: 0,
        }
    }

    pub fn complete_backoff(&mut self, base: f64) {
        self.backoff_base = base;
    }

    pub fn limit_branching(&mut self, max_split: usize) {
        self.max_branching = Some(max_split);
    }

    pub fn limit_queue(&mut self, max_queue: usize) {
        self.max_queue_size = Some(max_queue);
    }

    fn is_halting(&self) -> bool {
        if let Some(deadline) = self.deadline {
            Instant::now() >= deadline
        } else {
            false
        }
    }

    /// Find the shortest proof of the given CNF formula.
    ///
    /// Return the shortest proof discovered by the solver and the proof size lower bound
    pub fn run(&mut self, formula: &[Clause], deadline: Option<Instant>) -> ProofMinimizerState {
        self.deadline = deadline;
        // Initiate the branch-and-bound loop with the "subproblem"
        // corresponding to the original formula (i.e., no proof steps done yet)
        let mut search_state = ProofMinimizerState::default();
        let root_subproblem =
            Subproblem::from_axioms(formula, self.partition_order, self.max_branching);
        let root_key = -(root_subproblem.timestamp as i64);
        self.subproblems.push(root_subproblem, (0, root_key));
        // Counter for the subproblem identifiers
        let mut ticker: usize = 1;
        // Exponential backoff
        let mut remaining_subproblems: usize = 0;
        let mut next_remaining_subproblems: f64 = 1.0;
        loop {
            if self.is_halting() {
                info!("Time out");
                self.subproblems.clear();
                break;
            }
            // Subproblem queue has suggested some next subproblem, process it
            if let Some((mut next, _)) = self.subproblems.pop_min() {
                debug!("Inbound subproblem:");
                debug!("{next}");
                // Try pruning the subproblem against the new incumbent proof and
                // dominance cache entries
                if self.prune_subproblem(&mut next, &search_state) {
                    continue;
                }
                // Enumerate all clause sets that can be derived immediately
                // and add them in the branch-and-bound queue.
                let mut expand_limit = self.max_branching;
                for mut subproblem in next.expand(
                    self.partition_order,
                    self.max_branching,
                    &self.complete_args,
                    self.rng.clone(),
                ) {
                    subproblem.timestamp = ticker;
                    debug!("-- {subproblem}");
                    if let Some(expand_limit) = expand_limit {
                        debug!(
                            "{} / {} subproblems to introduce in the queue",
                            expand_limit,
                            self.max_branching.unwrap_or(0)
                        );
                        if expand_limit == 0 {
                            break;
                        }
                    }
                    if self.is_halting() {
                        info!("Time out");
                        self.subproblems.clear();
                        break;
                    }
                    if self.prune_subproblem(&mut subproblem, &search_state) {
                        debug!("Pruned before adding to the queue");
                        continue;
                    }
                    ticker += 1;
                    if let Some(expand_limit) = &mut expand_limit {
                        *expand_limit -= 1;
                    }
                    // Update the best known proof with a completion call
                    let mut add_to_queue = true;
                    if remaining_subproblems == 0 {
                        remaining_subproblems = next_remaining_subproblems.ceil() as usize - 1;
                        debug!("Reset the completion counter to {remaining_subproblems}");
                        next_remaining_subproblems *= self.backoff_base;
                        let subproblem_proof =
                            self.complete(&mut subproblem, &mut *self.rng.clone().borrow_mut());
                        // Ignore any subproblems that fix all axioms yet fail to yield a proof
                        add_to_queue = !subproblem.has_all_axioms || subproblem_proof.is_some();
                        if let Some(subproblem_proof) = subproblem_proof {
                            search_state.show_proof(subproblem_proof);
                        }
                    } else {
                        remaining_subproblems -= 1;
                    }
                    if add_to_queue {
                        let bound_key = subproblem.bound_cache().unwrap_or(0);
                        let timestamp_key = -(subproblem.timestamp as i64);
                        self.subproblems
                            .push(subproblem, (bound_key, timestamp_key));
                        debug!("Pushed into the queue with the bound key {bound_key}");
                    } else {
                        debug!("Pruned due to a SAT axiom layer");
                    }
                }
                // Mark this subproblem as explored
                self.num_nodes += 1;
                // Update the queue size statistics
                self.peak_queue = self.peak_queue.max(self.subproblems.len());
                // Prune the subproblems that have become bounded against the incumbent proof
                if let Some(proof_len) = search_state.proof_length() {
                    let subproblems_before_retain = self.subproblems.len();
                    while self
                        .subproblems
                        .peek_max()
                        .and_then(|(subproblem, _)| {
                            subproblem.bound_cache().map(|bound| bound >= proof_len)
                        })
                        .unwrap_or(false)
                    {
                        let (out_subproblem, key) =
                            self.subproblems.pop_max().expect("empty queue");
                        debug!(
                            "Pruned subproblem #{} by post-hoc bounding with key {key:?}",
                            out_subproblem.timestamp
                        );
                    }
                    let subproblems_after_retain = self.subproblems.len();
                    self.num_bounded += subproblems_before_retain - subproblems_after_retain;
                }
                // Dispose of subproblems that exceed the queue size
                if let Some(queue_size) = self.max_queue_size {
                    while self.subproblems.len() > queue_size {
                        let (out_subproblem, key) =
                            self.subproblems.pop_max().expect("empty queue");
                        debug!(
                            "Pruned subproblem #{} by queue size limit with key {key:?}",
                            out_subproblem.timestamp
                        );
                    }
                }
                // Refresh the lower bound
                let new_best_bound = self
                    .subproblems
                    .peek_min()
                    .and_then(|(x, _)| x.bound_cache())
                    .unwrap_or(0);
                search_state.show_bound(new_best_bound);
                // Check if lower and upper bounds coincide and halt if that is the case
                if search_state.is_optimal() {
                    self.num_bounded += self.subproblems.len();
                    self.subproblems.clear();
                }
            } else {
                // Pop can only return None if the queue has already been empty
                search_state.collapse();
                info!("Exhausted the subproblem queue");
                break;
            }
        }
        // The branch-and-bound loop has terminated, output statistics and shut down
        self.report();
        search_state
    }

    fn report(&mut self) {
        info!(
            "Search tree statistics:\n- opened: {} nodes,\n- pruned by dominance: {} nodes,\n- pruned by bounding: {} nodes,\n- forgotten: {} nodes",
            self.num_nodes, self.num_dominated, self.num_bounded, self.num_forgotten
        );
        info!("Peak queue size: {}", self.peak_queue);
        if let Some(stats) = self.bounding_durations.aggregate() {
            info!("Bounding call statistics:\n{stats}");
        }
        if let Some(stats) = self.completion_durations.borrow_mut().aggregate() {
            info!("Completion call statistics:\n{stats}");
        }
        info!("MUS bounder report:");
        self.mus_engine.report();
    }

    fn prune_subproblem(
        &mut self,
        subproblem: &mut Subproblem,
        search_state: &ProofMinimizerState,
    ) -> bool {
        if let Some(proof_len) = search_state.proof_length() {
            // Prune by bound: if a lower bound of the subproblem
            // is at least the size of the best known proof,
            // then this subproblem cannot improve the best proof
            let bounding = self.bound_subproblem(subproblem, proof_len);
            if matches!(bounding, BoundingDecision::Prune) {
                self.num_bounded += 1;
                debug!(
                    "Pruned subproblem #{} by lower bound {}",
                    subproblem.timestamp,
                    subproblem.bound_cache().unwrap_or(0)
                );
                true
            } else {
                false
            }
        } else {
            let bounding = self.bound_subproblem(subproblem, usize::MAX);
            if matches!(bounding, BoundingDecision::Prune) {
                unreachable!("Pruned by bound without a valid proof");
            } else {
                false
            }
        }
    }

    /// Complete the subproblem proof and store the elapsed time
    fn complete(&self, subproblem: &mut Subproblem, rng: &mut Rng) -> Option<Proof> {
        let start_time = Instant::now();
        let proof = subproblem.complete(&self.complete_args, rng);
        self.completion_durations
            .borrow_mut()
            .push(start_time.elapsed());
        proof
    }

    /// Bound the subproblem proof length and store the elapsed time
    fn bound_subproblem(&mut self, subproblem: &mut Subproblem, target: usize) -> BoundingDecision {
        let bound = if let Some(bound) = subproblem.bound_cache() {
            bound
        } else {
            let start_time = Instant::now();
            let bound = subproblem.bound(&mut self.mus_engine);
            self.bounding_durations.push(start_time.elapsed());
            bound
        };
        if bound >= target {
            BoundingDecision::Prune
        } else {
            BoundingDecision::Keep
        }
    }
}

#[derive(Debug, Clone, Copy)]
enum BoundingDecision {
    Prune,
    Keep,
}
