//! this example positions buffers adjacently and uses an ASYMMETRIC coil winding.
//! v.s. the fourth prototype, it changes the couplings in an attempt to reduce unwanted
//! clock -> mem2 coupling

use coremem::{Driver, mat, meas, SpirvDriver};
use coremem::geom::{region, Cube, Dilate, Memoize, Meters, Region, Spiral, SwapYZ, Torus, Translate, Wrap};
use coremem::render::CsvRenderer;
use coremem::stim::{CurlStimulus, Gated, Sinusoid1, TimeVarying1 as _};
use coremem::sim::units::{Seconds, Frame, Time as _};
use log::{info, warn};

#[allow(unused)]
use coremem::geom::{Coord as _, Region as _};

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
enum PulseType {
    Square,
    Sine
}

/// Return just the extrema of some collection
fn extrema(mut meas: Vec<f32>) -> Vec<f32> {
    let mut i = 0;
    while i + 2 < meas.len() {
        let (prev, cur, next) = (meas[i], meas[i+1], meas[i+2]);
        if (prev <= cur && cur <= next) || (prev >= cur && cur >= next) {
            meas.remove(i+1);
        } else {
            i += 1;
        }
    }
    meas
}

/// Return the (signed) peak magnitude and stable value
fn significa(meas: Vec<f32>) -> (f32, f32) {
    let peak = meas.iter().max_by(|a, b| a.abs().partial_cmp(&b.abs()).unwrap()).copied().unwrap_or_default();
    let stable = meas.last().copied().unwrap_or_default();
    (peak, stable)
}

#[derive(Copy, Clone, Debug)]
struct Params {
    dry_run: bool,
    feat_size: f32,
    buffer_xy: f32,
    buffer_z: f32,
    boundary_xy: f32,
    boundary_z: f32,

    ferro_major: f32,
    ferro_minor: f32,
    ferro_buffer: f32, // horizontal space between ferros,
    wire_minor: f32,
    wire_wrap_minor: f32, // 2x wire_wrap_minor + feat_size must be < ferro_buffer
    wire_set_major: f32,
    wire_wrap_dilation: f32,
    wire_wrap_iters: usize,

    wraps1: f32,
    wrap1_coverage: f32,
    wraps2: f32,
    wrap2_coverage: f32,

    wire_conductivity: f32,

    peak_set_current: f32,
    peak_clock_current: f32,
    set_duration: f32,
    clock_duration: f32,
    clock_type: PulseType,
    pre_time: f32, // how long between set and clock
    post_time: f32, // how long to wait after the clock

    dump_frames: Option<u64>,
}

struct Geometries {
    dim: Meters,
    ferro1_region: Torus,
    ferro2_region: Torus,
    set1_region: Torus,
    set2_region: Torus,
    coupling_region: region::Union,
    coupling_wire_top: Cube,
    coupling_wire_bot: Cube,
    wrap1_len: f32,
    wrap2_len: f32,
}

#[derive(Clone, Debug, Default)]
struct Results {
    m1_peak: f32,
    m2_peak: f32,
    m1_stable: f32,
    m2_stable: f32,
    h1_peak: f32,
    h2_peak: f32,
    h1_stable: f32,
    h2_stable: f32,
    iset_peak: f32,
    icoupling_peak: f32,
    peak_m_ratio: f32,
    stable_m_ratio: f32,
    t: f32,
}

fn derive_geometries(p: Params) -> Option<Geometries> {
    use std::f32::consts::PI;
    let feat_sizes = Meters::new(p.feat_size, p.feat_size, p.feat_size);
    let width = 4.0*p.ferro_major + 2.0*(p.buffer_xy + p.boundary_xy + p.wire_set_major + p.wire_minor) + p.ferro_buffer;
    let height = 2.0*(p.ferro_major + p.ferro_minor + 4.0*p.wire_wrap_minor + 12.0*p.feat_size + p.buffer_xy + p.boundary_xy);
    let depth = 2.0*(p.wire_set_major.max(4.0*p.wire_wrap_minor + p.ferro_minor + p.feat_size) + p.wire_minor + p.buffer_z + p.boundary_z);
    let dim = Meters::new(width, height, depth);

    let ferro1_center = Meters::new(
        p.buffer_xy + p.boundary_xy + p.wire_set_major + p.wire_minor + p.ferro_major,
        p.buffer_xy + p.boundary_xy + p.ferro_major + p.ferro_minor + 4.0*p.wire_wrap_minor + 12.0*p.feat_size,
        0.5*depth,
        // buffer_z + boundary_z + wire_set_major + wire_minor
    );
    let ferro2_center = ferro1_center + Meters::new(2.0*p.ferro_major + p.ferro_buffer, 0.0, 0.0);
    let ferro_center = (ferro1_center + ferro2_center)*0.5;

    // reserve the left/right locations for the SET wires.
    let set1_center = ferro1_center - Meters::new_x(p.ferro_major);
    let set2_center = ferro2_center + Meters::new_x(p.ferro_major);

    let ferro1_region = Torus::new_xy(ferro1_center, p.ferro_major, p.ferro_minor);
    let ferro2_region = Torus::new_xy(ferro2_center, p.ferro_major, p.ferro_minor);

    let set1_region = Torus::new_xz(set1_center, p.wire_set_major, p.wire_minor);
    let set2_region = Torus::new_xz(set2_center, p.wire_set_major, p.wire_minor);

    let wrap1_rate = 2.0*p.wrap1_coverage/p.wraps1;
    let coupling_region1 = Memoize::new(Dilate::new(
        Wrap::new_about(
            Translate::new(
                SwapYZ::new(region::and(
                    Spiral::new(p.ferro_minor + 2.0*p.wire_wrap_minor + p.feat_size, p.wire_wrap_minor, wrap1_rate),
                    Cube::new(Meters::new(-1.0, -1.0, -p.wrap1_coverage), Meters::new(1.0, 1.0, p.wrap1_coverage))
                )),
                ferro1_center + Meters::new(1.0*p.ferro_major, 0.0, 0.0),
            ),
            1.0, // one half-rev => y=1.0
            ferro1_center,
        ),
        p.wire_wrap_dilation,
        p.wire_wrap_dilation / (p.wire_wrap_iters as f32),
    ));

    let wrap2_rate = 2.0*p.wrap2_coverage/p.wraps2;
    let coupling_region2 = Memoize::new(Dilate::new(
        Wrap::new_about(
            Translate::new(
                SwapYZ::new(region::and_not(
                    Spiral::new(p.ferro_minor + 2.0*p.wire_wrap_minor + p.feat_size, p.wire_wrap_minor, wrap2_rate),
                    Cube::new(Meters::new(-1.0, -1.0, -1.0 + p.wrap2_coverage), Meters::new(1.0, 1.0, 1.0 - p.wrap2_coverage))
                )),
                ferro2_center + Meters::new_x(p.ferro_major),
            ),
            1.0, // one half-rev => y=1.0
            ferro2_center,
        ),
        p.wire_wrap_dilation,
        p.wire_wrap_dilation / (p.wire_wrap_iters as f32),
    ));

    let coupling_wire_top = Cube::new_centered(
        ferro_center - Meters::new_y(p.ferro_major + 4.0*p.wire_wrap_minor + 12.0*p.feat_size),
        Meters::new(p.ferro_buffer + 4.0*p.ferro_major + 4.0*p.feat_size, 2.0*p.feat_size, 2.0*p.feat_size)
    );
    let coupling_wire_bot = Cube::new_centered(
        ferro_center + Meters::new_y(p.ferro_major + 4.0*p.wire_wrap_minor + 12.0*p.feat_size),
        Meters::new(p.ferro_buffer + 4.0*p.ferro_major + 4.0*p.feat_size, 2.0*p.feat_size, 2.0*p.feat_size)
    );
    let wrap1_top = ferro1_center + Meters::new_x(p.ferro_major).rotate_z(-p.wrap1_coverage*PI);
    let wrap1_bot = ferro1_center + Meters::new_x(p.ferro_major).rotate_z(p.wrap1_coverage*PI);
    let wrap2_top = ferro2_center + Meters::new_x(p.ferro_major).rotate_z((1.0+p.wrap2_coverage)*PI);
    let wrap2_bot = ferro2_center + Meters::new_x(p.ferro_major).rotate_z((1.0-p.wrap2_coverage)*PI);
    let coupling_stub_top_left = Cube::new_including_negatives(
        wrap1_top + feat_sizes*2.0,
        wrap1_top.with_y(coupling_wire_top.bot()) - feat_sizes*2.0,
    );
    let coupling_stub_bot_left = Cube::new_including_negatives(
        wrap1_bot - feat_sizes*2.0,
        wrap1_bot.with_y(coupling_wire_bot.top()) + feat_sizes*2.0,
    );
    let coupling_stub_top_right = Cube::new_including_negatives(
        wrap2_top + feat_sizes*2.0,
        wrap2_top.with_y(coupling_wire_top.bot()) - feat_sizes*2.0,
    );
    let coupling_stub_bot_right = Cube::new_including_negatives(
        wrap2_bot - feat_sizes*2.0,
        wrap2_bot.with_y(coupling_wire_bot.top()) + feat_sizes*2.0,
    );
    let coupling_stubs = region::Union::new()
        .with(coupling_stub_top_left.clone())
        .with(coupling_stub_top_right.clone())
        .with(coupling_stub_bot_left.clone())
        .with(coupling_stub_bot_right.clone())
    ;

    let coupling_wires = region::Union::new()
        .with(coupling_wire_top.clone())
        .with(coupling_wire_bot.clone())
        .with(coupling_stubs.clone())
    ;

    let coupling_region = region::Union::new()
        .with(coupling_region1.clone())
        .with(coupling_region2.clone())
        .with(coupling_wires.clone())
    ;

    let wrap1_with_coupling = region::union(
        coupling_region1.clone(), coupling_wires.clone()
    );
    let wrap2_with_coupling = region::union(
        coupling_region2.clone(), coupling_wires.clone()
    );

    // show that the coupling top/bot wires are connected through the wrapping
    if !region::is_connected(
        &wrap1_with_coupling, 
        coupling_wire_top.center(),
        coupling_wire_bot.center(),
        p.feat_size,
    ) {
        warn!("wrap1 not connected for params: {:?}", p);
        return None;
    }
    if !region::is_connected(
        &wrap2_with_coupling, 
        coupling_wire_top.center(),
        coupling_wire_bot.center(),
        p.feat_size,
    ) {
        warn!("wrap2 not connected for params: {:?}", p);
        return None;
    }

    let wrap1_len = region::distance_to(
        &wrap1_with_coupling,
        coupling_stub_top_left.center().to_index(p.feat_size),
        coupling_stub_bot_left.center().to_index(p.feat_size),
        p.feat_size,
    ).unwrap();
    let wrap2_len = region::distance_to(
        &wrap2_with_coupling,
        coupling_stub_top_right.center().to_index(p.feat_size),
        coupling_stub_bot_right.center().to_index(p.feat_size),
        p.feat_size,
    ).unwrap();
    info!("wrap lengths: {}, {}", wrap1_len, wrap2_len);

    Some(Geometries {
        dim,
        ferro1_region,
        ferro2_region,
        set1_region,
        set2_region,
        coupling_region,
        coupling_wire_top,
        coupling_wire_bot,
        wrap1_len,
        wrap2_len,
    })
}

fn run_sim(id: u32, p: Params, g: Geometries) -> Results {
    info!("run_sim {}: {:?}", id, p);
    let m_to_um = |m: f32| (m * 1e6).round() as u32;
    let feat_vol = p.feat_size * p.feat_size * p.feat_size;

    // mu_r=881.33, starting at H=25 to H=75.
    let ferro_mat = mat::MHPgram::new(25.0, 881.33, 44000.0);
    // let ferro_mat = mat::db::conductor(wire_conductivity);
    let wire_mat = mat::db::conductor(p.wire_conductivity);

    let mut driver: Driver<_> = Driver::new_spirv(g.dim, p.feat_size);
    driver.set_steps_per_stim(1000);
    driver.fill_region(&g.ferro1_region, ferro_mat);
    driver.fill_region(&g.ferro2_region, ferro_mat);
    driver.fill_region(&g.set1_region, wire_mat);
    driver.fill_region(&g.set2_region, wire_mat);
    driver.fill_region(&g.coupling_region, wire_mat);

    info!("boundary: {}um; {}um", m_to_um(p.boundary_xy), m_to_um(p.boundary_z));
    info!("size: {:?}", g.dim);
    info!("ferro1: {:?}", g.ferro1_region.center());
    info!("ferro2: {:?}", g.ferro2_region.center());
    driver.add_classical_boundary(Meters::new(p.boundary_xy, p.boundary_xy, p.boundary_z));

    // assert!(driver.test_region_filled(&g.ferro1_region, ferro_mat));
    // assert!(driver.test_region_filled(&g.ferro2_region, ferro_mat));
    // assert!(driver.test_region_filled(&g.set1_region, wire_mat));
    // assert!(driver.test_region_filled(&g.set2_region, wire_mat));
    // assert!(driver.test_region_filled(&g.coupling_region, wire_mat));

    let add_drive_sine_pulse = |driver: &mut SpirvDriver, region: &Torus, start: f32, duration: f32, amp: f32| {
        let wave = Sinusoid1::from_wavelength(amp, duration * 2.0)
            .half_cycle()
            .shifted(start);
        driver.add_stimulus(CurlStimulus::new(
            region.clone(),
            wave.clone(),
            region.center(),
            region.axis()
        ));
    };

    let add_drive_square_pulse = |driver: &mut SpirvDriver, region: &Torus, start: f32, duration: f32, amp: f32| {
        let wave = Gated::new(amp, start, start+duration);
        driver.add_stimulus(CurlStimulus::new(
            region.clone(),
            wave.clone(),
            region.center(),
            region.axis()
        ));
    };

    let add_drive_step = |driver: &mut SpirvDriver, region: &Torus, start: f32, amp: f32| {
        add_drive_square_pulse(driver, region, start, 1.0, amp);
    };

    let add_drive_pulse = |ty: PulseType, driver: &mut SpirvDriver, region: &Torus, start: f32, duration: f32, amp: f32| {
        match ty {
            PulseType::Sine => add_drive_sine_pulse(driver, region, start, duration, amp),
            PulseType::Square => add_drive_square_pulse(driver, region, start, duration, amp),
        }
    };

    // J=\sigma E
    // dJ/dt = \sigma dE/dT
    // dE/dt = dJ/dt / \sigma
    // dE/dt = dI/dt / (A*\sigma)
    // if I = k*sin(w t) then dE/dt = k*w cos(w t) / (A*\sigma)
    // i.e. dE/dt is proportional to I/(A*\sigma), multiplied by w (or, divided by wavelength)
    let peak_set = p.peak_set_current / feat_vol / (g.set1_region.cross_section() * p.wire_conductivity);
    let peak_clock = p.peak_clock_current / feat_vol / (g.set1_region.cross_section() * p.wire_conductivity);

    // SET cores
    add_drive_sine_pulse(&mut driver, &g.set1_region, 0.01*p.set_duration, p.set_duration, -peak_set);
    add_drive_sine_pulse(&mut driver, &g.set2_region, 0.01*p.set_duration, p.set_duration, peak_set);

    // CLEAR core1
    add_drive_pulse(p.clock_type, &mut driver, &g.set1_region, p.set_duration + p.pre_time, p.clock_duration, peak_clock);
    // add_drive_step(&mut driver, &set1_region, set_duration + pre_time, peak_clock);

    let duration = Seconds(p.set_duration + p.pre_time + p.clock_duration + p.post_time)
        .to_frame(driver.timestep())
        .round_up(32000);

    driver.add_measurement(meas::Volume::new("mem1", g.ferro1_region.clone()));
    driver.add_measurement(meas::MagneticLoop::new("mem1", g.ferro1_region.clone()));

    driver.add_measurement(meas::Volume::new("mem2", g.ferro2_region.clone()));
    driver.add_measurement(meas::MagneticLoop::new("mem2", g.ferro2_region.clone()));

    driver.add_measurement(meas::CurrentLoop::new("set1", g.set1_region.clone()));
    driver.add_measurement(meas::Power::new("set1", g.set1_region.clone()));
    driver.add_measurement(meas::CurrentLoop::new("set2", g.set2_region.clone()));

    // driver.add_measurement(meas::CurrentLoop::new("coupling1", coupling_region1.clone()));
    // driver.add_measurement(meas::CurrentLoop::new("coupling2", coupling_region2.clone()));
    driver.add_measurement(meas::Current::new("couplingtop", g.coupling_wire_top.clone()));
    driver.add_measurement(meas::Current::new("couplingbot", g.coupling_wire_bot.clone()));

    let base = format!("buffer5-{}", id);
    let prefix = format!("out/{}/{}-{}-{}setmA-{}setps-{}clkmA-{}clkps-{}um-{}ferromaj-{}:{}wraps-{}:{}cov-{:?}clk",
        base,
        base,
        *driver.size(),
        (p.peak_set_current * 1e3).round() as i64,
        (p.set_duration * 1e12).round() as i64,
        (p.peak_clock_current * 1e3).round() as i64,
        (p.clock_duration * 1e12).round() as i64,
        (p.feat_size * 1e6).round() as i64,
        p.ferro_major,
        p.wraps1,
        p.wraps2,
        p.wrap1_coverage,
        p.wrap2_coverage,
        p.clock_type,
    );

    if p.dry_run {
        info!("bailing (dry run): {}", prefix);
        return Results::default();
    }

    let _ = std::fs::create_dir_all(&prefix);

    driver.add_state_file(&*format!("{}/state.bc", prefix), 16000);
    driver.add_serializer_renderer(&*format!("{}/frame-", prefix), 32000, p.dump_frames);
    let meas_csv = format!("{}/meas.csv", prefix);
    let meas_sparse_csv = format!("{}/meas-sparse.csv", prefix);
    driver.add_csv_renderer(&*meas_csv, 200, None);
    driver.add_csv_renderer(&*meas_sparse_csv, 8000, None);

    driver.step_until(duration);
    
    let (m1_peak, m1_stable) = significa(CsvRenderer::new(&*meas_sparse_csv).read_column_as_f32("M(mem1)"));
    let (m2_peak, m2_stable) = significa(CsvRenderer::new(&*meas_sparse_csv).read_column_as_f32("M(mem2)"));
    let (h1_peak, h1_stable) = significa(CsvRenderer::new(&*meas_sparse_csv).read_column_as_f32("H(mem1)"));
    let (h2_peak, h2_stable) = significa(CsvRenderer::new(&*meas_sparse_csv).read_column_as_f32("H(mem2)"));
    let (iset_peak, _iset_stable) = significa(CsvRenderer::new(&*meas_sparse_csv).read_column_as_f32("I(set1)"));
    let (icoupling_peak, _icoupling_stable) = significa(CsvRenderer::new(&*meas_sparse_csv).read_column_as_f32("Imag/cell(couplingtop)"));
    let res = Results {
        m1_peak,
        m2_peak,
        m1_stable,
        m2_stable,
        h1_peak,
        h2_peak,
        h1_stable,
        h2_stable,
        iset_peak,
        icoupling_peak,
        peak_m_ratio: m2_peak / m1_peak,
        stable_m_ratio: m2_stable / m1_stable,
        t: driver.time(),
    };
    std::fs::write(
        format!("{}/results.txt", prefix),
        format!("{:#?}\n", res),
    ).unwrap();
    info!("completed sim: {}", prefix);
    res
}


fn main() {
    coremem::init_logging();
    let i = 62;
    let dry_run = false;
    let mut variants = Vec::new();
    for ferro_major in [1360e-6, 2320e-6, 1680e-6] {
        // 10e-9 is enough to see m2 peak, 275e-9 for it to stabilize
        for post_time in [10e-9, 275e-9] {
            for peak_clock_current in [400.0, 100.0, 1600.0, 25.0] {
                for clock_type in [PulseType::Square, PulseType::Sine] {
                    for wrap1_density in [1.0, 0.5, 0.2] {
                        for wrap2_density in [1.0, 0.5, 0.2, 0.0] {
                            for (wrap1_cov, wrap2_cov) in [
                                (0.8, 0.8),
                                (0.8, 0.25),
                                (0.25, 0.8),
                            ] {
                                variants.push((
                                    peak_clock_current,
                                    post_time,
                                    clock_type,
                                    ferro_major,
                                    wrap1_cov,
                                    wrap2_cov,
                                    wrap1_density,
                                    wrap2_density
                                ));
                            }
                        }
                    }
                }
            }
        }
    }

    for (peak_clock_current, post_time, clock_type, ferro_major, wrap1_coverage, wrap2_coverage, wrap1_density, wrap2_density) in variants {
        info!("{}A {}s {}m {}:{}cov {}:{}density", peak_clock_current, post_time, ferro_major, wrap1_coverage, wrap2_coverage, wrap1_density, wrap2_density);
        let base_params = Params {
            dry_run,
            feat_size: 40e-6f32,

            buffer_xy: 160e-6,
            buffer_z: 160e-6,
            boundary_xy: 320e-6,
            boundary_z: 320e-6,

            ferro_major,
            ferro_minor: 60e-6,
            ferro_buffer: 1320e-6,
            wire_minor: 40e-6,
            wire_wrap_minor: 50e-6,
            wire_set_major: 200e-6,
            wire_wrap_dilation: 50e-6,
            wire_wrap_iters: 3,

            wraps1: 4.0,
            wraps2: 4.0,
            wrap1_coverage,
            wrap2_coverage,

            wire_conductivity: 5e6f32,

            peak_set_current: 60.0,
            peak_clock_current,
            set_duration: 0e-9,
            clock_duration: 25e-9,
            clock_type,
            pre_time: 1e-9,
            post_time,

            dump_frames: Some(256000),
        };
        let wraps1_max = (1..25)
            .into_iter()
            .filter_map(|wraps1| {
                let params = Params {
                    wraps1: (wraps1 * 4) as f32,
                    ..base_params
                };
                let geoms = derive_geometries(params.clone())?;
                Some((params, geoms))
            }).max_by(|(_p1, geoms1), (_p2, geoms2)|
                geoms1.wrap1_len.partial_cmp(&geoms2.wrap1_len).unwrap()
            ).unwrap()
            .0.wraps1;
        let wraps1 = (wraps1_max * wrap1_density / 4.0).round().max(1.0) * 4.0;

        let wraps2_max = (1..25)
            .into_iter()
            .filter_map(|wraps2| {
                let params = Params {
                    wraps2: (wraps2 * 4) as f32,
                    ..base_params
                };
                let geoms = derive_geometries(params.clone())?;
                Some((params, geoms))
            }).max_by(|(_p1, geoms1), (_p2, geoms2)|
                geoms1.wrap2_len.partial_cmp(&geoms2.wrap2_len).unwrap()
            ).unwrap()
            .0.wraps2;
        let wraps2 = (wraps2_max * wrap2_density / 4.0).round().max(1.0) * 4.0;
        let params = Params { wraps1, wraps2, ..base_params };
        let geoms = derive_geometries(params.clone()).unwrap();
        run_sim(i, params, geoms);
    }
}