use std::collections::HashSet;

type Index = i16;

fn remvove_literal(cnf: &mut Vec<Vec<Index>>, l: Index) {
    cnf.retain(|c| c.binary_search(&l).is_err());

    for c in cnf {
        if let Ok(i) = c.binary_search(&(-l)) {
            c.remove(i);
        }
    }
}

fn unit_propagate(cnf: &mut Vec<Vec<Index>>, solution: &mut Vec<Index>) {
    if let Some(i) = cnf
        .iter()
        .rev()
        .position(|c| c.len() == 1)
        .map(|i| cnf.len() - i - 1)
    {
        let l = cnf[i][0];
        solution.push(l);
        cnf.swap_remove(i);
        remvove_literal(cnf, l);
        unit_propagate(cnf, solution);
    }
}

fn pure_literal_assign(cnf: &mut Vec<Vec<Index>>, solution: &mut Vec<Index>) {
    let mut literals: HashSet<Index> = HashSet::new();
    for c in cnf.iter() {
        literals.extend(c.iter());
    }

    for &l in &literals {
        if !literals.contains(&(-l)) {
            solution.push(l);
            remvove_literal(cnf, l);
        }
    }
}

pub fn dpll(cnf: &[Vec<Index>], solution: &mut Vec<Index>) -> bool {
    let prev_solution = solution.len();
    let mut own_cnf = cnf.to_owned();

    unit_propagate(&mut own_cnf, solution);
    pure_literal_assign(&mut own_cnf, solution);

    if own_cnf.is_empty() {
        return true;
    }

    if own_cnf.iter().any(|c| c.is_empty()) {
        solution.resize(prev_solution, 0);
        return false;
    }

    let literal = own_cnf[0][0];

    own_cnf.push(vec![literal]);
    if dpll(&own_cnf, solution) {
        return true;
    }

    own_cnf.last_mut().unwrap()[0] = -literal;
    if dpll(&own_cnf, solution) {
        return true;
    }

    solution.resize(prev_solution, 0);
    false
}

#[cfg(test)]
mod test {
    use super::dpll;

    #[test]
    fn wikipedia() {
        let cnf = vec![
            vec![-1, 2, 3],
            vec![1, 3, 4],
            vec![-4, 1, 3],
            vec![-3, 1, 4],
            vec![-4, -3, 1],
            vec![-3, -2, 4],
            vec![-3, -1, 2],
            vec![-2, -1, 3],
        ];

        let mut solution = Vec::new();

        assert!(dpll(&cnf, &mut solution));
    }
}