legume-numeric 0.8.11

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. Shared by `Minibatches::shuffle_minibatch` (CPU path,
/// `usize`) and the device-resident loaders (`u32`).
pub(crate) fn bootstrap_indices<I>(n: usize, ntot: usize) -> Vec<I>
where
    I: TryFrom<usize> + Send,
    <I as TryFrom<usize>>::Error: std::fmt::Debug,
{
    use rand_distr::{Distribution, Uniform};
    let unif = Uniform::new(0usize, n).expect("unif [0 .. n)");
    (0..ntot)
        .into_par_iter()
        .map_init(rand::rng, |rng, _| {
            I::try_from(unif.sample(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.
    pub fn shuffle_minibatch(&mut self, batch_size: usize) {
        let mut rng = rand::rng();
        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);

        self.chunks = (0..nbatch)
            .par_bridge()
            .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)
    }
}