//! A module for the core image structs and traits

/// A struct representing an image
#[derive(Debug, Clone, PartialEq)]
pub struct Image<T: Number> {
    info: ImageInfo,
    data: Vec<T>,
}

/// A struct representing a part of an image
#[derive(Debug, Clone)]
pub struct SubImage<'a, T: Number> {
    info: ImageInfo,
    data: Vec<&'a [T]>,
}

/// A struct containing image information
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct ImageInfo {
    pub width: u32,
    pub height: u32,
    pub channels: u8,
    pub alpha: bool,
}

/// A trait for valid image channel types
pub trait Number:
std::marker::Copy
+ std::fmt::Display
+ std::cmp::PartialEq
+ std::ops::Add<Output=Self>
+ std::ops::Sub<Output=Self>
+ std::ops::Mul<Output=Self>
+ std::ops::Div<Output=Self>
+ std::ops::AddAssign
+ std::ops::SubAssign
+ std::ops::MulAssign
+ std::ops::DivAssign
+ From<u8>
    where Self: std::marker::Sized {}

impl<T> Number for T
    where T:
    std::marker::Copy
    + std::fmt::Display
    + std::cmp::PartialEq
    + std::ops::Add<Output=T>
    + std::ops::Sub<Output=T>
    + std::ops::Mul<Output=T>
    + std::ops::Div<Output=T>
    + std::ops::AddAssign
    + std::ops::SubAssign
    + std::ops::MulAssign
    + std::ops::DivAssign
    + From<u8> {}

/// A trait for a base image
pub trait BaseImage<T: Number> {
    /// Returns the image information
    fn info(&self) -> ImageInfo;

    /// Returns a PixelSlice representing the pixel located at `(x, y)`
    fn get_pixel(&self, x: u32, y: u32) -> &[T];
}

/// A trait for image pixels
pub trait Pixel<T: Number> {
    /// Returns the last channel of the pixel
    fn alpha(&self) -> T;

    /// Returns the last channel of the pixel
    fn channels_without_alpha(&self) -> &[T];

    /// Applies function `f` to each channel
    fn map_all<S: Number, F>(&self, f: F) -> Vec<S>
        where F: Fn(T) -> S;

    /// Applies function `f` to each channel except the last channel, and applies
    /// function `g` to the alpha channel
    fn map_alpha<S: Number, F, G>(&self, f: F, g: G) -> Vec<S>
        where F: Fn(T) -> S,
              G: Fn(T) -> S;

    /// Applies function `f` to each channel
    fn apply<F>(&mut self, f: F)
        where F: Fn(T) -> T;

    /// Applies function `f` to each channel except the last channel, and applies
    /// function `g` to the alpha channel
    fn apply_alpha<F, G>(&mut self, f: F, g: G)
        where F: Fn(T) -> T,
              G: Fn(T) -> T;

    /// Returns true if all channel values are zero
    fn is_black(&self) -> bool;

    /// Returns true if all channel values except the last channel is zero
    fn is_black_alpha(&self) -> bool;
}

impl ImageInfo {
    /// Creates a new ImageInfo
    pub fn new(width: u32, height: u32, channels: u8, alpha: bool) -> Self {
        ImageInfo { width, height, channels, alpha }
    }

    /// Returns the width and height of the image
    pub fn wh(&self) -> (u32, u32) {
        (self.width, self.height)
    }

    /// Returns the width, height, and channels of the image
    pub fn whc(&self) -> (u32, u32, u8) {
        (self.width, self.height, self.channels)
    }

    /// Returns the width, height, channels, and alpha of the image
    pub fn whca(&self) -> (u32, u32, u8, bool) {
        (self.width, self.height, self.channels, self.alpha)
    }

    /// Returns the number of non alpha channels in the image
    pub fn channels_non_alpha(&self) -> u8 {
        if self.alpha {
            self.channels - 1
        } else {
            self.channels
        }
    }

    /// Returns the size of the image (width * height)
    pub fn size(&self) -> u32 {
        self.width * self.height
    }

    /// Returns the full size of the image (width * height * channels)
    pub fn full_size(&self) -> u32 {
        self.width * self.height * (self.channels as u32)
    }
}

impl std::fmt::Display for ImageInfo {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "width: {}\nheight: {}\nchannels: {}\nalpha: {}", self.width, self.height, self.channels, self.alpha)
    }
}

impl<T: Number> Image<T> {
    /// Creates a new `Image<T>`
    pub fn new(width: u32, height: u32, channels: u8, alpha: bool, data: &[T]) -> Self {
        Image {
            info: ImageInfo{ width, height, channels, alpha },
            data: data.to_vec(),
        }
    }

    /// Creates an `Image<T>` populated with zeroes
    pub fn blank(info: ImageInfo) -> Self {
        Image {
            info,
            data: vec![0.into(); info.full_size() as usize],
        }
    }

    /// Creates an empty `Image<T>`
    pub fn empty(info: ImageInfo) -> Self {
        Image {
            info,
            data: Vec::with_capacity(info.full_size() as usize),
        }
    }

    /// Returns the 1d index corresponding to the 2d `(x, y)` indices
    pub fn index(&self, x: u32, y: u32) -> usize {
        ((y * self.info.width + x) * self.info.channels as u32) as usize
    }

    /// Returns all data as a slice
    pub fn data(&self) -> &[T] {
        &self.data[..]
    }

    /// Returns all data as a mutable slice
    pub fn data_mut(&mut self) -> &mut [T] {
        &mut self.data[..]
    }

    /// Returns a PixelSliceMut representing the pixel located at `(x, y`
    pub fn get_pixel_mut(&mut self, x: u32, y: u32) -> &mut [T] {
        if x >= self.info.width {
            panic!("index out of bounds: the width is {}, but the x index is {}", self.info.width, x)
        }
        if y >= self.info.height {
            panic!("index out of bounds: the height is {}, but the y index is {}", self.info.height, y)
        }

        let start = self.index(x, y);
        &mut self.data[start..(start + self.info.channels as usize)]
    }

    /// Returns a `SubImage<T>` representing the part of the image of width `width` and height
    /// `height`, with upper left hand corner located at `(x, y)`
    pub fn get_subimage(&self, x: u32, y: u32, width: u32, height: u32) -> SubImage<T> {
        let mut data = Vec::new();

        for i in x..(x + width) {
            for j in y..(y + height) {
                data.push(self.get_pixel(i, j));
            }
        }

        SubImage::new(width, height, self.info.channels, self.info.alpha, data)
    }

    /// Returns a `SubImage<T>` representing the row or column of pixels of length `size` centered at
    /// `(x, y)`. If `is_vert` is `true`, returns the column; otherwise, returns the row.
    /// Uses clamp padding for edge pixels (edge pixels are repeated indefinitely)
    pub fn get_neighborhood_1d(&self, x: u32, y: u32, size: u32, is_vert: bool) -> SubImage<T> {
        let mut data = Vec::new();

        if is_vert {
            let start_y = (y as i32) - (size as i32) / 2;

            for i in 0..size {
                let mut curr_y = start_y + (i as i32);
                if curr_y < 0 || curr_y >= self.info.height as i32 { curr_y = y as i32 };

                data.push(self.get_pixel(x, curr_y as u32));
            }

            SubImage::new(1, size, self.info.channels, self.info.alpha, data)
        } else {
            let start_x = (x as i32) - (size as i32) / 2;

            for i in 0..size {
                let mut curr_x = start_x + (i as i32);
                if curr_x < 0 || curr_x >= self.info.width as i32 { curr_x = x as i32 };

                data.push(self.get_pixel(curr_x as u32, y));
            }

            SubImage::new(size, 1, self.info.channels, self.info.alpha, data)
        }
    }

    /// Returns a `SubImage<T>` representing the "square" of pixels of side length `size` centered
    /// at `(x, y)`. Uses clamp padding for edge pixels (edge pixels are repeated indefinitely)
    pub fn get_neighborhood_2d(&self, x: u32, y: u32, size: u32) -> SubImage<T> {
        let start_x = (x as i32) - (size as i32) / 2;
        let start_y = (y as i32) - (size as i32) / 2;

        let mut data = Vec::new();
        for i in 0..size {
            for j in 0..size {
                let mut curr_x = start_x + (j as i32);
                let mut curr_y = start_y + (i as i32);

                if curr_x < 0 || curr_x >= self.info.width as i32 { curr_x = x as i32 };
                if curr_y < 0 || curr_y >= self.info.height as i32 { curr_y = y as i32 };

                data.push(self.get_pixel(curr_x as u32, curr_y as u32));
            }
        }

        SubImage::new(size, size, self.info.channels, self.info.alpha, data)
    }

    /// Replaces the pixel located at `(x, y)` with `pixel`
    pub fn set_pixel(&mut self, x: u32, y: u32, pixel: &[T]) {
        if pixel.len() != self.info.channels as usize {
            panic!("invalid pixel length: the number of channels is {}, \
                but the pixel length is {}", self.info.channels, pixel.len())
        }

        let start = self.index(x, y);
        for i in 0..(self.info.channels as usize) {
            self.data[i + start] = pixel[i];
        }
    }

    /// Replaces the pixel at index `index` with `pixel`
    pub fn set_pixel_indexed(&mut self, index: usize, pixel: &[T]) {
        let start = index * self.info.channels as usize;
        for i in 0..(self.info.channels as usize) {
            self.data[start + i] = pixel[i];
        }
    }

    /// Applies function `f` to each pixel
    pub fn map_pixels<S: Number, F>(&self, f: F) -> Image<S>
        where F: Fn(&[T]) -> Vec<S> {
        let mut  data= Vec::new();

        for i in 0..(self.info.size() as usize) {
            data.append(&mut f(&self[i]));
        }

        let channels = (data.len() as u32 / self.info.size()) as u8;

        Image {
            info: ImageInfo {
                width: self.info.width,
                height: self.info.height,
                channels,
                alpha: self.info.alpha
            },
            data,
        }
    }

    /// If `alpha`, applies function `f` to the non-alpha portion of each pixel and applies
    /// function `g` to the alpha channel of each pixel; otherwise, applies function `f` to
    /// each pixel
    pub fn map_pixels_if_alpha<S: Number, F, G>(&self, f: F, g: G) -> Image<S>
        where F: Fn(&[T]) -> Vec<S>,
              G: Fn(T) -> S {
        if !self.info.alpha {
            return self.map_pixels(f);
        }

        let mut data = Vec::new();
        for i in 0..(self.info.size() as usize) {
            let mut v = f(self[i].channels_without_alpha());
            v.push(g(self[i].alpha()));
            data.append(&mut v);
        }

        let channels = (data.len() as u32 / self.info.size()) as u8;

        Image {
            info: ImageInfo {
                width: self.info.width,
                height: self.info.height,
                channels,
                alpha: self.info.alpha
            },
            data,
        }
    }

    /// Applies function `f` to each channel of each pixel
    pub fn map_channels<S: Number, F>(&self, f: F) -> Image<S>
        where F: Fn(T) -> S {
        let mut data = Vec::new();

        for i in 0..(self.info.full_size() as usize) {
            data.push(f(self.data[i]));
        }

        Image {
            info: self.info,
            data,
        }
    }

    /// If `alpha`, applies function `f` to each non-alpha channel of each pixel and
    /// applies function `g` to the alpha channel of each pixel;
    /// otherwise, applies function `f` to each channel of each pixel
    pub fn map_channels_if_alpha<S: Number, F, G>(&self, f: F, g: G) -> Image<S>
        where F: Fn(T) -> S,
              G: Fn(T) -> S {
        if !self.info.alpha {
            return self.map_channels(f);
        }

        let mut data = Vec::new();
        for i in 0..(self.info.size() as usize) {
            data.append(&mut self[i].map_alpha(&f, &g));
        }

        Image {
            info: self.info,
            data,
        }
    }

    /// Applies function `f` to each pixel
    pub fn apply_pixels<F>(&mut self, f: F)
        where F: Fn(&[T]) -> Vec<T> {
        for i in 0..self.info.size() as usize {
            self.set_pixel_indexed(i, f(&self[i]).as_slice());
        }
    }

    /// If `alpha`, applies function `f` to the non-alpha portion of each pixel and
    /// applies function `g` to the alpha channel of each pixel;
    /// otherwise, applies function `f` to each pixel
    pub fn apply_pixels_if_alpha<F, G>(&mut self, f: F, g: G)
        where F: Fn(&[T]) -> Vec<T>,
              G: Fn(T) -> T {
        if !self.info.alpha {
            self.apply_pixels(f);
            return;
        }

        for i in 0..(self.info.size() as usize) {
            let mut v = f(self[i].channels_without_alpha());
            v.push(g(self[i].alpha()));
            self.set_pixel_indexed(i, v.as_slice());
        }
    }

    /// Applies function `f` to each channel of each pixel
    pub fn apply_channels<F>(&mut self, f: F)
        where F: Fn(T) -> T {
        for i in 0..(self.info.full_size() as usize) {
            self.data[i] = f(self.data[i]);
        }
    }

    /// If `alpha`, applies function `f` to each non-alpha channel of each pixel and
    /// applies function `g` to the alpha channel of each pixel;
    /// otherwise, applies function `f` to each channel of each pixel
    pub fn apply_channels_if_alpha<F, G>(&mut self, f: F, g: G)
        where F: Fn(T) -> T,
              G: Fn(T) -> T {
        if !self.info.alpha {
            self.apply_channels(f);
            return;
        }

        for i in 0..(self.info.size() as usize) {
            self[i].apply_alpha(&f, &g);
        }
    }

    /// Applies function `f` to each channel of index `index` of each pixel. Modifies `self`
    pub fn edit_channel<F>(&mut self, f: F, index: usize)
        where F: Fn(T) -> T {
        for i in (index..(self.info.full_size() as usize)).step_by(self.info.channels as usize) {
            self.data[i] = f(self.data[i]);
        }
    }
}

impl<T: Number> BaseImage<T> for Image<T> {
    fn info(&self) -> ImageInfo {
        self.info
    }

    fn get_pixel(&self, x: u32, y: u32) -> &[T] {
        if x >= self.info.width {
            panic!("index out of bounds: the width is {}, but the x index is {}", self.info.width, x)
        }
        if y >= self.info.height {
            panic!("index out of bounds: the height is {}, but the y index is {}", self.info.height, y)
        }

        &self[(y * self.info.width + x) as usize]
    }
}

impl<T: Number> std::ops::Index<usize> for Image<T> {
    type Output = [T];

    fn index(&self, i: usize) -> &Self::Output {
        if i >= self.info.size() as usize {
            panic!("index out of bounds: the len is {}, but the index is {}", self.info.size(), i)
        }

        let start = i * (self.info.channels as usize);
        &self.data[start..(start + self.info.channels as usize)]
    }
}

impl<T: Number> std::ops::IndexMut<usize> for Image<T> {
    fn index_mut(&mut self, i: usize) -> &mut Self::Output {
        if i >= self.info.size() as usize {
            panic!("index out of bounds: the len is {}, but the index is {}", self.info.size(), i)
        }

        let start = i * (self.info.channels as usize);
        &mut self.data[start..(start + self.info.channels as usize)]
    }
}

impl<'a, T: Number> SubImage<'a, T> {
    /// Creates a new `SubImage<T>`
    pub fn new(width: u32, height: u32, channels: u8, alpha: bool, data: Vec<&'a [T]>) -> Self {
        SubImage {
            info: ImageInfo { width, height, channels, alpha },
            data,
        }
    }

    pub fn data(&self) -> &[&[T]] {
        &self.data[..]
    }
}

impl<T: Number> BaseImage<T> for SubImage<'_, T> {
    fn info(&self) -> ImageInfo {
        self.info
    }

    fn get_pixel(&self, x: u32, y: u32) -> &[T] {
        &self[(y * self.info.width + x) as usize]
    }
}

impl<T: Number> std::ops::Index<usize> for SubImage<'_, T> {
    type Output = [T];

    fn index(&self, i: usize) -> &Self::Output {
        self.data[i]
    }
}

impl<T: Number> Pixel<T> for [T] {
    fn alpha(&self) -> T {
        self[self.len()-1]
    }

    fn channels_without_alpha(&self) -> &[T] {
        &self[..(self.len()-1)]
    }

    fn map_all<S: Number, F>(&self, f: F) -> Vec<S>
        where F: Fn(T) -> S {
        let mut channels_out = Vec::new();

        for channel in self.iter() {
            channels_out.push(f(*channel));
        }

        channels_out
    }

    fn map_alpha<S: Number, F, G>(&self, f: F, g: G) -> Vec<S>
        where F: Fn(T) -> S,
              G: Fn(T) -> S {
        let mut channels_out = Vec::new();

        for channel in self.channels_without_alpha().iter() {
            channels_out.push(f(*channel));
        }

        channels_out.push(g(self.alpha()));
        channels_out
    }

    fn apply<F>(&mut self, f: F)
        where F: Fn(T) -> T {
        for i in 0..self.len() {
            self[i] = f(self[i]);
        }
    }

    fn apply_alpha<F, G>(&mut self, f: F, g: G)
        where F: Fn(T) -> T,
              G: Fn(T) -> T {
        for i in 0..(self.len() - 1) {
            self[i] = f(self[i]);
        }

        self[self.len()-1] = g(self.alpha());
    }

    fn is_black(&self) -> bool {
        for channel in self.iter() {
            if *channel != 0.into() {
                return false;
            }
        }

        true
    }

    fn is_black_alpha(&self) -> bool {
        for channel in self.channels_without_alpha().iter() {
            if *channel != 0.into() {
                return false;
            }
        }

        true
    }
}
