Skip to main content

sc_neurocore_engine/
lib.rs

1// SPDX-License-Identifier: AGPL-3.0-or-later
2// Commercial license available
3// © Concepts 1996–2026 Miroslav Šotek. All rights reserved.
4// © Code 2020–2026 Miroslav Šotek. All rights reserved.
5// ORCID: 0009-0009-3560-0851
6// Contact: www.anulum.li | protoscience@anulum.li
7// SC-NeuroCore — Engine Crate Root
8
9#![allow(
10    clippy::useless_conversion,
11    clippy::needless_range_loop,
12    clippy::too_many_arguments,
13    deprecated
14)]
15
16use numpy::{
17    IntoPyArray, PyArray1, PyArray2, PyArrayMethods, PyReadonlyArray1, PyReadonlyArray2,
18    PyReadwriteArray1, PyUntypedArrayMethods,
19};
20use pyo3::exceptions::PyValueError;
21use pyo3::prelude::*;
22use pyo3::types::PyDict;
23use pyo3::IntoPyObject;
24
25pub mod adc_to_spike;
26pub mod analysis;
27pub mod attention;
28pub mod bitstream;
29pub mod brunel;
30pub mod connectome;
31pub mod conv;
32pub mod cordiv;
33pub mod cortical_column;
34pub mod cortical_inject;
35pub mod dna;
36pub mod ei_network;
37pub mod encoder;
38pub mod evo;
39pub mod fault;
40pub mod fusion;
41#[cfg(feature = "gpu")]
42pub mod gpu;
43pub mod grad;
44pub mod graph;
45pub mod ir;
46pub mod layer;
47pub mod lgssm;
48pub mod network_runner;
49pub mod neuron;
50pub mod neurons;
51pub mod optimizer;
52pub mod partition;
53pub mod phi;
54pub mod photonic;
55pub mod ping;
56pub mod predictive_coding;
57pub mod pyo3_neurons;
58pub mod quantum;
59pub mod rall_dendrite;
60pub mod recorder;
61pub mod recurrent;
62pub mod rk4_neurons;
63pub mod sc_inference;
64pub mod scpn;
65pub mod simd;
66pub mod sobol;
67pub mod supervisor;
68pub mod synapses;
69pub mod topology;
70pub mod wilson_cowan;
71pub mod wong_wang;
72
73// ── HDC / VSA PyO3 wrapper ───────────────────────────────────────────
74
75#[pyclass(
76    name = "BitStreamTensor",
77    module = "sc_neurocore_engine.sc_neurocore_engine"
78)]
79pub struct PyBitStreamTensor {
80    inner: bitstream::BitStreamTensor,
81}
82
83#[pymethods]
84impl PyBitStreamTensor {
85    /// Create a random binary vector of `dimension` bits.
86    #[new]
87    #[pyo3(signature = (dimension=10000, seed=0xACE1))]
88    fn new(dimension: usize, seed: u64) -> Self {
89        use rand::SeedableRng;
90        let mut rng = rand_xoshiro::Xoshiro256PlusPlus::seed_from_u64(seed);
91        let data = bitstream::bernoulli_packed(0.5, dimension, &mut rng);
92        Self {
93            inner: bitstream::BitStreamTensor::from_words(data, dimension),
94        }
95    }
96
97    /// Create from pre-packed u64 words.
98    #[staticmethod]
99    fn from_packed(data: Vec<u64>, length: usize) -> PyResult<Self> {
100        if length == 0 {
101            return Err(pyo3::exceptions::PyValueError::new_err(
102                "bitstream length must be > 0",
103            ));
104        }
105        Ok(Self {
106            inner: bitstream::BitStreamTensor::from_words(data, length),
107        })
108    }
109
110    /// In-place XOR (HDC bind).
111    fn xor_inplace(&mut self, other: &PyBitStreamTensor) {
112        self.inner.xor_inplace(&other.inner);
113    }
114
115    /// XOR returning a new tensor (HDC bind).
116    fn xor(&self, other: &PyBitStreamTensor) -> PyBitStreamTensor {
117        PyBitStreamTensor {
118            inner: self.inner.xor(&other.inner),
119        }
120    }
121
122    /// Cyclic right rotation by `shift` bits (HDC permute).
123    fn rotate_right(&mut self, shift: usize) {
124        self.inner.rotate_right(shift);
125    }
126
127    /// Normalized Hamming distance (0.0 = identical, 1.0 = opposite).
128    fn hamming_distance(&self, other: &PyBitStreamTensor) -> f32 {
129        self.inner.hamming_distance(&other.inner)
130    }
131
132    /// Majority-vote bundle of multiple tensors.
133    #[staticmethod]
134    fn bundle(vectors: Vec<PyRef<'_, PyBitStreamTensor>>) -> PyBitStreamTensor {
135        let refs: Vec<&bitstream::BitStreamTensor> = vectors.iter().map(|v| &v.inner).collect();
136        PyBitStreamTensor {
137            inner: bitstream::BitStreamTensor::bundle(&refs),
138        }
139    }
140
141    /// Count of set bits.
142    fn popcount(&self) -> u64 {
143        bitstream::popcount(&self.inner)
144    }
145
146    /// Packed u64 words (read-only copy).
147    #[getter]
148    fn data(&self) -> Vec<u64> {
149        self.inner.data.clone()
150    }
151
152    /// Logical bit length.
153    #[getter]
154    fn length(&self) -> usize {
155        self.inner.length
156    }
157
158    fn __len__(&self) -> usize {
159        self.inner.length
160    }
161
162    fn __repr__(&self) -> String {
163        format!(
164            "BitStreamTensor(length={}, popcount={})",
165            self.inner.length,
166            bitstream::popcount(&self.inner)
167        )
168    }
169}
170
171#[pyclass(
172    name = "StdpSynapse",
173    module = "sc_neurocore_engine.sc_neurocore_engine"
174)]
175pub struct StdpSynapse {
176    inner: synapses::StdpSynapse,
177}
178
179#[pymethods]
180impl StdpSynapse {
181    #[new]
182    #[pyo3(signature = (initial_weight, data_width=16, fraction=8))]
183    fn new(initial_weight: i16, data_width: u32, fraction: u32) -> Self {
184        Self {
185            inner: synapses::StdpSynapse::new(initial_weight, data_width, fraction),
186        }
187    }
188
189    #[allow(clippy::too_many_arguments)]
190    #[pyo3(signature = (pre_spike, post_spike, a_plus=16, a_minus=-16, decay=250, w_min=0, w_max=32767))]
191    fn step(
192        &mut self,
193        pre_spike: bool,
194        post_spike: bool,
195        a_plus: i16,
196        a_minus: i16,
197        decay: i16,
198        w_min: i16,
199        w_max: i16,
200    ) {
201        let params = synapses::StdpParams {
202            a_plus,
203            a_minus,
204            decay,
205            w_min,
206            w_max,
207        };
208        self.inner.step(pre_spike, post_spike, &params);
209    }
210
211    #[getter]
212    fn weight(&self) -> i16 {
213        self.inner.weight
214    }
215
216    #[setter]
217    fn set_weight(&mut self, value: i16) {
218        self.inner.weight = value;
219    }
220}
221
222// ── Brunel Network PyO3 wrapper ──────────────────────────────────────
223
224#[pyclass(
225    name = "BrunelNetwork",
226    module = "sc_neurocore_engine.sc_neurocore_engine"
227)]
228pub struct PyBrunelNetwork {
229    inner: brunel::BrunelNetwork,
230}
231
232#[pymethods]
233impl PyBrunelNetwork {
234    #[new]
235    #[pyo3(signature = (
236        n_neurons,
237        w_indptr,
238        w_indices,
239        w_data,
240        leak_k,
241        gain_k,
242        ext_lambda,
243        ext_weight_fp,
244        data_width=16,
245        fraction=8,
246        v_rest=0,
247        v_reset=0,
248        v_threshold=256,
249        refractory_period=2,
250        seed=42
251    ))]
252    #[allow(clippy::too_many_arguments)]
253    fn new(
254        n_neurons: usize,
255        w_indptr: PyReadonlyArray1<'_, i64>,
256        w_indices: PyReadonlyArray1<'_, i64>,
257        w_data: PyReadonlyArray1<'_, i16>,
258        leak_k: i16,
259        gain_k: i16,
260        ext_lambda: f64,
261        ext_weight_fp: i16,
262        data_width: u32,
263        fraction: u32,
264        v_rest: i16,
265        v_reset: i16,
266        v_threshold: i16,
267        refractory_period: i32,
268        seed: u64,
269    ) -> PyResult<Self> {
270        let indptr = w_indptr
271            .as_slice()
272            .map_err(|e| PyValueError::new_err(format!("Cannot read w_indptr: {e}")))?;
273        let indices = w_indices
274            .as_slice()
275            .map_err(|e| PyValueError::new_err(format!("Cannot read w_indices: {e}")))?;
276        let data = w_data
277            .as_slice()
278            .map_err(|e| PyValueError::new_err(format!("Cannot read w_data: {e}")))?;
279
280        let row_offsets: Vec<usize> = indptr.iter().map(|&v| v as usize).collect();
281        let col_indices: Vec<usize> = indices.iter().map(|&v| v as usize).collect();
282        let values: Vec<i16> = data.to_vec();
283
284        let inner = brunel::BrunelNetwork::new(
285            n_neurons,
286            row_offsets,
287            col_indices,
288            values,
289            data_width,
290            fraction,
291            v_rest,
292            v_reset,
293            v_threshold,
294            refractory_period,
295            leak_k,
296            gain_k,
297            ext_lambda,
298            ext_weight_fp,
299            seed,
300        )
301        .map_err(PyValueError::new_err)?;
302
303        Ok(Self { inner })
304    }
305
306    fn run<'py>(&mut self, py: Python<'py>, n_steps: usize) -> Bound<'py, PyArray1<u32>> {
307        let counts = self.inner.run(n_steps);
308        counts.into_pyarray(py)
309    }
310}
311
312// ── NetworkRunner PyO3 wrapper ────────────────────────────────────────
313
314#[pyclass(
315    name = "NetworkRunner",
316    module = "sc_neurocore_engine.sc_neurocore_engine"
317)]
318pub struct PyNetworkRunner {
319    inner: network_runner::NetworkRunner,
320}
321
322#[pymethods]
323impl PyNetworkRunner {
324    #[new]
325    fn new() -> Self {
326        Self {
327            inner: network_runner::NetworkRunner::new(),
328        }
329    }
330
331    fn add_population(&mut self, model: &str, n: usize) -> PyResult<usize> {
332        let pop = network_runner::create_population(model, n).map_err(PyValueError::new_err)?;
333        Ok(self.inner.add_population(pop))
334    }
335
336    #[pyo3(signature = (src, tgt, row_offsets, col_indices, values, delay=0))]
337    fn add_projection(
338        &mut self,
339        src: usize,
340        tgt: usize,
341        row_offsets: Vec<i64>,
342        col_indices: Vec<i64>,
343        values: Vec<f64>,
344        delay: usize,
345    ) {
346        let ro: Vec<usize> = row_offsets.iter().map(|&x| x as usize).collect();
347        let ci: Vec<usize> = col_indices.iter().map(|&x| x as usize).collect();
348        let proj = network_runner::ProjectionRunner::new(src, tgt, ro, ci, values, delay);
349        self.inner.add_projection(proj);
350    }
351
352    fn step_population<'py>(
353        &mut self,
354        py: Python<'py>,
355        pop_index: usize,
356        currents: PyReadonlyArray1<'py, f64>,
357    ) -> PyResult<Py<PyAny>> {
358        let currents = currents.as_slice()?;
359        let (spikes, voltages) = self
360            .inner
361            .step_population_with_currents(pop_index, currents)
362            .map_err(PyValueError::new_err)?;
363        let dict = PyDict::new(py);
364        dict.set_item("spikes", spikes.into_pyarray(py))?;
365        dict.set_item("voltages", voltages.into_pyarray(py))?;
366        Ok(dict.into_any().unbind())
367    }
368
369    fn run<'py>(&mut self, py: Python<'py>, n_steps: usize) -> PyResult<Py<PyAny>> {
370        let results = self.inner.run(n_steps);
371        let dict = PyDict::new(py);
372        let spike_counts: Vec<u64> = results.spike_counts.iter().map(|&c| c as u64).collect();
373        dict.set_item("spike_counts", spike_counts.into_pyarray(py))?;
374        let spike_data: Vec<Py<PyArray1<u64>>> = results
375            .spike_data
376            .into_iter()
377            .map(|v: Vec<u64>| v.into_pyarray(py).unbind())
378            .collect();
379        dict.set_item("spike_data", spike_data)?;
380        let voltages: Vec<Py<PyArray1<f64>>> = results
381            .voltages
382            .into_iter()
383            .map(|v: Vec<f64>| v.into_pyarray(py).unbind())
384            .collect();
385        dict.set_item("voltages", voltages)?;
386        Ok(dict.into_any().unbind())
387    }
388
389    #[staticmethod]
390    fn supported_models() -> Vec<&'static str> {
391        network_runner::supported_models()
392    }
393}
394
395// ── Batch model simulate PyO3 wrapper ────────────────────────────────
396
397/// Run a named neuron model for n_steps with a current trace, returning
398/// voltage trace + spike indices. Entire simulation in Rust.
399#[pyfunction]
400fn py_batch_simulate<'py>(
401    py: Python<'py>,
402    model_name: &str,
403    n_steps: usize,
404    current_trace: PyReadonlyArray1<'py, f64>,
405) -> PyResult<Py<PyAny>> {
406    let mut neuron = network_runner::create_neuron(model_name).map_err(PyValueError::new_err)?;
407    let currents = current_trace.as_slice()?;
408    let steps = n_steps.min(currents.len());
409
410    let mut voltages = vec![0.0f64; steps];
411    let mut spikes: Vec<u64> = Vec::new();
412
413    for t in 0..steps {
414        let fired = neuron.step(currents[t]);
415        voltages[t] = neuron.soma_voltage();
416        if fired != 0 {
417            spikes.push(t as u64);
418        }
419    }
420
421    let d = PyDict::new(py);
422    d.set_item("voltages", voltages.into_pyarray(py))?;
423    d.set_item("spikes", spikes.into_pyarray(py))?;
424    d.set_item("n_steps", steps)?;
425    Ok(d.into_any().unbind())
426}
427
428// ── E-I Network PyO3 wrapper ───────────────────────────────────────
429
430#[pyfunction]
431#[pyo3(signature = (
432    n_exc=80, n_inh=20,
433    w_ee=0.1, w_ei=0.4, w_ie=0.1, w_ii=0.4,
434    p_conn=0.2, ext_rate=5.0,
435    duration=200.0, dt=0.1, seed=42
436))]
437fn py_simulate_ei_network<'py>(
438    py: Python<'py>,
439    n_exc: usize,
440    n_inh: usize,
441    w_ee: f64,
442    w_ei: f64,
443    w_ie: f64,
444    w_ii: f64,
445    p_conn: f64,
446    ext_rate: f64,
447    duration: f64,
448    dt: f64,
449    seed: u64,
450) -> PyResult<Py<PyAny>> {
451    let r = ei_network::simulate_ei(
452        n_exc, n_inh, w_ee, w_ei, w_ie, w_ii, p_conn, ext_rate, duration, dt, seed,
453    );
454    let n_spikes = r.spike_times.len();
455    let d = PyDict::new(py);
456    d.set_item("spike_times", r.spike_times.into_pyarray(py))?;
457    d.set_item(
458        "spike_neurons",
459        r.spike_neurons
460            .iter()
461            .map(|&x| x as i64)
462            .collect::<Vec<_>>()
463            .into_pyarray(py),
464    )?;
465    d.set_item("n_exc", r.n_exc)?;
466    d.set_item("n_inh", r.n_inh)?;
467    d.set_item("n_total", r.n_exc + r.n_inh)?;
468    d.set_item("n_spikes", n_spikes)?;
469    d.set_item("rate_time", r.rate_time.into_pyarray(py))?;
470    d.set_item("exc_rates", r.exc_rates.into_pyarray(py))?;
471    d.set_item("inh_rates", r.inh_rates.into_pyarray(py))?;
472    d.set_item("duration", duration)?;
473    d.set_item("dt", dt)?;
474    d.set_item("mean_exc_rate", r.mean_exc_rate)?;
475    d.set_item("mean_inh_rate", r.mean_inh_rate)?;
476    Ok(d.into_any().unbind())
477}
478
479// ── Wong-Wang 2006 decision unit — batch PyO3 wrapper ─────────────────
480
481/// Simulate a single Wong-Wang unit for `stim1.len()` steps and return
482/// the per-step state traces + firing-rate traces.
483///
484/// Parity contract with `sc_neurocore.neurons.models.wong_wang.WongWangUnit`:
485/// the caller pre-draws `2 * n_steps` N(0, 1) samples so Python-side
486/// `numpy.random` ordering is preserved bit-for-bit.
487///
488/// Returns a dict with keys: `s1`, `s2`, `r1`, `r2`
489/// (1-D float64 arrays of length `n_steps`) and the final scalar
490/// `s1_final`, `s2_final`.
491#[pyfunction]
492#[pyo3(signature = (
493    s1_init, s2_init,
494    tau_s, gamma, j_n, j_cross, i_0, sigma, dt,
495    stim1, stim2, xi,
496))]
497#[allow(clippy::too_many_arguments)]
498fn py_wong_wang_simulate<'py>(
499    py: Python<'py>,
500    s1_init: f64,
501    s2_init: f64,
502    tau_s: f64,
503    gamma: f64,
504    j_n: f64,
505    j_cross: f64,
506    i_0: f64,
507    sigma: f64,
508    dt: f64,
509    stim1: PyReadonlyArray1<'py, f64>,
510    stim2: PyReadonlyArray1<'py, f64>,
511    xi: PyReadonlyArray1<'py, f64>,
512) -> PyResult<Py<PyAny>> {
513    let stim1 = stim1.as_slice()?;
514    let stim2 = stim2.as_slice()?;
515    let xi = xi.as_slice()?;
516    let n = stim1.len();
517    if stim2.len() != n {
518        return Err(PyValueError::new_err(format!(
519            "stim1 and stim2 length mismatch: {n} vs {}",
520            stim2.len()
521        )));
522    }
523    if xi.len() != 2 * n {
524        return Err(PyValueError::new_err(format!(
525            "xi length must be 2 * n_steps ({}): got {}",
526            2 * n,
527            xi.len()
528        )));
529    }
530    let mut s1o = vec![0.0_f64; n];
531    let mut s2o = vec![0.0_f64; n];
532    let mut r1o = vec![0.0_f64; n];
533    let mut r2o = vec![0.0_f64; n];
534    let (s1_final, s2_final) = wong_wang::simulate(
535        s1_init, s2_init, tau_s, gamma, j_n, j_cross, i_0, sigma, dt, stim1, stim2, xi, &mut s1o,
536        &mut s2o, &mut r1o, &mut r2o,
537    );
538    let d = PyDict::new(py);
539    d.set_item("s1", s1o.into_pyarray(py))?;
540    d.set_item("s2", s2o.into_pyarray(py))?;
541    d.set_item("r1", r1o.into_pyarray(py))?;
542    d.set_item("r2", r2o.into_pyarray(py))?;
543    d.set_item("s1_final", s1_final)?;
544    d.set_item("s2_final", s2_final)?;
545    Ok(d.into_any().unbind())
546}
547
548// ── DCLS-max Q8.8 tent kernel — batch PyO3 wrapper ───────────────────
549
550/// Batched DCLS-max triangular (tent) contraction in bit-true Q8.8 arithmetic.
551///
552/// Parity contract with `sc_neurocore.scpn.dcls_tent_kernel`: this Rust path,
553/// the Mojo, Julia and Go backends, and the Python floor all return
554/// bit-identical arrays because the kernel is exact integer arithmetic.
555///
556/// `spikes` and `weights_q88` are row-major `n_channels * n_taps`; `centres_q88`
557/// and `sigmas_q88` carry one learnable `(centre, sigma)` per output channel.
558///
559/// Returns a dict with keys `outputs_q88` (int16), `accumulators_q16_16`
560/// (int32), `overflow` (bool), `active_tap_counts` (int64) and `max_gates_q88`
561/// (int16), each a 1-D array of length `n_channels`.
562#[pyfunction]
563#[pyo3(signature = (spikes, weights_q88, centres_q88, sigmas_q88, n_taps))]
564fn py_dcls_max_forward_batch_q88<'py>(
565    py: Python<'py>,
566    spikes: PyReadonlyArray1<'py, u8>,
567    weights_q88: PyReadonlyArray1<'py, i16>,
568    centres_q88: PyReadonlyArray1<'py, i16>,
569    sigmas_q88: PyReadonlyArray1<'py, i16>,
570    n_taps: usize,
571) -> PyResult<Py<PyAny>> {
572    let result = scpn::dcls_max_forward_batch_q88(
573        spikes.as_slice()?,
574        weights_q88.as_slice()?,
575        centres_q88.as_slice()?,
576        sigmas_q88.as_slice()?,
577        n_taps,
578    )
579    .map_err(|e| PyValueError::new_err(e.to_string()))?;
580    let active_tap_counts: Vec<i64> = result.active_tap_counts.iter().map(|&c| c as i64).collect();
581    let d = PyDict::new(py);
582    d.set_item("outputs_q88", result.outputs_q88.into_pyarray(py))?;
583    d.set_item(
584        "accumulators_q16_16",
585        result.accumulators_q16_16.into_pyarray(py),
586    )?;
587    d.set_item("overflow", result.overflow.into_pyarray(py))?;
588    d.set_item("active_tap_counts", active_tap_counts.into_pyarray(py))?;
589    d.set_item("max_gates_q88", result.max_gates_q88.into_pyarray(py))?;
590    Ok(d.into_any().unbind())
591}
592
593// ── SC inference over pre-packed weights — PyO3 wrapper ──────────────
594
595/// Stochastic forward pass over caller-owned packed weight bitstreams.
596///
597/// Parity contract with `sc_neurocore.accel.sc_forward`: this Rust path and the
598/// NumPy fallback return bit-identical results for a fixed seed because the input
599/// encoder is the deterministic 16-bit LFSR comparator.
600///
601/// `weights_packed` is row-major `n_out * n_in * n_words` (`n_words =
602/// ceil(length / 64)`); `input_probs` is `n_in` float64 in `[0, 1]`. Returns an
603/// `n_out` float64 array, the AND-then-popcount estimate of
604/// `weights @ input_probs` divided by `length`.
605#[pyfunction]
606#[pyo3(signature = (weights_packed, n_out, n_in, n_words, input_probs, length, seed))]
607#[allow(clippy::too_many_arguments)]
608fn py_sc_forward_packed<'py>(
609    py: Python<'py>,
610    weights_packed: PyReadonlyArray1<'py, u64>,
611    n_out: usize,
612    n_in: usize,
613    n_words: usize,
614    input_probs: PyReadonlyArray1<'py, f64>,
615    length: usize,
616    seed: u64,
617) -> PyResult<Bound<'py, PyArray1<f64>>> {
618    let outputs = sc_inference::sc_forward_packed(
619        weights_packed.as_slice()?,
620        n_out,
621        n_in,
622        n_words,
623        input_probs.as_slice()?,
624        length,
625        seed,
626    )
627    .map_err(|e| PyValueError::new_err(e.to_string()))?;
628    Ok(outputs.into_pyarray(py))
629}
630
631// ── ADC-to-spike decimating rate-code — per-window PyO3 wrapper ───────
632
633/// Encode raw ADC samples into per-window spike rate codes.
634///
635/// Parity contract with `sc_neurocore.sensors.adc_to_spike_kernel`: this Rust
636/// path and the Julia, Go, Mojo and Python backends return bit-identical arrays
637/// because the per-window quantise/average/rate-code arithmetic is exact integer.
638///
639/// `signed_input` is `0` for offset-binary or `1` for two's-complement ADC
640/// samples. Returns a dict with `window_values_q` (int32), `spike_counts` (int32)
641/// and `polarities` (bool), each of length `samples.len() / decimation`.
642#[pyfunction]
643#[pyo3(signature = (
644    samples, adc_width, q_int, q_frac, decimation, signed_input, threshold_q,
645))]
646#[allow(clippy::too_many_arguments)]
647fn py_adc_to_spike_windows<'py>(
648    py: Python<'py>,
649    samples: PyReadonlyArray1<'py, i64>,
650    adc_width: u32,
651    q_int: u32,
652    q_frac: u32,
653    decimation: u32,
654    signed_input: i64,
655    threshold_q: i64,
656) -> PyResult<Py<PyAny>> {
657    let result = adc_to_spike::adc_to_spike_windows(
658        samples.as_slice()?,
659        adc_width,
660        q_int,
661        q_frac,
662        decimation,
663        signed_input != 0,
664        threshold_q,
665    )
666    .map_err(|e| PyValueError::new_err(e.to_string()))?;
667    let d = PyDict::new(py);
668    d.set_item("window_values_q", result.window_values_q.into_pyarray(py))?;
669    d.set_item("spike_counts", result.spike_counts.into_pyarray(py))?;
670    d.set_item("polarities", result.polarities.into_pyarray(py))?;
671    Ok(d.into_any().unbind())
672}
673
674// ── Mixed-precision Q8.8 × Q16.16 dense MAC — batch PyO3 wrapper ──────
675
676/// Batched integer mixed-precision Q8.8 × Q16.16 dense MAC.
677///
678/// Parity contract with `sc_neurocore.compiler.mixed_dense_kernel`: this Rust
679/// path and the Julia, Go, Mojo and Python backends return bit-identical arrays
680/// because the integer branch (divisor equal to the Q8.8 weight scale) is exact.
681///
682/// `weights_q88` is row-major `n_outputs * n_inputs`; `inputs_q1616` is row-major
683/// `n_batch * n_inputs`. Returns a dict with `outputs_q1616` (int32), `overflow`
684/// (bool) and `underflow` (bool), each a 1-D array of length `n_batch * n_outputs`.
685#[pyfunction]
686#[pyo3(signature = (weights_q88, inputs_q1616, n_outputs, n_inputs))]
687fn py_mixed_dense_forward_batch_q88_q1616<'py>(
688    py: Python<'py>,
689    weights_q88: PyReadonlyArray1<'py, i16>,
690    inputs_q1616: PyReadonlyArray1<'py, i32>,
691    n_outputs: usize,
692    n_inputs: usize,
693) -> PyResult<Py<PyAny>> {
694    let result = crate::ir::qformat::mixed_dense_forward_batch_q88_q1616(
695        weights_q88.as_slice()?,
696        inputs_q1616.as_slice()?,
697        n_outputs,
698        n_inputs,
699    )
700    .map_err(|e| PyValueError::new_err(e.to_string()))?;
701    let d = PyDict::new(py);
702    d.set_item("outputs_q1616", result.outputs_q1616.into_pyarray(py))?;
703    d.set_item("overflow", result.overflow.into_pyarray(py))?;
704    d.set_item("underflow", result.underflow.into_pyarray(py))?;
705    Ok(d.into_any().unbind())
706}
707
708// ── Wilson-Cowan 1972 E/I rate model — batch PyO3 wrapper ─────────────
709
710/// Simulate a single Wilson-Cowan E/I unit for `ext_input.len()` steps
711/// and return per-step `e`, `i` traces plus final scalars.
712///
713/// Returns a dict with keys: `e`, `i` (1-D float64 arrays of length
714/// `n_steps`) + the final scalars `e_final`, `i_final`.
715#[pyfunction]
716#[pyo3(signature = (
717    e_init, i_init,
718    w_ee, w_ei, w_ie, w_ii,
719    tau_e, tau_i,
720    a, theta, dt,
721    ext_input,
722))]
723#[allow(clippy::too_many_arguments)]
724fn py_wilson_cowan_simulate<'py>(
725    py: Python<'py>,
726    e_init: f64,
727    i_init: f64,
728    w_ee: f64,
729    w_ei: f64,
730    w_ie: f64,
731    w_ii: f64,
732    tau_e: f64,
733    tau_i: f64,
734    a: f64,
735    theta: f64,
736    dt: f64,
737    ext_input: PyReadonlyArray1<'py, f64>,
738) -> PyResult<Py<PyAny>> {
739    let ext = ext_input.as_slice()?;
740    let n = ext.len();
741    let mut e_out = vec![0.0_f64; n];
742    let mut i_out = vec![0.0_f64; n];
743    let (e_final, i_final) = wilson_cowan::simulate(
744        e_init, i_init, w_ee, w_ei, w_ie, w_ii, tau_e, tau_i, a, theta, dt, ext, &mut e_out,
745        &mut i_out,
746    );
747    let d = PyDict::new(py);
748    d.set_item("e", e_out.into_pyarray(py))?;
749    d.set_item("i", i_out.into_pyarray(py))?;
750    d.set_item("e_final", e_final)?;
751    d.set_item("i_final", i_final)?;
752    Ok(d.into_any().unbind())
753}
754
755/// SC-NeuroCore ─ High-Performance Rust Engine
756
757#[pymodule]
758fn sc_neurocore_engine(m: &Bound<'_, PyModule>) -> PyResult<()> {
759    m.add("__version__", env!("CARGO_PKG_VERSION"))?;
760    m.add_function(wrap_pyfunction!(simd_tier, m)?)?;
761    m.add_function(wrap_pyfunction!(set_num_threads, m)?)?;
762    m.add_function(wrap_pyfunction!(pack_bitstream, m)?)?;
763    m.add_function(wrap_pyfunction!(unpack_bitstream, m)?)?;
764    m.add_function(wrap_pyfunction!(popcount, m)?)?;
765    m.add_function(wrap_pyfunction!(pack_bitstream_numpy, m)?)?;
766    m.add_function(wrap_pyfunction!(popcount_numpy, m)?)?;
767    m.add_function(wrap_pyfunction!(unpack_bitstream_numpy, m)?)?;
768    m.add_function(wrap_pyfunction!(batch_lif_run, m)?)?;
769    m.add_function(wrap_pyfunction!(batch_lif_run_multi, m)?)?;
770    m.add_function(wrap_pyfunction!(batch_lif_run_varying, m)?)?;
771    m.add_function(wrap_pyfunction!(batch_encode, m)?)?;
772    m.add_function(wrap_pyfunction!(batch_encode_numpy, m)?)?;
773    m.add_function(wrap_pyfunction!(py_wong_wang_simulate, m)?)?;
774    m.add_function(wrap_pyfunction!(py_dcls_max_forward_batch_q88, m)?)?;
775    m.add_function(wrap_pyfunction!(py_mixed_dense_forward_batch_q88_q1616, m)?)?;
776    m.add_function(wrap_pyfunction!(py_adc_to_spike_windows, m)?)?;
777    m.add_function(wrap_pyfunction!(py_sc_forward_packed, m)?)?;
778    m.add_function(wrap_pyfunction!(py_wilson_cowan_simulate, m)?)?;
779    m.add_class::<Lfsr16>()?;
780    m.add_class::<BitstreamEncoder>()?;
781    m.add_class::<FixedPointLif>()?;
782    m.add_class::<DenseLayer>()?;
783    #[cfg(feature = "gpu")]
784    m.add_class::<gpu::PyGpuDenseLayer>()?;
785    m.add_class::<StdpSynapse>()?;
786    m.add_class::<PySurrogateLif>()?;
787    m.add_class::<PyDifferentiableDenseLayer>()?;
788    m.add_class::<PyStochasticAttention>()?;
789    m.add_class::<PyStochasticGraphLayer>()?;
790    m.add_class::<PyKuramotoSolver>()?;
791    m.add_class::<PySCPNMetrics>()?;
792    m.add_class::<PyBitStreamTensor>()?;
793    m.add_class::<PyBrunelNetwork>()?;
794    m.add_class::<PyIzhikevich>()?;
795    m.add_class::<PyBitstreamAverager>()?;
796    m.add_class::<PyScGraph>()?;
797    m.add_class::<PyScGraphBuilder>()?;
798    m.add_function(wrap_pyfunction!(ir_verify, m)?)?;
799    m.add_function(wrap_pyfunction!(ir_print, m)?)?;
800    m.add_function(wrap_pyfunction!(ir_parse, m)?)?;
801    m.add_function(wrap_pyfunction!(ir_emit_sv, m)?)?;
802    m.add_class::<PyAdExNeuron>()?;
803    m.add_class::<PyExpIFNeuron>()?;
804    m.add_class::<PyLapicqueNeuron>()?;
805    pyo3_neurons::register_neuron_classes(m)?;
806    m.add_class::<PyNetworkRunner>()?;
807    m.add_class::<supervisor::PySpikingControllerPool>()?;
808    m.add_function(wrap_pyfunction!(py_simulate_ei_network, m)?)?;
809    m.add_function(wrap_pyfunction!(py_batch_simulate, m)?)?;
810    m.add_function(wrap_pyfunction!(rk4_neurons::py_rk4_neuron_simulate, m)?)?;
811    m.add_function(wrap_pyfunction!(py_cordiv, m)?)?;
812    m.add_function(wrap_pyfunction!(py_adaptive_length, m)?)?;
813    m.add_function(wrap_pyfunction!(py_prediction_error, m)?)?;
814    m.add_function(wrap_pyfunction!(py_predict_xor_ema, m)?)?;
815    m.add_function(wrap_pyfunction!(py_predict_xor_lfsr, m)?)?;
816    m.add_function(wrap_pyfunction!(py_recover_xor_ema, m)?)?;
817    m.add_function(wrap_pyfunction!(py_recover_xor_lfsr, m)?)?;
818    m.add_function(wrap_pyfunction!(py_phi_star, m)?)?;
819    m.add_class::<PyCorticalColumn>()?;
820    m.add_class::<PyRallDendrite>()?;
821    // Analysis functions (P0-A: spike_stats)
822    m.add_function(wrap_pyfunction!(py_spike_times, m)?)?;
823    m.add_function(wrap_pyfunction!(py_isi, m)?)?;
824    m.add_function(wrap_pyfunction!(py_firing_rate, m)?)?;
825    m.add_function(wrap_pyfunction!(py_spike_count, m)?)?;
826    m.add_function(wrap_pyfunction!(py_bin_spike_train, m)?)?;
827    m.add_function(wrap_pyfunction!(py_instantaneous_rate, m)?)?;
828    m.add_function(wrap_pyfunction!(py_psth, m)?)?;
829    m.add_function(wrap_pyfunction!(py_cv_isi, m)?)?;
830    m.add_function(wrap_pyfunction!(py_cv2, m)?)?;
831    m.add_function(wrap_pyfunction!(py_local_variation, m)?)?;
832    m.add_function(wrap_pyfunction!(py_fano_factor, m)?)?;
833    m.add_function(wrap_pyfunction!(py_lempel_ziv_complexity, m)?)?;
834    m.add_function(wrap_pyfunction!(py_permutation_entropy, m)?)?;
835    m.add_function(wrap_pyfunction!(py_hurst_exponent, m)?)?;
836    m.add_function(wrap_pyfunction!(py_approximate_entropy, m)?)?;
837    m.add_function(wrap_pyfunction!(py_sample_entropy, m)?)?;
838    // correlation
839    m.add_function(wrap_pyfunction!(py_cross_correlation, m)?)?;
840    m.add_function(wrap_pyfunction!(py_pairwise_correlation, m)?)?;
841    m.add_function(wrap_pyfunction!(py_event_synchronization, m)?)?;
842    m.add_function(wrap_pyfunction!(py_spike_train_coherence, m)?)?;
843    m.add_function(wrap_pyfunction!(py_spike_time_tiling_coefficient, m)?)?;
844    m.add_function(wrap_pyfunction!(py_covariance_matrix, m)?)?;
845    m.add_function(wrap_pyfunction!(py_autocorrelation_time, m)?)?;
846    m.add_function(wrap_pyfunction!(py_noise_correlation, m)?)?;
847    m.add_function(wrap_pyfunction!(py_signal_correlation, m)?)?;
848    m.add_function(wrap_pyfunction!(py_spike_count_covariance, m)?)?;
849    m.add_function(wrap_pyfunction!(py_joint_psth, m)?)?;
850    m.add_function(wrap_pyfunction!(py_coincidence_index, m)?)?;
851    // distance
852    m.add_function(wrap_pyfunction!(py_van_rossum_distance, m)?)?;
853    m.add_function(wrap_pyfunction!(py_victor_purpura_distance, m)?)?;
854    m.add_function(wrap_pyfunction!(py_isi_distance, m)?)?;
855    m.add_function(wrap_pyfunction!(py_spike_distance, m)?)?;
856    m.add_function(wrap_pyfunction!(py_spike_sync, m)?)?;
857    m.add_function(wrap_pyfunction!(py_spike_sync_profile, m)?)?;
858    m.add_function(wrap_pyfunction!(py_spike_profile, m)?)?;
859    m.add_function(wrap_pyfunction!(py_isi_profile, m)?)?;
860    m.add_function(wrap_pyfunction!(py_adaptive_spike_distance, m)?)?;
861    m.add_function(wrap_pyfunction!(py_schreiber_similarity, m)?)?;
862    m.add_function(wrap_pyfunction!(py_hunter_milton_similarity, m)?)?;
863    m.add_function(wrap_pyfunction!(py_earth_movers_distance, m)?)?;
864    m.add_function(wrap_pyfunction!(py_multi_neuron_victor_purpura, m)?)?;
865    m.add_function(wrap_pyfunction!(py_spike_distance_matrix, m)?)?;
866    // information
867    m.add_function(wrap_pyfunction!(py_mutual_information, m)?)?;
868    m.add_function(wrap_pyfunction!(py_transfer_entropy, m)?)?;
869    m.add_function(wrap_pyfunction!(py_spike_train_entropy, m)?)?;
870    m.add_function(wrap_pyfunction!(py_noise_entropy, m)?)?;
871    m.add_function(wrap_pyfunction!(py_stimulus_specific_information, m)?)?;
872    m.add_function(wrap_pyfunction!(py_kozachenko_leonenko_mi, m)?)?;
873    // causality
874    m.add_function(wrap_pyfunction!(py_pairwise_granger_causality, m)?)?;
875    m.add_function(wrap_pyfunction!(py_conditional_granger_causality, m)?)?;
876    m.add_function(wrap_pyfunction!(py_spectral_granger_causality, m)?)?;
877    m.add_function(wrap_pyfunction!(py_partial_directed_coherence, m)?)?;
878    m.add_function(wrap_pyfunction!(py_directed_transfer_function, m)?)?;
879    // decoding
880    m.add_function(wrap_pyfunction!(py_population_vector_decode, m)?)?;
881    m.add_function(wrap_pyfunction!(py_bayesian_decode, m)?)?;
882    m.add_function(wrap_pyfunction!(py_maximum_likelihood_decode, m)?)?;
883    m.add_function(wrap_pyfunction!(py_linear_discriminant_decode, m)?)?;
884    m.add_function(wrap_pyfunction!(py_naive_bayes_decode, m)?)?;
885    // neural_decoders (P1)
886    m.add_function(wrap_pyfunction!(py_tokenise_spikes, m)?)?;
887    m.add_function(wrap_pyfunction!(py_sinusoidal_position_encode, m)?)?;
888    m.add_function(wrap_pyfunction!(py_scaled_dot_product_attention, m)?)?;
889    m.add_function(wrap_pyfunction!(py_gaussian_attention, m)?)?;
890    m.add_function(wrap_pyfunction!(py_infonce_loss, m)?)?;
891    // network
892    m.add_function(wrap_pyfunction!(py_functional_connectivity, m)?)?;
893    m.add_function(wrap_pyfunction!(py_unitary_events, m)?)?;
894    m.add_function(wrap_pyfunction!(py_cell_assembly_detection, m)?)?;
895    m.add_function(wrap_pyfunction!(py_synfire_chain_detection, m)?)?;
896    // surrogates
897    m.add_function(wrap_pyfunction!(py_surrogate_isi_shuffle, m)?)?;
898    m.add_function(wrap_pyfunction!(py_surrogate_dither, m)?)?;
899    m.add_function(wrap_pyfunction!(py_homogeneous_poisson, m)?)?;
900    m.add_function(wrap_pyfunction!(py_gamma_process, m)?)?;
901    m.add_function(wrap_pyfunction!(py_compound_poisson_process, m)?)?;
902    m.add_function(wrap_pyfunction!(py_surrogate_joint_isi, m)?)?;
903    m.add_function(wrap_pyfunction!(py_surrogate_bin_shuffling, m)?)?;
904    m.add_function(wrap_pyfunction!(py_surrogate_spike_train_shifting, m)?)?;
905    // temporal
906    m.add_function(wrap_pyfunction!(py_burst_detection, m)?)?;
907    m.add_function(wrap_pyfunction!(py_first_spike_latency, m)?)?;
908    m.add_function(wrap_pyfunction!(py_response_onset, m)?)?;
909    m.add_function(wrap_pyfunction!(py_change_point_detection, m)?)?;
910    // patterns
911    m.add_function(wrap_pyfunction!(py_spike_directionality, m)?)?;
912    m.add_function(wrap_pyfunction!(py_spike_train_order, m)?)?;
913    m.add_function(wrap_pyfunction!(py_cubic_higher_order, m)?)?;
914    // spectral
915    m.add_function(wrap_pyfunction!(py_power_spectrum, m)?)?;
916    // waveform
917    m.add_function(wrap_pyfunction!(py_waveform_width, m)?)?;
918    m.add_function(wrap_pyfunction!(py_waveform_amplitude, m)?)?;
919    m.add_function(wrap_pyfunction!(py_waveform_repolarization_slope, m)?)?;
920    m.add_function(wrap_pyfunction!(py_waveform_recovery_slope, m)?)?;
921    m.add_function(wrap_pyfunction!(py_waveform_halfwidth, m)?)?;
922    m.add_function(wrap_pyfunction!(py_waveform_pt_ratio, m)?)?;
923    // point_process
924    m.add_function(wrap_pyfunction!(py_conditional_intensity, m)?)?;
925    m.add_function(wrap_pyfunction!(py_isi_hazard_function, m)?)?;
926    m.add_function(wrap_pyfunction!(py_isi_survivor_function, m)?)?;
927    m.add_function(wrap_pyfunction!(py_renewal_density, m)?)?;
928    // stimulus
929    m.add_function(wrap_pyfunction!(py_spike_triggered_average, m)?)?;
930    m.add_function(wrap_pyfunction!(py_spike_triggered_covariance, m)?)?;
931    m.add_function(wrap_pyfunction!(py_spatial_information, m)?)?;
932    m.add_function(wrap_pyfunction!(py_place_field_detection, m)?)?;
933    m.add_function(wrap_pyfunction!(py_tuning_curve, m)?)?;
934    // lfp
935    m.add_function(wrap_pyfunction!(py_phase_locking_value, m)?)?;
936    m.add_function(wrap_pyfunction!(py_spike_field_coherence, m)?)?;
937    m.add_function(wrap_pyfunction!(py_spike_phase_histogram, m)?)?;
938    // sorting_quality
939    m.add_function(wrap_pyfunction!(py_isolation_distance, m)?)?;
940    m.add_function(wrap_pyfunction!(py_l_ratio, m)?)?;
941    m.add_function(wrap_pyfunction!(py_silhouette_score, m)?)?;
942    m.add_function(wrap_pyfunction!(py_d_prime, m)?)?;
943    m.add_function(wrap_pyfunction!(py_isi_violation_rate, m)?)?;
944    m.add_function(wrap_pyfunction!(py_presence_ratio, m)?)?;
945    m.add_function(wrap_pyfunction!(py_amplitude_cutoff, m)?)?;
946    m.add_function(wrap_pyfunction!(py_snr, m)?)?;
947    m.add_function(wrap_pyfunction!(py_nn_hit_rate, m)?)?;
948    m.add_function(wrap_pyfunction!(py_drift_metric, m)?)?;
949    // dimensionality
950    m.add_function(wrap_pyfunction!(py_spike_train_pca, m)?)?;
951    m.add_function(wrap_pyfunction!(py_demixed_pca, m)?)?;
952    m.add_function(wrap_pyfunction!(py_factor_analysis, m)?)?;
953    m.add_function(wrap_pyfunction!(py_pca_components, m)?)?;
954    m.add_function(wrap_pyfunction!(py_demixed_components, m)?)?;
955    m.add_function(wrap_pyfunction!(py_factor_loadings, m)?)?;
956    // gpfa
957    m.add_function(wrap_pyfunction!(py_gpfa, m)?)?;
958    m.add_function(wrap_pyfunction!(py_gpfa_em, m)?)?;
959    m.add_function(wrap_pyfunction!(py_gpfa_transform, m)?)?;
960    // spade
961    m.add_function(wrap_pyfunction!(py_spade_detect, m)?)?;
962    // dna
963    m.add_function(wrap_pyfunction!(py_dna_design_sequence, m)?)?;
964    m.add_function(wrap_pyfunction!(py_dna_design_orthogonal_set, m)?)?;
965    m.add_function(wrap_pyfunction!(py_dna_check_cross_hybridization, m)?)?;
966    m.add_function(wrap_pyfunction!(py_dna_simulate_kinetics, m)?)?;
967    m.add_function(wrap_pyfunction!(py_dna_detect_hairpins, m)?)?;
968    // Quantum annealing acceleration
969    m.add_function(wrap_pyfunction!(py_qa_ising_energy, m)?)?;
970    m.add_function(wrap_pyfunction!(py_qa_batch_ising_energy, m)?)?;
971    m.add_function(wrap_pyfunction!(py_qa_simulated_annealing, m)?)?;
972    m.add_function(wrap_pyfunction!(py_qa_gauge_transform, m)?)?;
973    m.add_function(wrap_pyfunction!(py_qa_generate_gauges, m)?)?;
974    m.add_function(wrap_pyfunction!(py_qa_greedy_partition, m)?)?;
975    // Photonic NoC acceleration
976    m.add_function(wrap_pyfunction!(py_ph_route_waveguides, m)?)?;
977    m.add_function(wrap_pyfunction!(py_ph_mzi_transfer_matrix, m)?)?;
978    m.add_function(wrap_pyfunction!(py_ph_cascade_mzi, m)?)?;
979    m.add_function(wrap_pyfunction!(py_ph_analyze_crosstalk, m)?)?;
980    m.add_function(wrap_pyfunction!(py_ph_analyze_power_budget, m)?)?;
981    m.add_function(wrap_pyfunction!(py_ph_analyze_crosstalk_bank, m)?)?;
982    m.add_function(wrap_pyfunction!(py_ph_analyze_crosstalk_pairs, m)?)?;
983    // SC-Optimizer acceleration
984    m.add_function(wrap_pyfunction!(py_opt_sa_search, m)?)?;
985    m.add_function(wrap_pyfunction!(py_opt_extract_pareto, m)?)?;
986    // Evolutionary substrate acceleration
987    m.add_function(wrap_pyfunction!(py_evo_batch_mutate, m)?)?;
988    m.add_function(wrap_pyfunction!(py_evo_batch_fitness, m)?)?;
989    m.add_function(wrap_pyfunction!(py_evo_batch_crossover, m)?)?;
990    m.add_function(wrap_pyfunction!(py_evo_diversity, m)?)?;
991    m.add_function(wrap_pyfunction!(py_evo_novelty, m)?)?;
992    m.add_function(wrap_pyfunction!(py_evo_tournament, m)?)?;
993    // LGSSM Kalman filter (predictive_model)
994    m.add_function(wrap_pyfunction!(py_lgssm_kalman_filter, m)?)?;
995    m.add_function(wrap_pyfunction!(py_ollivier_ricci_curvature, m)?)?;
996    m.add_function(wrap_pyfunction!(py_cazelles_map_simulate, m)?)?;
997    m.add_function(wrap_pyfunction!(py_courage_nekorkin_map_simulate, m)?)?;
998    m.add_function(wrap_pyfunction!(py_mckean_simulate, m)?)?;
999    m.add_function(wrap_pyfunction!(py_wilson_hr_simulate, m)?)?;
1000    m.add_function(wrap_pyfunction!(py_pernarowski_simulate, m)?)?;
1001    m.add_function(wrap_pyfunction!(py_terman_wang_simulate, m)?)?;
1002    m.add_function(wrap_pyfunction!(py_mihalas_niebur_simulate, m)?)?;
1003    m.add_function(wrap_pyfunction!(py_glif_simulate, m)?)?;
1004    m.add_function(wrap_pyfunction!(py_rulkov_map_simulate, m)?)?;
1005    m.add_function(wrap_pyfunction!(py_ibarz_tanaka_map_simulate, m)?)?;
1006    m.add_function(wrap_pyfunction!(py_medvedev_map_simulate, m)?)?;
1007    m.add_function(wrap_pyfunction!(py_ermentrout_kopell_map_simulate, m)?)?;
1008    m.add_function(wrap_pyfunction!(py_fitzhugh_nagumo_simulate, m)?)?;
1009    m.add_function(wrap_pyfunction!(py_hindmarsh_rose_simulate, m)?)?;
1010    m.add_function(wrap_pyfunction!(py_fitzhugh_rinzel_simulate, m)?)?;
1011    m.add_function(wrap_pyfunction!(py_izhikevich2007_simulate, m)?)?;
1012    // Byte-level fault injection (parity with FaultInjector.inject)
1013    m.add_function(wrap_pyfunction!(py_inject_bitflip_u8, m)?)?;
1014    m.add_function(wrap_pyfunction!(py_inject_stuck_at_0_u8, m)?)?;
1015    m.add_function(wrap_pyfunction!(py_inject_stuck_at_1_u8, m)?)?;
1016    m.add_function(wrap_pyfunction!(py_inject_dropout_u8, m)?)?;
1017    m.add_function(wrap_pyfunction!(py_inject_gaussian_u8, m)?)?;
1018    // Hierarchical partitioner KL refine
1019    m.add_function(wrap_pyfunction!(py_kl_refine, m)?)?;
1020    // PINGCircuit per-step kernel
1021    m.add_function(wrap_pyfunction!(py_ping_step, m)?)?;
1022    // CorticalColumn block-CSR per-row-parallel spmv
1023    m.add_function(wrap_pyfunction!(py_parallel_csr_spmv_add, m)?)?;
1024    m.add_function(wrap_pyfunction!(py_parallel_csr_multi_spmv_add, m)?)?;
1025    Ok(())
1026}
1027
1028// ── CORDIV + adaptive length ─────────────────────────────────────────
1029
1030#[pyfunction]
1031fn py_cordiv(
1032    py: Python<'_>,
1033    numerator: PyReadonlyArray1<'_, u8>,
1034    denominator: PyReadonlyArray1<'_, u8>,
1035) -> PyResult<Py<PyArray1<u8>>> {
1036    let num = numerator.as_slice()?;
1037    let den = denominator.as_slice()?;
1038    let result = cordiv::cordiv(num, den);
1039    Ok(result.into_pyarray(py).into())
1040}
1041
1042#[pyfunction]
1043fn py_adaptive_length(epsilon: f64, confidence: f64) -> usize {
1044    cordiv::adaptive_length_hoeffding(epsilon, confidence)
1045}
1046
1047// ── Predictive coding ────────────────────────────────────────────────
1048
1049#[pyfunction]
1050fn py_prediction_error(
1051    _py: Python<'_>,
1052    predicted: PyReadonlyArray1<'_, u64>,
1053    actual: PyReadonlyArray1<'_, u64>,
1054    length: usize,
1055) -> PyResult<f64> {
1056    let pred = predicted
1057        .as_slice()
1058        .map_err(|e| PyValueError::new_err(format!("predicted array must be contiguous: {e}")))?;
1059    let act = actual
1060        .as_slice()
1061        .map_err(|e| PyValueError::new_err(format!("actual array must be contiguous: {e}")))?;
1062    Ok(predictive_coding::prediction_error_packed(
1063        pred, act, length,
1064    ))
1065}
1066
1067// ── Spike codec predictors (EMA + LFSR) ─────────────────────────────
1068
1069#[pyfunction]
1070fn py_predict_xor_ema(
1071    _py: Python<'_>,
1072    spikes: PyReadonlyArray1<'_, i8>,
1073    n_channels: usize,
1074    alpha: f64,
1075    threshold: f64,
1076) -> PyResult<(Py<PyArray1<i8>>, usize)> {
1077    let data = spikes
1078        .as_slice()
1079        .map_err(|e| PyValueError::new_err(format!("spikes must be contiguous: {e}")))?;
1080    let (errors, correct) =
1081        predictive_coding::predict_and_xor_ema(data, n_channels, alpha, threshold);
1082    Ok((PyArray1::from_vec(_py, errors).into(), correct))
1083}
1084
1085#[pyfunction]
1086fn py_recover_xor_ema(
1087    _py: Python<'_>,
1088    errors: PyReadonlyArray1<'_, i8>,
1089    n_channels: usize,
1090    alpha: f64,
1091    threshold: f64,
1092) -> PyResult<Py<PyArray1<i8>>> {
1093    let data = errors
1094        .as_slice()
1095        .map_err(|e| PyValueError::new_err(format!("errors must be contiguous: {e}")))?;
1096    let spikes = predictive_coding::xor_and_recover_ema(data, n_channels, alpha, threshold);
1097    Ok(PyArray1::from_vec(_py, spikes).into())
1098}
1099
1100#[pyfunction]
1101fn py_predict_xor_lfsr(
1102    _py: Python<'_>,
1103    spikes: PyReadonlyArray1<'_, i8>,
1104    n_channels: usize,
1105    alpha_q8: i32,
1106    seed: u16,
1107) -> PyResult<(Py<PyArray1<i8>>, usize)> {
1108    let data = spikes
1109        .as_slice()
1110        .map_err(|e| PyValueError::new_err(format!("spikes must be contiguous: {e}")))?;
1111    let (errors, correct) =
1112        predictive_coding::predict_and_xor_lfsr(data, n_channels, alpha_q8, seed);
1113    Ok((PyArray1::from_vec(_py, errors).into(), correct))
1114}
1115
1116#[pyfunction]
1117fn py_recover_xor_lfsr(
1118    _py: Python<'_>,
1119    errors: PyReadonlyArray1<'_, i8>,
1120    n_channels: usize,
1121    alpha_q8: i32,
1122    seed: u16,
1123) -> PyResult<Py<PyArray1<i8>>> {
1124    let data = errors
1125        .as_slice()
1126        .map_err(|e| PyValueError::new_err(format!("errors must be contiguous: {e}")))?;
1127    let spikes = predictive_coding::xor_and_recover_lfsr(data, n_channels, alpha_q8, seed);
1128    Ok(PyArray1::from_vec(_py, spikes).into())
1129}
1130
1131// ── Phi* ─────────────────────────────────────────────────────────────
1132
1133#[pyfunction]
1134fn py_phi_star(_py: Python<'_>, data: PyReadonlyArray2<'_, f64>, tau: usize) -> PyResult<f64> {
1135    if !data.is_c_contiguous() {
1136        return Err(PyValueError::new_err(
1137            "py_phi_star requires C-contiguous array input",
1138        ));
1139    }
1140    let shape = data.shape();
1141    let n_channels = shape[0];
1142    let n_timesteps = shape[1];
1143    let flat = data.as_slice().map_err(|e| {
1144        PyValueError::new_err(format!("py_phi_star requires C-contiguous array: {e}"))
1145    })?;
1146    let channels: Vec<Vec<f64>> = (0..n_channels)
1147        .map(|i| flat[i * n_timesteps..(i + 1) * n_timesteps].to_vec())
1148        .collect();
1149    Ok(phi::phi_star(&channels, tau))
1150}
1151
1152// ── Cortical column ──────────────────────────────────────────────────
1153
1154#[pyclass(
1155    name = "CorticalColumnRust",
1156    module = "sc_neurocore_engine.sc_neurocore_engine"
1157)]
1158pub struct PyCorticalColumn {
1159    inner: cortical_column::CorticalColumnRust,
1160}
1161
1162#[pymethods]
1163impl PyCorticalColumn {
1164    #[new]
1165    fn new(n: usize, tau: f64, dt: f64, threshold: f64, w_exc: f64, w_inh: f64, seed: u64) -> Self {
1166        Self {
1167            inner: cortical_column::CorticalColumnRust::new(
1168                n, tau, dt, threshold, w_exc, w_inh, seed,
1169            ),
1170        }
1171    }
1172
1173    fn step<'py>(
1174        &mut self,
1175        py: Python<'py>,
1176        thalamic_input: PyReadonlyArray1<'py, f64>,
1177    ) -> PyResult<Py<PyDict>> {
1178        let input = thalamic_input.as_slice()?;
1179        let spikes = self.inner.step(input);
1180        let dict = PyDict::new(py);
1181        let names = ["l4", "l23_exc", "l23_inh", "l5", "l6"];
1182        for (i, name) in names.iter().enumerate() {
1183            dict.set_item(*name, spikes[i].clone().into_pyarray(py))?;
1184        }
1185        Ok(dict.into())
1186    }
1187
1188    fn reset(&mut self) {
1189        self.inner.reset();
1190    }
1191}
1192
1193// ── Rall dendrite ────────────────────────────────────────────────────
1194
1195#[pyclass(
1196    name = "RallDendriteRust",
1197    module = "sc_neurocore_engine.sc_neurocore_engine"
1198)]
1199pub struct PyRallDendrite {
1200    inner: rall_dendrite::RallDendriteRust,
1201}
1202
1203#[pymethods]
1204impl PyRallDendrite {
1205    #[new]
1206    fn new(n_branches: usize, branch_length: usize, tau: f64, coupling: f64, dt: f64) -> Self {
1207        Self {
1208            inner: rall_dendrite::RallDendriteRust::new(
1209                n_branches,
1210                branch_length,
1211                tau,
1212                coupling,
1213                dt,
1214            ),
1215        }
1216    }
1217
1218    fn step(&mut self, branch_inputs: PyReadonlyArray1<'_, f64>) -> PyResult<f64> {
1219        let inputs = branch_inputs.as_slice()?;
1220        Ok(self.inner.step(inputs))
1221    }
1222
1223    fn reset(&mut self) {
1224        self.inner.reset();
1225    }
1226
1227    #[getter]
1228    fn soma_v(&self) -> f64 {
1229        self.inner.soma_v
1230    }
1231}
1232
1233/// Returns the highest SIMD tier available on this CPU.
1234#[pyfunction]
1235fn simd_tier() -> &'static str {
1236    #[cfg(target_arch = "x86_64")]
1237    {
1238        if is_x86_feature_detected!("avx512vpopcntdq") {
1239            return "avx512-vpopcntdq";
1240        }
1241        if is_x86_feature_detected!("avx512bw") {
1242            return "avx512bw";
1243        }
1244        if is_x86_feature_detected!("avx512f") {
1245            return "avx512f";
1246        }
1247        if is_x86_feature_detected!("avx2") {
1248            return "avx2";
1249        }
1250        if is_x86_feature_detected!("popcnt") {
1251            return "popcnt";
1252        }
1253    }
1254    #[cfg(target_arch = "aarch64")]
1255    {
1256        return "neon";
1257    }
1258    "portable"
1259}
1260
1261/// Set the number of threads in the global rayon thread pool.
1262///
1263/// Must be called before any parallel operation.
1264/// Passing 0 uses rayon's default (number of CPU cores).
1265#[pyfunction]
1266fn set_num_threads(n: usize) -> PyResult<()> {
1267    if n == 0 {
1268        return Ok(());
1269    }
1270    rayon::ThreadPoolBuilder::new()
1271        .num_threads(n)
1272        .build_global()
1273        .map_err(|e| PyValueError::new_err(format!("Cannot set thread pool: {e}")))
1274}
1275
1276#[pyfunction]
1277fn pack_bitstream(py: Python<'_>, bits: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
1278    if let Ok(rows) = bits.extract::<Vec<Vec<u8>>>() {
1279        let packed_rows: Vec<Vec<u64>> = rows.iter().map(|row| bitstream::pack(row).data).collect();
1280        return Ok(packed_rows
1281            .into_pyobject(py)
1282            .map_err(|e| PyValueError::new_err(e.to_string()))?
1283            .into_any()
1284            .unbind());
1285    }
1286
1287    let flat = bits
1288        .extract::<Vec<u8>>()
1289        .map_err(|_| PyValueError::new_err("Expected a 1-D or 2-D array of uint8 bits."))?;
1290    Ok(bitstream::pack(&flat)
1291        .data
1292        .into_pyobject(py)
1293        .map_err(|e| PyValueError::new_err(e.to_string()))?
1294        .into_any()
1295        .unbind())
1296}
1297
1298#[pyfunction]
1299#[pyo3(signature = (packed, original_length, original_shape=None))]
1300fn unpack_bitstream(
1301    py: Python<'_>,
1302    packed: &Bound<'_, PyAny>,
1303    original_length: usize,
1304    original_shape: Option<(usize, usize)>,
1305) -> PyResult<Py<PyAny>> {
1306    if let Ok(rows) = packed.extract::<Vec<Vec<u64>>>() {
1307        let batch = rows.len();
1308        let per_batch_len = if let Some((expected_batch, length)) = original_shape {
1309            if expected_batch != batch {
1310                return Err(PyValueError::new_err(format!(
1311                    "original_shape batch {} does not match packed batch {}.",
1312                    expected_batch, batch
1313                )));
1314            }
1315            length
1316        } else {
1317            original_length.checked_div(batch).unwrap_or(0)
1318        };
1319
1320        let unpacked_rows: Vec<Vec<u8>> = rows
1321            .into_iter()
1322            .map(|row| {
1323                bitstream::unpack(&bitstream::BitStreamTensor::from_words(row, per_batch_len))
1324            })
1325            .collect();
1326        return Ok(unpacked_rows
1327            .into_pyobject(py)
1328            .map_err(|e| PyValueError::new_err(e.to_string()))?
1329            .into_any()
1330            .unbind());
1331    }
1332
1333    let words = packed.extract::<Vec<u64>>().map_err(|_| {
1334        PyValueError::new_err("Expected packed uint64 words as 1-D or 2-D sequence.")
1335    })?;
1336    let tensor = bitstream::BitStreamTensor::from_words(words, original_length);
1337    Ok(bitstream::unpack(&tensor)
1338        .into_pyobject(py)
1339        .map_err(|e| PyValueError::new_err(e.to_string()))?
1340        .into_any()
1341        .unbind())
1342}
1343
1344#[pyfunction]
1345fn popcount(packed: &Bound<'_, PyAny>) -> PyResult<u64> {
1346    if let Ok(rows) = packed.extract::<Vec<Vec<u64>>>() {
1347        return Ok(rows
1348            .iter()
1349            .map(|row| simd::popcount_dispatch(row))
1350            .sum::<u64>());
1351    }
1352
1353    let words = packed.extract::<Vec<u64>>().map_err(|_| {
1354        PyValueError::new_err("Expected packed uint64 words as 1-D or 2-D sequence.")
1355    })?;
1356    Ok(simd::popcount_dispatch(&words))
1357}
1358
1359/// Pack a 1-D numpy uint8 array into packed u64 words, returning a numpy array.
1360/// Zero-copy input, single-allocation output.
1361#[pyfunction]
1362fn pack_bitstream_numpy<'py>(
1363    py: Python<'py>,
1364    bits: PyReadonlyArray1<'py, u8>,
1365) -> PyResult<Bound<'py, PyArray1<u64>>> {
1366    let slice = bits
1367        .as_slice()
1368        .map_err(|e| PyValueError::new_err(format!("Cannot read numpy array: {e}")))?;
1369    let tensor = simd::pack_dispatch(slice);
1370    Ok(tensor.data.into_pyarray(py))
1371}
1372
1373/// Popcount on a numpy uint64 array — zero-copy input.
1374#[pyfunction]
1375fn popcount_numpy(packed: PyReadonlyArray1<'_, u64>) -> PyResult<u64> {
1376    let words = packed
1377        .as_slice()
1378        .map_err(|e| PyValueError::new_err(format!("Cannot read numpy array: {e}")))?;
1379    Ok(simd::popcount_dispatch(words))
1380}
1381
1382/// Unpack a numpy uint64 array back to a numpy uint8 array.
1383#[pyfunction]
1384fn unpack_bitstream_numpy<'py>(
1385    py: Python<'py>,
1386    packed: PyReadonlyArray1<'py, u64>,
1387    original_length: usize,
1388) -> PyResult<Bound<'py, PyArray1<u8>>> {
1389    let words = packed
1390        .as_slice()
1391        .map_err(|e| PyValueError::new_err(format!("Cannot read numpy array: {e}")))?;
1392    let tensor = bitstream::BitStreamTensor::from_words(words.to_vec(), original_length);
1393    let bits = bitstream::unpack(&tensor);
1394    Ok(bits.into_pyarray(py))
1395}
1396
1397/// Run a LIF neuron for N steps with constant inputs.
1398///
1399/// Returns (spikes: ndarray[i32], voltages: ndarray[i16]).
1400#[pyfunction]
1401#[pyo3(signature = (
1402    n_steps,
1403    leak_k,
1404    gain_k,
1405    i_t,
1406    noise_in=0,
1407    data_width=16,
1408    fraction=8,
1409    v_rest=0,
1410    v_reset=0,
1411    v_threshold=256,
1412    refractory_period=2
1413))]
1414#[allow(clippy::too_many_arguments)]
1415fn batch_lif_run<'py>(
1416    py: Python<'py>,
1417    n_steps: usize,
1418    leak_k: i16,
1419    gain_k: i16,
1420    i_t: i16,
1421    noise_in: i16,
1422    data_width: u32,
1423    fraction: u32,
1424    v_rest: i16,
1425    v_reset: i16,
1426    v_threshold: i16,
1427    refractory_period: i32,
1428) -> (Bound<'py, PyArray1<i32>>, Bound<'py, PyArray1<i16>>) {
1429    let mut lif = neuron::FixedPointLif::new(
1430        data_width,
1431        fraction,
1432        v_rest,
1433        v_reset,
1434        v_threshold,
1435        refractory_period,
1436    );
1437    let spikes_arr = PyArray1::<i32>::zeros(py, n_steps, false);
1438    let voltages_arr = PyArray1::<i16>::zeros(py, n_steps, false);
1439
1440    // SAFETY: Arrays are newly allocated and contiguous.
1441    let spikes_slice = unsafe {
1442        spikes_arr
1443            .as_slice_mut()
1444            .expect("newly allocated spikes array must be contiguous")
1445    };
1446    // SAFETY: Arrays are newly allocated and contiguous.
1447    let voltages_slice = unsafe {
1448        voltages_arr
1449            .as_slice_mut()
1450            .expect("newly allocated voltages array must be contiguous")
1451    };
1452
1453    for i in 0..n_steps {
1454        let (s, v) = lif.step(leak_k, gain_k, i_t, noise_in);
1455        spikes_slice[i] = s;
1456        voltages_slice[i] = v;
1457    }
1458
1459    (spikes_arr, voltages_arr)
1460}
1461
1462/// Run N independent LIF neurons in parallel, each with its own constant input.
1463///
1464/// Returns (spikes: ndarray[i32, (n_neurons, n_steps)],
1465///          voltages: ndarray[i16, (n_neurons, n_steps)]).
1466#[pyfunction]
1467#[pyo3(signature = (
1468    n_neurons,
1469    n_steps,
1470    leak_k,
1471    gain_k,
1472    currents,
1473    data_width=16,
1474    fraction=8,
1475    v_rest=0,
1476    v_reset=0,
1477    v_threshold=256,
1478    refractory_period=2
1479))]
1480#[allow(clippy::too_many_arguments)]
1481#[allow(clippy::type_complexity)]
1482fn batch_lif_run_multi<'py>(
1483    py: Python<'py>,
1484    n_neurons: usize,
1485    n_steps: usize,
1486    leak_k: i16,
1487    gain_k: i16,
1488    currents: PyReadonlyArray1<'py, i16>,
1489    data_width: u32,
1490    fraction: u32,
1491    v_rest: i16,
1492    v_reset: i16,
1493    v_threshold: i16,
1494    refractory_period: i32,
1495) -> PyResult<(Bound<'py, PyArray2<i32>>, Bound<'py, PyArray2<i16>>)> {
1496    use rayon::prelude::*;
1497
1498    let curr_slice = currents
1499        .as_slice()
1500        .map_err(|e| PyValueError::new_err(format!("Cannot read currents: {e}")))?;
1501    if curr_slice.len() != n_neurons {
1502        return Err(PyValueError::new_err(format!(
1503            "currents length {} does not match n_neurons {}.",
1504            curr_slice.len(),
1505            n_neurons
1506        )));
1507    }
1508
1509    let spikes_arr = PyArray2::<i32>::zeros(py, [n_neurons, n_steps], false);
1510    let voltages_arr = PyArray2::<i16>::zeros(py, [n_neurons, n_steps], false);
1511
1512    if n_neurons == 0 || n_steps == 0 {
1513        return Ok((spikes_arr, voltages_arr));
1514    }
1515
1516    // SAFETY: Arrays are newly allocated and contiguous.
1517    let spikes_flat = unsafe {
1518        spikes_arr
1519            .as_slice_mut()
1520            .expect("newly allocated spikes array must be contiguous")
1521    };
1522    // SAFETY: Arrays are newly allocated and contiguous.
1523    let voltages_flat = unsafe {
1524        voltages_arr
1525            .as_slice_mut()
1526            .expect("newly allocated voltages array must be contiguous")
1527    };
1528
1529    spikes_flat
1530        .par_chunks_mut(n_steps)
1531        .zip(voltages_flat.par_chunks_mut(n_steps))
1532        .zip(curr_slice.par_iter().copied())
1533        .for_each(|((spike_row, voltage_row), i_t)| {
1534            let mut lif = neuron::FixedPointLif::new(
1535                data_width,
1536                fraction,
1537                v_rest,
1538                v_reset,
1539                v_threshold,
1540                refractory_period,
1541            );
1542            for step in 0..n_steps {
1543                let (s, v) = lif.step(leak_k, gain_k, i_t, 0);
1544                spike_row[step] = s;
1545                voltage_row[step] = v;
1546            }
1547        });
1548
1549    Ok((spikes_arr, voltages_arr))
1550}
1551
1552/// Run a LIF neuron for N steps with per-step current and optional noise arrays.
1553#[pyfunction]
1554#[pyo3(signature = (
1555    leak_k,
1556    gain_k,
1557    currents,
1558    noises=None,
1559    data_width=16,
1560    fraction=8,
1561    v_rest=0,
1562    v_reset=0,
1563    v_threshold=256,
1564    refractory_period=2
1565))]
1566#[allow(clippy::too_many_arguments)]
1567#[allow(clippy::type_complexity)]
1568fn batch_lif_run_varying<'py>(
1569    py: Python<'py>,
1570    leak_k: i16,
1571    gain_k: i16,
1572    currents: PyReadonlyArray1<'py, i16>,
1573    noises: Option<PyReadonlyArray1<'py, i16>>,
1574    data_width: u32,
1575    fraction: u32,
1576    v_rest: i16,
1577    v_reset: i16,
1578    v_threshold: i16,
1579    refractory_period: i32,
1580) -> PyResult<(Bound<'py, PyArray1<i32>>, Bound<'py, PyArray1<i16>>)> {
1581    let curr_slice = currents
1582        .as_slice()
1583        .map_err(|e| PyValueError::new_err(format!("Cannot read currents: {e}")))?;
1584    let noise_slice: Option<&[i16]> = match noises.as_ref() {
1585        Some(n) => Some(
1586            n.as_slice()
1587                .map_err(|e| PyValueError::new_err(format!("Cannot read noises: {e}")))?,
1588        ),
1589        None => None,
1590    };
1591
1592    let n_steps = curr_slice.len();
1593    if let Some(ns) = noise_slice {
1594        if ns.len() != n_steps {
1595            return Err(PyValueError::new_err(format!(
1596                "noises length {} does not match currents length {}.",
1597                ns.len(),
1598                n_steps
1599            )));
1600        }
1601    }
1602
1603    let mut lif = neuron::FixedPointLif::new(
1604        data_width,
1605        fraction,
1606        v_rest,
1607        v_reset,
1608        v_threshold,
1609        refractory_period,
1610    );
1611    let spikes_arr = PyArray1::<i32>::zeros(py, n_steps, false);
1612    let voltages_arr = PyArray1::<i16>::zeros(py, n_steps, false);
1613
1614    // SAFETY: Arrays are newly allocated and contiguous.
1615    let spikes_slice = unsafe {
1616        spikes_arr
1617            .as_slice_mut()
1618            .expect("newly allocated spikes array must be contiguous")
1619    };
1620    // SAFETY: Arrays are newly allocated and contiguous.
1621    let voltages_slice = unsafe {
1622        voltages_arr
1623            .as_slice_mut()
1624            .expect("newly allocated voltages array must be contiguous")
1625    };
1626
1627    for i in 0..n_steps {
1628        let noise_in = noise_slice.map_or(0, |ns| ns[i]);
1629        let (s, v) = lif.step(leak_k, gain_k, curr_slice[i], noise_in);
1630        spikes_slice[i] = s;
1631        voltages_slice[i] = v;
1632    }
1633
1634    Ok((spikes_arr, voltages_arr))
1635}
1636
1637/// Bernoulli-encode a numpy float64 array into packed bitstream words.
1638///
1639/// Returns nested packed words with shape (n_probs, ceil(length / 64)).
1640#[pyfunction]
1641#[pyo3(signature = (probs, length=1024, seed=0xACE1))]
1642fn batch_encode<'py>(
1643    _py: Python<'py>,
1644    probs: PyReadonlyArray1<'py, f64>,
1645    length: usize,
1646    seed: u64,
1647) -> PyResult<Vec<Vec<u64>>> {
1648    let prob_slice = probs
1649        .as_slice()
1650        .map_err(|e| PyValueError::new_err(format!("Cannot read probs: {e}")))?;
1651    let words = length.div_ceil(64);
1652
1653    use rand::SeedableRng;
1654    let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(seed);
1655
1656    let packed: Vec<Vec<u64>> = prob_slice
1657        .iter()
1658        .map(|&p| {
1659            let mut data = bitstream::bernoulli_packed(p, length, &mut rng);
1660            data.resize(words, 0);
1661            data
1662        })
1663        .collect();
1664
1665    Ok(packed)
1666}
1667
1668/// Bernoulli-encode a numpy float64 array into a 2-D numpy uint64 array.
1669///
1670/// Returns shape `(n_probs, ceil(length / 64))`.
1671#[pyfunction]
1672#[pyo3(signature = (probs, length=1024, seed=0xACE1))]
1673fn batch_encode_numpy<'py>(
1674    py: Python<'py>,
1675    probs: PyReadonlyArray1<'py, f64>,
1676    length: usize,
1677    seed: u64,
1678) -> PyResult<Bound<'py, PyArray2<u64>>> {
1679    use rayon::prelude::*;
1680
1681    let prob_slice = probs
1682        .as_slice()
1683        .map_err(|e| PyValueError::new_err(format!("Cannot read probs: {e}")))?;
1684    let words = length.div_ceil(64);
1685    let n_probs = prob_slice.len();
1686
1687    let rows: Vec<Vec<u64>> = prob_slice
1688        .par_iter()
1689        .enumerate()
1690        .map(|(idx, &p)| {
1691            use rand::SeedableRng;
1692
1693            let prob_seed = seed.wrapping_add(idx as u64);
1694            let mut rng = rand_xoshiro::Xoshiro256PlusPlus::seed_from_u64(prob_seed);
1695            let mut row = bitstream::bernoulli_packed_simd(p, length, &mut rng);
1696            row.resize(words, 0);
1697            row
1698        })
1699        .collect();
1700
1701    let mut flat = Vec::with_capacity(n_probs * words);
1702    for row in &rows {
1703        flat.extend_from_slice(row);
1704    }
1705
1706    let arr = ndarray::Array2::from_shape_vec((n_probs, words), flat)
1707        .map_err(|e| PyValueError::new_err(format!("Shape construction failed: {e}")))?;
1708    Ok(arr.into_pyarray(py))
1709}
1710
1711#[pyclass(module = "sc_neurocore_engine.sc_neurocore_engine")]
1712pub struct Lfsr16 {
1713    inner: encoder::Lfsr16,
1714    seed_init: u16,
1715}
1716
1717#[pymethods]
1718impl Lfsr16 {
1719    #[new]
1720    #[pyo3(signature = (seed=0xACE1))]
1721    fn new(seed: u16) -> PyResult<Self> {
1722        if seed == 0 {
1723            return Err(PyValueError::new_err("LFSR seed must be non-zero."));
1724        }
1725        Ok(Self {
1726            inner: encoder::Lfsr16::new(seed),
1727            seed_init: seed,
1728        })
1729    }
1730
1731    fn step(&mut self) -> u16 {
1732        self.inner.step()
1733    }
1734
1735    #[getter]
1736    fn reg(&self) -> u16 {
1737        self.inner.reg
1738    }
1739
1740    #[getter]
1741    fn width(&self) -> u32 {
1742        self.inner.width
1743    }
1744
1745    #[pyo3(signature = (seed=None))]
1746    fn reset(&mut self, seed: Option<u16>) -> PyResult<()> {
1747        let next = seed.unwrap_or(self.seed_init);
1748        if next == 0 {
1749            return Err(PyValueError::new_err("LFSR seed must be non-zero."));
1750        }
1751        self.inner = encoder::Lfsr16::new(next);
1752        self.seed_init = next;
1753        Ok(())
1754    }
1755}
1756
1757#[pyclass(module = "sc_neurocore_engine.sc_neurocore_engine")]
1758pub struct BitstreamEncoder {
1759    inner: encoder::BitstreamEncoder,
1760    seed_init: u16,
1761}
1762
1763#[pymethods]
1764impl BitstreamEncoder {
1765    #[new]
1766    #[pyo3(signature = (data_width=16, seed=0xACE1))]
1767    fn new(data_width: u32, seed: u16) -> PyResult<Self> {
1768        if seed == 0 {
1769            return Err(PyValueError::new_err("LFSR seed must be non-zero."));
1770        }
1771        Ok(Self {
1772            inner: encoder::BitstreamEncoder::new(data_width, seed),
1773            seed_init: seed,
1774        })
1775    }
1776
1777    fn step(&mut self, x_value: u16) -> u8 {
1778        self.inner.step(x_value)
1779    }
1780
1781    #[getter]
1782    fn data_width(&self) -> u32 {
1783        self.inner.data_width
1784    }
1785
1786    #[getter]
1787    fn reg(&self) -> u16 {
1788        self.inner.lfsr.reg
1789    }
1790
1791    #[pyo3(signature = (seed=None))]
1792    fn reset(&mut self, seed: Option<u16>) -> PyResult<()> {
1793        let next = seed.unwrap_or(self.seed_init);
1794        if next == 0 {
1795            return Err(PyValueError::new_err("LFSR seed must be non-zero."));
1796        }
1797        self.inner.reset(Some(next));
1798        self.seed_init = next;
1799        Ok(())
1800    }
1801}
1802
1803#[pyclass(module = "sc_neurocore_engine.sc_neurocore_engine")]
1804pub struct FixedPointLif {
1805    inner: neuron::FixedPointLif,
1806}
1807
1808#[pymethods]
1809impl FixedPointLif {
1810    #[new]
1811    #[pyo3(signature = (
1812        data_width=16,
1813        fraction=8,
1814        v_rest=0,
1815        v_reset=0,
1816        v_threshold=256,
1817        refractory_period=2
1818    ))]
1819    fn new(
1820        data_width: u32,
1821        fraction: u32,
1822        v_rest: i16,
1823        v_reset: i16,
1824        v_threshold: i16,
1825        refractory_period: i32,
1826    ) -> Self {
1827        Self {
1828            inner: neuron::FixedPointLif::new(
1829                data_width,
1830                fraction,
1831                v_rest,
1832                v_reset,
1833                v_threshold,
1834                refractory_period,
1835            ),
1836        }
1837    }
1838
1839    #[pyo3(signature = (leak_k, gain_k, i_t, noise_in=0))]
1840    fn step(&mut self, leak_k: i16, gain_k: i16, i_t: i16, noise_in: i16) -> (i32, i16) {
1841        self.inner.step(leak_k, gain_k, i_t, noise_in)
1842    }
1843
1844    fn reset(&mut self) {
1845        self.inner.reset();
1846    }
1847
1848    fn reset_state(&mut self) {
1849        self.reset();
1850    }
1851
1852    fn get_state(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
1853        let dict = PyDict::new(py);
1854        dict.set_item("v", self.inner.v)?;
1855        dict.set_item("refractory_counter", self.inner.refractory_counter)?;
1856        Ok(dict.into_any().unbind())
1857    }
1858}
1859
1860// ── Izhikevich PyO3 wrapper ─────────────────────────────────────
1861
1862#[pyclass(
1863    name = "Izhikevich",
1864    module = "sc_neurocore_engine.sc_neurocore_engine"
1865)]
1866pub struct PyIzhikevich {
1867    inner: neuron::Izhikevich,
1868}
1869
1870#[pymethods]
1871impl PyIzhikevich {
1872    #[new]
1873    #[pyo3(signature = (a=0.02, b=0.2, c=-65.0, d=8.0, dt=1.0))]
1874    fn new(a: f64, b: f64, c: f64, d: f64, dt: f64) -> Self {
1875        Self {
1876            inner: neuron::Izhikevich::new(a, b, c, d, dt),
1877        }
1878    }
1879
1880    fn step(&mut self, current: f64) -> i32 {
1881        self.inner.step(current)
1882    }
1883
1884    fn reset(&mut self) {
1885        self.inner.reset();
1886    }
1887
1888    fn reset_state(&mut self) {
1889        self.reset();
1890    }
1891
1892    fn get_state(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
1893        let dict = PyDict::new(py);
1894        dict.set_item("v", self.inner.v)?;
1895        dict.set_item("u", self.inner.u)?;
1896        Ok(dict.into_any().unbind())
1897    }
1898}
1899
1900// ── BitstreamAverager PyO3 wrapper ──────────────────────────────
1901
1902#[pyclass(
1903    name = "BitstreamAverager",
1904    module = "sc_neurocore_engine.sc_neurocore_engine"
1905)]
1906pub struct PyBitstreamAverager {
1907    inner: neuron::BitstreamAverager,
1908}
1909
1910#[pymethods]
1911impl PyBitstreamAverager {
1912    #[new]
1913    #[pyo3(signature = (window=1024))]
1914    fn new(window: usize) -> Self {
1915        Self {
1916            inner: neuron::BitstreamAverager::new(window),
1917        }
1918    }
1919
1920    fn push(&mut self, bit: u8) {
1921        self.inner.push(bit);
1922    }
1923
1924    fn estimate(&self) -> f64 {
1925        self.inner.estimate()
1926    }
1927
1928    fn reset(&mut self) {
1929        self.inner.reset();
1930    }
1931
1932    #[getter]
1933    fn window(&self) -> usize {
1934        self.inner.window()
1935    }
1936}
1937
1938// ── AdEx, ExpIF, Lapicque PyO3 wrappers ────────────────────────
1939
1940#[pyclass(
1941    name = "AdExNeuron",
1942    module = "sc_neurocore_engine.sc_neurocore_engine"
1943)]
1944#[derive(Clone)]
1945pub struct PyAdExNeuron {
1946    inner: neuron::AdExNeuron,
1947}
1948
1949#[pymethods]
1950impl PyAdExNeuron {
1951    #[new]
1952    fn new() -> Self {
1953        Self {
1954            inner: neuron::AdExNeuron::new(),
1955        }
1956    }
1957    fn step(&mut self, current: f64) -> i32 {
1958        self.inner.step(current)
1959    }
1960    fn reset(&mut self) {
1961        self.inner.reset();
1962    }
1963    fn get_state(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
1964        let d = PyDict::new(py);
1965        d.set_item("v", self.inner.v)?;
1966        d.set_item("w", self.inner.w)?;
1967        Ok(d.into_any().unbind())
1968    }
1969}
1970
1971#[pyclass(
1972    name = "ExpIFNeuron",
1973    module = "sc_neurocore_engine.sc_neurocore_engine"
1974)]
1975#[derive(Clone)]
1976pub struct PyExpIFNeuron {
1977    inner: neuron::ExpIfNeuron,
1978}
1979
1980#[pymethods]
1981impl PyExpIFNeuron {
1982    #[new]
1983    fn new() -> Self {
1984        Self {
1985            inner: neuron::ExpIfNeuron::new(),
1986        }
1987    }
1988    fn step(&mut self, current: f64) -> i32 {
1989        self.inner.step(current)
1990    }
1991    fn reset(&mut self) {
1992        self.inner.reset();
1993    }
1994    fn get_state(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
1995        let d = PyDict::new(py);
1996        d.set_item("v", self.inner.v)?;
1997        Ok(d.into_any().unbind())
1998    }
1999}
2000
2001#[pyclass(
2002    name = "LapicqueNeuron",
2003    module = "sc_neurocore_engine.sc_neurocore_engine"
2004)]
2005#[derive(Clone)]
2006pub struct PyLapicqueNeuron {
2007    inner: neuron::LapicqueNeuron,
2008}
2009
2010#[pymethods]
2011impl PyLapicqueNeuron {
2012    #[new]
2013    #[pyo3(signature = (tau=20.0, resistance=1.0, threshold=1.0, dt=1.0))]
2014    fn new(tau: f64, resistance: f64, threshold: f64, dt: f64) -> Self {
2015        Self {
2016            inner: neuron::LapicqueNeuron::new(tau, resistance, threshold, dt),
2017        }
2018    }
2019    fn step(&mut self, current: f64) -> i32 {
2020        self.inner.step(current)
2021    }
2022    fn reset(&mut self) {
2023        self.inner.reset();
2024    }
2025    fn get_state(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
2026        let d = PyDict::new(py);
2027        d.set_item("v", self.inner.v)?;
2028        Ok(d.into_any().unbind())
2029    }
2030}
2031
2032// ── DenseLayer ──────────────────────────────────────────────────
2033
2034#[pyclass(module = "sc_neurocore_engine.sc_neurocore_engine")]
2035pub struct DenseLayer {
2036    inner: layer::DenseLayer,
2037}
2038
2039#[pymethods]
2040impl DenseLayer {
2041    #[new]
2042    #[pyo3(signature = (n_inputs, n_neurons, length=1024, seed=24301))]
2043    fn new(n_inputs: usize, n_neurons: usize, length: usize, seed: u64) -> Self {
2044        Self {
2045            inner: layer::DenseLayer::new(n_inputs, n_neurons, length, seed),
2046        }
2047    }
2048
2049    fn get_weights(&self) -> Vec<Vec<f64>> {
2050        self.inner.get_weights()
2051    }
2052
2053    fn set_weights(&mut self, weights: Vec<Vec<f64>>) -> PyResult<()> {
2054        self.inner
2055            .set_weights(weights)
2056            .map_err(PyValueError::new_err)
2057    }
2058
2059    fn refresh_packed_weights(&mut self) {
2060        self.inner.refresh_packed_weights();
2061    }
2062
2063    #[pyo3(signature = (input_values, seed=44257))]
2064    fn forward(&self, input_values: Vec<f64>, seed: u64) -> PyResult<Vec<f64>> {
2065        self.inner
2066            .forward(&input_values, seed)
2067            .map_err(PyValueError::new_err)
2068    }
2069
2070    #[pyo3(signature = (input_values, seed=44257))]
2071    fn forward_fast(&self, input_values: Vec<f64>, seed: u64) -> PyResult<Vec<f64>> {
2072        self.inner
2073            .forward_fast(&input_values, seed)
2074            .map_err(PyValueError::new_err)
2075    }
2076
2077    /// Dense forward accepting numpy input and returning numpy output.
2078    ///
2079    /// This performs parallel encoding + parallel compute in one FFI call.
2080    #[pyo3(signature = (input_values, seed=44257))]
2081    fn forward_numpy<'py>(
2082        &self,
2083        py: Python<'py>,
2084        input_values: PyReadonlyArray1<'py, f64>,
2085        seed: u64,
2086    ) -> PyResult<Bound<'py, PyArray1<f64>>> {
2087        let slice = input_values
2088            .as_slice()
2089            .map_err(|e| PyValueError::new_err(format!("Cannot read input array: {e}")))?;
2090        let out = self
2091            .inner
2092            .forward_numpy_inner(slice, seed)
2093            .map_err(PyValueError::new_err)?;
2094        Ok(out.into_pyarray(py))
2095    }
2096
2097    /// Dense forward for a batch of input samples in one FFI call.
2098    ///
2099    /// `inputs` must be a contiguous float64 array of shape (n_samples, n_inputs).
2100    /// Returns float64 array of shape (n_samples, n_neurons).
2101    #[pyo3(signature = (inputs, seed=44257))]
2102    fn forward_batch_numpy<'py>(
2103        &self,
2104        py: Python<'py>,
2105        inputs: PyReadonlyArray2<'py, f64>,
2106        seed: u64,
2107    ) -> PyResult<Bound<'py, PyArray2<f64>>> {
2108        let shape = inputs.shape();
2109        let n_samples = shape[0];
2110        let n_inputs = shape[1];
2111        if n_inputs != self.inner.n_inputs {
2112            return Err(PyValueError::new_err(format!(
2113                "Expected {} input features, got {}.",
2114                self.inner.n_inputs, n_inputs
2115            )));
2116        }
2117
2118        let flat_inputs = inputs
2119            .as_slice()
2120            .map_err(|e| PyValueError::new_err(format!("Array not contiguous: {e}")))?;
2121        let out = PyArray2::<f64>::zeros(py, [n_samples, self.inner.n_neurons], false);
2122        // SAFETY: Newly allocated numpy arrays are contiguous.
2123        let out_slice = unsafe {
2124            out.as_slice_mut()
2125                .expect("newly allocated output array must be contiguous")
2126        };
2127
2128        self.inner
2129            .forward_batch_into(flat_inputs, n_samples, seed, out_slice)
2130            .map_err(PyValueError::new_err)?;
2131        Ok(out)
2132    }
2133
2134    /// Forward pass with pre-packed input bitstreams.
2135    ///
2136    /// Accepts either:
2137    /// - 2-D numpy array of dtype uint64 with shape (n_inputs, words)
2138    /// - list[list[int]]
2139    fn forward_prepacked(&self, packed_inputs: &Bound<'_, PyAny>) -> PyResult<Vec<f64>> {
2140        if let Ok(arr) = packed_inputs.extract::<PyReadonlyArray2<u64>>() {
2141            let view = arr.as_array();
2142            let rows: Vec<Vec<u64>> = (0..view.nrows()).map(|i| view.row(i).to_vec()).collect();
2143            return self
2144                .inner
2145                .forward_prepacked(&rows)
2146                .map_err(PyValueError::new_err);
2147        }
2148
2149        let rows = packed_inputs.extract::<Vec<Vec<u64>>>().map_err(|_| {
2150            PyValueError::new_err(
2151                "packed_inputs must be a 2-D numpy uint64 array or list[list[int]].",
2152            )
2153        })?;
2154        self.inner
2155            .forward_prepacked(&rows)
2156            .map_err(PyValueError::new_err)
2157    }
2158
2159    /// Dense forward with pre-packed numpy 2-D input (true zero-copy).
2160    ///
2161    /// Accepts a contiguous numpy uint64 array of shape (n_inputs, words).
2162    /// This avoids all row-copying that the `forward_prepacked` method does.
2163    #[pyo3(signature = (packed_inputs,))]
2164    fn forward_prepacked_numpy<'py>(
2165        &self,
2166        py: Python<'py>,
2167        packed_inputs: PyReadonlyArray2<'py, u64>,
2168    ) -> PyResult<Bound<'py, PyArray1<f64>>> {
2169        let shape = packed_inputs.shape();
2170        let n_inputs = shape[0];
2171        let words = shape[1];
2172        let flat = packed_inputs
2173            .as_slice()
2174            .map_err(|e| PyValueError::new_err(format!("Array not contiguous: {e}")))?;
2175        let out = self
2176            .inner
2177            .forward_prepacked_2d(flat, n_inputs, words)
2178            .map_err(PyValueError::new_err)?;
2179        Ok(out.into_pyarray(py))
2180    }
2181}
2182
2183fn parse_surrogate(name: &str, k: Option<f32>) -> PyResult<grad::SurrogateType> {
2184    let normalized = name.to_ascii_lowercase().replace('-', "_");
2185    match normalized.as_str() {
2186        "fast_sigmoid" => Ok(grad::SurrogateType::FastSigmoid {
2187            k: k.unwrap_or(25.0),
2188        }),
2189        "superspike" | "super_spike" => Ok(grad::SurrogateType::SuperSpike {
2190            k: k.unwrap_or(100.0),
2191        }),
2192        "arctan" | "arc_tan" => Ok(grad::SurrogateType::ArcTan { k: k.unwrap_or(10.0) }),
2193        "straightthrough" | "straight_through" | "ste" => Ok(grad::SurrogateType::StraightThrough),
2194        _ => Err(PyValueError::new_err(format!(
2195            "Unknown surrogate '{}'. Use one of: fast_sigmoid, superspike, arctan, straight_through.",
2196            name
2197        ))),
2198    }
2199}
2200
2201fn extract_matrix_f64(data: &Bound<'_, PyAny>, name: &str) -> PyResult<(Vec<f64>, usize, usize)> {
2202    if let Ok(rows) = data.extract::<Vec<Vec<f64>>>() {
2203        if rows.is_empty() {
2204            return Err(PyValueError::new_err(format!(
2205                "{} must not be an empty matrix.",
2206                name
2207            )));
2208        }
2209        let row_count = rows.len();
2210        let cols = rows[0].len();
2211        if cols == 0 {
2212            return Err(PyValueError::new_err(format!(
2213                "{} must not have zero columns.",
2214                name
2215            )));
2216        }
2217        if rows.iter().any(|r| r.len() != cols) {
2218            return Err(PyValueError::new_err(format!(
2219                "{} must be a rectangular matrix.",
2220                name
2221            )));
2222        }
2223        let out = rows.into_iter().flatten().collect::<Vec<f64>>();
2224        return Ok((out, row_count, cols));
2225    }
2226
2227    if let Ok(flat) = data.extract::<Vec<f64>>() {
2228        if flat.is_empty() {
2229            return Err(PyValueError::new_err(format!(
2230                "{} must not be an empty vector.",
2231                name
2232            )));
2233        }
2234        let cols = flat.len();
2235        return Ok((flat, 1, cols));
2236    }
2237
2238    Err(PyValueError::new_err(format!(
2239        "{} must be a 1-D or 2-D float array.",
2240        name
2241    )))
2242}
2243
2244fn reshape_flat_to_rows(flat: Vec<f64>, rows: usize, cols: usize) -> Vec<Vec<f64>> {
2245    let mut out = Vec::with_capacity(rows);
2246    for i in 0..rows {
2247        out.push(flat[i * cols..(i + 1) * cols].to_vec());
2248    }
2249    out
2250}
2251
2252#[pyclass(
2253    name = "SurrogateLif",
2254    module = "sc_neurocore_engine.sc_neurocore_engine"
2255)]
2256pub struct PySurrogateLif {
2257    inner: grad::SurrogateLif,
2258}
2259
2260#[pymethods]
2261impl PySurrogateLif {
2262    #[new]
2263    #[pyo3(signature = (
2264        data_width=16,
2265        fraction=8,
2266        v_rest=0,
2267        v_reset=0,
2268        v_threshold=256,
2269        refractory_period=2,
2270        surrogate="fast_sigmoid",
2271        k=None
2272    ))]
2273    #[allow(clippy::too_many_arguments)]
2274    fn new(
2275        data_width: u32,
2276        fraction: u32,
2277        v_rest: i16,
2278        v_reset: i16,
2279        v_threshold: i16,
2280        refractory_period: i32,
2281        surrogate: &str,
2282        k: Option<f32>,
2283    ) -> PyResult<Self> {
2284        let surrogate = parse_surrogate(surrogate, k)?;
2285        Ok(Self {
2286            inner: grad::SurrogateLif::new(
2287                data_width,
2288                fraction,
2289                v_rest,
2290                v_reset,
2291                v_threshold,
2292                refractory_period,
2293                surrogate,
2294            ),
2295        })
2296    }
2297
2298    #[pyo3(signature = (leak_k, gain_k, i_t, noise_in=0))]
2299    fn forward(&mut self, leak_k: i16, gain_k: i16, i_t: i16, noise_in: i16) -> (i32, i16) {
2300        self.inner.forward(leak_k, gain_k, i_t, noise_in)
2301    }
2302
2303    fn backward(&mut self, grad_output: f32) -> PyResult<f32> {
2304        self.inner
2305            .backward(grad_output)
2306            .map_err(PyValueError::new_err)
2307    }
2308
2309    fn clear_trace(&mut self) {
2310        self.inner.clear_trace();
2311    }
2312
2313    fn reset(&mut self) {
2314        self.inner.reset();
2315    }
2316
2317    fn trace_len(&self) -> usize {
2318        self.inner.trace_len()
2319    }
2320}
2321
2322#[pyclass(
2323    name = "DifferentiableDenseLayer",
2324    module = "sc_neurocore_engine.sc_neurocore_engine"
2325)]
2326pub struct PyDifferentiableDenseLayer {
2327    inner: grad::DifferentiableDenseLayer,
2328}
2329
2330#[pymethods]
2331impl PyDifferentiableDenseLayer {
2332    #[new]
2333    #[pyo3(signature = (
2334        n_inputs,
2335        n_neurons,
2336        length=1024,
2337        seed=24301,
2338        surrogate="fast_sigmoid",
2339        k=None
2340    ))]
2341    fn new(
2342        n_inputs: usize,
2343        n_neurons: usize,
2344        length: usize,
2345        seed: u64,
2346        surrogate: &str,
2347        k: Option<f32>,
2348    ) -> PyResult<Self> {
2349        let surrogate = parse_surrogate(surrogate, k)?;
2350        Ok(Self {
2351            inner: grad::DifferentiableDenseLayer::new(
2352                n_inputs, n_neurons, length, seed, surrogate,
2353            ),
2354        })
2355    }
2356
2357    fn get_weights(&self) -> Vec<Vec<f64>> {
2358        self.inner.layer.get_weights()
2359    }
2360
2361    #[pyo3(signature = (input_values, seed=44257))]
2362    fn forward(&mut self, input_values: Vec<f64>, seed: u64) -> PyResult<Vec<f64>> {
2363        self.inner
2364            .forward(&input_values, seed)
2365            .map_err(PyValueError::new_err)
2366    }
2367
2368    fn backward(&self, grad_output: Vec<f64>) -> PyResult<(Vec<f64>, Vec<Vec<f64>>)> {
2369        self.inner
2370            .backward(&grad_output)
2371            .map_err(PyValueError::new_err)
2372    }
2373
2374    fn update_weights(&mut self, weight_grads: Vec<Vec<f64>>, lr: f64) -> PyResult<()> {
2375        if weight_grads.len() != self.inner.layer.n_neurons {
2376            return Err(PyValueError::new_err(format!(
2377                "Expected {} grad rows, got {}.",
2378                self.inner.layer.n_neurons,
2379                weight_grads.len()
2380            )));
2381        }
2382        if weight_grads
2383            .iter()
2384            .any(|row| row.len() != self.inner.layer.n_inputs)
2385        {
2386            return Err(PyValueError::new_err(format!(
2387                "Expected each grad row to have length {}.",
2388                self.inner.layer.n_inputs
2389            )));
2390        }
2391        self.inner.update_weights(&weight_grads, lr);
2392        Ok(())
2393    }
2394
2395    fn clear_cache(&mut self) {
2396        self.inner.clear_cache();
2397    }
2398}
2399
2400#[pyclass(
2401    name = "StochasticAttention",
2402    module = "sc_neurocore_engine.sc_neurocore_engine"
2403)]
2404pub struct PyStochasticAttention {
2405    inner: attention::StochasticAttention,
2406}
2407
2408#[pymethods]
2409impl PyStochasticAttention {
2410    #[new]
2411    #[pyo3(signature = (dim_k, temperature=None))]
2412    fn new(dim_k: usize, temperature: Option<f64>) -> Self {
2413        Self {
2414            inner: match temperature {
2415                Some(t) => attention::StochasticAttention::with_temperature(dim_k, t),
2416                None => attention::StochasticAttention::new(dim_k),
2417            },
2418        }
2419    }
2420
2421    fn forward_softmax(
2422        &self,
2423        q: &Bound<'_, PyAny>,
2424        k: &Bound<'_, PyAny>,
2425        v: &Bound<'_, PyAny>,
2426    ) -> PyResult<Vec<Vec<f64>>> {
2427        let (q_data, q_rows, q_cols) = extract_matrix_f64(q, "Q")?;
2428        let (k_data, k_rows, k_cols) = extract_matrix_f64(k, "K")?;
2429        let (v_data, v_rows, v_cols) = extract_matrix_f64(v, "V")?;
2430
2431        let out = self
2432            .inner
2433            .forward_softmax(
2434                &q_data, q_rows, q_cols, &k_data, k_rows, k_cols, &v_data, v_rows, v_cols,
2435            )
2436            .map_err(PyValueError::new_err)?;
2437
2438        Ok(reshape_flat_to_rows(out, q_rows, v_cols))
2439    }
2440
2441    #[pyo3(signature = (q, k, v, n_heads))]
2442    fn forward_multihead_softmax(
2443        &self,
2444        q: &Bound<'_, PyAny>,
2445        k: &Bound<'_, PyAny>,
2446        v: &Bound<'_, PyAny>,
2447        n_heads: usize,
2448    ) -> PyResult<Vec<Vec<f64>>> {
2449        let (q_data, q_rows, q_cols) = extract_matrix_f64(q, "Q")?;
2450        let (k_data, k_rows, k_cols) = extract_matrix_f64(k, "K")?;
2451        let (v_data, v_rows, v_cols) = extract_matrix_f64(v, "V")?;
2452
2453        let out = self
2454            .inner
2455            .forward_multihead_softmax(
2456                &q_data, q_rows, q_cols, &k_data, k_rows, k_cols, &v_data, v_rows, v_cols, n_heads,
2457            )
2458            .map_err(PyValueError::new_err)?;
2459
2460        let out_cols = v_cols;
2461        Ok(reshape_flat_to_rows(out, q_rows, out_cols))
2462    }
2463
2464    fn forward(
2465        &self,
2466        q: &Bound<'_, PyAny>,
2467        k: &Bound<'_, PyAny>,
2468        v: &Bound<'_, PyAny>,
2469    ) -> PyResult<Vec<Vec<f64>>> {
2470        let (q_data, q_rows, q_cols) = extract_matrix_f64(q, "Q")?;
2471        let (k_data, k_rows, k_cols) = extract_matrix_f64(k, "K")?;
2472        let (v_data, v_rows, v_cols) = extract_matrix_f64(v, "V")?;
2473
2474        let out = self
2475            .inner
2476            .forward(
2477                &q_data, q_rows, q_cols, &k_data, k_rows, k_cols, &v_data, v_rows, v_cols,
2478            )
2479            .map_err(PyValueError::new_err)?;
2480
2481        Ok(reshape_flat_to_rows(out, q_rows, v_cols))
2482    }
2483
2484    #[pyo3(signature = (q, k, v, length=1024, seed=44257))]
2485    fn forward_sc(
2486        &self,
2487        q: &Bound<'_, PyAny>,
2488        k: &Bound<'_, PyAny>,
2489        v: &Bound<'_, PyAny>,
2490        length: usize,
2491        seed: u64,
2492    ) -> PyResult<Vec<Vec<f64>>> {
2493        let (q_data, q_rows, q_cols) = extract_matrix_f64(q, "Q")?;
2494        let (k_data, k_rows, k_cols) = extract_matrix_f64(k, "K")?;
2495        let (v_data, v_rows, v_cols) = extract_matrix_f64(v, "V")?;
2496
2497        let out = self
2498            .inner
2499            .forward_sc(
2500                &q_data, q_rows, q_cols, &k_data, k_rows, k_cols, &v_data, v_rows, v_cols, length,
2501                seed,
2502            )
2503            .map_err(PyValueError::new_err)?;
2504
2505        Ok(reshape_flat_to_rows(out, q_rows, v_cols))
2506    }
2507
2508    #[pyo3(signature = (q, k, v, n_heads))]
2509    fn forward_multihead(
2510        &self,
2511        q: &Bound<'_, PyAny>,
2512        k: &Bound<'_, PyAny>,
2513        v: &Bound<'_, PyAny>,
2514        n_heads: usize,
2515    ) -> PyResult<Vec<Vec<f64>>> {
2516        let (q_data, q_rows, q_cols) = extract_matrix_f64(q, "Q")?;
2517        let (k_data, k_rows, k_cols) = extract_matrix_f64(k, "K")?;
2518        let (v_data, v_rows, v_cols) = extract_matrix_f64(v, "V")?;
2519
2520        let out = self
2521            .inner
2522            .forward_multihead(
2523                &q_data, q_rows, q_cols, &k_data, k_rows, k_cols, &v_data, v_rows, v_cols, n_heads,
2524            )
2525            .map_err(PyValueError::new_err)?;
2526
2527        let out_cols = v_cols;
2528        Ok(reshape_flat_to_rows(out, q_rows, out_cols))
2529    }
2530}
2531
2532#[pyclass(
2533    name = "StochasticGraphLayer",
2534    module = "sc_neurocore_engine.sc_neurocore_engine"
2535)]
2536pub struct PyStochasticGraphLayer {
2537    inner: graph::StochasticGraphLayer,
2538}
2539
2540#[pymethods]
2541impl PyStochasticGraphLayer {
2542    #[new]
2543    #[pyo3(signature = (adj_matrix, n_features, seed=42))]
2544    fn new(adj_matrix: &Bound<'_, PyAny>, n_features: usize, seed: u64) -> PyResult<Self> {
2545        let (adj_flat, n_rows, n_cols) = extract_matrix_f64(adj_matrix, "adj_matrix")?;
2546        if n_rows != n_cols {
2547            return Err(PyValueError::new_err(format!(
2548                "adj_matrix must be square, got {}x{}.",
2549                n_rows, n_cols
2550            )));
2551        }
2552        Ok(Self {
2553            inner: graph::StochasticGraphLayer::new(adj_flat, n_rows, n_features, seed),
2554        })
2555    }
2556
2557    /// Construct from CSR arrays (row_offsets, col_indices, values).
2558    #[staticmethod]
2559    #[pyo3(signature = (row_offsets, col_indices, values, n_nodes, n_features, seed=42))]
2560    fn from_sparse(
2561        row_offsets: Vec<usize>,
2562        col_indices: Vec<usize>,
2563        values: Vec<f64>,
2564        n_nodes: usize,
2565        n_features: usize,
2566        seed: u64,
2567    ) -> PyResult<Self> {
2568        let csr = graph::CsrMatrix::new(row_offsets, col_indices, values, n_nodes, n_nodes)
2569            .map_err(PyValueError::new_err)?;
2570        let inner = graph::StochasticGraphLayer::new_sparse(csr, n_features, seed)
2571            .map_err(PyValueError::new_err)?;
2572        Ok(Self { inner })
2573    }
2574
2575    /// Dense adjacency with automatic CSR conversion if density < threshold.
2576    #[staticmethod]
2577    #[pyo3(signature = (adj_matrix, n_features, seed=42, density_threshold=0.3))]
2578    fn from_dense_auto(
2579        adj_matrix: &Bound<'_, PyAny>,
2580        n_features: usize,
2581        seed: u64,
2582        density_threshold: f64,
2583    ) -> PyResult<Self> {
2584        let (adj_flat, n_rows, n_cols) = extract_matrix_f64(adj_matrix, "adj_matrix")?;
2585        if n_rows != n_cols {
2586            return Err(PyValueError::new_err(format!(
2587                "adj_matrix must be square, got {}x{}.",
2588                n_rows, n_cols
2589            )));
2590        }
2591        Ok(Self {
2592            inner: graph::StochasticGraphLayer::from_dense_auto(
2593                adj_flat,
2594                n_rows,
2595                n_features,
2596                seed,
2597                density_threshold,
2598            ),
2599        })
2600    }
2601
2602    fn is_sparse(&self) -> bool {
2603        self.inner.is_sparse()
2604    }
2605
2606    fn forward(&self, node_features: &Bound<'_, PyAny>) -> PyResult<Vec<Vec<f64>>> {
2607        let (x_flat, x_rows, x_cols) = extract_matrix_f64(node_features, "node_features")?;
2608        if x_rows != self.inner.n_nodes || x_cols != self.inner.n_features {
2609            return Err(PyValueError::new_err(format!(
2610                "Expected node_features shape ({}, {}), got ({}, {}).",
2611                self.inner.n_nodes, self.inner.n_features, x_rows, x_cols
2612            )));
2613        }
2614        let out = self.inner.forward(&x_flat).map_err(PyValueError::new_err)?;
2615        Ok(reshape_flat_to_rows(
2616            out,
2617            self.inner.n_nodes,
2618            self.inner.n_features,
2619        ))
2620    }
2621
2622    #[pyo3(signature = (node_features, length=1024, seed=44257))]
2623    fn forward_sc(
2624        &self,
2625        node_features: &Bound<'_, PyAny>,
2626        length: usize,
2627        seed: u64,
2628    ) -> PyResult<Vec<Vec<f64>>> {
2629        let (x_flat, x_rows, x_cols) = extract_matrix_f64(node_features, "node_features")?;
2630        if x_rows != self.inner.n_nodes || x_cols != self.inner.n_features {
2631            return Err(PyValueError::new_err(format!(
2632                "Expected node_features shape ({}, {}), got ({}, {}).",
2633                self.inner.n_nodes, self.inner.n_features, x_rows, x_cols
2634            )));
2635        }
2636        let out = self
2637            .inner
2638            .forward_sc(&x_flat, length, seed)
2639            .map_err(PyValueError::new_err)?;
2640        Ok(reshape_flat_to_rows(
2641            out,
2642            self.inner.n_nodes,
2643            self.inner.n_features,
2644        ))
2645    }
2646
2647    fn get_weights(&self) -> Vec<f64> {
2648        self.inner.get_weights()
2649    }
2650
2651    fn set_weights(&mut self, weights: Vec<f64>) -> PyResult<()> {
2652        self.inner
2653            .set_weights(weights)
2654            .map_err(PyValueError::new_err)
2655    }
2656}
2657
2658#[pyclass(
2659    name = "KuramotoSolver",
2660    module = "sc_neurocore_engine.sc_neurocore_engine"
2661)]
2662pub struct PyKuramotoSolver {
2663    inner: scpn::KuramotoSolver,
2664}
2665
2666fn validate_kuramoto_finite(name: &str, values: &[f64]) -> PyResult<()> {
2667    if values.iter().all(|value| value.is_finite()) {
2668        Ok(())
2669    } else {
2670        Err(PyValueError::new_err(format!(
2671            "{name} values must be finite"
2672        )))
2673    }
2674}
2675
2676fn validate_kuramoto_dt(dt: f64) -> PyResult<()> {
2677    if dt.is_finite() && dt > 0.0 {
2678        Ok(())
2679    } else {
2680        Err(PyValueError::new_err("dt must be finite and positive"))
2681    }
2682}
2683
2684fn validate_kuramoto_matrix_shape(
2685    name: &str,
2686    values_len: usize,
2687    rows: usize,
2688    cols: usize,
2689    n: usize,
2690) -> PyResult<()> {
2691    let is_absent = rows == 0 && cols == 0 && values_len == 0;
2692    let is_flat = rows == 1 && values_len == n * n;
2693    let is_square = rows == n && cols == n;
2694    if is_absent || is_flat || is_square {
2695        Ok(())
2696    } else {
2697        Err(PyValueError::new_err(format!(
2698            "{name} must be shape ({n}, {n}) or flat length {}",
2699            n * n
2700        )))
2701    }
2702}
2703
2704#[pymethods]
2705impl PyKuramotoSolver {
2706    #[getter]
2707    fn phases(&self) -> Vec<f64> {
2708        self.inner.phases.clone()
2709    }
2710    #[new]
2711    #[pyo3(signature = (omega, coupling, phases, noise_amp=0.1))]
2712    fn new(
2713        omega: Vec<f64>,
2714        coupling: &Bound<'_, PyAny>,
2715        phases: Vec<f64>,
2716        noise_amp: f64,
2717    ) -> PyResult<Self> {
2718        let n = omega.len();
2719        if n == 0 {
2720            return Err(PyValueError::new_err("omega must not be empty."));
2721        }
2722        if phases.len() != n {
2723            return Err(PyValueError::new_err(format!(
2724                "phases length mismatch: got {}, expected {}.",
2725                phases.len(),
2726                n
2727            )));
2728        }
2729        validate_kuramoto_finite("omega", &omega)?;
2730        validate_kuramoto_finite("initial_phases", &phases)?;
2731        if !(noise_amp.is_finite() && noise_amp >= 0.0) {
2732            return Err(PyValueError::new_err(
2733                "noise_amp must be finite and non-negative",
2734            ));
2735        }
2736
2737        let (coupling_flat, rows, cols) = extract_matrix_f64(coupling, "coupling")?;
2738        if rows == 1 {
2739            if coupling_flat.len() != n * n {
2740                return Err(PyValueError::new_err(format!(
2741                    "Flat coupling length mismatch: got {}, expected {}.",
2742                    coupling_flat.len(),
2743                    n * n
2744                )));
2745            }
2746        } else if rows != n || cols != n {
2747            return Err(PyValueError::new_err(format!(
2748                "coupling must be shape ({}, {}) or flat length {}, got ({}, {}).",
2749                n,
2750                n,
2751                n * n,
2752                rows,
2753                cols
2754            )));
2755        }
2756        validate_kuramoto_finite("coupling", &coupling_flat)?;
2757
2758        Ok(Self {
2759            inner: scpn::KuramotoSolver::new(omega, coupling_flat, phases, noise_amp),
2760        })
2761    }
2762
2763    #[pyo3(signature = (dt, seed=0))]
2764    fn step(&mut self, dt: f64, seed: u64) -> PyResult<f64> {
2765        validate_kuramoto_dt(dt)?;
2766        Ok(self.inner.step(dt, seed))
2767    }
2768
2769    #[pyo3(signature = (n_steps, dt, seed=0))]
2770    fn run(&mut self, n_steps: usize, dt: f64, seed: u64) -> PyResult<Vec<f64>> {
2771        validate_kuramoto_dt(dt)?;
2772        Ok(self.inner.run(n_steps, dt, seed))
2773    }
2774
2775    fn set_field_pressure(&mut self, f: f64) -> PyResult<()> {
2776        if !f.is_finite() {
2777            return Err(PyValueError::new_err("field_pressure must be finite"));
2778        }
2779        self.inner.set_field_pressure(f);
2780        Ok(())
2781    }
2782
2783    #[pyo3(signature = (
2784        dt,
2785        seed=0,
2786        W=None,
2787        sigma_g=0.0,
2788        h_munu=None,
2789        pgbo_weight=0.0,
2790    ))]
2791    #[allow(non_snake_case)]
2792    fn step_ssgf(
2793        &mut self,
2794        dt: f64,
2795        seed: u64,
2796        W: Option<&Bound<'_, PyAny>>,
2797        sigma_g: f64,
2798        h_munu: Option<&Bound<'_, PyAny>>,
2799        pgbo_weight: f64,
2800    ) -> PyResult<f64> {
2801        validate_kuramoto_dt(dt)?;
2802        if !sigma_g.is_finite() {
2803            return Err(PyValueError::new_err("sigma_g must be finite"));
2804        }
2805        if !pgbo_weight.is_finite() {
2806            return Err(PyValueError::new_err("pgbo_weight must be finite"));
2807        }
2808        let (w_flat, w_rows, w_cols) = match W {
2809            Some(w) => extract_matrix_f64(w, "W")?,
2810            None => (vec![], 0, 0),
2811        };
2812        let (h_flat, h_rows, h_cols) = match h_munu {
2813            Some(h) => extract_matrix_f64(h, "h_munu")?,
2814            None => (vec![], 0, 0),
2815        };
2816        validate_kuramoto_matrix_shape("W", w_flat.len(), w_rows, w_cols, self.inner.n)?;
2817        validate_kuramoto_matrix_shape("h_munu", h_flat.len(), h_rows, h_cols, self.inner.n)?;
2818        validate_kuramoto_finite("w_flat", &w_flat)?;
2819        validate_kuramoto_finite("h_flat", &h_flat)?;
2820        Ok(self
2821            .inner
2822            .step_ssgf(dt, seed, &w_flat, sigma_g, &h_flat, pgbo_weight))
2823    }
2824
2825    #[pyo3(signature = (
2826        n_steps,
2827        dt,
2828        seed=0,
2829        W=None,
2830        sigma_g=0.0,
2831        h_munu=None,
2832        pgbo_weight=0.0,
2833    ))]
2834    #[allow(clippy::too_many_arguments, non_snake_case)]
2835    fn run_ssgf(
2836        &mut self,
2837        n_steps: usize,
2838        dt: f64,
2839        seed: u64,
2840        W: Option<&Bound<'_, PyAny>>,
2841        sigma_g: f64,
2842        h_munu: Option<&Bound<'_, PyAny>>,
2843        pgbo_weight: f64,
2844    ) -> PyResult<Vec<f64>> {
2845        validate_kuramoto_dt(dt)?;
2846        if !sigma_g.is_finite() {
2847            return Err(PyValueError::new_err("sigma_g must be finite"));
2848        }
2849        if !pgbo_weight.is_finite() {
2850            return Err(PyValueError::new_err("pgbo_weight must be finite"));
2851        }
2852        let (w_flat, w_rows, w_cols) = match W {
2853            Some(w) => extract_matrix_f64(w, "W")?,
2854            None => (vec![], 0, 0),
2855        };
2856        let (h_flat, h_rows, h_cols) = match h_munu {
2857            Some(h) => extract_matrix_f64(h, "h_munu")?,
2858            None => (vec![], 0, 0),
2859        };
2860        validate_kuramoto_matrix_shape("W", w_flat.len(), w_rows, w_cols, self.inner.n)?;
2861        validate_kuramoto_matrix_shape("h_munu", h_flat.len(), h_rows, h_cols, self.inner.n)?;
2862        validate_kuramoto_finite("w_flat", &w_flat)?;
2863        validate_kuramoto_finite("h_flat", &h_flat)?;
2864        Ok(self
2865            .inner
2866            .run_ssgf(n_steps, dt, seed, &w_flat, sigma_g, &h_flat, pgbo_weight))
2867    }
2868
2869    fn order_parameter(&self) -> f64 {
2870        self.inner.order_parameter()
2871    }
2872
2873    fn apply_phases(&mut self, phases: Vec<f64>) -> PyResult<()> {
2874        if phases.len() != self.inner.n {
2875            return Err(PyValueError::new_err(format!(
2876                "phases length mismatch: got {}, expected {}.",
2877                phases.len(),
2878                self.inner.n
2879            )));
2880        }
2881        validate_kuramoto_finite("phases", &phases)?;
2882        self.inner.set_phases(phases);
2883        Ok(())
2884    }
2885
2886    fn set_phases(&mut self, phases: Vec<f64>) -> PyResult<()> {
2887        self.apply_phases(phases)
2888    }
2889
2890    #[setter(phases)]
2891    fn set_phases_attr(&mut self, phases: Vec<f64>) -> PyResult<()> {
2892        self.apply_phases(phases)
2893    }
2894
2895    fn set_coupling(&mut self, coupling: &Bound<'_, PyAny>) -> PyResult<()> {
2896        let n = self.inner.n;
2897        let (coupling_flat, rows, cols) = extract_matrix_f64(coupling, "coupling")?;
2898        if rows == 1 {
2899            if coupling_flat.len() != n * n {
2900                return Err(PyValueError::new_err(format!(
2901                    "Flat coupling length mismatch: got {}, expected {}.",
2902                    coupling_flat.len(),
2903                    n * n
2904                )));
2905            }
2906        } else if rows != n || cols != n {
2907            return Err(PyValueError::new_err(format!(
2908                "coupling must be shape ({}, {}) or flat length {}, got ({}, {}).",
2909                n,
2910                n,
2911                n * n,
2912                rows,
2913                cols
2914            )));
2915        }
2916        validate_kuramoto_finite("coupling", &coupling_flat)?;
2917        self.inner.set_coupling(coupling_flat);
2918        Ok(())
2919    }
2920}
2921
2922#[pyclass(
2923    name = "SCPNMetrics",
2924    module = "sc_neurocore_engine.sc_neurocore_engine"
2925)]
2926pub struct PySCPNMetrics;
2927
2928#[pymethods]
2929impl PySCPNMetrics {
2930    #[new]
2931    fn new() -> Self {
2932        Self
2933    }
2934
2935    #[staticmethod]
2936    fn global_coherence(weights: [f64; 7], metrics: [f64; 7]) -> f64 {
2937        scpn::SCPNMetrics::global_coherence(&weights, &metrics)
2938    }
2939
2940    #[staticmethod]
2941    fn consciousness_index(phases_l4: Vec<f64>, glyph_l7: [f64; 6]) -> f64 {
2942        scpn::SCPNMetrics::consciousness_index(&phases_l4, &glyph_l7)
2943    }
2944}
2945
2946// IR bridge
2947
2948#[pyclass(name = "ScGraph", module = "sc_neurocore_engine.sc_neurocore_engine")]
2949pub struct PyScGraph {
2950    inner: ir::graph::ScGraph,
2951}
2952
2953#[pymethods]
2954impl PyScGraph {
2955    /// Number of operations in the graph.
2956    fn len(&self) -> usize {
2957        self.inner.len()
2958    }
2959
2960    fn __len__(&self) -> usize {
2961        self.inner.len()
2962    }
2963
2964    /// Whether the graph is empty.
2965    fn is_empty(&self) -> bool {
2966        self.inner.is_empty()
2967    }
2968
2969    /// Graph name.
2970    #[getter]
2971    fn name(&self) -> &str {
2972        &self.inner.name
2973    }
2974
2975    /// Number of input ports.
2976    fn num_inputs(&self) -> usize {
2977        self.inner.inputs().len()
2978    }
2979
2980    /// Number of output ports.
2981    fn num_outputs(&self) -> usize {
2982        self.inner.outputs().len()
2983    }
2984
2985    fn __repr__(&self) -> String {
2986        format!("ScGraph('{}', ops={})", self.inner.name, self.inner.len())
2987    }
2988}
2989
2990#[pyclass(
2991    name = "ScGraphBuilder",
2992    module = "sc_neurocore_engine.sc_neurocore_engine"
2993)]
2994pub struct PyScGraphBuilder {
2995    inner: Option<ir::builder::ScGraphBuilder>,
2996}
2997
2998impl PyScGraphBuilder {
2999    fn builder_mut(&mut self) -> PyResult<&mut ir::builder::ScGraphBuilder> {
3000        self.inner
3001            .as_mut()
3002            .ok_or_else(|| PyValueError::new_err("Builder already consumed by build()."))
3003    }
3004}
3005
3006#[pymethods]
3007impl PyScGraphBuilder {
3008    #[new]
3009    fn new(name: String) -> Self {
3010        Self {
3011            inner: Some(ir::builder::ScGraphBuilder::new(name)),
3012        }
3013    }
3014
3015    /// Add a typed input port. Returns value ID.
3016    fn input(&mut self, name: &str, ty: &str) -> PyResult<u32> {
3017        let sc_type = parse_sc_type(ty)?;
3018        Ok(self.builder_mut()?.input(name, sc_type).0)
3019    }
3020
3021    /// Add an output port forwarding a value.
3022    fn output(&mut self, name: &str, source_id: u32) -> PyResult<u32> {
3023        Ok(self
3024            .builder_mut()?
3025            .output(name, ir::graph::ValueId(source_id))
3026            .0)
3027    }
3028
3029    /// Add a float constant.
3030    fn constant_f64(&mut self, value: f64, ty: &str) -> PyResult<u32> {
3031        let sc_type = parse_sc_type(ty)?;
3032        Ok(self
3033            .builder_mut()?
3034            .constant(ir::graph::ScConst::F64(value), sc_type)
3035            .0)
3036    }
3037
3038    /// Add an integer constant.
3039    fn constant_i64(&mut self, value: i64, ty: &str) -> PyResult<u32> {
3040        let sc_type = parse_sc_type(ty)?;
3041        Ok(self
3042            .builder_mut()?
3043            .constant(ir::graph::ScConst::I64(value), sc_type)
3044            .0)
3045    }
3046
3047    /// Add a Bernoulli encode operation.
3048    fn encode(&mut self, prob_id: u32, length: usize, seed: u64) -> PyResult<u32> {
3049        let seed = u16::try_from(seed)
3050            .map_err(|_| PyValueError::new_err(format!("Seed out of range for u16: {seed}")))?;
3051        Ok(self
3052            .builder_mut()?
3053            .encode(ir::graph::ValueId(prob_id), length, seed)
3054            .0)
3055    }
3056
3057    /// Add a bitwise AND (SC multiply).
3058    fn bitwise_and(&mut self, lhs_id: u32, rhs_id: u32) -> PyResult<u32> {
3059        Ok(self
3060            .builder_mut()?
3061            .bitwise_and(ir::graph::ValueId(lhs_id), ir::graph::ValueId(rhs_id))
3062            .0)
3063    }
3064
3065    /// Add a popcount operation.
3066    fn popcount(&mut self, input_id: u32) -> PyResult<u32> {
3067        Ok(self.builder_mut()?.popcount(ir::graph::ValueId(input_id)).0)
3068    }
3069
3070    /// Add a LIF neuron step.
3071    #[pyo3(signature = (
3072        current_id,
3073        leak_id,
3074        gain_id,
3075        noise_id,
3076        data_width=16,
3077        fraction=8,
3078        v_rest=0,
3079        v_reset=0,
3080        v_threshold=256,
3081        refractory_period=2
3082    ))]
3083    #[allow(clippy::too_many_arguments)]
3084    fn lif_step(
3085        &mut self,
3086        current_id: u32,
3087        leak_id: u32,
3088        gain_id: u32,
3089        noise_id: u32,
3090        data_width: u32,
3091        fraction: u32,
3092        v_rest: i64,
3093        v_reset: i64,
3094        v_threshold: i64,
3095        refractory_period: u32,
3096    ) -> PyResult<u32> {
3097        let params = ir::graph::LifParams {
3098            data_width,
3099            fraction,
3100            v_rest,
3101            v_reset,
3102            v_threshold,
3103            refractory_period,
3104        };
3105        Ok(self
3106            .builder_mut()?
3107            .lif_step(
3108                ir::graph::ValueId(current_id),
3109                ir::graph::ValueId(leak_id),
3110                ir::graph::ValueId(gain_id),
3111                ir::graph::ValueId(noise_id),
3112                params,
3113            )
3114            .0)
3115    }
3116
3117    /// Add a dense layer forward pass.
3118    #[pyo3(signature = (
3119        inputs_id,
3120        weights_id,
3121        leak_id,
3122        gain_id,
3123        n_inputs=3,
3124        n_neurons=7,
3125        data_width=16,
3126        stream_length=1024,
3127        seed_base=0xACE1u64,
3128        y_min=0,
3129        y_max=65535
3130    ))]
3131    #[allow(clippy::too_many_arguments)]
3132    fn dense_forward(
3133        &mut self,
3134        inputs_id: u32,
3135        weights_id: u32,
3136        leak_id: u32,
3137        gain_id: u32,
3138        n_inputs: usize,
3139        n_neurons: usize,
3140        data_width: u32,
3141        stream_length: usize,
3142        seed_base: u64,
3143        y_min: i64,
3144        y_max: i64,
3145    ) -> PyResult<u32> {
3146        let input_seed_base = u16::try_from(seed_base).map_err(|_| {
3147            PyValueError::new_err(format!("seed_base out of range for u16: {seed_base}"))
3148        })?;
3149        let params = ir::graph::DenseParams {
3150            n_inputs,
3151            n_neurons,
3152            data_width,
3153            stream_length,
3154            input_seed_base,
3155            weight_seed_base: input_seed_base.wrapping_add(1),
3156            y_min,
3157            y_max,
3158        };
3159        Ok(self
3160            .builder_mut()?
3161            .dense_forward(
3162                ir::graph::ValueId(inputs_id),
3163                ir::graph::ValueId(weights_id),
3164                ir::graph::ValueId(leak_id),
3165                ir::graph::ValueId(gain_id),
3166                params,
3167            )
3168            .0)
3169    }
3170
3171    /// Add a scale (multiply by constant factor) operation.
3172    fn scale(&mut self, input_id: u32, factor: f64) -> PyResult<u32> {
3173        Ok(self
3174            .builder_mut()?
3175            .scale(ir::graph::ValueId(input_id), factor)
3176            .0)
3177    }
3178
3179    /// Add an offset (add constant) operation.
3180    fn offset(&mut self, input_id: u32, offset_val: f64) -> PyResult<u32> {
3181        Ok(self
3182            .builder_mut()?
3183            .offset(ir::graph::ValueId(input_id), offset_val)
3184            .0)
3185    }
3186
3187    /// Add a divide-by-constant operation.
3188    fn div_const(&mut self, input_id: u32, divisor: u64) -> PyResult<u32> {
3189        Ok(self
3190            .builder_mut()?
3191            .div_const(ir::graph::ValueId(input_id), divisor)
3192            .0)
3193    }
3194
3195    /// Consume the builder and return a graph.
3196    fn build(&mut self) -> PyResult<PyScGraph> {
3197        let builder = self
3198            .inner
3199            .take()
3200            .ok_or_else(|| PyValueError::new_err("Builder already consumed by build()."))?;
3201        Ok(PyScGraph {
3202            inner: builder.build(),
3203        })
3204    }
3205}
3206
3207/// Verify an IR graph. Returns None on success, or a list of error strings.
3208#[pyfunction]
3209fn ir_verify(graph: PyRef<'_, PyScGraph>) -> Option<Vec<String>> {
3210    match ir::verify::verify(&graph.inner) {
3211        Ok(()) => None,
3212        Err(errors) => Some(errors.iter().map(|e| e.to_string()).collect()),
3213    }
3214}
3215
3216/// Print an IR graph to its stable text format.
3217#[pyfunction]
3218fn ir_print(graph: PyRef<'_, PyScGraph>) -> String {
3219    ir::printer::print(&graph.inner)
3220}
3221
3222/// Parse an IR graph from text format.
3223#[pyfunction]
3224fn ir_parse(text: &str) -> PyResult<PyScGraph> {
3225    ir::parser::parse(text)
3226        .map(|graph| PyScGraph { inner: graph })
3227        .map_err(|e| PyValueError::new_err(e.to_string()))
3228}
3229
3230/// Emit SystemVerilog from an IR graph.
3231#[pyfunction]
3232fn ir_emit_sv(graph: PyRef<'_, PyScGraph>) -> PyResult<String> {
3233    ir::emit_sv::emit(&graph.inner).map_err(PyValueError::new_err)
3234}
3235
3236/// Parse a Python type string into ScType.
3237///
3238/// Accepted formats: "bool", "rate", "u32", "u64", "i16", "i32",
3239/// "bitstream", "bitstream<1024>", "fixed<16,8>", "vec<bool,7>".
3240fn parse_sc_type(s: &str) -> PyResult<ir::graph::ScType> {
3241    let s = s.trim();
3242    let lower = s.to_ascii_lowercase();
3243    match lower.as_str() {
3244        "bool" => Ok(ir::graph::ScType::Bool),
3245        "rate" => Ok(ir::graph::ScType::Rate),
3246        "u32" => Ok(ir::graph::ScType::UInt { width: 32 }),
3247        "u64" => Ok(ir::graph::ScType::UInt { width: 64 }),
3248        "i16" => Ok(ir::graph::ScType::SInt { width: 16 }),
3249        "i32" => Ok(ir::graph::ScType::SInt { width: 32 }),
3250        "bitstream" => Ok(ir::graph::ScType::Bitstream { length: 0 }),
3251        _ => {
3252            if let Some(width) = lower.strip_prefix('u') {
3253                if let Ok(width) = width.parse::<u32>() {
3254                    return Ok(ir::graph::ScType::UInt { width });
3255                }
3256            }
3257            if let Some(width) = lower.strip_prefix('i') {
3258                if let Ok(width) = width.parse::<u32>() {
3259                    return Ok(ir::graph::ScType::SInt { width });
3260                }
3261            }
3262            if let Some(inner) = lower
3263                .strip_prefix("bitstream<")
3264                .and_then(|r| r.strip_suffix('>'))
3265            {
3266                let length = inner.parse::<usize>().map_err(|_| {
3267                    PyValueError::new_err(format!("Invalid bitstream length: '{inner}'"))
3268                })?;
3269                return Ok(ir::graph::ScType::Bitstream { length });
3270            }
3271            if let Some(inner) = lower
3272                .strip_prefix("fixed<")
3273                .and_then(|r| r.strip_suffix('>'))
3274            {
3275                let parts: Vec<&str> = inner.split(',').collect();
3276                if parts.len() != 2 {
3277                    return Err(PyValueError::new_err(format!(
3278                        "fixed type needs 2 params: '{s}'"
3279                    )));
3280                }
3281                let width = parts[0].trim().parse::<u32>().map_err(|_| {
3282                    PyValueError::new_err(format!("Invalid fixed width: '{}'", parts[0]))
3283                })?;
3284                let frac = parts[1].trim().parse::<u32>().map_err(|_| {
3285                    PyValueError::new_err(format!("Invalid fixed frac: '{}'", parts[1]))
3286                })?;
3287                return Ok(ir::graph::ScType::FixedPoint { width, frac });
3288            }
3289            if let Some(inner) = lower.strip_prefix("vec<").and_then(|r| r.strip_suffix('>')) {
3290                if let Some(comma_pos) = inner.rfind(',') {
3291                    let inner_ty_str = &inner[..comma_pos];
3292                    let count_str = inner[comma_pos + 1..].trim();
3293                    let inner_ty = parse_sc_type(inner_ty_str)?;
3294                    let count = count_str.parse::<usize>().map_err(|_| {
3295                        PyValueError::new_err(format!("Invalid vec count: '{count_str}'"))
3296                    })?;
3297                    return Ok(ir::graph::ScType::Vec {
3298                        element: Box::new(inner_ty),
3299                        count,
3300                    });
3301                }
3302            }
3303            Err(PyValueError::new_err(format!("Unknown IR type: '{s}'")))
3304        }
3305    }
3306}
3307
3308// ── Analysis PyO3 wrappers (P0-A: spike_stats) ─────────────────────
3309
3310#[pyfunction]
3311#[pyo3(signature = (binary_train, dt=0.001))]
3312fn py_spike_times(
3313    py: Python<'_>,
3314    binary_train: PyReadonlyArray1<'_, i32>,
3315    dt: f64,
3316) -> Py<PyArray1<f64>> {
3317    let data = binary_train.as_slice().unwrap();
3318    analysis::basic::spike_times(data, dt)
3319        .into_pyarray(py)
3320        .into()
3321}
3322
3323#[pyfunction]
3324#[pyo3(signature = (binary_train, dt=0.001))]
3325fn py_isi(py: Python<'_>, binary_train: PyReadonlyArray1<'_, i32>, dt: f64) -> Py<PyArray1<f64>> {
3326    let data = binary_train.as_slice().unwrap();
3327    analysis::basic::isi(data, dt).into_pyarray(py).into()
3328}
3329
3330#[pyfunction]
3331#[pyo3(signature = (binary_train, dt=0.001))]
3332fn py_firing_rate(binary_train: PyReadonlyArray1<'_, i32>, dt: f64) -> f64 {
3333    let data = binary_train.as_slice().unwrap();
3334    analysis::basic::firing_rate(data, dt)
3335}
3336
3337#[pyfunction]
3338fn py_spike_count(binary_train: PyReadonlyArray1<'_, i32>) -> i64 {
3339    let data = binary_train.as_slice().unwrap();
3340    analysis::basic::spike_count(data)
3341}
3342
3343#[pyfunction]
3344#[pyo3(signature = (binary_train, bin_size=10))]
3345fn py_bin_spike_train(
3346    py: Python<'_>,
3347    binary_train: PyReadonlyArray1<'_, i32>,
3348    bin_size: usize,
3349) -> Py<PyArray1<i64>> {
3350    let data = binary_train.as_slice().unwrap();
3351    analysis::basic::bin_spike_train(data, bin_size)
3352        .into_pyarray(py)
3353        .into()
3354}
3355
3356#[pyfunction]
3357#[pyo3(signature = (binary_train, dt=0.001, kernel="gaussian", sigma_ms=10.0))]
3358fn py_instantaneous_rate(
3359    py: Python<'_>,
3360    binary_train: PyReadonlyArray1<'_, f64>,
3361    dt: f64,
3362    kernel: &str,
3363    sigma_ms: f64,
3364) -> Py<PyArray1<f64>> {
3365    let data = binary_train.as_slice().unwrap();
3366    analysis::rate::instantaneous_rate(data, dt, kernel, sigma_ms)
3367        .into_pyarray(py)
3368        .into()
3369}
3370
3371#[pyfunction]
3372#[pyo3(signature = (trials, bin_ms=10.0, dt=0.001))]
3373fn py_psth(
3374    py: Python<'_>,
3375    trials: Vec<PyReadonlyArray1<'_, f64>>,
3376    bin_ms: f64,
3377    dt: f64,
3378) -> (Py<PyArray1<f64>>, Py<PyArray1<f64>>) {
3379    let vecs: Vec<Vec<f64>> = trials
3380        .iter()
3381        .map(|t| t.as_slice().unwrap().to_vec())
3382        .collect();
3383    let (rates, centers) = analysis::rate::psth(&vecs, bin_ms, dt);
3384    (
3385        rates.into_pyarray(py).into(),
3386        centers.into_pyarray(py).into(),
3387    )
3388}
3389
3390#[pyfunction]
3391#[pyo3(signature = (binary_train, dt=0.001))]
3392fn py_cv_isi(binary_train: PyReadonlyArray1<'_, i32>, dt: f64) -> f64 {
3393    let data = binary_train.as_slice().unwrap();
3394    analysis::variability::cv_isi(data, dt)
3395}
3396
3397#[pyfunction]
3398#[pyo3(signature = (binary_train, dt=0.001))]
3399fn py_cv2(binary_train: PyReadonlyArray1<'_, i32>, dt: f64) -> f64 {
3400    let data = binary_train.as_slice().unwrap();
3401    analysis::variability::cv2(data, dt)
3402}
3403
3404#[pyfunction]
3405#[pyo3(signature = (binary_train, dt=0.001))]
3406fn py_local_variation(binary_train: PyReadonlyArray1<'_, i32>, dt: f64) -> f64 {
3407    let data = binary_train.as_slice().unwrap();
3408    analysis::variability::local_variation(data, dt)
3409}
3410
3411#[pyfunction]
3412#[pyo3(signature = (binary_train, window_ms=50.0, dt=0.001))]
3413fn py_fano_factor(binary_train: PyReadonlyArray1<'_, i32>, window_ms: f64, dt: f64) -> f64 {
3414    let data = binary_train.as_slice().unwrap();
3415    analysis::variability::fano_factor(data, window_ms, dt)
3416}
3417
3418#[pyfunction]
3419fn py_lempel_ziv_complexity(binary_train: PyReadonlyArray1<'_, i32>) -> f64 {
3420    let data = binary_train.as_slice().unwrap();
3421    analysis::variability::lempel_ziv_complexity(data)
3422}
3423
3424#[pyfunction]
3425#[pyo3(signature = (binary_train, order=3, delay=1))]
3426fn py_permutation_entropy(
3427    binary_train: PyReadonlyArray1<'_, i32>,
3428    order: usize,
3429    delay: usize,
3430) -> f64 {
3431    let data = binary_train.as_slice().unwrap();
3432    analysis::variability::permutation_entropy(data, order, delay)
3433}
3434
3435#[pyfunction]
3436#[pyo3(signature = (binary_train, min_window=10))]
3437fn py_hurst_exponent(binary_train: PyReadonlyArray1<'_, i32>, min_window: usize) -> f64 {
3438    let data = binary_train.as_slice().unwrap();
3439    analysis::variability::hurst_exponent(data, min_window)
3440}
3441
3442#[pyfunction]
3443#[pyo3(signature = (binary_train, m=2, r_factor=0.2))]
3444fn py_approximate_entropy(binary_train: PyReadonlyArray1<'_, i32>, m: usize, r_factor: f64) -> f64 {
3445    let data = binary_train.as_slice().unwrap();
3446    analysis::variability::approximate_entropy(data, m, r_factor)
3447}
3448
3449#[pyfunction]
3450#[pyo3(signature = (binary_train, m=2, r_factor=0.2))]
3451fn py_sample_entropy(binary_train: PyReadonlyArray1<'_, i32>, m: usize, r_factor: f64) -> f64 {
3452    let data = binary_train.as_slice().unwrap();
3453    analysis::variability::sample_entropy(data, m, r_factor)
3454}
3455
3456// ── Correlation PyO3 wrappers (P0-A: spike_stats/correlation) ────
3457
3458#[pyfunction]
3459#[pyo3(signature = (train_a, train_b, max_lag_ms=50.0, dt=0.001))]
3460fn py_cross_correlation(
3461    py: Python<'_>,
3462    train_a: PyReadonlyArray1<'_, i32>,
3463    train_b: PyReadonlyArray1<'_, i32>,
3464    max_lag_ms: f64,
3465    dt: f64,
3466) -> (Py<PyArray1<f64>>, Py<PyArray1<f64>>) {
3467    let a = train_a.as_slice().unwrap();
3468    let b = train_b.as_slice().unwrap();
3469    let (cc, lags) = analysis::correlation::cross_correlation(a, b, max_lag_ms, dt);
3470    (cc.into_pyarray(py).into(), lags.into_pyarray(py).into())
3471}
3472
3473#[pyfunction]
3474#[pyo3(signature = (trains, dt=0.001))]
3475fn py_pairwise_correlation(
3476    py: Python<'_>,
3477    trains: Vec<PyReadonlyArray1<'_, i32>>,
3478    dt: f64,
3479) -> Py<PyArray2<f64>> {
3480    let vecs: Vec<Vec<i32>> = trains
3481        .iter()
3482        .map(|t| t.as_slice().unwrap().to_vec())
3483        .collect();
3484    let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
3485    let mat = analysis::correlation::pairwise_correlation(&refs, dt);
3486    let n = mat.len();
3487    let flat: Vec<f64> = mat.into_iter().flatten().collect();
3488    numpy::PyArray2::from_vec2(py, &flat.chunks(n).map(|c| c.to_vec()).collect::<Vec<_>>())
3489        .unwrap()
3490        .into()
3491}
3492
3493#[pyfunction]
3494#[pyo3(signature = (train_a, train_b, dt=0.001, tau_ms=5.0))]
3495fn py_event_synchronization(
3496    train_a: PyReadonlyArray1<'_, i32>,
3497    train_b: PyReadonlyArray1<'_, i32>,
3498    dt: f64,
3499    tau_ms: f64,
3500) -> f64 {
3501    let a = train_a.as_slice().unwrap();
3502    let b = train_b.as_slice().unwrap();
3503    analysis::correlation::event_synchronization(a, b, dt, tau_ms)
3504}
3505
3506#[pyfunction]
3507#[pyo3(signature = (train_a, train_b, dt=0.001))]
3508fn py_spike_train_coherence(
3509    py: Python<'_>,
3510    train_a: PyReadonlyArray1<'_, i32>,
3511    train_b: PyReadonlyArray1<'_, i32>,
3512    dt: f64,
3513) -> (Py<PyArray1<f64>>, Py<PyArray1<f64>>) {
3514    let a = train_a.as_slice().unwrap();
3515    let b = train_b.as_slice().unwrap();
3516    let (coh, freqs) = analysis::correlation::spike_train_coherence(a, b, dt);
3517    (coh.into_pyarray(py).into(), freqs.into_pyarray(py).into())
3518}
3519
3520#[pyfunction]
3521#[pyo3(signature = (train_a, train_b, dt=0.001, delta_ms=5.0))]
3522fn py_spike_time_tiling_coefficient(
3523    train_a: PyReadonlyArray1<'_, i32>,
3524    train_b: PyReadonlyArray1<'_, i32>,
3525    dt: f64,
3526    delta_ms: f64,
3527) -> f64 {
3528    let a = train_a.as_slice().unwrap();
3529    let b = train_b.as_slice().unwrap();
3530    analysis::correlation::spike_time_tiling_coefficient(a, b, dt, delta_ms)
3531}
3532
3533#[pyfunction]
3534#[pyo3(signature = (trains, bin_size=10))]
3535fn py_covariance_matrix(
3536    py: Python<'_>,
3537    trains: Vec<PyReadonlyArray1<'_, i32>>,
3538    bin_size: usize,
3539) -> Py<PyArray2<f64>> {
3540    let vecs: Vec<Vec<i32>> = trains
3541        .iter()
3542        .map(|t| t.as_slice().unwrap().to_vec())
3543        .collect();
3544    let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
3545    let mat = analysis::correlation::covariance_matrix(&refs, bin_size);
3546    let n = mat.len();
3547    let rows: Vec<Vec<f64>> = mat.into_iter().collect();
3548    numpy::PyArray2::from_vec2(py, &rows)
3549        .unwrap_or_else(|_| numpy::PyArray2::zeros(py, [n, n], false))
3550        .into()
3551}
3552
3553#[pyfunction]
3554#[pyo3(signature = (binary_train, dt=0.001, max_lag_ms=100.0))]
3555fn py_autocorrelation_time(
3556    binary_train: PyReadonlyArray1<'_, i32>,
3557    dt: f64,
3558    max_lag_ms: f64,
3559) -> f64 {
3560    let data = binary_train.as_slice().unwrap();
3561    analysis::correlation::autocorrelation_time(data, dt, max_lag_ms)
3562}
3563
3564#[pyfunction]
3565#[pyo3(signature = (trains, bin_size=50))]
3566fn py_noise_correlation(
3567    py: Python<'_>,
3568    trains: Vec<PyReadonlyArray1<'_, i32>>,
3569    bin_size: usize,
3570) -> Py<PyArray2<f64>> {
3571    let vecs: Vec<Vec<i32>> = trains
3572        .iter()
3573        .map(|t| t.as_slice().unwrap().to_vec())
3574        .collect();
3575    let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
3576    let mat = analysis::correlation::noise_correlation(&refs, bin_size);
3577    let rows: Vec<Vec<f64>> = mat.into_iter().collect();
3578    let n = rows.len();
3579    numpy::PyArray2::from_vec2(py, &rows)
3580        .unwrap_or_else(|_| numpy::PyArray2::zeros(py, [n, n], false))
3581        .into()
3582}
3583
3584#[pyfunction]
3585#[pyo3(signature = (trains, bin_size=50))]
3586fn py_signal_correlation(
3587    py: Python<'_>,
3588    trains: Vec<PyReadonlyArray1<'_, i32>>,
3589    bin_size: usize,
3590) -> Py<PyArray2<f64>> {
3591    let vecs: Vec<Vec<i32>> = trains
3592        .iter()
3593        .map(|t| t.as_slice().unwrap().to_vec())
3594        .collect();
3595    let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
3596    let mat = analysis::correlation::signal_correlation(&refs, bin_size);
3597    let rows: Vec<Vec<f64>> = mat.into_iter().collect();
3598    let n = rows.len();
3599    numpy::PyArray2::from_vec2(py, &rows)
3600        .unwrap_or_else(|_| numpy::PyArray2::zeros(py, [n, n], false))
3601        .into()
3602}
3603
3604#[pyfunction]
3605#[pyo3(signature = (trains, window=50))]
3606fn py_spike_count_covariance(
3607    py: Python<'_>,
3608    trains: Vec<PyReadonlyArray1<'_, i32>>,
3609    window: usize,
3610) -> Py<PyArray2<f64>> {
3611    let vecs: Vec<Vec<i32>> = trains
3612        .iter()
3613        .map(|t| t.as_slice().unwrap().to_vec())
3614        .collect();
3615    let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
3616    let mat = analysis::correlation::spike_count_covariance(&refs, window);
3617    let rows: Vec<Vec<f64>> = mat.into_iter().collect();
3618    let n = rows.len();
3619    numpy::PyArray2::from_vec2(py, &rows)
3620        .unwrap_or_else(|_| numpy::PyArray2::zeros(py, [n, n], false))
3621        .into()
3622}
3623
3624#[pyfunction]
3625#[pyo3(signature = (train_a, train_b, bin_size=10))]
3626fn py_joint_psth(
3627    py: Python<'_>,
3628    train_a: PyReadonlyArray1<'_, i32>,
3629    train_b: PyReadonlyArray1<'_, i32>,
3630    bin_size: usize,
3631) -> Py<PyArray2<f64>> {
3632    let a = train_a.as_slice().unwrap();
3633    let b = train_b.as_slice().unwrap();
3634    let (flat, n) = analysis::correlation::joint_psth(a, b, bin_size);
3635    if n == 0 {
3636        return numpy::PyArray2::zeros(py, [0, 0], false).into();
3637    }
3638    let rows: Vec<Vec<f64>> = flat.chunks(n).map(|c| c.to_vec()).collect();
3639    numpy::PyArray2::from_vec2(py, &rows).unwrap().into()
3640}
3641
3642#[pyfunction]
3643#[pyo3(signature = (train_a, train_b, dt=0.001, delta_ms=2.0))]
3644fn py_coincidence_index(
3645    train_a: PyReadonlyArray1<'_, i32>,
3646    train_b: PyReadonlyArray1<'_, i32>,
3647    dt: f64,
3648    delta_ms: f64,
3649) -> f64 {
3650    let a = train_a.as_slice().unwrap();
3651    let b = train_b.as_slice().unwrap();
3652    analysis::correlation::coincidence_index(a, b, dt, delta_ms)
3653}
3654
3655// ── Distance PyO3 wrappers (P0-A: spike_stats/distance) ─────────
3656
3657#[pyfunction]
3658#[pyo3(signature = (train_a, train_b, dt=0.001, tau_ms=10.0))]
3659fn py_van_rossum_distance(
3660    train_a: PyReadonlyArray1<'_, i32>,
3661    train_b: PyReadonlyArray1<'_, i32>,
3662    dt: f64,
3663    tau_ms: f64,
3664) -> f64 {
3665    let a = train_a.as_slice().unwrap();
3666    let b = train_b.as_slice().unwrap();
3667    analysis::distance::van_rossum_distance(a, b, dt, tau_ms)
3668}
3669
3670#[pyfunction]
3671#[pyo3(signature = (times_a, times_b, cost_per_s=1000.0))]
3672fn py_victor_purpura_distance(
3673    times_a: PyReadonlyArray1<'_, f64>,
3674    times_b: PyReadonlyArray1<'_, f64>,
3675    cost_per_s: f64,
3676) -> f64 {
3677    let a = times_a.as_slice().unwrap();
3678    let b = times_b.as_slice().unwrap();
3679    analysis::distance::victor_purpura_distance(a, b, cost_per_s)
3680}
3681
3682#[pyfunction]
3683#[pyo3(signature = (train_a, train_b, dt=0.001))]
3684fn py_isi_distance(
3685    train_a: PyReadonlyArray1<'_, i32>,
3686    train_b: PyReadonlyArray1<'_, i32>,
3687    dt: f64,
3688) -> f64 {
3689    let a = train_a.as_slice().unwrap();
3690    let b = train_b.as_slice().unwrap();
3691    analysis::distance::isi_distance(a, b, dt)
3692}
3693
3694#[pyfunction]
3695#[pyo3(signature = (times_a, times_b, t_start=0.0, t_end=1.0))]
3696fn py_spike_distance(
3697    times_a: PyReadonlyArray1<'_, f64>,
3698    times_b: PyReadonlyArray1<'_, f64>,
3699    t_start: f64,
3700    t_end: f64,
3701) -> f64 {
3702    let a = times_a.as_slice().unwrap();
3703    let b = times_b.as_slice().unwrap();
3704    analysis::distance::spike_distance(a, b, t_start, t_end)
3705}
3706
3707#[pyfunction]
3708#[pyo3(signature = (times_a, times_b, t_start=0.0, t_end=1.0))]
3709fn py_spike_sync(
3710    times_a: PyReadonlyArray1<'_, f64>,
3711    times_b: PyReadonlyArray1<'_, f64>,
3712    t_start: f64,
3713    t_end: f64,
3714) -> f64 {
3715    let a = times_a.as_slice().unwrap();
3716    let b = times_b.as_slice().unwrap();
3717    analysis::distance::spike_sync(a, b, t_start, t_end)
3718}
3719
3720#[pyfunction]
3721#[pyo3(signature = (times_a, times_b, n_bins=50, t_start=0.0, t_end=1.0))]
3722fn py_spike_sync_profile(
3723    py: Python<'_>,
3724    times_a: PyReadonlyArray1<'_, f64>,
3725    times_b: PyReadonlyArray1<'_, f64>,
3726    n_bins: usize,
3727    t_start: f64,
3728    t_end: f64,
3729) -> Py<PyArray1<f64>> {
3730    let a = times_a.as_slice().unwrap();
3731    let b = times_b.as_slice().unwrap();
3732    analysis::distance::spike_sync_profile(a, b, n_bins, t_start, t_end)
3733        .into_pyarray(py)
3734        .into()
3735}
3736
3737#[pyfunction]
3738#[pyo3(signature = (times_a, times_b, n_bins=50, t_start=0.0, t_end=1.0))]
3739fn py_spike_profile(
3740    py: Python<'_>,
3741    times_a: PyReadonlyArray1<'_, f64>,
3742    times_b: PyReadonlyArray1<'_, f64>,
3743    n_bins: usize,
3744    t_start: f64,
3745    t_end: f64,
3746) -> Py<PyArray1<f64>> {
3747    let a = times_a.as_slice().unwrap();
3748    let b = times_b.as_slice().unwrap();
3749    analysis::distance::spike_profile(a, b, n_bins, t_start, t_end)
3750        .into_pyarray(py)
3751        .into()
3752}
3753
3754#[pyfunction]
3755#[pyo3(signature = (binary_train_a, binary_train_b, dt=0.001, n_bins=50))]
3756fn py_isi_profile(
3757    py: Python<'_>,
3758    binary_train_a: PyReadonlyArray1<'_, i32>,
3759    binary_train_b: PyReadonlyArray1<'_, i32>,
3760    dt: f64,
3761    n_bins: usize,
3762) -> Py<PyArray1<f64>> {
3763    let a = binary_train_a.as_slice().unwrap();
3764    let b = binary_train_b.as_slice().unwrap();
3765    analysis::distance::isi_profile(a, b, dt, n_bins)
3766        .into_pyarray(py)
3767        .into()
3768}
3769
3770#[pyfunction]
3771#[pyo3(signature = (times_a, times_b, t_start=0.0, t_end=1.0, cost=0.0))]
3772fn py_adaptive_spike_distance(
3773    times_a: PyReadonlyArray1<'_, f64>,
3774    times_b: PyReadonlyArray1<'_, f64>,
3775    t_start: f64,
3776    t_end: f64,
3777    cost: f64,
3778) -> f64 {
3779    let a = times_a.as_slice().unwrap();
3780    let b = times_b.as_slice().unwrap();
3781    analysis::distance::adaptive_spike_distance(a, b, t_start, t_end, cost)
3782}
3783
3784#[pyfunction]
3785#[pyo3(signature = (train_a, train_b, dt=0.001, sigma_ms=5.0))]
3786fn py_schreiber_similarity(
3787    train_a: PyReadonlyArray1<'_, i32>,
3788    train_b: PyReadonlyArray1<'_, i32>,
3789    dt: f64,
3790    sigma_ms: f64,
3791) -> f64 {
3792    let a = train_a.as_slice().unwrap();
3793    let b = train_b.as_slice().unwrap();
3794    analysis::distance::schreiber_similarity(a, b, dt, sigma_ms)
3795}
3796
3797#[pyfunction]
3798#[pyo3(signature = (times_a, times_b, dt_max=0.01))]
3799fn py_hunter_milton_similarity(
3800    times_a: PyReadonlyArray1<'_, f64>,
3801    times_b: PyReadonlyArray1<'_, f64>,
3802    dt_max: f64,
3803) -> f64 {
3804    let a = times_a.as_slice().unwrap();
3805    let b = times_b.as_slice().unwrap();
3806    analysis::distance::hunter_milton_similarity(a, b, dt_max)
3807}
3808
3809#[pyfunction]
3810#[pyo3(signature = (times_a, times_b, t_start=0.0, t_end=1.0, n_bins=100))]
3811fn py_earth_movers_distance(
3812    times_a: PyReadonlyArray1<'_, f64>,
3813    times_b: PyReadonlyArray1<'_, f64>,
3814    t_start: f64,
3815    t_end: f64,
3816    n_bins: usize,
3817) -> f64 {
3818    let a = times_a.as_slice().unwrap();
3819    let b = times_b.as_slice().unwrap();
3820    analysis::distance::earth_movers_distance(a, b, t_start, t_end, n_bins)
3821}
3822
3823#[pyfunction]
3824#[pyo3(signature = (spike_times_list, cost_per_s=1000.0))]
3825fn py_multi_neuron_victor_purpura(
3826    py: Python<'_>,
3827    spike_times_list: Vec<PyReadonlyArray1<'_, f64>>,
3828    cost_per_s: f64,
3829) -> Py<PyArray2<f64>> {
3830    let vecs: Vec<Vec<f64>> = spike_times_list
3831        .iter()
3832        .map(|t| t.as_slice().unwrap().to_vec())
3833        .collect();
3834    let refs: Vec<&[f64]> = vecs.iter().map(|v| v.as_slice()).collect();
3835    let mat = analysis::distance::multi_neuron_victor_purpura(&refs, cost_per_s);
3836    let rows: Vec<Vec<f64>> = mat.into_iter().collect();
3837    let n = rows.len();
3838    numpy::PyArray2::from_vec2(py, &rows)
3839        .unwrap_or_else(|_| numpy::PyArray2::zeros(py, [n, n], false))
3840        .into()
3841}
3842
3843#[pyfunction]
3844#[pyo3(signature = (spike_times_list, metric="spike_distance", t_start=0.0, t_end=1.0))]
3845fn py_spike_distance_matrix(
3846    py: Python<'_>,
3847    spike_times_list: Vec<PyReadonlyArray1<'_, f64>>,
3848    metric: &str,
3849    t_start: f64,
3850    t_end: f64,
3851) -> Py<PyArray2<f64>> {
3852    let vecs: Vec<Vec<f64>> = spike_times_list
3853        .iter()
3854        .map(|t| t.as_slice().unwrap().to_vec())
3855        .collect();
3856    let refs: Vec<&[f64]> = vecs.iter().map(|v| v.as_slice()).collect();
3857    let mat = analysis::distance::spike_distance_matrix(&refs, metric, t_start, t_end);
3858    let rows: Vec<Vec<f64>> = mat.into_iter().collect();
3859    let n = rows.len();
3860    numpy::PyArray2::from_vec2(py, &rows)
3861        .unwrap_or_else(|_| numpy::PyArray2::zeros(py, [n, n], false))
3862        .into()
3863}
3864
3865// ── Information PyO3 wrappers (P0-A: spike_stats/information) ────
3866
3867#[pyfunction]
3868#[pyo3(signature = (train_a, train_b, bin_size=10))]
3869fn py_mutual_information(
3870    train_a: PyReadonlyArray1<'_, i32>,
3871    train_b: PyReadonlyArray1<'_, i32>,
3872    bin_size: usize,
3873) -> f64 {
3874    let a = train_a.as_slice().unwrap();
3875    let b = train_b.as_slice().unwrap();
3876    analysis::information::mutual_information(a, b, bin_size)
3877}
3878
3879#[pyfunction]
3880#[pyo3(signature = (source, target, bin_size=10, lag=1))]
3881fn py_transfer_entropy(
3882    source: PyReadonlyArray1<'_, i32>,
3883    target: PyReadonlyArray1<'_, i32>,
3884    bin_size: usize,
3885    lag: usize,
3886) -> f64 {
3887    let s = source.as_slice().unwrap();
3888    let t = target.as_slice().unwrap();
3889    analysis::information::transfer_entropy(s, t, bin_size, lag)
3890}
3891
3892#[pyfunction]
3893#[pyo3(signature = (binary_train, bin_size=10, word_length=4))]
3894fn py_spike_train_entropy(
3895    binary_train: PyReadonlyArray1<'_, i32>,
3896    bin_size: usize,
3897    word_length: usize,
3898) -> f64 {
3899    let data = binary_train.as_slice().unwrap();
3900    analysis::information::spike_train_entropy(data, bin_size, word_length)
3901}
3902
3903#[pyfunction]
3904#[pyo3(signature = (binary_train, n_trials=10, bin_size=10, word_length=4))]
3905fn py_noise_entropy(
3906    binary_train: PyReadonlyArray1<'_, i32>,
3907    n_trials: usize,
3908    bin_size: usize,
3909    word_length: usize,
3910) -> f64 {
3911    let data = binary_train.as_slice().unwrap();
3912    analysis::information::noise_entropy(data, n_trials, bin_size, word_length)
3913}
3914
3915#[pyfunction]
3916fn py_stimulus_specific_information(
3917    spike_counts: PyReadonlyArray1<'_, f64>,
3918    stimulus_ids: PyReadonlyArray1<'_, i64>,
3919) -> f64 {
3920    let counts = spike_counts.as_slice().unwrap();
3921    let ids = stimulus_ids.as_slice().unwrap();
3922    analysis::information::stimulus_specific_information(counts, ids)
3923}
3924
3925#[pyfunction]
3926#[pyo3(signature = (x, y, k=3))]
3927fn py_kozachenko_leonenko_mi(
3928    x: PyReadonlyArray1<'_, f64>,
3929    y: PyReadonlyArray1<'_, f64>,
3930    k: usize,
3931) -> f64 {
3932    let xd = x.as_slice().unwrap();
3933    let yd = y.as_slice().unwrap();
3934    analysis::information::kozachenko_leonenko_mi(xd, yd, k)
3935}
3936
3937// ── Causality PyO3 wrappers (P0-A: spike_stats/causality) ───────
3938
3939#[pyfunction]
3940#[pyo3(signature = (source, target, bin_size=10, order=5))]
3941fn py_pairwise_granger_causality(
3942    source: PyReadonlyArray1<'_, i32>,
3943    target: PyReadonlyArray1<'_, i32>,
3944    bin_size: usize,
3945    order: usize,
3946) -> f64 {
3947    let s = source.as_slice().unwrap();
3948    let t = target.as_slice().unwrap();
3949    analysis::causality::pairwise_granger_causality(s, t, bin_size, order)
3950}
3951
3952#[pyfunction]
3953#[pyo3(signature = (source, target, condition, bin_size=10, order=5))]
3954fn py_conditional_granger_causality(
3955    source: PyReadonlyArray1<'_, i32>,
3956    target: PyReadonlyArray1<'_, i32>,
3957    condition: PyReadonlyArray1<'_, i32>,
3958    bin_size: usize,
3959    order: usize,
3960) -> f64 {
3961    let s = source.as_slice().unwrap();
3962    let t = target.as_slice().unwrap();
3963    let c = condition.as_slice().unwrap();
3964    analysis::causality::conditional_granger_causality(s, t, c, bin_size, order)
3965}
3966
3967#[pyfunction]
3968#[pyo3(signature = (trains, bin_size=10, order=5, n_freqs=64))]
3969fn py_spectral_granger_causality(
3970    py: Python<'_>,
3971    trains: Vec<PyReadonlyArray1<'_, i32>>,
3972    bin_size: usize,
3973    order: usize,
3974    n_freqs: usize,
3975) -> (Py<PyArray1<f64>>, usize, usize) {
3976    let vecs: Vec<Vec<i32>> = trains
3977        .iter()
3978        .map(|t| t.as_slice().unwrap().to_vec())
3979        .collect();
3980    let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
3981    let (gc, d) = analysis::causality::spectral_granger_causality(&refs, bin_size, order, n_freqs);
3982    (gc.into_pyarray(py).into(), d, n_freqs)
3983}
3984
3985#[pyfunction]
3986#[pyo3(signature = (trains, bin_size=10, order=5, n_freqs=64))]
3987fn py_partial_directed_coherence(
3988    py: Python<'_>,
3989    trains: Vec<PyReadonlyArray1<'_, i32>>,
3990    bin_size: usize,
3991    order: usize,
3992    n_freqs: usize,
3993) -> (Py<PyArray1<f64>>, usize, usize) {
3994    let vecs: Vec<Vec<i32>> = trains
3995        .iter()
3996        .map(|t| t.as_slice().unwrap().to_vec())
3997        .collect();
3998    let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
3999    let (pdc, d) = analysis::causality::partial_directed_coherence(&refs, bin_size, order, n_freqs);
4000    (pdc.into_pyarray(py).into(), d, n_freqs)
4001}
4002
4003#[pyfunction]
4004#[pyo3(signature = (trains, bin_size=10, order=5, n_freqs=64))]
4005fn py_directed_transfer_function(
4006    py: Python<'_>,
4007    trains: Vec<PyReadonlyArray1<'_, i32>>,
4008    bin_size: usize,
4009    order: usize,
4010    n_freqs: usize,
4011) -> (Py<PyArray1<f64>>, usize, usize) {
4012    let vecs: Vec<Vec<i32>> = trains
4013        .iter()
4014        .map(|t| t.as_slice().unwrap().to_vec())
4015        .collect();
4016    let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
4017    let (dtf, d) = analysis::causality::directed_transfer_function(&refs, bin_size, order, n_freqs);
4018    (dtf.into_pyarray(py).into(), d, n_freqs)
4019}
4020
4021// ── Decoding PyO3 wrappers (P0-A: spike_stats/decoding) ─────────
4022
4023#[pyfunction]
4024#[pyo3(signature = (trains, preferred_directions, window=50))]
4025fn py_population_vector_decode(
4026    py: Python<'_>,
4027    trains: Vec<PyReadonlyArray1<'_, i32>>,
4028    preferred_directions: PyReadonlyArray1<'_, f64>,
4029    window: usize,
4030) -> Py<PyArray1<f64>> {
4031    let vecs: Vec<Vec<i32>> = trains
4032        .iter()
4033        .map(|t| t.as_slice().unwrap().to_vec())
4034        .collect();
4035    let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
4036    let dirs = preferred_directions.as_slice().unwrap();
4037    analysis::decoding::population_vector_decode(&refs, dirs, window)
4038        .into_pyarray(py)
4039        .into()
4040}
4041
4042#[pyfunction]
4043#[pyo3(signature = (spike_counts, tuning_rates, n_stimuli, n_neurons, prior=None))]
4044fn py_bayesian_decode(
4045    spike_counts: PyReadonlyArray1<'_, f64>,
4046    tuning_rates: PyReadonlyArray1<'_, f64>,
4047    n_stimuli: usize,
4048    n_neurons: usize,
4049    prior: Option<PyReadonlyArray1<'_, f64>>,
4050) -> usize {
4051    let counts = spike_counts.as_slice().unwrap();
4052    let rates = tuning_rates.as_slice().unwrap();
4053    let p: Vec<f64> = prior
4054        .map(|p| p.as_slice().unwrap().to_vec())
4055        .unwrap_or_default();
4056    analysis::decoding::bayesian_decode(counts, rates, n_stimuli, n_neurons, &p)
4057}
4058
4059#[pyfunction]
4060fn py_maximum_likelihood_decode(
4061    spike_counts: PyReadonlyArray1<'_, f64>,
4062    tuning_rates: PyReadonlyArray1<'_, f64>,
4063    n_stimuli: usize,
4064    n_neurons: usize,
4065) -> usize {
4066    let counts = spike_counts.as_slice().unwrap();
4067    let rates = tuning_rates.as_slice().unwrap();
4068    analysis::decoding::maximum_likelihood_decode(counts, rates, n_stimuli, n_neurons)
4069}
4070
4071#[pyfunction]
4072fn py_linear_discriminant_decode(
4073    train_data: PyReadonlyArray1<'_, f64>,
4074    n_samples: usize,
4075    n_features: usize,
4076    labels: PyReadonlyArray1<'_, i64>,
4077    test_point: PyReadonlyArray1<'_, f64>,
4078) -> i64 {
4079    let data = train_data.as_slice().unwrap();
4080    let lbl = labels.as_slice().unwrap();
4081    let tp = test_point.as_slice().unwrap();
4082    analysis::decoding::linear_discriminant_decode(data, n_samples, n_features, lbl, tp)
4083}
4084
4085#[pyfunction]
4086fn py_naive_bayes_decode(
4087    train_data: PyReadonlyArray1<'_, f64>,
4088    n_samples: usize,
4089    n_features: usize,
4090    labels: PyReadonlyArray1<'_, i64>,
4091    test_point: PyReadonlyArray1<'_, f64>,
4092) -> i64 {
4093    let data = train_data.as_slice().unwrap();
4094    let lbl = labels.as_slice().unwrap();
4095    let tp = test_point.as_slice().unwrap();
4096    analysis::decoding::naive_bayes_decode(data, n_samples, n_features, lbl, tp)
4097}
4098
4099// ── Neural decoder PyO3 wrappers (P1: neural_decoders) ─────────
4100
4101#[pyfunction]
4102#[pyo3(signature = (trains, dt=1.0))]
4103fn py_tokenise_spikes(
4104    py: Python<'_>,
4105    trains: Vec<PyReadonlyArray1<'_, i32>>,
4106    dt: f64,
4107) -> (Py<PyArray1<i64>>, Py<PyArray1<f64>>) {
4108    let vecs: Vec<Vec<i32>> = trains
4109        .iter()
4110        .map(|t| t.as_slice().unwrap().to_vec())
4111        .collect();
4112    let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
4113    let tokens = analysis::neural_decoders::tokenise_spikes(&refs, dt);
4114    let uids: Vec<i64> = tokens.iter().map(|t| t.0 as i64).collect();
4115    let times: Vec<f64> = tokens.iter().map(|t| t.1).collect();
4116    (uids.into_pyarray(py).into(), times.into_pyarray(py).into())
4117}
4118
4119#[pyfunction]
4120fn py_sinusoidal_position_encode(
4121    py: Python<'_>,
4122    timestamps: PyReadonlyArray1<'_, f64>,
4123    d_model: usize,
4124) -> Py<PyArray1<f64>> {
4125    let ts = timestamps.as_slice().unwrap();
4126    analysis::neural_decoders::sinusoidal_position_encode(ts, d_model)
4127        .into_pyarray(py)
4128        .into()
4129}
4130
4131#[pyfunction]
4132fn py_scaled_dot_product_attention(
4133    py: Python<'_>,
4134    queries: PyReadonlyArray1<'_, f64>,
4135    keys: PyReadonlyArray1<'_, f64>,
4136    values: PyReadonlyArray1<'_, f64>,
4137    nq: usize,
4138    nk: usize,
4139    d: usize,
4140) -> Py<PyArray1<f64>> {
4141    let q = queries.as_slice().unwrap();
4142    let k = keys.as_slice().unwrap();
4143    let v = values.as_slice().unwrap();
4144    analysis::neural_decoders::scaled_dot_product_attention(q, k, v, nq, nk, d)
4145        .into_pyarray(py)
4146        .into()
4147}
4148
4149#[pyfunction]
4150fn py_gaussian_attention(
4151    py: Python<'_>,
4152    queries: PyReadonlyArray1<'_, f64>,
4153    keys: PyReadonlyArray1<'_, f64>,
4154    values: PyReadonlyArray1<'_, f64>,
4155    nq: usize,
4156    nk: usize,
4157    d: usize,
4158    sigma: f64,
4159) -> Py<PyArray1<f64>> {
4160    let q = queries.as_slice().unwrap();
4161    let k = keys.as_slice().unwrap();
4162    let v = values.as_slice().unwrap();
4163    analysis::neural_decoders::gaussian_attention(q, k, v, nq, nk, d, sigma)
4164        .into_pyarray(py)
4165        .into()
4166}
4167
4168#[pyfunction]
4169fn py_infonce_loss(
4170    anchors: PyReadonlyArray1<'_, f64>,
4171    positives: PyReadonlyArray1<'_, f64>,
4172    n: usize,
4173    d: usize,
4174    temperature: f64,
4175) -> f64 {
4176    let a = anchors.as_slice().unwrap();
4177    let p = positives.as_slice().unwrap();
4178    analysis::neural_decoders::infonce_loss(a, p, n, d, temperature)
4179}
4180
4181// ── Network PyO3 wrappers (P0-A: spike_stats/network) ───────────
4182
4183#[pyfunction]
4184#[pyo3(signature = (trains, max_lag_ms=20.0, dt=0.001))]
4185fn py_functional_connectivity(
4186    py: Python<'_>,
4187    trains: Vec<PyReadonlyArray1<'_, i32>>,
4188    max_lag_ms: f64,
4189    dt: f64,
4190) -> Py<PyArray2<f64>> {
4191    let vecs: Vec<Vec<i32>> = trains
4192        .iter()
4193        .map(|t| t.as_slice().unwrap().to_vec())
4194        .collect();
4195    let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
4196    let mat = analysis::network::functional_connectivity(&refs, max_lag_ms, dt);
4197    let n = refs.len();
4198    let rows: Vec<Vec<f64>> = mat.chunks(n).map(|c| c.to_vec()).collect();
4199    numpy::PyArray2::from_vec2(py, &rows)
4200        .unwrap_or_else(|_| numpy::PyArray2::zeros(py, [n, n], false))
4201        .into()
4202}
4203
4204#[pyfunction]
4205#[pyo3(signature = (trains, bin_size=5, alpha=0.05))]
4206fn py_unitary_events(
4207    py: Python<'_>,
4208    trains: Vec<PyReadonlyArray1<'_, i32>>,
4209    bin_size: usize,
4210    alpha: f64,
4211) -> Py<PyArray1<i64>> {
4212    let vecs: Vec<Vec<i32>> = trains
4213        .iter()
4214        .map(|t| t.as_slice().unwrap().to_vec())
4215        .collect();
4216    let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
4217    let result = analysis::network::unitary_events(&refs, bin_size, alpha);
4218    let as_i64: Vec<i64> = result.into_iter().map(|v| v as i64).collect();
4219    as_i64.into_pyarray(py).into()
4220}
4221
4222#[pyfunction]
4223#[pyo3(signature = (trains, bin_size=5, threshold=2.0))]
4224fn py_cell_assembly_detection(
4225    trains: Vec<PyReadonlyArray1<'_, i32>>,
4226    bin_size: usize,
4227    threshold: f64,
4228) -> Vec<Vec<usize>> {
4229    let vecs: Vec<Vec<i32>> = trains
4230        .iter()
4231        .map(|t| t.as_slice().unwrap().to_vec())
4232        .collect();
4233    let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
4234    analysis::network::cell_assembly_detection(&refs, bin_size, threshold)
4235}
4236
4237#[pyfunction]
4238#[pyo3(signature = (trains, dt=0.001, max_delay_ms=20.0, min_chain_length=3))]
4239fn py_synfire_chain_detection(
4240    trains: Vec<PyReadonlyArray1<'_, i32>>,
4241    dt: f64,
4242    max_delay_ms: f64,
4243    min_chain_length: usize,
4244) -> Vec<Vec<usize>> {
4245    let vecs: Vec<Vec<i32>> = trains
4246        .iter()
4247        .map(|t| t.as_slice().unwrap().to_vec())
4248        .collect();
4249    let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
4250    analysis::network::synfire_chain_detection(&refs, dt, max_delay_ms, min_chain_length)
4251}
4252
4253// ── Surrogates PyO3 wrappers (P0-A: spike_stats/surrogates) ─────
4254
4255#[pyfunction]
4256#[pyo3(signature = (binary_train, seed=0))]
4257fn py_surrogate_isi_shuffle(
4258    py: Python<'_>,
4259    binary_train: PyReadonlyArray1<'_, i32>,
4260    seed: u64,
4261) -> Py<PyArray1<i32>> {
4262    let data = binary_train.as_slice().unwrap();
4263    analysis::surrogates::surrogate_isi_shuffle(data, seed)
4264        .into_pyarray(py)
4265        .into()
4266}
4267
4268#[pyfunction]
4269#[pyo3(signature = (binary_train, dither_ms=5.0, dt=0.001, seed=0))]
4270fn py_surrogate_dither(
4271    py: Python<'_>,
4272    binary_train: PyReadonlyArray1<'_, i32>,
4273    dither_ms: f64,
4274    dt: f64,
4275    seed: u64,
4276) -> Py<PyArray1<i32>> {
4277    let data = binary_train.as_slice().unwrap();
4278    analysis::surrogates::surrogate_dither(data, dither_ms, dt, seed)
4279        .into_pyarray(py)
4280        .into()
4281}
4282
4283#[pyfunction]
4284#[pyo3(signature = (rate_hz, duration_s, dt=0.001, seed=0))]
4285fn py_homogeneous_poisson(
4286    py: Python<'_>,
4287    rate_hz: f64,
4288    duration_s: f64,
4289    dt: f64,
4290    seed: u64,
4291) -> Py<PyArray1<f64>> {
4292    analysis::surrogates::homogeneous_poisson(rate_hz, duration_s, dt, seed)
4293        .into_pyarray(py)
4294        .into()
4295}
4296
4297#[pyfunction]
4298#[pyo3(signature = (rate_hz, shape, duration_s, dt=0.001, seed=0))]
4299fn py_gamma_process(
4300    py: Python<'_>,
4301    rate_hz: f64,
4302    shape: f64,
4303    duration_s: f64,
4304    dt: f64,
4305    seed: u64,
4306) -> Py<PyArray1<f64>> {
4307    analysis::surrogates::gamma_process(rate_hz, shape, duration_s, dt, seed)
4308        .into_pyarray(py)
4309        .into()
4310}
4311
4312#[pyfunction]
4313#[pyo3(signature = (rate_hz, burst_mean, duration_s, dt=0.001, seed=0))]
4314fn py_compound_poisson_process(
4315    py: Python<'_>,
4316    rate_hz: f64,
4317    burst_mean: f64,
4318    duration_s: f64,
4319    dt: f64,
4320    seed: u64,
4321) -> Py<PyArray1<f64>> {
4322    analysis::surrogates::compound_poisson_process(rate_hz, burst_mean, duration_s, dt, seed)
4323        .into_pyarray(py)
4324        .into()
4325}
4326
4327#[pyfunction]
4328#[pyo3(signature = (binary_train, seed=0))]
4329fn py_surrogate_joint_isi(
4330    py: Python<'_>,
4331    binary_train: PyReadonlyArray1<'_, i32>,
4332    seed: u64,
4333) -> Py<PyArray1<i32>> {
4334    let data = binary_train.as_slice().unwrap();
4335    analysis::surrogates::surrogate_joint_isi(data, seed)
4336        .into_pyarray(py)
4337        .into()
4338}
4339
4340#[pyfunction]
4341#[pyo3(signature = (binary_train, bin_size=10, seed=0))]
4342fn py_surrogate_bin_shuffling(
4343    py: Python<'_>,
4344    binary_train: PyReadonlyArray1<'_, i32>,
4345    bin_size: usize,
4346    seed: u64,
4347) -> Py<PyArray1<i32>> {
4348    let data = binary_train.as_slice().unwrap();
4349    analysis::surrogates::surrogate_bin_shuffling(data, bin_size, seed)
4350        .into_pyarray(py)
4351        .into()
4352}
4353
4354#[pyfunction]
4355#[pyo3(signature = (binary_train, max_shift=50, seed=0))]
4356fn py_surrogate_spike_train_shifting(
4357    py: Python<'_>,
4358    binary_train: PyReadonlyArray1<'_, i32>,
4359    max_shift: usize,
4360    seed: u64,
4361) -> Py<PyArray1<i32>> {
4362    let data = binary_train.as_slice().unwrap();
4363    analysis::surrogates::surrogate_spike_train_shifting(data, max_shift, seed)
4364        .into_pyarray(py)
4365        .into()
4366}
4367
4368// ── Temporal PyO3 wrappers (P0-A: spike_stats/temporal) ─────────
4369
4370#[pyfunction]
4371#[pyo3(signature = (binary_train, dt=0.001, max_isi_ms=10.0, min_spikes=3))]
4372fn py_burst_detection(
4373    binary_train: PyReadonlyArray1<'_, i32>,
4374    dt: f64,
4375    max_isi_ms: f64,
4376    min_spikes: usize,
4377) -> Vec<(f64, f64, usize)> {
4378    let data = binary_train.as_slice().unwrap();
4379    analysis::temporal::burst_detection(data, dt, max_isi_ms, min_spikes)
4380}
4381
4382#[pyfunction]
4383#[pyo3(signature = (binary_train, dt=0.001))]
4384fn py_first_spike_latency(binary_train: PyReadonlyArray1<'_, i32>, dt: f64) -> f64 {
4385    let data = binary_train.as_slice().unwrap();
4386    analysis::temporal::first_spike_latency(data, dt)
4387}
4388
4389#[pyfunction]
4390#[pyo3(signature = (binary_train, baseline_steps=100, dt=0.001, threshold_sigma=3.0))]
4391fn py_response_onset(
4392    binary_train: PyReadonlyArray1<'_, i32>,
4393    baseline_steps: usize,
4394    dt: f64,
4395    threshold_sigma: f64,
4396) -> f64 {
4397    let data = binary_train.as_slice().unwrap();
4398    analysis::temporal::response_onset(data, baseline_steps, dt, threshold_sigma)
4399}
4400
4401#[pyfunction]
4402#[pyo3(signature = (binary_train, bin_size=50, threshold=3.0))]
4403fn py_change_point_detection(
4404    py: Python<'_>,
4405    binary_train: PyReadonlyArray1<'_, i32>,
4406    bin_size: usize,
4407    threshold: f64,
4408) -> Py<PyArray1<i64>> {
4409    let data = binary_train.as_slice().unwrap();
4410    let cps = analysis::temporal::change_point_detection(data, bin_size, threshold);
4411    let as_i64: Vec<i64> = cps.into_iter().map(|v| v as i64).collect();
4412    as_i64.into_pyarray(py).into()
4413}
4414
4415// ── Patterns PyO3 wrappers (P0-A: spike_stats/patterns) ─────────
4416
4417#[pyfunction]
4418#[pyo3(signature = (times_a, times_b, t_start=0.0, t_end=1.0))]
4419fn py_spike_directionality(
4420    times_a: PyReadonlyArray1<'_, f64>,
4421    times_b: PyReadonlyArray1<'_, f64>,
4422    t_start: f64,
4423    t_end: f64,
4424) -> f64 {
4425    let a = times_a.as_slice().unwrap();
4426    let b = times_b.as_slice().unwrap();
4427    analysis::patterns::spike_directionality(a, b, t_start, t_end)
4428}
4429
4430#[pyfunction]
4431#[pyo3(signature = (times_list, t_start=0.0, t_end=1.0))]
4432fn py_spike_train_order(
4433    py: Python<'_>,
4434    times_list: Vec<PyReadonlyArray1<'_, f64>>,
4435    t_start: f64,
4436    t_end: f64,
4437) -> Py<PyArray2<f64>> {
4438    let vecs: Vec<Vec<f64>> = times_list
4439        .iter()
4440        .map(|t| t.as_slice().unwrap().to_vec())
4441        .collect();
4442    let refs: Vec<&[f64]> = vecs.iter().map(|v| v.as_slice()).collect();
4443    let mat = analysis::patterns::spike_train_order(&refs, t_start, t_end);
4444    let n = refs.len();
4445    let rows: Vec<Vec<f64>> = mat.chunks(n).map(|c| c.to_vec()).collect();
4446    numpy::PyArray2::from_vec2(py, &rows)
4447        .unwrap_or_else(|_| numpy::PyArray2::zeros(py, [n, n], false))
4448        .into()
4449}
4450
4451#[pyfunction]
4452#[pyo3(signature = (binary_train, dt=0.001, max_lag=20))]
4453fn py_cubic_higher_order(
4454    py: Python<'_>,
4455    binary_train: PyReadonlyArray1<'_, i32>,
4456    dt: f64,
4457    max_lag: usize,
4458) -> Py<PyArray2<f64>> {
4459    let data = binary_train.as_slice().unwrap();
4460    let c3 = analysis::patterns::cubic_higher_order(data, dt, max_lag);
4461    let rows: Vec<Vec<f64>> = c3.chunks(max_lag).map(|c| c.to_vec()).collect();
4462    numpy::PyArray2::from_vec2(py, &rows)
4463        .unwrap_or_else(|_| numpy::PyArray2::zeros(py, [max_lag, max_lag], false))
4464        .into()
4465}
4466
4467// ── Spectral PyO3 wrappers (P0-A: spike_stats/spectral) ─────────
4468
4469#[pyfunction]
4470#[pyo3(signature = (binary_train, dt=0.001))]
4471fn py_power_spectrum(
4472    py: Python<'_>,
4473    binary_train: PyReadonlyArray1<'_, i32>,
4474    dt: f64,
4475) -> (Py<PyArray1<f64>>, Py<PyArray1<f64>>) {
4476    let data = binary_train.as_slice().unwrap();
4477    let (psd, freqs) = analysis::spectral::power_spectrum(data, dt);
4478    (psd.into_pyarray(py).into(), freqs.into_pyarray(py).into())
4479}
4480
4481// ── Waveform PyO3 wrappers (P0-A: spike_stats/waveform) ─────────
4482
4483#[pyfunction]
4484#[pyo3(signature = (waveform, dt=3.3333333333333335e-05))]
4485fn py_waveform_width(waveform: PyReadonlyArray1<'_, f64>, dt: f64) -> f64 {
4486    analysis::waveform::waveform_width(waveform.as_slice().unwrap(), dt)
4487}
4488
4489#[pyfunction]
4490fn py_waveform_amplitude(waveform: PyReadonlyArray1<'_, f64>) -> f64 {
4491    analysis::waveform::waveform_amplitude(waveform.as_slice().unwrap())
4492}
4493
4494#[pyfunction]
4495#[pyo3(signature = (waveform, dt=3.3333333333333335e-05))]
4496fn py_waveform_repolarization_slope(waveform: PyReadonlyArray1<'_, f64>, dt: f64) -> f64 {
4497    analysis::waveform::waveform_repolarization_slope(waveform.as_slice().unwrap(), dt)
4498}
4499
4500#[pyfunction]
4501#[pyo3(signature = (waveform, dt=3.3333333333333335e-05))]
4502fn py_waveform_recovery_slope(waveform: PyReadonlyArray1<'_, f64>, dt: f64) -> f64 {
4503    analysis::waveform::waveform_recovery_slope(waveform.as_slice().unwrap(), dt)
4504}
4505
4506#[pyfunction]
4507#[pyo3(signature = (waveform, dt=3.3333333333333335e-05))]
4508fn py_waveform_halfwidth(waveform: PyReadonlyArray1<'_, f64>, dt: f64) -> f64 {
4509    analysis::waveform::waveform_halfwidth(waveform.as_slice().unwrap(), dt)
4510}
4511
4512#[pyfunction]
4513fn py_waveform_pt_ratio(waveform: PyReadonlyArray1<'_, f64>) -> f64 {
4514    analysis::waveform::waveform_pt_ratio(waveform.as_slice().unwrap())
4515}
4516
4517// ── Point process PyO3 wrappers (P0-A: spike_stats/point_process) ──
4518
4519#[pyfunction]
4520#[pyo3(signature = (binary_train, dt=0.001, window_ms=50.0))]
4521fn py_conditional_intensity(
4522    py: Python<'_>,
4523    binary_train: PyReadonlyArray1<'_, i32>,
4524    dt: f64,
4525    window_ms: f64,
4526) -> Py<PyArray1<f64>> {
4527    let data = binary_train.as_slice().unwrap();
4528    analysis::point_process::conditional_intensity(data, dt, window_ms)
4529        .into_pyarray(py)
4530        .into()
4531}
4532
4533#[pyfunction]
4534#[pyo3(signature = (binary_train, dt=0.001, bins=30))]
4535fn py_isi_hazard_function(
4536    py: Python<'_>,
4537    binary_train: PyReadonlyArray1<'_, i32>,
4538    dt: f64,
4539    bins: usize,
4540) -> (Py<PyArray1<f64>>, Py<PyArray1<f64>>) {
4541    let data = binary_train.as_slice().unwrap();
4542    let (hazard, centres) = analysis::point_process::isi_hazard_function(data, dt, bins);
4543    (
4544        hazard.into_pyarray(py).into(),
4545        centres.into_pyarray(py).into(),
4546    )
4547}
4548
4549#[pyfunction]
4550#[pyo3(signature = (binary_train, dt=0.001, bins=30))]
4551fn py_isi_survivor_function(
4552    py: Python<'_>,
4553    binary_train: PyReadonlyArray1<'_, i32>,
4554    dt: f64,
4555    bins: usize,
4556) -> (Py<PyArray1<f64>>, Py<PyArray1<f64>>) {
4557    let data = binary_train.as_slice().unwrap();
4558    let (surv, centres) = analysis::point_process::isi_survivor_function(data, dt, bins);
4559    (
4560        surv.into_pyarray(py).into(),
4561        centres.into_pyarray(py).into(),
4562    )
4563}
4564
4565#[pyfunction]
4566#[pyo3(signature = (binary_train, dt=0.001, bins=30))]
4567fn py_renewal_density(
4568    py: Python<'_>,
4569    binary_train: PyReadonlyArray1<'_, i32>,
4570    dt: f64,
4571    bins: usize,
4572) -> (Py<PyArray1<f64>>, Py<PyArray1<f64>>) {
4573    let data = binary_train.as_slice().unwrap();
4574    let (dens, centres) = analysis::point_process::renewal_density(data, dt, bins);
4575    (
4576        dens.into_pyarray(py).into(),
4577        centres.into_pyarray(py).into(),
4578    )
4579}
4580
4581// ── Stimulus PyO3 wrappers (P0-A: spike_stats/stimulus) ─────────
4582
4583#[pyfunction]
4584#[pyo3(signature = (stimulus, binary_train, window_steps=50))]
4585fn py_spike_triggered_average(
4586    py: Python<'_>,
4587    stimulus: PyReadonlyArray1<'_, f64>,
4588    binary_train: PyReadonlyArray1<'_, i32>,
4589    window_steps: usize,
4590) -> Py<PyArray1<f64>> {
4591    analysis::stimulus::spike_triggered_average(
4592        stimulus.as_slice().unwrap(),
4593        binary_train.as_slice().unwrap(),
4594        window_steps,
4595    )
4596    .into_pyarray(py)
4597    .into()
4598}
4599
4600#[pyfunction]
4601#[pyo3(signature = (stimulus, binary_train, window_steps=50))]
4602fn py_spike_triggered_covariance(
4603    py: Python<'_>,
4604    stimulus: PyReadonlyArray1<'_, f64>,
4605    binary_train: PyReadonlyArray1<'_, i32>,
4606    window_steps: usize,
4607) -> Py<PyArray2<f64>> {
4608    let cov = analysis::stimulus::spike_triggered_covariance(
4609        stimulus.as_slice().unwrap(),
4610        binary_train.as_slice().unwrap(),
4611        window_steps,
4612    );
4613    let rows: Vec<Vec<f64>> = cov.chunks(window_steps).map(|c| c.to_vec()).collect();
4614    numpy::PyArray2::from_vec2(py, &rows)
4615        .unwrap_or_else(|_| numpy::PyArray2::zeros(py, [window_steps, window_steps], false))
4616        .into()
4617}
4618
4619#[pyfunction]
4620#[pyo3(signature = (binary_train, positions, n_bins=20, dt=0.001))]
4621fn py_spatial_information(
4622    binary_train: PyReadonlyArray1<'_, i32>,
4623    positions: PyReadonlyArray1<'_, f64>,
4624    n_bins: usize,
4625    dt: f64,
4626) -> f64 {
4627    analysis::stimulus::spatial_information(
4628        binary_train.as_slice().unwrap(),
4629        positions.as_slice().unwrap(),
4630        n_bins,
4631        dt,
4632    )
4633}
4634
4635#[pyfunction]
4636#[pyo3(signature = (binary_train, positions, n_bins=50, threshold_std=2.0, dt=0.001))]
4637fn py_place_field_detection(
4638    binary_train: PyReadonlyArray1<'_, i32>,
4639    positions: PyReadonlyArray1<'_, f64>,
4640    n_bins: usize,
4641    threshold_std: f64,
4642    dt: f64,
4643) -> Vec<(f64, f64)> {
4644    analysis::stimulus::place_field_detection(
4645        binary_train.as_slice().unwrap(),
4646        positions.as_slice().unwrap(),
4647        n_bins,
4648        threshold_std,
4649        dt,
4650    )
4651}
4652
4653#[pyfunction]
4654#[pyo3(signature = (binary_train, stimulus_values, n_bins=20, dt=0.001))]
4655fn py_tuning_curve(
4656    py: Python<'_>,
4657    binary_train: PyReadonlyArray1<'_, i32>,
4658    stimulus_values: PyReadonlyArray1<'_, f64>,
4659    n_bins: usize,
4660    dt: f64,
4661) -> (Py<PyArray1<f64>>, Py<PyArray1<f64>>) {
4662    let (rates, centres) = analysis::stimulus::tuning_curve(
4663        binary_train.as_slice().unwrap(),
4664        stimulus_values.as_slice().unwrap(),
4665        n_bins,
4666        dt,
4667    );
4668    (
4669        rates.into_pyarray(py).into(),
4670        centres.into_pyarray(py).into(),
4671    )
4672}
4673
4674// ── LFP PyO3 wrappers (P0-A: spike_stats/lfp) ─────────────────
4675
4676#[pyfunction]
4677fn py_phase_locking_value(
4678    binary_train: PyReadonlyArray1<'_, i32>,
4679    lfp_signal: PyReadonlyArray1<'_, f64>,
4680) -> f64 {
4681    analysis::lfp::phase_locking_value(
4682        binary_train.as_slice().unwrap(),
4683        lfp_signal.as_slice().unwrap(),
4684    )
4685}
4686
4687#[pyfunction]
4688#[pyo3(signature = (binary_train, lfp_signal, dt=0.001))]
4689fn py_spike_field_coherence(
4690    py: Python<'_>,
4691    binary_train: PyReadonlyArray1<'_, i32>,
4692    lfp_signal: PyReadonlyArray1<'_, f64>,
4693    dt: f64,
4694) -> (Py<PyArray1<f64>>, Py<PyArray1<f64>>) {
4695    let (sfc, freqs) = analysis::lfp::spike_field_coherence(
4696        binary_train.as_slice().unwrap(),
4697        lfp_signal.as_slice().unwrap(),
4698        dt,
4699    );
4700    (sfc.into_pyarray(py).into(), freqs.into_pyarray(py).into())
4701}
4702
4703#[pyfunction]
4704#[pyo3(signature = (binary_train, lfp_signal, n_bins=36))]
4705fn py_spike_phase_histogram(
4706    py: Python<'_>,
4707    binary_train: PyReadonlyArray1<'_, i32>,
4708    lfp_signal: PyReadonlyArray1<'_, f64>,
4709    n_bins: usize,
4710) -> (Py<PyArray1<i64>>, Py<PyArray1<f64>>) {
4711    let (hist, centres) = analysis::lfp::spike_phase_histogram(
4712        binary_train.as_slice().unwrap(),
4713        lfp_signal.as_slice().unwrap(),
4714        n_bins,
4715    );
4716    (
4717        hist.into_pyarray(py).into(),
4718        centres.into_pyarray(py).into(),
4719    )
4720}
4721
4722// ── Sorting quality PyO3 wrappers (P0-A: spike_stats/sorting_quality)
4723
4724#[pyfunction]
4725fn py_isolation_distance(
4726    cluster: PyReadonlyArray2<'_, f64>,
4727    noise: PyReadonlyArray2<'_, f64>,
4728) -> f64 {
4729    let c_shape = cluster.shape();
4730    let n_shape = noise.shape();
4731    let d = c_shape[1];
4732    let c_data: Vec<f64> = cluster.as_slice().unwrap().to_vec();
4733    let n_data: Vec<f64> = noise.as_slice().unwrap().to_vec();
4734    analysis::sorting_quality::isolation_distance(&c_data, c_shape[0], &n_data, n_shape[0], d)
4735}
4736
4737#[pyfunction]
4738fn py_l_ratio(cluster: PyReadonlyArray2<'_, f64>, noise: PyReadonlyArray2<'_, f64>) -> f64 {
4739    let c_shape = cluster.shape();
4740    let n_shape = noise.shape();
4741    let d = c_shape[1];
4742    let c_data: Vec<f64> = cluster.as_slice().unwrap().to_vec();
4743    let n_data: Vec<f64> = noise.as_slice().unwrap().to_vec();
4744    analysis::sorting_quality::l_ratio(&c_data, c_shape[0], &n_data, n_shape[0], d)
4745}
4746
4747#[pyfunction]
4748fn py_silhouette_score(
4749    features: PyReadonlyArray2<'_, f64>,
4750    labels: PyReadonlyArray1<'_, i64>,
4751) -> f64 {
4752    let shape = features.shape();
4753    let f_data: Vec<f64> = features.as_slice().unwrap().to_vec();
4754    let l_data: Vec<i64> = labels.as_slice().unwrap().to_vec();
4755    analysis::sorting_quality::silhouette_score(&f_data, shape[0], shape[1], &l_data)
4756}
4757
4758#[pyfunction]
4759fn py_d_prime(cluster_a: PyReadonlyArray2<'_, f64>, cluster_b: PyReadonlyArray2<'_, f64>) -> f64 {
4760    let a_shape = cluster_a.shape();
4761    let b_shape = cluster_b.shape();
4762    let d = a_shape[1];
4763    let a_data: Vec<f64> = cluster_a.as_slice().unwrap().to_vec();
4764    let b_data: Vec<f64> = cluster_b.as_slice().unwrap().to_vec();
4765    analysis::sorting_quality::d_prime(&a_data, a_shape[0], &b_data, b_shape[0], d)
4766}
4767
4768#[pyfunction]
4769#[pyo3(signature = (binary_train, dt=0.001, refractory_ms=1.5))]
4770fn py_isi_violation_rate(
4771    binary_train: PyReadonlyArray1<'_, i32>,
4772    dt: f64,
4773    refractory_ms: f64,
4774) -> f64 {
4775    analysis::sorting_quality::isi_violation_rate(
4776        binary_train.as_slice().unwrap(),
4777        dt,
4778        refractory_ms,
4779    )
4780}
4781
4782#[pyfunction]
4783#[pyo3(signature = (binary_train, n_bins=100))]
4784fn py_presence_ratio(binary_train: PyReadonlyArray1<'_, i32>, n_bins: usize) -> f64 {
4785    analysis::sorting_quality::presence_ratio(binary_train.as_slice().unwrap(), n_bins)
4786}
4787
4788#[pyfunction]
4789#[pyo3(signature = (amplitudes, bins=100))]
4790fn py_amplitude_cutoff(amplitudes: PyReadonlyArray1<'_, f64>, bins: usize) -> f64 {
4791    analysis::sorting_quality::amplitude_cutoff(amplitudes.as_slice().unwrap(), bins)
4792}
4793
4794#[pyfunction]
4795fn py_snr(waveforms: PyReadonlyArray2<'_, f64>) -> f64 {
4796    let shape = waveforms.shape();
4797    let data: Vec<f64> = waveforms.as_slice().unwrap().to_vec();
4798    analysis::sorting_quality::snr(&data, shape[0], shape[1])
4799}
4800
4801#[pyfunction]
4802#[pyo3(signature = (cluster, noise, k=4))]
4803fn py_nn_hit_rate(
4804    cluster: PyReadonlyArray2<'_, f64>,
4805    noise: PyReadonlyArray2<'_, f64>,
4806    k: usize,
4807) -> f64 {
4808    let c_shape = cluster.shape();
4809    let n_shape = noise.shape();
4810    let d = c_shape[1];
4811    let c_data: Vec<f64> = cluster.as_slice().unwrap().to_vec();
4812    let n_data: Vec<f64> = noise.as_slice().unwrap().to_vec();
4813    analysis::sorting_quality::nn_hit_rate(&c_data, c_shape[0], &n_data, n_shape[0], d, k)
4814}
4815
4816#[pyfunction]
4817#[pyo3(signature = (waveforms, timestamps, n_bins=10))]
4818fn py_drift_metric(
4819    waveforms: PyReadonlyArray2<'_, f64>,
4820    timestamps: PyReadonlyArray1<'_, f64>,
4821    n_bins: usize,
4822) -> f64 {
4823    let shape = waveforms.shape();
4824    let data: Vec<f64> = waveforms.as_slice().unwrap().to_vec();
4825    let ts: Vec<f64> = timestamps.as_slice().unwrap().to_vec();
4826    analysis::sorting_quality::drift_metric(&data, shape[0], shape[1], &ts, n_bins)
4827}
4828
4829// ── Dimensionality PyO3 wrappers (P0-A: spike_stats/dimensionality)
4830
4831#[pyfunction]
4832#[pyo3(signature = (trains, n_components=3, bin_size=10))]
4833fn py_spike_train_pca(
4834    py: Python<'_>,
4835    trains: Vec<PyReadonlyArray1<'_, i32>>,
4836    n_components: usize,
4837    bin_size: usize,
4838) -> (Py<PyArray1<f64>>, Py<PyArray1<f64>>) {
4839    let vecs: Vec<Vec<i32>> = trains
4840        .iter()
4841        .map(|t| t.as_slice().unwrap().to_vec())
4842        .collect();
4843    let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
4844    let (proj, expl) = analysis::dimensionality::spike_train_pca(&refs, n_components, bin_size);
4845    (proj.into_pyarray(py).into(), expl.into_pyarray(py).into())
4846}
4847
4848#[pyfunction]
4849#[pyo3(signature = (conditions, n_components=3, bin_size=10))]
4850fn py_demixed_pca(
4851    py: Python<'_>,
4852    conditions: Vec<Vec<PyReadonlyArray1<'_, i32>>>,
4853    n_components: usize,
4854    bin_size: usize,
4855) -> (Py<PyArray1<f64>>, Py<PyArray1<f64>>) {
4856    let vecs: Vec<Vec<Vec<i32>>> = conditions
4857        .iter()
4858        .map(|cond| {
4859            cond.iter()
4860                .map(|t| t.as_slice().unwrap().to_vec())
4861                .collect()
4862        })
4863        .collect();
4864    let refs: Vec<Vec<&[i32]>> = vecs
4865        .iter()
4866        .map(|cond| cond.iter().map(|v| v.as_slice()).collect())
4867        .collect();
4868    let (proj, expl) = analysis::dimensionality::demixed_pca(&refs, n_components, bin_size);
4869    (proj.into_pyarray(py).into(), expl.into_pyarray(py).into())
4870}
4871
4872#[pyfunction]
4873#[pyo3(signature = (trains, n_factors=3, bin_size=10, n_iter=50))]
4874fn py_factor_analysis(
4875    py: Python<'_>,
4876    trains: Vec<PyReadonlyArray1<'_, i32>>,
4877    n_factors: usize,
4878    bin_size: usize,
4879    n_iter: usize,
4880) -> (Py<PyArray1<f64>>, Py<PyArray1<f64>>) {
4881    let vecs: Vec<Vec<i32>> = trains
4882        .iter()
4883        .map(|t| t.as_slice().unwrap().to_vec())
4884        .collect();
4885    let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
4886    let (loadings, psi) =
4887        analysis::dimensionality::factor_analysis(&refs, n_factors, bin_size, n_iter);
4888    (
4889        loadings.into_pyarray(py).into(),
4890        psi.into_pyarray(py).into(),
4891    )
4892}
4893
4894// Matrix-input wrappers: the caller bins and mean-centres once, so every backend
4895// (NumPy / Rust / Julia / Go / Mojo) shares an identical input matrix and the
4896// outputs agree to floating-point round-off.
4897
4898#[pyfunction]
4899#[pyo3(signature = (mat, n_components=3))]
4900fn py_pca_components(
4901    py: Python<'_>,
4902    mat: PyReadonlyArray2<'_, f64>,
4903    n_components: usize,
4904) -> (Py<PyArray1<f64>>, Py<PyArray1<f64>>) {
4905    let shape = mat.shape();
4906    let data: Vec<f64> = mat.as_slice().unwrap().to_vec();
4907    let (proj, expl) =
4908        analysis::dimensionality::pca_from_centered(&data, shape[0], shape[1], n_components);
4909    (proj.into_pyarray(py).into(), expl.into_pyarray(py).into())
4910}
4911
4912#[pyfunction]
4913#[pyo3(signature = (mean_mat, n_components=3))]
4914fn py_demixed_components(
4915    py: Python<'_>,
4916    mean_mat: PyReadonlyArray2<'_, f64>,
4917    n_components: usize,
4918) -> (Py<PyArray1<f64>>, Py<PyArray1<f64>>) {
4919    let shape = mean_mat.shape();
4920    let data: Vec<f64> = mean_mat.as_slice().unwrap().to_vec();
4921    let (proj, expl) =
4922        analysis::dimensionality::demixed_from_centered(&data, shape[0], shape[1], n_components);
4923    (proj.into_pyarray(py).into(), expl.into_pyarray(py).into())
4924}
4925
4926#[pyfunction]
4927#[pyo3(signature = (mat, n_factors=3, n_iter=50))]
4928fn py_factor_loadings(
4929    py: Python<'_>,
4930    mat: PyReadonlyArray2<'_, f64>,
4931    n_factors: usize,
4932    n_iter: usize,
4933) -> (Py<PyArray1<f64>>, Py<PyArray1<f64>>) {
4934    let shape = mat.shape();
4935    let data: Vec<f64> = mat.as_slice().unwrap().to_vec();
4936    let (loadings, psi) =
4937        analysis::dimensionality::fa_from_centered(&data, shape[0], shape[1], n_factors, n_iter);
4938    (
4939        loadings.into_pyarray(py).into(),
4940        psi.into_pyarray(py).into(),
4941    )
4942}
4943
4944// ── GPFA PyO3 wrappers (P0-A: spike_stats/gpfa) ─────────────────
4945
4946#[pyfunction]
4947#[pyo3(signature = (trains, n_latents=3, bin_ms=20.0, dt=0.001, max_iter=50, tol=1e-4, seed=42))]
4948fn py_gpfa<'py>(
4949    py: Python<'py>,
4950    trains: Vec<PyReadonlyArray1<'py, i32>>,
4951    n_latents: usize,
4952    bin_ms: f64,
4953    dt: f64,
4954    max_iter: usize,
4955    tol: f64,
4956    seed: u64,
4957) -> PyResult<Bound<'py, PyDict>> {
4958    let vecs: Vec<Vec<i32>> = trains
4959        .iter()
4960        .map(|t| t.as_slice().unwrap().to_vec())
4961        .collect();
4962    let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
4963    let result = analysis::gpfa::gpfa(&refs, n_latents, bin_ms, dt, max_iter, tol, seed);
4964
4965    let dict = PyDict::new(py);
4966    dict.set_item("trajectories", result.trajectories.into_pyarray(py))?;
4967    dict.set_item("C", result.c.into_pyarray(py))?;
4968    dict.set_item("d", result.d.into_pyarray(py))?;
4969    dict.set_item("R", result.r.into_pyarray(py))?;
4970    dict.set_item("tau", result.tau.into_pyarray(py))?;
4971    dict.set_item("log_likelihoods", result.log_likelihoods.into_pyarray(py))?;
4972    dict.set_item("n_latents", result.n_latents)?;
4973    dict.set_item("n_bins", result.n_bins)?;
4974    dict.set_item("n_neurons", result.n_neurons)?;
4975    Ok(dict)
4976}
4977
4978/// Run the GPFA EM loop from a caller-supplied deterministic initialisation.
4979///
4980/// Parity contract with `sc_neurocore.analysis.spike_stats.gpfa.gpfa_em`: identical
4981/// inputs (the PCA init computed once in Python) produce the same trajectories,
4982/// parameters and exact-marginal log-likelihoods up to floating-point round-off.
4983#[pyfunction]
4984#[pyo3(signature = (y, n_neurons, n_bins, c0, d0, r0_diag, tau, n_latents, max_iter, tol))]
4985#[allow(clippy::too_many_arguments, clippy::type_complexity)]
4986fn py_gpfa_em<'py>(
4987    py: Python<'py>,
4988    y: PyReadonlyArray1<'py, f64>,
4989    n_neurons: usize,
4990    n_bins: usize,
4991    c0: PyReadonlyArray1<'py, f64>,
4992    d0: PyReadonlyArray1<'py, f64>,
4993    r0_diag: PyReadonlyArray1<'py, f64>,
4994    tau: PyReadonlyArray1<'py, f64>,
4995    n_latents: usize,
4996    max_iter: usize,
4997    tol: f64,
4998) -> PyResult<(
4999    Bound<'py, PyArray1<f64>>,
5000    Bound<'py, PyArray1<f64>>,
5001    Bound<'py, PyArray1<f64>>,
5002    Bound<'py, PyArray1<f64>>,
5003    Vec<f64>,
5004)> {
5005    let (x_post, c, d, r, log_liks) = analysis::gpfa::gpfa_em_from_init(
5006        y.as_slice()?,
5007        c0.as_slice()?,
5008        d0.as_slice()?,
5009        r0_diag.as_slice()?,
5010        tau.as_slice()?,
5011        n_neurons,
5012        n_bins,
5013        n_latents,
5014        max_iter,
5015        tol,
5016    );
5017    Ok((
5018        x_post.into_pyarray(py),
5019        c.into_pyarray(py),
5020        d.into_pyarray(py),
5021        r.into_pyarray(py),
5022        log_liks,
5023    ))
5024}
5025
5026#[pyfunction]
5027#[pyo3(signature = (new_trains, c, d, r, tau, n_latents, bin_ms=20.0, dt=0.001))]
5028fn py_gpfa_transform(
5029    py: Python<'_>,
5030    new_trains: Vec<PyReadonlyArray1<'_, i32>>,
5031    c: PyReadonlyArray1<'_, f64>,
5032    d: PyReadonlyArray1<'_, f64>,
5033    r: PyReadonlyArray1<'_, f64>,
5034    tau: PyReadonlyArray1<'_, f64>,
5035    n_latents: usize,
5036    bin_ms: f64,
5037    dt: f64,
5038) -> Py<PyArray1<f64>> {
5039    let vecs: Vec<Vec<i32>> = new_trains
5040        .iter()
5041        .map(|t| t.as_slice().unwrap().to_vec())
5042        .collect();
5043    let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
5044    let proj = analysis::gpfa::gpfa_transform(
5045        &refs,
5046        c.as_slice().unwrap(),
5047        d.as_slice().unwrap(),
5048        r.as_slice().unwrap(),
5049        tau.as_slice().unwrap(),
5050        n_latents,
5051        bin_ms,
5052        dt,
5053    );
5054    proj.into_pyarray(py).into()
5055}
5056
5057// ── SPADE PyO3 wrappers (P0-A: spike_stats/spade) ─────────────
5058
5059#[pyfunction]
5060#[pyo3(signature = (trains, bin_ms=5.0, dt=0.001, min_support=3, max_pattern_size=5, n_surrogates=100, alpha=0.05, seed=42))]
5061fn py_spade_detect<'py>(
5062    py: Python<'py>,
5063    trains: Vec<PyReadonlyArray1<'py, i32>>,
5064    bin_ms: f64,
5065    dt: f64,
5066    min_support: usize,
5067    max_pattern_size: usize,
5068    n_surrogates: usize,
5069    alpha: f64,
5070    seed: u64,
5071) -> PyResult<Vec<Bound<'py, PyDict>>> {
5072    let vecs: Vec<Vec<i32>> = trains
5073        .iter()
5074        .map(|t| t.as_slice().unwrap().to_vec())
5075        .collect();
5076    let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
5077    let results = analysis::spade::spade_detect(
5078        &refs,
5079        bin_ms,
5080        dt,
5081        min_support,
5082        max_pattern_size,
5083        n_surrogates,
5084        alpha,
5085        seed,
5086    );
5087    let mut dicts = Vec::new();
5088    for pat in results {
5089        let dict = PyDict::new(py);
5090        dict.set_item(
5091            "neurons",
5092            pat.neurons.iter().map(|&n| n as i64).collect::<Vec<_>>(),
5093        )?;
5094        dict.set_item("lags", pat.lags.clone())?;
5095        dict.set_item("count", pat.count as i64)?;
5096        dict.set_item("p_value", pat.p_value)?;
5097        dicts.push(dict);
5098    }
5099    Ok(dicts)
5100}
5101
5102// ── DNA acceleration PyO3 wrappers ───────────────────────────────────
5103
5104#[pyfunction]
5105fn py_dna_design_sequence(_py: Python<'_>, length: usize, seed: u64) -> String {
5106    String::from_utf8(dna::design_sequence(length, seed)).unwrap_or_default()
5107}
5108
5109#[pyfunction]
5110fn py_dna_design_orthogonal_set(
5111    _py: Python<'_>,
5112    count: usize,
5113    length: usize,
5114    seed: u64,
5115) -> Vec<String> {
5116    dna::design_orthogonal_set(count, length, seed)
5117        .into_iter()
5118        .map(|s| String::from_utf8(s).unwrap_or_default())
5119        .collect()
5120}
5121
5122#[pyfunction]
5123fn py_dna_check_cross_hybridization<'py>(
5124    py: Python<'py>,
5125    sequences: Vec<String>,
5126    threshold: usize,
5127) -> PyResult<Py<PyAny>> {
5128    let seqs: Vec<Vec<u8>> = sequences.into_iter().map(|s| s.into_bytes()).collect();
5129    let flags = dna::check_cross_hybridization(&seqs, threshold);
5130    let result: Vec<Py<PyAny>> = flags
5131        .into_iter()
5132        .map(|(i, j, score)| {
5133            let d = PyDict::new(py);
5134            d.set_item("strand_a", i).unwrap();
5135            d.set_item("strand_b", j).unwrap();
5136            d.set_item("score", score).unwrap();
5137            d.into_any().unbind()
5138        })
5139        .collect();
5140    Ok(result.into_pyobject(py)?.into())
5141}
5142
5143#[pyfunction]
5144#[pyo3(signature = (
5145    gate_types, gate_inputs, gate_outputs, gate_thresholds, gate_leaks,
5146    input_names, input_concs, duration_s=1800.0, dt=1.0,
5147    k_hyb=3e5, k_disp=1.0, temperature_c=37.0, use_rk4=true
5148))]
5149#[allow(clippy::too_many_arguments)]
5150fn py_dna_simulate_kinetics<'py>(
5151    py: Python<'py>,
5152    gate_types: Vec<String>,
5153    gate_inputs: Vec<Vec<String>>,
5154    gate_outputs: Vec<String>,
5155    gate_thresholds: Vec<f64>,
5156    gate_leaks: Vec<f64>,
5157    input_names: Vec<String>,
5158    input_concs: Vec<f64>,
5159    duration_s: f64,
5160    dt: f64,
5161    k_hyb: f64,
5162    k_disp: f64,
5163    temperature_c: f64,
5164    use_rk4: bool,
5165) -> PyResult<Py<PyAny>> {
5166    let gates: Vec<dna::DnaGateSpec> = gate_types
5167        .iter()
5168        .zip(gate_inputs.iter())
5169        .zip(gate_outputs.iter())
5170        .zip(gate_thresholds.iter())
5171        .zip(gate_leaks.iter())
5172        .map(|((((gt, gi), go), th), lk)| {
5173            let gate_type = match gt.to_uppercase().as_str() {
5174                "AND" => dna::DnaGateType::And,
5175                "OR" => dna::DnaGateType::Or,
5176                "NOT" => dna::DnaGateType::Not,
5177                "THRESHOLD" => dna::DnaGateType::Threshold,
5178                "MUX" => dna::DnaGateType::Mux,
5179                "AMPLIFIER" => dna::DnaGateType::Amplifier,
5180                "BUFFER" => dna::DnaGateType::Buffer,
5181                "NAND" => dna::DnaGateType::Nand,
5182                "XOR" => dna::DnaGateType::Xor,
5183                _ => dna::DnaGateType::And,
5184            };
5185            dna::DnaGateSpec {
5186                gate_type,
5187                input_names: gi.clone(),
5188                output_name: go.clone(),
5189                threshold: *th,
5190                leak_rate: *lk,
5191            }
5192        })
5193        .collect();
5194
5195    let mut inputs = std::collections::HashMap::new();
5196    for (name, conc) in input_names.into_iter().zip(input_concs) {
5197        inputs.insert(name, conc);
5198    }
5199
5200    let config = dna::KineticConfig {
5201        k_hyb,
5202        k_disp,
5203        temperature_c,
5204        max_conc: 200.0,
5205        use_rk4,
5206    };
5207
5208    let result = dna::simulate_kinetics(&gates, &inputs, duration_s, dt, &config);
5209
5210    let dict = PyDict::new(py);
5211    for (key, trace) in result {
5212        dict.set_item(key, trace.into_pyarray(py))?;
5213    }
5214    Ok(dict.into_any().unbind())
5215}
5216
5217#[pyfunction]
5218#[pyo3(signature = (sequence, min_stem=4, min_loop=3))]
5219fn py_dna_detect_hairpins(
5220    _py: Python<'_>,
5221    sequence: &str,
5222    min_stem: usize,
5223    min_loop: usize,
5224) -> Vec<(usize, usize, usize)> {
5225    dna::detect_hairpins(sequence.as_bytes(), min_stem, min_loop)
5226}
5227
5228// ── Quantum Annealing PyO3 Wrappers ──────────────────────────────────
5229
5230/// Compute Ising energy for a spin configuration.
5231///
5232/// Args:
5233///     h_indices, h_values: linear biases (parallel arrays)
5234///     j_i, j_j, j_values: quadratic couplings (parallel arrays)
5235///     spins: spin configuration (+1/-1)
5236///     offset: constant energy offset
5237#[pyfunction]
5238#[pyo3(signature = (h_indices, h_values, j_i, j_j, j_values, spins, offset=0.0))]
5239fn py_qa_ising_energy(
5240    _py: Python<'_>,
5241    h_indices: Vec<usize>,
5242    h_values: Vec<f64>,
5243    j_i: Vec<usize>,
5244    j_j: Vec<usize>,
5245    j_values: Vec<f64>,
5246    spins: Vec<i8>,
5247    offset: f64,
5248) -> f64 {
5249    let h: Vec<(usize, f64)> = h_indices.into_iter().zip(h_values).collect();
5250    let j: Vec<((usize, usize), f64)> = j_i.into_iter().zip(j_j).zip(j_values).collect();
5251    quantum::ising_energy(&h, &j, &spins, offset)
5252}
5253
5254/// Batch compute Ising energies for many configurations.
5255#[pyfunction]
5256#[pyo3(signature = (h_indices, h_values, j_i, j_j, j_values, configs, offset=0.0))]
5257fn py_qa_batch_ising_energy(
5258    _py: Python<'_>,
5259    h_indices: Vec<usize>,
5260    h_values: Vec<f64>,
5261    j_i: Vec<usize>,
5262    j_j: Vec<usize>,
5263    j_values: Vec<f64>,
5264    configs: Vec<Vec<i8>>,
5265    offset: f64,
5266) -> Vec<f64> {
5267    let h: Vec<(usize, f64)> = h_indices.into_iter().zip(h_values).collect();
5268    let j: Vec<((usize, usize), f64)> = j_i.into_iter().zip(j_j).zip(j_values).collect();
5269    quantum::batch_ising_energy(&h, &j, &configs, offset)
5270}
5271
5272/// Run simulated annealing on an Ising model (Rust-accelerated).
5273///
5274/// Returns dict with best_spins, best_energy, energies, samples.
5275#[pyfunction]
5276#[pyo3(signature = (h_indices, h_values, j_i, j_j, j_values, n_qubits, offset=0.0, n_sweeps=1000, num_reads=10, beta_start=0.1, beta_end=10.0, seed=42))]
5277fn py_qa_simulated_annealing<'py>(
5278    py: Python<'py>,
5279    h_indices: Vec<usize>,
5280    h_values: Vec<f64>,
5281    j_i: Vec<usize>,
5282    j_j: Vec<usize>,
5283    j_values: Vec<f64>,
5284    n_qubits: usize,
5285    offset: f64,
5286    n_sweeps: usize,
5287    num_reads: usize,
5288    beta_start: f64,
5289    beta_end: f64,
5290    seed: u64,
5291) -> PyResult<Py<PyAny>> {
5292    let h: Vec<(usize, f64)> = h_indices.into_iter().zip(h_values).collect();
5293    let j: Vec<((usize, usize), f64)> = j_i.into_iter().zip(j_j).zip(j_values).collect();
5294
5295    let (best_spins, best_energy, all_energies, all_samples) = quantum::simulated_annealing(
5296        &h, &j, n_qubits, offset, n_sweeps, num_reads, beta_start, beta_end, seed,
5297    );
5298
5299    let dict = PyDict::new(py);
5300    dict.set_item("best_spins", best_spins)?;
5301    dict.set_item("best_energy", best_energy)?;
5302    dict.set_item("energies", all_energies)?;
5303    dict.set_item("samples", all_samples)?;
5304    dict.set_item("n_sweeps", n_sweeps)?;
5305    dict.set_item("num_reads", num_reads)?;
5306    dict.set_item("backend", "rust")?;
5307    Ok(dict.into_any().unbind())
5308}
5309
5310/// Apply gauge transform to Ising biases and couplings.
5311#[pyfunction]
5312#[allow(clippy::type_complexity)] // tuple shape mirrors Python's QUBO format
5313fn py_qa_gauge_transform(
5314    _py: Python<'_>,
5315    h_indices: Vec<usize>,
5316    h_values: Vec<f64>,
5317    j_i: Vec<usize>,
5318    j_j: Vec<usize>,
5319    j_values: Vec<f64>,
5320    gauge: Vec<i8>,
5321) -> (Vec<(usize, f64)>, Vec<((usize, usize), f64)>) {
5322    let h: Vec<(usize, f64)> = h_indices.into_iter().zip(h_values).collect();
5323    let j: Vec<((usize, usize), f64)> = j_i.into_iter().zip(j_j).zip(j_values).collect();
5324    quantum::gauge_transform(&h, &j, &gauge)
5325}
5326
5327/// Generate random gauge vectors.
5328#[pyfunction]
5329#[pyo3(signature = (n_qubits, n_gauges=10, seed=42))]
5330fn py_qa_generate_gauges(
5331    _py: Python<'_>,
5332    n_qubits: usize,
5333    n_gauges: usize,
5334    seed: u64,
5335) -> Vec<Vec<i8>> {
5336    quantum::generate_gauges(n_qubits, n_gauges, seed)
5337}
5338
5339/// Greedy graph partitioning for problem decomposition.
5340#[pyfunction]
5341#[pyo3(signature = (n_qubits, j_i, j_j, j_values, max_partition_size=64))]
5342fn py_qa_greedy_partition(
5343    _py: Python<'_>,
5344    n_qubits: usize,
5345    j_i: Vec<usize>,
5346    j_j: Vec<usize>,
5347    j_values: Vec<f64>,
5348    max_partition_size: usize,
5349) -> Vec<Vec<usize>> {
5350    let j: Vec<((usize, usize), f64)> = j_i.into_iter().zip(j_j).zip(j_values).collect();
5351    quantum::greedy_partition(n_qubits, &j, max_partition_size)
5352}
5353
5354// ── Photonic NoC PyO3 Wrappers ───────────────────────────────────────
5355
5356/// Route waveguides on a mesh topology (Rust-accelerated).
5357#[pyfunction]
5358#[pyo3(signature = (adjacency_flat, n, pitch_um=250.0, loss_db_per_cm=2.0))]
5359fn py_ph_route_waveguides<'py>(
5360    py: Python<'py>,
5361    adjacency_flat: Vec<f64>,
5362    n: usize,
5363    pitch_um: f64,
5364    loss_db_per_cm: f64,
5365) -> PyResult<Py<PyAny>> {
5366    let result = photonic::route_waveguides(&adjacency_flat, n, pitch_um, loss_db_per_cm);
5367
5368    let dict = PyDict::new(py);
5369    let sources: Vec<usize> = result.iter().map(|r| r.source).collect();
5370    let targets: Vec<usize> = result.iter().map(|r| r.target).collect();
5371    let lengths: Vec<f64> = result.iter().map(|r| r.length_um).collect();
5372    let losses: Vec<f64> = result.iter().map(|r| r.loss_db).collect();
5373    let crossings: Vec<usize> = result.iter().map(|r| r.n_crossings).collect();
5374
5375    dict.set_item("sources", sources)?;
5376    dict.set_item("targets", targets)?;
5377    dict.set_item("lengths_um", lengths)?;
5378    dict.set_item("losses_db", losses)?;
5379    dict.set_item("crossings", crossings)?;
5380    dict.set_item("n_segments", result.len())?;
5381    dict.set_item("backend", "rust")?;
5382    Ok(dict.into_any().unbind())
5383}
5384
5385/// Compute MZI 2×2 transfer matrix for a given phase.
5386#[pyfunction]
5387fn py_ph_mzi_transfer_matrix(_py: Python<'_>, phase_rad: f64) -> Vec<f64> {
5388    photonic::mzi_transfer_matrix(phase_rad).to_vec()
5389}
5390
5391/// Cascade multiple MZI stages via matrix multiplication.
5392#[pyfunction]
5393fn py_ph_cascade_mzi(_py: Python<'_>, phases: Vec<f64>) -> Vec<f64> {
5394    photonic::cascade_mzi(&phases).to_vec()
5395}
5396
5397/// Analyze WDM crosstalk (Rust-accelerated).
5398#[pyfunction]
5399#[pyo3(signature = (channel_ids, wavelengths, bandwidths, powers, adjacent_xt_db=-25.0))]
5400fn py_ph_analyze_crosstalk<'py>(
5401    py: Python<'py>,
5402    channel_ids: Vec<usize>,
5403    wavelengths: Vec<f64>,
5404    bandwidths: Vec<f64>,
5405    powers: Vec<f64>,
5406    adjacent_xt_db: f64,
5407) -> PyResult<Py<PyAny>> {
5408    let channels: Vec<(usize, f64, f64, f64)> = channel_ids
5409        .into_iter()
5410        .zip(wavelengths)
5411        .zip(bandwidths)
5412        .zip(powers)
5413        .map(|(((id, wl), bw), p)| (id, wl, bw, p))
5414        .collect();
5415
5416    let result = photonic::analyze_crosstalk(&channels, adjacent_xt_db);
5417
5418    let dict = PyDict::new(py);
5419    let ids: Vec<usize> = result.iter().map(|r| r.channel_id).collect();
5420    let xts: Vec<f64> = result.iter().map(|r| r.crosstalk_db).collect();
5421    let osnrs: Vec<f64> = result.iter().map(|r| r.osnr_db).collect();
5422    let adjs: Vec<usize> = result.iter().map(|r| r.n_adjacent).collect();
5423
5424    dict.set_item("channel_ids", ids)?;
5425    dict.set_item("crosstalk_db", xts)?;
5426    dict.set_item("osnr_db", osnrs)?;
5427    dict.set_item("n_adjacent", adjs)?;
5428    dict.set_item("backend", "rust")?;
5429    Ok(dict.into_any().unbind())
5430}
5431
5432/// Analyze optical power budget (Rust-accelerated).
5433#[pyfunction]
5434#[pyo3(signature = (wg_sources, wg_targets, wg_losses, laser_power_dbm=0.0, detector_sensitivity_dbm=-20.0))]
5435fn py_ph_analyze_power_budget<'py>(
5436    py: Python<'py>,
5437    wg_sources: Vec<usize>,
5438    wg_targets: Vec<usize>,
5439    wg_losses: Vec<f64>,
5440    laser_power_dbm: f64,
5441    detector_sensitivity_dbm: f64,
5442) -> PyResult<Py<PyAny>> {
5443    let wgs: Vec<(usize, usize, f64)> = wg_sources
5444        .into_iter()
5445        .zip(wg_targets)
5446        .zip(wg_losses)
5447        .map(|((s, t), l)| (s, t, l))
5448        .collect();
5449
5450    let result =
5451        photonic::analyze_power_budget(&wgs, &[], laser_power_dbm, detector_sensitivity_dbm);
5452
5453    let dict = PyDict::new(py);
5454    let margins: Vec<f64> = result.iter().map(|r| r.margin_db).collect();
5455    let passed: Vec<bool> = result.iter().map(|r| r.passed).collect();
5456    let total_losses: Vec<f64> = result.iter().map(|r| r.total_loss_db).collect();
5457
5458    dict.set_item("margins_db", margins)?;
5459    dict.set_item("passed", passed)?;
5460    dict.set_item("total_losses_db", total_losses)?;
5461    dict.set_item("n_paths", result.len())?;
5462    dict.set_item("backend", "rust")?;
5463    Ok(dict.into_any().unbind())
5464}
5465
5466/// Geometric crosstalk analysis for a uniform bank of parallel waveguides.
5467#[pyfunction]
5468#[pyo3(signature = (num_waveguides, gap_nm, coupling_length_um, wavelength_nm=1550.0, core_index=3.48, cladding_index=1.45))]
5469fn py_ph_analyze_crosstalk_bank<'py>(
5470    py: Python<'py>,
5471    num_waveguides: usize,
5472    gap_nm: f64,
5473    coupling_length_um: f64,
5474    wavelength_nm: f64,
5475    core_index: f64,
5476    cladding_index: f64,
5477) -> PyResult<Py<PyAny>> {
5478    let r = photonic::analyze_crosstalk_bank(
5479        num_waveguides,
5480        gap_nm,
5481        coupling_length_um,
5482        wavelength_nm,
5483        core_index,
5484        cladding_index,
5485    );
5486    let dict = PyDict::new(py);
5487    dict.set_item("num_waveguides", r.num_waveguides)?;
5488    dict.set_item("num_pairs", r.num_near_pairs + r.num_far_pairs)?;
5489    dict.set_item("num_near_pairs", r.num_near_pairs)?;
5490    dict.set_item("num_far_pairs", r.num_far_pairs)?;
5491    dict.set_item("gap_nm", r.gap_nm)?;
5492    dict.set_item("coupling_length_um", r.coupling_length_um)?;
5493    dict.set_item("adjacent_coupling_ratio", r.adjacent_coupling_ratio)?;
5494    dict.set_item("adjacent_isolation_db", r.adjacent_isolation_db)?;
5495    dict.set_item("next_nearest_coupling_ratio", r.next_nearest_coupling_ratio)?;
5496    dict.set_item("next_nearest_isolation_db", r.next_nearest_isolation_db)?;
5497    dict.set_item("worst_isolation_db", r.worst_isolation_db)?;
5498    dict.set_item("mean_coupling_ratio", r.mean_coupling_ratio)?;
5499    dict.set_item("max_coupling_ratio", r.max_coupling_ratio)?;
5500    dict.set_item("crosstalk_safe", r.crosstalk_safe)?;
5501    dict.set_item("backend", "rust")?;
5502    Ok(dict.into_any().unbind())
5503}
5504
5505/// Per-pair geometric crosstalk for arbitrary waveguide geometry.
5506/// `pairs_a[i]`, `pairs_b[i]`, `gaps_nm[i]`, `lengths_um[i]` describe pair i.
5507/// Evaluated in parallel via Rayon — the O(N²) analysis path.
5508#[pyfunction]
5509#[pyo3(signature = (pairs_a, pairs_b, gaps_nm, lengths_um, wavelength_nm=1550.0, core_index=3.48, cladding_index=1.45))]
5510fn py_ph_analyze_crosstalk_pairs<'py>(
5511    py: Python<'py>,
5512    pairs_a: Vec<usize>,
5513    pairs_b: Vec<usize>,
5514    gaps_nm: Vec<f64>,
5515    lengths_um: Vec<f64>,
5516    wavelength_nm: f64,
5517    core_index: f64,
5518    cladding_index: f64,
5519) -> PyResult<Py<PyAny>> {
5520    let n = pairs_a.len();
5521    if pairs_b.len() != n || gaps_nm.len() != n || lengths_um.len() != n {
5522        return Err(pyo3::exceptions::PyValueError::new_err(
5523            "pairs_a, pairs_b, gaps_nm, lengths_um must be equal length",
5524        ));
5525    }
5526    let pairs: Vec<(usize, usize, f64, f64)> = pairs_a
5527        .into_iter()
5528        .zip(pairs_b)
5529        .zip(gaps_nm)
5530        .zip(lengths_um)
5531        .map(|(((a, b), g), l)| (a, b, g, l))
5532        .collect();
5533    let results =
5534        photonic::analyze_crosstalk_pairs(&pairs, wavelength_nm, core_index, cladding_index);
5535
5536    let dict = PyDict::new(py);
5537    let idx_a: Vec<usize> = results.iter().map(|r| r.index_a).collect();
5538    let idx_b: Vec<usize> = results.iter().map(|r| r.index_b).collect();
5539    let gaps: Vec<f64> = results.iter().map(|r| r.gap_nm).collect();
5540    let lens: Vec<f64> = results.iter().map(|r| r.coupling_length_um).collect();
5541    let kappas: Vec<f64> = results
5542        .iter()
5543        .map(|r| r.coupling_coefficient_per_um)
5544        .collect();
5545    let ratios: Vec<f64> = results.iter().map(|r| r.coupling_ratio).collect();
5546    let isos: Vec<f64> = results.iter().map(|r| r.isolation_db).collect();
5547
5548    dict.set_item("pair_a", idx_a)?;
5549    dict.set_item("pair_b", idx_b)?;
5550    dict.set_item("gap_nm", gaps)?;
5551    dict.set_item("coupling_length_um", lens)?;
5552    dict.set_item("coupling_coefficient_per_um", kappas)?;
5553    dict.set_item("coupling_ratio", ratios)?;
5554    dict.set_item("isolation_db", isos)?;
5555    dict.set_item("num_pairs", n)?;
5556    dict.set_item("backend", "rust")?;
5557    Ok(dict.into_any().unbind())
5558}
5559
5560// ── SC-Optimizer PyO3 Wrappers ───────────────────────────────────────
5561
5562/// Run SA design-space search (Rust-accelerated).
5563///
5564/// mac_counts: per-layer MAC counts
5565/// weights: per-layer scoring weights
5566/// Returns dict with best config indices, score, pareto data
5567#[pyfunction]
5568#[pyo3(signature = (mac_counts, weights, max_luts, max_power, max_latency=0, t_init=1.0, t_min=0.001, alpha=0.95, max_iter=2000, seed=42))]
5569fn py_opt_sa_search<'py>(
5570    py: Python<'py>,
5571    mac_counts: Vec<i64>,
5572    weights: Vec<f64>,
5573    max_luts: i64,
5574    max_power: f64,
5575    max_latency: i64,
5576    t_init: f64,
5577    t_min: f64,
5578    alpha: f64,
5579    max_iter: usize,
5580    seed: u64,
5581) -> PyResult<Py<PyAny>> {
5582    let candidates: Vec<Vec<optimizer::Candidate>> = mac_counts
5583        .iter()
5584        .map(|&mc| optimizer::generate_candidates(mc))
5585        .collect();
5586
5587    let result = optimizer::simulated_annealing(
5588        &candidates,
5589        &weights,
5590        max_luts,
5591        max_power,
5592        max_latency,
5593        t_init,
5594        t_min,
5595        alpha,
5596        max_iter,
5597        seed,
5598    );
5599
5600    let dict = PyDict::new(py);
5601    match result {
5602        Some(r) => {
5603            // Extract details before moving best_config
5604            let mut luts_list = Vec::new();
5605            let mut power_list = Vec::new();
5606            let mut acc_list = Vec::new();
5607            for (i, &idx) in r.best_config.iter().enumerate() {
5608                let c = &candidates[i][idx];
5609                luts_list.push(c.luts);
5610                power_list.push(c.power);
5611                acc_list.push(c.accuracy);
5612            }
5613
5614            dict.set_item("best_config", r.best_config)?;
5615            dict.set_item("best_score", r.best_score)?;
5616            dict.set_item("pareto_luts", r.pareto_luts)?;
5617            dict.set_item("pareto_power", r.pareto_power)?;
5618            dict.set_item("pareto_score", r.pareto_score)?;
5619            dict.set_item("feasible", true)?;
5620            dict.set_item("layer_luts", luts_list)?;
5621            dict.set_item("layer_power", power_list)?;
5622            dict.set_item("layer_accuracy", acc_list)?;
5623        }
5624        None => {
5625            dict.set_item("feasible", false)?;
5626        }
5627    }
5628    dict.set_item("backend", "rust")?;
5629    Ok(dict.into_any().unbind())
5630}
5631
5632/// Extract Pareto frontier from (luts, power, score) arrays.
5633#[pyfunction]
5634fn py_opt_extract_pareto<'py>(
5635    py: Python<'py>,
5636    luts: Vec<i64>,
5637    power: Vec<f64>,
5638    score: Vec<f64>,
5639) -> PyResult<Py<PyAny>> {
5640    let indices = optimizer::extract_pareto(&luts, &power, &score);
5641    let dict = PyDict::new(py);
5642    let p_luts: Vec<i64> = indices.iter().map(|&i| luts[i]).collect();
5643    let p_power: Vec<f64> = indices.iter().map(|&i| power[i]).collect();
5644    let p_score: Vec<f64> = indices.iter().map(|&i| score[i]).collect();
5645    dict.set_item("indices", indices)?;
5646    dict.set_item("luts", p_luts)?;
5647    dict.set_item("power", p_power)?;
5648    dict.set_item("score", p_score)?;
5649    dict.set_item("backend", "rust")?;
5650    Ok(dict.into_any().unbind())
5651}
5652
5653// ── Evolutionary Substrate PyO3 Wrappers ─────────────────────────────
5654
5655/// Batch-mutate population weights (Rust-accelerated).
5656#[pyfunction]
5657#[pyo3(signature = (genomes, mutation_rate=0.1, mutation_scale=0.1, seed=42))]
5658fn py_evo_batch_mutate(
5659    _py: Python<'_>,
5660    mut genomes: Vec<Vec<f64>>,
5661    mutation_rate: f64,
5662    mutation_scale: f64,
5663    seed: u64,
5664) -> Vec<Vec<f64>> {
5665    evo::batch_mutate_weights(&mut genomes, mutation_rate, mutation_scale, seed);
5666    genomes
5667}
5668
5669/// Batch fitness evaluation (Rust-accelerated).
5670#[pyfunction]
5671fn py_evo_batch_fitness(
5672    _py: Python<'_>,
5673    genomes: Vec<Vec<f64>>,
5674    inputs: Vec<f64>,
5675    target: f64,
5676) -> Vec<f64> {
5677    evo::batch_evaluate_fitness(&genomes, &inputs, target)
5678}
5679
5680/// Batch uniform crossover (Rust-accelerated).
5681#[pyfunction]
5682#[pyo3(signature = (parents_a, parents_b, seed=42))]
5683fn py_evo_batch_crossover(
5684    _py: Python<'_>,
5685    parents_a: Vec<Vec<f64>>,
5686    parents_b: Vec<Vec<f64>>,
5687    seed: u64,
5688) -> Vec<Vec<f64>> {
5689    evo::batch_crossover(&parents_a, &parents_b, seed)
5690}
5691
5692/// Population diversity (mean pairwise L2 distance).
5693#[pyfunction]
5694fn py_evo_diversity(_py: Python<'_>, genomes: Vec<Vec<f64>>) -> f64 {
5695    evo::population_diversity(&genomes)
5696}
5697
5698/// Novelty scores against archive.
5699#[pyfunction]
5700#[pyo3(signature = (genomes, archive, k_nearest=5))]
5701fn py_evo_novelty(
5702    _py: Python<'_>,
5703    genomes: Vec<Vec<f64>>,
5704    archive: Vec<Vec<f64>>,
5705    k_nearest: usize,
5706) -> Vec<f64> {
5707    evo::novelty_scores(&genomes, &archive, k_nearest)
5708}
5709
5710/// Tournament selection.
5711#[pyfunction]
5712#[pyo3(signature = (fitness, n_select, tournament_size=3, seed=42))]
5713fn py_evo_tournament(
5714    _py: Python<'_>,
5715    fitness: Vec<f64>,
5716    n_select: usize,
5717    tournament_size: usize,
5718    seed: u64,
5719) -> Vec<usize> {
5720    evo::tournament_select(&fitness, n_select, tournament_size, seed)
5721}
5722
5723// ── LGSSM Kalman filter (predictive_model) ──────────────────────────
5724
5725/// Forward Kalman filter for a Linear Gaussian State-Space Model.
5726///
5727/// Parity contract with `sc_neurocore.world_model.predictive_model.KalmanFilter`:
5728/// for the same model parameters and observation sequence, the
5729/// returned (means, covariances, log_likelihood) must agree with
5730/// the Python implementation to within float64 round-off.
5731///
5732/// All matrices are passed as flat row-major Vec<f64>; the caller
5733/// supplies their shapes explicitly. Returns a dict with keys:
5734///   - "means": Vec<Vec<f64>> shape (T, d)
5735///   - "covariances": Vec<Vec<Vec<f64>>> shape (T, d, d)
5736///   - "pred_means": Vec<Vec<f64>> shape (T, d)
5737///   - "pred_covariances": Vec<Vec<Vec<f64>>> shape (T, d, d)
5738///   - "log_likelihood": f64
5739///   - "backend": "rust"
5740#[pyfunction]
5741#[pyo3(signature = (
5742    obs_flat, controls_flat, t_len, p_dim, m_dim,
5743    a_flat, b_flat, c_flat, d_flat, q_flat, r_flat,
5744    mu_0, sigma_0_flat, d_dim,
5745))]
5746#[allow(clippy::too_many_arguments)]
5747fn py_lgssm_kalman_filter<'py>(
5748    py: Python<'py>,
5749    obs_flat: Vec<f64>,
5750    controls_flat: Vec<f64>,
5751    t_len: usize,
5752    p_dim: usize,
5753    m_dim: usize,
5754    a_flat: Vec<f64>,
5755    b_flat: Vec<f64>,
5756    c_flat: Vec<f64>,
5757    d_flat: Vec<f64>,
5758    q_flat: Vec<f64>,
5759    r_flat: Vec<f64>,
5760    mu_0: Vec<f64>,
5761    sigma_0_flat: Vec<f64>,
5762    d_dim: usize,
5763) -> PyResult<Py<PyAny>> {
5764    use ndarray::Array1;
5765    use ndarray::Array2;
5766
5767    let to_2d = |flat: &[f64], rows: usize, cols: usize| -> Array2<f64> {
5768        Array2::from_shape_vec((rows, cols), flat.to_vec()).expect("shape")
5769    };
5770    let obs = to_2d(&obs_flat, t_len, p_dim);
5771    let controls = to_2d(&controls_flat, t_len, m_dim);
5772    let a = to_2d(&a_flat, d_dim, d_dim);
5773    let b = to_2d(&b_flat, d_dim, m_dim);
5774    let c = to_2d(&c_flat, p_dim, d_dim);
5775    let d = to_2d(&d_flat, p_dim, m_dim);
5776    let q = to_2d(&q_flat, d_dim, d_dim);
5777    let r = to_2d(&r_flat, p_dim, p_dim);
5778    let mu_0_arr = Array1::from(mu_0);
5779    let sigma_0 = to_2d(&sigma_0_flat, d_dim, d_dim);
5780
5781    let result = lgssm::kalman_filter(
5782        obs.view(),
5783        controls.view(),
5784        a.view(),
5785        b.view(),
5786        c.view(),
5787        d.view(),
5788        q.view(),
5789        r.view(),
5790        mu_0_arr.view(),
5791        sigma_0.view(),
5792    );
5793
5794    // Convert to Python-friendly nested Vec
5795    let means: Vec<Vec<f64>> = (0..t_len)
5796        .map(|t| (0..d_dim).map(|i| result.means[(t, i)]).collect())
5797        .collect();
5798    let covs: Vec<Vec<Vec<f64>>> = (0..t_len)
5799        .map(|t| {
5800            (0..d_dim)
5801                .map(|i| (0..d_dim).map(|j| result.covariances[(t, i, j)]).collect())
5802                .collect()
5803        })
5804        .collect();
5805    let pred_means: Vec<Vec<f64>> = (0..t_len)
5806        .map(|t| (0..d_dim).map(|i| result.pred_means[(t, i)]).collect())
5807        .collect();
5808    let pred_covs: Vec<Vec<Vec<f64>>> = (0..t_len)
5809        .map(|t| {
5810            (0..d_dim)
5811                .map(|i| {
5812                    (0..d_dim)
5813                        .map(|j| result.pred_covariances[(t, i, j)])
5814                        .collect()
5815                })
5816                .collect()
5817        })
5818        .collect();
5819
5820    let dict = PyDict::new(py);
5821    dict.set_item("means", means)?;
5822    dict.set_item("covariances", covs)?;
5823    dict.set_item("pred_means", pred_means)?;
5824    dict.set_item("pred_covariances", pred_covs)?;
5825    dict.set_item("log_likelihood", result.log_likelihood)?;
5826    dict.set_item("backend", "rust")?;
5827    Ok(dict.into_any().unbind())
5828}
5829
5830/// Discrete Ollivier-Ricci curvature between two nodes of a coupling graph.
5831///
5832/// Parity contract with `sc_neurocore.math.topology.ollivier_ricci_curvature`:
5833/// for the same `knm` and `(i, j)`, the Rust value agrees with the Python
5834/// value to within float64 round-off.
5835///
5836/// `knm_flat` is the row-major `n x n` coupling matrix. Raises `ValueError`
5837/// on a malformed shape, a non-finite or negative entry, or an out-of-range
5838/// index, mirroring the Python validation.
5839#[pyfunction]
5840#[pyo3(signature = (knm_flat, n, i, j))]
5841fn py_ollivier_ricci_curvature(knm_flat: Vec<f64>, n: usize, i: usize, j: usize) -> PyResult<f64> {
5842    use topology::CurvatureError;
5843    match topology::ollivier_ricci_curvature(&knm_flat, n, i, j) {
5844        Ok(value) => Ok(value),
5845        Err(CurvatureError::BadShape) => Err(PyValueError::new_err(
5846            "knm must be a square coupling matrix with at least one node",
5847        )),
5848        Err(CurvatureError::BadValue) => Err(PyValueError::new_err(
5849            "knm must contain only finite, non-negative values",
5850        )),
5851        Err(CurvatureError::BadIndex) => Err(PyValueError::new_err(
5852            "node index out of range for coupling graph",
5853        )),
5854        Err(CurvatureError::Infeasible) => {
5855            Err(PyValueError::new_err("transport problem is infeasible"))
5856        }
5857    }
5858}
5859
5860/// N-step Cazelles-Courbage-Rabinovich (2001) bursting-map simulation.
5861///
5862/// Parity contract with `sc_neurocore.neurons.models.cazelles_map.CazellesMapNeuron.simulate`:
5863/// for the same parameters and constant input the returned `x` trace, spike
5864/// count, and final `(x, y)` state are bit-identical to the Python reference
5865/// (the map is exact floating-point arithmetic, no transcendental functions).
5866#[pyfunction]
5867#[pyo3(signature = (x0, y0, a, epsilon, sigma, x_threshold, n_steps, current))]
5868#[allow(clippy::too_many_arguments)]
5869fn py_cazelles_map_simulate<'py>(
5870    py: Python<'py>,
5871    x0: f64,
5872    y0: f64,
5873    a: f64,
5874    epsilon: f64,
5875    sigma: f64,
5876    x_threshold: f64,
5877    n_steps: usize,
5878    current: f64,
5879) -> (Bound<'py, PyArray1<f64>>, i64, f64, f64) {
5880    let mut neuron = crate::neurons::CazellesMapNeuron {
5881        x: x0,
5882        y: y0,
5883        a,
5884        epsilon,
5885        sigma,
5886        x_threshold,
5887    };
5888    let (trace, spikes) = neuron.simulate(n_steps, current);
5889    // Return the trace as a NumPy array directly; marshalling a multi-million
5890    // element Vec<f64> into a Python list would dominate the wall-clock.
5891    (trace.into_pyarray(py), spikes, neuron.x, neuron.y)
5892}
5893
5894/// N-step Courbage-Nekorkin-Vdovin (2007) discontinuous spiking-map simulation.
5895///
5896/// Parity contract with
5897/// `sc_neurocore.neurons.models.courage_nekorkin_map.CourageNekorkinMapNeuron.simulate`:
5898/// for the same parameters and constant input the returned `x` trace,
5899/// upward-crossing spike count, and final `(x, y)` state are bit-identical to
5900/// the Python reference (the map is exact floating-point arithmetic — additions,
5901/// multiplications, one division for the breakpoints, and a piecewise/Heaviside
5902/// branch, no transcendental functions).
5903#[pyfunction]
5904#[pyo3(signature = (x0, y0, m0, m1, a, d, j, beta, eps, x_threshold, n_steps, current))]
5905#[allow(clippy::too_many_arguments)]
5906fn py_courage_nekorkin_map_simulate<'py>(
5907    py: Python<'py>,
5908    x0: f64,
5909    y0: f64,
5910    m0: f64,
5911    m1: f64,
5912    a: f64,
5913    d: f64,
5914    j: f64,
5915    beta: f64,
5916    eps: f64,
5917    x_threshold: f64,
5918    n_steps: usize,
5919    current: f64,
5920) -> (Bound<'py, PyArray1<f64>>, i64, f64, f64) {
5921    let mut neuron = crate::neurons::CourageNekorkinMapNeuron {
5922        x: x0,
5923        y: y0,
5924        m0,
5925        m1,
5926        a,
5927        d,
5928        j,
5929        beta,
5930        eps,
5931        x_threshold,
5932    };
5933    let (trace, spikes) = neuron.simulate(n_steps, current);
5934    (trace.into_pyarray(py), spikes, neuron.x, neuron.y)
5935}
5936
5937/// Parity contract with `sc_neurocore.neurons.models.rulkov_map.RulkovMapNeuron.simulate`:
5938/// for the same parameters and constant input the returned `x` trace, upward-
5939/// crossing spike count, and final `(x, y)` state are bit-identical to the
5940/// Python reference (the map is exact floating-point arithmetic — one division,
5941/// additions and multiplications, no transcendental functions).
5942#[pyfunction]
5943#[pyo3(signature = (x0, y0, alpha, sigma, mu, x_threshold, n_steps, current))]
5944#[allow(clippy::too_many_arguments)]
5945fn py_rulkov_map_simulate<'py>(
5946    py: Python<'py>,
5947    x0: f64,
5948    y0: f64,
5949    alpha: f64,
5950    sigma: f64,
5951    mu: f64,
5952    x_threshold: f64,
5953    n_steps: usize,
5954    current: f64,
5955) -> (Bound<'py, PyArray1<f64>>, i64, f64, f64) {
5956    let mut neuron = crate::neurons::RulkovMapNeuron {
5957        x: x0,
5958        y: y0,
5959        alpha,
5960        sigma,
5961        mu,
5962        x_threshold,
5963    };
5964    let (trace, spikes) = neuron.simulate(n_steps, current);
5965    (trace.into_pyarray(py), spikes, neuron.x, neuron.y)
5966}
5967
5968/// Parity contract with
5969/// `sc_neurocore.neurons.models.ibarz_tanaka_map.IbarzTanakaMapNeuron.simulate`:
5970/// for the same parameters and constant input the returned `x` trace (already
5971/// reset on spiking steps), spike count, and final `(x, y)` state are
5972/// bit-identical to the Python reference (the map is exact floating-point
5973/// arithmetic — one division, additions and multiplications, no transcendental
5974/// functions).
5975#[pyfunction]
5976#[pyo3(signature = (x0, y0, alpha, beta, mu, sigma, x_threshold, x_reset, n_steps, current))]
5977#[allow(clippy::too_many_arguments)]
5978fn py_ibarz_tanaka_map_simulate<'py>(
5979    py: Python<'py>,
5980    x0: f64,
5981    y0: f64,
5982    alpha: f64,
5983    beta: f64,
5984    mu: f64,
5985    sigma: f64,
5986    x_threshold: f64,
5987    x_reset: f64,
5988    n_steps: usize,
5989    current: f64,
5990) -> (Bound<'py, PyArray1<f64>>, i64, f64, f64) {
5991    let mut neuron = crate::neurons::IbarzTanakaMapNeuron {
5992        x: x0,
5993        y: y0,
5994        alpha,
5995        beta,
5996        mu,
5997        sigma,
5998        x_threshold,
5999        x_reset,
6000    };
6001    let (trace, spikes) = neuron.simulate(n_steps, current);
6002    (trace.into_pyarray(py), spikes, neuron.x, neuron.y)
6003}
6004
6005/// Parity contract with
6006/// `sc_neurocore.neurons.models.medvedev_map.MedvedevMapNeuron.simulate`: for
6007/// the same parameters and constant input the returned `x` trace, upward-
6008/// crossing spike count, and final `x` state are bit-identical to the Python
6009/// reference (the map is exact floating-point arithmetic — a multiply, an add,
6010/// and a fold into `[0, 1)`; `f64::rem_euclid(1.0)` equals Python's `x % 1.0`
6011/// bit-for-bit). This is a one-dimensional map, so there is no `y` state.
6012#[pyfunction]
6013#[pyo3(signature = (x0, alpha, beta, x_threshold, n_steps, current))]
6014fn py_medvedev_map_simulate<'py>(
6015    py: Python<'py>,
6016    x0: f64,
6017    alpha: f64,
6018    beta: f64,
6019    x_threshold: f64,
6020    n_steps: usize,
6021    current: f64,
6022) -> (Bound<'py, PyArray1<f64>>, i64, f64) {
6023    let mut neuron = crate::neurons::MedvedevMapNeuron {
6024        x: x0,
6025        alpha,
6026        beta,
6027        x_threshold,
6028    };
6029    let (trace, spikes) = neuron.simulate(n_steps, current);
6030    (trace.into_pyarray(py), spikes, neuron.x)
6031}
6032
6033/// Parity contract with
6034/// `sc_neurocore.neurons.models.ermentrout_kopell_map_neuron.ErmentroutKopellMapNeuron.simulate`:
6035/// for the same parameters and constant input the returned `theta` trace,
6036/// upward-crossing spike count, and final `theta` state match the Python
6037/// reference bit-for-bit on a shared libm (the only transcendental is `cos`,
6038/// and the non-chaotic phase flow does not amplify ULP differences). This is a
6039/// one-dimensional phase map, so there is no second state.
6040#[pyfunction]
6041#[pyo3(signature = (theta0, dt, gain, theta_threshold, n_steps, current))]
6042fn py_ermentrout_kopell_map_simulate<'py>(
6043    py: Python<'py>,
6044    theta0: f64,
6045    dt: f64,
6046    gain: f64,
6047    theta_threshold: f64,
6048    n_steps: usize,
6049    current: f64,
6050) -> (Bound<'py, PyArray1<f64>>, i64, f64) {
6051    let mut neuron = crate::neurons::ErmentroutKopellMapNeuron {
6052        theta: theta0,
6053        dt,
6054        gain,
6055        theta_threshold,
6056    };
6057    let (trace, spikes) = neuron.simulate(n_steps, current);
6058    (trace.into_pyarray(py), spikes, neuron.theta)
6059}
6060
6061/// N-step McKean (1970) piecewise-linear FitzHugh-Nagumo caricature simulation.
6062///
6063/// Parity contract with
6064/// `sc_neurocore.neurons.models.mckean.McKeanNeuron.simulate`: for the same
6065/// parameters and constant input the returned `v` trace, upward-`v_peak`-crossing
6066/// spike count, and final `(v, w)` state are bit-identical to the Python RK4
6067/// reference (the piecewise-linear right-hand side is exact arithmetic —
6068/// additions, multiplications and branch selection, no transcendental functions —
6069/// and a two-dimensional autonomous flow cannot be chaotic).
6070#[pyfunction]
6071#[pyo3(signature = (v0, w0, a, epsilon, gamma, dt, v_peak, n_steps, current))]
6072#[allow(clippy::too_many_arguments)]
6073fn py_mckean_simulate<'py>(
6074    py: Python<'py>,
6075    v0: f64,
6076    w0: f64,
6077    a: f64,
6078    epsilon: f64,
6079    gamma: f64,
6080    dt: f64,
6081    v_peak: f64,
6082    n_steps: usize,
6083    current: f64,
6084) -> (Bound<'py, PyArray1<f64>>, i64, f64, f64) {
6085    let mut neuron = crate::neurons::McKeanNeuron {
6086        v: v0,
6087        w: w0,
6088        a,
6089        epsilon,
6090        gamma,
6091        dt,
6092        v_peak,
6093    };
6094    let (trace, spikes) = neuron.simulate(n_steps, current);
6095    (trace.into_pyarray(py), spikes, neuron.v, neuron.w)
6096}
6097
6098/// N-step Wilson (1999) polynomial cortical-neuron simulation.
6099///
6100/// Parity contract with
6101/// `sc_neurocore.neurons.models.wilson_hr.WilsonHRNeuron.simulate`: for the same
6102/// parameters and constant input the returned `v` trace (already hard-reset to
6103/// `-0.7` on spiking steps), spike count, and final `(v, r)` state are
6104/// bit-identical to the Python RK4 reference (the right-hand side is exact
6105/// polynomial arithmetic — no transcendental functions).
6106#[pyfunction]
6107#[pyo3(signature = (v0, r0, tau_r, v_peak, dt, n_steps, current))]
6108#[allow(clippy::too_many_arguments)]
6109fn py_wilson_hr_simulate<'py>(
6110    py: Python<'py>,
6111    v0: f64,
6112    r0: f64,
6113    tau_r: f64,
6114    v_peak: f64,
6115    dt: f64,
6116    n_steps: usize,
6117    current: f64,
6118) -> (Bound<'py, PyArray1<f64>>, i64, f64, f64) {
6119    let mut neuron = crate::neurons::WilsonHRNeuron {
6120        v: v0,
6121        r: r0,
6122        tau_r,
6123        v_peak,
6124        dt,
6125    };
6126    let (trace, spikes) = neuron.simulate(n_steps, current);
6127    (trace.into_pyarray(py), spikes, neuron.v, neuron.r)
6128}
6129
6130/// N-step Pernarowski (1994) pancreatic beta-cell burster simulation.
6131///
6132/// Parity contract with
6133/// `sc_neurocore.neurons.models.pernarowski.PernarowskiNeuron.simulate`: for the
6134/// same parameters and constant input the returned `v` trace, upward-crossing
6135/// spike count, and final `(v, w, z)` state are bit-identical to the Python RK4
6136/// reference (the cubic uses `v.powi(3)` = `v*v*v`, matching the Python `v*v*v`;
6137/// no transcendental functions).
6138#[pyfunction]
6139#[pyo3(signature = (v0, w0, z0, alpha, beta, eps1, eps2, gamma, dt, v_threshold, n_steps, current))]
6140#[allow(clippy::too_many_arguments)]
6141fn py_pernarowski_simulate<'py>(
6142    py: Python<'py>,
6143    v0: f64,
6144    w0: f64,
6145    z0: f64,
6146    alpha: f64,
6147    beta: f64,
6148    eps1: f64,
6149    eps2: f64,
6150    gamma: f64,
6151    dt: f64,
6152    v_threshold: f64,
6153    n_steps: usize,
6154    current: f64,
6155) -> (Bound<'py, PyArray1<f64>>, i64, f64, f64, f64) {
6156    let mut neuron = crate::neurons::PernarowskiNeuron {
6157        v: v0,
6158        w: w0,
6159        z: z0,
6160        alpha,
6161        beta,
6162        eps1,
6163        eps2,
6164        gamma,
6165        dt,
6166        v_threshold,
6167    };
6168    let (trace, spikes) = neuron.simulate(n_steps, current);
6169    (trace.into_pyarray(py), spikes, neuron.v, neuron.w, neuron.z)
6170}
6171
6172/// N-step Terman-Wang (LEGION) relaxation-oscillator simulation.
6173///
6174/// Parity contract with
6175/// `sc_neurocore.neurons.models.terman_wang.TermanWangOscillator.simulate`: for
6176/// the same parameters and constant input the returned `v` trace, upward-crossing
6177/// spike count, and final `(v, w)` state match the Python RK4 reference. The cubic
6178/// uses `v.powi(3)` = `v*v*v` (matching the Python `v*v*v`); the `tanh` gating
6179/// resolves to the same glibc symbol as Python on Linux, so this backend is
6180/// bit-identical there (Julia/Go/Mojo use their own libm `tanh`, ULP-bounded, and
6181/// the two-dimensional relaxation oscillator is non-chaotic so it does not amplify).
6182#[pyfunction]
6183#[pyo3(signature = (v0, w0, alpha, beta, epsilon, rho, dt, v_peak, n_steps, current))]
6184#[allow(clippy::too_many_arguments)]
6185fn py_terman_wang_simulate<'py>(
6186    py: Python<'py>,
6187    v0: f64,
6188    w0: f64,
6189    alpha: f64,
6190    beta: f64,
6191    epsilon: f64,
6192    rho: f64,
6193    dt: f64,
6194    v_peak: f64,
6195    n_steps: usize,
6196    current: f64,
6197) -> (Bound<'py, PyArray1<f64>>, i64, f64, f64) {
6198    let mut neuron = crate::neurons::TermanWangOscillator {
6199        v: v0,
6200        w: w0,
6201        alpha,
6202        beta,
6203        epsilon,
6204        rho,
6205        dt,
6206        v_peak,
6207    };
6208    let (trace, spikes) = neuron.simulate(n_steps, current);
6209    (trace.into_pyarray(py), spikes, neuron.v, neuron.w)
6210}
6211
6212/// Parity contract with
6213/// `sc_neurocore.neurons.models.mihalas_niebur.MihalasNieburNeuron.simulate`:
6214/// for the same parameters and constant input the returned `v` trace, total
6215/// spike count, and final `(v, theta, i1, i2)` state are bit-identical to the
6216/// Python RK4 reference. The Mihalas-Niebur 2009 generalised integrate-and-fire
6217/// model has a purely linear right-hand side (no transcendental functions), so
6218/// every RK4 stage is exact arithmetic and the trace reproduces to the last bit
6219/// across Rust/Julia/Go (Mojo FMA-fuses, validated non-amplifying).
6220#[pyfunction]
6221#[pyo3(signature = (
6222    v0, theta0, i1_0, i2_0, v_rest, v_reset, theta_reset, theta_inf,
6223    tau_v, tau_theta, tau_1, tau_2, a, b, r1, r2, dt, n_steps, current
6224))]
6225#[allow(clippy::too_many_arguments)]
6226fn py_mihalas_niebur_simulate<'py>(
6227    py: Python<'py>,
6228    v0: f64,
6229    theta0: f64,
6230    i1_0: f64,
6231    i2_0: f64,
6232    v_rest: f64,
6233    v_reset: f64,
6234    theta_reset: f64,
6235    theta_inf: f64,
6236    tau_v: f64,
6237    tau_theta: f64,
6238    tau_1: f64,
6239    tau_2: f64,
6240    a: f64,
6241    b: f64,
6242    r1: f64,
6243    r2: f64,
6244    dt: f64,
6245    n_steps: usize,
6246    current: f64,
6247) -> (Bound<'py, PyArray1<f64>>, i64, f64, f64, f64, f64) {
6248    let mut neuron = crate::neurons::MihalasNieburNeuron {
6249        v: v0,
6250        theta: theta0,
6251        i1: i1_0,
6252        i2: i2_0,
6253        v_rest,
6254        v_reset,
6255        theta_reset,
6256        theta_inf,
6257        tau_v,
6258        tau_theta,
6259        tau_1,
6260        tau_2,
6261        a,
6262        b,
6263        r1,
6264        r2,
6265        dt,
6266    };
6267    let (trace, spikes) = neuron.simulate(n_steps, current);
6268    (
6269        trace.into_pyarray(py),
6270        spikes,
6271        neuron.v,
6272        neuron.theta,
6273        neuron.i1,
6274        neuron.i2,
6275    )
6276}
6277
6278/// Parity contract with
6279/// `sc_neurocore.neurons.models.glif.GLIFNeuron.simulate`: for the same
6280/// parameters and constant input the returned `v` trace, total spike count, and
6281/// final `(v, theta, i_asc1, i_asc2)` state are bit-identical to the Python RK4
6282/// reference. The Allen GLIF5 model has a purely linear right-hand side (no
6283/// transcendental functions), so every RK4 stage is exact arithmetic and the
6284/// trace reproduces to the last bit across Rust/Julia/Go (Mojo FMA-fuses,
6285/// validated non-amplifying).
6286#[pyfunction]
6287#[pyo3(signature = (
6288    v0, theta0, theta_inf, i_asc1_0, i_asc2_0, v_rest, v_reset, tau_m, tau_theta,
6289    tau_asc1, tau_asc2, a_theta, delta_theta, r_asc1, r_asc2, resistance, dt,
6290    n_steps, current
6291))]
6292#[allow(clippy::too_many_arguments)]
6293fn py_glif_simulate<'py>(
6294    py: Python<'py>,
6295    v0: f64,
6296    theta0: f64,
6297    theta_inf: f64,
6298    i_asc1_0: f64,
6299    i_asc2_0: f64,
6300    v_rest: f64,
6301    v_reset: f64,
6302    tau_m: f64,
6303    tau_theta: f64,
6304    tau_asc1: f64,
6305    tau_asc2: f64,
6306    a_theta: f64,
6307    delta_theta: f64,
6308    r_asc1: f64,
6309    r_asc2: f64,
6310    resistance: f64,
6311    dt: f64,
6312    n_steps: usize,
6313    current: f64,
6314) -> (Bound<'py, PyArray1<f64>>, i64, f64, f64, f64, f64) {
6315    let mut neuron = crate::neurons::GLIFNeuron {
6316        v: v0,
6317        theta: theta0,
6318        theta_inf,
6319        i_asc1: i_asc1_0,
6320        i_asc2: i_asc2_0,
6321        v_rest,
6322        v_reset,
6323        tau_m,
6324        tau_theta,
6325        tau_asc1,
6326        tau_asc2,
6327        a_theta,
6328        delta_theta,
6329        r_asc1,
6330        r_asc2,
6331        resistance,
6332        dt,
6333    };
6334    let (trace, spikes) = neuron.simulate(n_steps, current);
6335    (
6336        trace.into_pyarray(py),
6337        spikes,
6338        neuron.v,
6339        neuron.theta,
6340        neuron.i_asc1,
6341        neuron.i_asc2,
6342    )
6343}
6344
6345/// Parity contract with
6346/// `sc_neurocore.neurons.models.fitzhugh_nagumo.FitzHughNagumoNeuron.simulate`:
6347/// for the same parameters and constant input the returned `v` trace, upward-
6348/// crossing spike count, and final `(v, w)` state are bit-identical to the
6349/// Python RK4 reference (the right-hand side is exact arithmetic — a cube
6350/// `v.powi(3)` = `v*v*v`, additions and multiplications, no transcendental
6351/// functions — and a two-dimensional flow cannot be chaotic).
6352#[pyfunction]
6353#[pyo3(signature = (v0, w0, a, b, epsilon, dt, v_threshold, n_steps, current))]
6354#[allow(clippy::too_many_arguments)]
6355fn py_fitzhugh_nagumo_simulate<'py>(
6356    py: Python<'py>,
6357    v0: f64,
6358    w0: f64,
6359    a: f64,
6360    b: f64,
6361    epsilon: f64,
6362    dt: f64,
6363    v_threshold: f64,
6364    n_steps: usize,
6365    current: f64,
6366) -> (Bound<'py, PyArray1<f64>>, i64, f64, f64) {
6367    let mut neuron = crate::neurons::FitzHughNagumoNeuron {
6368        v: v0,
6369        w: w0,
6370        a,
6371        b,
6372        epsilon,
6373        dt,
6374        v_threshold,
6375    };
6376    let (trace, spikes) = neuron.simulate(n_steps, current);
6377    (trace.into_pyarray(py), spikes, neuron.v, neuron.w)
6378}
6379
6380/// Parity contract with
6381/// `sc_neurocore.neurons.models.hindmarsh_rose.HindmarshRoseNeuron.simulate`:
6382/// for the same parameters and constant input the returned `x` trace, upward-
6383/// crossing spike count, and final `(x, y, z)` state are bit-identical to the
6384/// Python RK4 reference (the right-hand side is exact arithmetic — `x.powi(3)`
6385/// = `x*x*x`, `x.powi(2)` = `x*x`, no transcendental functions — so even the
6386/// chaotic bursting trace reproduces exactly).
6387#[pyfunction]
6388#[pyo3(signature = (x0, y0, z0, b, r, s, x_rest, dt, x_threshold, n_steps, current))]
6389#[allow(clippy::too_many_arguments)]
6390fn py_hindmarsh_rose_simulate<'py>(
6391    py: Python<'py>,
6392    x0: f64,
6393    y0: f64,
6394    z0: f64,
6395    b: f64,
6396    r: f64,
6397    s: f64,
6398    x_rest: f64,
6399    dt: f64,
6400    x_threshold: f64,
6401    n_steps: usize,
6402    current: f64,
6403) -> (Bound<'py, PyArray1<f64>>, i64, f64, f64, f64) {
6404    let mut neuron = crate::neurons::HindmarshRoseNeuron {
6405        x: x0,
6406        y: y0,
6407        z: z0,
6408        b,
6409        r,
6410        s,
6411        x_rest,
6412        dt,
6413        x_threshold,
6414    };
6415    let (trace, spikes) = neuron.simulate(n_steps, current);
6416    (trace.into_pyarray(py), spikes, neuron.x, neuron.y, neuron.z)
6417}
6418
6419/// Parity contract with
6420/// `sc_neurocore.neurons.models.fitzhugh_rinzel.FitzHughRinzelNeuron.simulate`:
6421/// for the same parameters and constant input the returned `v` trace, upward-
6422/// crossing spike count, and final `(v, w, y)` state are bit-identical to the
6423/// Python RK4 reference (the right-hand side is exact arithmetic — `v.powi(3)`
6424/// = `v*v*v`, additions and multiplications, no transcendental functions).
6425#[pyfunction]
6426#[pyo3(signature = (v0, w0, y0, a, b, c, d, delta, mu, dt, v_threshold, n_steps, current))]
6427#[allow(clippy::too_many_arguments)]
6428fn py_fitzhugh_rinzel_simulate<'py>(
6429    py: Python<'py>,
6430    v0: f64,
6431    w0: f64,
6432    y0: f64,
6433    a: f64,
6434    b: f64,
6435    c: f64,
6436    d: f64,
6437    delta: f64,
6438    mu: f64,
6439    dt: f64,
6440    v_threshold: f64,
6441    n_steps: usize,
6442    current: f64,
6443) -> (Bound<'py, PyArray1<f64>>, i64, f64, f64, f64) {
6444    let mut neuron = crate::neurons::FitzHughRinzelNeuron {
6445        v: v0,
6446        w: w0,
6447        y: y0,
6448        a,
6449        b,
6450        c,
6451        d,
6452        delta,
6453        mu,
6454        dt,
6455        v_threshold,
6456    };
6457    let (trace, spikes) = neuron.simulate(n_steps, current);
6458    (trace.into_pyarray(py), spikes, neuron.v, neuron.w, neuron.y)
6459}
6460
6461/// Parity contract with
6462/// `sc_neurocore.neurons.models.izhikevich2007.Izhikevich2007Neuron.simulate`:
6463/// for the same parameters and constant input the returned `v` trace, spike
6464/// count, and final `(v, u)` state are bit-identical to the Python RK4 reference
6465/// (the NeuroML right-hand side `k (v-vr)(v-vt)/C` is exact arithmetic — products,
6466/// a sum and a division, no transcendental functions).
6467#[pyfunction]
6468#[pyo3(signature = (v0, u0, cap, k, vr, vt, vpeak, a, b, c, d, dt, n_steps, current))]
6469#[allow(clippy::too_many_arguments)]
6470fn py_izhikevich2007_simulate<'py>(
6471    py: Python<'py>,
6472    v0: f64,
6473    u0: f64,
6474    cap: f64,
6475    k: f64,
6476    vr: f64,
6477    vt: f64,
6478    vpeak: f64,
6479    a: f64,
6480    b: f64,
6481    c: f64,
6482    d: f64,
6483    dt: f64,
6484    n_steps: usize,
6485    current: f64,
6486) -> (Bound<'py, PyArray1<f64>>, i64, f64, f64) {
6487    let mut neuron = crate::rk4_neurons::Izhikevich2007Rk4 {
6488        v: v0,
6489        u: u0,
6490        cap,
6491        k,
6492        vr,
6493        vt,
6494        vpeak,
6495        a,
6496        b,
6497        c,
6498        d,
6499        dt,
6500    };
6501    let (trace, spikes) = neuron.simulate(n_steps, current);
6502    (trace.into_pyarray(py), spikes, neuron.v, neuron.u)
6503}
6504
6505// ── Byte-level fault injection (PyO3) ──
6506//
6507// Each function consumes a numpy bool/uint8 array (read-only), copies
6508// into a Vec<u8>, applies the inject kernel in pure Rust, and returns
6509// (corrupted_array, num_affected). RNG seed is per-call so back-to-back
6510// invocations are reproducible at the API level.
6511
6512#[pyfunction]
6513fn py_inject_bitflip_u8<'py>(
6514    py: Python<'py>,
6515    bitstream: PyReadonlyArray1<'_, u8>,
6516    ber: f64,
6517    seed: u64,
6518) -> PyResult<(Py<PyArray1<u8>>, u64)> {
6519    let mut buf = bitstream.as_slice()?.to_vec();
6520    let n = fault::inject_bitflip_u8(&mut buf, ber, seed);
6521    Ok((buf.into_pyarray(py).into(), n))
6522}
6523
6524#[pyfunction]
6525fn py_inject_stuck_at_0_u8<'py>(
6526    py: Python<'py>,
6527    bitstream: PyReadonlyArray1<'_, u8>,
6528    ber: f64,
6529    seed: u64,
6530) -> PyResult<(Py<PyArray1<u8>>, u64)> {
6531    let mut buf = bitstream.as_slice()?.to_vec();
6532    let n = fault::inject_stuck_at_0_u8(&mut buf, ber, seed);
6533    Ok((buf.into_pyarray(py).into(), n))
6534}
6535
6536#[pyfunction]
6537fn py_inject_stuck_at_1_u8<'py>(
6538    py: Python<'py>,
6539    bitstream: PyReadonlyArray1<'_, u8>,
6540    ber: f64,
6541    seed: u64,
6542) -> PyResult<(Py<PyArray1<u8>>, u64)> {
6543    let mut buf = bitstream.as_slice()?.to_vec();
6544    let n = fault::inject_stuck_at_1_u8(&mut buf, ber, seed);
6545    Ok((buf.into_pyarray(py).into(), n))
6546}
6547
6548#[pyfunction]
6549fn py_inject_dropout_u8<'py>(
6550    py: Python<'py>,
6551    bitstream: PyReadonlyArray1<'_, u8>,
6552    ber: f64,
6553    seed: u64,
6554) -> PyResult<(Py<PyArray1<u8>>, u64)> {
6555    let mut buf = bitstream.as_slice()?.to_vec();
6556    let n = fault::inject_dropout_u8(&mut buf, ber, seed);
6557    Ok((buf.into_pyarray(py).into(), n))
6558}
6559
6560#[pyfunction]
6561fn py_inject_gaussian_u8<'py>(
6562    py: Python<'py>,
6563    bitstream: PyReadonlyArray1<'_, u8>,
6564    ber: f64,
6565    seed: u64,
6566) -> PyResult<(Py<PyArray1<u8>>, u64)> {
6567    let mut buf = bitstream.as_slice()?.to_vec();
6568    let n = fault::inject_gaussian_u8(&mut buf, ber, seed);
6569    Ok((buf.into_pyarray(py).into(), n))
6570}
6571
6572// ── Hierarchical partitioner — KL refine (PyO3) ──
6573//
6574// Caller passes flat numpy arrays (CSR adjacency + flat scc weights +
6575// flat vertex_weights + initial part_map). The kernel mutates a copy
6576// of part_map in-place and returns (new_part_map, num_moves).
6577
6578#[pyfunction]
6579#[pyo3(signature = (
6580    adj_offsets, adj_neighbours, adj_scc_abs, vertex_weights,
6581    part_map, parts_concat, parts_offsets,
6582    n_parts, kl_iterations, correlation_penalty,
6583))]
6584#[allow(clippy::too_many_arguments)]
6585fn py_kl_refine<'py>(
6586    py: Python<'py>,
6587    adj_offsets: PyReadonlyArray1<'_, i64>,
6588    adj_neighbours: PyReadonlyArray1<'_, i32>,
6589    adj_scc_abs: PyReadonlyArray1<'_, f64>,
6590    vertex_weights: PyReadonlyArray1<'_, f64>,
6591    part_map: PyReadonlyArray1<'_, i32>,
6592    parts_concat: PyReadonlyArray1<'_, i32>,
6593    parts_offsets: PyReadonlyArray1<'_, i64>,
6594    n_parts: i32,
6595    kl_iterations: i32,
6596    correlation_penalty: f64,
6597) -> PyResult<(Py<PyArray1<i32>>, u64)> {
6598    let mut pm = part_map.as_slice()?.to_vec();
6599    let moves = partition::kl_refine(
6600        adj_offsets.as_slice()?,
6601        adj_neighbours.as_slice()?,
6602        adj_scc_abs.as_slice()?,
6603        vertex_weights.as_slice()?,
6604        &mut pm,
6605        parts_concat.as_slice()?,
6606        parts_offsets.as_slice()?,
6607        n_parts,
6608        kl_iterations,
6609        correlation_penalty,
6610    );
6611    Ok((pm.into_pyarray(py).into(), moves))
6612}
6613
6614// ── PINGCircuit per-step kernel ─────────────────────────────────────
6615//
6616// Mirrors `PINGCircuit.step()` in
6617// `src/sc_neurocore/network/gamma_oscillation.py`. Caller hands over
6618// the per-instance state arrays + this-step `xi` noise samples (drawn
6619// on the Python side from the per-instance RNG so seed determinism is
6620// preserved), and the kernel writes the new state in-place plus a
6621// boolean spike vector per population. Returns `(n_e_spikes,
6622// n_i_spikes)` so the caller can do the synaptic-conductance update
6623// step (which needs the cross-population summary, not the per-cell
6624// detail).
6625
6626#[pyfunction]
6627#[pyo3(signature = (
6628    v_e, g_ampa_e, g_gaba_e, refrac_e, i_drive_e, xi_e, spikes_e_out,
6629    v_i, g_ampa_i, g_gaba_i, refrac_i, i_drive_i, xi_i, spikes_i_out,
6630    e_l, e_ampa, e_gaba, g_l, c_m, v_threshold, v_reset, t_refrac,
6631    tau_ampa, tau_gaba, sigma_e, sigma_i, dt,
6632))]
6633#[allow(clippy::too_many_arguments)]
6634fn py_ping_step<'py>(
6635    _py: Python<'py>,
6636    v_e: PyReadwriteArray1<'_, f64>,
6637    g_ampa_e: PyReadwriteArray1<'_, f64>,
6638    g_gaba_e: PyReadwriteArray1<'_, f64>,
6639    refrac_e: PyReadwriteArray1<'_, f64>,
6640    i_drive_e: PyReadonlyArray1<'_, f64>,
6641    xi_e: PyReadonlyArray1<'_, f64>,
6642    spikes_e_out: PyReadwriteArray1<'_, u8>,
6643    v_i: PyReadwriteArray1<'_, f64>,
6644    g_ampa_i: PyReadwriteArray1<'_, f64>,
6645    g_gaba_i: PyReadwriteArray1<'_, f64>,
6646    refrac_i: PyReadwriteArray1<'_, f64>,
6647    i_drive_i: PyReadonlyArray1<'_, f64>,
6648    xi_i: PyReadonlyArray1<'_, f64>,
6649    spikes_i_out: PyReadwriteArray1<'_, u8>,
6650    e_l: f64,
6651    e_ampa: f64,
6652    e_gaba: f64,
6653    g_l: f64,
6654    c_m: f64,
6655    v_threshold: f64,
6656    v_reset: f64,
6657    t_refrac: f64,
6658    tau_ampa: f64,
6659    tau_gaba: f64,
6660    sigma_e: f64,
6661    sigma_i: f64,
6662    dt: f64,
6663) -> PyResult<(u32, u32)> {
6664    let mut v_e = v_e;
6665    let mut g_ampa_e = g_ampa_e;
6666    let mut g_gaba_e = g_gaba_e;
6667    let mut refrac_e = refrac_e;
6668    let mut spikes_e_out = spikes_e_out;
6669    let mut v_i = v_i;
6670    let mut g_ampa_i = g_ampa_i;
6671    let mut g_gaba_i = g_gaba_i;
6672    let mut refrac_i = refrac_i;
6673    let mut spikes_i_out = spikes_i_out;
6674    let (ne, ni) = ping::step_kernel(
6675        v_e.as_slice_mut()?,
6676        g_ampa_e.as_slice_mut()?,
6677        g_gaba_e.as_slice_mut()?,
6678        refrac_e.as_slice_mut()?,
6679        i_drive_e.as_slice()?,
6680        xi_e.as_slice()?,
6681        spikes_e_out.as_slice_mut()?,
6682        v_i.as_slice_mut()?,
6683        g_ampa_i.as_slice_mut()?,
6684        g_gaba_i.as_slice_mut()?,
6685        refrac_i.as_slice_mut()?,
6686        i_drive_i.as_slice()?,
6687        xi_i.as_slice()?,
6688        spikes_i_out.as_slice_mut()?,
6689        e_l,
6690        e_ampa,
6691        e_gaba,
6692        g_l,
6693        c_m,
6694        v_threshold,
6695        v_reset,
6696        t_refrac,
6697        tau_ampa,
6698        tau_gaba,
6699        sigma_e,
6700        sigma_i,
6701        dt,
6702    );
6703    Ok((ne, ni))
6704}
6705
6706// ── CorticalColumn block-CSR spmv (per-row-parallel) ────────────────
6707//
6708// `y += W @ x` where `W` is a CSR matrix described by
6709// `(indptr, indices, data)`. Rows are processed in parallel via
6710// rayon — bit-identical to scipy single-threaded for matching
6711// inputs because the per-row reduction is local. Used by
6712// `CorticalColumn._inject_block(dt)` once per `(source-type, bin)`
6713// pair, replacing scipy's single-threaded csr_matvec for that step.
6714
6715#[pyfunction]
6716#[pyo3(signature = (indptr, indices, data, x, y))]
6717fn py_parallel_csr_spmv_add(
6718    indptr: PyReadonlyArray1<'_, i32>,
6719    indices: PyReadonlyArray1<'_, i32>,
6720    data: PyReadonlyArray1<'_, f64>,
6721    x: PyReadonlyArray1<'_, f64>,
6722    y: PyReadwriteArray1<'_, f64>,
6723) -> PyResult<()> {
6724    let mut y = y;
6725    cortical_inject::parallel_csr_spmv_add(
6726        indptr.as_slice()?,
6727        indices.as_slice()?,
6728        data.as_slice()?,
6729        x.as_slice()?,
6730        y.as_slice_mut()?,
6731    );
6732    Ok(())
6733}
6734
6735// Batched multi-spmv. The Python side passes lists of numpy arrays
6736// (one per block) for indptrs / indices / data / xs and a single
6737// mutable output. ONE FFI call per step replaces N (= n_delay_bins
6738// for E + same for I, typically 10) per-block calls. At scale=0.1
6739// the per-call savings amortise over a 600 ms simulation.
6740#[pyfunction]
6741#[pyo3(signature = (indptrs, indices_list, data_list, xs, y))]
6742fn py_parallel_csr_multi_spmv_add(
6743    indptrs: Vec<PyReadonlyArray1<'_, i32>>,
6744    indices_list: Vec<PyReadonlyArray1<'_, i32>>,
6745    data_list: Vec<PyReadonlyArray1<'_, f64>>,
6746    xs: Vec<PyReadonlyArray1<'_, f64>>,
6747    y: PyReadwriteArray1<'_, f64>,
6748) -> PyResult<()> {
6749    let mut y = y;
6750    let indptr_slices: Vec<&[i32]> = indptrs
6751        .iter()
6752        .map(|a| a.as_slice())
6753        .collect::<Result<_, _>>()?;
6754    let indices_slices: Vec<&[i32]> = indices_list
6755        .iter()
6756        .map(|a| a.as_slice())
6757        .collect::<Result<_, _>>()?;
6758    let data_slices: Vec<&[f64]> = data_list
6759        .iter()
6760        .map(|a| a.as_slice())
6761        .collect::<Result<_, _>>()?;
6762    let x_slices: Vec<&[f64]> = xs.iter().map(|a| a.as_slice()).collect::<Result<_, _>>()?;
6763    cortical_inject::parallel_csr_multi_spmv_add(
6764        &indptr_slices,
6765        &indices_slices,
6766        &data_slices,
6767        &x_slices,
6768        y.as_slice_mut()?,
6769    );
6770    Ok(())
6771}