//! Exhaustive small census of covering and irreducible modulus sets.
//! This is a generator/research program; accepted witnesses are rechecked by
//! the separate implementation in `check_answer/`.

use rayon::prelude::*;
use std::env;
use std::sync::atomic::{AtomicU64, Ordering};

fn gcd(mut a: usize, mut b: usize) -> usize {
    while b != 0 {
        (a, b) = (b, a % b);
    }
    a
}

fn lcm(a: usize, b: usize) -> Option<usize> {
    (a / gcd(a, b)).checked_mul(b)
}

fn period(ns: &[usize]) -> Option<usize> {
    ns.iter().try_fold(1, |p, &n| lcm(p, n))
}

#[derive(Default)]
struct SearchStats {
    nodes: u64,
}

struct State<'a> {
    ns: &'a [usize],
    p: usize,
    counts: Vec<u8>,
    used: Vec<bool>,
    chosen: Vec<Option<usize>>,
    stats: SearchStats,
}

impl State<'_> {
    fn add(&mut self, i: usize, r: usize) -> usize {
        let mut gain = 0;
        for x in (r..self.p).step_by(self.ns[i]) {
            gain += usize::from(self.counts[x] == 0);
            self.counts[x] += 1;
        }
        gain
    }

    fn remove(&mut self, i: usize, r: usize) {
        for x in (r..self.p).step_by(self.ns[i]) {
            self.counts[x] -= 1;
        }
    }

    fn solve(&mut self, missing: usize) -> bool {
        self.stats.nodes += 1;
        if missing == 0 {
            return true;
        }
        let capacity: usize = self
            .ns
            .iter()
            .enumerate()
            .filter(|(i, _)| !self.used[*i])
            .map(|(_, n)| self.p / n)
            .sum();
        if capacity < missing {
            return false;
        }
        let x = self.counts.iter().position(|&c| c == 0).unwrap();
        let mut options: Vec<(usize, usize)> = self
            .ns
            .iter()
            .enumerate()
            .filter(|(i, _)| !self.used[*i])
            .map(|(i, &n)| {
                let r = x % n;
                let gain = (r..self.p)
                    .step_by(n)
                    .filter(|&y| self.counts[y] == 0)
                    .count();
                (i, gain)
            })
            .collect();
        options.sort_unstable_by_key(|&(_, gain)| std::cmp::Reverse(gain));
        for (i, _) in options {
            let r = x % self.ns[i];
            self.used[i] = true;
            self.chosen[i] = Some(r);
            let gain = self.add(i, r);
            if self.solve(missing - gain) {
                return true;
            }
            self.remove(i, r);
            self.used[i] = false;
            self.chosen[i] = None;
        }
        false
    }
}

fn find_cover(ns: &[usize]) -> Option<(Vec<usize>, u64)> {
    if ns.is_empty() {
        return None;
    }
    let p = period(ns)?;
    if ns.iter().map(|&n| p / n).sum::<usize>() < p {
        return None;
    }
    let mut s = State {
        ns,
        p,
        counts: vec![0; p],
        used: vec![false; ns.len()],
        chosen: vec![None; ns.len()],
        stats: SearchStats::default(),
    };
    // Translate a hypothetical cover so that its first class is 0 mod ns[0].
    s.used[0] = true;
    s.chosen[0] = Some(0);
    let gain = s.add(0, 0);
    if s.solve(p - gain) {
        Some((s.chosen.into_iter().map(|r| r.unwrap_or(0)).collect(), s.stats.nodes))
    } else {
        None
    }
}

fn irreducible(ns: &[usize]) -> Option<(Vec<usize>, u64)> {
    let (residues, mut nodes) = find_cover(ns)?;
    for omitted in 0..ns.len() {
        let sub: Vec<usize> = ns
            .iter()
            .enumerate()
            .filter(|(i, _)| *i != omitted)
            .map(|(_, &n)| n)
            .collect();
        if let Some((_, subnodes)) = find_cover(&sub) {
            nodes += subnodes;
            return None;
        }
    }
    Some((residues, nodes))
}

fn combinations(lo: usize, hi: usize, k: usize) -> Vec<Vec<usize>> {
    fn rec(next: usize, hi: usize, left: usize, acc: &mut Vec<usize>, out: &mut Vec<Vec<usize>>) {
        if left == 0 {
            out.push(acc.clone());
            return;
        }
        if next > hi || hi - next + 1 < left {
            return;
        }
        for n in next..=hi - left + 1 {
            acc.push(n);
            rec(n + 1, hi, left - 1, acc, out);
            acc.pop();
        }
    }
    let mut out = Vec::new();
    rec(lo, hi, k, &mut Vec::with_capacity(k), &mut out);
    out
}

fn main() {
    let args: Vec<String> = env::args().collect();
    if args.len() != 3 {
        eprintln!("usage: {} K MAX_MODULUS", args[0]);
        std::process::exit(2);
    }
    let k: usize = args[1].parse().expect("K must be a positive integer");
    let max_n: usize = args[2].parse().expect("MAX_MODULUS must be an integer");
    let candidates = combinations(2, max_n, k);
    eprintln!("testing {} size-{k} sets from [2,{max_n}]", candidates.len());
    let total_nodes = AtomicU64::new(0);
    let mut answers: Vec<_> = candidates
        .par_iter()
        .filter_map(|ns| {
            let result = irreducible(ns);
            if let Some((_, nodes)) = &result {
                total_nodes.fetch_add(*nodes, Ordering::Relaxed);
            }
            result.map(|(a, nodes)| (ns.clone(), a, nodes))
        })
        .collect();
    answers.sort_unstable_by(|x, y| x.0.cmp(&y.0));
    for (ns, residues, nodes) in &answers {
        println!(
            "{} | {} | nodes={nodes}",
            ns.iter().map(ToString::to_string).collect::<Vec<_>>().join(","),
            residues.iter().map(ToString::to_string).collect::<Vec<_>>().join(",")
        );
    }
    eprintln!(
        "irreducible_count={} accepted_search_nodes={}",
        answers.len(),
        total_nodes.load(Ordering::Relaxed)
    );
}
