use std::{collections::{BTreeMap, BTreeSet}, ops::Deref};

use dimacs::{Clause, Lit};
use log::{trace, warn};

use crate::dimacs_ext::{to_i64, ToDimacs};

#[derive(Eq, PartialEq, Clone)]
struct ClauseKey(Clause);

impl Deref for ClauseKey {
    type Target = Clause;

    fn deref(&self) -> &Self::Target { &self.0 }
}

impl PartialOrd for ClauseKey {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for ClauseKey {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        let lhs_key = self.lits().iter().map(to_i64);
        let rhs_key = other.lits().iter().map(to_i64);
        lhs_key.cmp(rhs_key)
    }
}

#[derive(Debug)]
pub enum LratDerivation {
    Delete,
    Derive {
        clause: Clause,
        is_rat: bool,
        chain_length: usize
    }
}

pub struct LratParser {
    cnf: BTreeSet<ClauseKey>,
    clause_map: BTreeMap<usize, Clause>,
    used_axioms: BTreeSet<ClauseKey>,
    visited: BTreeMap<ClauseKey, usize>,
    rat_lemmas: BTreeSet<ClauseKey>,
    derived: usize,
}

impl From<&[Clause]> for LratParser {
    fn from(cnf: &[Clause]) -> Self {
        let clause_map = BTreeMap::from_iter(cnf.iter().cloned().enumerate().map(|(index, clause)| (index + 1, clause)));
        let cnf_dedup = BTreeSet::from_iter(cnf.iter().cloned().map(ClauseKey));
        let visited = BTreeMap::from_iter(cnf.iter().cloned().map(|clause| (ClauseKey(clause), 1_usize)));
        LratParser { cnf: cnf_dedup, visited, used_axioms: BTreeSet::default(), rat_lemmas: BTreeSet::default(), derived: 0, clause_map }
    }
}

impl LratParser {
    pub fn parse<'a>(&'a mut self, input: impl std::io::BufRead + 'a) -> impl Iterator<Item = Result<LratDerivation, String>> + 'a {
        let mut proceed = true;
        input.lines()
            .flat_map(|line| line.map(|line| {
                // Extract the new clause index
                let mut tokens = line.split_whitespace().collect::<Vec<_>>();
                tokens.reverse();
                let new_index = tokens.pop().ok_or(line.clone())?.parse::<usize>().map_err(|_| line.clone())?;
                // Skip the line if this is the deletion line
                if tokens.last().ok_or(line.clone())? == &"d" {
                    return Ok(LratDerivation::Delete);
                }
                // Parse the clause literals
                let mut lits = vec![];
                loop {
                    let next_lit = tokens.pop().ok_or(line.clone())?.parse::<i64>().map_err(|_| line.clone())?;
                    if next_lit == 0 { break; }
                    else {
                        lits.push(Lit::from_i64(next_lit));
                    }
                }
                let clause = Clause::from_vec(lits);
                // Parse the dependency clauses, short-circuit on negative indices
                // (as those indicate a RAT derivation)
                let mut is_rat = false;
                let mut chain_length = 0;
                let mut deps = vec![];
                loop {
                    let next_dep = tokens.pop().ok_or(line.clone())?.parse::<i64>().map_err(|_| line.clone())?;
                    match next_dep {
                        0 => { break; },
                        rat_index if rat_index < 0 => {
                            is_rat = true;
                            chain_length = 0;
                            break;
                        }
                        _ => {
                            deps.push(next_dep as usize);
                        }
                    }
                }
                // Starting with the last mentioned clause,
                // resolve the running clause with the last remaining clause
                // and increment its visited counter and the chain length by 1
                let mut res_clause = self.clause_map.get(&deps.pop().ok_or(line.clone())?).ok_or(line.clone())?.clone();
                self.use_axiom(res_clause.clone());
                deps.reverse();
                for next_index in deps {
                    let next_premise = self.clause_map.get(&next_index).ok_or(line.clone())?.clone();
                    self.use_axiom(next_premise.clone());
                    res_clause = resolve(&res_clause, &next_premise).ok_or(line.clone())?;
                    let n_visits = self.visited.entry(ClauseKey(res_clause.clone())).or_default();
                    if *n_visits == 0 {
                        chain_length += 1;
                    } else {
                        trace!("Duplicated derivation of {} in the justification of {}",
                            res_clause.to_dimacs(), clause.to_dimacs());
                    }
                    *n_visits += 1;
                }
                assert!(res_clause.lits().iter().all(|lit| clause.lits().contains(lit)));
                if clause.lits().iter().any(|lit| !res_clause.lits().contains(lit)) {
                    warn!("Derived clause {} is subsumed by the resolution chain result {}",
                        clause.to_dimacs(), res_clause.to_dimacs());
                }
                self.derived += chain_length;
                self.clause_map.insert(new_index, clause.clone());
                Ok(LratDerivation::Derive { clause, is_rat, chain_length })
            }))
            .take_while(move |out| {
                let proceed_prev = proceed;
                proceed &= out.is_ok();
                proceed_prev
            })
    }

    pub fn count_axioms(&self) -> usize {
        assert!(self.used_axioms.iter().all(|axiom| self.cnf.contains(axiom)));
        self.used_axioms.len()
    }

    pub fn count_derived(&self) -> usize {
        self.derived
    }

    pub fn count_rat_steps(&self) -> usize {
        self.rat_lemmas.len()
    }

    pub fn repetitions(&self) -> impl Iterator<Item = (Clause, usize)> + '_ {
        self.visited.iter().filter_map(|(key, &count)| {
            if count > 1 {
                Some((key.0.clone(), count - 1))
            } else {
                assert_ne!(count, 0);
                None
            }
        })
    }

    pub fn purge(&mut self) {
        self.visited = BTreeMap::from_iter(self.cnf.iter().cloned().map(|clause| (clause, 1_usize)));
        self.rat_lemmas.clear();
        self.derived = 0;
    }

    fn use_axiom(&mut self, clause: Clause) {
        if self.cnf.contains(&ClauseKey(clause.clone())) {
            self.used_axioms.insert(ClauseKey(clause));
        }
    }
}

fn resolve(lhs: &Clause, rhs: &Clause) -> Option<Clause> {
    // Find a literal to resolve upon
    let resolve_lit = lhs
        .lits()
        .iter()
        .filter(|lit| {
            rhs.lits()
                .iter()
                .any(|rhs_lit| rhs_lit.var() == lit.var() && rhs_lit.sign() != lit.sign())
        })
        .collect::<Vec<_>>();
    if let [resolve_lit] = resolve_lit[..] {
        // Merge the literals (except for the resolving one) and return
        let mut lits = lhs
            .lits()
            .iter()
            .filter(|lit| lit.var() != resolve_lit.var())
            .chain(
                rhs.lits()
                    .iter()
                    .filter(|lit| lit.var() != resolve_lit.var()),
            )
            .cloned()
            .collect::<Vec<_>>();
        lits.sort_by_key(|lit| lit.var().to_u64());
        lits.dedup();
        Some(Clause::from_vec(lits))
    } else {
        // If non-unique, return None
        None
    }
}
