rustsim 0.0.1

//! Authoritative columnar phase runtime.
//!
//! [`ColumnarRuntime`] is the production-oriented execution path for large
//! populations. Unlike [`rustsim_core::standard::StandardModel`], which keeps
//! agents in an AoS store and steps them one at a time, this runtime owns a
//! persistent [`DeviceSoaStore`] and treats the columnar buffers as the source
//! of truth. AoS stores are upload/download boundaries only.
//!
//! The runtime executes explicit ordered phases. CPU phases operate directly on
//! resident SoA columns; CUDA-resident phases launch kernels against persistent
//! device buffers and only synchronize back to host when a CPU phase or explicit
//! download requires it.

use crate::device_store::{DeviceSoaMutationError, DeviceSoaStore, DeviceSoaStoreF64};
#[cfg(feature = "cuda")]
use crate::device_store::{SoaColumnSchema, SoaColumnType};
use crate::interaction::{DeviceSpatialConfig2D, DeviceSpatialError, DeviceSpatialIndex2D};
use rustsim_core::soa::{SoaExtractable, SoaExtractableF64};
use rustsim_core::store::AgentStore;
use rustsim_core::types::AgentId;
use thiserror::Error;

/// Errors returned by [`ColumnarRuntime`].
#[derive(Debug, Error)]
pub enum ColumnarRuntimeError {
    /// A caller supplied an invalid runtime configuration.
    #[error("invalid columnar runtime configuration: {0}")]
    InvalidConfig(String),

    /// A lifecycle mutation would make the authoritative SoA state invalid.
    #[error("columnar runtime mutation failed: {0}")]
    Mutation(#[from] DeviceSoaMutationError),

    /// A spatial interaction index or phase was configured incorrectly.
    #[error("columnar runtime spatial interaction failed: {0}")]
    Spatial(#[from] DeviceSpatialError),

    /// CUDA initialization, launch, or synchronization failed.
    #[cfg(feature = "cuda")]
    #[error("cuda resident runtime error: {0}")]
    Cuda(String),
}

/// Owned `f32` SoA rows used by production lifecycle commands.
#[derive(Debug, Clone, PartialEq)]
pub struct ColumnarAgentBatch {
    /// Agent IDs in the same row order as every column.
    pub ids: Vec<AgentId>,
    /// Column data in SoA layout. `columns[c][i]` belongs to `ids[i]`.
    pub columns: Vec<Vec<f32>>,
}

impl ColumnarAgentBatch {
    /// Build an owned insert payload.
    pub fn new(ids: Vec<AgentId>, columns: Vec<Vec<f32>>) -> Self {
        Self { ids, columns }
    }

    fn column_refs(&self) -> Vec<&[f32]> {
        self.columns.iter().map(Vec::as_slice).collect()
    }
}

/// Owned `f64` SoA rows used by double-precision lifecycle commands.
#[derive(Debug, Clone, PartialEq)]
pub struct ColumnarAgentBatchF64 {
    /// Agent IDs in the same row order as every column.
    pub ids: Vec<AgentId>,
    /// Column data in SoA layout. `columns[c][i]` belongs to `ids[i]`.
    pub columns: Vec<Vec<f64>>,
}

impl ColumnarAgentBatchF64 {
    /// Build an owned insert payload.
    pub fn new(ids: Vec<AgentId>, columns: Vec<Vec<f64>>) -> Self {
        Self { ids, columns }
    }

    fn column_refs(&self) -> Vec<&[f64]> {
        self.columns.iter().map(Vec::as_slice).collect()
    }
}

/// Authoritative lifecycle mutation for a columnar runtime.
#[derive(Debug, Clone, PartialEq)]
pub enum ColumnarLifecycleCommand {
    /// Append new rows to the authoritative SoA state.
    Insert(ColumnarAgentBatch),
    /// Remove existing agent IDs from the authoritative SoA state.
    Remove(Vec<AgentId>),
    /// Validate then remove and insert as one runtime mutation.
    Replace {
        /// Existing IDs to remove before appending replacement rows.
        remove_ids: Vec<AgentId>,
        /// New rows to append after removals are validated.
        insert: ColumnarAgentBatch,
    },
}

/// Double-precision lifecycle mutation for a columnar runtime.
#[derive(Debug, Clone, PartialEq)]
pub enum ColumnarLifecycleCommandF64 {
    /// Append new rows to the authoritative SoA state.
    Insert(ColumnarAgentBatchF64),
    /// Remove existing agent IDs from the authoritative SoA state.
    Remove(Vec<AgentId>),
    /// Validate then remove and insert as one runtime mutation.
    Replace {
        /// Existing IDs to remove before appending replacement rows.
        remove_ids: Vec<AgentId>,
        /// New rows to append after removals are validated.
        insert: ColumnarAgentBatchF64,
    },
}

/// Summary returned after a lifecycle command mutates authoritative SoA state.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ColumnarLifecycleReport {
    /// Number of agents removed by the command.
    pub removed: usize,
    /// Number of agents inserted by the command.
    pub inserted: usize,
    /// Active agent count after the command.
    pub agent_count: usize,
}

/// Declarative ABI contract for a CUDA-resident kernel phase.
#[cfg(feature = "cuda")]
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CudaKernelAbi {
    /// Number of SoA columns the kernel expects before the trailing row count.
    pub columns: usize,
    /// Precision expected for every column.
    pub column_type: SoaColumnType,
}

#[cfg(feature = "cuda")]
impl CudaKernelAbi {
    /// ABI for kernels receiving all columns as `float*` plus a trailing row count.
    pub const fn f32_columns(columns: usize) -> Self {
        Self {
            columns,
            column_type: SoaColumnType::F32,
        }
    }

    fn validate(&self, schema: &[SoaColumnSchema]) -> Result<(), ColumnarRuntimeError> {
        if schema.len() != self.columns {
            return Err(ColumnarRuntimeError::InvalidConfig(format!(
                "CUDA kernel ABI expects {} columns but runtime has {}",
                self.columns,
                schema.len()
            )));
        }
        for (column, entry) in schema.iter().enumerate() {
            if entry.column_type != self.column_type {
                return Err(ColumnarRuntimeError::InvalidConfig(format!(
                    "CUDA kernel ABI expects {:?} at column {} but runtime schema has {:?}",
                    self.column_type, column, entry.column_type
                )));
            }
        }
        Ok(())
    }
}

/// Timing for a single executed phase.
#[derive(Debug, Clone)]
pub struct ColumnarPhaseTiming {
    /// Phase index in the runtime's ordered phase list.
    pub phase_index: usize,
    /// Human-readable phase name.
    pub name: &'static str,
    /// Backend used for this phase.
    pub backend: ColumnarPhaseBackend,
    /// Wall-clock time in microseconds.
    ///
    /// For CUDA-resident phases this is launch-submission time; call
    /// [`ColumnarRuntime::sync_to_host`] or [`ColumnarRuntime::download`] when
    /// host-visible completion is required.
    pub elapsed_us: u128,
}

/// Timing for one complete columnar runtime step.
#[derive(Debug, Clone)]
pub struct ColumnarStepTiming {
    /// Step index before the step was advanced.
    pub step_index: u64,
    /// Number of agents processed.
    pub agent_count: usize,
    /// Per-phase timings.
    pub phases: Vec<ColumnarPhaseTiming>,
    /// Total host-side wall-clock time for the step.
    pub total_us: u128,
}

/// Backend used to execute a phase.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ColumnarPhaseBackend {
    /// Single-threaded CPU phase over resident SoA columns.
    Cpu,
    /// Rayon chunked CPU phase over resident SoA columns.
    #[cfg(feature = "rayon")]
    Rayon,
    /// CUDA-resident kernel launch.
    #[cfg(feature = "cuda")]
    CudaResident,
}

/// Runtime whose authoritative state is a persistent SoA store.
pub struct ColumnarRuntime {
    store: DeviceSoaStore,
    phases: Vec<ColumnarPhase>,
    step_index: u64,
    #[cfg(feature = "cuda")]
    cuda_ptx: Option<String>,
    #[cfg(feature = "cuda")]
    cuda_initialized: bool,
}

type CpuPhaseFn = dyn FnMut(&mut [Vec<f32>], usize);
type CpuPhaseFnF64 = dyn FnMut(&mut [Vec<f64>], usize);
type SpatialCpuPhaseFn = dyn FnMut(&DeviceSpatialIndex2D, &mut [Vec<f32>], usize);
type SpatialCpuPhaseFnF64 = dyn FnMut(&DeviceSpatialIndex2D, &mut [Vec<f64>], usize);

enum ColumnarPhase {
    Cpu {
        name: &'static str,
        function: Box<CpuPhaseFn>,
    },
    SpatialCpu2D {
        name: &'static str,
        config: DeviceSpatialConfig2D,
        function: Box<SpatialCpuPhaseFn>,
    },
    #[cfg(feature = "rayon")]
    Rayon {
        name: &'static str,
        chunk_size: usize,
        function: Box<dyn Fn(usize, &mut [&mut [f32]]) + Send + Sync>,
    },
    #[cfg(feature = "cuda")]
    CudaResident {
        name: &'static str,
        kernel_name: String,
        block_size: u32,
    },
}

impl ColumnarRuntime {
    /// Build a columnar runtime by uploading an existing AoS [`AgentStore`].
    ///
    /// After construction, the [`DeviceSoaStore`] inside this runtime is the
    /// authoritative simulation state. Call [`download`](Self::download) only
    /// at observation, checkpoint, or interop boundaries.
    pub fn upload<A, S>(store: &S) -> Self
    where
        A: SoaExtractable,
        S: AgentStore<A>,
    {
        Self::from_device_store(DeviceSoaStore::upload::<A, S>(store))
    }

    /// Build a runtime from an existing persistent [`DeviceSoaStore`].
    pub fn from_device_store(store: DeviceSoaStore) -> Self {
        Self {
            store,
            phases: Vec::new(),
            step_index: 0,
            #[cfg(feature = "cuda")]
            cuda_ptx: None,
            #[cfg(feature = "cuda")]
            cuda_initialized: false,
        }
    }

    /// Consume the runtime and return its authoritative store.
    pub fn into_device_store(self) -> DeviceSoaStore {
        self.store
    }

    /// Replace the authoritative store and step index while preserving phases.
    ///
    /// This is the restore boundary used by the production control plane:
    /// configured CPU/rayon/CUDA phases remain attached to the runtime, while
    /// the simulation state is rewound to a checkpoint. CUDA-resident state is
    /// reinitialized lazily on the next CUDA phase.
    pub fn restore_device_store(&mut self, store: DeviceSoaStore, step_index: u64) {
        self.store = store;
        self.step_index = step_index;
        #[cfg(feature = "cuda")]
        {
            self.cuda_initialized = false;
        }
    }

    /// Access the authoritative store.
    pub fn device_store(&self) -> &DeviceSoaStore {
        &self.store
    }

    /// Mutably access the authoritative store.
    ///
    /// When CUDA residency is active, prefer the runtime's phase and scatter
    /// methods so host/device dirty flags stay coordinated.
    pub fn device_store_mut(&mut self) -> &mut DeviceSoaStore {
        #[cfg(feature = "cuda")]
        self.store.mark_host_dirty();
        &mut self.store
    }

    /// Current step index.
    pub fn step_index(&self) -> u64 {
        self.step_index
    }

    /// Number of active agents.
    pub fn agent_count(&self) -> usize {
        self.store.agent_count()
    }

    /// Number of SoA columns.
    pub fn num_columns(&self) -> usize {
        self.store.num_columns()
    }

    /// Agent IDs in current row order.
    pub fn ids(&self) -> &[AgentId] {
        self.store.ids()
    }

    /// Return the current row for `id`, if present.
    pub fn row_of(&self, id: AgentId) -> Option<usize> {
        self.store.row_of(id)
    }

    /// Whether the authoritative runtime state contains `id`.
    pub fn contains_id(&self, id: AgentId) -> bool {
        self.store.contains_id(id)
    }

    /// Smallest non-negative [`AgentId`] not currently present in the runtime.
    pub fn next_available_id(&self) -> AgentId {
        self.store.next_available_id()
    }

    /// Column names supplied by [`SoaExtractable`].
    pub fn column_names(&self) -> &[&'static str] {
        self.store.column_names()
    }

    /// Add a serial CPU phase.
    pub fn add_cpu_phase(
        &mut self,
        name: &'static str,
        function: impl FnMut(&mut [Vec<f32>], usize) + 'static,
    ) {
        self.phases.push(ColumnarPhase::Cpu {
            name,
            function: Box::new(function),
        });
    }

    /// Add a 2-D spatial CPU phase over authoritative SoA state.
    ///
    /// Each step rebuilds a row-ordered spatial index from the configured
    /// coordinate columns, then invokes `function(index, columns, n)`. This is
    /// the production replacement for AoS-side neighbor/message systems when a
    /// run is owned by `ColumnarRuntime`.
    pub fn add_spatial_cpu_phase_2d(
        &mut self,
        name: &'static str,
        config: DeviceSpatialConfig2D,
        function: impl FnMut(&DeviceSpatialIndex2D, &mut [Vec<f32>], usize) + 'static,
    ) -> Result<(), ColumnarRuntimeError> {
        DeviceSpatialIndex2D::build_f32(&self.store, config)?;
        self.phases.push(ColumnarPhase::SpatialCpu2D {
            name,
            config,
            function: Box::new(function),
        });
        Ok(())
    }

    /// Build a 2-D spatial index from the current authoritative SoA columns.
    pub fn spatial_index_2d(
        &mut self,
        config: DeviceSpatialConfig2D,
    ) -> Result<DeviceSpatialIndex2D, ColumnarRuntimeError> {
        self.sync_to_host()?;
        Ok(DeviceSpatialIndex2D::build_f32(&self.store, config)?)
    }

    /// Add a chunked Rayon CPU phase.
    #[cfg(feature = "rayon")]
    pub fn add_rayon_phase(
        &mut self,
        name: &'static str,
        chunk_size: usize,
        function: impl Fn(usize, &mut [&mut [f32]]) + Send + Sync + 'static,
    ) -> Result<(), ColumnarRuntimeError> {
        if chunk_size == 0 {
            return Err(ColumnarRuntimeError::InvalidConfig(
                "chunk_size must be positive".to_string(),
            ));
        }
        self.phases.push(ColumnarPhase::Rayon {
            name,
            chunk_size,
            function: Box::new(function),
        });
        Ok(())
    }

    /// Set the PTX module used by subsequent CUDA-resident phases.
    #[cfg(feature = "cuda")]
    pub fn set_cuda_ptx(&mut self, ptx_source: impl Into<String>) {
        self.cuda_ptx = Some(ptx_source.into());
        self.cuda_initialized = false;
        self.store.release_cuda();
    }

    /// Add a CUDA-resident kernel phase.
    ///
    /// Kernels are launched against the runtime's persistent device columns.
    /// Consecutive CUDA phases are submitted as a single resident sequence,
    /// with no host/device transfer between phases.
    #[cfg(feature = "cuda")]
    pub fn add_cuda_resident_phase(
        &mut self,
        name: &'static str,
        kernel_name: impl Into<String>,
        block_size: u32,
    ) -> Result<(), ColumnarRuntimeError> {
        if block_size == 0 {
            return Err(ColumnarRuntimeError::InvalidConfig(
                "block_size must be positive".to_string(),
            ));
        }
        self.phases.push(ColumnarPhase::CudaResident {
            name,
            kernel_name: kernel_name.into(),
            block_size,
        });
        Ok(())
    }

    /// Add a CUDA-resident phase after validating its declared ABI against the runtime schema.
    #[cfg(feature = "cuda")]
    pub fn add_cuda_resident_phase_checked(
        &mut self,
        name: &'static str,
        kernel_name: impl Into<String>,
        block_size: u32,
        abi: CudaKernelAbi,
    ) -> Result<(), ColumnarRuntimeError> {
        abi.validate(self.store.schema())?;
        self.add_cuda_resident_phase(name, kernel_name, block_size)
    }

    /// Number of configured phases.
    pub fn phase_count(&self) -> usize {
        self.phases.len()
    }

    /// Execute one ordered phase step.
    pub fn step(&mut self) -> Result<ColumnarStepTiming, ColumnarRuntimeError> {
        let total_start = std::time::Instant::now();
        let step_index = self.step_index;
        let agent_count = self.agent_count();
        let mut timings = Vec::with_capacity(self.phases.len());

        let mut phase_index = 0;
        while phase_index < self.phases.len() {
            match &mut self.phases[phase_index] {
                ColumnarPhase::Cpu { name, function } => {
                    #[cfg(feature = "cuda")]
                    self.store
                        .sync_to_host()
                        .map_err(ColumnarRuntimeError::Cuda)?;

                    let phase_start = std::time::Instant::now();
                    function(self.store.columns_mut(), agent_count);
                    #[cfg(feature = "cuda")]
                    self.store.mark_host_dirty();
                    timings.push(ColumnarPhaseTiming {
                        phase_index,
                        name,
                        backend: ColumnarPhaseBackend::Cpu,
                        elapsed_us: phase_start.elapsed().as_micros(),
                    });
                    phase_index += 1;
                }
                ColumnarPhase::SpatialCpu2D {
                    name,
                    config,
                    function,
                } => {
                    #[cfg(feature = "cuda")]
                    self.store
                        .sync_to_host()
                        .map_err(ColumnarRuntimeError::Cuda)?;

                    let phase_start = std::time::Instant::now();
                    let index = DeviceSpatialIndex2D::build_f32(&self.store, *config)?;
                    function(&index, self.store.columns_mut(), agent_count);
                    #[cfg(feature = "cuda")]
                    self.store.mark_host_dirty();
                    timings.push(ColumnarPhaseTiming {
                        phase_index,
                        name,
                        backend: ColumnarPhaseBackend::Cpu,
                        elapsed_us: phase_start.elapsed().as_micros(),
                    });
                    phase_index += 1;
                }
                #[cfg(feature = "rayon")]
                ColumnarPhase::Rayon {
                    name,
                    chunk_size,
                    function,
                } => {
                    #[cfg(feature = "cuda")]
                    self.store
                        .sync_to_host()
                        .map_err(ColumnarRuntimeError::Cuda)?;

                    let phase_start = std::time::Instant::now();
                    {
                        let mut slices: Vec<&mut [f32]> = self
                            .store
                            .columns_mut()
                            .iter_mut()
                            .map(|column| column.as_mut_slice())
                            .collect();
                        let function_ref: &(dyn Fn(usize, &mut [&mut [f32]]) + Send + Sync) =
                            &**function;
                        crate::parallel::par_apply_chunks_multi(
                            &mut slices,
                            *chunk_size,
                            |start, chunk| function_ref(start, chunk),
                        );
                    }
                    #[cfg(feature = "cuda")]
                    self.store.mark_host_dirty();
                    timings.push(ColumnarPhaseTiming {
                        phase_index,
                        name,
                        backend: ColumnarPhaseBackend::Rayon,
                        elapsed_us: phase_start.elapsed().as_micros(),
                    });
                    phase_index += 1;
                }
                #[cfg(feature = "cuda")]
                ColumnarPhase::CudaResident { .. } => {
                    let start_index = phase_index;
                    let mut names = Vec::new();
                    let mut kernels = Vec::new();
                    let mut block_size = None;

                    while phase_index < self.phases.len() {
                        match &self.phases[phase_index] {
                            ColumnarPhase::CudaResident {
                                name,
                                kernel_name,
                                block_size: phase_block_size,
                            } => {
                                if let Some(expected) = block_size {
                                    if expected != *phase_block_size {
                                        break;
                                    }
                                } else {
                                    block_size = Some(*phase_block_size);
                                }
                                names.push(*name);
                                kernels.push(kernel_name.clone());
                                phase_index += 1;
                            }
                            _ => break,
                        }
                    }

                    let elapsed = self.run_cuda_sequence(&kernels, block_size.unwrap_or(256))?;
                    for (offset, (name, elapsed_us)) in names.into_iter().zip(elapsed).enumerate() {
                        timings.push(ColumnarPhaseTiming {
                            phase_index: start_index + offset,
                            name,
                            backend: ColumnarPhaseBackend::CudaResident,
                            elapsed_us,
                        });
                    }
                }
            }
        }

        self.step_index = self.step_index.saturating_add(1);
        Ok(ColumnarStepTiming {
            step_index,
            agent_count,
            phases: timings,
            total_us: total_start.elapsed().as_micros(),
        })
    }

    /// Execute `steps` ordered phase steps.
    pub fn run(&mut self, steps: usize) -> Result<Vec<ColumnarStepTiming>, ColumnarRuntimeError> {
        let mut timings = Vec::with_capacity(steps);
        for _ in 0..steps {
            timings.push(self.step()?);
        }
        Ok(timings)
    }

    /// Synchronize device-resident data back to host columns when CUDA is active.
    pub fn sync_to_host(&mut self) -> Result<(), ColumnarRuntimeError> {
        #[cfg(feature = "cuda")]
        {
            self.store
                .sync_to_host()
                .map_err(ColumnarRuntimeError::Cuda)?;
        }
        Ok(())
    }

    /// Download the authoritative columnar state into an AoS [`AgentStore`].
    pub fn download<A, S>(&mut self, store: &S) -> Result<(), ColumnarRuntimeError>
    where
        A: SoaExtractable,
        S: AgentStore<A>,
    {
        self.sync_to_host()?;
        self.store.download::<A, S>(store);
        Ok(())
    }

    /// Remove agents from the authoritative SoA state.
    pub fn scatter_remove(&mut self, dead_ids: &[AgentId]) -> Result<(), ColumnarRuntimeError> {
        self.sync_to_host()?;
        self.store.try_scatter_remove(dead_ids)?;
        #[cfg(feature = "cuda")]
        self.store.mark_host_dirty();
        Ok(())
    }

    /// Append agents to the authoritative SoA state.
    pub fn scatter_insert(
        &mut self,
        new_ids: &[AgentId],
        new_columns: &[&[f32]],
    ) -> Result<(), ColumnarRuntimeError> {
        self.sync_to_host()?;
        self.store.try_scatter_insert(new_ids, new_columns)?;
        #[cfg(feature = "cuda")]
        self.store.mark_host_dirty();
        Ok(())
    }

    /// Apply an authoritative lifecycle command to the production SoA runtime.
    ///
    /// This is the preferred mutation boundary for production runs: all IDs
    /// and column shapes are validated before state changes, CUDA-resident
    /// state is synchronized once, and the returned report records the new
    /// active population size.
    pub fn apply_lifecycle(
        &mut self,
        command: ColumnarLifecycleCommand,
    ) -> Result<ColumnarLifecycleReport, ColumnarRuntimeError> {
        self.sync_to_host()?;
        let (removed, inserted) = match command {
            ColumnarLifecycleCommand::Insert(batch) => {
                let refs = batch.column_refs();
                let inserted = batch.ids.len();
                self.store.try_scatter_insert(&batch.ids, &refs)?;
                (0, inserted)
            }
            ColumnarLifecycleCommand::Remove(ids) => {
                let removed = ids.len();
                self.store.try_scatter_remove(&ids)?;
                (removed, 0)
            }
            ColumnarLifecycleCommand::Replace { remove_ids, insert } => {
                let refs = insert.column_refs();
                let removed = remove_ids.len();
                let inserted = insert.ids.len();
                self.store
                    .try_scatter_replace(&remove_ids, &insert.ids, &refs)?;
                (removed, inserted)
            }
        };
        #[cfg(feature = "cuda")]
        self.store.mark_host_dirty();
        Ok(ColumnarLifecycleReport {
            removed,
            inserted,
            agent_count: self.agent_count(),
        })
    }

    /// Apply several lifecycle commands in order.
    pub fn apply_lifecycle_batch<I>(
        &mut self,
        commands: I,
    ) -> Result<Vec<ColumnarLifecycleReport>, ColumnarRuntimeError>
    where
        I: IntoIterator<Item = ColumnarLifecycleCommand>,
    {
        let mut reports = Vec::new();
        for command in commands {
            reports.push(self.apply_lifecycle(command)?);
        }
        Ok(reports)
    }

    #[cfg(feature = "cuda")]
    fn run_cuda_sequence(
        &mut self,
        kernels: &[String],
        block_size: u32,
    ) -> Result<Vec<u128>, ColumnarRuntimeError> {
        if kernels.is_empty() {
            return Ok(Vec::new());
        }
        let ptx = self.cuda_ptx.as_deref().ok_or_else(|| {
            ColumnarRuntimeError::InvalidConfig(
                "set_cuda_ptx must be called before CUDA resident phases execute".to_string(),
            )
        })?;
        if !self.cuda_initialized || !self.store.has_cuda_resident() {
            self.store
                .init_cuda(ptx, &kernels[0])
                .map_err(ColumnarRuntimeError::Cuda)?;
            self.cuda_initialized = true;
        }

        let kernel_refs: Vec<&str> = kernels.iter().map(String::as_str).collect();
        self.store
            .run_sequence_cuda_resident(&kernel_refs, block_size)
            .map(|items| items.into_iter().map(|(_, elapsed)| elapsed).collect())
            .map_err(ColumnarRuntimeError::Cuda)
    }
}

/// Authoritative double-precision columnar runtime.
pub struct ColumnarRuntimeF64 {
    store: DeviceSoaStoreF64,
    phases: Vec<ColumnarPhaseF64>,
    step_index: u64,
}

enum ColumnarPhaseF64 {
    Cpu {
        name: &'static str,
        function: Box<CpuPhaseFnF64>,
    },
    SpatialCpu2D {
        name: &'static str,
        config: DeviceSpatialConfig2D,
        function: Box<SpatialCpuPhaseFnF64>,
    },
}

impl ColumnarRuntimeF64 {
    /// Build a double-precision columnar runtime by uploading an AoS store.
    pub fn upload<A, S>(store: &S) -> Self
    where
        A: SoaExtractableF64,
        S: AgentStore<A>,
    {
        Self::from_device_store(DeviceSoaStoreF64::upload::<A, S>(store))
    }

    /// Build a runtime from an existing persistent f64 SoA store.
    pub fn from_device_store(store: DeviceSoaStoreF64) -> Self {
        Self {
            store,
            phases: Vec::new(),
            step_index: 0,
        }
    }

    /// Consume the runtime and return its authoritative store.
    pub fn into_device_store(self) -> DeviceSoaStoreF64 {
        self.store
    }

    /// Replace the authoritative store and step index while preserving phases.
    pub fn restore_device_store(&mut self, store: DeviceSoaStoreF64, step_index: u64) {
        self.store = store;
        self.step_index = step_index;
    }

    /// Access the authoritative f64 store.
    pub fn device_store(&self) -> &DeviceSoaStoreF64 {
        &self.store
    }

    /// Mutably access the authoritative f64 store.
    pub fn device_store_mut(&mut self) -> &mut DeviceSoaStoreF64 {
        &mut self.store
    }

    /// Current step index.
    pub fn step_index(&self) -> u64 {
        self.step_index
    }

    /// Number of active agents.
    pub fn agent_count(&self) -> usize {
        self.store.agent_count()
    }

    /// Number of SoA columns.
    pub fn num_columns(&self) -> usize {
        self.store.num_columns()
    }

    /// Agent IDs in current row order.
    pub fn ids(&self) -> &[AgentId] {
        self.store.ids()
    }

    /// Return the current row for `id`, if present.
    pub fn row_of(&self, id: AgentId) -> Option<usize> {
        self.store.row_of(id)
    }

    /// Whether the authoritative runtime state contains `id`.
    pub fn contains_id(&self, id: AgentId) -> bool {
        self.store.contains_id(id)
    }

    /// Smallest non-negative [`AgentId`] not currently present in the runtime.
    pub fn next_available_id(&self) -> AgentId {
        self.store.next_available_id()
    }

    /// Column names supplied by `SoaExtractableF64`.
    pub fn column_names(&self) -> &[&'static str] {
        self.store.column_names()
    }

    /// Add a serial CPU phase over f64 columns.
    pub fn add_cpu_phase(
        &mut self,
        name: &'static str,
        function: impl FnMut(&mut [Vec<f64>], usize) + 'static,
    ) {
        self.phases.push(ColumnarPhaseF64::Cpu {
            name,
            function: Box::new(function),
        });
    }

    /// Add a 2-D spatial CPU phase over authoritative f64 SoA state.
    pub fn add_spatial_cpu_phase_2d(
        &mut self,
        name: &'static str,
        config: DeviceSpatialConfig2D,
        function: impl FnMut(&DeviceSpatialIndex2D, &mut [Vec<f64>], usize) + 'static,
    ) -> Result<(), ColumnarRuntimeError> {
        DeviceSpatialIndex2D::build_f64(&self.store, config)?;
        self.phases.push(ColumnarPhaseF64::SpatialCpu2D {
            name,
            config,
            function: Box::new(function),
        });
        Ok(())
    }

    /// Build a 2-D spatial index from the current authoritative f64 SoA columns.
    pub fn spatial_index_2d(
        &self,
        config: DeviceSpatialConfig2D,
    ) -> Result<DeviceSpatialIndex2D, ColumnarRuntimeError> {
        Ok(DeviceSpatialIndex2D::build_f64(&self.store, config)?)
    }

    /// Number of configured phases.
    pub fn phase_count(&self) -> usize {
        self.phases.len()
    }

    /// Execute one ordered double-precision phase step.
    pub fn step(&mut self) -> Result<ColumnarStepTiming, ColumnarRuntimeError> {
        let total_start = std::time::Instant::now();
        let step_index = self.step_index;
        let agent_count = self.agent_count();
        let mut timings = Vec::with_capacity(self.phases.len());

        for (phase_index, phase) in self.phases.iter_mut().enumerate() {
            match phase {
                ColumnarPhaseF64::Cpu { name, function } => {
                    let phase_start = std::time::Instant::now();
                    function(self.store.columns_mut(), agent_count);
                    timings.push(ColumnarPhaseTiming {
                        phase_index,
                        name,
                        backend: ColumnarPhaseBackend::Cpu,
                        elapsed_us: phase_start.elapsed().as_micros(),
                    });
                }
                ColumnarPhaseF64::SpatialCpu2D {
                    name,
                    config,
                    function,
                } => {
                    let phase_start = std::time::Instant::now();
                    let index = DeviceSpatialIndex2D::build_f64(&self.store, *config)?;
                    function(&index, self.store.columns_mut(), agent_count);
                    timings.push(ColumnarPhaseTiming {
                        phase_index,
                        name,
                        backend: ColumnarPhaseBackend::Cpu,
                        elapsed_us: phase_start.elapsed().as_micros(),
                    });
                }
            }
        }

        self.step_index = self.step_index.saturating_add(1);
        Ok(ColumnarStepTiming {
            step_index,
            agent_count,
            phases: timings,
            total_us: total_start.elapsed().as_micros(),
        })
    }

    /// Execute `steps` ordered phase steps.
    pub fn run(&mut self, steps: usize) -> Result<Vec<ColumnarStepTiming>, ColumnarRuntimeError> {
        let mut timings = Vec::with_capacity(steps);
        for _ in 0..steps {
            timings.push(self.step()?);
        }
        Ok(timings)
    }

    /// No-op parity method with the CUDA-enabled f32 runtime.
    pub fn sync_to_host(&mut self) -> Result<(), ColumnarRuntimeError> {
        Ok(())
    }

    /// Download the authoritative f64 columnar state into an AoS `AgentStore`.
    pub fn download<A, S>(&mut self, store: &S) -> Result<(), ColumnarRuntimeError>
    where
        A: SoaExtractableF64,
        S: AgentStore<A>,
    {
        self.store.download::<A, S>(store);
        Ok(())
    }

    /// Remove agents from the authoritative f64 SoA state.
    pub fn scatter_remove(&mut self, dead_ids: &[AgentId]) -> Result<(), ColumnarRuntimeError> {
        self.store.try_scatter_remove(dead_ids)?;
        Ok(())
    }

    /// Append agents to the authoritative f64 SoA state.
    pub fn scatter_insert(
        &mut self,
        new_ids: &[AgentId],
        new_columns: &[&[f64]],
    ) -> Result<(), ColumnarRuntimeError> {
        self.store.try_scatter_insert(new_ids, new_columns)?;
        Ok(())
    }

    /// Apply an authoritative lifecycle command to the double-precision runtime.
    pub fn apply_lifecycle(
        &mut self,
        command: ColumnarLifecycleCommandF64,
    ) -> Result<ColumnarLifecycleReport, ColumnarRuntimeError> {
        let (removed, inserted) = match command {
            ColumnarLifecycleCommandF64::Insert(batch) => {
                let refs = batch.column_refs();
                let inserted = batch.ids.len();
                self.store.try_scatter_insert(&batch.ids, &refs)?;
                (0, inserted)
            }
            ColumnarLifecycleCommandF64::Remove(ids) => {
                let removed = ids.len();
                self.store.try_scatter_remove(&ids)?;
                (removed, 0)
            }
            ColumnarLifecycleCommandF64::Replace { remove_ids, insert } => {
                let refs = insert.column_refs();
                let removed = remove_ids.len();
                let inserted = insert.ids.len();
                self.store
                    .try_scatter_replace(&remove_ids, &insert.ids, &refs)?;
                (removed, inserted)
            }
        };
        Ok(ColumnarLifecycleReport {
            removed,
            inserted,
            agent_count: self.agent_count(),
        })
    }

    /// Apply several double-precision lifecycle commands in order.
    pub fn apply_lifecycle_batch<I>(
        &mut self,
        commands: I,
    ) -> Result<Vec<ColumnarLifecycleReport>, ColumnarRuntimeError>
    where
        I: IntoIterator<Item = ColumnarLifecycleCommandF64>,
    {
        let mut reports = Vec::new();
        for command in commands {
            reports.push(self.apply_lifecycle(command)?);
        }
        Ok(reports)
    }
}