sat.rs

//! The algorithm based on the ideas
//! from the [article](https://habr.com/ru/post/433330/)
#![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> {
        // 1. Каждый блок, объявленный в строке или столбце обязан появиться хотя-бы в одной позиции.
        // Этому соответствует клоз вида (X1 V X2 V… XN),
        // где X1, X2… XN — все возможные позиции данного блока в строке или столбце.
        at_least_one(positions.vars_iter())
    }

    fn block_once_clauses(positions: &BlockPositions) -> impl Iterator<Item = Vec<Lit>> {
        // 2. Каждый блок в строке или столбце должен появиться не более одного раза.
        // Этому соответствует множество клозов вида (not Xi) V (not Xj),
        // где Xi, Xj (i != j) — все возможные позиции данного блока в строке или столбце.
        let vars: Vec<_> = positions.vars_iter().collect();
        at_most_one(&vars)
    }

    fn non_overlap_clauses(positions: &LinePositions) -> impl Iterator<Item = Vec<Lit>> {
        // 3. Правильный порядок блоков. Поскольку необходимо поддерживать правильный порядок расположения блоков,
        // а также исключить их пересечение, необходимо добавить клозы, вида (not Xi) V (not Xj),
        // где Xi, Xj — переменные, соответствующие позициям разных блоков, которые имеют неправильный порядок или пересекаются.
        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 {
                        // conflict encoding (!block1_position V !block2_position)
                        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>> {
        // 4. Окрашенная клетка должна содержаться внутри хотя бы одного блока, позиция которого включает данную клетку.
        // Этому соответствует клоз вида ((not Yk) V X1 V X2… XN), где Yk — переменная, соответствующая клетке,
        // а X1, X2… XN — переменные, соответствующие позициям блоков, содержащих данную клетку.
        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);

                // support encoding (!color_var V position1 V position2 V ...)
                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>> {
        // 5. Каждая пустая клетка не должна содержаться ни в одной возможной позиции ни одного блока.
        // Этому соответствует множество клозов вида Yi V (not Xj), где Yi — переменная, соответствующая клетке,
        // а Xj — переменная, соответствующая одной позиции какого-либо блока, содержащая данную клетку.
        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)
                    // conflict encoding (!white V !block_position === color V !block_position)
                    .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 {
                        // blank cell
                        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>> + '_ {
        // (X -> Y) == (-X \/ Y)
        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 {
            // 'not blank' equivalent to 'at least one color is set'
            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 {
                    // TODO
                    info!(
                        "Not found impact var for color {:?} in point {:?}. Available vars: {:?}",
                        color, point, impact_vars
                    );
                    vec![]
                }
            } else {
                // 'blank impact' means
                // (probe V -impact_color1) /\ (probe V -impact_color2) /\ ...
                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()
        })
    }
}