use crate::geom::{Meters, Region, Torus, Vec3, WorldRegion};
use crate::real::{Real as _, ToFloat as _};
use crate::sim::SampleableSim;
use dyn_clone::{self, DynClone};
use indexmap::IndexMap;
use serde::{Serialize, Deserialize};

// TODO: remove this Clone and Send requirement? Have Measurements be shared by-reference across
// threads? i.e. Sync, and no Clone
#[typetag::serde(tag = "type")]
pub trait AbstractMeasurement: Send + Sync + DynClone {
    fn eval(&self, state: &dyn SampleableSim) -> String;
    fn key_value(&self, state: &dyn SampleableSim) -> IndexMap<String, String>;
}
dyn_clone::clone_trait_object!(AbstractMeasurement);

pub fn eval_multiple_kv(state: &dyn SampleableSim, meas: &[Box<dyn AbstractMeasurement>]) -> IndexMap<String, String> {
    let mut r = IndexMap::new();
    for m in meas {
        let other = m.key_value(state);
        r.extend(other.into_iter());
    }
    r
}

#[derive(Clone, Serialize, Deserialize)]
pub struct Time;

#[typetag::serde]
impl AbstractMeasurement for Time {
    fn eval(&self, state: &dyn SampleableSim) -> String {
        format!("{:.3e}s (step {})", state.time(), state.step_no())
    }
    fn key_value(&self, state: &dyn SampleableSim) -> IndexMap<String, String> {
        [
            ("step".to_string(), state.step_no().to_string()),
            ("time".to_string(), state.time().to_string()),
        ].into_iter().collect()
    }
}

#[derive(Clone, Serialize, Deserialize)]
pub struct Meta;

#[typetag::serde]
impl AbstractMeasurement for Meta {
    fn eval(&self, state: &dyn SampleableSim) -> String {
        format!("{}x{}x{} feat: {:.1e}m", state.width(), state.height(), state.depth(), state.feature_size())
    }
    fn key_value(&self, state: &dyn SampleableSim) -> IndexMap<String, String> {
        [
            ("width".to_string(), state.width().to_string()),
            ("height".to_string(), state.height().to_string()),
            ("depth".to_string(), state.depth().to_string()),
            ("feature_size".to_string(), state.feature_size().to_string()),
        ].into_iter().collect()
    }
}

#[derive(Clone, Serialize, Deserialize)]
pub struct Label(pub String);

impl Label {
    pub fn new<S: Into<String>>(s: S) -> Self {
        Self(s.into())
    }
}

#[typetag::serde]
impl AbstractMeasurement for Label {
    fn eval(&self, _state: &dyn SampleableSim) -> String {
        self.0.clone()
    }
    fn key_value(&self, _state: &dyn SampleableSim) -> IndexMap<String, String> {
        [
            (self.0.clone(), self.0.clone()),
        ].into_iter().collect()
    }
}

#[derive(Clone, Serialize, Deserialize)]
pub struct Volume {
    name: String,
    region: Box<dyn Region>,
}

impl Volume {
    pub fn new<R: Region + 'static>(name: &str, r: R) -> Self {
        Self {
            name: name.into(),
            region: Box::new(r)
        }
    }
    /// Returns the volume of the region, in units of um^3
    fn data(&self, state: &dyn SampleableSim) -> f32 {
        let feat_um = state.feature_size() as f64 * 1e6;
        
        (state.volume_of_region(&*self.region) as f64 * feat_um * feat_um * feat_um) as f32
    }
}

#[typetag::serde]
impl AbstractMeasurement for Volume {
    fn eval(&self, state: &dyn SampleableSim) -> String {
        format!("Vol({}): {:.2e} um^3",
            self.name,
            self.data(state),
        )
    }
    fn key_value(&self, state: &dyn SampleableSim) -> IndexMap<String, String> {
        [
            (format!("Vol({})", self.name), self.data(state).to_string()),
        ].into_iter().collect()
    }
}

#[derive(Clone, Serialize, Deserialize)]
pub struct Current {
    name: String,
    region: Box<dyn Region>,
}

impl Current {
    pub fn new<R: Region + 'static>(name: &str, r: R) -> Self {
        Self {
            name: name.into(),
            region: Box::new(r)
        }
    }
    fn data(&self, state: &dyn SampleableSim) -> (f32, Vec3<f32>) {
        let FieldSample(volume, current_mag, current_vec) = state.map_sum_over_enumerated(&*self.region, |coord: Meters, _cell| {
            let current = state.current(coord);
            FieldSample(1, current.mag().cast(), current.cast())
        });
        let mean_current_mag = current_mag.to_f32() / (f32::from_primitive(volume));
        let mean_current_vec = current_vec.cast::<f32>() / (f32::from_primitive(volume));
        (mean_current_mag, mean_current_vec.cast())
    }
}

#[derive(Default)]
struct FieldSample(u32, f64, Vec3<f64>);

impl std::iter::Sum for FieldSample {
    fn sum<I>(iter: I) -> Self
        where I: Iterator<Item = Self>
    {
        let mut s = FieldSample::default();
        for FieldSample(a, b, c) in iter {
            s = FieldSample(s.0 + a, s.1 + b, s.2 + c);
        }
        s
    }
}

#[derive(Default)]
struct FieldSamples<T>(T);

impl std::iter::Sum for FieldSamples<[FieldSample; 2]> {
    fn sum<I>(iter: I) -> Self
        where I: Iterator<Item = Self>
    {
        let mut s = Self::default();
        for p in iter {
            s.0[0] = FieldSample(s.0[0].0 + p.0[0].0, s.0[0].1 + p.0[0].1, s.0[0].2 + p.0[0].2);
            s.0[1] = FieldSample(s.0[1].0 + p.0[1].0, s.0[1].1 + p.0[1].1, s.0[1].2 + p.0[1].2);
        }
        s
    }
}

impl std::iter::Sum for FieldSamples<[FieldSample; 3]> {
    fn sum<I>(iter: I) -> Self
        where I: Iterator<Item = Self>
    {
        let mut s = Self::default();
        for p in iter {
            s.0[0] = FieldSample(s.0[0].0 + p.0[0].0, s.0[0].1 + p.0[0].1, s.0[0].2 + p.0[0].2);
            s.0[1] = FieldSample(s.0[1].0 + p.0[1].0, s.0[1].1 + p.0[1].1, s.0[1].2 + p.0[1].2);
            s.0[2] = FieldSample(s.0[2].0 + p.0[2].0, s.0[2].1 + p.0[2].1, s.0[2].2 + p.0[2].2);
        }
        s
    }
}

#[typetag::serde]
impl AbstractMeasurement for Current {
    fn eval(&self, state: &dyn SampleableSim) -> String {
        let (mean_current_mag, mean_current_vec) = self.data(state);
        format!("I/cell({}): {:.2e} {:.2e}",
            self.name,
            mean_current_mag,
            mean_current_vec)
    }
    fn key_value(&self, state: &dyn SampleableSim) -> IndexMap<String, String> {
        let (mean_current_mag, mean_current_vec) = self.data(state);
        [
            (format!("Imag/cell({})", self.name), mean_current_mag.to_string()),
            (format!("I/cell({})", self.name), mean_current_vec.to_string()),
        ].into_iter().collect()
    }
}

/// Measures the current directed around a closed loop
#[derive(Clone, Serialize, Deserialize)]
pub struct CurrentLoop {
    name: String,
    region: Torus
}

impl CurrentLoop {
    pub fn new(name: &str, r: Torus) -> Self {
        Self {
            name: name.into(),
            region: r,
        }
    }
    fn data(&self, state: &dyn SampleableSim) -> f32 {
        let FieldSample(volume, directed_current, _current_vec) = state.map_sum_over_enumerated(&self.region, |coord: Meters, _cell| {
            let normal = self.region.axis();
            let to_coord = *coord - *self.region.center();
            let tangent = normal.cross(to_coord).norm();
            let current = state.current(coord);
            let directed_current = current.dot(tangent.cast());
            FieldSample(1, directed_current.cast(), current.cast())
        });
        let mean_directed_current = directed_current.cast::<f32>() / f32::from_primitive(volume);
        let cross_section = self.region.cross_section() / (state.feature_size() * state.feature_size());
        let cross_sectional_current = mean_directed_current * cross_section;
        cross_sectional_current
    }
}

#[typetag::serde]
impl AbstractMeasurement for CurrentLoop {
    fn eval(&self, state: &dyn SampleableSim) -> String {
        let cross_sectional_current = self.data(state);
        format!("I({}): {:.2e}", self.name, cross_sectional_current)
    }
    fn key_value(&self, state: &dyn SampleableSim) -> IndexMap<String, String> {
        let cross_sectional_current = self.data(state);
        [
            (format!("I({})", self.name), cross_sectional_current.to_string()),
        ].into_iter().collect()
    }
}


/// Measures the M, B field directed around a closed loop
#[derive(Clone, Serialize, Deserialize)]
pub struct MagneticLoop {
    name: String,
    region: Torus
}

impl MagneticLoop {
    pub fn new(name: &str, r: Torus) -> Self {
        Self {
            name: name.into(),
            region: r,
        }
    }
    fn data(&self, state: &dyn SampleableSim) -> (f32, f32, f32) {
        let FieldSamples([
            FieldSample(volume, directed_m, _m_vec),
            FieldSample(_, directed_b, _b_vec),
            FieldSample(_, directed_h, _h_vec),
        ]) = state.map_sum_over_enumerated(&self.region, |coord: Meters, cell|
        {
            let normal = self.region.axis();
            let to_coord = *coord - *self.region.center();
            let tangent = normal.cross(to_coord).norm();

            let m = cell.m();
            let directed_m = m.dot(tangent.cast());
            let b = cell.b();
            let directed_b = b.dot(tangent.cast());
            let h = cell.h();
            let directed_h = h.dot(tangent.cast());
            FieldSamples([
                FieldSample(1, directed_m.cast(), m.cast()),
                FieldSample(1, directed_b.cast(), b.cast()),
                FieldSample(1, directed_h.cast(), h.cast()),
            ])
        });
        // let cross_section = self.region.cross_section() / (state.feature_size() * state.feature_size());
        let mean_directed_m = directed_m.cast::<f32>() / f32::from_primitive(volume);
        // let cross_sectional_m = mean_directed_m * cross_section;
        let mean_directed_b = directed_b.cast::<f32>() / f32::from_primitive(volume);
        // let cross_sectional_b = mean_directed_b * cross_section;
        let mean_directed_h = directed_h.cast::<f32>() / f32::from_primitive(volume);
        // let cross_sectional_h = mean_directed_h * cross_section;
        // format!(
        //     "M({}): {:.2e}; B({}): {:.2e}; H({}): {:.2e}",
        //     self.name, cross_sectional_m,
        //     self.name, cross_sectional_b,
        //     self.name, cross_sectional_h,
        // )
        (mean_directed_m, mean_directed_b, mean_directed_h)
    }
}

#[typetag::serde]
impl AbstractMeasurement for MagneticLoop {
    fn eval(&self, state: &dyn SampleableSim) -> String {
        let (mean_directed_m, mean_directed_b, mean_directed_h) = self.data(state);
        format!(
            "M({}): {:.2e}; B({}): {:.2e}; H({}): {:.2e}",
            self.name, mean_directed_m,
            self.name, mean_directed_b,
            self.name, mean_directed_h,
        )
    }
    fn key_value(&self, state: &dyn SampleableSim) -> IndexMap<String, String> {
        let (mean_directed_m, mean_directed_b, mean_directed_h) = self.data(state);
        [
            (format!("M({})", self.name), mean_directed_m.to_string()),
            (format!("B({})", self.name), mean_directed_b.to_string()),
            (format!("H({})", self.name), mean_directed_h.to_string()),
        ].into_iter().collect()
    }
}

/// mean M over a region
#[derive(Clone, Serialize, Deserialize)]
pub struct MagneticFlux {
    name: String,
    region: Box<dyn Region>,
}

impl MagneticFlux {
    pub fn new<R: Region + 'static>(name: &str, r: R) -> Self {
        Self {
            name: name.into(),
            region: Box::new(r)
        }
    }
    fn data(&self, state: &dyn SampleableSim) -> Vec3<f32> {
        let FieldSample(volume, _directed_mag, mag_vec) = state.map_sum_over(&*self.region, |cell| {
            let b = cell.b();
            let mag = b.mag();
            FieldSample(1, mag.cast(), b.cast())
        });
        let mean_mag = mag_vec.cast() / f32::from_primitive(volume);
        mean_mag
    }
}

#[typetag::serde]
impl AbstractMeasurement for MagneticFlux {
    fn eval(&self, state: &dyn SampleableSim) -> String {
        let mean_mag = self.data(state);
        format!("Bavg({}): {:.2e}", self.name, mean_mag)
    }
    fn key_value(&self, state: &dyn SampleableSim) -> IndexMap<String, String> {
        let mean_mag = self.data(state);
        [
            (format!("Bavg({})", self.name), mean_mag.to_string()),
        ].into_iter().collect()
    }
}

/// mean B over a region
#[derive(Clone, Serialize, Deserialize)]
pub struct Magnetization {
    name: String,
    region: Box<dyn Region>,
}

impl Magnetization {
    pub fn new<R: Region + 'static>(name: &str, r: R) -> Self {
        Self {
            name: name.into(),
            region: Box::new(r)
        }
    }
    fn data(&self, state: &dyn SampleableSim) -> Vec3<f32> {
        let FieldSample(volume, _directed_mag, mag_vec) = state.map_sum_over(&*self.region, |cell| {
            let m = cell.m();
            let mag = m.mag();
            FieldSample(1, mag.cast(), m.cast())
        });
        let mean_mag = mag_vec.cast() / f32::from_primitive(volume);
        mean_mag
    }
}

#[typetag::serde]
impl AbstractMeasurement for Magnetization {
    fn eval(&self, state: &dyn SampleableSim) -> String {
        let mean_mag = self.data(state);
        format!("Mavg({}): {:.2e}", self.name, mean_mag)
    }
    fn key_value(&self, state: &dyn SampleableSim) -> IndexMap<String, String> {
        let mean_mag = self.data(state);
        [
            (format!("Mavg({})", self.name), mean_mag.to_string()),
        ].into_iter().collect()
    }
}

fn loc(v: Meters) -> String {
    format!("{:.0} um", *v * f32::from_primitive(1_000_000))
}

/// M
#[derive(Clone, Serialize, Deserialize)]
pub struct MagnetizationAt(pub Meters);

#[typetag::serde]
impl AbstractMeasurement for MagnetizationAt {
    fn eval(&self, state: &dyn SampleableSim) -> String {
        let m = state.sample(self.0).m();
        format!("M{}: {:.2e}", loc(self.0), m)
    }
    fn key_value(&self, state: &dyn SampleableSim) -> IndexMap<String, String> {
        let m = state.sample(self.0).m();
        [
            (format!("M{}", loc(self.0)), m.to_string()),
        ].into_iter().collect()
    }
}

/// B
#[derive(Clone, Serialize, Deserialize)]
pub struct MagneticFluxAt(pub Meters);

#[typetag::serde]
impl AbstractMeasurement for MagneticFluxAt {
    fn eval(&self, state: &dyn SampleableSim) -> String {
        let b = state.sample(self.0).b();
        format!("B{}: {:.2e}", loc(self.0), b)
    }
    fn key_value(&self, state: &dyn SampleableSim) -> IndexMap<String, String> {
        let b = state.sample(self.0).b();
        [
            (format!("B{}", loc(self.0)), b.to_string()),
        ].into_iter().collect()
    }
}

/// H
#[derive(Clone, Serialize, Deserialize)]
pub struct MagneticStrengthAt(pub Meters);

#[typetag::serde]
impl AbstractMeasurement for MagneticStrengthAt {
    fn eval(&self, state: &dyn SampleableSim) -> String {
        let h = state.sample(self.0).h();
        format!("H{}: {:.2e}", loc(self.0), h)
    }
    fn key_value(&self, state: &dyn SampleableSim) -> IndexMap<String, String> {
        let h = state.sample(self.0).h();
        [
            (format!("H{}", loc(self.0)), h.to_string()),
        ].into_iter().collect()
    }
}

#[derive(Clone, Serialize, Deserialize)]
pub struct ElectricField(pub Meters);

#[typetag::serde]
impl AbstractMeasurement for ElectricField {
    fn eval(&self, state: &dyn SampleableSim) -> String {
        let e = state.sample(self.0).e();
        format!("E{}: {:.2e}", loc(self.0), e)
    }
    fn key_value(&self, state: &dyn SampleableSim) -> IndexMap<String, String> {
        let e = state.sample(self.0).e();
        [
            (format!("E{}", loc(self.0)), e.to_string()),
        ].into_iter().collect()
    }
}

#[derive(Clone, Serialize, Deserialize)]
pub struct Energy {
    name: String,
    region: Box<dyn Region>,
}

impl Energy {
    pub fn world() -> Self {
        Self::new("World", WorldRegion)
    }
    pub fn new<R: Region + 'static>(name: &str, region: R) -> Self {
        Self {
            name: name.into(),
            region: Box::new(region),
        }
    }
    fn data(&self, state: &dyn SampleableSim) -> f32 {
        // Potential energy stored in a E/M field:
        //   https://en.wikipedia.org/wiki/Magnetic_energy
        //   https://en.wikipedia.org/wiki/Electric_potential_energy#Energy_stored_in_an_electrostatic_field_distribution
        //   TODO: consider the M field? https://en.wikipedia.org/wiki/Potential_energy#Magnetic_potential_energy
        //   U(B) = 1/2 \int H . B dV
        //   U(E) = 1/2 \int E . D dV
        #[allow(non_snake_case)]
        let dV = state.feature_volume();
        let e = f64::from_primitive(0.5 * dV) * state.map_sum_over(&*self.region, |cell| {
            // E . D = E . (E + P) = E.E since we don't model polarization fields
            cell.h().dot(cell.b()).to_f64() + cell.e().mag_sq().to_f64()
        });
        e.cast()
    }
}

#[typetag::serde]
impl AbstractMeasurement for Energy {
    fn eval(&self, state: &dyn SampleableSim) -> String {
        let e = self.data(state);
        format!("U({}): {:.2e}", self.name, e)
    }
    fn key_value(&self, state: &dyn SampleableSim) -> IndexMap<String, String> {
        let e = self.data(state);
        [
            (format!("U({})", self.name), e.to_string()),
        ].into_iter().collect()
    }
}

#[derive(Clone, Serialize, Deserialize)]
pub struct Power {
    name: String,
    region: Box<dyn Region>
}

impl Power {
    pub fn world() -> Self {
        Self::new("World", WorldRegion)
    }
    pub fn new<R: Region + 'static>(name: &str, region: R) -> Self {
        Self {
            name: name.into(),
            region: Box::new(region),
        }
    }
    fn data(&self, state: &dyn SampleableSim) -> f32 {
        // Power is P = IV = A*J*V = L^2*J.(LE) = L^3 J.E
        // where L is feature size.
        #[allow(non_snake_case)]
        let dV = state.feature_volume();
        let power = f64::from_primitive(dV) * state.map_sum_over(&*self.region, |cell| {
            cell.current_density().dot(cell.e()).to_f64()
        });
        power.cast()
    }
}

#[typetag::serde]
impl AbstractMeasurement for Power {
    fn eval(&self, state: &dyn SampleableSim) -> String {
        let power = self.data(state);
        format!("P({}): {:.2e}", self.name, power)
    }
    fn key_value(&self, state: &dyn SampleableSim) -> IndexMap<String, String> {
        let power = self.data(state);
        [
            (format!("P({})", self.name), power.to_string()),
        ].into_iter().collect()
    }
}