use coremem::{Driver, mat, meas, SpirvDriver};
use coremem::geom::{Meters, Torus};
use coremem::sim::units::Seconds;
use coremem::stim::{CurlStimulus, Sinusoid, TimeVarying as _};

fn main() {
    coremem::init_logging();
    let feat_size = 40e-6f32;
    let depth = 1600e-6;
    let buffer_x = 240e-6;
    let buffer_y = 240e-6;
    // let buffer_z = 240e-6;
    let ferro_major = 320e-6;
    let ferro_minor = 60e-6;
    let ferro_buffer = 60e-6;
    let wire_minor = 40e-6;
    let wire_major = 160e-6;
    let wire_coupling_major = 280e-6; // (ferro_minor*4 + ferro_buffer)/2 + wire_minor = 190
    let peak_current = 1.5e6;
    let current_duration = 24e-9; // half-wavelength of the sine wave
    // let current_break = 0.2e-9; // time between 'set' pulse and 'clear' pulse
    let drive_conductivity = 5e6f32;
    let sense_conductivity = 5e3f32;
    // let coupling_offsets = [0.0];
    let coupling_offsets = [-1.5*wire_minor, 1.5*wire_minor];

    let m_to_um = |m: f32| (m * 1e6).round() as u32;
    let half_depth = depth * 0.5;

    let base = "buffer-22";

    let boundary_xy = 500e-6;
    let boundary_z = 300e-6;
    let ferro_top_mid = boundary_xy + buffer_y + wire_minor + wire_major;
    let ferro_center_y = ferro_top_mid + ferro_major;
    let ferro_bot_mid = ferro_center_y + ferro_major;
    let height = ferro_bot_mid + wire_major + wire_minor + buffer_y + boundary_xy;
    let ferro1_left_edge = boundary_xy + buffer_x;
    let ferro1_center = ferro1_left_edge + ferro_minor + ferro_major;
    let ferro1_right_edge = ferro1_center + ferro_major + ferro_minor;
    let ferro2_left_edge = ferro1_right_edge + ferro_buffer;
    let ferro2_center = ferro2_left_edge + ferro_minor + ferro_major;
    let ferro2_right_edge = ferro2_center + ferro_major + ferro_minor;
    let width = ferro2_right_edge + wire_major + wire_minor + buffer_x + boundary_xy;

    let ferro1_region = Torus::new_xy(Meters::new(ferro1_center, ferro_center_y, half_depth), ferro_major, ferro_minor);
    let ferro2_region = Torus::new_xy(Meters::new(ferro2_center, ferro_center_y, half_depth), ferro_major, ferro_minor);
    // SET connected to top of ferro1
    let set_region = Torus::new_yz(Meters::new(ferro1_center, ferro_center_y - ferro_major, half_depth), wire_major, wire_minor);
    // RESET connected to bottom of ferro1
    let reset_region = Torus::new_yz(Meters::new(ferro1_center, ferro_center_y + ferro_major, half_depth), wire_major, wire_minor);
    let coupling_regions: Vec<_> = coupling_offsets.iter().map(|off| {
        Torus::new_xz(Meters::new(0.5*(ferro1_center + ferro2_center), ferro_center_y + off, half_depth), wire_coupling_major, wire_minor)
    }).collect();
    // SENSE connected to right of ferro2
    let sense_region = Torus::new_xz(Meters::new(ferro2_center + ferro_major, ferro_center_y, half_depth), wire_major, wire_minor);
    // TODO: make sure none of the regions overlap

    let mut driver: SpirvDriver = Driver::new_spirv(Meters::new(width, height, depth), feat_size);
    driver.set_steps_per_stim(1000);
    //driver.fill_region(&ferro1_region, mat::db::linear_iron());
    // Original, 3R1-LIKE ferromagnet (only a vague likeness), sr-latch-8:
    // driver.fill_region(&ferro1_region, mat::MBPgram::new(-0.3899, 0.3900, 310_000.0));
    // driver.fill_region(&ferro2_region, mat::MBPgram::new(-0.3899, 0.3900, 310_000.0));
    // sr-latch-9; dead spot from B=[-0.03, 0.03]. This will help us see if the math is H-triggered
    // or B-triggered
    // driver.fill_region(&ferro1_region, mat::MBPgram::new(-0.3300, 0.3900, 310_000.0));
    // driver.fill_region(&ferro2_region, mat::MBPgram::new(-0.3300, 0.3900, 310_000.0));
    // mu_r=881.33, starting at H=25 to H=75.
    driver.fill_region(&ferro1_region, mat::MHPgram::new(25.0, 881.33, 44000.0));
    driver.fill_region(&ferro2_region, mat::MHPgram::new(25.0, 881.33, 44000.0));
    driver.fill_region(&set_region, mat::IsomorphicConductor::new(drive_conductivity));
    driver.fill_region(&reset_region, mat::IsomorphicConductor::new(drive_conductivity));
    for r in &coupling_regions {
        driver.fill_region(r, mat::IsomorphicConductor::new(drive_conductivity));
    }
    driver.fill_region(&sense_region, mat::IsomorphicConductor::new(sense_conductivity));

    println!("boundary: {}um; {}um", m_to_um(boundary_xy), m_to_um(boundary_z));
    driver.add_classical_boundary(Meters::new(boundary_xy, boundary_xy, boundary_z));

    assert!(driver.test_region_filled(&ferro1_region, mat::MHPgram::new(25.0, 881.33, 44000.0)));
    assert!(driver.test_region_filled(&ferro2_region, mat::MHPgram::new(25.0, 881.33, 44000.0)));
    assert!(driver.test_region_filled(&set_region, mat::IsomorphicConductor::new(drive_conductivity)));
    assert!(driver.test_region_filled(&reset_region, mat::IsomorphicConductor::new(drive_conductivity)));
    for r in &coupling_regions {
        assert!(driver.test_region_filled(r, mat::IsomorphicConductor::new(drive_conductivity)));
    }
    assert!(driver.test_region_filled(&sense_region, mat::IsomorphicConductor::new(sense_conductivity)));

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

    // 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_stim1 = peak_current/current_duration / (set_region.cross_section() * drive_conductivity);

    // SET
    add_drive_pulse(&set_region, 0.01*current_duration, current_duration, peak_stim1);

    let duration = 20.0*current_duration;

    for (i, r) in coupling_regions.iter().enumerate() {
        driver.add_measurement(meas::CurrentLoop::new(&*format!("coupling{}", i), r.clone()));
        driver.add_measurement(meas::Current::new(&*format!("coupling{}", i), r.clone()));
    }
    driver.add_measurement(meas::CurrentLoop::new("sense", sense_region.clone()));
    driver.add_measurement(meas::Current::new("sense", sense_region.clone()));

    driver.add_measurement(meas::Volume::new("mem1", ferro1_region.clone()));
    driver.add_measurement(meas::MagneticLoop::new("mem1", ferro1_region.clone()));
    driver.add_measurement(meas::Magnetization::new("mem1", ferro1_region.clone()));
    driver.add_measurement(meas::MagneticFlux::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("set", set_region.clone()));
    driver.add_measurement(meas::Current::new("set", set_region.clone()));
    driver.add_measurement(meas::Power::new("set", set_region.clone()));
    driver.add_measurement(meas::CurrentLoop::new("reset", reset_region.clone()));

    let prefix = format!("out/{}/{}-{}mA-{}ps-{}um-{}xcoupling",
        base,
        base,
        (peak_current * 1e3).round() as i64,
        (current_duration * 1e12).round() as i64,
        (feat_size * 1e6).round() as i64,
        coupling_regions.len(),
    );
    let _ = std::fs::create_dir_all(&prefix);

    driver.add_state_file(&*format!("{}/state.bc", prefix), 9600);
    // driver.add_serializer_renderer(&*format!("{}/frame-", prefix), 36000);
    driver.add_csv_renderer(&*format!("{}/meas.csv", prefix), 200, None);
    driver.add_csv_renderer(&*format!("{}/meas-sparse.csv", prefix), 8000, None);

    driver.step_until(Seconds(duration));
}