#![allow(clippy::option_if_let_else)]
use std::{
collections::{HashMap, HashSet},
iter::{from_fn, once},
ops::{Deref, Range},
};
use log::{debug, info, warn};
use varisat::{solver::Solver, CnfFormula, ExtendFormula, Lit, Var};
use crate::{
block::{base::color::ColorId, Block, Color, Description},
board::Point,
solver::probing::Impact,
utils::{pair_combinations, product, rc::ReadRc},
};
#[derive(Debug, Clone)]
struct Position {
var: Var,
range: Range<usize>,
}
impl Position {
fn var_if_included(&self, point: usize) -> Option<Var> {
if self.range.contains(&point) {
Some(self.var)
} else {
None
}
}
}
#[derive(Debug, Clone)]
struct BlockPositions {
index: usize,
color: ColorId,
vec: Vec<Position>,
}
impl Deref for BlockPositions {
type Target = [Position];
fn deref(&self) -> &Self::Target {
&self.vec
}
}
impl BlockPositions {
fn vars_iter(&self) -> impl Iterator<Item = Var> + '_ {
self.iter().map(|pos| pos.var)
}
}
#[derive(Debug)]
struct LinePositions(Vec<BlockPositions>);
impl Deref for LinePositions {
type Target = [BlockPositions];
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl LinePositions {
fn covering_vars(&self, color: ColorId, point: usize) -> Vec<Var> {
self.iter()
.filter(|block_pos| block_pos.color == color)
.flat_map(move |block_pos| {
block_pos
.iter()
.filter_map(move |pos| pos.var_if_included(point))
})
.collect()
}
}
#[derive(Debug)]
pub struct ClauseGenerator<B>
where
B: Block,
{
columns_vars: Vec<LinePositions>,
rows_vars: Vec<LinePositions>,
cell_vars: Vec<Vec<HashMap<ColorId, Var>>>,
cells: Vec<B::Color>,
width: usize,
height: usize,
}
fn at_least_one(vars: impl Iterator<Item = Var>) -> Vec<Lit> {
vars.map(Var::positive).collect()
}
fn at_most_one(vars: &[Var]) -> impl Iterator<Item = Vec<Lit>> + 'static {
let pairs = pair_combinations(vars);
pairs
.into_iter()
.map(|(f, s)| vec![f.negative(), s.negative()])
}
impl<B> ClauseGenerator<B>
where
B: Block,
B::Color: std::fmt::Debug,
{
const BLACK_COLOR: ColorId = 0;
pub fn with_clues(
columns: &[ReadRc<Description<B>>],
rows: &[ReadRc<Description<B>>],
cells: Vec<B::Color>,
) -> Self {
let mut block_colors: HashSet<_> = rows.iter().flat_map(|row| row.colors()).collect();
if block_colors.is_empty() {
let _ = block_colors.insert(Self::BLACK_COLOR);
}
let width = columns.len();
let height = rows.len();
let mut formula = CnfFormula::new();
let columns_vars = Self::clues_vars(columns, height, &mut formula);
let rows_vars = Self::clues_vars(rows, width, &mut formula);
let clues_vars = formula.var_count();
let cell_vars = cells
.chunks(width)
.map(|row| {
row.iter()
.map(|_cell| {
block_colors
.iter()
.map(|&color_id| (color_id, formula.new_var()))
.collect()
})
.collect()
})
.collect();
let vars_total = formula.var_count();
warn!(
"Vars: {} (clues: {}, cells: {})",
vars_total,
clues_vars,
vars_total - clues_vars
);
Self {
columns_vars,
rows_vars,
cell_vars,
cells,
width,
height,
}
}
fn clues_vars_count(&self) -> usize {
let col_vars: usize = self
.columns_vars
.iter()
.flat_map(|col| col.iter().map(|block| block.len()))
.sum();
let row_vars: usize = self
.rows_vars
.iter()
.flat_map(|col| col.iter().map(|block| block.len()))
.sum();
col_vars + row_vars
}
fn clues_vars(
clues: &[ReadRc<Description<B>>],
line_length: usize,
formula: &mut CnfFormula,
) -> Vec<LinePositions> {
clues
.iter()
.map(|clue| {
let positions = clue.positions_number(line_length);
LinePositions(
clue.block_starts()
.iter()
.zip(&clue.vec)
.enumerate()
.map(|(index, (&start, block))| BlockPositions {
index,
color: Self::get_id(block.color()).expect("Block without color!"),
vec: formula
.new_var_iter(positions)
.zip(start..)
.map(|(var, start)| Position {
var,
range: start..start + block.size(),
})
.collect(),
})
.collect(),
)
})
.collect()
}
fn block_positions_clause(positions: &BlockPositions) -> Vec<Lit> {
at_least_one(positions.vars_iter())
}
fn block_once_clauses(positions: &BlockPositions) -> impl Iterator<Item = Vec<Lit>> {
let vars: Vec<_> = positions.vars_iter().collect();
at_most_one(&vars)
}
fn non_overlap_clauses(positions: &LinePositions) -> impl Iterator<Item = Vec<Lit>> {
assert!(!positions.is_empty());
let pairs = pair_combinations(positions);
pairs.into_iter().flat_map(|(block1, block2)| {
assert!(block1.index < block2.index);
let pairs = product(&block1, &block2);
pairs.into_iter().filter_map(
move |(
Position {
range: range1,
var: var1,
},
Position {
range: range2,
var: var2,
},
)| {
let conflict = if block1.color == block2.color {
range1.end >= range2.start
} else {
range1.end > range2.start
};
if conflict {
Some(vec![var1.negative(), var2.negative()])
} else {
None
}
},
)
})
}
fn covering_positions(&self, cell_point: Point, color_id: ColorId) -> (Vec<Var>, Vec<Var>) {
let (x, y) = (cell_point.x, cell_point.y);
let column = &self.columns_vars[x];
let row = &self.rows_vars[y];
let column_vars = column.covering_vars(color_id, y);
let row_vars = row.covering_vars(color_id, x);
(column_vars, row_vars)
}
fn cell_color_clauses(&self, cell_point: Point) -> Vec<Vec<Lit>> {
let point_vars = self.get_vars(cell_point);
point_vars
.iter()
.flat_map(|(&color_id, color_var)| {
let (column_vars, row_vars) = self.covering_positions(cell_point, color_id);
let column_clause = column_vars
.into_iter()
.map(Var::positive)
.chain(once(color_var.negative()))
.collect();
let row_clause = row_vars
.into_iter()
.map(Var::positive)
.chain(once(color_var.negative()))
.collect();
once(column_clause).chain(once(row_clause))
})
.collect()
}
fn cell_space_clauses(&self, cell_point: Point) -> Vec<Vec<Lit>> {
let point_vars = self.get_vars(cell_point);
point_vars
.iter()
.flat_map(|(&color_id, color_var)| {
let (column_vars, row_vars) = self.covering_positions(cell_point, color_id);
column_vars
.into_iter()
.chain(row_vars)
.map(move |block_position_var| {
vec![block_position_var.negative(), color_var.positive()]
})
})
.collect()
}
fn get_id(color: B::Color) -> Option<ColorId> {
if color == Color::blank() {
return None;
}
color.as_color_id().or(Some(Self::BLACK_COLOR))
}
fn point_once_clauses(&self, cell_point: Point) -> impl Iterator<Item = Vec<Lit>> {
let point_vars = self.get_vars(cell_point);
let values: Vec<_> = point_vars.values().copied().collect();
at_most_one(&values)
}
fn precomputed_cells_clauses(&self) -> Vec<Lit> {
if self.cells.is_empty() {
return Vec::new();
}
self.cell_vars
.concat()
.iter()
.zip(&self.cells)
.flat_map(|(vars, &cell)| {
if cell.is_solved() {
let color_id = Self::get_id(cell);
if let Some(color_id) = color_id {
let var = vars.get(&color_id).expect("Solved color should be present");
vec![var.positive()]
} else {
vars.values().map(|var| var.negative()).collect()
}
} else {
let colors: Vec<_> = cell
.variants()
.into_iter()
.filter_map(Self::get_id)
.collect();
vars.iter()
.filter_map(|(color, var)| {
if colors.contains(color) {
None
} else {
Some(var.negative())
}
})
.collect()
}
})
.collect()
}
fn clauses(&self) -> impl Iterator<Item = Vec<Lit>> + '_ {
let columns_positions = self
.columns_vars
.iter()
.flat_map(|line_positions| line_positions.iter().map(Self::block_positions_clause));
let rows_positions = self
.rows_vars
.iter()
.flat_map(|line_positions| line_positions.iter().map(Self::block_positions_clause));
let columns_once_positions = self
.columns_vars
.iter()
.flat_map(|line_positions| line_positions.iter().flat_map(Self::block_once_clauses));
let rows_once_positions = self
.rows_vars
.iter()
.flat_map(|line_positions| line_positions.iter().flat_map(Self::block_once_clauses));
let non_overlap_columns = self
.columns_vars
.iter()
.filter(|line_pos| line_pos.len() > 1)
.flat_map(Self::non_overlap_clauses);
let non_overlap_rows = self
.rows_vars
.iter()
.filter(|line_pos| line_pos.len() > 1)
.flat_map(Self::non_overlap_clauses);
let all_points = product(
&(0..self.width).collect::<Vec<_>>(),
&(0..self.height).collect::<Vec<_>>(),
);
let color_clauses = all_points
.clone()
.into_iter()
.flat_map(move |(x, y)| self.cell_color_clauses(Point::new(x, y)));
let space_clauses = all_points
.clone()
.into_iter()
.flat_map(move |(x, y)| self.cell_space_clauses(Point::new(x, y)));
let point_once_clauses = all_points
.into_iter()
.flat_map(move |(x, y)| self.point_once_clauses(Point::new(x, y)));
let fixed_points = self
.precomputed_cells_clauses()
.into_iter()
.map(|lit| vec![lit]);
columns_positions
.chain(rows_positions)
.chain(columns_once_positions)
.chain(rows_once_positions)
.chain(non_overlap_columns)
.chain(non_overlap_rows)
.chain(color_clauses)
.chain(space_clauses)
.chain(point_once_clauses)
.chain(fixed_points)
}
fn get_formula(&self) -> CnfFormula {
debug!("{:#?}", self);
let mut formula = CnfFormula::new();
for clause in self.clauses() {
formula.add_clause(&clause);
}
formula
}
fn get_vars(&self, point: Point) -> &HashMap<ColorId, Var> {
self.cell_vars
.get(point.y)
.and_then(|row| row.get(point.x))
.expect("Cannot get vars for point")
}
fn generate_implications(
&self,
probe: Point,
probe_color: B::Color,
impact: Vec<(Point, B::Color)>,
) -> impl Iterator<Item = Vec<Lit>> + '_ {
let probe_vars = self.get_vars(probe);
let probe_id = Self::get_id(probe_color);
let probe_stmt = if let Some(probe_id) = probe_id {
let probe_var = probe_vars.get(&probe_id).expect("Not found probe var");
vec![probe_var.negative()]
} else {
probe_vars.values().map(|var| var.positive()).collect()
};
impact.into_iter().flat_map(move |(point, color)| {
let impact_vars = self.get_vars(point);
let impact_id = Self::get_id(color);
if let Some(impact_id) = impact_id {
if let Some(impact_var) = impact_vars.get(&impact_id) {
let res = probe_stmt
.clone()
.into_iter()
.chain(once(impact_var.positive()))
.collect();
debug!(
"Probe ({:?}, {:?}) -> ({:?}, {:?}):\n{:?}",
probe, probe_color, point, color, res
);
vec![res]
} else {
info!(
"Not found impact var for color {:?} in point {:?}. Available vars: {:?}",
color, point, impact_vars
);
vec![]
}
} else {
impact_vars
.values()
.map(|impact_var| {
probe_stmt
.clone()
.into_iter()
.chain(once(impact_var.negative()))
.collect()
})
.collect()
}
})
}
pub fn run(
&self,
probing_impact: Impact<B>,
solutions_number: Option<usize>,
) -> impl Iterator<Item = Vec<B::Color>> {
let mut formula = self.get_formula();
let total_impact_clauses: usize = probing_impact
.into_iter()
.map(|impact| {
let (probe, probe_color, impact, _) = impact.into_tuple();
self.generate_implications(probe, probe_color, impact)
.inspect(|clause| {
formula.add_clause(clause);
})
.count()
})
.sum();
warn!("Add {} impact clauses", total_impact_clauses);
warn!("Total clauses: {}", formula.len());
for (i, clause) in formula.iter().enumerate() {
info!("{}. {:?}", i, clause);
}
let block_vars = self.clues_vars_count();
let cell_vars = self.cell_vars.clone();
let mut solver = Solver::new();
solver.add_formula(&formula);
let mut found = 0;
from_fn(move || {
if let Some(solutions_number) = solutions_number {
if found >= solutions_number {
return None;
}
}
let _ = solver.solve().unwrap();
solver.model().map(|model| {
found += 1;
let cells = &model[block_vars..];
let colored_cells: Vec<_> = cells
.iter()
.filter_map(|cell| {
if cell.is_positive() {
Some(cell.var())
} else {
None
}
})
.collect();
let fixed_clause: Vec<_> =
colored_cells.iter().map(|&var| var.negative()).collect();
solver.add_clause(&fixed_clause);
colored_cells
})
})
.map(move |colored_vars| {
cell_vars
.iter()
.flat_map(|row| row.iter())
.map(|cell_map| {
let color_id = cell_map.iter().find_map(|(&color_id, var)| {
if colored_vars.contains(var) {
Some(color_id)
} else {
None
}
});
match color_id {
None => B::Color::blank(),
Some(Self::BLACK_COLOR) => B::Color::from_color_ids(&[]),
Some(color_id) => B::Color::from_color_ids(&[color_id]),
}
})
.collect()
})
}
}