use crate::geom::{Coord, Cube, Index, InvertedRegion, Meters, Region};
use crate::mat;
use crate::types::real::Real;
use crate::types::vec::{Vec3, Vec3u};
use crate::stim::{AbstractStimulus, NoopStimulus};
use ndarray::Zip;
use rayon::prelude::*;
use serde::{Serialize, Deserialize};
use std::iter::Sum;

pub mod legacy;
pub mod spirv;
pub mod units;

pub use legacy::StaticSim;


pub trait MaterialSim: GenericSim {
    type Material: PartialEq;

    fn put_material<C: Coord, M: Into<Self::Material>>(&mut self, pos: C, mat: M);
    // XXX: would ideally return by-ref, but some backends need to return a handle instead
    fn get_material<C: Coord>(&self, pos: C) -> Self::Material;

    fn fill_region_using<C, Reg, F, M>(&mut self, region: &Reg, f: F)
    where
        Reg: Region,
        F: Fn(C) -> M,
        C: Coord,
        M: Into<Self::Material>
    {
        for z in 0..self.depth() {
            for y in 0..self.height() {
                for x in 0..self.width() {
                    let loc = Index((x, y, z).into());
                    let meters = loc.to_meters(self.feature_size());
                    if region.contains(meters) {
                        self.put_material(loc, f(C::from_either(loc, meters)));
                    }
                }
            }
        }
    }

    fn fill_region<Reg: Region, M: Into<Self::Material> + Clone>(&mut self, region: &Reg, mat: M) {
        self.fill_region_using(region, |_idx: Index| mat.clone());
    }

    fn examine_region<C, Reg, F>(&self, region: &Reg, mut f: F)
    where
        Reg: Region,
        F: FnMut(C, &Self::Material),
        C: Coord
    {
        for z in 0..self.depth() {
            for y in 0..self.height() {
                for x in 0..self.width() {
                    let loc = Index((x, y, z).into());
                    let meters = loc.to_meters(self.feature_size());
                    if region.contains(meters) {
                        f(C::from_either(loc, meters), &self.get_material(loc));
                    }
                }
            }
        }
    }

    /// Return true if the given region is filled exclusively with the provided material.
    fn test_region_filled<Reg: Region, M: Into<Self::Material> + Clone>(&self, region: &Reg, mat: M) -> bool {
        let mut all = true;
        self.examine_region(region, |_idx: Index, m: &Self::Material| {
            all = all && m == &mat.clone().into();
        });
        all
    }

    /// Fill the boundary, where `thickness` describes how far the boundary extends in each
    /// direction, and `f` takes a vec where each coordinate represents how far into the boundary
    /// the location being queried is, in each direction.
    /// e.g. `f((1.0, 0.0, 0.2))` means the location being queried is at either extreme end on the
    /// x axis, is not inside the y axis boundary, and is 20% of the way from the onset of the z
    /// boundary to the edge of the z world.
    fn fill_boundary_using<C, F, M>(&mut self, thickness: C, f: F)
    where
        C: Coord,
        F: Fn(Vec3<f32>) -> M,
        M: Into<Self::Material>,
    {
        // TODO: maybe this function belongs on the Driver?
        let feat = self.feature_size();
        let upper_left = thickness.to_index(feat);
        let size = self.size();
        let lower_right = size - upper_left - Index::new(1, 1, 1);
        let region = InvertedRegion::new(Cube::new(upper_left.to_meters(feat), lower_right.to_meters(feat)));
        self.fill_region_using(&region, |loc: Index| {
            let depth_x = if loc.x() < upper_left.x() {
                (upper_left.x() - loc.x()) as f32 / upper_left.x() as f32
            } else if loc.x() > lower_right.x() {
                (loc.x() - lower_right.x()) as f32 / upper_left.x() as f32
            } else {
                0.0
            };
            let depth_y = if loc.y() < upper_left.y() {
                (upper_left.y() - loc.y()) as f32 / upper_left.y() as f32
            } else if loc.y() > lower_right.y() {
                (loc.y() - lower_right.y()) as f32 / upper_left.y() as f32
            } else {
                0.0
            };
            let depth_z = if loc.z() < upper_left.z() {
                (upper_left.z() - loc.z()) as f32 / upper_left.z() as f32
            } else if loc.z() > lower_right.z() {
                (loc.z() - lower_right.z()) as f32 / upper_left.z() as f32
            } else {
                0.0
            };
            // println!("{} {}", loc, Vec3::new(depth_x, depth_y, depth_z));
            f(Vec3::new(depth_x, depth_y, depth_z))
        });
    }
}

// XXX the Send/Sync bounds here could be removed with some refactoring
pub trait GenericSim: SampleableSim {
    fn step(&mut self) {
        // XXX: try not to exercise this path! NoopStimulus is probably a lot of waste.
        self.step_multiple(1, &NoopStimulus);
    }
    fn step_multiple<S: AbstractStimulus>(&mut self, num_steps: u32, s: &S);

    /// DEPRECATED. Use stimulus instead
    fn impulse_e_meters(&mut self, pos: Meters, amount: Vec3<f32>);
    /// DEPRECATED. Use stimulus instead
    fn impulse_h_meters(&mut self, pos: Meters, amount: Vec3<f32>) {
        self.impulse_b_meters(pos, amount * f32::mu0());
    }
    /// DEPRECATED. Use stimulus instead
    fn impulse_b_meters(&mut self, pos: Meters, amount: Vec3<f32>) {
        self.impulse_h_meters(pos, amount * f32::mu0_inv());
    }

    fn impulse_e<C: Coord>(&mut self, pos: C, amt: Vec3<f32>) {
        self.impulse_e_meters(pos.to_meters(self.feature_size()), amt)
    }

    fn impulse_h<C: Coord>(&mut self, pos: C, amt: Vec3<f32>) {
        self.impulse_h_meters(pos.to_meters(self.feature_size()), amt)
    }

    fn impulse_b<C: Coord>(&mut self, pos: C, amt: Vec3<f32>) {
        self.impulse_b_meters(pos.to_meters(self.feature_size()), amt)
    }

    fn impulse_ex<C: Coord>(&mut self, c: C, ex: f32) {
        self.impulse_e(c, Vec3::new_x(ex));
    }
    fn impulse_ey<C: Coord>(&mut self, c: C, ey: f32) {
        self.impulse_e(c, Vec3::new_y(ey));
    }
    fn impulse_ez<C: Coord>(&mut self, c: C, ez: f32) {
        self.impulse_e(c, Vec3::new_z(ez));
    }

    fn impulse_hx<C: Coord>(&mut self, c: C, hx: f32) {
        self.impulse_h(c, Vec3::new_x(hx));
    }
    fn impulse_hy<C: Coord>(&mut self, c: C, hy: f32) {
        self.impulse_h(c, Vec3::new_y(hy));
    }
    fn impulse_hz<C: Coord>(&mut self, c: C, hz: f32) {
        self.impulse_h(c, Vec3::new_z(hz));
    }

    fn impulse_bx<C: Coord>(&mut self, c: C, bx: f32) {
        self.impulse_b(c, Vec3::new_x(bx));
    }
    fn impulse_by<C: Coord>(&mut self, c: C, by: f32) {
        self.impulse_b(c, Vec3::new_y(by));
    }
    fn impulse_bz<C: Coord>(&mut self, c: C, bz: f32) {
        self.impulse_b(c, Vec3::new_z(bz));
    }
}

/// Conceptually, one cell looks like this (in 2d):
///
/// +-------.-------+
/// |       Ex      |
/// |               |
/// |               |
/// .Ey     .Bz     .
/// |               |
/// |               |
/// |               |
/// +-------.-------+
///
/// Where the right hand rule indicates that positive Bz is pointing into the page, away from the
/// reader.
///
/// The dot on bottom is Ex of the cell at (x, y+1) and the dot on the right is the Ey of the cell at
/// (x+1, y). The `+` only indicates the corner of the cell -- nothing of interest is measured at
/// the pluses.
///
/// To adapt this to 3d, we keep the above cross-section, and then add another cross-section at a
/// half cell depth into the page:
/// 3d version (there may be multiple solutions):
///
/// Ez------.-------+
/// |       By      |
/// |               |
/// |               |
/// .Bx             .
/// |               |
/// |               |
/// |               |
/// +-------.-------+
///
/// Altogether, each cell is a cube, with the field vectors lying on the surface of the 1/8th of
/// the cube closest to the asterisk, at (0, 0, 0)
///
/// +------------+------------+
/// |\            \            \
/// | \            \            \
/// |  \            \            \
/// |   \ Ez         \ By         \
/// |    +------------+------------+
/// |    |\            \            \                  z
/// +    | \            \            \                  \
/// |\   |  \            \            \                  \
/// | \  |   \            \ Ex         \                  \
/// |  \ |    *------------+------------+                  *--- x
/// |   \|Bx  |            |            |                  |
/// |    +    |            |            |                  |
/// |    |\   |            |            |                  |
/// +    | \  |            |            |                  y
///  \   |  \ |            |            |
///   \  |   \|Ey          |Bz          |
///    \ |    +------------+------------+
///     \|    |            |            |
///      +    |            |            |
///       \   |            |            |
///        \  |            |            |
///         \ |            |            |
///          \|            |            |
///           +------------+------------+
///
#[derive(Clone, Default, Serialize, Deserialize)]
pub struct Sample {
    state: CellStateWithM<f32>,
    conductivity: Vec3<f32>,
}

#[derive(Copy, Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct CellStateWithM<R> {
    e: Vec3<R>,
    h: Vec3<R>,
    m: Vec3<R>,
}

impl<R: Real> CellStateWithM<R> {
    pub fn e(&self) -> Vec3<R> {
        self.e
    }
    pub fn h(&self) -> Vec3<R> {
        self.h
    }
    pub fn m(&self) -> Vec3<R> {
        self.m
    }
    pub fn b(&self) -> Vec3<R> {
        (self.h() + self.m()) * R::mu0()
    }
}

impl<'a> dyn SampleableSim + 'a {
    pub fn get<C: Coord>(&self, at: C) -> Sample {
        self.sample(at.to_meters(self.feature_size()))
    }

    /// Apply `F` to each Cell, and sum the results.
    pub fn map_sum<F, Ret>(&self, f: F) -> Ret
    where
          F: Fn(&Sample) -> Ret + Sync,
          Ret: Sum<Ret> + Send,
    {
        self.map_sum_enumerated(|_at: Index, cell| f(cell))
    }

    pub fn map_sum_enumerated<C, F, Ret>(&self, f: F) -> Ret
    where C: Coord,
          F: Fn(C, &Sample) -> Ret + Sync,
          Ret: Sum<Ret> + Send,
    {
        let (w, h, d) = (self.width(), self.height(), self.depth());
        (0..d).into_par_iter().map(
            |z| (0..h).into_par_iter().map_with(z,
            |&mut z, y| (0..w).into_par_iter().map_with((z, y),
            |&mut (z, y), x|
        {
            let at = Index(Vec3u::new(x, y, z));
            f(C::from_index(at, self.feature_size()), &self.get(at))
        }))).flatten().flatten().sum()
    }

    pub fn volume_of_region<R: Region + ?Sized>(&self, region: &R) -> u32 {
        self.map_sum_over(region, |_| 1)
    }

    /// Apply `F` to each Cell, and sum the results.
    pub fn map_sum_over<F, Ret, Reg>(&self, region: &Reg, f: F) -> Ret
    where
          F: Fn(&Sample) -> Ret + Sync,
          Ret: Sum<Ret> + Default + Send,
          Reg: Region + ?Sized
    {
        self.map_sum_over_enumerated(region, |_at: Index, cell| f(cell))
    }

    pub fn map_sum_over_enumerated<C, F, Ret, Reg>(&self, region: &Reg, f: F) -> Ret
    where C: Coord,
          F: Fn(C, &Sample) -> Ret + Sync,
          Ret: Sum<Ret> + Default + Send,
          Reg: Region + ?Sized,
    {
        let (w, h, d) = (self.width(), self.height(), self.depth());
        (0..d).into_par_iter().map(
            |z| (0..h).into_par_iter().map_with(z,
            |&mut z, y| (0..w).into_par_iter().map_with((z, y),
            |&mut (z, y), x|
        {
            let at = Index(Vec3u::new(x, y, z));
            let meters = at.to_meters(self.feature_size());
            if region.contains(meters) {
                f(C::from_index(at, self.feature_size()), &self.get(at))
            } else {
                Default::default()
            }
        }))).flatten().flatten().sum()
    }

    pub fn current<C: Coord>(&self, c: C) -> Vec3<f32> {
        self.get(c).current_density() * self.feature_size() * self.feature_size()
    }
}

// XXX the Send/Sync bounds here could be removed with some refactoring
pub trait SampleableSim: Send + Sync {
    fn sample(&self, pos: Meters) -> Sample;

    fn size(&self) -> Index;
    fn feature_size(&self) -> f32;
    fn feature_volume(&self) -> f32 {
        let f = self.feature_size();
        f*f*f
    }
    fn volume(&self) -> f32 {
        let s = self.size().to_meters(self.feature_size());
        s.x() * s.y() * s.z()
    }
    fn timestep(&self) -> f32;

    fn step_no(&self) -> u64;

    fn width(&self) -> u32 {
        self.size().x()
    }
    fn height(&self) -> u32 {
        self.size().y()
    }
    fn depth(&self) -> u32 {
        self.size().z()
    }
    fn time(&self) -> f32 {
        self.timestep() * self.step_no() as f32
    }

    /// Take a "snapshot" of the simulation, dropping all material-specific information.
    fn to_static(&self) -> StaticSim {
        let mut state = legacy::SimState::new(self.size(), self.feature_size());
        state.step_no = self.step_no();
        Zip::from(ndarray::indices_of(&state.cells)).and(&mut state.e).and(&mut state.h).par_map_assign_into(
            &mut state.cells,
            |(z, y, x), e, h| {
                let idx = Index((x as u32, y as u32, z as u32).into());
                let cell = self.sample(idx.to_meters(self.feature_size()));
                *e = cell.e().cast();
                *h = cell.h().cast();
                mat::Static {
                    conductivity: cell.conductivity().cast(),
                    m: cell.m().cast(),
                }
            });
        state
    }
}