//! Ratio's matrix data library

use std::marker::PhantomData;
use std::ops::{Add, AddAssign, Range, Sub, SubAssign};

use snafu::prelude::*;

use crate::traits::{Contains2D, Dimensions2D, Domain2D, Matrix2D, MatrixMut, MatrixRef};

pub mod traits;

#[cfg(test)]
mod tests;

/// Matrix dimensions don't agree. Values ({n_values}) modulo columns ({n_cols}) should be 0.
#[derive(Clone, Debug, Snafu)]
pub struct MatrixDimensionError {
    /// Number of values.
    pub n_values: usize,
    /// Invalid number of columns in a row.
    pub n_cols: usize,
}

/// Matrix view out of bounds.
#[derive(Clone, Debug, Snafu)]
pub struct MatrixViewOutOfBoundsError {
    /// Input data dimensions.
    pub dimensions: MatrixRange,
    /// Invalid range.
    pub range: MatrixRange,
}

/// Matrix cell access out of bounds.
#[derive(Clone, Debug, Snafu)]
pub struct CoordinatesOutOfBoundsError {
    /// Matrix data area.
    pub range: MatrixRange,
    /// Attempted access coordinates.
    pub coordinates: MatrixCoordinates,
}

/// Two-dimensional matrix.
#[derive(Clone, Debug, Default, PartialEq)]
#[cfg_attr(
    feature = "serde",
    derive(serde::Serialize, serde::Deserialize),
    serde(default, rename_all = "camelCase")
)]
pub struct Matrix<C> {
    /// Matrix values as a single vector.
    pub values: Vec<C>,
    /// Matrix dimensions.
    pub dimensions: MatrixDimensions,
}
impl<C> Matrix<C> {
    /// Get the vector index based on a set of coordinates.
    pub fn index<T: Into<MatrixCoordinates>>(
        &self,
        coordinates: T,
    ) -> Result<usize, CoordinatesOutOfBoundsError> {
        let x: MatrixCoordinates = coordinates.into();
        if self.range().contains(&x) {
            Ok(x.row * self.n_cols() + x.col)
        } else {
            Err(CoordinatesOutOfBoundsError {
                coordinates: x,
                range: self.range(),
            })
        }
    }
    /// Slice this matrix to these rows and columns.
    pub fn slice<R: Into<MatrixRange>>(
        &self,
        range: R,
    ) -> Result<MatrixSlice<'_, C>, MatrixSliceError> {
        let range: MatrixRange = range.into();
        if self.range().contains(&range) {
            Ok(MatrixSlice {
                values: &self.values,
                input_dimensions: self.dimensions,
                view_range: range,
            })
        } else {
            MatrixViewOutOfBoundsSnafu {
                range,
                dimensions: self.dimensions,
            }
            .fail()
            .context(SliceRangeOutOfBoundsSnafu)
        }
    }
}
impl<C> Matrix2D<C> for Matrix<C> {
    fn new(values: Vec<C>, n_cols: usize) -> Result<Self, MatrixDimensionError> {
        let n_values = values.len();
        let dimensions = MatrixDimensions::new(n_values / n_cols, n_cols);
        if dimensions.len() == n_values {
            Ok(Self { values, dimensions })
        } else {
            MatrixDimensionSnafu { n_values, n_cols }.fail()
        }
    }
}
impl<C> MatrixRef<C> for Matrix<C> {
    fn get<T: Into<MatrixCoordinates>>(
        &self,
        coordinates: T,
    ) -> Result<&C, CoordinatesOutOfBoundsError> {
        let x: MatrixCoordinates = coordinates.into();
        let index = self.index(x)?;
        self.values.get(index).context(CoordinatesOutOfBoundsSnafu {
            range: self.range(),
            coordinates: x,
        })
    }
}
impl<C> MatrixMut<C> for Matrix<C> {
    fn set<T: Into<MatrixCoordinates>>(
        &mut self,
        coordinates: T,
        value: C,
    ) -> Result<C, CoordinatesOutOfBoundsError> {
        let coordinates: MatrixCoordinates = coordinates.into();
        let index = self.index(coordinates)?;
        if let Some(entry) = self.values.get_mut(index) {
            Ok(std::mem::replace(entry, value))
        } else {
            Err(CoordinatesOutOfBoundsError {
                range: self.range(),
                coordinates,
            })
        }
    }
}
impl<C> Domain2D for Matrix<C> {
    fn row_range(&self) -> Range<usize> {
        0..self.dimensions.n_rows
    }

    fn col_range(&self) -> Range<usize> {
        0..self.dimensions.n_cols
    }
}

/// Matrix dimensions don't agree. Values ({n_values}) modulo columns ({n_cols}) should be 0.
#[derive(Clone, Debug, Snafu)]
pub enum MatrixSliceError {
    SliceInputDimension { source: MatrixDimensionError },
    SliceRangeOutOfBounds { source: MatrixViewOutOfBoundsError },
}

/// Slice of a matrix.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct MatrixSlice<'a, C> {
    /// Input matrix values as a single slice reference.
    pub values: &'a [C],
    /// Dimensions of the input data.
    pub input_dimensions: MatrixDimensions,
    /// Slice or view range with respect to the input data.
    pub view_range: MatrixRange,
}
impl<'a, C> MatrixSlice<'a, C> {
    /// Create a new slice from a matrix. Takes a reference to the input values' slice and the
    /// number of columns in each row as well as the intended range of the slice.
    pub fn new<R: Into<MatrixRange>>(
        values: &'a [C],
        n_cols: usize,
        range: R,
    ) -> Result<Self, MatrixSliceError> {
        let n_values = values.len();
        let n_rows = n_values / n_cols;
        let input_dimensions = MatrixDimensions::new(n_rows, n_cols);
        if values.len() != input_dimensions.len() {
            return MatrixDimensionSnafu { n_values, n_cols }
                .fail()
                .context(SliceInputDimensionSnafu);
        }
        let domain: MatrixRange = input_dimensions.into();
        let range: MatrixRange = range.into();
        if domain.contains(&range) {
            Ok(Self {
                values,
                input_dimensions,
                view_range: range,
            })
        } else {
            MatrixViewOutOfBoundsSnafu {
                dimensions: input_dimensions,
                range,
            }
            .fail()
            .context(SliceRangeOutOfBoundsSnafu)
        }
    }

    /// Get the index for item access in the reference slice.
    /// The provided coordinates should be with respect to the slice's start!
    pub fn index<T: Into<MatrixCoordinates>>(
        &self,
        coordinates: T,
    ) -> Result<usize, CoordinatesOutOfBoundsError> {
        let x: MatrixCoordinates = coordinates.into();
        let x = x + self.view_range.start;
        if self.view_range.contains(&x) {
            Ok(x.row * self.input_dimensions.n_cols + x.col)
        } else {
            Err(CoordinatesOutOfBoundsError {
                coordinates: x,
                range: self.view_range,
            })
        }
    }
}

impl<'a, C> MatrixSlice<'a, C>
where
    C: Clone + Default,
{
    /// Clone the data into this slice into a matrix of it's own size.
    pub fn to_matrix(&self) -> Matrix<C> {
        Matrix::new(
            self.row_range()
                .flat_map(|row| {
                    self.col_range().filter_map(move |col| {
                        self.get(MatrixCoordinates::new(row, col)).ok().cloned()
                    })
                })
                .collect(),
            self.view_range.col_range().end,
        )
        .unwrap_or_default()
    }
}
impl<'a, C> MatrixRef<C> for MatrixSlice<'a, C> {
    fn get<T: Into<MatrixCoordinates>>(
        &self,
        coordinates: T,
    ) -> Result<&C, CoordinatesOutOfBoundsError> {
        let coordinates: MatrixCoordinates = coordinates.into();
        let index = self.index(coordinates)?;
        self.values.get(index).context(CoordinatesOutOfBoundsSnafu {
            range: self.view_range,
            coordinates,
        })
    }
}
impl<'a, C> Domain2D for MatrixSlice<'a, C> {
    fn row_range(&self) -> Range<usize> {
        0..(self.view_range.end.row - self.view_range.start.row)
    }
    fn col_range(&self) -> Range<usize> {
        0..(self.view_range.end.col - self.view_range.start.col)
    }
}

/// Matrix cell.
#[derive(Clone, Debug, PartialEq)]
pub struct MatrixCellRef<'a, M: MatrixRef<C> + ?Sized, C> {
    /// Matrix to pull values from.
    pub matrix: &'a M,
    /// Coordinates of this cell.
    pub coordinates: MatrixCoordinates,
    /// Phantom data marker for the cell value type.
    _phantom: PhantomData<C>,
}
impl<'a, M: MatrixRef<C> + ?Sized, C> MatrixCellRef<'a, M, C> {
    /// Create a new cell reference using the given coordinates.
    pub fn new<T: Into<MatrixCoordinates>>(matrix: &'a M, coordinates: T) -> Self {
        Self {
            matrix,
            coordinates: coordinates.into(),
            _phantom: PhantomData,
        }
    }
}

/// Location in a 2D matrix.
#[derive(Copy, Clone, Debug, PartialEq)]
#[cfg_attr(
    feature = "serde",
    derive(serde::Serialize, serde::Deserialize),
    serde(rename_all = "camelCase")
)]
pub enum MatrixLocation {
    /// Entire row in a matrix.
    Row(usize),
    /// Entire column in a matrix.
    Col(usize),
    /// Specific coordinates in a matrix.
    Coordinates(MatrixCoordinates),
    /// 2D area in a matrix.
    Range(MatrixRange),
}
impl Default for MatrixLocation {
    fn default() -> Self {
        Self::Coordinates(Default::default())
    }
}
impl From<MatrixCoordinates> for MatrixLocation {
    fn from(value: MatrixCoordinates) -> Self {
        Self::Coordinates(value)
    }
}
impl MatrixLocation {
    /// Create a new optional location from an optional row and column coordinate.
    pub fn from_coords(row: Option<usize>, col: Option<usize>) -> Option<Self> {
        match (row, col) {
            (Some(row), Some(col)) => Some(Self::Coordinates(MatrixCoordinates::new(row, col))),
            (Some(row), None) => Some(Self::Row(row)),
            (None, Some(col)) => Some(Self::Col(col)),
            (None, None) => None,
        }
    }
    /// Row or starting row of this location.
    pub fn row(&self) -> Option<usize> {
        match self {
            Self::Row(row) => Some(*row),
            Self::Col(_) => None,
            Self::Coordinates(coord) => Some(coord.row),
            Self::Range(area) => Some(area.row_range().start),
        }
    }
    /// Column or starting column of this location.
    pub fn col(&self) -> Option<usize> {
        match self {
            Self::Row(_) => None,
            Self::Col(col) => Some(*col),
            Self::Coordinates(coord) => Some(coord.col),
            Self::Range(area) => Some(area.col_range().start),
        }
    }
}
impl Domain2D for MatrixLocation {
    fn row_range(&self) -> Range<usize> {
        match self {
            Self::Coordinates(coordinates) => Range {
                start: coordinates.row,
                end: coordinates.row + 1,
            },
            Self::Row(row) => Range {
                start: *row,
                end: row + 1,
            },
            Self::Col(_) => Range { start: 0, end: 0 },
            Self::Range(range) => range.row_range(),
        }
    }
    fn col_range(&self) -> Range<usize> {
        match self {
            Self::Coordinates(coordinates) => Range {
                start: coordinates.col,
                end: coordinates.col + 1,
            },
            Self::Row(_) => Range { start: 0, end: 0 },
            Self::Col(col) => Range {
                start: *col,
                end: col + 1,
            },
            Self::Range(range) => range.col_range(),
        }
    }
}

/// Coordinate in a 2D matrix.
#[derive(Copy, Clone, Debug, Default, PartialEq)]
#[cfg_attr(
    feature = "serde",
    derive(serde::Serialize, serde::Deserialize),
    serde(default, rename_all = "camelCase")
)]
pub struct MatrixCoordinates {
    /// Coordinate row.
    pub row: usize,
    /// Coordinate column.
    pub col: usize,
}
impl MatrixCoordinates {
    /// Create a new matrix coordinate instance.
    pub fn new(row: usize, col: usize) -> Self {
        Self { row, col }
    }
}
impl From<(usize, usize)> for MatrixCoordinates {
    fn from(value: (usize, usize)) -> Self {
        Self {
            row: value.0,
            col: value.1,
        }
    }
}
impl Add<MatrixCoordinates> for MatrixCoordinates {
    type Output = Self;
    fn add(self, rhs: MatrixCoordinates) -> Self::Output {
        MatrixCoordinates {
            row: self.row + rhs.row,
            col: self.col + rhs.col,
        }
    }
}
impl AddAssign<MatrixCoordinates> for MatrixCoordinates {
    fn add_assign(&mut self, rhs: MatrixCoordinates) {
        self.row += rhs.row;
        self.col += rhs.col;
    }
}
impl Sub<MatrixCoordinates> for MatrixCoordinates {
    type Output = Self;
    fn sub(self, rhs: MatrixCoordinates) -> Self::Output {
        MatrixCoordinates {
            row: self.row - rhs.row,
            col: self.col - rhs.col,
        }
    }
}
impl SubAssign<MatrixCoordinates> for MatrixCoordinates {
    fn sub_assign(&mut self, rhs: MatrixCoordinates) {
        self.row -= rhs.row;
        self.col -= rhs.col;
    }
}

/// 2D matrix dimensions.
#[derive(Copy, Clone, Debug, Default, PartialEq)]
#[cfg_attr(
    feature = "serde",
    derive(serde::Serialize, serde::Deserialize),
    serde(default, rename_all = "camelCase")
)]
pub struct MatrixDimensions {
    /// Coordinate row.
    pub n_rows: usize,
    /// Coordinate column.
    pub n_cols: usize,
}
impl MatrixDimensions {
    /// Create a new matrix dimensions instance.
    pub fn new(n_rows: usize, n_cols: usize) -> Self {
        Self { n_rows, n_cols }
    }
}
impl From<(usize, usize)> for MatrixDimensions {
    fn from(value: (usize, usize)) -> Self {
        Self {
            n_rows: value.0,
            n_cols: value.1,
        }
    }
}
impl Dimensions2D for MatrixDimensions {
    fn n_rows(&self) -> usize {
        self.n_rows
    }
    fn n_cols(&self) -> usize {
        self.n_cols
    }
}
impl Add<MatrixDimensions> for MatrixDimensions {
    type Output = Self;
    fn add(self, rhs: MatrixDimensions) -> Self::Output {
        MatrixDimensions {
            n_rows: self.n_rows + rhs.n_rows,
            n_cols: self.n_cols + rhs.n_cols,
        }
    }
}
impl AddAssign<MatrixDimensions> for MatrixDimensions {
    fn add_assign(&mut self, rhs: MatrixDimensions) {
        self.n_rows += rhs.n_rows;
        self.n_cols += rhs.n_cols;
    }
}
impl Sub<MatrixDimensions> for MatrixDimensions {
    type Output = Self;
    fn sub(self, rhs: MatrixDimensions) -> Self::Output {
        MatrixDimensions {
            n_rows: self.n_rows - rhs.n_rows,
            n_cols: self.n_cols - rhs.n_cols,
        }
    }
}
impl SubAssign<MatrixDimensions> for MatrixDimensions {
    fn sub_assign(&mut self, rhs: MatrixDimensions) {
        self.n_rows -= rhs.n_rows;
        self.n_cols -= rhs.n_cols;
    }
}

/// 2D range in a matrix.
#[derive(Copy, Clone, Debug, Default, PartialEq)]
#[cfg_attr(
    feature = "serde",
    derive(serde::Serialize, serde::Deserialize),
    serde(default, rename_all = "camelCase")
)]
pub struct MatrixRange {
    /// Starting coordinates of the range (inclusive).
    pub start: MatrixCoordinates,
    /// Ending coordinates of the range (exclusive).
    pub end: MatrixCoordinates,
}
impl MatrixRange {
    /// Create a new matrix area instance.
    pub fn new(rows: Range<usize>, cols: Range<usize>) -> Self {
        Self {
            start: MatrixCoordinates::new(rows.start, cols.start),
            end: MatrixCoordinates::new(rows.end, cols.end),
        }
    }
}
impl Domain2D for MatrixRange {
    /// Row range of this area.
    fn row_range(&self) -> Range<usize> {
        self.start.row..self.end.row
    }

    /// Column range of this area.
    fn col_range(&self) -> Range<usize> {
        self.start.col..self.end.col
    }
}
impl From<(Range<usize>, Range<usize>)> for MatrixRange {
    fn from(value: (Range<usize>, Range<usize>)) -> Self {
        Self::new(value.0, value.1)
    }
}
impl From<((usize, usize), (usize, usize))> for MatrixRange {
    /// Create a range from two sets of coordinates (start, end).
    fn from(value: ((usize, usize), (usize, usize))) -> Self {
        Self {
            start: value.0.into(),
            end: value.1.into(),
        }
    }
}
impl From<MatrixDimensions> for MatrixRange {
    fn from(value: MatrixDimensions) -> Self {
        Self::new(0..value.n_rows, 0..value.n_cols)
    }
}
impl Add<MatrixCoordinates> for MatrixRange {
    type Output = Self;
    fn add(self, rhs: MatrixCoordinates) -> Self::Output {
        Self {
            start: self.start + rhs,
            end: self.end + rhs,
        }
    }
}
impl AddAssign<MatrixCoordinates> for MatrixRange {
    fn add_assign(&mut self, rhs: MatrixCoordinates) {
        self.start += rhs;
        self.end += rhs;
    }
}
impl Sub<MatrixCoordinates> for MatrixRange {
    type Output = Self;
    fn sub(self, rhs: MatrixCoordinates) -> Self::Output {
        Self {
            start: self.start - rhs,
            end: self.end - rhs,
        }
    }
}
impl SubAssign<MatrixCoordinates> for MatrixRange {
    fn sub_assign(&mut self, rhs: MatrixCoordinates) {
        self.start -= rhs;
        self.end -= rhs;
    }
}