use crate::diagnostics::SyncDiagnostics;
use crate::geom::{Coord, Cube, Index, InvertedRegion, Region};
use crate::cross::mat::{GenericMaterial, Material};
use crate::cross::real::Real;
use crate::cross::step::SimMeta;
use crate::cross::vec::{Vec3, Vec3u};
use crate::stim::{Stimulus, NoopStimulus};
use rayon::prelude::*;
use serde::{Serialize, Deserialize};
use std::iter::Sum;

pub mod spirv;
pub mod units;

use spirv::{CpuBackend, SpirvSim};

pub type GenericSim<R> = SpirvSim<R, GenericMaterial<R>, CpuBackend>;


/// 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          |
///    \ |    +------------+------------+
///     \|    |            |            |
///      +    |            |            |
///       \   |            |            |
///        \  |            |            |
///         \ |            |            |
///          \|            |            |
///           +------------+------------+
///
pub struct Sample<'a, R, M> {
    fields: Fields<R>,
    material: &'a M,
}

impl<'a, R: Real, M> Sample<'a, R, M> {
    pub fn fields(&self) -> Fields<R> {
        self.fields
    }
    pub fn material(&self) -> &'a M {
        self.material
    }

    pub fn e(&self) -> Vec3<R> {
        self.fields.e()
    }
    pub fn ex(&self) -> R {
        self.e().x()
    }
    pub fn ey(&self) -> R {
        self.e().y()
    }
    pub fn ez(&self) -> R {
        self.e().z()
    }

    pub fn h(&self) -> Vec3<R> {
        self.fields.h()
    }
    pub fn hx(&self) -> R {
        self.h().x()
    }
    pub fn hy(&self) -> R {
        self.h().y()
    }
    pub fn hz(&self) -> R {
        self.h().z()
    }

    pub fn b(&self) -> Vec3<R> {
        self.fields.b()
    }
    pub fn bx(&self) -> R {
        self.b().x()
    }
    pub fn by(&self) -> R {
        self.b().y()
    }
    pub fn bz(&self) -> R {
        self.b().z()
    }

    pub fn m(&self) -> Vec3<R> {
        self.fields.m()
    }
}

impl<'a, R: Real, M: Material<R>> Sample<'a, R, M> {
    pub fn conductivity(&self) -> Vec3<R> {
        self.material.conductivity()
    }

    pub fn current_density(&self) -> Vec3<R> {
        // TODO: justify/derive this.
        // i guess the actual current density is the gradient of E multiplied by conductivity,
        // and that when summed over a loop the negative terms from the neighbors are canceled
        // when we grab their current density?
        //
        // or maybe it's $V = \int{E . dl} = IR$, so $I = \int{\sigma E . dl}$
        // in which case this might not be "current density", but something related.
        let conductivity = self.conductivity();
        self.e().elem_mul(conductivity)
    }
}

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

impl<R: Real> Fields<R> {
    pub fn new(e: Vec3<R>, h: Vec3<R>, m: Vec3<R>) -> Self {
        Self { e, h, m }
    }
    pub fn cast<R2: Real>(&self) -> Fields<R2> {
        Fields {
            e: self.e.cast(),
            h: self.h.cast(),
            m: self.m.cast(),
        }
    }
    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()
    }
    pub fn with_material<'a, M>(self, material: &'a M) -> Sample<'a, R, M> {
        Sample {
            fields: self,
            material,
        }
    }
}

// TODO: the Sync bound here could be removed with some refactoring
pub trait AbstractSim: Sync {
    type Real: Real;
    type Material: Material<Self::Real>;

    // TODO: should return SimMeta<Self::Real>?
    fn meta(&self) -> SimMeta<f32>;
    fn step_no(&self) -> u64;
    fn fields_at_index(&self, pos: Index) -> Fields<Self::Real>;
    fn get_material_index(&self, at: Index) -> &Self::Material;
    fn put_material_index(&mut self, at: Index, m: Self::Material);
    fn step_multiple<S: Stimulus<Self::Real>>(&mut self, num_steps: u32, s: &S);

    /// convert to something which is maximally generalizable.
    /// the result might then be saved to disc for later readback, where having that consistent
    /// serialization format is more important than perf.
    fn to_generic(&self) -> GenericSim<Self::Real>;
    fn use_diagnostics(&mut self, _diag: SyncDiagnostics) {
        // optional
    }


    //--- HELPER METHODS below (derived) ---//

    fn get_material<C: Coord>(&self, pos: C) -> &Self::Material {
        self.get_material_index(pos.to_index(self.feature_size()))
    }
    fn put_material<C: Coord, M: Into<Self::Material>>(&mut self, pos: C, mat: M) {
        self.put_material_index(pos.to_index(self.feature_size()), mat.into())
    }

    fn sample<'a, C: Coord>(&'a self, pos: C) -> Sample<'a, Self::Real, Self::Material> {
        self.fields_at_index(pos.to_index(self.feature_size()))
            .with_material(self.get_material(pos))
    }

    fn step(&mut self) {
        // XXX: try not to exercise this path! NoopStimulus is probably a lot of waste.
        self.step_multiple(1, &NoopStimulus);
    }

    fn size(&self) -> Index {
        Index(self.meta().dim())
    }
    fn feature_size(&self) -> f32 {
        self.meta().feature_size()
    }
    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 {
        self.meta().time_step()
    }

    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
    }

    // TODO: these should all live off-trait as some sort of `SimExt` thing.

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

    fn map_sum_enumerated<C, F, Ret>(&self, f: F) -> Ret
    where C: Coord,
          F: Fn(C, &Sample<'_, Self::Real, Self::Material>) -> 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.sample(at))
        }))).flatten().flatten().sum()
    }

    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.
    fn map_sum_over<F, Ret, Reg>(&self, region: &Reg, f: F) -> Ret
    where
          F: Fn(&Sample<'_, Self::Real, Self::Material>) -> Ret + Sync,
          Ret: Sum<Ret> + Default + Send,
          Reg: Region + ?Sized
    {
        self.map_sum_over_enumerated(region, |_at: Index, cell| f(cell))
    }

    fn map_sum_over_enumerated<C, F, Ret, Reg>(&self, region: &Reg, f: F) -> Ret
    where C: Coord,
          F: Fn(C, &Sample<'_, Self::Real, Self::Material>) -> 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.sample(at))
            } else {
                Default::default()
            }
        }))).flatten().flatten().sum()
    }

    /// returns the directed current at `c`, in `A / m^2`
    fn current_density<C: Coord>(&self, c: C) -> Vec3<f32> {
        self.sample(c).current_density().cast::<f32>()
    }

    /// returns the directed current at `c` in absolute units, `A`, or rather, `A` per cell, since
    /// this looks at just a single cell. you probably want to use `current_density`.
    fn current<C: Coord>(&self, c: C) -> Vec3<f32> {
        self.current_density(c) * self.feature_size() * self.feature_size()
    }

    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>(&self, region: &Reg, mat: M) -> bool
    where
        M: Into<Self::Material> + Clone,
        Self::Material: PartialEq,
    {
        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))
        });
    }
}