//! The device backend (native targets).

use std::collections::HashMap;

use fieldcraft_gpu::select::{self, Policy};
use fieldcraft_lbm::backend::{Backend, Fields};
use fieldcraft_lbm::lattice::{C, Q, W};
use fieldcraft_lbm::setup::MAX_LATTICE_SPEED;
use fieldcraft_lbm::{Error, Flow, Group, Result, Rule, Target};
use fieldcraft_model::ThermalCondition;
use wgpu::util::DeviceExt;

const SHADER: &str = include_str!("lbm.wgsl");
const WG: u32 = 256;
const MAX_GROUPS_X: u32 = 32768;
const KIND_PULL: u32 = 0;
const KIND_BOUNCE: u32 = 1;
const KIND_REFLECT: u32 = 2;
const KIND_OPEN: u32 = 3;
/// The scalar lattice's directions (the first seven of D3Q19) and weights.
const QT: usize = 7;
const WT: [f64; QT] = [0.25, 0.125, 0.125, 0.125, 0.125, 0.125, 0.125];

fn gpu_error(e: impl std::fmt::Display) -> Error {
    Error::Invalid(format!("GPU: {e}"))
}

/// Steps per submission: the stability check and a cancel take effect between submissions.
const MAX_CHUNK: u64 = 64;

/// How long one wait for the device may take before the solve gives up on it.
const DEVICE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(120);

/// A bounded chunk's dispatches, shared by the normal and timestamped paths.
type Dispatch = (usize, &'static str, (u32, u32));

/// A vector of `len` copies of `value`, or [`Error::Limit`] when memory runs out (the host
/// arrays are sized by the grid: 16 M cells is over a gigabyte of populations).
fn filled<T: Clone>(len: usize, value: T, what: &str) -> Result<Vec<T>> {
    let mut v = Vec::new();
    v.try_reserve_exact(len).map_err(|_| Error::Limit(format!("not enough memory for the {what} ({len} entries)")))?;
    v.resize(len, value);
    Ok(v)
}

/// An empty vector with room for `len`, or [`Error::Limit`].
fn reserved<T>(len: usize, what: &str) -> Result<Vec<T>> {
    let mut v = Vec::new();
    v.try_reserve_exact(len).map_err(|_| Error::Limit(format!("not enough memory for the {what} ({len} entries)")))?;
    Ok(v)
}

/// Workgroup counts (x, y) covering `count` items.
fn groups(count: usize) -> Result<(u32, u32)> {
    let total = u32::try_from(count.div_ceil(WG as usize).max(1)).map_err(|_| Error::Limit("too many cells for one dispatch".into()))?;
    let gx = total.min(MAX_GROUPS_X);
    let gy = total.div_ceil(gx);
    if gy > 65535 {
        return Err(Error::Limit("too many cells for one dispatch".into()));
    }
    Ok((gx, gy))
}

fn bytes_u32(v: &[u32]) -> Vec<u8> {
    v.iter().flat_map(|x| x.to_le_bytes()).collect()
}

fn bytes_f32(v: &[f32]) -> Vec<u8> {
    v.iter().flat_map(|x| x.to_le_bytes()).collect()
}

/// Register this backend with `fieldcraft_lbm` (an application does this once at start), so
/// `solve.fluid` with `backend: "gpu"` steps on the device `FIELDCRAFT_GPU` picks.
pub fn register() {
    fieldcraft_lbm::register_backend(Box::new(|flow, _setup| GpuFlow::from_env(flow).map(|g| Box::new(g) as Box<dyn Backend>)));
}

/// How the momentum populations stream (see `lbm.wgsl`).
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum Streaming {
    /// In place: measured faster on every device tried (RTX 4090, GB10, M4 Max; #563).
    #[default]
    Auto,
    /// One population buffer, collided and streamed in place (the AA pattern).
    InPlace,
    /// Two population buffers: collide in place, then pull into the other (A–B).
    TwoBuffers,
}

/// How a [`GpuFlow`] runs; the defaults are what a solve uses.
#[derive(Clone, Copy, Debug, Default)]
pub struct Options {
    pub streaming: Streaming,
    /// Time every kernel with timestamp queries ([`GpuFlow::kernel_times`]). Each dispatch then
    /// gets its own pass and every chunk waits for its timestamps: for measurements, not solves.
    pub profile: bool,
}

/// GPU time per kernel, summed over the dispatches since the flow was created (with
/// [`Options::profile`]).
#[derive(Clone, Debug, Default, PartialEq)]
pub struct KernelTime {
    pub dispatches: u64,
    pub seconds: f64,
}

/// A prepared flow on the device.
pub struct GpuFlow {
    device: wgpu::Device,
    queue: wgpu::Queue,
    /// With [`Options::profile`]: the GPU time of each kernel so far.
    kernel_times: Option<std::cell::RefCell<std::collections::BTreeMap<&'static str, KernelTime>>>,
    /// How long creating the compute pipelines (compiling the shader) took.
    pipeline_seconds: f64,
    aa: bool,
    adapter: String,
    pipelines: HashMap<&'static str, wgpu::ComputePipeline>,
    /// AB swaps populations; AA keeps one population buffer and changes its orientation.
    /// Both modes swap the separate scalar lattice buffers.
    bind: [wgpu::BindGroup; 2],
    /// Which bind group's `f` holds the state.
    current: usize,
    n: usize,
    cells: (u32, u32),
    link_count: usize,
    links_groups: (u32, u32),
    open_count: usize,
    open_groups: (u32, u32),
    group_count: usize,
    force_groups: (u32, u32),
    stats: wgpu::Buffer,
    moments: wgpu::Buffer,
    forces: wgpu::Buffer,
    partials: wgpu::Buffer,
    partial_count: usize,
    /// The scalar populations remain A–B with heat; and T₀ they are shifted by.
    thermal: Option<[wgpu::Buffer; 2]>,
    initial_temperature: f64,
    fluid: Vec<bool>,
    fluid_cells: usize,
    groups: Vec<Group>,
    empty_obstacles: Vec<String>,
    steps_done: u64,
    /// The momentum (xyz) and heat (w) each group received in the last step, once read back
    /// (zero before the first step).
    exchange_cache: Option<Vec<[f64; 4]>>,
    /// Each group's share of the momentum from the populations' weights, which the kernel leaves
    /// out (`link_forces` sums the deviations): added to what it reads back.
    weight_forces: Vec<[f64; 3]>,
}

impl GpuFlow {
    /// `flow` on the device `FIELDCRAFT_GPU` picks.
    pub fn from_env(flow: &Flow) -> Result<GpuFlow> {
        let policy = Policy::from_env().map_err(gpu_error)?;
        GpuFlow::new(flow, policy)
    }

    /// `flow` (its grid, rules and initial state: everything the step needs) on the device
    /// `policy` picks.
    pub fn new(flow: &Flow, policy: Policy) -> Result<GpuFlow> {
        GpuFlow::with_options(flow, policy, Options::default())
    }

    /// [`GpuFlow::new`] with `options`.
    pub fn with_options(flow: &Flow, policy: Policy, options: Options) -> Result<GpuFlow> {
        let Some((adapter, tier)) = select::adapter(policy).map_err(gpu_error)? else {
            return Err(Error::Invalid(format!("{}=cpu: no GPU device; solve with `backend: \"cpu\"`", select::ENV)));
        };
        let info = adapter.get_info();
        let profile = options.profile;
        let features = if profile { wgpu::Features::TIMESTAMP_QUERY } else { wgpu::Features::empty() };
        if !adapter.features().contains(features) {
            return Err(gpu_error("profiling needs timestamp queries, which this adapter does not offer"));
        }
        let have = adapter.limits();
        let mut limits = wgpu::Limits::default().using_resolution(have.clone());
        limits.max_buffer_size = have.max_buffer_size;
        limits.max_storage_buffer_binding_size = have.max_storage_buffer_binding_size;
        limits.max_storage_buffers_per_shader_stage = have.max_storage_buffers_per_shader_stage.max(limits.max_storage_buffers_per_shader_stage);
        let (device, queue) = pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor {
            label: Some("fieldcraft lbm"),
            required_features: features,
            required_limits: limits,
            ..Default::default()
        }))
        .map_err(gpu_error)?;
        let adapter_name = format!("{}{}", info.name, if tier == select::Tier::Software { " (software)" } else { "" });

        let aa = match options.streaming {
            Streaming::Auto | Streaming::InPlace => true,
            Streaming::TwoBuffers => false,
        };
        let n = flow.grid().len();
        let elem = 4u64;
        let pop_bytes = elem * (Q * n) as u64;
        let cap = device.limits().max_storage_buffer_binding_size.min(device.limits().max_buffer_size);
        if pop_bytes > cap {
            return Err(Error::Limit(format!(
                "{n} cells need {} MB of populations per buffer; this GPU allows {} MB per buffer: coarsen the spacing",
                pop_bytes >> 20,
                cap >> 20
            )));
        }

        // The state and the flags.
        // Populations as deviations from their weights (see lbm.wgsl); solid cells stay zero.
        let mut f = filled(Q * n, 0.0f32, "populations")?;
        for q in 0..Q {
            let w = fieldcraft_lbm::lattice::W.get(q).copied().unwrap_or(0.0);
            for cell in 0..n {
                if let (Some(slot), true) = (f.get_mut(q * n + cell), flow.is_fluid(cell)) {
                    *slot = (flow.population(q, cell) - w) as f32;
                }
            }
        }
        // The scalar populations as deviations from w_q T₀ (see lbm.wgsl).
        let thermal = flow.thermal();
        let t0 = thermal.map_or(0.0, |t| t.initial());
        let mut tpop = filled(if thermal.is_some() { QT * n } else { 0 }, 0.0f32, "temperature populations")?;
        if thermal.is_some() {
            for q in 0..QT {
                let w = WT.get(q).copied().unwrap_or(0.0) * t0;
                for cell in 0..n {
                    if let (Some(slot), true) = (tpop.get_mut(q * n + cell), flow.is_fluid(cell)) {
                        *slot = (flow.thermal_population(q, cell) - w) as f32;
                    }
                }
            }
        }
        let mut fluid = filled(n, false, "cell flags")?;
        let mut fluid_u32 = filled(n, 0u32, "cell flags")?;
        for cell in 0..n {
            if let (Some(b), Some(u)) = (fluid.get_mut(cell), fluid_u32.get_mut(cell)) {
                *b = flow.is_fluid(cell);
                *u = u32::from(*b);
            }
        }

        // The links, by cell, and the ones with a force by group.
        let links = flow.links();
        let mut link_u = reserved(4 * links.len(), "links")?;
        let mut link_f = reserved(links.len(), "links")?;
        let mut by_group: Vec<(u32, u32)> = reserved(links.len(), "links by group")?;
        let group_count = flow.groups().len();
        let mut weight_forces = filled(group_count, [0.0f64; 3], "force groups")?;
        let c = |q: u32| C.get(q as usize).copied().unwrap_or([0; 3]).map(f64::from);
        let w = |q: u32| W.get(q as usize).copied().unwrap_or(0.0);
        for (l, link) in links.iter().enumerate() {
            let (kind, from, group, q2, du) = match link.rule {
                Rule::Pull { from } => (KIND_PULL, from, 0, 0, 0.0),
                Rule::Bounce { group, du } => (KIND_BOUNCE, 0, group, 0, du),
                Rule::Reflect { from, q, group } => (KIND_REFLECT, from, group, u32::from(q), 0.0),
                Rule::Open => (KIND_OPEN, 0, 0, 0, 0.0),
            };
            // What the weights add to the link's exchange (see `link_forces`).
            let weight = match link.rule {
                Rule::Bounce { .. } => c(u32::from(link.q)).map(|x| -2.0 * w(u32::from(link.q)) * x),
                Rule::Reflect { .. } => std::array::from_fn(|a| (c(q2)[a] - c(u32::from(link.q))[a]) * w(q2)),
                _ => [0.0; 3],
            };
            if matches!(link.rule, Rule::Bounce { .. } | Rule::Reflect { .. }) {
                by_group.push((group, l as u32));
                if let Some(g) = weight_forces.get_mut(group as usize) {
                    for (s, x) in g.iter_mut().zip(weight) {
                        *s += x;
                    }
                }
            }
            link_u.extend_from_slice(&[u32::from(link.q) | (kind << 8) | (q2 << 16), from, group, link.cell]);
            link_f.push(du as f32);
        }
        by_group.sort();
        let mut group_start = filled(group_count + 1, 0u32, "force groups")?;
        for (group, _) in &by_group {
            if let Some(next) = group_start.get_mut(*group as usize + 1) {
                *next += 1;
            }
        }
        for g in 0..group_count {
            let prev = group_start.get(g).copied().unwrap_or(0);
            if let Some(next) = group_start.get_mut(g + 1) {
                *next += prev;
            }
        }
        let by_group: Vec<u32> = by_group.iter().map(|(_, l)| *l).collect();
        // Each group's thermal condition (held temperatures shifted by T₀).
        let mut group_t = reserved(4 * group_count, "thermal conditions")?;
        for g in 0..group_count {
            let (kind, value) = match thermal.and_then(|t| t.groups().get(g)).copied().unwrap_or_default() {
                ThermalCondition::Adiabatic {} => (0.0, 0.0),
                ThermalCondition::Temperature { value } => (1.0, (value - t0) as f32),
                ThermalCondition::Flux { value } => (2.0, value as f32),
            };
            group_t.extend_from_slice(&[kind, value, 0.0, 0.0]);
        }

        // The open cells.
        let open = flow.open_cells();
        let mut open_u = reserved(4 * open.len(), "open cells")?;
        let mut open_f = reserved(4 * open.len(), "open cells")?;
        for oc in open {
            let (kind, target) = match oc.target {
                Target::Velocity(u) => (0u32, [u[0] as f32, u[1] as f32, u[2] as f32]),
                Target::Density(rho) => (1u32, [rho as f32, 0.0, 0.0]),
            };
            open_u.extend_from_slice(&[oc.cell, oc.axis as u32, kind, 0]);
            open_f.extend_from_slice(&[target[0], target[1], target[2], oc.inward as f32]);
        }
        let mut open_t = reserved(4 * open.len(), "open cell temperatures")?;
        for k in 0..open.len() {
            match thermal.and_then(|t| t.open().get(k)).copied().flatten() {
                Some(value) => open_t.extend_from_slice(&[1.0, (value - t0) as f32, 0.0, 0.0]),
                None => open_t.extend_from_slice(&[0.0; 4]),
            }
        }

        let (omega_plus, omega_minus) = flow.omegas();
        let acc = flow.acceleration();
        let forced = acc != [0.0; 3];
        let cells = groups(n)?;
        let links_groups = groups(links.len())?;
        let open_groups = groups(open.len())?;
        let force_groups = groups((group_count * WG as usize).max(1))?;
        let partial_count = (cells.0 as usize) * (cells.1 as usize);
        let mut params = Vec::with_capacity(64);
        let u = |v: usize| u32::try_from(v).map_err(|_| Error::Limit("too many cells for the GPU solver".into()));
        for v in [
            u(n)?,
            u(flow.grid().cells[0])?,
            u(flow.grid().cells[1])?,
            u(flow.grid().cells[2])?,
            cells.0,
            u(links.len())?,
            u(open.len())?,
            u(group_count)?,
        ] {
            params.extend_from_slice(&v.to_le_bytes());
        }
        for v in [
            omega_plus as f32,
            omega_minus as f32,
            (1.0 - 0.5 * omega_plus) as f32,
            (1.0 - 0.5 * omega_minus) as f32,
            acc[0] as f32,
            acc[1] as f32,
            acc[2] as f32,
            if forced { 1.0 } else { 0.0 },
        ] {
            params.extend_from_slice(&v.to_le_bytes());
        }
        let buoyancy = thermal.map_or([0.0; 3], |t| t.buoyancy());
        for v in [
            buoyancy[0] as f32,
            buoyancy[1] as f32,
            buoyancy[2] as f32,
            if buoyancy != [0.0; 3] { 1.0 } else { 0.0 },
            thermal.map_or(0.0, |t| t.omega()) as f32,
            t0 as f32,
            if thermal.is_some() { 1.0 } else { 0.0 },
            0.0,
        ] {
            params.extend_from_slice(&v.to_le_bytes());
        }

        // Buffers, with the device's out-of-memory errors caught.
        let scopes = [wgpu::ErrorFilter::OutOfMemory, wgpu::ErrorFilter::Validation].map(|f| device.push_error_scope(f));
        let storage = |label: &str, bytes: &[u8]| {
            device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some(label),
                contents: if bytes.is_empty() { &[0u8; 16][..] } else { bytes },
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::COPY_DST,
            })
        };
        let empty = |label: &str, size: u64| {
            device.create_buffer(&wgpu::BufferDescriptor {
                label: Some(label),
                size: size.max(16),
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            })
        };
        let params = [0u32, 1].map(|phase| {
            let mut bytes = params.clone();
            bytes.extend_from_slice(&bytes_u32(&[phase, 0, 0, 0]));
            device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("lbm params"),
                contents: &bytes,
                usage: wgpu::BufferUsages::UNIFORM,
            })
        });
        let fluid_buf = storage("lbm fluid", &bytes_u32(&fluid_u32));
        // AA reuses binding g for boundary scratch only, never another full population field.
        let pops = [storage("lbm f", &bytes_f32(&f)), empty("lbm g / AA boundary scratch", if aa { 4 * links.len() as u64 } else { pop_bytes })];
        let link_u = storage("lbm links", &bytes_u32(&link_u));
        let link_f = storage("lbm link du", &bytes_f32(&link_f));
        let open_u = storage("lbm open cells", &bytes_u32(&open_u));
        let open_f = storage("lbm open targets", &bytes_f32(&open_f));
        let moments = empty("lbm moments", 16 * n as u64);
        let stats = storage("lbm stats", &bytes_u32(&[0, 0, 0, 0]));
        let contrib = empty("lbm link forces", 16 * links.len() as u64);
        let by_group = storage("lbm links by group", &bytes_u32(&by_group));
        let group_start = storage("lbm group start", &bytes_u32(&group_start));
        let forces = empty("lbm forces", 16 * group_count as u64);
        let partials = empty("lbm partials", 4 * partial_count as u64);
        let therm_bytes = elem * (QT * n) as u64;
        let therm = if thermal.is_some() {
            [storage("lbm t", &bytes_f32(&tpop)), empty("lbm t2", therm_bytes)]
        } else {
            [empty("lbm t", 16), empty("lbm t2", 16)]
        };
        let group_t = storage("lbm thermal conditions", &bytes_f32(&group_t));
        let open_t = storage("lbm open temperatures", &bytes_f32(&open_t));

        let module = device
            .create_shader_module(wgpu::ShaderModuleDescriptor { label: Some("fieldcraft lbm"), source: wgpu::ShaderSource::Wgsl(SHADER.into()) });
        let entry = |binding: u32, ty: wgpu::BufferBindingType| wgpu::BindGroupLayoutEntry {
            binding,
            visibility: wgpu::ShaderStages::COMPUTE,
            ty: wgpu::BindingType::Buffer { ty, has_dynamic_offset: false, min_binding_size: None },
            count: None,
        };
        let read = wgpu::BufferBindingType::Storage { read_only: true };
        let write = wgpu::BufferBindingType::Storage { read_only: false };
        let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("fieldcraft lbm"),
            entries: &[
                entry(0, wgpu::BufferBindingType::Uniform),
                entry(1, read),
                entry(2, write),
                entry(3, write),
                entry(5, read),
                entry(6, read),
                entry(7, read),
                entry(8, read),
                entry(9, write),
                entry(10, write),
                entry(11, write),
                entry(12, read),
                entry(13, read),
                entry(14, write),
                entry(15, write),
                entry(16, write),
                entry(17, write),
                entry(18, read),
                entry(19, read),
            ],
        });
        let bind = [0usize, 1].map(|k| {
            device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("fieldcraft lbm"),
                layout: &layout,
                entries: &[
                    wgpu::BindGroupEntry { binding: 0, resource: params[k].as_entire_binding() },
                    wgpu::BindGroupEntry { binding: 1, resource: fluid_buf.as_entire_binding() },
                    wgpu::BindGroupEntry { binding: 2, resource: pops[if aa { 0 } else { k }].as_entire_binding() },
                    wgpu::BindGroupEntry { binding: 3, resource: pops[if aa { 1 } else { 1 - k }].as_entire_binding() },
                    wgpu::BindGroupEntry { binding: 5, resource: link_u.as_entire_binding() },
                    wgpu::BindGroupEntry { binding: 6, resource: link_f.as_entire_binding() },
                    wgpu::BindGroupEntry { binding: 7, resource: open_u.as_entire_binding() },
                    wgpu::BindGroupEntry { binding: 8, resource: open_f.as_entire_binding() },
                    wgpu::BindGroupEntry { binding: 9, resource: moments.as_entire_binding() },
                    wgpu::BindGroupEntry { binding: 10, resource: stats.as_entire_binding() },
                    wgpu::BindGroupEntry { binding: 11, resource: contrib.as_entire_binding() },
                    wgpu::BindGroupEntry { binding: 12, resource: by_group.as_entire_binding() },
                    wgpu::BindGroupEntry { binding: 13, resource: group_start.as_entire_binding() },
                    wgpu::BindGroupEntry { binding: 14, resource: forces.as_entire_binding() },
                    wgpu::BindGroupEntry { binding: 15, resource: partials.as_entire_binding() },
                    wgpu::BindGroupEntry { binding: 16, resource: therm[k].as_entire_binding() },
                    wgpu::BindGroupEntry { binding: 17, resource: therm[1 - k].as_entire_binding() },
                    wgpu::BindGroupEntry { binding: 18, resource: group_t.as_entire_binding() },
                    wgpu::BindGroupEntry { binding: 19, resource: open_t.as_entire_binding() },
                ],
            })
        });
        let pl = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("fieldcraft lbm"),
            bind_group_layouts: &[Some(&layout)],
            immediate_size: 0,
        });
        let mut pipelines = HashMap::new();
        let pipeline_start = std::time::Instant::now();
        for name in [
            "collide",
            "stream",
            "stream_links",
            "stream_links_exchange",
            "scatter_links",
            "open_boundaries",
            "link_forces",
            "group_forces",
            "cell_moments",
            "mass",
            "collide_thermal",
            "stream_thermal",
            "stream_links_thermal",
            "open_thermal",
            "thermal_check",
        ] {
            let p = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: Some(name),
                layout: Some(&pl),
                module: &module,
                entry_point: Some(name),
                compilation_options: wgpu::PipelineCompilationOptions { constants: &[("AA_MODE", if aa { 1.0 } else { 0.0 })], ..Default::default() },
                cache: None,
            });
            pipelines.insert(name, p);
        }
        for scope in scopes.into_iter().rev() {
            if let Some(e) = pollster::block_on(scope.pop()) {
                return Err(gpu_error(e));
            }
        }
        let pipeline_seconds = pipeline_start.elapsed().as_secs_f64();
        Ok(GpuFlow {
            device,
            queue,
            kernel_times: profile.then(Default::default),
            pipeline_seconds,
            aa,
            adapter: adapter_name,
            pipelines,
            bind,
            current: 0,
            n,
            cells,
            link_count: links.len(),
            links_groups,
            open_count: open.len(),
            open_groups,
            group_count,
            force_groups,
            stats,
            moments,
            forces,
            partials,
            partial_count,
            thermal: thermal.map(|_| therm),
            initial_temperature: t0,
            fluid_cells: flow.fluid_cells(),
            fluid,
            groups: flow.groups().to_vec(),
            empty_obstacles: flow.empty_obstacles().to_vec(),
            steps_done: 0,
            exchange_cache: Some(vec![[0.0; 4]; group_count]),
            weight_forces,
        })
    }

    /// The adapter's name.
    pub fn adapter_name(&self) -> &str {
        &self.adapter
    }

    /// Whether the populations stream in place (see [`Streaming`]).
    pub fn streams_in_place(&self) -> bool {
        self.aa
    }

    /// How long creating the compute pipelines took, in seconds.
    pub fn pipeline_seconds(&self) -> f64 {
        self.pipeline_seconds
    }

    /// With [`Options::profile`], the GPU time of each kernel so far, by kernel name.
    pub fn kernel_times(&self) -> Option<std::collections::BTreeMap<&'static str, KernelTime>> {
        self.kernel_times.as_ref().map(|t| t.borrow().clone())
    }

    fn pipeline(&self, name: &str) -> Result<&wgpu::ComputePipeline> {
        self.pipelines.get(name).ok_or_else(|| Error::Invalid(format!("GPU: no kernel {name}")))
    }

    fn dispatch(&self, pass: &mut wgpu::ComputePass<'_>, bind: usize, name: &str, (gx, gy): (u32, u32)) -> Result<()> {
        pass.set_bind_group(0, self.bind.get(bind).ok_or_else(|| Error::Invalid("GPU: no bind group".into()))?, &[]);
        pass.set_pipeline(self.pipeline(name)?);
        pass.dispatch_workgroups(gx, gy, 1);
        Ok(())
    }

    /// Timestamped passes are diagnostic: their extra pass boundaries and readback are not
    /// used for throughput measurements. Normally every chunk remains one compute pass.
    fn submit(&self, label: &str, commands: &[Dispatch]) -> Result<()> {
        let mut enc = self.device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some(label) });
        let Some(times) = &self.kernel_times else {
            {
                let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor { label: Some(label), timestamp_writes: None });
                for &(bind, name, groups) in commands {
                    self.dispatch(&mut pass, bind, name, groups)?;
                }
            }
            self.queue.submit([enc.finish()]);
            return Ok(());
        };
        let count = u32::try_from(commands.len() * 2).map_err(|_| gpu_error("too many timestamp queries"))?;
        if count == 0 || count > wgpu::QUERY_SET_MAX_QUERIES {
            return Err(gpu_error("timestamp query count is outside the device limit"));
        }
        let query = self.device.create_query_set(&wgpu::QuerySetDescriptor { label: Some(label), ty: wgpu::QueryType::Timestamp, count });
        let result = self.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("lbm timestamps"),
            size: u64::from(count) * 8,
            usage: wgpu::BufferUsages::QUERY_RESOLVE | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        });
        for (i, &(bind, name, groups)) in commands.iter().enumerate() {
            let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor {
                label: Some(name),
                timestamp_writes: Some(wgpu::ComputePassTimestampWrites {
                    query_set: &query,
                    beginning_of_pass_write_index: Some(2 * i as u32),
                    end_of_pass_write_index: Some(2 * i as u32 + 1),
                }),
            });
            self.dispatch(&mut pass, bind, name, groups)?;
        }
        enc.resolve_query_set(&query, 0..count, &result, 0);
        self.queue.submit([enc.finish()]);
        let bytes = self.read(&result, u64::from(count) * 8)?;
        let period = f64::from(self.queue.get_timestamp_period());
        let mut totals = times.borrow_mut();
        for ((_, name, _), pair) in commands.iter().zip(bytes.as_chunks::<16>().0) {
            let [begin, end] = pair.as_chunks::<8>().0 else { return Err(gpu_error("invalid timestamp pair")) };
            let ticks = u64::from_le_bytes(*end).checked_sub(u64::from_le_bytes(*begin)).ok_or_else(|| gpu_error("invalid GPU timestamps"))?;
            let total = totals.entry(name).or_default();
            total.dispatches += 1;
            total.seconds += ticks as f64 * period * 1e-9;
        }
        Ok(())
    }

    /// Copy `size` bytes of `src` back to the host.
    fn read(&self, src: &wgpu::Buffer, size: u64) -> Result<Vec<u8>> {
        let size = size.max(16);
        let staging = self.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("lbm readback"),
            size,
            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let mut enc = self.device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("lbm readback") });
        enc.copy_buffer_to_buffer(src, 0, &staging, 0, size);
        self.queue.submit([enc.finish()]);
        let slice = staging.slice(..);
        let (tx, rx) = std::sync::mpsc::channel();
        slice.map_async(wgpu::MapMode::Read, move |r| {
            let _ = tx.send(r);
        });
        // A bounded wait: a hung or lost device is an error, never a hang.
        match self.device.poll(wgpu::PollType::Wait { submission_index: None, timeout: Some(DEVICE_TIMEOUT) }) {
            Ok(_) => {}
            Err(e) => return Err(gpu_error(format!("the device did not answer within {} s: {e}", DEVICE_TIMEOUT.as_secs()))),
        }
        rx.try_recv().map_err(|_| gpu_error("the readback did not complete"))?.map_err(gpu_error)?;
        let bytes = slice.get_mapped_range().map_err(gpu_error)?.to_vec();
        staging.unmap();
        Ok(bytes)
    }

    fn read_f32(&self, src: &wgpu::Buffer, count: usize) -> Result<Vec<f32>> {
        let bytes = self.read(src, 4 * count as u64)?;
        Ok(bytes.as_chunks::<4>().0.iter().take(count).map(|b| f32::from_le_bytes(*b)).collect())
    }

    /// The momentum (xyz) and heat (w) each group received in the last step.
    fn exchange(&mut self) -> Result<&[[f64; 4]]> {
        if self.exchange_cache.is_none() {
            let v = self.read_f32(&self.forces, 4 * self.group_count)?;
            self.exchange_cache = Some(
                v.as_chunks::<4>()
                    .0
                    .iter()
                    .zip(&self.weight_forces)
                    .map(|(g, w)| [f64::from(g[0]) + w[0], f64::from(g[1]) + w[1], f64::from(g[2]) + w[2], f64::from(g[3])])
                    .collect(),
            );
        }
        Ok(self.exchange_cache.as_deref().unwrap_or_default())
    }

    /// The moments of the current state into the moments buffer.
    fn compute_moments(&self) -> Result<()> {
        self.submit("lbm moments", &[(self.current, "cell_moments", self.cells)])
    }
}

impl GpuFlow {
    /// One submission of `steps` (≤ MAX_CHUNK) time steps, then the stability check.
    fn advance_chunk(&mut self, steps: u64) -> Result<()> {
        self.queue.write_buffer(&self.stats, 0, &bytes_u32(&[0, 0, 0, 0]));
        let mut commands = Vec::with_capacity(8 * MAX_CHUNK as usize + 3);
        for step in 0..steps {
            commands.push((self.current, "collide", self.cells));
            if !self.aa {
                commands.push((self.current, "stream", self.cells));
            }
            if self.link_count > 0 {
                // The last step's momentum exchange is the one read back.
                let exchange = self.aa && step + 1 == steps;
                commands.push((self.current, if exchange { "stream_links_exchange" } else { "stream_links" }, self.links_groups));
                if self.aa {
                    commands.push((self.current, "scatter_links", self.links_groups));
                }
            }
            if self.open_count > 0 {
                commands.push((self.current, "open_boundaries", self.open_groups));
            }
            if self.thermal.is_some() {
                commands.push((self.current, "collide_thermal", self.cells));
                commands.push((self.current, "stream_thermal", self.cells));
                if self.link_count > 0 {
                    commands.push((self.current, "stream_links_thermal", self.links_groups));
                }
                if self.open_count > 0 {
                    commands.push((self.current, "open_thermal", self.open_groups));
                }
            }
            self.current ^= 1;
        }
        // The momentum (and heat) exchange of the last step, from its post-collision
        // populations.
        if self.link_count > 0 && self.group_count > 0 {
            commands.push((self.current ^ 1, "link_forces", self.links_groups));
            commands.push((self.current ^ 1, "group_forces", self.force_groups));
        }
        if self.thermal.is_some() {
            commands.push((self.current, "thermal_check", self.cells));
        }
        self.submit("lbm steps", &commands)?;
        self.steps_done += steps;
        self.exchange_cache = None;
        let stats = self.read(&self.stats, 16)?;
        let word = |k: usize| stats.get(4 * k..4 * k + 4).and_then(|b| b.try_into().ok()).map(u32::from_le_bytes).unwrap_or(0);
        let max_u2 = f32::from_bits(word(0));
        if word(1) != 0 || !(f64::from(max_u2) <= MAX_LATTICE_SPEED * MAX_LATTICE_SPEED) {
            return Err(Error::Unstable(format!(
                "the flow went unstable by time step {}: the lattice speed reached {:.3} (limit {MAX_LATTICE_SPEED}) or a density stopped being a positive number; lower the time step, refine the spacing or check the boundary conditions",
                self.steps_done,
                max_u2.sqrt()
            )));
        }
        if word(2) != 0 {
            return Err(Error::Unstable(format!(
                "the heat went unstable by time step {}: a temperature stopped being a number; lower the time step, refine the spacing or check the heat fluxes",
                self.steps_done
            )));
        }
        Ok(())
    }
}

impl Backend for GpuFlow {
    fn name(&self) -> String {
        format!("gpu: {}, {} streaming", self.adapter, if self.aa { "in-place" } else { "two-buffer" })
    }

    fn advance(&mut self, steps: u64) -> Result<()> {
        // Submissions of at most MAX_CHUNK steps: the stability check runs after each, so a
        // run that blows up is caught within that many steps whatever the caller asks for.
        let mut left = steps;
        while left > 0 {
            let chunk = left.min(MAX_CHUNK);
            self.advance_chunk(chunk)?;
            left -= chunk;
        }
        Ok(())
    }

    fn steps_done(&self) -> u64 {
        self.steps_done
    }

    fn fields(&mut self) -> Result<Fields> {
        self.compute_moments()?;
        let m = self.read_f32(&self.moments, 4 * self.n)?;
        let density = m.as_chunks::<4>().0.iter().map(|v| f64::from(v[3])).collect();
        let velocity = m.as_chunks::<4>().0.iter().map(|v| [f64::from(v[0]), f64::from(v[1]), f64::from(v[2])]).collect();
        // The temperature from the scalar populations, summed on the host: T₀ + Σ t (T₀ in
        // solid cells, whose populations stay zero).
        let temperature = match &self.thermal {
            None => Vec::new(),
            Some(therm) => {
                let t = self.read_f32(therm.get(self.current).ok_or_else(|| Error::Invalid("GPU: no thermal buffer".into()))?, QT * self.n)?;
                (0..self.n)
                    .map(|cell| self.initial_temperature + (0..QT).map(|q| t.get(q * self.n + cell).copied().map_or(0.0, f64::from)).sum::<f64>())
                    .collect()
            }
        };
        Ok(Fields { density, velocity, temperature })
    }

    fn heat(&mut self) -> Result<Vec<f64>> {
        if self.thermal.is_none() {
            return Ok(Vec::new());
        }
        Ok(self.exchange()?.iter().map(|g| g[3]).collect())
    }

    fn mass(&mut self) -> Result<f64> {
        self.compute_moments()?;
        self.submit("lbm mass", &[(self.current, "mass", self.cells)])?;
        let partials = self.read_f32(&self.partials, self.partial_count)?;
        // Summed in a fixed order on the host, so the result is deterministic on this device.
        Ok(partials.iter().map(|v| f64::from(*v)).sum())
    }

    fn forces(&mut self) -> Result<Vec<[f64; 3]>> {
        Ok(self.exchange()?.iter().map(|g| [g[0], g[1], g[2]]).collect())
    }

    fn groups(&self) -> &[Group] {
        &self.groups
    }

    fn fluid_cells(&self) -> usize {
        self.fluid_cells
    }

    fn is_fluid(&self, cell: usize) -> bool {
        self.fluid.get(cell).copied().unwrap_or(false)
    }

    fn empty_obstacles(&self) -> &[String] {
        &self.empty_obstacles
    }
}