//! Cardinality networks and sorting networks for at-most-k constraints.
//!
//! This module implements various encodings for cardinality constraints:
//! - Sorting networks (bitonic sort, odd-even merge sort)
//! - Cardinality networks (simplified sorting for at-most-k)
//! - Direct encoding, grouped encoding, bimander encoding
//!
//! Reference: Z3's `util/sorting_network.h`
//! Based on:
//! - Batcher (1968): "Sorting Networks and their Applications"
//! - Eén & Sörensson (2006): "Translating Pseudo-Boolean Constraints into SAT"
//! - Abío et al. (2013): "A Parametric Approach for Smaller and Better Encodings of Cardinality Constraints" (CP 2013)

use oxiz_sat::{Lit, Var};
use smallvec::SmallVec;

/// Clause generated by cardinality network encoding
#[derive(Debug, Clone)]
pub struct CardinalityClause {
    pub lits: SmallVec<[Lit; 4]>,
}

impl CardinalityClause {
    pub fn new(lits: impl IntoIterator<Item = Lit>) -> Self {
        Self {
            lits: lits.into_iter().collect(),
        }
    }
}

/// Type of cardinality network encoding
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CardinalityNetworkEncoding {
    /// Direct encoding (good for small k)
    Direct,
    /// Sorting network (full sort)
    Sorting,
    /// Cardinality network (optimized for at-most-k)
    Cardinality,
    /// Grouped encoding
    Grouped,
    /// Bimander encoding (hybrid)
    Bimander,
    /// Odd-even merge sort
    OddEvenMerge,
}

/// Sorting network builder.
///
/// This implements a bitonic sorting network that can sort literals.
/// The network is built as a circuit of comparators.
#[derive(Debug)]
pub struct SortingNetwork {
    /// Generated clauses
    clauses: Vec<CardinalityClause>,
    /// Next variable ID
    next_var: u32,
    /// Encoding type
    #[allow(dead_code)]
    encoding: CardinalityNetworkEncoding,
}

impl SortingNetwork {
    /// Create a new sorting network
    pub fn new(next_var: u32, encoding: CardinalityNetworkEncoding) -> Self {
        Self {
            clauses: Vec::new(),
            next_var,
            encoding,
        }
    }

    /// Get the next available variable ID
    pub fn next_var(&self) -> u32 {
        self.next_var
    }

    /// Build a full sorting network for inputs
    ///
    /// Returns the output literals in sorted order (smallest to largest when viewed as bits)
    pub fn build_sorting_network(&mut self, inputs: &[Lit]) -> Vec<Lit> {
        if inputs.is_empty() {
            return Vec::new();
        }
        if inputs.len() == 1 {
            return vec![inputs[0]];
        }

        // Pad to power of 2 with false literals
        let mut padded: Vec<Lit> = inputs.to_vec();
        while !Self::is_power_of_two(padded.len()) {
            // Use a constant false variable
            let false_var = Var(self.next_var);
            self.next_var += 1;
            // Add unit clause: ~false_var (making it always false)
            self.clauses
                .push(CardinalityClause::new([Lit::neg(false_var)]));
            padded.push(Lit::pos(false_var));
        }

        // Build bitonic sorting network
        let mut output = padded;
        let len = output.len();
        self.bitonic_sort(&mut output, 0, len, true);
        output
    }

    /// Bitonic sort implementation (recursive)
    fn bitonic_sort(&mut self, arr: &mut [Lit], start: usize, len: usize, ascending: bool) {
        if len <= 1 {
            return;
        }

        let mid = len / 2;

        // Sort first half ascending
        self.bitonic_sort(arr, start, mid, true);
        // Sort second half descending
        self.bitonic_sort(arr, start + mid, len - mid, false);

        // Merge the two halves
        self.bitonic_merge(arr, start, len, ascending);
    }

    /// Bitonic merge
    fn bitonic_merge(&mut self, arr: &mut [Lit], start: usize, len: usize, ascending: bool) {
        if len <= 1 {
            return;
        }

        let mid = Self::greatest_power_of_two_less_than(len);

        // Compare-exchange elements that are mid apart
        for i in start..(start + len - mid) {
            self.compare_exchange(arr, i, i + mid, ascending);
        }

        // Recursively merge
        self.bitonic_merge(arr, start, mid, ascending);
        if start + mid < arr.len() {
            self.bitonic_merge(arr, start + mid, len - mid, ascending);
        }
    }

    /// Compare and exchange two elements
    fn compare_exchange(&mut self, arr: &mut [Lit], i: usize, j: usize, ascending: bool) {
        if i >= arr.len() || j >= arr.len() {
            return;
        }

        let a = arr[i];
        let b = arr[j];

        // Create fresh variables for min and max
        let min_var = Var(self.next_var);
        self.next_var += 1;
        let max_var = Var(self.next_var);
        self.next_var += 1;

        let min_lit = Lit::pos(min_var);
        let max_lit = Lit::pos(max_var);

        // min = a AND b
        // min => a: ~min | a
        self.clauses
            .push(CardinalityClause::new([Lit::neg(min_var), a]));
        // min => b: ~min | b
        self.clauses
            .push(CardinalityClause::new([Lit::neg(min_var), b]));
        // a AND b => min: ~a | ~b | min
        self.clauses
            .push(CardinalityClause::new([a.negate(), b.negate(), min_lit]));

        // max = a OR b
        // a => max: ~a | max
        self.clauses
            .push(CardinalityClause::new([a.negate(), max_lit]));
        // b => max: ~b | max
        self.clauses
            .push(CardinalityClause::new([b.negate(), max_lit]));
        // max => a OR b: ~max | a | b
        self.clauses
            .push(CardinalityClause::new([Lit::neg(max_var), a, b]));

        // Update array based on sort direction
        if ascending {
            arr[i] = min_lit;
            arr[j] = max_lit;
        } else {
            arr[i] = max_lit;
            arr[j] = min_lit;
        }
    }

    /// Build a cardinality network (optimized for at-most-k)
    ///
    /// This is more efficient than a full sorting network when we only need
    /// to enforce at-most-k constraint.
    pub fn build_cardinality_network(&mut self, inputs: &[Lit], k: usize) -> Vec<Lit> {
        if inputs.is_empty() || k >= inputs.len() {
            return inputs.to_vec();
        }

        // For small k, use direct encoding
        if k == 1 {
            return self.build_direct_amo(inputs);
        }

        // For larger k, use simplified sorting network
        // Only compute the first k+1 outputs
        self.build_partial_sorting_network(inputs, k + 1)
    }

    /// Build a direct at-most-one encoding (returns single output representing "at least one")
    fn build_direct_amo(&mut self, inputs: &[Lit]) -> Vec<Lit> {
        // Pairwise encoding: for all pairs, ~x_i | ~x_j
        for i in 0..inputs.len() {
            for j in (i + 1)..inputs.len() {
                self.clauses.push(CardinalityClause::new([
                    inputs[i].negate(),
                    inputs[j].negate(),
                ]));
            }
        }

        // Return a literal representing "at least one is true"
        let at_least_one_var = Var(self.next_var);
        self.next_var += 1;
        let at_least_one = Lit::pos(at_least_one_var);

        // at_least_one <=> x_1 | x_2 | ... | x_n
        // at_least_one => x_1 | ... | x_n: ~at_least_one | x_1 | ... | x_n
        let mut clause_lits = SmallVec::<[Lit; 8]>::new();
        clause_lits.push(Lit::neg(at_least_one_var));
        clause_lits.extend(inputs.iter().copied());
        self.clauses.push(CardinalityClause::new(clause_lits));

        // x_i => at_least_one for each i
        for &input in inputs {
            self.clauses
                .push(CardinalityClause::new([input.negate(), at_least_one]));
        }

        vec![at_least_one]
    }

    /// Build a partial sorting network (only outputs the first k elements)
    fn build_partial_sorting_network(&mut self, inputs: &[Lit], k: usize) -> Vec<Lit> {
        // For simplicity, use merge-based approach
        // Split inputs into pairs, merge pairs, etc.
        let mut current_level: Vec<Vec<Lit>> = inputs.iter().map(|&lit| vec![lit]).collect();

        while current_level.len() > 1 {
            let mut next_level = Vec::new();

            let mut i = 0;
            while i + 1 < current_level.len() {
                let merged = self.merge_sorted(&current_level[i], &current_level[i + 1], k);
                next_level.push(merged);
                i += 2;
            }

            // Handle odd element
            if i < current_level.len() {
                next_level.push(current_level[i].clone());
            }

            current_level = next_level;
        }

        if current_level.is_empty() {
            Vec::new()
        } else {
            let mut result = current_level[0].clone();
            result.truncate(k);
            result
        }
    }

    /// Merge two sorted lists (odd-even merge)
    fn merge_sorted(&mut self, a: &[Lit], b: &[Lit], max_out: usize) -> Vec<Lit> {
        if a.is_empty() {
            return b.iter().copied().take(max_out).collect();
        }
        if b.is_empty() {
            return a.iter().copied().take(max_out).collect();
        }

        // Simple odd-even merge
        let n = (a.len() + b.len()).min(max_out);
        let mut result = vec![Lit::pos(Var(0)); n]; // Placeholder

        // Interleave and compare
        let mut ai = 0;
        let mut bi = 0;
        let mut out_idx = 0;

        while out_idx < n && (ai < a.len() || bi < b.len()) {
            if ai >= a.len() {
                result[out_idx] = b[bi];
                bi += 1;
            } else if bi >= b.len() {
                result[out_idx] = a[ai];
                ai += 1;
            } else {
                // Create comparator
                let min_var = Var(self.next_var);
                self.next_var += 1;
                let min_lit = Lit::pos(min_var);

                // min = a[ai] AND b[bi]
                self.clauses
                    .push(CardinalityClause::new([Lit::neg(min_var), a[ai]]));
                self.clauses
                    .push(CardinalityClause::new([Lit::neg(min_var), b[bi]]));
                self.clauses.push(CardinalityClause::new([
                    a[ai].negate(),
                    b[bi].negate(),
                    min_lit,
                ]));

                result[out_idx] = min_lit;
                ai += 1;
                bi += 1;
            }
            out_idx += 1;
        }

        result.truncate(n);
        result
    }

    /// Take the generated clauses
    pub fn take_clauses(&mut self) -> Vec<CardinalityClause> {
        std::mem::take(&mut self.clauses)
    }

    /// Get the number of clauses generated
    pub fn num_clauses(&self) -> usize {
        self.clauses.len()
    }

    /// Check if a number is a power of two
    fn is_power_of_two(n: usize) -> bool {
        n != 0 && (n & (n - 1)) == 0
    }

    /// Get the greatest power of two less than n
    fn greatest_power_of_two_less_than(n: usize) -> usize {
        if n <= 1 {
            return 1;
        }
        let mut k = 1;
        while k * 2 < n {
            k *= 2;
        }
        k
    }
}

/// Build an at-most-k constraint using a cardinality network
///
/// Returns clauses and optionally an assumption literal
pub fn encode_at_most_k_cardinality_network(
    inputs: &[Lit],
    k: usize,
    encoding: CardinalityNetworkEncoding,
    next_var: u32,
) -> (Vec<CardinalityClause>, Option<Lit>, u32) {
    let mut network = SortingNetwork::new(next_var, encoding);

    match encoding {
        CardinalityNetworkEncoding::Direct if k == 1 => {
            let _outputs = network.build_direct_amo(inputs);
            let clauses = network.take_clauses();
            // The constraint is: at most one is true, which means "at least one" can be false
            // We already encoded this in the pairwise constraints
            (clauses, None, network.next_var())
        }
        CardinalityNetworkEncoding::Sorting => {
            let outputs = network.build_sorting_network(inputs);
            let clauses = network.take_clauses();
            // at-most-k means: outputs[k+1] must be false (if it exists)
            let assumption = if k + 1 < outputs.len() {
                Some(outputs[k + 1].negate())
            } else {
                None
            };
            (clauses, assumption, network.next_var())
        }
        CardinalityNetworkEncoding::Cardinality
        | CardinalityNetworkEncoding::Direct
        | CardinalityNetworkEncoding::OddEvenMerge => {
            let outputs = network.build_cardinality_network(inputs, k);
            let clauses = network.take_clauses();
            // at-most-k means: outputs[k] must be false (if it exists)
            let assumption = if k < outputs.len() {
                Some(outputs[k].negate())
            } else {
                None
            };
            (clauses, assumption, network.next_var())
        }
        CardinalityNetworkEncoding::Grouped | CardinalityNetworkEncoding::Bimander => {
            // For now, fall back to cardinality network
            // These are more complex hybrid encodings
            let outputs = network.build_cardinality_network(inputs, k);
            let clauses = network.take_clauses();
            let assumption = if k < outputs.len() {
                Some(outputs[k].negate())
            } else {
                None
            };
            (clauses, assumption, network.next_var())
        }
    }
}

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

    fn var(v: u32) -> Var {
        Var(v)
    }

    fn pos(v: u32) -> Lit {
        Lit::pos(var(v))
    }

    #[test]
    fn test_sorting_network_empty() {
        let mut network = SortingNetwork::new(10, CardinalityNetworkEncoding::Sorting);
        let outputs = network.build_sorting_network(&[]);
        assert!(outputs.is_empty());
    }

    #[test]
    fn test_sorting_network_single() {
        let mut network = SortingNetwork::new(10, CardinalityNetworkEncoding::Sorting);
        let inputs = vec![pos(0)];
        let outputs = network.build_sorting_network(&inputs);
        assert_eq!(outputs.len(), 1);
        assert_eq!(outputs[0], pos(0));
    }

    #[test]
    fn test_sorting_network_two() {
        let mut network = SortingNetwork::new(10, CardinalityNetworkEncoding::Sorting);
        let inputs = vec![pos(0), pos(1)];
        let outputs = network.build_sorting_network(&inputs);
        assert_eq!(outputs.len(), 2);
        // Should have generated comparison clauses
        assert!(network.num_clauses() > 0);
    }

    #[test]
    fn test_sorting_network_four() {
        let mut network = SortingNetwork::new(10, CardinalityNetworkEncoding::Sorting);
        let inputs = vec![pos(0), pos(1), pos(2), pos(3)];
        let outputs = network.build_sorting_network(&inputs);
        assert_eq!(outputs.len(), 4);
        assert!(network.num_clauses() > 0);
    }

    #[test]
    fn test_direct_amo() {
        let mut network = SortingNetwork::new(10, CardinalityNetworkEncoding::Direct);
        let inputs = vec![pos(0), pos(1), pos(2)];
        let outputs = network.build_direct_amo(&inputs);
        assert_eq!(outputs.len(), 1); // One output representing "at least one"
        // Should have C(3,2) = 3 pairwise clauses + definition clauses
        assert!(network.num_clauses() >= 3);
    }

    #[test]
    fn test_cardinality_network() {
        let mut network = SortingNetwork::new(10, CardinalityNetworkEncoding::Cardinality);
        let inputs = vec![pos(0), pos(1), pos(2), pos(3)];
        let outputs = network.build_cardinality_network(&inputs, 2);
        // Should produce outputs for checking at-most-2
        assert!(!outputs.is_empty());
        assert!(network.num_clauses() > 0);
    }

    #[test]
    fn test_encode_at_most_k_direct() {
        let inputs = vec![pos(0), pos(1), pos(2)];
        let (clauses, _, _) = encode_at_most_k_cardinality_network(
            &inputs,
            1,
            CardinalityNetworkEncoding::Direct,
            10,
        );
        assert!(clauses.len() >= 3); // At least pairwise clauses
    }

    #[test]
    fn test_encode_at_most_k_sorting() {
        let inputs = vec![pos(0), pos(1), pos(2), pos(3)];
        let (clauses, assumption, _) = encode_at_most_k_cardinality_network(
            &inputs,
            2,
            CardinalityNetworkEncoding::Sorting,
            10,
        );
        assert!(!clauses.is_empty());
        assert!(assumption.is_some());
    }

    #[test]
    fn test_power_of_two() {
        assert!(SortingNetwork::is_power_of_two(1));
        assert!(SortingNetwork::is_power_of_two(2));
        assert!(SortingNetwork::is_power_of_two(4));
        assert!(SortingNetwork::is_power_of_two(8));
        assert!(!SortingNetwork::is_power_of_two(3));
        assert!(!SortingNetwork::is_power_of_two(5));
        assert!(!SortingNetwork::is_power_of_two(0));
    }

    #[test]
    fn test_greatest_power_of_two() {
        assert_eq!(SortingNetwork::greatest_power_of_two_less_than(7), 4);
        assert_eq!(SortingNetwork::greatest_power_of_two_less_than(8), 4);
        assert_eq!(SortingNetwork::greatest_power_of_two_less_than(9), 8);
        assert_eq!(SortingNetwork::greatest_power_of_two_less_than(15), 8);
        assert_eq!(SortingNetwork::greatest_power_of_two_less_than(16), 8);
    }
}
