//! 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;

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

/// 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
}

#[derive(Copy, Clone, Debug)]
struct Params {
    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,
    wraps2: f32,

    wire_conductivity: f32,

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

    dump_frames: Option<u64>,
}

#[derive(Clone, Debug)]
struct Results {
    m1: Vec<f32>,
    m2: Vec<f32>,
    h1: Vec<f32>,
    h2: Vec<f32>,
    iset1: Vec<f32>,
    icoupling: Vec<f32>,
}

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

    let base = format!("buffer5-{}", id);

    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 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 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, 1.0/p.wraps1),
                    Cube::new(Meters::new(-1.0, -1.0, -0.5), Meters::new(1.0, 1.0, 0.5))
                )),
                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 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, 1.0/p.wraps2),
                    Cube::new(Meters::new(-1.0, -1.0, -0.5), Meters::new(1.0, 1.0, 0.5))
                )),
                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 + 2.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 + 2.0*p.ferro_major + 4.0*p.feat_size, 2.0*p.feat_size, 2.0*p.feat_size)
    );
    let coupling_stub_top_left = Cube::new_including_negatives(
        coupling_wire_top.bot_left_out(),
        coupling_wire_top.bot_left_out() + Meters::new(4.0*p.feat_size, 10.0*p.feat_size, 2.0*p.feat_size)
    );
    let coupling_stub_top_right = Cube::new_including_negatives(
        coupling_wire_top.bot_right_out(),
        coupling_wire_top.bot_right_out() + Meters::new(-4.0*p.feat_size, 10.0*p.feat_size, 2.0*p.feat_size)
    );
    let coupling_stub_bot_left = Cube::new_including_negatives(
        coupling_wire_bot.top_left_out(),
        coupling_wire_bot.top_left_out() + Meters::new(4.0*p.feat_size, -10.0*p.feat_size, 2.0*p.feat_size)
    );
    let coupling_stub_bot_right = Cube::new_including_negatives(
        coupling_wire_bot.top_right_out(),
        coupling_wire_bot.top_right_out() + Meters::new(-4.0*p.feat_size, -10.0*p.feat_size, 2.0*p.feat_size)
    );
    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
    assert!(region::is_connected(
        &wrap1_with_coupling, 
        coupling_wire_top.center(),
        coupling_wire_bot.center(),
        p.feat_size,
    ), "{:?}", p);
    assert!(region::is_connected(
        &wrap2_with_coupling, 
        coupling_wire_top.center(),
        coupling_wire_bot.center(),
        p.feat_size,
    ), "{:?}", p);

    assert!(wrap1_with_coupling.contains(
        coupling_stub_top_left.center().to_index(p.feat_size).to_meters(p.feat_size),
    ));
    assert!(wrap1_with_coupling.contains(
        coupling_stub_bot_left.center().to_index(p.feat_size).to_meters(p.feat_size),
    ));
    assert!(wrap2_with_coupling.contains(
        coupling_stub_top_right.center().to_index(p.feat_size).to_meters(p.feat_size),
    ));
    assert!(wrap2_with_coupling.contains(
        coupling_stub_bot_right.center().to_index(p.feat_size).to_meters(p.feat_size),
    ));
    info!("wrap1 length: {}", 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());
    info!("wrap2 length: {}", 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());

    // 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(Meters::new(width, height, depth), p.feat_size);
    driver.set_steps_per_stim(1000);
    driver.fill_region(&ferro1_region, ferro_mat);
    driver.fill_region(&ferro2_region, ferro_mat);
    driver.fill_region(&set1_region, wire_mat);
    driver.fill_region(&set2_region, wire_mat);
    driver.fill_region(&coupling_region, wire_mat);

    info!("boundary: {}um; {}um", m_to_um(p.boundary_xy), m_to_um(p.boundary_z));
    info!("size: {}, {}, {}", width, height, depth);
    info!("ferro1: {:?}", ferro1_center);
    info!("ferro2: {:?}", ferro2_center);
    driver.add_classical_boundary(Meters::new(p.boundary_xy, p.boundary_xy, p.boundary_z));

    assert!(driver.test_region_filled(&ferro1_region, ferro_mat));
    assert!(driver.test_region_filled(&ferro2_region, ferro_mat));
    assert!(driver.test_region_filled(&set1_region, wire_mat));
    assert!(driver.test_region_filled(&set2_region, wire_mat));
    assert!(driver.test_region_filled(&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);
    };

    // 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 / (set1_region.cross_section() * p.wire_conductivity);
    let peak_clock = p.peak_clock_current / feat_vol / (set1_region.cross_section() * p.wire_conductivity);

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

    // CLEAR core1
    add_drive_square_pulse(&mut driver, &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", ferro1_region.clone()));
    driver.add_measurement(meas::MagneticLoop::new("mem1", ferro1_region.clone()));

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

    driver.add_measurement(meas::CurrentLoop::new("set1", set1_region.clone()));
    driver.add_measurement(meas::Power::new("set1", set1_region.clone()));
    driver.add_measurement(meas::CurrentLoop::new("set2", 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", coupling_wire_top.clone()));
    driver.add_measurement(meas::Current::new("couplingbot", coupling_wire_bot.clone()));

    let prefix = format!("out/{}/{}-{}-{}setmA-{}setps-{}clkmA-{}clkps-{}um-{}ferromaj-{}:{}wraps",
        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,
    );
    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 res = Results {
        m1: extrema(CsvRenderer::new(&*meas_sparse_csv).read_column_as_f32("M(mem1)")),
        m2: extrema(CsvRenderer::new(&*meas_sparse_csv).read_column_as_f32("M(mem2)")),
        h1: extrema(CsvRenderer::new(&*meas_sparse_csv).read_column_as_f32("H(mem1)")),
        h2: extrema(CsvRenderer::new(&*meas_sparse_csv).read_column_as_f32("H(mem2)")),
        iset1: extrema(CsvRenderer::new(&*meas_sparse_csv).read_column_as_f32("I(set1)")),
        icoupling: extrema(CsvRenderer::new(&*meas_sparse_csv).read_column_as_f32("Imag/cell(couplingtop)")),
    };
    std::fs::write(
        format!("{}/results.txt", prefix),
        format!("{:?}", res),
    ).unwrap();
    res
}


fn main() {
    coremem::init_logging();
    for (i, dump_frames, wraps1, wraps2) in [
        // (38, Some(1000000), 20.0, 12.0),
        // (39, Some(1000000), 20.0, 8.0),
        // (51, Some(1000000), 16.0, 16.0),
        // (41, Some(1000000), 20.0, 20.0),
        // // (42, Some(1000000), 24.0, 12.0), // FAILS auto connection check
        // // (43, Some(1000000), 28.0, 12.0),
        // // (44, Some(1000000), 32.0, 12.0),
        // // (45, Some(1000000), 24.0, 16.0),
        // // (46, Some(1000000), 24.0, 20.0),
        (47, Some(1000000), 24.0, 16.0),
        (48, Some(1000000), 24.0, 20.0), // passes auto connection check
        (49, Some(1000000), 28.0, 16.0), // passes auto connection check
        (50, Some(1000000), 28.0, 20.0), // passes auto connection check
        // // (51, Some(1000000), 32.0, 16.0), // FAILS auto connection check
    ].iter().copied() {
        run_sim(i, Params {
            feat_size: 40e-6f32,

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

            ferro_major: 1360e-6,
            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,
            wraps2,

            wire_conductivity: 5e6f32,

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

            dump_frames,
        });
    }
}