legume-numeric 0.9.0

use crate::matrix::rand_util::mix_seed;
use crate::matrix::traits::ConvertMatOps;
use anyhow::anyhow;
use candle_core::{Device, Tensor};
use rand::prelude::SliceRandom;
use rayon::prelude::*;

/// Upload a matrix to `device` as a contiguous `[N, D]` tensor.
/// `to_tensor` returns a transposed view; `contiguous()` makes it
/// layout-compatible with `index_select`.
pub(crate) fn upload_to_device<D: ConvertMatOps>(x: &D, device: &Device) -> anyhow::Result<Tensor> {
    Ok(x.to_tensor(device)?.contiguous()?)
}

/// Upload a column-major `[D, P]` matrix as the `[P, D]` tensor it already is.
///
/// nalgebra stores `[D, P]` column-major — element `(d, p)` at `d + p·D` — which
/// is byte for byte the row-major layout of `[P, D]`. So a caller holding the
/// gene-by-sample matrix needs no transpose at all to get sample-by-gene rows:
/// not the host `DMatrix::transpose()` (a strided copy of the whole matrix), and
/// not the device `contiguous()` that [`upload_to_device`] pays to undo
/// `to_tensor`'s transposed view. The two transposes were cancelling.
pub(crate) fn upload_columns_as_rows(
    x: &nalgebra::DMatrix<f32>,
    device: &Device,
) -> anyhow::Result<Tensor> {
    Ok(Tensor::from_slice(
        x.as_slice(),
        (x.ncols(), x.nrows()),
        device,
    )?)
}

/// Bootstrap-sample `ntot` indices from `[0, n)` with replacement, in
/// parallel, drawn from `seed` alone, so a run replays whatever the thread
/// count. Shared by
/// `Minibatches::shuffle_minibatch` (CPU path, `usize`) and the
/// device-resident loaders (`u32`).
pub(crate) fn bootstrap_indices<I>(n: usize, ntot: usize, seed: u64) -> Vec<I>
where
    I: TryFrom<usize> + Send,
    <I as TryFrom<usize>>::Error: std::fmt::Debug,
{
    use rand::{rngs::StdRng, SeedableRng};
    use rand_distr::{Distribution, Uniform};
    const CHUNK: usize = 1 << 16;
    let unif = Uniform::new(0usize, n).expect("unif [0 .. n)");
    // One sub-stream per chunk, collected in order, so the draw is fixed
    // whatever rayon's schedule.
    (0..ntot.div_ceil(CHUNK))
        .into_par_iter()
        .flat_map_iter(|ci| {
            let mut rng = StdRng::seed_from_u64(mix_seed(seed, ci as u64 + 1));
            (ci * CHUNK..((ci + 1) * CHUNK).min(ntot)).map(move |_| {
                I::try_from(unif.sample(&mut rng)).expect("index fits in target type")
            })
        })
        .collect()
}

/// Minibatch index bookkeeping.
///
/// Despite the `shuffle_minibatch` name, [`Minibatches::shuffle_minibatch`]
/// **bootstrap-resamples with replacement** — `chunks` is a bootstrap
/// *cover*, not a partition: within one call some indices appear more than
/// once and others not at all. (`samples` is shuffled but currently
/// unused by `chunks` construction — kept for callers that read it.)
pub struct Minibatches {
    pub samples: Vec<usize>,
    pub chunks: Vec<Vec<usize>>,
}

impl Minibatches {
    /// Bootstrap-resample the data into `chunks` of `batch_size` indices.
    ///
    /// Not a partition: indices are drawn *with replacement* via
    /// `bootstrap_indices`, so one pass over `chunks` is a bootstrap
    /// cover: some samples repeat, some are skipped. Every draw comes from
    /// `seed`, and the chunks keep their order, so a run replays.
    pub fn shuffle_minibatch(&mut self, batch_size: usize, seed: u64) {
        use rand::{rngs::StdRng, SeedableRng};
        let mut rng = StdRng::seed_from_u64(seed);
        self.samples.shuffle(&mut rng);

        let nbatch = (self.size() + batch_size) / batch_size;
        let ntot = nbatch * batch_size;

        let indexes: Vec<usize> = bootstrap_indices(self.size(), ntot, mix_seed(seed, 1));

        self.chunks = (0..nbatch)
            .into_par_iter()
            .map(|b| {
                let lb = b * batch_size;
                let ub = (b + 1) * batch_size;
                (lb..ub).map(|i| indexes[i]).collect()
            })
            .collect::<Vec<Vec<usize>>>();
    }

    pub fn size(&self) -> usize {
        self.samples.len()
    }
}

pub fn take_lb_ub(
    lb: usize,
    ub: usize,
    target_device: &Device,
    data_vec: Option<&Vec<Tensor>>,
) -> anyhow::Result<Option<Tensor>> {
    if let Some(data_vec) = data_vec {
        if lb > ub || ub > data_vec.len() {
            return Err(anyhow!(
                "check lb {}, ub {} vs. ntot {}",
                lb,
                ub,
                data_vec.len()
            ));
        }
        if lb == ub {
            return Ok(None);
        }

        let chunk = Tensor::cat(
            &(lb..ub).map(|i| data_vec[i].clone()).collect::<Vec<_>>(),
            0,
        )?;
        Ok(Some(chunk.to_device(target_device)?))
    } else {
        Ok(None)
    }
}

pub fn copy_shuffled(
    samples: &[usize],
    data: Option<&Vec<Tensor>>,
    shuffled_data: Option<&mut Vec<Tensor>>,
) -> anyhow::Result<()> {
    if let (Some(data), Some(shuffled)) = (data, shuffled_data) {
        let chunk: Vec<Tensor> = samples.iter().map(|&i| data[i].clone()).collect();
        let x = Tensor::cat(&chunk, 0)?;
        shuffled.push(x);
    }
    Ok(())
}

pub fn take_shuffled(
    batch_idx: usize,
    target_device: &Device,
    data_vec: Option<&Vec<Tensor>>,
) -> anyhow::Result<Option<Tensor>> {
    if let Some(data_vec) = data_vec {
        if data_vec.len() <= batch_idx {
            Err(anyhow!(
                "invalid index = {} vs. total # = {}",
                batch_idx,
                data_vec.len()
            ))
        } else {
            Ok(Some(data_vec[batch_idx].to_device(target_device)?))
        }
    } else {
        // if the data vector doesn't exist
        Ok(None)
    }
}