nonogrid 0.7.3

use std::{
    fmt,
    ops::{Add, Sub},
};

use hashbrown::HashSet;
use log::debug;

use crate::{
    block::base::{
        color::{ColorId, ColorPalette},
        Block, Color,
    },
    utils::{from_two_powers, two_powers},
};

#[derive(Debug, PartialEq, Eq, Hash, Default, Copy, Clone)]
pub struct MultiColor(pub ColorId);

impl Color for MultiColor {
    fn blank() -> Self {
        Self(ColorPalette::WHITE_ID)
    }

    fn is_solved(self) -> bool {
        self.0.is_power_of_two()
    }

    fn memoize_rate() -> bool {
        true
    }

    /// Calculate the rate of the given cell.
    /// The formula is like that:
    ///   `rate = (N - n) / (N - 1)`, where
    ///    N = full puzzle color set
    ///    n = current color set for given cell,
    ///
    ///    in particular:
    ///    a) when the cell is completely unsolved
    ///       rate = (N - N) / (N - 1) = 0
    ///    b) when the cell is solved
    ///       rate = (N - 1) / (N - 1) = 1
    fn solution_rate(self, all_colors: &[ColorId]) -> f64 {
        let all_colors: HashSet<_> = all_colors.iter().copied().collect();
        let cell_colors = self.variants_as_ids();
        let current_size = cell_colors.intersection(&all_colors).count();

        if current_size == 0 {
            return 0.0;
        }
        if current_size == 1 {
            return 1.0;
        }

        let full_size = all_colors.len();
        let rate = full_size - current_size;
        #[allow(clippy::cast_precision_loss)]
        let normalized_rate = rate as f64 / (full_size - 1) as f64;
        assert!((0.0..=1.0).contains(&normalized_rate));

        normalized_rate
    }

    fn variants(self) -> Vec<Self> {
        two_powers(self.0).map(Self).collect()
    }

    fn as_color_id(self) -> Option<ColorId> {
        Some(self.0)
    }

    fn from_color_ids(ids: &[ColorId]) -> Self {
        Self(from_two_powers(ids))
    }
}

impl MultiColor {
    fn variants_as_ids(self) -> HashSet<ColorId> {
        two_powers(self.0).collect()
    }
}

impl Add for MultiColor {
    type Output = Self;

    fn add(self, rhs: Self) -> Self::Output {
        rhs
    }
}

impl Sub for MultiColor {
    type Output = Result<Self, String>;

    fn sub(self, rhs: Self) -> Self::Output {
        if self.is_solved() {
            return Err(format!(
                "Cannot unset {:?} from already set cell {:?}",
                rhs, self
            ));
        }

        let colors = self.variants_as_ids();
        let bad_state = rhs.variants_as_ids();
        debug!("Previous state: {:?}", colors);
        debug!("Bad state: {:?}", bad_state);

        let new_value: HashSet<_> = colors.difference(&bad_state).copied().collect();

        if !new_value.is_empty() && new_value.is_subset(&colors) {
            let new_state: Vec<_> = new_value.into_iter().collect();
            debug!("New state: {:?}", new_state);
            Ok(Self(from_two_powers(&new_state)))
        } else {
            Err(format!(
                "Cannot unset the colors {:?} from {:?}",
                bad_state, colors
            ))
        }
    }
}

impl fmt::Display for MultiColor {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let colors = self.variants_as_ids();
        if colors.len() == 1 {
            let color = colors
                .into_iter()
                .next()
                .expect("There should be a value: the length is 1");
            write!(f, "{}", color)
        } else {
            write!(f, "?")
        }
    }
}

#[derive(Debug, PartialEq, Eq, Hash, Default, Clone, Copy)]
pub struct ColoredBlock {
    size: usize,
    color: ColorId,
}

impl ColoredBlock {
    pub const fn from_size_and_color(size: usize, color: ColorId) -> Self {
        Self { size, color }
    }
}

impl Block for ColoredBlock {
    type Color = MultiColor;

    fn from_size_and_color(size: usize, color: Option<ColorId>) -> Self {
        let color = color.expect("Color not provided for ColoredBlock");
        Self { size, color }
    }

    fn partial_sums(desc: &[Self]) -> Vec<usize> {
        desc.iter()
            .scan(None, |acc_block: &mut Option<Self>, block| {
                let prev_sum = acc_block.map_or(0, |acc_block| {
                    // 1 cell is for a minimal gap between the previous run of blocks
                    // and the current block
                    let gap_size = if acc_block.color() == block.color() {
                        1
                    } else {
                        0
                    };
                    acc_block.size() + gap_size
                });

                let current = prev_sum + block.size();
                *acc_block = Some(Self::from_size_and_color(current, block.color));
                Some(current)
            })
            .collect()
    }

    fn size(self) -> usize {
        self.size
    }

    fn color(self) -> Self::Color {
        MultiColor(self.color)
    }
}

impl fmt::Display for ColoredBlock {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.size)
    }
}

#[cfg(test)]
mod tests {
    use crate::block::{Block, Description};

    use super::ColoredBlock;

    #[test]
    fn partial_sums_empty() {
        let d = Description::new(vec![]);
        assert_eq!(ColoredBlock::partial_sums(&d.vec), Vec::<usize>::new());
    }

    #[test]
    fn partial_sums_single() {
        let d = Description::new(vec![ColoredBlock::from_size_and_color(5, 1)]);
        assert_eq!(ColoredBlock::partial_sums(&d.vec), vec![5]);
    }

    #[test]
    fn check_partial_sums() {
        let d = Description::new(vec![
            ColoredBlock::from_size_and_color(1, 1),
            ColoredBlock::from_size_and_color(2, 1),
            ColoredBlock::from_size_and_color(3, 2),
        ]);
        assert_eq!(ColoredBlock::partial_sums(&d.vec), vec![1, 4, 7]);
    }
}