ndbioimage 0.2.0

use crate::error::Error;
use crate::stats::MinMax;
use ndarray::{Array, Dimension, Ix2, SliceInfo, SliceInfoElem};
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use serde_with::{DeserializeAs, SerializeAs};
use std::collections::HashMap;
use std::fmt::{Display, Formatter};
use std::hash::{Hash, Hasher};
use std::ops::Index;
use strum::{AsRefStr, Display, EnumString};

/// a trait to find axis indices from any object
pub trait Ax {
    /// C: 0, Z: 1, T: 2, Y: 3, X: 4
    fn n(&self) -> usize;

    /// the indices of axes in self.axes, which always has all of CZTYX
    fn pos(&self, axes: &[Axis], slice: &[SliceInfoElem]) -> Result<usize, Error>;

    /// the indices of axes in self.axes, which always has all of CZTYX, but skip axes with an operation
    fn pos_op(
        &self,
        axes: &[Axis],
        slice: &[SliceInfoElem],
        op_axes: &[Axis],
    ) -> Result<usize, Error>;
}

/// Enum for CZTYX axes or a new axis
#[derive(
    Clone, Copy, Debug, Eq, Ord, PartialOrd, Serialize, Deserialize, EnumString, AsRefStr, Display,
)]
#[strum(ascii_case_insensitive)]
pub enum Axis {
    /// channel
    C,
    /// z slice
    Z,
    /// time
    T,
    /// y coordinate
    Y,
    /// x coordinate
    X,
    #[strum(serialize = "N")]
    /// a new axis created by slicing, can be operated on but has no reader data
    New,
}

impl Hash for Axis {
    fn hash<H: Hasher>(&self, state: &mut H) {
        (*self as usize).hash(state);
    }
}

impl Ax for Axis {
    fn n(&self) -> usize {
        *self as usize
    }

    fn pos(&self, axes: &[Axis], _slice: &[SliceInfoElem]) -> Result<usize, Error> {
        if let Some(pos) = axes.iter().position(|a| a == self) {
            Ok(pos)
        } else {
            Err(Error::AxisNotFound(
                format!("{:?}", self),
                format!("{:?}", axes),
            ))
        }
    }

    fn pos_op(
        &self,
        axes: &[Axis],
        _slice: &[SliceInfoElem],
        _op_axes: &[Axis],
    ) -> Result<usize, Error> {
        self.pos(axes, _slice)
    }
}

impl Ax for usize {
    fn n(&self) -> usize {
        *self
    }

    fn pos(&self, _axes: &[Axis], slice: &[SliceInfoElem]) -> Result<usize, Error> {
        let idx: Vec<_> = slice
            .iter()
            .enumerate()
            .filter_map(|(i, s)| if s.is_index() { None } else { Some(i) })
            .collect();
        Ok(idx[*self])
    }

    fn pos_op(
        &self,
        axes: &[Axis],
        slice: &[SliceInfoElem],
        op_axes: &[Axis],
    ) -> Result<usize, Error> {
        let idx: Vec<_> = axes
            .iter()
            .zip(slice.iter())
            .enumerate()
            .filter_map(|(i, (ax, s))| {
                if s.is_index() | op_axes.contains(ax) {
                    None
                } else {
                    Some(i)
                }
            })
            .collect();
        debug_assert!(*self < idx.len(), "self: {}, idx: {:?}", self, idx);
        Ok(idx[*self])
    }
}

#[derive(
    Clone, Debug, Serialize, Deserialize, EnumString, AsRefStr, Display, PartialEq, Eq, Hash,
)]
#[strum(ascii_case_insensitive)]
/// an operation to reduce an axis
pub enum Operation {
    /// take the max along the axis
    Max,
    /// take the min along the axis
    Min,
    /// take the sum along the axis
    Sum,
    /// take the mean along the axis
    Mean,
}

impl Operation {
    pub(crate) fn operate<T, D>(
        &self,
        array: Array<T, D>,
        axis: usize,
    ) -> Result<<Array<T, D> as MinMax>::Output, Error>
    where
        D: Dimension,
        Array<T, D>: MinMax,
    {
        match self {
            Operation::Max => array.max(axis),
            Operation::Min => array.min(axis),
            Operation::Sum => array.sum(axis),
            Operation::Mean => array.mean(axis),
        }
    }
}

impl PartialEq for Axis {
    fn eq(&self, other: &Self) -> bool {
        (*self as u8) == (*other as u8)
    }
}

pub(crate) fn slice_info<D: Dimension>(
    info: &[SliceInfoElem],
) -> Result<SliceInfo<&[SliceInfoElem], Ix2, D>, Error> {
    match info.try_into() {
        Ok(slice) => Ok(slice),
        Err(err) => Err(Error::TryInto(err.to_string())),
    }
}

#[derive(Serialize, Deserialize)]
#[serde(remote = "SliceInfoElem")]
pub(crate) enum SliceInfoElemDef {
    Slice {
        start: isize,
        end: Option<isize>,
        step: isize,
    },
    Index(isize),
    NewAxis,
}

impl SerializeAs<SliceInfoElem> for SliceInfoElemDef {
    fn serialize_as<S>(source: &SliceInfoElem, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        SliceInfoElemDef::serialize(source, serializer)
    }
}

impl<'de> DeserializeAs<'de, SliceInfoElem> for SliceInfoElemDef {
    fn deserialize_as<D>(deserializer: D) -> Result<SliceInfoElem, D::Error>
    where
        D: Deserializer<'de>,
    {
        SliceInfoElemDef::deserialize(deserializer)
    }
}

#[derive(Clone, Debug)]
pub(crate) struct Slice {
    start: isize,
    end: isize,
    step: isize,
}

impl Slice {
    pub(crate) fn new(start: isize, end: isize, step: isize) -> Self {
        Self { start, end, step }
    }

    pub(crate) fn empty() -> Self {
        Self {
            start: 0,
            end: 0,
            step: 1,
        }
    }
}

impl Iterator for Slice {
    type Item = isize;

    fn next(&mut self) -> Option<Self::Item> {
        if self.end - self.start >= self.step {
            let r = self.start;
            self.start += self.step;
            Some(r)
        } else {
            None
        }
    }
}

impl IntoIterator for &Slice {
    type Item = isize;
    type IntoIter = Slice;

    fn into_iter(self) -> Self::IntoIter {
        self.clone()
    }
}

/// the size of a view along the axes it contains
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct Shape {
    /// size c (# channels)
    pub c: usize,
    /// size z (# slices)
    pub z: usize,
    /// size t (# time/frames)
    pub t: usize,
    /// size y (vertical)
    pub y: usize,
    /// size x (horizontal)
    pub x: usize,
    /// the axes that are present in the view, in order
    pub order: Vec<Axis>,
}

impl Default for Shape {
    fn default() -> Self {
        Self {
            c: 1,
            z: 1,
            t: 1,
            y: 1,
            x: 1,
            order: vec![Axis::C, Axis::Z, Axis::T, Axis::Y, Axis::X],
        }
    }
}

impl Display for Shape {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        let order = self
            .order
            .iter()
            .map(|i| i.to_string())
            .collect::<Vec<_>>()
            .join("");
        let shape = self
            .order
            .iter()
            .map(|i| format!("{}", self[i]))
            .collect::<Vec<_>>()
            .join(", ");
        write!(f, "{}: {}", order, shape)
    }
}

impl Index<Axis> for Shape {
    type Output = usize;

    fn index(&self, ax: Axis) -> &Self::Output {
        &self[&ax]
    }
}

impl Index<&Axis> for Shape {
    type Output = usize;

    fn index(&self, ax: &Axis) -> &Self::Output {
        match ax {
            Axis::C => &self.c,
            Axis::Z => &self.z,
            Axis::T => &self.t,
            Axis::Y => &self.y,
            Axis::X => &self.x,
            Axis::New => &1,
        }
    }
}

impl Index<usize> for Shape {
    type Output = usize;

    fn index(&self, dim: usize) -> &Self::Output {
        &self[self.order[dim % self.order.len()]]
    }
}

impl Index<&usize> for Shape {
    type Output = usize;

    fn index(&self, dim: &usize) -> &Self::Output {
        &self[self.order[dim % self.order.len()]]
    }
}

impl From<Shape> for Vec<usize> {
    fn from(shape: Shape) -> Self {
        shape.order.iter().map(|axis| shape[axis]).collect()
    }
}

impl From<Shape> for HashMap<Axis, usize> {
    fn from(shape: Shape) -> Self {
        shape
            .order
            .iter()
            .map(|axis| (*axis, shape[axis]))
            .collect()
    }
}

/// iterator over the sizes along the axes of a [`Shape`]
pub struct ShapeIter {
    shape: Shape,
    index: usize,
}

impl Iterator for ShapeIter {
    type Item = usize;
    fn next(&mut self) -> Option<Self::Item> {
        let r = self.shape[self.shape.order.get(self.index)?];
        self.index += 1;
        Some(r)
    }
}

/// iterator over the sizes along the axes of a borrowed [`Shape`]
pub struct ShapeIterBorrow<'a> {
    shape: &'a Shape,
    index: usize,
}

impl<'a> Iterator for ShapeIterBorrow<'a> {
    type Item = &'a usize;
    fn next(&mut self) -> Option<Self::Item> {
        let r = &self.shape[self.shape.order.get(self.index)?];
        self.index += 1;
        Some(r)
    }
}

impl IntoIterator for Shape {
    type Item = usize;
    type IntoIter = ShapeIter;

    fn into_iter(self) -> Self::IntoIter {
        ShapeIter {
            shape: self,
            index: 0,
        }
    }
}

impl Shape {
    /// create an empty shape with all sizes set to 1 and no axes
    pub fn new() -> Self {
        Self {
            c: 1,
            z: 1,
            t: 1,
            y: 1,
            x: 1,
            order: vec![],
        }
    }

    /// iterate over the sizes along the axes, in order
    pub fn iter(&self) -> ShapeIterBorrow<'_> {
        ShapeIterBorrow {
            shape: self,
            index: 0,
        }
    }

    /// the number of axes in the shape
    pub fn len(&self) -> usize {
        self.order.len()
    }

    /// whether the shape has no axes
    pub fn is_empty(&self) -> bool {
        self.order.is_empty()
    }

    /// the sizes along the axes, in order
    pub fn to_vec(&self) -> Vec<usize> {
        self.order.iter().map(|axis| self[axis]).collect()
    }

    /// the sizes along the axes as a map from axis to size
    pub fn to_hashmap(&self) -> HashMap<Axis, usize> {
        let mut map = HashMap::new();
        for axis in self.order.iter() {
            map.insert(*axis, self[axis]);
        }
        map
    }

    /// set the size of an axis
    pub fn set_axis(&mut self, axis: &Axis, value: usize) {
        match axis {
            Axis::C => self.c = value,
            Axis::Z => self.z = value,
            Axis::T => self.t = value,
            Axis::Y => self.y = value,
            Axis::X => self.x = value,
            Axis::New => (),
        }
    }
}