pub use coord_2d::{Coord, Size};
use grid_2d::Grid;
use image::{DynamicImage, Rgba, RgbaImage};
use rand::{Rng, SeedableRng};
use std::num::NonZeroU32;
use wfc::orientation::OrientationTable;
pub use wfc::orientation::{self, Orientation};
use wfc::overlapping::{OverlappingPatterns, Pattern};
use wfc::retry as wfc_retry;
pub use wfc::wrap;
pub use wfc::ForbidNothing;
use wfc::*;
pub use wrap::WrapXY;

pub mod retry {
    #[cfg(feature = "parallel")]
    pub use super::wfc_retry::ParNumTimes;
    pub use super::wfc_retry::RetryOwn as Retry;
    pub use super::wfc_retry::{Forever, NumTimes};

    pub trait ImageRetry: Retry {
        type ImageReturn;
        #[doc(hidden)]
        fn image_return(
            r: Self::Return,
            image_patterns: &super::ImagePatterns,
        ) -> Self::ImageReturn;
    }
}

pub struct ImagePatterns {
    overlapping_patterns: OverlappingPatterns<Rgba<u8>>,
    empty_colour: Rgba<u8>,
}

impl ImagePatterns {
    pub fn new(
        image: &DynamicImage,
        pattern_size: NonZeroU32,
        orientations: &[Orientation],
    ) -> Self {
        let rgba_image = image.to_rgba();
        let size = Size::new(rgba_image.width(), rgba_image.height());
        let grid = Grid::new_fn(size, |Coord { x, y }| {
            *rgba_image.get_pixel(x as u32, y as u32)
        });
        let overlapping_patterns =
            OverlappingPatterns::new(grid, pattern_size, orientations);
        Self {
            overlapping_patterns,
            empty_colour: Rgba([0, 0, 0, 0]),
        }
    }

    pub fn set_empty_colour(&mut self, empty_colour: Rgba<u8>) {
        self.empty_colour = empty_colour;
    }

    pub fn image_from_wave(&self, wave: &Wave) -> DynamicImage {
        let size = wave.grid().size();
        let mut rgba_image = RgbaImage::new(size.width(), size.height());
        wave.grid().enumerate().for_each(|(Coord { x, y }, cell)| {
            let colour = match cell.chosen_pattern_id() {
                Ok(pattern_id) => {
                    *self.overlapping_patterns.pattern_top_left_value(pattern_id)
                }
                Err(_) => self.empty_colour,
            };
            rgba_image.put_pixel(x as u32, y as u32, colour);
        });
        DynamicImage::ImageRgba8(rgba_image)
    }

    pub fn weighted_average_colour<'a>(&self, cell: &'a WaveCellRef<'a>) -> Rgba<u8> {
        use wfc::EnumerateCompatiblePatternWeights::*;
        match cell.enumerate_compatible_pattern_weights() {
            MultipleCompatiblePatternsWithoutWeights | NoCompatiblePattern => {
                self.empty_colour
            }
            SingleCompatiblePatternWithoutWeight(pattern_id) => {
                *self.overlapping_patterns.pattern_top_left_value(pattern_id)
            }
            CompatiblePatternsWithWeights(iter) => {
                let (r, g, b, a) = iter
                    .map(|(pattern_id, weight)| {
                        let &Rgba([r, g, b, a]) =
                            self.overlapping_patterns.pattern_top_left_value(pattern_id);
                        (
                            r as u32 * weight,
                            g as u32 * weight,
                            b as u32 * weight,
                            a as u32 * weight,
                        )
                    })
                    .fold(
                        (0, 0, 0, 0),
                        |(acc_r, acc_g, acc_b, acc_a), (r, g, b, a)| {
                            (acc_r + r, acc_g + g, acc_b + b, acc_a + a)
                        },
                    );
                let total_weight = cell.sum_compatible_pattern_weight();
                Rgba([
                    (r / total_weight) as u8,
                    (g / total_weight) as u8,
                    (b / total_weight) as u8,
                    (a / total_weight) as u8,
                ])
            }
        }
    }

    pub fn grid(&self) -> &Grid<Rgba<u8>> {
        self.overlapping_patterns.grid()
    }

    pub fn id_grid(&self) -> Grid<OrientationTable<PatternId>> {
        self.overlapping_patterns.id_grid()
    }

    pub fn id_grid_original_orientation(&self) -> Grid<PatternId> {
        self.overlapping_patterns.id_grid_original_orientation()
    }

    pub fn pattern(&self, pattern_id: PatternId) -> &Pattern {
        self.overlapping_patterns.pattern(pattern_id)
    }

    pub fn pattern_mut(&mut self, pattern_id: PatternId) -> &mut Pattern {
        self.overlapping_patterns.pattern_mut(pattern_id)
    }

    pub fn global_stats(&self) -> GlobalStats {
        self.overlapping_patterns.global_stats()
    }

    pub fn collapse_wave_retrying<W, F, RT, R>(
        &self,
        output_size: Size,
        wrap: W,
        forbid: F,
        retry: RT,
        rng: &mut R,
    ) -> RT::Return
    where
        W: Wrap,
        F: ForbidPattern + Send + Sync + Clone,
        RT: retry::Retry,
        R: Rng + Send + Sync + Clone,
    {
        let global_stats = self.global_stats();
        let run = RunOwn::new_wrap_forbid(output_size, &global_stats, wrap, forbid, rng);
        run.collapse_retrying(retry, rng)
    }
}

impl retry::ImageRetry for retry::Forever {
    type ImageReturn = DynamicImage;
    fn image_return(
        r: Self::Return,
        image_patterns: &ImagePatterns,
    ) -> Self::ImageReturn {
        image_patterns.image_from_wave(&r)
    }
}

impl retry::ImageRetry for retry::NumTimes {
    type ImageReturn = Result<DynamicImage, PropagateError>;
    fn image_return(
        r: Self::Return,
        image_patterns: &ImagePatterns,
    ) -> Self::ImageReturn {
        match r {
            Ok(r) => Ok(image_patterns.image_from_wave(&r)),
            Err(e) => Err(e),
        }
    }
}

#[cfg(feature = "parallel")]
impl retry::ImageRetry for retry::ParNumTimes {
    type ImageReturn = Result<DynamicImage, PropagateError>;
    fn image_return(
        r: Self::Return,
        image_patterns: &ImagePatterns,
    ) -> Self::ImageReturn {
        match r {
            Ok(r) => Ok(image_patterns.image_from_wave(&r)),
            Err(e) => Err(e),
        }
    }
}

pub fn generate_image_with_rng<W, F, IR, R>(
    image: &DynamicImage,
    pattern_size: NonZeroU32,
    output_size: Size,
    orientations: &[Orientation],
    wrap: W,
    forbid: F,
    retry: IR,
    rng: &mut R,
) -> IR::ImageReturn
where
    W: Wrap,
    F: ForbidPattern + Send + Sync + Clone,
    IR: retry::ImageRetry,
    R: Rng + Send + Sync + Clone,
{
    let image_patterns = ImagePatterns::new(image, pattern_size, orientations);
    IR::image_return(
        image_patterns.collapse_wave_retrying(output_size, wrap, forbid, retry, rng),
        &image_patterns,
    )
}

pub fn generate_image<W, F, IR>(
    image: &DynamicImage,
    pattern_size: NonZeroU32,
    output_size: Size,
    orientations: &[Orientation],
    wrap: W,
    forbid: F,
    retry: IR,
) -> IR::ImageReturn
where
    W: Wrap,
    F: ForbidPattern + Send + Sync + Clone,
    IR: retry::ImageRetry,
{
    generate_image_with_rng(
        image,
        pattern_size,
        output_size,
        orientations,
        wrap,
        forbid,
        retry,
        &mut rand::rngs::StdRng::from_entropy(),
    )
}