use naga::GlobalVariable;
use num::Complex;
use std::borrow::Cow;
use std::num::NonZeroU64;
use std::sync::mpsc::{Receiver, Sender};
use wgpu::naga;

#[cfg(feature = "gpu")]
#[macro_export]
macro_rules! include_wgsl {
    ($filename:literal) => {
        include_str!(concat!(env!("OUT_DIR"), "/shaders/", $filename))
    };
}

// Magic to allow converting a primitive type, a Vec, or a slice into a slice of bytes.
pub trait Buffer<T> {
    fn as_bytes(&self) -> &[T];
}
impl<T: bytemuck::Pod> Buffer<T> for T {
    fn as_bytes(&self) -> &[T] {
        std::slice::from_ref(self)
    }
}
impl<T> Buffer<T> for Vec<T> {
    fn as_bytes(&self) -> &[T] {
        self.as_slice()
    }
}
impl<T> Buffer<T> for &[T] {
    fn as_bytes(&self) -> &[T] {
        self
    }
}

pub struct BufferBuilder<'a> {
    label: Option<String>,
    size: u64,
    usage: wgpu::BufferUsages,
    mapped_at_creation: bool,
    gpu: &'a Gpu,
}

impl<'a> BufferBuilder<'a> {
    fn new<T: bytemuck::Pod>(data: &impl Buffer<T>, gpu: &'a Gpu) -> Self {
        Self::new_by_sz(std::mem::size_of_val(data.as_bytes()), gpu)
    }

    fn new_by_sz(sz: usize, gpu: &'a Gpu) -> Self {
        Self {
            label: None,
            size: wgpu::util::align_to(sz, 16) as u64,
            usage: wgpu::BufferUsages::empty(),
            mapped_at_creation: false,
            gpu: gpu,
        }
    }

    pub fn copy_dst(&mut self) -> &mut Self {
        self.usage |= wgpu::BufferUsages::COPY_DST;
        self
    }

    pub fn copy_src(&mut self) -> &mut Self {
        self.usage |= wgpu::BufferUsages::COPY_SRC;
        self
    }

    pub fn map_read(&mut self) -> &mut Self {
        self.usage |= wgpu::BufferUsages::MAP_READ;
        self
    }

    pub fn query_resolve(&mut self) -> &mut Self {
        self.usage |= wgpu::BufferUsages::QUERY_RESOLVE;
        self
    }

    pub fn storage(&mut self) -> &mut Self {
        self.usage |= wgpu::BufferUsages::STORAGE;
        self
    }

    pub fn uniform(&mut self) -> &mut Self {
        self.usage |= wgpu::BufferUsages::UNIFORM;
        self
    }

    pub fn build(&self) -> wgpu::Buffer {
        self.gpu.device.create_buffer(&wgpu::BufferDescriptor {
            label: self.label.as_deref(),
            size: self.size,
            usage: self.usage,
            mapped_at_creation: self.mapped_at_creation,
        })
    }
}

pub struct GpuBuilder {
    required_features: wgpu::Features,
}

impl GpuBuilder {
    pub fn new() -> Self {
        Self {
            required_features: wgpu::Features::empty(),
        }
    }

    pub fn enable_profiling(&mut self) -> &mut Self {
        self.required_features |= wgpu::Features::TIMESTAMP_QUERY;
        self
    }

    pub fn build(&self) -> anyhow::Result<Gpu> {
        let instance = wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle());
        let adapter =
            pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions::default()))?;
        println!("Got adapter: {:#?}", adapter.get_info());

        let supported = adapter.limits();
        let mut limits = wgpu::Limits::downlevel_defaults();
        limits.max_compute_invocations_per_workgroup =
            supported.max_compute_invocations_per_workgroup;
        limits.max_compute_workgroup_storage_size = supported.max_compute_workgroup_storage_size;
        println!(
            "Max compute invocations per workgroup: {}",
            limits.max_compute_invocations_per_workgroup
        );

        let (device, queue) =
            pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor {
                label: None,
                required_features: self.required_features,
                required_limits: limits,
                #[cfg(target_arch = "wasm32")]
                default_queue: wgpu::QueueDescriptor { label: None },
                experimental_features: wgpu::ExperimentalFeatures::disabled(),
                memory_hints: wgpu::MemoryHints::MemoryUsage,
                trace: wgpu::Trace::Off,
            }))?;

        let mut timestamps = None;
        if self
            .required_features
            .contains(wgpu::Features::TIMESTAMP_QUERY)
        {
            timestamps = Some(device.create_query_set(&wgpu::QuerySetDescriptor {
                label: None,
                ty: wgpu::QueryType::Timestamp,
                count: 2,
            }));
        }

        Ok(Gpu {
            adapter,
            device,
            queue,
            timestamps,
        })
    }
}

pub struct Gpu {
    adapter: wgpu::Adapter,
    device: wgpu::Device,
    queue: wgpu::Queue,
    timestamps: Option<wgpu::QuerySet>,
}

impl Gpu {
    pub fn buffer<T: bytemuck::Pod>(&self, data: &impl Buffer<T>) -> BufferBuilder<'_> {
        BufferBuilder::new(data, self)
    }

    pub fn buffer_by_sz(&self, sz: usize) -> BufferBuilder<'_> {
        BufferBuilder::new_by_sz(sz, self)
    }
}

fn get_entry_point_name(module: &naga::Module) -> anyhow::Result<String> {
    if module.entry_points.len() != 1 {
        anyhow::bail!(
            "Expected exactly one entry point, found {}",
            module.entry_points.len()
        );
    }
    for entry in module.entry_points.iter() {
        return Ok(entry.name.clone());
    }
    anyhow::bail!("Could not find entry point");
}

fn get_entry_point_workgroup_size(module: &naga::Module, name: &str) -> anyhow::Result<[u32; 3]> {
    for entry_point in module.entry_points.iter() {
        if entry_point.name != name {
            continue;
        }
        return Ok(entry_point.workgroup_size);
    }
    anyhow::bail!("Could not find entry point named {}", name);
}

fn for_each_global_variable<T>(
    module: &naga::Module,
    module_info: &naga::valid::ModuleInfo,
    mut f: impl FnMut(
        naga::Handle<GlobalVariable>,
        &naga::GlobalVariable,
        &naga::ResourceBinding,
        &wgpu::ShaderStages,
        bool,
    ) -> anyhow::Result<Option<T>>,
) -> anyhow::Result<Vec<T>> {
    let mut result = Vec::new();
    for (global_handle, global) in module.global_variables.iter() {
        let Some(binding) = &global.binding else {
            continue;
        };

        let can_write = match global.space {
            naga::AddressSpace::Storage { access } => {
                access.contains(naga::StorageAccess::STORE)
                    || access.contains(naga::StorageAccess::ATOMIC)
            }
            naga::AddressSpace::Uniform => false,
            _ => {
                eprintln!(
                    "Skipping global {:?} with address space {:?}",
                    global.name, global.space
                );
                continue;
            }
        };
        let mut stages = wgpu::ShaderStages::empty();
        for (entry_idx, entry) in module.entry_points.iter().enumerate() {
            let usage = module_info.get_entry_point(entry_idx)[global_handle];
            if usage.is_empty() {
                continue;
            }

            stages |= match entry.stage {
                naga::ShaderStage::Vertex => wgpu::ShaderStages::VERTEX,
                naga::ShaderStage::Fragment => wgpu::ShaderStages::FRAGMENT,
                naga::ShaderStage::Compute => wgpu::ShaderStages::COMPUTE,
                _ => continue,
            }
        }
        if stages.is_empty() {
            eprintln!(
                "Skipping bound global {:?} which is not referenced by any stage",
                global.name
            );
            continue;
        }
        if let Some(value) = f(global_handle, global, binding, &stages, can_write)? {
            result.push(value);
        }
    }
    Ok(result)
}

fn get_bind_group_layout(
    gpu: &Gpu,
    module: &naga::Module,
    module_info: &naga::valid::ModuleInfo,
    module_layout: &naga::proc::Layouter,
) -> anyhow::Result<wgpu::BindGroupLayout> {
    let mut entries = Vec::<wgpu::BindGroupLayoutEntry>::new();
    for (global_handle, global) in module.global_variables.iter() {
        let Some(binding) = &global.binding else {
            continue;
        };

        let can_write = match global.space {
            naga::AddressSpace::Storage { access } => {
                access.contains(naga::StorageAccess::STORE)
                    || access.contains(naga::StorageAccess::ATOMIC)
            }
            naga::AddressSpace::Uniform => false,
            _ => {
                eprintln!(
                    "Skipping global {:?} with address space {:?}",
                    global.name, global.space
                );
                continue;
            }
        };

        let mut stages = wgpu::ShaderStages::empty();
        for (entry_idx, entry) in module.entry_points.iter().enumerate() {
            let usage = module_info.get_entry_point(entry_idx)[global_handle];
            if usage.is_empty() {
                continue;
            }

            stages |= match entry.stage {
                naga::ShaderStage::Vertex => wgpu::ShaderStages::VERTEX,
                naga::ShaderStage::Fragment => wgpu::ShaderStages::FRAGMENT,
                naga::ShaderStage::Compute => wgpu::ShaderStages::COMPUTE,
                _ => continue,
            }
        }
        if stages.is_empty() {
            eprintln!(
                "Skipping bound global {:?} which is not referenced by any stage",
                global.name
            );
            continue;
        }

        let global_ty = &module.types[global.ty];
        let entry = match &module.types[global.ty].inner {
            naga::ir::TypeInner::Array { .. } => {
                let min_size = module_layout[global.ty].size;
                wgpu::BindGroupLayoutEntry {
                    binding: binding.binding,
                    visibility: stages,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Storage {
                            read_only: !can_write,
                        },
                        min_binding_size: NonZeroU64::new(min_size as u64),
                        has_dynamic_offset: false,
                    },
                    count: None,
                }
            }
            naga::ir::TypeInner::Struct { .. } => {
                let min_size = module_layout[global.ty].size;
                wgpu::BindGroupLayoutEntry {
                    binding: binding.binding,
                    visibility: stages,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Uniform,
                        min_binding_size: NonZeroU64::new(min_size as u64),
                        has_dynamic_offset: false,
                    },
                    count: None,
                }
            }
            _ => {
                anyhow::bail!(format!(
                    "Unsupported global type: {:?} (name={:?})",
                    global_ty.name, global_ty.inner
                ))
            }
        };
        entries.push(entry);
    }
    let layout = wgpu::BindGroupLayoutDescriptor {
        label: None,
        entries: &entries,
    };
    Ok(gpu.device.create_bind_group_layout(&layout))
}

pub struct Shader {
    bind_group_layout: wgpu::BindGroupLayout,
    shader_module: wgpu::ShaderModule,
    entry_point_name: String,
    entry_point_workgroup_size: [u32; 3],
}

impl Shader {
    pub fn new(gpu: &Gpu, shader: &str) -> anyhow::Result<Self> {
        let module = naga::front::wgsl::parse_str(shader)?;
        let module_info = naga::valid::Validator::new(
            naga::valid::ValidationFlags::all(),
            naga::valid::Capabilities::all(),
        )
        .validate(&module)?;
        let mut module_layout = naga::proc::Layouter::default();
        module_layout.update(module.to_ctx())?;

        Self::ensure_compute_compatibility(gpu, &module)?;
        let bind_group_layout = get_bind_group_layout(gpu, &module, &module_info, &module_layout)?;
        let entry_point_name = get_entry_point_name(&module)?;
        let entry_point_workgroup_size =
            get_entry_point_workgroup_size(&module, &entry_point_name)?;

        let module_descriptor = wgpu::ShaderModuleDescriptor {
            label: None,
            source: wgpu::ShaderSource::Naga(Cow::Owned(module.clone())),
        };
        let shader_module = gpu.device.create_shader_module(module_descriptor.clone());

        Ok(Self {
            bind_group_layout,
            // naga_module: module,
            // naga_module_info: module_info,
            shader_module,
            entry_point_name,
            entry_point_workgroup_size,
            // module_layout,
        })
    }

    fn ensure_compute_compatibility(gpu: &Gpu, module: &naga::Module) -> anyhow::Result<()> {
        let mut has_compute_stage = false;
        for entry_point in module.entry_points.iter() {
            if entry_point.stage == naga::ShaderStage::Compute {
                has_compute_stage = true;
            }
        }
        if !has_compute_stage {
            return Ok(());
        }
        let downlevel_caps = gpu.adapter.get_downlevel_capabilities();
        if !downlevel_caps
            .flags
            .contains(wgpu::DownlevelFlags::COMPUTE_SHADERS)
        {
            anyhow::bail!(
                "GPU {} does not support compute shaders.",
                gpu.adapter.get_info().name
            )
        }
        Ok(())
    }
}

pub struct MaterialPass {
    buffers: Vec<wgpu::Buffer>,
    bind_group: wgpu::BindGroup,
}

pub struct MaterialPassBuilder {
    buffers: Vec<wgpu::Buffer>,
    bind_group: Option<wgpu::BindGroup>,
}

impl MaterialPassBuilder {
    pub fn new() -> Self {
        Self {
            buffers: Vec::new(),
            bind_group: None,
        }
    }

    pub fn buffers(mut self, buffers: Vec<wgpu::Buffer>) -> Self {
        self.buffers = buffers;
        self
    }

    // Must be run after buffers(), etc.
    pub fn bind_group(mut self, gpu: &Gpu, shader: &Shader) -> Self {
        let mut descriptor_entries = Vec::<wgpu::BindGroupEntry>::new();
        let entries =
            self.buffers
                .iter()
                .enumerate()
                .fold(Vec::new(), |mut entries, (index, buf)| {
                    entries.push(wgpu::BindGroupEntry {
                        binding: index as u32,
                        resource: buf.as_entire_binding(),
                    });
                    entries
                });
        descriptor_entries.extend(entries);
        let descriptor = wgpu::BindGroupDescriptor {
            label: None,
            layout: &shader.bind_group_layout,
            entries: &descriptor_entries,
        };
        let bind_group = gpu.device.create_bind_group(&descriptor);
        self.bind_group = Some(bind_group);
        self
    }

    pub fn build(self) -> anyhow::Result<MaterialPass> {
        if self.bind_group.is_none() {
            anyhow::bail!("bind_group not set");
        }
        Ok(MaterialPass {
            buffers: self.buffers,
            bind_group: self.bind_group.unwrap(),
        })
    }
}

pub struct MaterialInstance {
    passes: Vec<MaterialPass>,
    readback: Option<wgpu::Buffer>,
}

impl MaterialInstance {
    pub fn new(passes: Vec<MaterialPass>, readback: Option<wgpu::Buffer>) -> Self {
        Self {
            passes: passes,
            readback: readback,
        }
    }
}

// A batch of Material work to be submitted together. This allows us to send a
// batch of work while another is still processing on the GPU.
pub struct MaterialBatch {
    pub instances: Vec<MaterialInstance>,
    ready_tx: Sender<(usize, Result<(), wgpu::BufferAsyncError>)>,
    ready_rx: Receiver<(usize, Result<(), wgpu::BufferAsyncError>)>,
    // Optional timestamp (ts) buffers for profiling.
    timing: Option<MaterialTiming>,
    // Whether the batch is currently processing. This is set true as soon as a batch is reserved with get_batch().
    // The caller is responsible for setting this false
    pub busy: bool,
    // How many instances in the batch are currently processing. Set to
    // batch_size (n_instancces) when the batch is reserved The caller may use
    // this to track completion as each instance in the batch completes.
    pub n_busy: usize,
}

impl MaterialBatch {
    // Use MaterialBuilder instead. It handles creating batches.
    fn new(
        gpu: &Gpu,
        mut passes_generator: impl FnMut() -> (Vec<MaterialPass>, Option<wgpu::Buffer>),
        n_instances: usize,
        enable_timing: bool,
    ) -> Self {
        let mut instances = Vec::<MaterialInstance>::new();
        instances.reserve(n_instances);
        for _ in 0..n_instances {
            let (passes, readback) = passes_generator();
            instances.push(MaterialInstance::new(passes, readback));
        }
        let (ready_tx, ready_rx) = std::sync::mpsc::channel();
        let mut timing = None;
        if enable_timing {
            std::hint::cold_path();
            timing = Some(MaterialTiming::new(gpu));
        }
        Self {
            instances: instances,
            ready_tx: ready_tx,
            ready_rx: ready_rx,
            timing: timing,
            busy: false,
            n_busy: 0,
        }
    }

    pub fn enable_profiling(mut self, gpu: &Gpu) -> Self {
        self.timing = Some(MaterialTiming::new(gpu));
        self
    }
}

pub struct MaterialTiming {
    ts_buf: wgpu::Buffer,
    ts_readback: wgpu::Buffer,
    ready_tx: Sender<Result<(), wgpu::BufferAsyncError>>,
    ready_rx: Receiver<Result<(), wgpu::BufferAsyncError>>,
    busy: bool,
}

impl MaterialTiming {
    pub fn new(gpu: &Gpu) -> Self {
        let ts_buf = gpu.buffer_by_sz(16).query_resolve().copy_src().build();
        let ts_readback = gpu.buffer_by_sz(16).copy_dst().map_read().build();
        let (ready_tx, ready_rx) = std::sync::mpsc::channel();
        Self {
            ts_buf: ts_buf,
            ts_readback: ts_readback,
            ready_tx: ready_tx,
            ready_rx: ready_rx,
            busy: false,
        }
    }
}

pub struct Material {
    pub batches: Vec<MaterialBatch>,
    pub batch_size: usize,
}

impl Material {
    pub fn get_batch(&mut self) -> Option<&mut MaterialBatch> {
        let batch = self.batches.iter_mut().find(|batch| {
            if batch.busy {
                return false;
            }
            if let Some(timing) = &batch.timing {
                std::hint::cold_path();
                if timing.busy {
                    return false;
                }
            }
            true
        })?;
        // Found a batch.
        batch.busy = true;
        if let Some(timing) = &mut batch.timing {
            std::hint::cold_path();
            timing.busy = true;
        }
        batch.n_busy = self.batch_size;
        Some(batch)
    }
}

pub struct MaterialBuilder {
    batch_size: usize, // instances per batch
    n_batches: usize,
    enable_timing: bool,
}

impl MaterialBuilder {
    pub fn new() -> Self {
        Self {
            batch_size: 1,
            n_batches: 1,
            enable_timing: false,
        }
    }

    pub fn batch_size(mut self, batch_size: usize) -> Self {
        self.batch_size = batch_size;
        self
    }

    pub fn n_batches(mut self, n_batches: usize) -> Self {
        self.n_batches = n_batches;
        self
    }

    pub fn enable_timing(mut self) -> Self {
        self.enable_timing = true;
        self
    }

    pub fn build(
        self,
        gpu: &Gpu,
        mut passes_generator: impl FnMut() -> (Vec<MaterialPass>, Option<wgpu::Buffer>),
    ) -> Material {
        let mut batches = Vec::<MaterialBatch>::new();
        batches.reserve(self.n_batches);
        for _ in 0..self.n_batches {
            batches.push(MaterialBatch::new(
                gpu,
                &mut passes_generator,
                self.batch_size,
                self.enable_timing,
            ));
        }
        Material {
            batches: batches,
            batch_size: self.batch_size,
        }
    }
}

pub fn get_pipeline(gpu: &Gpu, shader: &Shader) -> wgpu::ComputePipeline {
    let pipeline_layout = gpu
        .device
        .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: None,
            bind_group_layouts: &[Some(&shader.bind_group_layout)],
            immediate_size: 0,
        });
    // TODO we should use a builder. Don't want to skip
    // zero-initializing memory in general case.
    let compilation_options = wgpu::PipelineCompilationOptions {
        zero_initialize_workgroup_memory: false,
        ..Default::default()
    };
    gpu.device
        .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: None,
            layout: Some(&pipeline_layout),
            module: &shader.shader_module,
            entry_point: Some(&shader.entry_point_name),
            compilation_options: compilation_options,
            cache: None,
        })
}

pub fn submit_dft_batch(
    gpu: &Gpu,
    shader: &Shader,
    material: &mut Material,
    pipeline: &wgpu::ComputePipeline,
    data_batch: &[Vec<Complex<f32>>],
) -> anyhow::Result<usize> {
    assert_eq!(data_batch.len(), material.batch_size);

    let Some(batch) = material.get_batch() else {
        return Ok(0);
    };

    // Submit data to the GPU for each input in the batch.
    for (data, instance) in data_batch.iter().zip(batch.instances.iter()) {
        gpu.queue.write_buffer(
            &instance.passes[0].buffers[0],
            0,
            bytemuck::cast_slice(data),
        );
    }

    let mut encoder = gpu
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
    let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
        label: None,
        timestamp_writes: None,
    });
    compute_pass.set_pipeline(pipeline);
    let shader_local_size = shader
        .entry_point_workgroup_size
        .iter()
        .fold(1, |acc, &x| acc * x) as usize;
    for idx in 0..data_batch.len() {
        compute_pass.set_bind_group(0, &batch.instances[idx].passes[0].bind_group, &[]);
        compute_pass.dispatch_workgroups(
            data_batch[idx].len().div_ceil(shader_local_size) as u32,
            1,
            1,
        );
    }
    drop(compute_pass);

    for idx in 0..data_batch.len() {
        let slot = &batch.instances[idx];
        encoder.copy_buffer_to_buffer(
            &slot.passes[0].buffers[1],
            0,
            &slot.readback.as_ref().unwrap(),
            0,
            slot.passes[0].buffers[1].size(),
        );
        let ready_tx = batch.ready_tx.clone();
        encoder.map_buffer_on_submit(
            &slot.readback.as_ref().unwrap(),
            wgpu::MapMode::Read,
            ..,
            move |result| {
                ready_tx.send((idx, result)).ok();
            },
        );
    }

    gpu.queue.submit([encoder.finish()]);
    Ok(data_batch.len())
}

#[allow(dead_code)]
fn drain_dft(
    gpu: &Gpu,
    material: &mut Material,
) -> anyhow::Result<Vec<(usize, Vec<Complex<f32>>)>> {
    gpu.device.poll(wgpu::PollType::Poll)?;
    let mut results = Vec::new();
    for batch in &mut material.batches {
        while let Ok((instance_id, map_result)) = batch.ready_rx.try_recv() {
            let slot = &mut batch.instances[instance_id];
            let data = slot.readback.as_ref().unwrap().get_mapped_range(..)?;
            let result = bytemuck::cast_slice(&data).to_vec();
            drop(data);

            slot.readback.as_ref().unwrap().unmap();
            batch.n_busy -= 1;
            if batch.n_busy == 0 {
                batch.busy = false;
            }

            map_result?;
            results.push((instance_id, result));
        }
    }
    Ok(results)
}

#[allow(dead_code)]
fn submit_fft_v0_batch(
    gpu: &Gpu,
    shader: &Shader,
    material: &mut Material,
    pipeline: &wgpu::ComputePipeline,
    data_batch: &[Vec<Complex<f32>>],
) -> anyhow::Result<usize> {
    assert_eq!(data_batch.len(), material.batch_size);

    let Some(batch) = material.get_batch() else {
        return Ok(0);
    };

    // Submit data to the GPU for each input in the batch.
    for (data, instance) in data_batch.iter().zip(batch.instances.iter()) {
        gpu.queue.write_buffer(
            &instance.passes[0].buffers[0],
            0,
            bytemuck::cast_slice(data),
        );
    }

    let mut encoder = gpu
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
    let timestamp_writes = if let Some(timestamps) = &gpu.timestamps {
        std::hint::cold_path();
        Some(wgpu::ComputePassTimestampWrites {
            query_set: &timestamps,
            beginning_of_pass_write_index: Some(0),
            end_of_pass_write_index: Some(1),
        })
    } else {
        None
    };

    let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
        label: None,
        timestamp_writes: timestamp_writes,
    });
    compute_pass.set_pipeline(pipeline);
    let shader_local_size = shader
        .entry_point_workgroup_size
        .iter()
        .fold(1, |acc, &x| acc * x) as usize;
    for (data, instance) in data_batch.iter().zip(batch.instances.iter()) {
        for pass in &instance.passes {
            compute_pass.set_bind_group(0, &pass.bind_group, &[]);
            // Each instance computes two outputs, so we only need to dispatch len/2 instances.
            compute_pass.dispatch_workgroups(
                (data.len() / 2).div_ceil(shader_local_size) as u32,
                1,
                1,
            );
        }
    }
    drop(compute_pass);

    for (instance_id, instance) in batch.instances.iter().enumerate() {
        encoder.copy_buffer_to_buffer(
            &instance.passes.last().unwrap().buffers[1],
            0,
            &instance.readback.as_ref().unwrap(),
            0,
            instance.passes.last().unwrap().buffers[1].size(),
        );
        let ready_tx = batch.ready_tx.clone();
        encoder.map_buffer_on_submit(
            &instance.readback.as_ref().unwrap(),
            wgpu::MapMode::Read,
            ..,
            move |result| {
                ready_tx.send((instance_id, result)).ok();
            },
        );
    }

    if let Some(timing) = &batch.timing {
        std::hint::cold_path();
        encoder.resolve_query_set(gpu.timestamps.as_ref().unwrap(), 0..2, &timing.ts_buf, 0);
        encoder.copy_buffer_to_buffer(&timing.ts_buf, 0, &timing.ts_readback, 0, 16);
        let ready_tx = timing.ready_tx.clone();
        encoder.map_buffer_on_submit(
            &timing.ts_readback,
            wgpu::MapMode::Read,
            ..,
            move |result| {
                ready_tx.send(result).ok();
            },
        );
    }

    gpu.queue.submit([encoder.finish()]);
    Ok(data_batch.len())
}

#[allow(dead_code)]
fn drain_fft_v0(
    gpu: &Gpu,
    material: &mut Material,
    mut cb: impl FnMut(usize, &[Complex<f32>]) -> anyhow::Result<()>,
) -> anyhow::Result<()> {
    gpu.device.poll(wgpu::PollType::Poll)?;
    for batch in &mut material.batches {
        if !batch.busy {
            continue;
        }
        while let Ok((instance_id, map_result)) = batch.ready_rx.try_recv() {
            let instance = &batch.instances[instance_id];
            map_result?;
            let data = instance.readback.as_ref().unwrap().get_mapped_range(..)?;
            let cb_result = cb(instance_id, bytemuck::cast_slice(&data));
            drop(data);

            instance.readback.as_ref().unwrap().unmap();
            batch.n_busy -= 1;
            if batch.n_busy == 0 {
                batch.busy = false;
            }

            cb_result?;
        }
    }
    Ok(())
}

#[allow(dead_code)]
fn submit_fft_v1_batch(
    gpu: &Gpu,
    shader: &Shader,
    material: &mut Material,
    pipeline: &wgpu::ComputePipeline,
    data_len: usize,                // length of 1 input
    data_batch: &Vec<Complex<f32>>, // batch of N inputs
) -> anyhow::Result<usize> {
    assert!(data_len > 0);
    assert_eq!(data_batch.len() % data_len, 0);

    let Some(batch) = material.get_batch() else {
        return Ok(0);
    };
    // We do not use instances since the input is pre-batched.
    assert_eq!(batch.instances.len(), 1);
    let instance = &batch.instances[0];

    // Submit data
    gpu.queue.write_buffer(
        &instance.passes[0].buffers[0],
        0,
        bytemuck::cast_slice(&data_batch),
    );

    let mut encoder = gpu
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
    let timestamp_writes = if let Some(timestamps) = &gpu.timestamps {
        std::hint::cold_path();
        Some(wgpu::ComputePassTimestampWrites {
            query_set: &timestamps,
            beginning_of_pass_write_index: Some(0),
            end_of_pass_write_index: Some(1),
        })
    } else {
        None
    };

    let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
        label: None,
        timestamp_writes: timestamp_writes,
    });
    compute_pass.set_pipeline(pipeline);
    let local_x = shader.entry_point_workgroup_size[0];
    for pass in &instance.passes {
        compute_pass.set_bind_group(0, &pass.bind_group, &[]);
        // Each instance computes two outputs, so we only need to dispatch len/2 instances.
        compute_pass.dispatch_workgroups(
            (data_len / 2).div_ceil(local_x as usize) as u32,
            (data_batch.len() / data_len) as u32,
            1,
        );
    }
    drop(compute_pass);

    encoder.copy_buffer_to_buffer(
        &instance.passes.last().unwrap().buffers[1],
        0,
        &instance.readback.as_ref().unwrap(),
        0,
        instance.passes.last().unwrap().buffers[1].size(),
    );
    let ready_tx = batch.ready_tx.clone();
    encoder.map_buffer_on_submit(
        &instance.readback.as_ref().unwrap(),
        wgpu::MapMode::Read,
        ..,
        move |result| {
            ready_tx.send((0, result)).ok();
        },
    );

    if let Some(timing) = &batch.timing {
        std::hint::cold_path();
        encoder.resolve_query_set(gpu.timestamps.as_ref().unwrap(), 0..2, &timing.ts_buf, 0);
        encoder.copy_buffer_to_buffer(&timing.ts_buf, 0, &timing.ts_readback, 0, 16);
        let ready_tx = timing.ready_tx.clone();
        encoder.map_buffer_on_submit(
            &timing.ts_readback,
            wgpu::MapMode::Read,
            ..,
            move |result| {
                ready_tx.send(result).ok();
            },
        );
    }

    gpu.queue.submit([encoder.finish()]);
    Ok(data_batch.len())
}

#[allow(dead_code)]
fn submit_fft_v2_batch(
    gpu: &Gpu,
    shader: &Shader,
    material: &mut Material,
    pipeline: &wgpu::ComputePipeline,
    data_len: usize,                // length of 1 input
    data_batch: &Vec<Complex<f32>>, // batch of N inputs
) -> anyhow::Result<usize> {
    assert!(data_len > 0);
    assert_eq!(data_batch.len() % data_len, 0);

    let Some(batch) = material.get_batch() else {
        return Ok(0);
    };
    // We do not use instances since the input is pre-batched.
    assert_eq!(batch.instances.len(), 1);
    let instance = &batch.instances[0];

    // Submit data
    gpu.queue.write_buffer(
        &instance.passes[0].buffers[0],
        0,
        bytemuck::cast_slice(&data_batch),
    );

    let mut encoder = gpu
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
    let timestamp_writes = if let Some(timestamps) = &gpu.timestamps {
        std::hint::cold_path();
        Some(wgpu::ComputePassTimestampWrites {
            query_set: &timestamps,
            beginning_of_pass_write_index: Some(0),
            end_of_pass_write_index: Some(1),
        })
    } else {
        None
    };

    let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
        label: None,
        timestamp_writes: timestamp_writes,
    });
    compute_pass.set_pipeline(pipeline);
    let local_x = shader.entry_point_workgroup_size[0];
    for pass in &instance.passes {
        compute_pass.set_bind_group(0, &pass.bind_group, &[]);
        // Each instance computes two outputs, so we only need to dispatch len/2 instances.
        compute_pass.dispatch_workgroups(1, 1, (data_batch.len() / data_len) as u32);
    }
    drop(compute_pass);

    encoder.copy_buffer_to_buffer(
        &instance.passes.last().unwrap().buffers[1],
        0,
        &instance.readback.as_ref().unwrap(),
        0,
        instance.passes.last().unwrap().buffers[1].size(),
    );
    let ready_tx = batch.ready_tx.clone();
    encoder.map_buffer_on_submit(
        &instance.readback.as_ref().unwrap(),
        wgpu::MapMode::Read,
        ..,
        move |result| {
            ready_tx.send((0, result)).ok();
        },
    );

    if let Some(timing) = &batch.timing {
        std::hint::cold_path();
        encoder.resolve_query_set(gpu.timestamps.as_ref().unwrap(), 0..2, &timing.ts_buf, 0);
        encoder.copy_buffer_to_buffer(&timing.ts_buf, 0, &timing.ts_readback, 0, 16);
        let ready_tx = timing.ready_tx.clone();
        encoder.map_buffer_on_submit(
            &timing.ts_readback,
            wgpu::MapMode::Read,
            ..,
            move |result| {
                ready_tx.send(result).ok();
            },
        );
    }

    gpu.queue.submit([encoder.finish()]);
    Ok(data_batch.len())
}

#[allow(dead_code)]
fn drain_fft_v0_timing(
    gpu: &Gpu,
    material: &mut Material,
    mut cb: impl FnMut(f64),
) -> anyhow::Result<()> {
    for batch in &mut material.batches {
        if let Some(timing) = &mut batch.timing {
            std::hint::cold_path();

            if !timing.busy {
                continue;
            }

            if let Ok(map_result) = timing.ready_rx.try_recv() {
                map_result?;

                let data = timing.ts_readback.get_mapped_range(..)?;
                let ticks: &[u64] = bytemuck::cast_slice(&data);
                let elapsed_ns = (ticks[1].wrapping_sub(ticks[0]) as f64)
                    * (gpu.queue.get_timestamp_period() as f64);

                drop(data);
                timing.ts_readback.unmap();
                cb(elapsed_ns);
                timing.busy = false;
            }
        }
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use crate::util::*;
    use std::time::Duration;
    use std::time::Instant;

    use num::Complex;
    use num::traits::Float;
    use rand::RngExt;
    use rand::SeedableRng;
    use rand::rngs::StdRng;

    // Cast one float type to another, truncating if needed.
    fn t0_to_t1<T0: Float, T1: Float>(val: T0) -> T1 {
        return num::cast(val).unwrap();
    }

    // Casts one Complex<T> to another.
    fn complex_to_t<T0: Float, T1: Float>(data: &[Complex<T0>]) -> Vec<Complex<T1>> {
        data.iter()
            .map(|c| Complex::new(t0_to_t1(c.re), t0_to_t1(c.im)))
            .collect()
    }

    fn evaluate_results(
        result_ref_64: &Vec<Complex<f64>>,
        result_cur_64: &Vec<Complex<f64>>,
        duration_dft: &std::time::Duration,
        duration_cur: &std::time::Duration,
        algo_name: &str,
    ) {
        assert_eq!(result_ref_64.len(), result_cur_64.len());

        println!("  Algorithm:  {}", algo_name);
        println!("    Duration:  {:?}", duration_cur);
        println!(
            "    Speedup:   {:?}",
            duration_dft.as_nanos() as f32 / duration_cur.as_nanos() as f32
        );

        let mut max_err: f64 = 0.0;
        let mut sum_err: f64 = 0.0;
        // TODO: median, 90p, 99p

        let result_64: Vec<Complex<f64>> = complex_to_t(&result_cur_64);

        for i in 0..result_ref_64.len() {
            max_err = max_err.max((result_ref_64[i] - result_64[i]).norm());
            sum_err += (result_ref_64[i] - result_64[i]).norm();
        }
        println!("    Max err:   {}", max_err);
        println!("    Avg err:   {}", sum_err / (result_ref_64.len() as f64));
    }

    fn evaluate_naive_dft(data: &[Complex<f32>], gpu: &super::Gpu) -> std::time::Duration {
        let shader = super::Shader::new(&gpu, include_wgsl!("naive_dft.wgsl")).unwrap();
        let mut result_dft = Vec::from(data);

        const BATCH_SIZE: usize = 128;
        let inputs = vec![Vec::from(data); BATCH_SIZE];

        let passes_gen = || -> (Vec<super::MaterialPass>, Option<wgpu::Buffer>) {
            let buffer_a = gpu.buffer(&data).storage().copy_dst().build();
            let buffer_b = gpu.buffer(&data).storage().copy_src().build();
            let buffers = vec![buffer_a, buffer_b];

            let passes = vec![
                super::MaterialPassBuilder::new()
                    .buffers(buffers)
                    .bind_group(&gpu, &shader)
                    .build()
                    .unwrap(),
            ];
            let readback = gpu.buffer(&data).copy_dst().map_read().build();
            (passes, Some(readback))
        };

        let mut material = super::MaterialBuilder::new()
            .batch_size(BATCH_SIZE)
            .n_batches(2)
            .build(gpu, passes_gen);
        let pipeline = super::get_pipeline(&gpu, &shader);

        let warmup = Instant::now() + Duration::from_secs(3);
        let end = warmup + Duration::from_secs(5);
        let mut n_iter: u32 = 0;
        let mut now = Instant::now();
        while now < end {
            // Fill the pipeline.
            if super::submit_dft_batch(&gpu, &shader, &mut material, &pipeline, &inputs).unwrap()
                > 0
            {
                continue;
            }
            // Drain if pipeline is full.
            let results = super::drain_dft(&gpu, &mut material).unwrap();
            for (_, result) in &results {
                result_dft = result.clone();
            }
            if now > warmup {
                n_iter += results.len() as u32;
            }
            now = Instant::now();
        }
        let duration_dft = (now - warmup) / n_iter;
        println!("Duration: {:?}", duration_dft);

        let mut result_ref: Vec<rustfft::num_complex::Complex<f64>> = complex_to_t(&data);
        let mut ref_planner = rustfft::FftPlanner::<f64>::new();
        let ref_fft = ref_planner.plan_fft_forward(result_ref.len());
        ref_fft.process(&mut result_ref);
        let result_ref_64: Vec<Complex<f64>> = complex_to_t(&result_ref);

        let result_dft_64 = complex_to_t::<f32, f64>(&result_dft);
        evaluate_results(
            &result_ref_64,
            &result_dft_64,
            &duration_dft,
            &duration_dft,
            "Naive DFT",
        );

        duration_dft
    }

    fn evaluate_fft_v0(
        data: &[Complex<f32>],
        gpu: &super::Gpu,
        duration_dft: &std::time::Duration,
    ) {
        let shader = super::Shader::new(&gpu, include_wgsl!("fft_v0_basic.wgsl")).unwrap();
        let mut result = Vec::from(data);
        const BATCH_SIZE: usize = 8;
        let inputs = vec![Vec::from(data); BATCH_SIZE];

        let twiddles = precompute_twiddles::<f32>(data.len());
        let twiddles_buffer = gpu.buffer(&twiddles).storage().copy_dst().build();
        gpu.queue.write_buffer(
            &twiddles_buffer,
            0,
            bytemuck::cast_slice(twiddles.as_slice()),
        );

        let passes_gen = || -> (Vec<super::MaterialPass>, Option<wgpu::Buffer>) {
            let buffer_a = gpu.buffer(&data).storage().copy_src().copy_dst().build();
            let buffer_b = gpu.buffer(&data).storage().copy_src().copy_dst().build();

            let mut stride = data.len() / 2;
            let mut big_n = 2;
            let num_passes = log_k_of::<2>(data.len());
            let mut passes = Vec::<super::MaterialPass>::new();
            passes.reserve(num_passes);
            for i in 0..num_passes {
                // TODO each pass needs its own uniforms, but each *instance* should share them
                #[repr(C)]
                #[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
                struct Uniforms {
                    stride: u32,
                    big: u32,
                }
                let uniforms = Uniforms {
                    stride: stride as u32,
                    big: big_n as u32,
                };
                let uniforms_buffer = gpu.buffer(&uniforms).uniform().copy_dst().build();
                gpu.queue
                    .write_buffer(&uniforms_buffer, 0, bytemuck::bytes_of(&uniforms));
                stride /= 2;
                big_n *= 2;

                let buffers = if i % 2 == 0 {
                    vec![
                        buffer_a.clone(),
                        buffer_b.clone(),
                        twiddles_buffer.clone(),
                        uniforms_buffer.clone(),
                    ]
                } else {
                    vec![
                        buffer_b.clone(),
                        buffer_a.clone(),
                        twiddles_buffer.clone(),
                        uniforms_buffer.clone(),
                    ]
                };
                passes.push(
                    super::MaterialPassBuilder::new()
                        .buffers(buffers)
                        .bind_group(&gpu, &shader)
                        .build()
                        .unwrap(),
                );
            }
            let readback = gpu.buffer(&data).copy_dst().map_read().build();
            (passes, Some(readback))
        };

        // Allocate twice as many buffers as each batch uses.
        // This lets us send data over while the GPU is still working.
        let mut material = super::MaterialBuilder::new()
            .batch_size(BATCH_SIZE)
            .n_batches(2)
            .enable_timing()
            .build(gpu, passes_gen);
        let pipeline = super::get_pipeline(&gpu, &shader);

        let warmup = Instant::now() + Duration::from_secs(3);
        let end = warmup + Duration::from_secs(5);
        let mut n_iter: u32 = 0;
        let mut sum_gpu_ns: f64 = 0.0;
        let mut now = Instant::now();
        while now < end {
            // Fill the pipeline.
            if super::submit_fft_v0_batch(&gpu, &shader, &mut material, &pipeline, &inputs).unwrap()
                > 0
            {
                continue;
            }
            // Drain if pipeline is full.
            let mut n_results = 0;
            super::drain_fft_v0(&gpu, &mut material, |_, data| {
                result.copy_from_slice(data);
                n_results += 1;
                Ok(())
            })
            .unwrap();
            let mut cur_sum_gpu_ns: f64 = 0.0;
            super::drain_fft_v0_timing(gpu, &mut material, |elapsed_ns| {
                cur_sum_gpu_ns += elapsed_ns;
            })
            .unwrap();
            if now > warmup {
                n_iter += n_results;
                sum_gpu_ns += cur_sum_gpu_ns;
            }
            now = Instant::now();
        }
        let duration = (now - warmup) / n_iter;
        println!("Duration (wall): {:?}", duration);

        let gpu_duration = Duration::from_secs_f64(sum_gpu_ns / (n_iter as f64 * 1e9));
        println!("Duration (GPU):  {:?}", gpu_duration);

        let mut result_ref: Vec<rustfft::num_complex::Complex<f64>> = complex_to_t(&data);
        let mut ref_planner = rustfft::FftPlanner::<f64>::new();
        let ref_fft = ref_planner.plan_fft_forward(result_ref.len());
        ref_fft.process(&mut result_ref);
        let result_ref_64: Vec<Complex<f64>> = complex_to_t(&result_ref);

        let result_64 = complex_to_t::<f32, f64>(&result);
        evaluate_results(
            &result_ref_64,
            &result_64,
            &duration_dft,
            &duration,
            "FFT v0",
        );
    }

    fn evaluate_fft_v1(
        data: &[Complex<f32>],
        gpu: &super::Gpu,
        duration_dft: &std::time::Duration,
    ) {
        let shader = super::Shader::new(&gpu, include_wgsl!("fft_v1_batched.wgsl")).unwrap();
        let mut result = Vec::from(data);
        const BATCH_SIZE: usize = 64;

        let mut input_batch = Vec::<Complex<f32>>::new();
        input_batch.reserve(data.len() * BATCH_SIZE);
        for _ in 0..BATCH_SIZE {
            input_batch.extend_from_slice(data);
        }

        let twiddles = precompute_twiddles::<f32>(data.len());
        let twiddles_buffer = gpu.buffer(&twiddles).storage().copy_dst().build();
        gpu.queue.write_buffer(
            &twiddles_buffer,
            0,
            bytemuck::cast_slice(twiddles.as_slice()),
        );

        let passes_gen = || -> (Vec<super::MaterialPass>, Option<wgpu::Buffer>) {
            let buffer_a = gpu
                .buffer(&input_batch)
                .storage()
                .copy_src()
                .copy_dst()
                .build();
            let buffer_b = gpu
                .buffer(&input_batch)
                .storage()
                .copy_src()
                .copy_dst()
                .build();

            let mut stride = data.len() / 2;
            let mut big_n = 2;
            let num_passes = log_k_of::<2>(data.len());
            let mut passes = Vec::<super::MaterialPass>::new();
            passes.reserve(num_passes);
            for i in 0..num_passes {
                // TODO each pass needs its own uniforms, but each *instance* should share them
                #[repr(C)]
                #[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
                struct Uniforms {
                    stride: u32,
                    big_n: u32,
                    input_len: u32,
                }
                let uniforms = Uniforms {
                    stride: stride as u32,
                    big_n: big_n as u32,
                    input_len: data.len() as u32,
                };
                let uniforms_buffer = gpu.buffer(&uniforms).uniform().copy_dst().build();
                gpu.queue
                    .write_buffer(&uniforms_buffer, 0, bytemuck::bytes_of(&uniforms));
                stride /= 2;
                big_n *= 2;

                let buffers = if i % 2 == 0 {
                    vec![
                        buffer_a.clone(),
                        buffer_b.clone(),
                        twiddles_buffer.clone(),
                        uniforms_buffer.clone(),
                    ]
                } else {
                    vec![
                        buffer_b.clone(),
                        buffer_a.clone(),
                        twiddles_buffer.clone(),
                        uniforms_buffer.clone(),
                    ]
                };
                passes.push(
                    super::MaterialPassBuilder::new()
                        .buffers(buffers)
                        .bind_group(&gpu, &shader)
                        .build()
                        .unwrap(),
                );
            }
            let readback = gpu.buffer(&input_batch).copy_dst().map_read().build();
            (passes, Some(readback))
        };

        // Allocate twice as many buffers as each batch uses.
        // This lets us send data over while the GPU is still working.
        let mut material = super::MaterialBuilder::new()
            .batch_size(1) // use 1 instance per batch since we're pre-batching our inputs
            .n_batches(2)
            .enable_timing()
            .build(gpu, passes_gen);
        let pipeline = super::get_pipeline(&gpu, &shader);

        let warmup = Instant::now() + Duration::from_secs(3);
        let end = warmup + Duration::from_secs(5);
        let mut n_iter: u32 = 0;
        let mut sum_gpu_ns: f64 = 0.0;
        let mut now = Instant::now();
        while now < end {
            // Fill the pipeline.
            if super::submit_fft_v1_batch(
                &gpu,
                &shader,
                &mut material,
                &pipeline,
                data.len(),
                &input_batch,
            )
            .unwrap()
                > 0
            {
                // Keep submitting as long as we're able to.
                continue;
            }

            // Drain if pipeline is full.
            let mut n_results = 0;
            let data_len = data.len();
            super::drain_fft_v0(&gpu, &mut material, |_, data_batch| {
                for i in 0..BATCH_SIZE {
                    result.copy_from_slice(&data_batch[i * data_len..(i + 1) * data_len]);
                }
                n_results += BATCH_SIZE as u32;
                Ok(())
            })
            .unwrap();
            let mut cur_sum_gpu_ns: f64 = 0.0;
            super::drain_fft_v0_timing(gpu, &mut material, |elapsed_ns| {
                cur_sum_gpu_ns += elapsed_ns;
            })
            .unwrap();
            if now > warmup {
                n_iter += n_results;
                sum_gpu_ns += cur_sum_gpu_ns;
            }
            now = Instant::now();
        }
        let duration = (now - warmup) / n_iter;
        println!("Duration (wall): {:?}", duration);

        let gpu_duration = Duration::from_secs_f64(sum_gpu_ns / (n_iter as f64 * 1e9));
        println!("Duration (GPU):  {:?}", gpu_duration);

        let mut result_ref: Vec<rustfft::num_complex::Complex<f64>> = complex_to_t(&data);
        let mut ref_planner = rustfft::FftPlanner::<f64>::new();
        let ref_fft = ref_planner.plan_fft_forward(result_ref.len());
        ref_fft.process(&mut result_ref);
        let result_ref_64: Vec<Complex<f64>> = complex_to_t(&result_ref);

        let result_64 = complex_to_t::<f32, f64>(&result);
        evaluate_results(
            &result_ref_64,
            &result_64,
            &duration_dft,
            &duration,
            "FFT v1",
        );
    }

    fn evaluate_fft_v2(
        data: &[Complex<f32>],
        gpu: &super::Gpu,
        duration_dft: &std::time::Duration,
        shader: super::Shader,
    ) {
        let mut result = Vec::from(data);
        const BATCH_SIZE: usize = 1024;

        let mut input_batch = Vec::<Complex<f32>>::new();
        input_batch.reserve(data.len() * BATCH_SIZE);
        for _ in 0..BATCH_SIZE {
            input_batch.extend_from_slice(data);
        }

        let twiddles = precompute_twiddles::<f32>(data.len());
        let twiddles_buffer = gpu.buffer(&twiddles).storage().copy_dst().build();
        gpu.queue.write_buffer(
            &twiddles_buffer,
            0,
            bytemuck::cast_slice(twiddles.as_slice()),
        );

        let passes_gen = || -> (Vec<super::MaterialPass>, Option<wgpu::Buffer>) {
            let buffer_a = gpu
                .buffer(&input_batch)
                .storage()
                .copy_src()
                .copy_dst()
                .build();
            let buffer_b = gpu
                .buffer(&input_batch)
                .storage()
                .copy_src()
                .copy_dst()
                .build();

            let mut passes = Vec::<super::MaterialPass>::new();

            // TODO each pass needs its own uniforms, but each *instance* should share them
            #[repr(C)]
            #[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
            struct Uniforms {
                input_len: u32,
                n_iter: u32,
            }
            let uniforms = Uniforms {
                input_len: data.len() as u32,
                n_iter: log_k_of::<4>(data.len()) as u32,
            };
            let uniforms_buffer = gpu.buffer(&uniforms).uniform().copy_dst().build();
            gpu.queue
                .write_buffer(&uniforms_buffer, 0, bytemuck::bytes_of(&uniforms));

            passes.push(
                super::MaterialPassBuilder::new()
                    .buffers(vec![
                        buffer_a,
                        buffer_b,
                        twiddles_buffer.clone(),
                        uniforms_buffer,
                    ])
                    .bind_group(&gpu, &shader)
                    .build()
                    .unwrap(),
            );
            let readback = gpu.buffer(&input_batch).copy_dst().map_read().build();
            (passes, Some(readback))
        };

        // Allocate twice as many buffers as each batch uses.
        // This lets us send data over while the GPU is still working.
        let mut material = super::MaterialBuilder::new()
            .batch_size(1) // use 1 instance per batch since we're pre-batching our inputs
            .n_batches(2)
            .enable_timing()
            .build(gpu, passes_gen);
        let pipeline = super::get_pipeline(&gpu, &shader);

        let warmup = Instant::now() + Duration::from_secs(3);
        let end = warmup + Duration::from_secs(5);
        let mut n_iter: u32 = 0;
        let mut sum_gpu_ns: f64 = 0.0;
        let mut now = Instant::now();
        while now < end {
            // Fill the pipeline.
            if super::submit_fft_v2_batch(
                &gpu,
                &shader,
                &mut material,
                &pipeline,
                data.len(),
                &input_batch,
            )
            .unwrap()
                > 0
            {
                // Keep submitting as long as we're able to.
                continue;
            }

            // Drain if pipeline is full.
            let mut n_results = 0;
            let data_len = data.len();
            super::drain_fft_v0(&gpu, &mut material, |_, data_batch| {
                for i in 0..BATCH_SIZE {
                    result.copy_from_slice(&data_batch[i * data_len..(i + 1) * data_len]);
                }
                n_results += BATCH_SIZE as u32;
                Ok(())
            })
            .unwrap();
            let mut cur_sum_gpu_ns: f64 = 0.0;
            super::drain_fft_v0_timing(gpu, &mut material, |elapsed_ns| {
                cur_sum_gpu_ns += elapsed_ns;
            })
            .unwrap();
            if now > warmup {
                n_iter += n_results;
                sum_gpu_ns += cur_sum_gpu_ns;
            }
            now = Instant::now();
        }
        let duration = (now - warmup) / n_iter;
        println!("Duration (wall): {:?}", duration);

        let gpu_duration = Duration::from_secs_f64(sum_gpu_ns / (n_iter as f64 * 1e9));
        println!("Duration (GPU):  {:?}", gpu_duration);

        let mut result_ref: Vec<rustfft::num_complex::Complex<f64>> = complex_to_t(&data);
        let mut ref_planner = rustfft::FftPlanner::<f64>::new();
        let ref_fft = ref_planner.plan_fft_forward(result_ref.len());
        ref_fft.process(&mut result_ref);
        let result_ref_64: Vec<Complex<f64>> = complex_to_t(&result_ref);

        let result_64 = complex_to_t::<f32, f64>(&result);
        evaluate_results(
            &result_ref_64,
            &result_64,
            &duration_dft,
            &duration,
            "FFT v2",
        );
    }

    fn evaluate_fft_v4(
        data: &[Complex<f32>],
        gpu: &super::Gpu,
        duration_dft: &std::time::Duration,
        shader: super::Shader,
    ) {
        let mut result = Vec::from(data);
        const BATCH_SIZE: usize = 2048;

        let mut input_batch = Vec::<Complex<f32>>::new();
        input_batch.reserve(data.len() * BATCH_SIZE);
        for _ in 0..BATCH_SIZE {
            input_batch.extend_from_slice(data);
        }

        let twiddles = precompute_twiddles::<f32>(data.len());
        let twiddles_buffer = gpu.buffer(&twiddles).storage().copy_dst().build();
        gpu.queue.write_buffer(
            &twiddles_buffer,
            0,
            bytemuck::cast_slice(twiddles.as_slice()),
        );

        let passes_gen = || -> (Vec<super::MaterialPass>, Option<wgpu::Buffer>) {
            let buffer_a = gpu
                .buffer(&input_batch)
                .storage()
                .copy_src()
                .copy_dst()
                .build();
            let buffer_b = gpu
                .buffer(&input_batch)
                .storage()
                .copy_src()
                .copy_dst()
                .build();

            let mut passes = Vec::<super::MaterialPass>::new();

            // TODO each pass needs its own uniforms, but each *instance* should share them
            #[repr(C)]
            #[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
            struct Uniforms {
                input_len: u32,
                n_iter: u32,
            }
            let uniforms = Uniforms {
                input_len: data.len() as u32,
                n_iter: log_k_of::<8>(data.len()) as u32,
            };
            let uniforms_buffer = gpu.buffer(&uniforms).uniform().copy_dst().build();
            gpu.queue
                .write_buffer(&uniforms_buffer, 0, bytemuck::bytes_of(&uniforms));

            passes.push(
                super::MaterialPassBuilder::new()
                    .buffers(vec![
                        buffer_a,
                        buffer_b,
                        twiddles_buffer.clone(),
                        uniforms_buffer,
                    ])
                    .bind_group(&gpu, &shader)
                    .build()
                    .unwrap(),
            );
            let readback = gpu.buffer(&input_batch).copy_dst().map_read().build();
            (passes, Some(readback))
        };

        // Allocate twice as many buffers as each batch uses.
        // This lets us send data over while the GPU is still working.
        let mut material = super::MaterialBuilder::new()
            .batch_size(1) // use 1 instance per batch since we're pre-batching our inputs
            .n_batches(4)
            .enable_timing()
            .build(gpu, passes_gen);
        let pipeline = super::get_pipeline(&gpu, &shader);

        let warmup = Instant::now() + Duration::from_secs(3);
        let end = warmup + Duration::from_secs(5);
        let mut n_iter: u32 = 0;
        let mut sum_gpu_ns: f64 = 0.0;
        let mut now = Instant::now();
        while now < end {
            // Fill the pipeline.
            if super::submit_fft_v2_batch(
                &gpu,
                &shader,
                &mut material,
                &pipeline,
                data.len(),
                &input_batch,
            )
            .unwrap()
                > 0
            {
                // Keep submitting as long as we're able to.
                continue;
            }

            // Drain if pipeline is full.
            let mut n_results = 0;
            let data_len = data.len();
            super::drain_fft_v0(&gpu, &mut material, |_, data_batch| {
                for i in 0..BATCH_SIZE {
                    result.copy_from_slice(&data_batch[i * data_len..(i + 1) * data_len]);
                }
                n_results += BATCH_SIZE as u32;
                Ok(())
            })
            .unwrap();
            let mut cur_sum_gpu_ns: f64 = 0.0;
            super::drain_fft_v0_timing(gpu, &mut material, |elapsed_ns| {
                cur_sum_gpu_ns += elapsed_ns;
            })
            .unwrap();
            if now > warmup {
                n_iter += n_results;
                sum_gpu_ns += cur_sum_gpu_ns;
            }
            now = Instant::now();
        }
        let duration = (now - warmup) / n_iter;
        println!("Duration (wall): {:?}", duration);

        let gpu_duration = Duration::from_secs_f64(sum_gpu_ns / (n_iter as f64 * 1e9));
        println!("Duration (GPU):  {:?}", gpu_duration);

        let mut result_ref: Vec<rustfft::num_complex::Complex<f64>> = complex_to_t(&data);
        let mut ref_planner = rustfft::FftPlanner::<f64>::new();
        let ref_fft = ref_planner.plan_fft_forward(result_ref.len());
        ref_fft.process(&mut result_ref);
        let result_ref_64: Vec<Complex<f64>> = complex_to_t(&result_ref);

        let result_64 = complex_to_t::<f32, f64>(&result);
        evaluate_results(
            &result_ref_64,
            &result_64,
            &duration_dft,
            &duration,
            "FFT v4",
        );
    }

    #[test]
    fn run_it() {
        let mut data_f64: Vec<Complex<f64>> = Vec::new();
        data_f64.resize((2 as usize).pow(12), Complex::ZERO);
        // Randomize data.
        let mut rng = StdRng::seed_from_u64(42);
        for i in 0..data_f64.len() {
            data_f64[i] = Complex::new(rng.random(), rng.random());
        }
        let data_f32: Vec<Complex<f32>> = complex_to_t::<f64, f32>(&data_f64);

        let gpu = super::GpuBuilder::new().enable_profiling().build().unwrap();

        //let duration_dft = evaluate_naive_dft(&data_f32, &gpu);
        let duration_dft = Duration::from_secs(0);
        //evaluate_fft_v0(&data_f32, &gpu, &duration_dft);
        //evaluate_fft_v1(&data_f32, &gpu, &duration_dft);
        //evaluate_fft_v2(&data_f32, &gpu, &duration_dft, super::Shader::new(&gpu, include_wgsl!("fft_v2_barrier.wgsl")).unwrap());
        /*
        evaluate_fft_v2(
            &data_f32,
            &gpu,
            &duration_dft,
            super::Shader::new(&gpu, include_wgsl!("fft_v3_groupshared.wgsl")).unwrap(),
        );
        */
        evaluate_fft_v4(
            &data_f32,
            &gpu,
            &duration_dft,
            super::Shader::new(&gpu, include_wgsl!("fft_v4_radix_8.wgsl")).unwrap(),
        );
    }
}