//! A generic QuadPrism-shaped dense container indexed by hex coordinates.

use std::convert::TryFrom;
use std::iter::{FusedIterator, Iterator};
use std::ops::{Index, IndexMut};

use crate::quad_prism::PointsAxial;
use crate::{Point, PointAxial, QuadPrism};

/// A generic QuadPrism-shaped dense container indexed by hex coordinates.
///
/// This type is only available with the `collections` feature.
#[derive(Clone, Eq, PartialEq, Debug)]
pub struct HexMap<T, I = i16> {
    data: Vec<T>,
    area: QuadPrism<I>,
    row_stride: usize,
    plane_stride: usize,
}

fn calculate_strides<I: Subscript>(area: &QuadPrism<I>) -> (usize, usize) {
    let size = area.size_axial();
    let row_stride = size.x.try_into().ok().expect("size should be non-negative");
    let plane_stride = row_stride * size.y.try_into().ok().expect("size should be non-negative");
    (row_stride, plane_stride)
}

/// Error indicating a size mismatch between the arguments for `HexMap::try_from_raw`.
#[derive(Clone, Debug, Error)]
#[error("expected a vec of size {area} from given area, got {vec}")]
pub struct RawSizeMismatch {
    area: usize,
    vec: usize,
}

/// Sealed trait for numeric types that can be used as indices.
///
/// Notably, this is implemented for all primitive integer types.
pub trait Subscript: private::Subscript {}
impl<T: private::Subscript> Subscript for T {}

impl<T, I: Subscript> HexMap<T, I> {
    /// Creates from a default value for all hexes. `T` must be `Clone`.
    ///
    /// # Panics
    ///
    /// If area has a negative volume.
    pub fn create(area: QuadPrism<I>, value: T) -> Self
    where
        T: Clone,
    {
        let cap = usize::try_from(area.volume()).expect("volume should be non-negative");
        let mut data = Vec::with_capacity(cap);
        data.resize(cap, value);
        let (row_stride, plane_stride) = calculate_strides(&area);
        HexMap {
            data,
            area,
            row_stride,
            plane_stride,
        }
    }

    /// Creates with an initializer function that takes axial points.
    ///
    /// # Panics
    ///
    /// If area has a negative volume.
    pub fn create_with_axial<F>(area: QuadPrism<I>, mut initializer: F) -> Self
    where
        F: FnMut(PointAxial<I>) -> T,
    {
        let cap = usize::try_from(area.volume()).expect("volume should be non-negative");
        let mut data = Vec::with_capacity(cap);
        let mut iter = area.points_axial();
        data.resize_with(cap, || {
            initializer(iter.next().expect("volume larger than count of points"))
        });
        assert!(iter.next().is_none(), "more points than volume");
        let (row_stride, plane_stride) = calculate_strides(&area);
        HexMap {
            data,
            area,
            row_stride,
            plane_stride,
        }
    }

    /// Creates with an initializer function that takes cube points.
    ///
    /// # Panics
    ///
    /// If area has a negative volume.
    pub fn create_with<F>(area: QuadPrism<I>, mut initializer: F) -> Self
    where
        F: FnMut(Point<I>) -> T,
    {
        Self::create_with_axial(area, |p| initializer(p.to_cube()))
    }

    /// Creates from an existing `HexMap`
    pub fn create_from_axial<U, F>(base: HexMap<U, I>, mut op: F) -> Self
    where
        F: FnMut(PointAxial<I>, U) -> T,
    {
        let data = base
            .area
            .points_axial()
            .zip(base.data.into_iter())
            .map(|(pt, item)| op(pt, item))
            .collect::<Vec<_>>();

        assert_eq!(Ok(data.len()), usize::try_from(base.area.volume()));

        HexMap {
            data,
            area: base.area,
            row_stride: base.row_stride,
            plane_stride: base.plane_stride,
        }
    }

    /// Creates from an existing `HexMap`
    pub fn create_from<U, F>(base: HexMap<U, I>, mut op: F) -> Self
    where
        F: FnMut(Point<I>, U) -> T,
    {
        Self::create_from_axial(base, |pt, data| op(pt.to_cube(), data))
    }

    /// Creates from reference of an existing `HexMap`
    pub fn create_from_ref_axial<U, F>(base: &HexMap<U, I>, mut op: F) -> Self
    where
        F: FnMut(PointAxial<I>, &U) -> T,
    {
        let data = base
            .area
            .points_axial()
            .zip(base.data.iter())
            .map(|(pt, item)| op(pt, item))
            .collect::<Vec<_>>();

        assert_eq!(Ok(data.len()), usize::try_from(base.area.volume()));

        HexMap {
            data,
            area: base.area,
            row_stride: base.row_stride,
            plane_stride: base.plane_stride,
        }
    }

    /// Creates from reference of an existing `HexMap`
    pub fn create_from_ref<U, F>(base: &HexMap<U, I>, mut op: F) -> Self
    where
        F: FnMut(Point<I>, &U) -> T,
    {
        Self::create_from_ref_axial(base, |pt, data| op(pt.to_cube(), data))
    }

    /// Creates from a raw data Vec.
    ///
    /// # Panics
    ///
    /// If the sizes of `area` and `data` mismatch.
    pub fn from_raw(area: QuadPrism<I>, data: Vec<T>) -> Self {
        Self::try_from_raw(area, data).expect("vec should be the exact size as area's volume")
    }

    /// Creates from a raw data Vec.
    ///
    /// # Errors
    ///
    /// If the sizes of `area` and `data` mismatch.
    pub fn try_from_raw(area: QuadPrism<I>, data: Vec<T>) -> Result<Self, RawSizeMismatch> {
        let cap = usize::try_from(area.volume()).expect("volume should be non-negative");
        if cap != data.len() {
            return Err(RawSizeMismatch {
                area: cap,
                vec: data.len(),
            });
        }
        let (row_stride, plane_stride) = calculate_strides(&area);
        Ok(HexMap {
            data,
            area,
            row_stride,
            plane_stride,
        })
    }

    /// Returns the `QuadPrism` that this `HexMap` covers.
    pub fn area(&self) -> &QuadPrism<I> {
        &self.area
    }

    fn to_idx(&self, point: PointAxial<I>) -> Option<usize> {
        if !self.area.contains_axial(&point) {
            return None;
        }
        // positive X, Y, W components guaranteed by the previous check
        let v = (point - self.area.low).cast::<usize>();
        Some(v.w * self.plane_stride + v.y * self.row_stride + v.x)
    }

    /// Returns a reference to a value, or `None` if the point is out of bounds.
    pub fn get_axial(&self, point: PointAxial<I>) -> Option<&T> {
        self.to_idx(point).map(|idx| {
            self.data
                .get(idx)
                .expect("bad index passed contained check")
        })
    }

    /// Returns a reference to a value, or `None` if the point is out of bounds.
    pub fn get(&self, point: Point<I>) -> Option<&T> {
        self.get_axial(point.into_axial())
    }

    /// Returns a mutable reference to a value, or `None` if the point is out of
    /// bounds.
    pub fn get_mut_axial(&mut self, point: PointAxial<I>) -> Option<&mut T> {
        self.to_idx(point).map(move |idx| {
            self.data
                .get_mut(idx)
                .expect("bad index passed contained check")
        })
    }

    /// Returns a mutable reference to a value, or `None` if the point is out of
    /// bounds.
    pub fn get_mut(&mut self, point: Point<I>) -> Option<&mut T> {
        self.get_mut_axial(point.into_axial())
    }

    /// Returns an iterator that yields axial coordinates and references to
    /// corresponding values.
    pub fn iter_axial(&self) -> IterAxial<T, I> {
        IterAxial {
            map: &self,
            points: self.area.points_axial(),
        }
    }

    /// Returns an iterator that yields cube coordinates and references to
    /// corresponding values.
    pub fn iter(&self) -> Iter<T, I> {
        Iter(self.iter_axial())
    }

    /// Returns an iterator that yields axial coordinates and mutable
    /// references to corresponding values.
    pub fn iter_mut_axial(&mut self) -> IterMutAxial<T, I> {
        // SAFETY: IterMutAxial only mutates `data`, not `area`.
        let area = unsafe { &*(&self.area as *const QuadPrism<I>) };
        IterMutAxial {
            map: self,
            points: area.points_axial(),
        }
    }

    /// Returns an iterator that yields cube coordinates and mutable
    /// references to corresponding values.
    pub fn iter_mut(&mut self) -> IterMut<T, I> {
        IterMut(self.iter_mut_axial())
    }
}

impl<T> Default for HexMap<T> {
    fn default() -> Self {
        Self::create_with(Default::default(), |_| {
            unreachable!("initializer should not actually be called when area is empty")
        })
    }
}

impl<T, I: Subscript> Index<PointAxial<I>> for HexMap<T, I> {
    type Output = T;

    fn index(&self, point: PointAxial<I>) -> &T {
        self.get_axial(point).expect("index out of bounds")
    }
}

impl<T, I: Subscript> Index<Point<I>> for HexMap<T, I> {
    type Output = T;

    fn index(&self, point: Point<I>) -> &T {
        self.get(point).expect("index out of bounds")
    }
}

impl<T, I: Subscript> IndexMut<PointAxial<I>> for HexMap<T, I> {
    fn index_mut(&mut self, point: PointAxial<I>) -> &mut T {
        self.get_mut_axial(point).expect("index out of bounds")
    }
}

impl<T, I: Subscript> IndexMut<Point<I>> for HexMap<T, I> {
    fn index_mut(&mut self, point: Point<I>) -> &mut T {
        self.get_mut(point).expect("index out of bounds")
    }
}

/// Iterator with axial coordinates.
#[derive(Debug)]
pub struct IterAxial<'a, T, I> {
    map: &'a HexMap<T, I>,
    points: PointsAxial<'a, I>,
}

impl<'a, T, I: Subscript> Iterator for IterAxial<'a, T, I> {
    type Item = (PointAxial<I>, &'a T);

    fn next(&mut self) -> Option<Self::Item> {
        self.points.next().map(|c| (c, &self.map[c]))
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.points.size_hint()
    }
}

impl<'a, T, I: Subscript> FusedIterator for IterAxial<'a, T, I> {}

/// Mutable iterator with axial coordinates.
#[derive(Debug)]
pub struct IterMutAxial<'a, T, I> {
    map: &'a mut HexMap<T, I>,
    points: PointsAxial<'a, I>,
}

impl<'a, T, I: Subscript> Iterator for IterMutAxial<'a, T, I> {
    type Item = (PointAxial<I>, &'a mut T);

    fn next(&mut self) -> Option<Self::Item> {
        self.points.next().map(|c| {
            // SAFETY: PointsAxial iterates through each point exactly once, and different points translate to different indices in the underlying vector.
            // It's impossible to obtain an aliased mut reference from this iterator.
            (c, unsafe { &mut *((&mut self.map[c]) as *mut T) })
        })
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.points.size_hint()
    }
}

impl<'a, T, I: Subscript> FusedIterator for IterMutAxial<'a, T, I> {}

/// Iterator with cube coordinates.
#[derive(Debug)]
pub struct Iter<'a, T, I>(IterAxial<'a, T, I>);

impl<'a, T, I: Subscript> Iterator for Iter<'a, T, I> {
    type Item = (Point<I>, &'a T);

    fn next(&mut self) -> Option<Self::Item> {
        self.0.next().map(|(c, v)| (c.to_cube(), v))
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.0.size_hint()
    }
}

impl<'a, T, I: Subscript> FusedIterator for Iter<'a, T, I> {}

/// Mutable iterator with cube coordinates.
#[derive(Debug)]
pub struct IterMut<'a, T, I>(IterMutAxial<'a, T, I>);

impl<'a, T, I: Subscript> Iterator for IterMut<'a, T, I> {
    type Item = (Point<I>, &'a mut T);

    fn next(&mut self) -> Option<Self::Item> {
        self.0.next().map(|(c, v)| (c.to_cube(), v))
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.0.size_hint()
    }
}

impl<'a, T, I: Subscript> FusedIterator for IterMut<'a, T, I> {}

mod private {
    pub trait Subscript:
        std::convert::TryInto<usize>
        + num_traits::PrimInt
        + std::ops::AddAssign
        + num_traits::AsPrimitive<usize>
        + num_traits::AsPrimitive<i64>
    {
    }

    impl<T> Subscript for T where
        T: std::convert::TryInto<usize>
            + num_traits::PrimInt
            + std::ops::AddAssign
            + num_traits::AsPrimitive<usize>
            + num_traits::AsPrimitive<i64>
    {
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use crate::{PointAxial as HA, VectorAxial as VA};

    #[test]
    fn can_access() {
        #[allow(clippy::explicit_counter_loop)]
        fn test_sanity(area: QuadPrism<i16>) {
            let mut cloned = HexMap::create(area, 42);
            let mut closure_axial =
                HexMap::create_with_axial(area, |p| p.to_cube().into_vector().length());
            let mut closure_cube = HexMap::create_with(area, |p| p.into_vector().length());

            assert_eq!(area, *cloned.area());
            assert_eq!(area, *closure_axial.area());
            assert_eq!(area, *closure_cube.area());

            {
                let mut i = 0;
                for pt in area.points() {
                    assert_eq!(42, cloned[pt]);
                    assert_eq!(pt.into_vector().length(), closure_axial[pt]);
                    assert_eq!(pt.into_vector().length(), closure_cube[pt]);
                    assert_eq!(Some(&42), cloned.get(pt));
                    assert_eq!(Some(&pt.into_vector().length()), closure_axial.get(pt));
                    assert_eq!(Some(&pt.into_vector().length()), closure_cube.get(pt));

                    i += 1;
                    cloned[pt] = i;
                    closure_axial[pt] = i * 2;
                    closure_cube[pt] = i * 3;
                }
            }

            {
                let mut i = 0;
                for pt in area.points_axial() {
                    i += 1;
                    assert_eq!(i, cloned[pt]);
                    assert_eq!(i * 2, closure_axial[pt]);
                    assert_eq!(i * 3, closure_cube[pt]);
                }
            }
        }

        test_sanity(QuadPrism::from_base_size_axial(
            HA::new(0, 0, 0),
            VA::new(0, 0, 0),
        ));
        test_sanity(QuadPrism::from_base_size_axial(
            HA::new(3, -4, 5),
            VA::new(-5, 4, -3),
        ));
        test_sanity(QuadPrism::from_base_size_axial(
            HA::new(0, 1, 2),
            VA::new(2, 1, 0),
        ));
        test_sanity(QuadPrism {
            low: HA::new(42, 42, 42),
            high: HA::new(-42, -42, -42),
        });
    }

    #[test]
    fn can_default() {
        let map = HexMap::<i32>::default();
        assert_eq!(0, map.area.volume());
        assert_eq!(0, map.data.len());
    }

    #[test]
    fn can_iter() {
        let area = QuadPrism::from_points_axial(HA::new(0, 1, 2), HA::new(3, 4, 5));
        let mut map = HexMap::create_with(area, |p| p.into_vector().length());

        // iteration order should match
        let mut points = area.points();
        for (pt, value) in map.iter() {
            assert_eq!(points.next(), Some(pt));
            assert_eq!(pt.into_vector().length(), *value);
        }

        let mut points = area.points();
        for (pt, value) in map.iter_mut() {
            assert_eq!(points.next(), Some(pt));
            *value *= 42;
        }

        let mut points = area.points();
        for (pt, value) in map.iter() {
            assert_eq!(points.next(), Some(pt));
            assert_eq!(pt.into_vector().length() * 42, *value);
        }
    }
}