// use spirv_std::RuntimeArray;
use spirv_std::glam::UVec3;
use crate::mat::Material;
use crate::support::{Array3, Array3Mut, ArrayHandle, ArrayHandleMut, Optional, UnsizedArray, UVec3Std, Vec3Std};

#[derive(Copy, Clone)]
pub struct SerializedSimMeta {
    // XXX this HAS to be the tuple version instead of the glam dim, else:
    //   OpLoad Pointer <id> '138[%138]' is not a logical pointer.
    pub dim: UVec3Std,
    pub inv_feature_size: f32,
    pub time_step: f32,
    pub feature_size: f32,
}

impl SerializedSimMeta {
    fn dim(&self) -> UVec3 {
        self.dim.into()
    }
}

/// Whatever data we received from the host in their call to step_h
pub struct SerializedStepH<'a, M> {
    meta: &'a SerializedSimMeta,
    stimulus_h: &'a UnsizedArray<Vec3Std>,
    material: &'a UnsizedArray<M>,
    e: &'a UnsizedArray<Vec3Std>,
    h: &'a mut UnsizedArray<Vec3Std>,
    m: &'a mut UnsizedArray<Vec3Std>,
}

impl<'a, M> SerializedStepH<'a, M> {
    pub fn new(
        meta: &'a SerializedSimMeta,
        stimulus_h: &'a UnsizedArray<Vec3Std>,
        material: &'a UnsizedArray<M>,
        e: &'a UnsizedArray<Vec3Std>,
        h: &'a mut UnsizedArray<Vec3Std>,
        m: &'a mut UnsizedArray<Vec3Std>,
    ) -> Self {
        Self {
            meta, stimulus_h, material, e, h, m
        }
    }

    pub fn index(self, idx: UVec3) -> StepHContext<'a, M> {
        let dim = self.meta.dim();
        let stim_h_matrix = Array3::new(self.stimulus_h, dim);
        let mat_matrix = Array3::new(self.material, dim);
        let e = Array3::new(self.e, dim);
        let h = Array3Mut::new(self.h, dim);
        let m = Array3Mut::new(self.m, dim);

        let in_e = VolumeSamplePos {
            mid: e.get(idx).unwrap(),
            xp1: e.get(idx + UVec3::X),
            yp1: e.get(idx + UVec3::Y),
            zp1: e.get(idx + UVec3::Z),
        };
        let out_h = h.into_mut_handle(idx);
        let out_m = m.into_mut_handle(idx);

        let mat = mat_matrix.into_handle(idx);

        StepHContext {
            inv_feature_size: self.meta.inv_feature_size,
            time_step: self.meta.time_step,
            stim_h: stim_h_matrix.get(idx).unwrap(),
            mat,
            in_e,
            out_h,
            out_m,
        }
    }
}

/// Whatever data we received from the host in their call to step_e
pub struct SerializedStepE<'a, M> {
    meta: &'a SerializedSimMeta,
    stimulus_e: &'a UnsizedArray<Vec3Std>,
    material: &'a UnsizedArray<M>,
    e: &'a mut UnsizedArray<Vec3Std>,
    h: &'a UnsizedArray<Vec3Std>,
}

impl<'a, M> SerializedStepE<'a, M> {
    pub fn new(
        meta: &'a SerializedSimMeta,
        stimulus_e: &'a UnsizedArray<Vec3Std>,
        material: &'a UnsizedArray<M>,
        e: &'a mut UnsizedArray<Vec3Std>,
        h: &'a UnsizedArray<Vec3Std>,
    ) -> Self {
        Self {
            meta, stimulus_e, material, e, h
        }
    }

    pub fn index(self, idx: UVec3) -> StepEContext<'a, M> {
        let dim = self.meta.dim();
        let stim_e_matrix = Array3::new(self.stimulus_e, dim);
        let mat_matrix = Array3::new(self.material, dim);
        let e = Array3Mut::new(self.e, dim);
        let h = Array3::new(self.h, dim);

        let xm1 = if idx.x == 0 {
            Optional::none()
        } else {
            h.get(idx - UVec3::X)
        };
        let ym1 = if idx.y == 0 {
            Optional::none()
        } else {
            h.get(idx - UVec3::Y)
        };
        let zm1 = if idx.z == 0 {
            Optional::none()
        } else {
            h.get(idx - UVec3::Z)
        };

        let in_h = VolumeSampleNeg {
            mid: h.get(idx).unwrap(),
            xm1,
            ym1,
            zm1,
        };
        let out_e = e.into_mut_handle(idx);

        let mat = mat_matrix.into_handle(idx);

        StepEContext {
            inv_feature_size: self.meta.inv_feature_size,
            time_step: self.meta.time_step,
            stim_e: stim_e_matrix.get(idx).unwrap(),
            mat,
            in_h,
            out_e,
        }
    }
}

/// Package the field vectors adjacent to some particular location.
/// Particular those at negative offsets from the midpoint.
/// This is used in step_e when looking at the H field deltas.
#[derive(Copy, Clone)]
struct VolumeSampleNeg {
    mid: Vec3Std,
    xm1: Optional<Vec3Std>,
    ym1: Optional<Vec3Std>,
    zm1: Optional<Vec3Std>,
}

impl VolumeSampleNeg {
    /// Calculate the delta in H values amongst this cell and its neighbors (left/up/out)
    fn delta_h(self) -> FieldDeltas {
        let mid = self.mid;
        // let (dfy_dx, dfz_dx) = self.xm1.map(|xm1| {
        //     (mid.y() - xm1.y(), mid.z() - xm1.z())
        // }).unwrap_or_default();

        // let (dfx_dy, dfz_dy) = self.ym1.map(|ym1| {
        //     (mid.x() - ym1.x(), mid.z() - ym1.z())
        // }).unwrap_or_default();

        // let (dfx_dz, dfy_dz) = self.zm1.map(|zm1| {
        //     (mid.x() - zm1.x(), mid.y() - zm1.y())
        // }).unwrap_or_default();

        let (dfy_dx, dfz_dx) = if self.xm1.is_some() {
            (mid.y() - self.xm1.unwrap().y(), mid.z() - self.xm1.unwrap().z())
        } else {
            (0.0, 0.0)
        };

        let (dfx_dy, dfz_dy) = if self.ym1.is_some() {
            (mid.x() - self.ym1.unwrap().x(), mid.z() - self.ym1.unwrap().z())
        } else {
            (0.0, 0.0)
        };

        let (dfx_dz, dfy_dz) = if self.zm1.is_some() {
            (mid.x() - self.zm1.unwrap().x(), mid.y() - self.zm1.unwrap().y())
        } else {
            (0.0, 0.0)
        };
        
        FieldDeltas {
            dfy_dx,
            dfz_dx,
            dfx_dy,
            dfz_dy,
            dfx_dz,
            dfy_dz,
        }
    }
}


/// Package the field vectors adjacent to some particular location.
/// Particular those at positive offsets from the midpoint.
/// This is used in step_h when looking at the E field deltas.
#[derive(Copy, Clone)]
struct VolumeSamplePos {
    mid: Vec3Std,
    xp1: Optional<Vec3Std>,
    yp1: Optional<Vec3Std>,
    zp1: Optional<Vec3Std>
}

impl VolumeSamplePos {
    /// Calculate the delta in E values amongst this cell and its neighbors (right/down/in)
    fn delta_e(self) -> FieldDeltas {
        let mid = self.mid;
        // let (dfy_dx, dfz_dx) = self.xp1.map(|xp1| {
        //     (xp1.y() - mid.y(), xp1.z() - mid.z())
        // }).unwrap_or_default();

        // let (dfx_dy, dfz_dy) = self.yp1.map(|yp1| {
        //     (yp1.x() - mid.x(), yp1.z() - mid.z())
        // }).unwrap_or_default();

        // let (dfx_dz, dfy_dz) = self.zp1.map(|zp1| {
        //     (zp1.x() - mid.x(), zp1.y() - mid.y())
        // }).unwrap_or_default();

        let (dfy_dx, dfz_dx) = if self.xp1.is_some() {
            (self.xp1.unwrap().y() - mid.y(), self.xp1.unwrap().z() - mid.z())
        } else {
            (0.0, 0.0)
        };

        let (dfx_dy, dfz_dy) = if self.yp1.is_some() {
            (self.yp1.unwrap().x() - mid.x(), self.yp1.unwrap().z() - mid.z())
        } else {
            (0.0, 0.0)
        };

        let (dfx_dz, dfy_dz) = if self.zp1.is_some() {
            (self.zp1.unwrap().x() - mid.x(), self.zp1.unwrap().y() - mid.y())
        } else {
            (0.0, 0.0)
        };
        
        FieldDeltas {
            dfy_dx,
            dfz_dx,
            dfx_dy,
            dfz_dy,
            dfx_dz,
            dfy_dz,
        }
    }
}

struct FieldDeltas {
    dfy_dx: f32,
    dfz_dx: f32,
    dfx_dy: f32,
    dfz_dy: f32,
    dfx_dz: f32,
    dfy_dz: f32,
}

impl FieldDeltas {
    fn nabla(self) -> Vec3Std {
        Vec3Std::new(
            self.dfz_dy - self.dfy_dz,
            self.dfx_dz - self.dfz_dx,
            self.dfy_dx - self.dfx_dy,
        )
    }
}

pub struct StepEContext<'a, M> {
    inv_feature_size: f32,
    time_step: f32,
    stim_e: Vec3Std,
    mat: ArrayHandle<'a, M>,
    /// Input field sampled near this location
    in_h: VolumeSampleNeg,
    /// Handle to the output field at one specific index.
    out_e: ArrayHandleMut<'a, Vec3Std>,
}

impl<'a, M: Material> StepEContext<'a, M> {
    pub fn step_e(mut self) {
        // const EPS0_INV: f32 = 112940906737.1361;
        const TWICE_EPS0: f32 = 1.7708375625626e-11;
        let deltas = self.in_h.delta_h();
        // \nabla x H
        let nabla_h = deltas.nabla() * self.inv_feature_size;
        // $\nabla x H = \epsilon_0 dE/dt + \sigma E$
        // no-conductivity version:
        //   let delta_e = nabla_h * (self.time_step * EPS0_INV);
        let sigma = self.mat.get_ref().conductivity();
        let e_prev = self.out_e.get();
        let delta_e = (nabla_h - e_prev.elem_mul(sigma)).elem_div(
            sigma*self.time_step + Vec3Std::uniform(TWICE_EPS0)
        )*(2.0*self.time_step);
        // println!("spirv-step_e delta_e: {:?}", delta_e);
        self.out_e.write(e_prev + delta_e + self.stim_e);
    }
}


pub struct StepHContext<'a, M> {
    inv_feature_size: f32,
    time_step: f32,
    stim_h: Vec3Std,
    mat: ArrayHandle<'a, M>,
    /// Input field sampled near this location
    in_e: VolumeSamplePos,
    /// Handle to the output field at one specific index.
    out_h: ArrayHandleMut<'a, Vec3Std>,
    out_m: ArrayHandleMut<'a, Vec3Std>,
}

impl<'a, M: Material> StepHContext<'a, M> {
    pub fn step_h(mut self) {
        const MU0: f32 = 1.2566370621219e-06;
        const MU0_INV: f32 = 795774.715025073;
        let deltas = self.in_e.delta_e();
        // println!("spirv-step_h delta_e_struct: {:?}", deltas);
        // \nabla x E
        let nabla_e = deltas.nabla() * self.inv_feature_size;
        // println!("spirv-step_h nabla_e: {:?}", nabla_e);
        let delta_b = nabla_e * (-self.time_step);

        // Relation between these is: B = mu0*(H + M)
        let old_h = self.out_h.get();
        let old_m = self.out_m.get();
        let old_b = (old_h + old_m) * MU0;

        let new_b = old_b + delta_b + self.stim_h * MU0;
        let mat = self.mat.get_ref();
        let new_m = mat.move_b_vec(old_m, new_b);
        let new_h = new_b * MU0_INV - new_m;
        // println!("spirv-step_h delta_h: {:?}", delta_h);
        self.out_h.write(new_h);
        self.out_m.write(new_m);
    }
}