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

fn main() {
    coremem::init_logging();
    let feat_size = 10e-6f32; // feature size
    let duration = 60.0e-9;
    let width = 3400e-6;
    let height = 1800e-6;
    let depth = 1400e-6;
    let buffer = 200e-6;
    let ferro_spacing = 600e-6;
    let ferro_major = 320e-6;
    let ferro_minor = 60e-6;
    let wire_minor = 40e-6;
    let wire_major = 160e-6;
    let peak_current1 = 7.5e6;
    //let peak_current2 = 5e5;
    let peak_current2 = 7.5e6;
    let current_duration = 1.0e-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 from_m = |m: f32| (m/feat_size).round() as u32;
    let m_to_um = |m: f32| (m * 1e6).round() as u32;
    let half_width = width * 0.5;
    let half_height = height * 0.5;
    let half_depth = depth * 0.5;

    let width_px = from_m(width);
    let height_px = from_m(height);
    let depth_px = from_m(depth);
    let size_px = Index((width_px, height_px, depth_px).into());
    let mut driver: SpirvDriver = Driver::new_spirv(size_px, feat_size);
    // let mut driver: Driver<SimState<f32, mat::GenericMaterial<f32>>> = Driver::new(size_px, feat_size);
    driver.set_steps_per_stim(1000);
    let base = "wrapped_torus-44-classical-boundary-spirv-100steps-stepfn-on-device-eh_reorder";

    let ferro1_center = half_width - ferro_major - 0.5*ferro_spacing;
    let ferro1_region = Torus::new_xy(Meters::new(ferro1_center, half_height, half_depth), ferro_major, ferro_minor);
    let drive1_region = Torus::new_xz(Meters::new(ferro1_center - ferro_major, half_height, half_depth), wire_major, wire_minor);
    let sense1_region = Torus::new_xz(Meters::new(ferro1_center + ferro_major, half_height, half_depth), wire_major, wire_minor);

    //driver.fill_region(&ferro1_region, mat::db::linear_iron());
    driver.fill_region(&ferro1_region, mat::MBPgram::new(-0.3899, 0.3900, 310_000.0));
    driver.fill_region(&drive1_region, mat::IsomorphicConductor::new(drive_conductivity));
    driver.fill_region(&sense1_region, mat::IsomorphicConductor::new(sense_conductivity));

    let ferro2_center = half_width + ferro_major + 0.5*ferro_spacing;
    let ferro2_region = Torus::new_xy(Meters::new(ferro2_center, half_height, half_depth), ferro_major, ferro_minor);
    let drive2_region = Torus::new_xz(Meters::new(ferro2_center - ferro_major, half_height, half_depth), wire_major, wire_minor);
    let sense2_region = Torus::new_xz(Meters::new(ferro2_center + ferro_major, half_height, half_depth), wire_major, wire_minor);

    driver.fill_region(&ferro2_region, mat::MBPgram::new(-0.3899, 0.3900, 310_000.0));
    driver.fill_region(&drive2_region, mat::IsomorphicConductor::new(drive_conductivity));
    driver.fill_region(&sense2_region, mat::IsomorphicConductor::new(sense_conductivity));

    let boundary_xy = ferro1_center - ferro_major - ferro_minor - buffer;
    let boundary_z = half_depth - wire_major - wire_minor - buffer;
    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));

    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 sin(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_current1/current_duration / (drive1_region.cross_section() * drive_conductivity);

    add_drive_pulse(&drive1_region, 0.0, current_duration, peak_stim1);
    // add_drive_pulse(&drive1_region, current_duration + current_break, current_duration, -4.0*peak_stim1);
    // add_drive_pulse(&drive1_region, 2.0*(current_duration + current_break), current_duration, -4.0*peak_stim1);
    // add_drive_pulse(&drive1_region, 3.0*(current_duration + current_break), current_duration, peak_stim1);

    let peak_stim2 = peak_current2/current_duration / (drive2_region.cross_section() * drive_conductivity);
    add_drive_pulse(&drive2_region, 0.0, current_duration, peak_stim2);
    // add_drive_pulse(&drive2_region, current_duration + current_break, current_duration, -4.0*peak_stim2);
    // add_drive_pulse(&drive2_region, 2.0*(current_duration + current_break), current_duration, -4.0*peak_stim2);
    // add_drive_pulse(&drive2_region, 3.0*(current_duration + current_break), current_duration, peak_stim2);

    driver.add_measurement(meas::CurrentLoop::new("sense1", sense1_region.clone()));
    driver.add_measurement(meas::Current::new("sense1", sense1_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::CurrentLoop::new("drive1", drive1_region.clone()));
    driver.add_measurement(meas::Current::new("drive1", drive1_region.clone()));
    driver.add_measurement(meas::Power::new("drive1", drive1_region.clone()));

    driver.add_measurement(meas::CurrentLoop::new("sense2", sense2_region.clone()));
    driver.add_measurement(meas::Current::new("sense2", sense2_region.clone()));
    driver.add_measurement(meas::MagneticLoop::new("mem2", ferro2_region.clone()));
    driver.add_measurement(meas::Magnetization::new("mem2", ferro2_region.clone()));
    driver.add_measurement(meas::MagneticFlux::new("mem2", ferro2_region.clone()));
    driver.add_measurement(meas::CurrentLoop::new("drive2", drive2_region.clone()));
    driver.add_measurement(meas::Current::new("drive2", drive2_region.clone()));
    driver.add_measurement(meas::Power::new("drive2", drive2_region.clone()));

    let prefix = format!("out/{}/{}-flt{}-{}-feat{}um-{}mA-{}ps--radii{}um-{}um-{}um-{}um",
        base,
        base,
        32,
        *size_px,
        m_to_um(feat_size),
        (peak_current1 * 1e3).round() as i64,
        (current_duration * 1e12).round() as i64,
        m_to_um(ferro_major),
        m_to_um(ferro_minor),
        m_to_um(wire_major),
        m_to_um(wire_minor),
    );
    let _ = std::fs::create_dir_all(&prefix);

    // driver.add_state_file(&*format!("{}/state.bc", prefix), 1000);
    // driver.add_serializer_renderer(&*format!("{}/frame-", prefix), 1000);
    driver.add_csv_renderer(&*format!("{}/meas.csv", prefix), 100, None);

    driver.step_until(Seconds(duration));
}