fem_2d 0.2.2

use super::MeshAccessError;
use super::ParaDir;
use json::{object, JsonValue};
use std::fmt;
use std::ops::AddAssign;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
/// Description of an `Elem`s h-Refinement levels in the u and v directions
///
/// These values facilitate the matching of basis functions among neighboring `Elem`s
pub struct HLevels {
    pub u: u8,
    pub v: u8,
}

impl HLevels {
    fn from(u: u8, v: u8) -> Self {
        Self { u, v }
    }

    /// Construct a new set of [HLevels] from an [HRef] of `self`
    pub fn refined(&self, refinement: HRef) -> Self {
        match refinement {
            HRef::T => Self::from(self.u + 1, self.v + 1),
            HRef::U(_) => Self::from(self.u + 1, self.v),
            HRef::V(_) => Self::from(self.u, self.v + 1),
        }
    }

    /// The ranking of the associated `Elem` according to the local `Edge`s (and their direction)
    pub fn edge_ranking(self, edge_dir: ParaDir) -> [u8; 2] {
        match edge_dir {
            ParaDir::U => [self.v, self.u],
            ParaDir::V => [self.u, self.v],
        }
    }

    /// The ranking of the associated `Elem` according to the local `Node`s
    pub fn node_ranking(self) -> [u8; 2] {
        [self.u, self.v]
    }
}

impl Default for HLevels {
    fn default() -> Self {
        Self { u: 0, v: 0 }
    }
}

#[cfg(feature = "json_export")]
impl From<HLevels> for JsonValue {
    fn from(h: HLevels) -> Self {
        object! {
            "u": h.u,
            "v": h.v,
        }
    }
}

/// The h-Refinement Type
///
/// h-refinements are used to introduce new `Elem`s into the mesh
/// * T: 4 new `Elem`s are introduced over a given `Elem`
/// * U(None): 2 new `Elem`s are introduced over a given `Elem` (improving resolution in the u-direction)
/// * V(None): 2 new `Elem`s are introduced over a given `Elem` (improving resolution in the v-direction)
/// * U(Some( 0 or 1 )): A U-type refinement, followed by a V-type refinement of the 0th or 1st resultant `Elem`
/// * V(Some( 0 or 1 )): A V-type refinement, followed by a U-type refinement of the 0th or 1st resultant `Elem`
///
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum HRef {
    /// isotropic
    T,
    /// anisotropic about the u-direction (with the option for the subsequent v-refinement of one child Elem)
    U(Option<usize>),
    /// anisotropic about the v-direction (with the option for the subsequent u-refinement of one child Elem)
    V(Option<usize>),
}

impl HRef {
    /// Basic T-Type refinement
    pub const fn t() -> Self {
        Self::T
    }

    /// Basic U-Type refinement
    pub const fn u() -> Self {
        Self::U(None)
    }

    /// Basic V-Type refinement
    pub const fn v() -> Self {
        Self::V(None)
    }

    /// Extended U-type refinement
    ///
    /// The 0th or 1st child generated by the U-Type refinement is then V-Type refined
    /// Returns an error if `child_idx` is not 0 or 1
    pub const fn u_extened(child_idx: u8) -> Result<Self, HRefError> {
        match child_idx {
            0 => Ok(Self::U(Some(0))),
            1 => Ok(Self::U(Some(1))),
            _ => Err(HRefError::BisectionIdxExceeded),
        }
    }

    /// Extended V-type refinement
    ///
    /// The 0th or 1st child generated by the V-Type refinement is then U-Type refined
    /// Returns an error if `child_idx` is not 0 or 1
    pub const fn v_extened(child_idx: u8) -> Result<Self, HRefError> {
        match child_idx {
            0 => Ok(Self::V(Some(0))),
            1 => Ok(Self::V(Some(1))),
            _ => Err(HRefError::BisectionIdxExceeded),
        }
    }

    pub(crate) fn indices_and_ids(
        &self,
        id_counter: &mut usize,
    ) -> Box<dyn Iterator<Item = (usize, usize)> + '_> {
        let starting_id = *id_counter;
        match self {
            Self::T => {
                *id_counter += 4;
                Box::new((0..4).map(move |idx| (idx, starting_id + idx)))
            }
            Self::U(_) | Self::V(_) => {
                *id_counter += 2;
                Box::new((0..2).map(move |idx| (idx, starting_id + idx)))
            }
        }
    }

    pub(crate) fn loc(&self, idx: usize) -> HRefLoc {
        match self {
            Self::T => match idx {
                0 => HRefLoc::SW,
                1 => HRefLoc::SE,
                2 => HRefLoc::NW,
                3 => HRefLoc::NE,
                _ => unreachable!(),
            },
            Self::U(_) => match idx {
                0 => HRefLoc::W,
                1 => HRefLoc::E,
                _ => unreachable!(),
            },
            Self::V(_) => match idx {
                0 => HRefLoc::S,
                1 => HRefLoc::N,
                _ => unreachable!(),
            },
        }
    }
}

impl AddAssign for HRef {
    fn add_assign(&mut self, rhs: Self) {
        match self {
            HRef::T => (),
            HRef::U(None) => match rhs {
                HRef::V(_) => *self = HRef::T,
                HRef::T => *self = HRef::T,
                _ => (),
            },
            HRef::V(None) => match rhs {
                HRef::U(_) => *self = HRef::T,
                HRef::T => *self = HRef::T,
                _ => (),
            },
            HRef::U(Some(idx)) => match rhs {
                HRef::U(Some(r_idx)) => {
                    if r_idx != *idx {
                        *self = HRef::T;
                    }
                }
                HRef::V(_) => *self = HRef::T,
                HRef::T => *self = HRef::T,
                _ => (),
            },
            HRef::V(Some(idx)) => match rhs {
                HRef::V(Some(r_idx)) => {
                    if r_idx != *idx {
                        *self = HRef::T;
                    }
                }
                HRef::U(_) => *self = HRef::T,
                HRef::T => *self = HRef::T,
                _ => (),
            },
        }
    }
}

impl fmt::Display for HRef {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "HRef ")?;
        match self {
            HRef::T => write!(f, "T"),
            HRef::U(None) => write!(f, "U"),
            HRef::V(None) => write!(f, "V"),
            HRef::U(Some(idx)) => write!(f, "U->V({})", idx),
            HRef::V(Some(idx)) => write!(f, "V->U({})", idx),
        }
    }
}

/// The location of a child `Elem` relative to its parent following an [HRef]
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum HRefLoc {
    /// South West: T(0)
    SW,
    /// South East: T(1)
    SE,
    /// North West: T(2)
    NW,
    /// North East: T(3)
    NE,
    /// West: U(0)
    W,
    /// East: U(1)
    E,
    /// South: V(0)
    S,
    /// North: V(1)
    N,
}

impl HRefLoc {
    /// Index of the child Elem relative to its parent according to the convention defined by the `Elem` Type
    pub fn index(&self) -> usize {
        match self {
            Self::SW => 0,
            Self::SE => 1,
            Self::NW => 2,
            Self::NE => 3,
            Self::W => 0,
            Self::E => 1,
            Self::S => 0,
            Self::N => 1,
        }
    }

    /// Given an input range in the parents parametric space, what sub space does the child `Elem` occupy?
    pub fn sub_range(&self, [u_range, v_range]: [[f64; 2]; 2]) -> [[f64; 2]; 2] {
        [self.sub_u_range(u_range), self.sub_v_range(v_range)]
    }

    /// Given an input range along the parent's `u` parametric axis, what sub-range does the child `Elem` occupy?
    pub fn sub_u_range(&self, [min_u, max_u]: [f64; 2]) -> [f64; 2] {
        let middle = (min_u + max_u) / 2.0;
        match self {
            Self::SW => [min_u, middle],
            Self::SE => [middle, max_u],
            Self::NW => [min_u, middle],
            Self::NE => [middle, max_u],
            Self::W => [min_u, middle],
            Self::E => [middle, max_u],
            Self::S => [min_u, max_u],
            Self::N => [min_u, max_u],
        }
    }

    /// Given an input range along the parent's `v` parametric axis, what sub-range does the child `Elem` occupy?
    pub fn sub_v_range(&self, [min_v, max_v]: [f64; 2]) -> [f64; 2] {
        let middle = (min_v + max_v) / 2.0;
        match self {
            Self::SW => [min_v, middle],
            Self::SE => [min_v, middle],
            Self::NW => [middle, max_v],
            Self::NE => [middle, max_v],
            Self::W => [min_v, max_v],
            Self::E => [min_v, max_v],
            Self::S => [min_v, middle],
            Self::N => [middle, max_v],
        }
    }
}

/// The Error Type for invalid h-refinements
#[derive(Debug)]
pub enum HRefError {
    // Errors caused by internal problems with the Mesh (Should never happen)
    MinEdgeLength(usize),
    ElemHasChildren(usize),
    EdgeHasChildren(usize),
    UninitializedElem(usize),
    EdgeOnEqualPoints(usize),
    BisectionIdxExceeded,
    // Public Errors
    ElemDoesNotExist(usize),
    ElemNotRefineable(usize),
    DuplicateElemIds,
}

impl std::error::Error for HRefError {}

impl fmt::Display for HRefError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Self::MinEdgeLength(edge_id) => write!(
                f,
                "h-refinement will result in Edge length below minimum; Cannot h-refine Edge {}!",
                edge_id
            ),
            Self::ElemHasChildren(elem_id) => {
                write!(f, "Elem {} already has children; Cannot h-refine!", elem_id)
            }
            Self::EdgeHasChildren(edge_id) => {
                write!(f, "Edge {} already has children; Cannot h-refine!", edge_id)
            }
            Self::UninitializedElem(elem_uninit_id) => write!(
                f,
                "ElemUninit {} was not fully initialized by the conclusion of h-refinement",
                elem_uninit_id
            ),
            Self::ElemDoesNotExist(elem_id) => write!(
                f,
                "Elem {} does not exist; cannot apply h-Refinement!",
                elem_id
            ),
            Self::ElemNotRefineable(elem_id) => write!(
                f,
                "Elem {} cannot be h-refined; it is either too small or has already been refined!",
                elem_id,
            ),
            Self::DuplicateElemIds => {
                write!(f, "Duplicate element ids in h-Refinement; Cannot h-Refine!")
            }
            Self::EdgeOnEqualPoints(elem_id) => write!(
                f,
                "Attempt to generate a child-Edge between two identical points over Elem {}!",
                elem_id
            ),
            Self::BisectionIdxExceeded => write!(f, "Extended refinement index must be 0 or 1!"),
        }
    }
}

impl From<MeshAccessError> for HRefError {
    fn from(err: MeshAccessError) -> Self {
        match err {
            MeshAccessError::ElemDoesNotExist(elem_id) => Self::ElemDoesNotExist(elem_id),
            _ => unreachable!(),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn h_ref_loc() {
        const EXPECTED_ANISO_COORDS: [[[f64; 2]; 2]; 5] = [
            [[-1.0, 1.0], [-1.0, 1.0]],
            [[-1.0, 1.0], [0.0, 1.0]],
            [[0.0, 1.0], [0.0, 1.0]],
            [[0.0, 1.0], [0.0, 0.5]],
            [[0.0, 0.5], [0.0, 0.5]],
        ];

        const EXPECTED_ISO_COORDS: [[[f64; 2]; 2]; 5] = [
            [[-1.0, 1.0], [-1.0, 1.0]],
            [[-1.0, 0.0], [-1.0, 0.0]],
            [[-0.5, 0.0], [-1.0, -0.5]],
            [[-0.5, -0.25], [-0.75, -0.5]],
            [[-0.375, -0.25], [-0.625, -0.5]],
        ];

        let loc_stack_aniso = vec![HRefLoc::N, HRefLoc::E, HRefLoc::S, HRefLoc::W];
        let loc_stack_iso = vec![HRefLoc::SW, HRefLoc::SE, HRefLoc::NW, HRefLoc::NE];

        let final_aniso =
            loc_stack_aniso
                .iter()
                .enumerate()
                .fold([[-1.0, 1.0], [-1.0, 1.0]], |acc, (i, loc)| {
                    for uv in 0..2 {
                        for mm in 0..2 {
                            assert!((acc[uv][mm] - EXPECTED_ANISO_COORDS[i][uv][mm]).abs() < 1e-14);
                        }
                    }
                    loc.sub_range(acc)
                });

        for uv in 0..2 {
            for mm in 0..2 {
                assert!(
                    (final_aniso[uv][mm] - EXPECTED_ANISO_COORDS.last().unwrap()[uv][mm]).abs()
                        < 1e-14
                );
            }
        }

        let final_iso =
            loc_stack_iso
                .iter()
                .enumerate()
                .fold([[-1.0, 1.0], [-1.0, 1.0]], |acc, (i, loc)| {
                    for uv in 0..2 {
                        for mm in 0..2 {
                            assert!((acc[uv][mm] - EXPECTED_ISO_COORDS[i][uv][mm]).abs() < 1e-14);
                        }
                    }
                    loc.sub_range(acc)
                });

        for uv in 0..2 {
            for mm in 0..2 {
                assert!(
                    (final_iso[uv][mm] - EXPECTED_ISO_COORDS.last().unwrap()[uv][mm]).abs() < 1e-14
                );
            }
        }
    }
}