use crate::real::*;
use crate::cross::vec::Vec3;
use crate::geom::{Meters, Region};
use rand;

/// field densities
#[derive(Clone, Copy, Debug, Default, PartialEq)]
pub struct Fields {
    pub e: Vec3<f32>,
    pub h: Vec3<f32>,
}

/// field magnitude densities (really, signed magnitude)
#[derive(Clone, Copy, Debug, Default, PartialEq)]
pub struct FieldMags {
    pub e: f32,
    pub h: f32,
}

pub trait AbstractStimulus: Sync {
    // TODO: using floats for time and position is not great.
    /// Return the (E, H) field which should be added PER-SECOND to the provided position/time.
    fn at(&self, t_sec: f32, pos: Meters) -> Fields;
}

// impl<T: AbstractStimulus> AbstractStimulus for &T {
//     fn at(&self, t_sec: f32, pos: Meters) -> Vec3 {
//         (*self).at(t_sec, pos)
//     }
// }

impl<T: AbstractStimulus> AbstractStimulus for Vec<T> {
    fn at(&self, t_sec: f32, pos: Meters) -> Fields {
        let mut fields = Fields::default();
        for i in self {
            let Fields { e, h } = i.at(t_sec, pos);
            fields.e += e;
            fields.h += h;
        }
        fields
    }
}

impl AbstractStimulus for Box<dyn AbstractStimulus> {
    fn at(&self, t_sec: f32, pos: Meters) -> Fields {
        (**self).at(t_sec, pos)
    }
}

pub struct NoopStimulus;
impl AbstractStimulus for NoopStimulus {
    fn at(&self, _t_sec: f32, _pos: Meters) -> Fields {
        Fields::default()
    }
}

pub struct UniformStimulus {
    e: Vec3<f32>,
    h: Vec3<f32>,
}

impl UniformStimulus {
    pub fn new(e: Vec3<f32>, h: Vec3<f32>) -> Self {
        Self { e, h }
    }
    pub fn new_e(e: Vec3<f32>) -> Self {
        Self::new(e, Vec3::zero())
    }
}

impl AbstractStimulus for UniformStimulus {
    fn at(&self, t_sec: f32, _pos: Meters) -> Fields {
        TimeVarying3::at(self, t_sec)
    }
}
impl TimeVarying for UniformStimulus {}
impl TimeVarying3 for UniformStimulus {
    fn at(&self, _t_sec: f32) -> Fields {
        Fields { e: self.e, h: self.h }
    }
}

pub struct RngStimulus {
    seed: u64,
    e_scale: f32,
    h_scale: f32,
}

impl RngStimulus {
    pub fn new(seed: u64) -> Self {
        Self { seed, e_scale: 1e15, h_scale: 1e15 }
    }
    pub fn new_e(seed: u64) -> Self {
        Self { seed, e_scale: 1e15, h_scale: 0.0 }
    }
    fn gen(&self, t_sec: f32, pos: Meters, scale: f32, salt: u64) -> Vec3<f32> {
        use rand::{Rng as _, SeedableRng as _};
        let seed = self.seed
            ^ (t_sec.to_bits() as u64)
            ^ ((pos.x().to_bits() as u64) << 8)
            ^ ((pos.y().to_bits() as u64) << 16)
            ^ ((pos.z().to_bits() as u64) << 24)
            ^ (salt << 32);
        let mut rng = rand::rngs::StdRng::seed_from_u64(seed);
        Vec3::new(
            rng.gen_range(-scale..=scale),
            rng.gen_range(-scale..=scale),
            rng.gen_range(-scale..=scale),
        )
    }
}

impl AbstractStimulus for RngStimulus {
    fn at(&self, t_sec: f32, pos: Meters) -> Fields {
        Fields {
            e: self.gen(t_sec, pos, self.e_scale, 0),
            h: self.gen(t_sec, pos, self.h_scale, 0x7de3)
        }
    }
}

/// Apply a time-varying stimulus uniformly across some region
#[derive(Clone)]
pub struct Stimulus<R, T> {
    region: R,
    stim: T,
}

impl<R, T> Stimulus<R, T> {
    pub fn new(region: R, stim: T) -> Self {
        Self {
            region, stim
        }
    }
}

impl<R: Region + Sync, T: TimeVarying3 + Sync> AbstractStimulus for Stimulus<R, T> {
    fn at(&self, t_sec: f32, pos: Meters) -> Fields {
        if self.region.contains(pos) {
            self.stim.at(t_sec)
        } else {
            Fields::default()
        }
    }
}

/// Apply a time-varying stimulus across some region.
/// The stimulus seen at each point is based on its angle about the specified ray.
#[derive(Clone)]
pub struct CurlStimulus<R, T> {
    region: R,
    stim: T,
    center: Meters,
    axis: Meters,
}

impl<R, T> CurlStimulus<R, T> {
    pub fn new(region: R, stim: T, center: Meters, axis: Meters) -> Self {
        Self { region, stim, center, axis }
    }
}

impl<R: Region + Sync, T: TimeVarying1 + Sync> AbstractStimulus for CurlStimulus<R, T> {
    fn at(&self, t_sec: f32, pos: Meters) -> Fields {
        if self.region.contains(pos) {
            let FieldMags { e, h } = self.stim.at(t_sec);
            let from_center_to_point = *pos - *self.center;
            // TODO: is this inverted?
            let rotational = from_center_to_point.cross(*self.axis);
            let impulse_e = rotational.with_mag(e.cast()).unwrap_or_default();
            let impulse_h = rotational.with_mag(h.cast()).unwrap_or_default();
            Fields {
                e: impulse_e,
                h: impulse_h,
            }
        } else {
            Fields::default()
        }
    }
}

pub trait TimeVarying: Sized {
    fn shifted(self, new_start: f32) -> Shifted<Self> {
        Shifted::new(self, new_start)
    }
    fn gated(self, from: f32, to: f32) -> Gated<Self> {
        Gated::new(self, from, to)
    }
}

pub trait TimeVarying1: TimeVarying {
    /// Retrieve the (E, H) impulse to apply PER-SECOND at the provided time (in seconds).
    fn at(&self, t_sec: f32) -> FieldMags;
}

pub trait TimeVarying3: TimeVarying {
    /// Retrieve the (E, H) impulse to apply PER-SECOND at the provided time (in seconds).
    fn at(&self, t_sec: f32) -> Fields;
}

// assumed to represent the E field
impl TimeVarying for f32 {}
impl TimeVarying1 for f32 {
    fn at(&self, _t_sec: f32) -> FieldMags {
        FieldMags { e: *self, h: 0.0 }
    }
}

/// E field which changes magnitude sinusoidally as a function of t
#[derive(Clone)]
pub struct Sinusoid<A> {
    amp: A,
    omega: f32,
}

pub type Sinusoid1 = Sinusoid<f32>;
pub type Sinusoid3 = Sinusoid<Vec3<f32>>;

impl<A> Sinusoid<A> {
    pub fn new(amp: A, freq: f32) -> Self {
        Self {
            amp,
            omega: freq * f32::two_pi(),
        }
    }
    pub fn from_wavelength(amp: A, lambda: f32) -> Self {
        Self::new(amp, 1.0/lambda)
    }
    pub fn freq(&self) -> f32 {
        self.omega / f32::two_pi()
    }
    pub fn wavelength(&self) -> f32 {
        1.0 / self.freq()
    }
    pub fn one_cycle(self) -> Gated<Self> {
        let wl = self.wavelength();
        Gated::new(self, 0.0, wl)
    }
    pub fn half_cycle(self) -> Gated<Self> {
        let wl = self.wavelength();
        Gated::new(self, 0.0, 0.5 * wl)
    }
}

impl<A> TimeVarying for Sinusoid<A> {}

impl TimeVarying1 for Sinusoid1 {
    fn at(&self, t_sec: f32) -> FieldMags {
        FieldMags {
            e: self.amp * (t_sec * self.omega).sin(),
            h: 0.0,
        }
    }
}

impl TimeVarying3 for Sinusoid3 {
    fn at(&self, t_sec: f32) -> Fields {
        Fields {
            e: self.amp * (t_sec * self.omega).sin(),
            h: Vec3::zero(),
        }
    }
}

/// E field with magnitude that decays exponentially over t.
#[derive(Clone)]
pub struct Exp<A> {
    amp: A,
    tau: f32,
}
pub type Exp1 = Exp<f32>;
pub type Exp3 = Exp<Vec3<f32>>;

impl<A> Exp<A> {
    pub fn new(amp: A, half_life: f32) -> Self {
        let tau = std::f32::consts::LN_2/half_life;
        Self { amp, tau }
    }

    pub fn new_at(amp: A, start: f32, half_life: f32) -> Shifted<Gated<Self>> {
        Self::new(amp, half_life)
            .gated(0.0, half_life*100.0)
            .shifted(start)
    }
}

impl<A> TimeVarying for Exp<A> {}

impl TimeVarying1 for Exp1 {
    fn at(&self, t_sec: f32) -> FieldMags {
        FieldMags {
            e: self.amp * (t_sec * -self.tau).exp(),
            h: 0.0,
        }
    }
}

impl TimeVarying3 for Exp3 {
    fn at(&self, t_sec: f32) -> Fields {
        Fields {
            e: self.amp * (t_sec * -self.tau).exp(),
            h: Vec3::zero(),
        }
    }
}

#[derive(Clone)]
pub struct Gated<T> {
    inner: T,
    start: f32,
    end: f32,
}

impl<T> Gated<T> {
    pub fn new(inner: T, start: f32, end: f32) -> Self {
        Self { inner, start, end }
    }
}

impl<T> TimeVarying for Gated<T> {}
impl<T: TimeVarying1> TimeVarying1 for Gated<T> {
    fn at(&self, t_sec: f32) -> FieldMags {
        if (self.start..self.end).contains(&t_sec) {
            self.inner.at(t_sec)
        } else {
            Default::default()
        }
    }
}

impl<T: TimeVarying3> TimeVarying3 for Gated<T> {
    fn at(&self, t_sec: f32) -> Fields {
        if (self.start..self.end).contains(&t_sec) {
            self.inner.at(t_sec)
        } else {
            Default::default()
        }
    }
}

#[derive(Clone)]
pub struct Shifted<T> {
    inner: T,
    start_at: f32,
}

impl<T> Shifted<T> {
    pub fn new(inner: T, start_at: f32) -> Self {
        Self { inner, start_at }
    }
}

impl<T> TimeVarying for Shifted<T> {}
impl<T: TimeVarying1> TimeVarying1 for Shifted<T> {
    fn at(&self, t_sec: f32) -> FieldMags {
        self.inner.at(t_sec - self.start_at)
    }
}

impl<T: TimeVarying3> TimeVarying3 for Shifted<T> {
    fn at(&self, t_sec: f32) -> Fields {
        self.inner.at(t_sec - self.start_at)
    }
}

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

    macro_rules! assert_approx_eq {
        ($x:expr, $e:expr, $h:expr) => {
            let x = $x;
            let e = $e;
            let h = $h;
            let diff_e = (x.e - e).mag();
            assert!(diff_e <= 0.001, "{:?} != {:?}", x, e);
            let diff_h = (x.h - h).mag();
            assert!(diff_h <= 0.001, "{:?} != {:?}", x, h);
        }
    }
    
    #[test]
    fn sinusoid3() {
        let s = Sinusoid3::new(Vec3::new(10.0, 1.0, -100.0), 1000.0);
        assert_eq!(s.at(0.0), Fields::default());
        assert_approx_eq!(s.at(0.00025),
            Vec3::new(10.0, 1.0, -100.0), Vec3::zero()
        );
        assert_approx_eq!(s.at(0.00050), Vec3::zero(), Vec3::zero());
        assert_approx_eq!(s.at(0.00075), Vec3::new(-10.0, -1.0, 100.0), Vec3::zero());
    }
    
    #[test]
    fn sinusoid3_from_wavelength() {
        let s = Sinusoid3::from_wavelength(Vec3::new(10.0, 1.0, -100.0), 0.001);
        assert_eq!(s.at(0.0), Fields::default());
        assert_approx_eq!(s.at(0.00025), Vec3::new(10.0, 1.0, -100.0), Vec3::zero());
        assert_approx_eq!(s.at(0.00050), Vec3::zero(), Vec3::zero());
        assert_approx_eq!(s.at(0.00075), Vec3::new(-10.0, -1.0, 100.0), Vec3::zero());
    }
}