//! Ratio's matrix data library

use std::marker::PhantomData;
use std::ops::Range;

use snafu::prelude::*;
mod traits;
pub use traits::*;

/// 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.
    n_values: usize,
    /// Invalid number of columns in a row.
    n_cols: usize,
}

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

/// Matrix cell access out of bounds.
#[derive(Clone, Debug, Snafu)]
pub struct CoordinatesOutOfBoundsError {
    /// Matrix data area.
    range: MatrixRange,
    /// Attempted access coordinates.
    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.
    values: Vec<C>,
    /// Matrix dimensions.
    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(&self, range: MatrixRange) -> Result<MatrixSlice<'_, C>, MatrixSliceError> {
        if self.range().contains(&range) {
            Ok(MatrixSlice {
                values: &self.values,
                dimensions: self.dimensions,
                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 {
            return 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)]
pub struct MatrixSlice<'a, C> {
    /// Input matrix values as a single slice reference.
    values: &'a [C],
    /// Dimensions of the input data.
    dimensions: MatrixDimensions,
    /// Slice or view range.
    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(
        values: &'a [C],
        n_cols: usize,
        range: MatrixRange,
    ) -> Result<Self, MatrixSliceError> {
        let n_values = values.len();
        let n_rows = n_values / n_cols;
        let dimensions = MatrixDimensions::new(n_rows, n_cols);
        if values.len() != dimensions.len() {
            return MatrixDimensionSnafu { n_values, n_cols }
                .fail()
                .context(SliceInputDimensionSnafu);
        }
        let domain: MatrixRange = dimensions.into();
        if domain.contains(&range) {
            Ok(Self {
                values,
                dimensions,
                range,
            })
        } else {
            MatrixViewOutOfBoundsSnafu { dimensions, range }
                .fail()
                .context(SliceRangeOutOfBoundsSnafu)
        }
    }

    /// Get the index for item access in the reference slice.
    /// The provided coordinates should be with respect to the original input data!
    pub fn index<T: Into<MatrixCoordinates>>(
        &self,
        coordinates: T,
    ) -> Result<usize, CoordinatesOutOfBoundsError> {
        let x: MatrixCoordinates = coordinates.into();
        let n_rows = self.dimensions.n_rows;
        let n_cols = self.dimensions.n_cols;
        if x.row < n_rows && x.col < n_cols {
            Ok(x.row * self.n_cols() + x.col)
        } else {
            Err(CoordinatesOutOfBoundsError {
                coordinates: x,
                range: self.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.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.range,
            coordinates,
        })
    }
}
impl<'a, C> Domain2D for MatrixSlice<'a, C> {
    fn row_range(&self) -> Range<usize> {
        self.range.row_range()
    }
    fn col_range(&self) -> Range<usize> {
        self.range.col_range()
    }
}

/// 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,
        }
    }
}

/// 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
    }
}

/// 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 bounds ((row lower, row upper), (col lower, col upper)).
    fn from(value: ((usize, usize), (usize, usize))) -> Self {
        Self {
            start: MatrixCoordinates::new(value.0 .0, value.1 .0),
            end: MatrixCoordinates::new(value.0 .1, value.1 .1),
        }
    }
}
impl From<MatrixDimensions> for MatrixRange {
    fn from(value: MatrixDimensions) -> Self {
        Self::new(0..value.n_rows, 0..value.n_cols)
    }
}