use std::collections::HashSet;
use std::sync::atomic::{AtomicBool, Ordering};
use crate::cyclotomic::{IsRing, Units};
use crate::enumerate::boundary::Boundary;
use crate::enumerate::canonical::CanonicalOps;
use crate::geom::rat::Rat;
use crate::enumerate::prune::Prunes;
use crate::enumerate::stats::DfsStats;
use crate::enumerate::stream::progress::{FLUSH_MASK, WorkerCell, odometer_fraction};
pub static STREAM_RAT_LINES: AtomicBool = AtomicBool::new(true);
#[allow(clippy::too_many_arguments)]
pub fn rat_enum_with<ZZ, B, Mk>(
mk: Mk,
max_steps: usize,
step: i8,
ops: CanonicalOps,
label: &str,
prefix: &str,
paranoid: bool,
prunes: &Prunes,
) -> (Vec<Vec<i8>>, DfsStats)
where
ZZ: IsRing,
B: Boundary<ZZ>,
Mk: Fn(&[i8]) -> B,
{
let mut result: HashSet<Vec<i8>> = HashSet::new();
let mut b = mk(&[]);
let mut stats = DfsStats::default();
println!("-------- {label} started --------");
if paranoid {
println!("paranoid: per-step fresh-snake cross-check enabled");
}
{
let mut record = hashset_recorder(&mut result);
rat_enum_step::<ZZ, B>(
&mut b,
max_steps,
step,
&mut record,
&mut stats,
ops,
paranoid,
prunes,
None,
usize::MAX,
&mut Vec::new(),
);
}
println!(
"-------- {label} completed --------\n{prefix}{} rats found",
result.len()
);
let mut result: Vec<Vec<i8>> = result.into_iter().collect();
result.sort_by_key(|x| x.len());
(result, stats)
}
pub fn hashset_recorder<'a>(set: &'a mut HashSet<Vec<i8>>) -> impl FnMut(&[i8]) + 'a {
move |seq: &[i8]| {
if set.insert(seq.to_vec()) && STREAM_RAT_LINES.load(Ordering::Relaxed) {
println!("RAT {seq:?}");
}
}
}
#[inline]
fn publish_progress<ZZ: IsRing>(cell: &WorkerCell, angles: &[i8], stats: &DfsStats, step: i8) {
let seed_len = cell.seed_len.load(Ordering::Relaxed) as usize;
let rel: &[i8] = angles.get(seed_len..).unwrap_or(&[]);
let ppm = ((odometer_fraction(rel, ZZ::hturn(), step) * 1_000_000.0) as u32).min(999_999);
cell.publish(stats.total(), stats.closed, angles.len() as u32, ppm);
}
#[inline]
fn maybe_publish<ZZ: IsRing>(
progress: Option<&WorkerCell>,
key: u64,
angles: &[i8],
stats: &DfsStats,
step: i8,
) {
if let Some(cell) = progress
&& key & FLUSH_MASK == 0
{
publish_progress::<ZZ>(cell, angles, stats, step);
}
}
#[allow(clippy::too_many_arguments)]
pub fn rat_enum_step<ZZ: IsRing, B: Boundary<ZZ>>(
b: &mut B,
max_steps: usize,
step: i8,
record: &mut dyn FnMut(&[i8]),
stats: &mut DfsStats,
ops: CanonicalOps,
paranoid: bool,
prunes: &Prunes,
progress: Option<&WorkerCell>,
split_depth: usize,
seeds: &mut Vec<Vec<i8>>,
) {
let depth = b.angles().len();
if depth >= max_steps {
return;
}
let remaining = (max_steps - depth) as i64;
for direction in ((-ZZ::hturn() + 1)..ZZ::hturn()).rev() {
if direction.rem_euclid(step) != 0 {
continue;
}
if !(ops.is_canonical)(b.angles(), direction) {
stats.canonical_skip += 1;
continue;
}
let new_pt =
b.offset() + <ZZ as Units>::unit(b.direction()) * <ZZ as Units>::unit(direction);
if !new_pt.is_zero() && !new_pt.within_radius(remaining) {
stats.too_far += 1;
continue;
}
if let Some(sp) = prunes.shadow_prune.as_deref()
&& !new_pt.is_zero()
&& !sp.allows_closure(&new_pt, remaining)
{
stats.shadow_skip += 1;
continue;
}
let remaining_after = (remaining as usize).saturating_sub(1);
if let Some(mp) = prunes.modular_prune.as_deref()
&& !mp.allows_closure(new_pt.int_coeffs_slice(), remaining_after)
{
stats.modular_skip += 1;
continue;
}
if let Some(ck) = prunes.closure_table_prune.as_deref()
&& remaining_after <= ck.max_l
{
let turn = ZZ::turn();
let new_facing = (b.direction() + direction).rem_euclid(turn);
let neg_facing = (-new_facing).rem_euclid(turn);
let target: ZZ = -(<ZZ as Units>::unit(neg_facing) * new_pt);
let key = (target.int_coeffs_slice().to_vec(), neg_facing);
if !ck.keys.contains(&key) {
stats.closure_table_skip += 1;
continue;
}
}
if !b.add(direction) {
stats.intersected += 1;
continue;
}
if paranoid {
b.paranoid_recheck();
}
if b.is_closed() {
stats.closed += 1;
maybe_publish::<ZZ>(progress, stats.closed, b.angles(), stats, step);
let r = {
let tmp = Rat::<ZZ>::from_slice_trusted(b.angles());
if tmp.chirality() > 0 {
tmp
} else {
tmp.reversed()
}
.canonical()
};
let seq = (ops.canonicalize)(r.seq());
record(&seq);
} else {
stats.recursed += 1;
maybe_publish::<ZZ>(progress, stats.recursed, b.angles(), stats, step);
if b.angles().len() >= split_depth {
seeds.push(b.angles().to_vec());
} else {
rat_enum_step::<ZZ, B>(
b,
max_steps,
step,
record,
stats,
ops,
paranoid,
prunes,
progress,
split_depth,
seeds,
);
}
}
b.pop();
}
}