scirs2-series 0.3.3

//! Latent Neural ODE for irregular time series
//!
//! The Latent ODE model (Chen et al., 2018; Rubanova et al., 2019) combines:
//!
//! 1. **ODE-RNN encoder** – reads (time, observation) pairs in reverse time order,
//!    updating a hidden state between observations via a local ODE solve.
//! 2. **Variational posterior** – maps the final hidden state to a distribution
//!    over initial latent positions `z₀`.
//! 3. **Neural ODE latent dynamics** – integrates `dz/dt = f_θ(z)` forward from `z₀`.
//! 4. **Decoder** – projects latent trajectory points back to observation space.
//!
//! All ODE solves use a fixed-step 4th-order Runge-Kutta integrator so that the
//! implementation is self-contained (no external ODE library needed).

use crate::error::{Result, TimeSeriesError};
use scirs2_core::ndarray::{Array1, Array2};
use scirs2_core::numeric::{Float, FromPrimitive};
use std::fmt::Debug;

// ---------------------------------------------------------------------------
// Small helpers
// ---------------------------------------------------------------------------

/// Sigmoid activation.
#[inline]
fn sigmoid<F: Float + FromPrimitive>(x: F) -> F {
    let one = F::from(1.0).expect("1.0");
    one / (one + (-x).exp())
}

/// Tanh activation (delegates to the trait method).
#[inline]
fn tanh_act<F: Float>(x: F) -> F {
    x.tanh()
}

/// Dense (fully-connected) layer: `y = tanh(W x + b)`.
fn dense_tanh<F: Float + FromPrimitive>(
    w: &Array2<F>,
    b: &Array1<F>,
    x: &Array1<F>,
) -> Array1<F> {
    let out_dim = w.nrows();
    let mut y = Array1::zeros(out_dim);
    for i in 0..out_dim {
        let mut s = b[i];
        for j in 0..x.len() {
            s = s + w[[i, j]] * x[j];
        }
        y[i] = tanh_act(s);
    }
    y
}

/// Dense layer without activation: `y = W x + b`.
fn dense_linear<F: Float + FromPrimitive>(
    w: &Array2<F>,
    b: &Array1<F>,
    x: &Array1<F>,
) -> Array1<F> {
    let out_dim = w.nrows();
    let mut y = Array1::zeros(out_dim);
    for i in 0..out_dim {
        let mut s = b[i];
        for j in 0..x.len() {
            s = s + w[[i, j]] * x[j];
        }
        y[i] = s;
    }
    y
}

// ---------------------------------------------------------------------------
// GRU cell used inside the ODE-RNN encoder
// ---------------------------------------------------------------------------

/// Minimal GRU cell operating on `Array1`.
///
/// Gates:
/// * reset gate  `r = σ(Wr [x; h] + br)`
/// * update gate `u = σ(Wu [x; h] + bu)`
/// * candidate   `n = tanh(Wn [x; r⊙h] + bn)`
/// * new hidden  `h' = (1-u)⊙h + u⊙n`
#[derive(Debug, Clone)]
pub struct GruCell<F: Float + Debug + FromPrimitive + Clone> {
    /// Input + hidden dim for concatenated input
    in_dim: usize,
    /// Hidden dimension
    hidden_dim: usize,
    /// Reset-gate weight matrix  [hidden × (in + hidden)]
    w_r: Array2<F>,
    /// Update-gate weight matrix [hidden × (in + hidden)]
    w_u: Array2<F>,
    /// Candidate weight matrix   [hidden × (in + hidden)]
    w_n: Array2<F>,
    /// Biases [hidden]
    b_r: Array1<F>,
    b_u: Array1<F>,
    b_n: Array1<F>,
}

impl<F: Float + Debug + FromPrimitive + Clone> GruCell<F> {
    /// Create a new GRU cell with pseudo-random weight initialisation.
    pub fn new(input_dim: usize, hidden_dim: usize, seed: u64) -> Self {
        let concat_dim = input_dim + hidden_dim;
        let std_dev = F::from(2.0 / (concat_dim + hidden_dim) as f64)
            .expect("std_dev")
            .sqrt();

        Self {
            in_dim: input_dim,
            hidden_dim,
            w_r: random_matrix(hidden_dim, concat_dim, std_dev, seed),
            w_u: random_matrix(hidden_dim, concat_dim, std_dev, seed.wrapping_add(1)),
            w_n: random_matrix(hidden_dim, concat_dim, std_dev, seed.wrapping_add(2)),
            b_r: Array1::zeros(hidden_dim),
            b_u: Array1::zeros(hidden_dim),
            b_n: Array1::zeros(hidden_dim),
        }
    }

    /// Forward step: returns next hidden state.
    pub fn forward(&self, x: &Array1<F>, h: &Array1<F>) -> Result<Array1<F>> {
        if x.len() != self.in_dim {
            return Err(TimeSeriesError::DimensionMismatch {
                expected: self.in_dim,
                actual: x.len(),
            });
        }
        if h.len() != self.hidden_dim {
            return Err(TimeSeriesError::DimensionMismatch {
                expected: self.hidden_dim,
                actual: h.len(),
            });
        }

        // Concatenate [x; h]
        let mut xh = Array1::zeros(self.in_dim + self.hidden_dim);
        for i in 0..self.in_dim {
            xh[i] = x[i];
        }
        for i in 0..self.hidden_dim {
            xh[self.in_dim + i] = h[i];
        }

        // Reset and update gates
        let r_raw = dense_linear(&self.w_r, &self.b_r, &xh);
        let u_raw = dense_linear(&self.w_u, &self.b_u, &xh);
        let mut r = Array1::zeros(self.hidden_dim);
        let mut u = Array1::zeros(self.hidden_dim);
        for i in 0..self.hidden_dim {
            r[i] = sigmoid(r_raw[i]);
            u[i] = sigmoid(u_raw[i]);
        }

        // Candidate: [x; r⊙h]
        let mut xrh = Array1::zeros(self.in_dim + self.hidden_dim);
        for i in 0..self.in_dim {
            xrh[i] = x[i];
        }
        for i in 0..self.hidden_dim {
            xrh[self.in_dim + i] = r[i] * h[i];
        }
        let n_raw = dense_linear(&self.w_n, &self.b_n, &xrh);

        // New hidden
        let one = F::from(1.0).expect("1.0");
        let mut h_new = Array1::zeros(self.hidden_dim);
        for i in 0..self.hidden_dim {
            h_new[i] = (one - u[i]) * h[i] + u[i] * tanh_act(n_raw[i]);
        }
        Ok(h_new)
    }
}

// ---------------------------------------------------------------------------
// Latent dynamics network  f_θ(z)
// ---------------------------------------------------------------------------

/// Two-layer MLP that parameterises the ODE vector field `dz/dt = f(z)`.
#[derive(Debug, Clone)]
pub struct LatentDynamics<F: Float + Debug + FromPrimitive + Clone> {
    latent_dim: usize,
    /// First hidden layer
    w1: Array2<F>,
    b1: Array1<F>,
    /// Second hidden layer (output, same size as latent)
    w2: Array2<F>,
    b2: Array1<F>,
}

impl<F: Float + Debug + FromPrimitive + Clone> LatentDynamics<F> {
    /// Create dynamics network with `hidden_dim` internal width.
    pub fn new(latent_dim: usize, hidden_dim: usize, seed: u64) -> Self {
        let std1 = F::from(2.0 / (latent_dim + hidden_dim) as f64)
            .expect("std")
            .sqrt();
        let std2 = F::from(2.0 / (hidden_dim + latent_dim) as f64)
            .expect("std")
            .sqrt();
        Self {
            latent_dim,
            w1: random_matrix(hidden_dim, latent_dim, std1, seed),
            b1: Array1::zeros(hidden_dim),
            w2: random_matrix(latent_dim, hidden_dim, std2, seed.wrapping_add(10)),
            b2: Array1::zeros(latent_dim),
        }
    }

    /// Evaluate `f(z)`.
    pub fn forward(&self, z: &Array1<F>) -> Result<Array1<F>> {
        if z.len() != self.latent_dim {
            return Err(TimeSeriesError::DimensionMismatch {
                expected: self.latent_dim,
                actual: z.len(),
            });
        }
        let h = dense_tanh(&self.w1, &self.b1, z);
        Ok(dense_tanh(&self.w2, &self.b2, &h))
    }
}

// ---------------------------------------------------------------------------
// Fixed-step RK4 integrator
// ---------------------------------------------------------------------------

/// Integrate `dz/dt = f(z)` from `t0` to `t1` with `n_steps` RK4 steps.
fn rk4_integrate<F: Float + Debug + FromPrimitive + Clone>(
    f: &dyn Fn(&Array1<F>) -> Result<Array1<F>>,
    z0: &Array1<F>,
    t0: F,
    t1: F,
    n_steps: usize,
) -> Result<Array1<F>> {
    if n_steps == 0 {
        return Ok(z0.clone());
    }
    let dim = z0.len();
    let h = (t1 - t0) / F::from(n_steps).expect("n_steps");
    let half = F::from(0.5).expect("0.5");
    let sixth = F::from(1.0 / 6.0).expect("1/6");

    let mut z = z0.clone();

    for _ in 0..n_steps {
        let k1 = f(&z)?;
        let z_k2: Array1<F> = {
            let mut tmp = Array1::zeros(dim);
            for i in 0..dim {
                tmp[i] = z[i] + half * h * k1[i];
            }
            tmp
        };
        let k2 = f(&z_k2)?;
        let z_k3: Array1<F> = {
            let mut tmp = Array1::zeros(dim);
            for i in 0..dim {
                tmp[i] = z[i] + half * h * k2[i];
            }
            tmp
        };
        let k3 = f(&z_k3)?;
        let z_k4: Array1<F> = {
            let mut tmp = Array1::zeros(dim);
            for i in 0..dim {
                tmp[i] = z[i] + h * k3[i];
            }
            tmp
        };
        let k4 = f(&z_k4)?;

        for i in 0..dim {
            z[i] = z[i] + h * sixth * (k1[i] + F::from(2.0).expect("2") * k2[i]
                + F::from(2.0).expect("2") * k3[i]
                + k4[i]);
        }
    }
    Ok(z)
}

// ---------------------------------------------------------------------------
// Decoder: latent → observation space
// ---------------------------------------------------------------------------

/// Linear decoder from latent space to observation space.
#[derive(Debug, Clone)]
pub struct Decoder<F: Float + Debug + FromPrimitive + Clone> {
    w: Array2<F>,
    b: Array1<F>,
    latent_dim: usize,
    obs_dim: usize,
}

impl<F: Float + Debug + FromPrimitive + Clone> Decoder<F> {
    /// Create a new Decoder with random Glorot initialisation.
    pub fn new(latent_dim: usize, obs_dim: usize, seed: u64) -> Self {
        let std_dev = F::from(2.0 / (latent_dim + obs_dim) as f64)
            .expect("std")
            .sqrt();
        Self {
            w: random_matrix(obs_dim, latent_dim, std_dev, seed),
            b: Array1::zeros(obs_dim),
            latent_dim,
            obs_dim,
        }
    }

    /// Decode a latent vector to observation space.
    pub fn decode(&self, z: &Array1<F>) -> Result<Array1<F>> {
        if z.len() != self.latent_dim {
            return Err(TimeSeriesError::DimensionMismatch {
                expected: self.latent_dim,
                actual: z.len(),
            });
        }
        Ok(dense_linear(&self.w, &self.b, z))
    }

    /// Observation dimension.
    pub fn obs_dim(&self) -> usize {
        self.obs_dim
    }
}

// ---------------------------------------------------------------------------
// Latent ODE model
// ---------------------------------------------------------------------------

/// Hyperparameters for the Latent ODE model.
#[derive(Debug, Clone)]
pub struct LatentODEConfig {
    /// Dimensionality of the observation vectors.
    pub obs_dim: usize,
    /// Dimensionality of the latent state `z`.
    pub latent_dim: usize,
    /// Hidden size of the GRU encoder.
    pub encoder_hidden: usize,
    /// Hidden size of the ODE dynamics MLP.
    pub ode_hidden: usize,
    /// Number of RK4 steps per unit time interval.
    pub ode_steps_per_unit: usize,
    /// Random seed for weight initialisation.
    pub seed: u64,
}

impl Default for LatentODEConfig {
    fn default() -> Self {
        Self {
            obs_dim: 1,
            latent_dim: 8,
            encoder_hidden: 16,
            ode_hidden: 32,
            ode_steps_per_unit: 10,
            seed: 42,
        }
    }
}

/// Latent Neural ODE model for irregular time series.
///
/// ## Model overview
///
/// ```text
///   Encoder (ODE-RNN, reverse time)
///     (t_n, y_n) → ... → (t_1, y_1) → h_0
///     ↓
///   z₀ = Linear(h_0)          (approximate posterior mean)
///     ↓
///   Neural ODE:  dz/dt = f_θ(z),  z(t₀) = z₀
///     ↓
///   z(t_i) for query times t_i
///     ↓
///   ŷ_i = Decoder(z(t_i))
/// ```
///
/// Weights are randomly initialised; to obtain a useful model the weights
/// must be trained externally (e.g. via autograd or numerical gradient
/// approximation).  The forward-pass machinery is fully implemented and
/// differentiable in principle.
#[derive(Debug, Clone)]
pub struct LatentODE<F: Float + Debug + FromPrimitive + Clone> {
    config: LatentODEConfig,
    /// GRU encoder cell  (input: obs_dim + 1 time delta, hidden: encoder_hidden)
    encoder_gru: GruCell<F>,
    /// Linear projection from encoder hidden → latent initial state z₀
    encoder_to_z0_w: Array2<F>,
    encoder_to_z0_b: Array1<F>,
    /// ODE dynamics network
    dynamics: LatentDynamics<F>,
    /// Decoder
    decoder: Decoder<F>,
}

impl<F: Float + Debug + FromPrimitive + Clone> LatentODE<F> {
    /// Construct a Latent ODE model with the given configuration.
    ///
    /// All weights are initialised with Xavier-like Glorot initialisation
    /// using the specified seed.
    pub fn new(config: LatentODEConfig) -> Result<Self> {
        if config.obs_dim == 0 {
            return Err(TimeSeriesError::InvalidParameter {
                name: "obs_dim".into(),
                message: "must be ≥ 1".into(),
            });
        }
        if config.latent_dim == 0 {
            return Err(TimeSeriesError::InvalidParameter {
                name: "latent_dim".into(),
                message: "must be ≥ 1".into(),
            });
        }

        let s = config.seed;
        // GRU input: observation vector + 1 time-delta scalar
        let gru_input_dim = config.obs_dim + 1;
        let encoder_gru = GruCell::new(gru_input_dim, config.encoder_hidden, s);

        let enc_std = F::from(
            (2.0 / (config.encoder_hidden + config.latent_dim) as f64).sqrt(),
        )
        .expect("std");
        let encoder_to_z0_w =
            random_matrix(config.latent_dim, config.encoder_hidden, enc_std, s.wrapping_add(100));
        let encoder_to_z0_b = Array1::zeros(config.latent_dim);

        let dynamics =
            LatentDynamics::new(config.latent_dim, config.ode_hidden, s.wrapping_add(200));
        let decoder = Decoder::new(config.latent_dim, config.obs_dim, s.wrapping_add(300));

        Ok(Self {
            config,
            encoder_gru,
            encoder_to_z0_w,
            encoder_to_z0_b,
            dynamics,
            decoder,
        })
    }

    // ------------------------------------------------------------------
    // ODE helper
    // ------------------------------------------------------------------

    fn integrate_latent(
        &self,
        z0: &Array1<F>,
        t_start: F,
        t_end: F,
    ) -> Result<Array1<F>> {
        let duration = (t_end - t_start).abs();
        let n_steps = ((duration
            * F::from(self.config.ode_steps_per_unit).expect("steps"))
        .ceil()
        .to_usize()
        .unwrap_or(1))
        .max(1);

        let dynamics_ref = &self.dynamics;
        rk4_integrate(
            &|z| dynamics_ref.forward(z),
            z0,
            t_start,
            t_end,
            n_steps,
        )
    }

    // ------------------------------------------------------------------
    // Public API
    // ------------------------------------------------------------------

    /// Encode an irregularly-sampled time series to an initial latent state z₀.
    ///
    /// # Arguments
    /// * `times` – observation timestamps (monotonically non-decreasing).
    /// * `observations` – list of observation vectors, one per timestamp.
    ///
    /// # Returns
    /// Initial latent state z₀ ∈ ℝ^`latent_dim`.
    pub fn encode(&self, times: &[F], observations: &[Array1<F>]) -> Result<Array1<F>> {
        if times.len() != observations.len() {
            return Err(TimeSeriesError::DimensionMismatch {
                expected: times.len(),
                actual: observations.len(),
            });
        }
        if times.is_empty() {
            return Err(TimeSeriesError::InvalidInput(
                "encode requires at least one observation".into(),
            ));
        }
        for obs in observations {
            if obs.len() != self.config.obs_dim {
                return Err(TimeSeriesError::DimensionMismatch {
                    expected: self.config.obs_dim,
                    actual: obs.len(),
                });
            }
        }

        let n = times.len();
        // Process observations in **reverse** time order (standard ODE-RNN)
        let mut h = Array1::zeros(self.config.encoder_hidden);

        for idx in (0..n).rev() {
            // Time delta to previous (future in forward time) observation
            let delta = if idx + 1 < n {
                times[idx + 1] - times[idx]
            } else {
                F::zero()
            };

            // ODE-solve h from t[idx+1] → t[idx] between GRU updates
            // (Approximation: for very small deltas we skip the ODE step)
            let abs_delta = delta.abs();
            if abs_delta > F::epsilon() {
                let steps = ((abs_delta
                    * F::from(self.config.ode_steps_per_unit).expect("steps"))
                .ceil()
                .to_usize()
                .unwrap_or(1))
                .max(1);
                let dynamics_ref = &self.dynamics;
                // We use the dynamics in hidden space as a proxy (in a full
                // implementation the encoder has its own ODE; here we share
                // the latent dynamics for simplicity).
                // Project h to latent, solve, project back is expensive; instead
                // we apply a simple exponential decay as the ODE between obs.
                let decay = F::from(0.1).expect("0.1");
                let dt = abs_delta / F::from(steps).expect("steps");
                for _ in 0..steps {
                    for v in h.iter_mut() {
                        *v = *v * (F::one() - decay * dt);
                    }
                }
            }

            // GRU update with [obs; delta]
            let mut gru_input = Array1::zeros(self.config.obs_dim + 1);
            for j in 0..self.config.obs_dim {
                gru_input[j] = observations[idx][j];
            }
            gru_input[self.config.obs_dim] = delta;

            h = self.encoder_gru.forward(&gru_input, &h)?;
        }

        // Project encoder hidden state → z₀
        Ok(dense_linear(
            &self.encoder_to_z0_w,
            &self.encoder_to_z0_b,
            &h,
        ))
    }

    /// Decode a latent trajectory: integrate from z₀ and decode at each query time.
    ///
    /// # Arguments
    /// * `z0` – initial latent state (output of `encode`).
    /// * `t_span` – query timestamps (monotonically non-decreasing).
    ///
    /// # Returns
    /// Decoded observation predictions, one per element in `t_span`.
    pub fn decode(&self, z0: &Array1<F>, t_span: &[F]) -> Result<Vec<Array1<F>>> {
        if z0.len() != self.config.latent_dim {
            return Err(TimeSeriesError::DimensionMismatch {
                expected: self.config.latent_dim,
                actual: z0.len(),
            });
        }
        if t_span.is_empty() {
            return Ok(Vec::new());
        }

        let mut predictions = Vec::with_capacity(t_span.len());
        let t0 = t_span[0];
        let mut z = self.integrate_latent(z0, F::zero(), t0)?;
        predictions.push(self.decoder.decode(&z)?);

        for i in 1..t_span.len() {
            z = self.integrate_latent(&z, t_span[i - 1], t_span[i])?;
            predictions.push(self.decoder.decode(&z)?);
        }

        Ok(predictions)
    }

    /// End-to-end prediction: encode observed (irregular) time series, then
    /// decode at future query times.
    ///
    /// # Arguments
    /// * `times` – timestamps of past observations.
    /// * `obs` – observed values at those timestamps.
    /// * `future_times` – future timestamps at which to predict.
    ///
    /// # Returns
    /// Predicted observations at each element of `future_times`.
    pub fn predict(
        &self,
        times: &[F],
        obs: &[Array1<F>],
        future_times: &[F],
    ) -> Result<Vec<Array1<F>>> {
        let z0 = self.encode(times, obs)?;
        self.decode(&z0, future_times)
    }

    /// Return the observation dimensionality of the model.
    pub fn obs_dim(&self) -> usize {
        self.config.obs_dim
    }

    /// Return the latent dimensionality of the model.
    pub fn latent_dim(&self) -> usize {
        self.config.latent_dim
    }

    /// Return a reference to the model configuration.
    pub fn config(&self) -> &LatentODEConfig {
        &self.config
    }
}

// ---------------------------------------------------------------------------
// Utility: weight initialisation
// ---------------------------------------------------------------------------

/// Generate a pseudo-random matrix using a linear-congruential generator.
fn random_matrix<F: Float + FromPrimitive>(
    rows: usize,
    cols: usize,
    std_dev: F,
    seed: u64,
) -> Array2<F> {
    let mut mat = Array2::zeros((rows, cols));
    let mut state = seed.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
    for i in 0..rows {
        for j in 0..cols {
            // xorshift64
            state ^= state << 13;
            state ^= state >> 7;
            state ^= state << 17;
            let fval = (state as f64) / (u64::MAX as f64) - 0.5; // [-0.5, 0.5]
            mat[[i, j]] = F::from(fval * 2.0).expect("rand_f") * std_dev;
        }
    }
    mat
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use scirs2_core::ndarray::array;

    fn make_model() -> LatentODE<f64> {
        LatentODE::new(LatentODEConfig {
            obs_dim: 2,
            latent_dim: 4,
            encoder_hidden: 8,
            ode_hidden: 8,
            ode_steps_per_unit: 5,
            seed: 0,
        })
        .expect("model construction")
    }

    #[test]
    fn test_encode_shape() {
        let model = make_model();
        let times = vec![0.0_f64, 0.5, 1.0, 1.7];
        let obs: Vec<Array1<f64>> = times.iter().map(|_| array![1.0, 0.5]).collect();
        let z0 = model.encode(&times, &obs).expect("encode");
        assert_eq!(z0.len(), 4);
    }

    #[test]
    fn test_decode_shape() {
        let model = make_model();
        let z0 = Array1::zeros(4_usize);
        let future = vec![0.0_f64, 0.25, 0.5, 0.75, 1.0];
        let preds = model.decode(&z0, &future).expect("decode");
        assert_eq!(preds.len(), 5);
        for p in &preds {
            assert_eq!(p.len(), 2);
        }
    }

    #[test]
    fn test_predict_pipeline() {
        let model = make_model();
        let times = vec![0.0_f64, 1.0, 2.5];
        let obs: Vec<Array1<f64>> = times.iter().map(|_| array![0.3, -0.1]).collect();
        let future = vec![3.0_f64, 4.0, 5.0];
        let preds = model.predict(&times, &obs, &future).expect("predict");
        assert_eq!(preds.len(), 3);
    }

    #[test]
    fn test_dimension_mismatch_error() {
        let model = make_model();
        let times = vec![0.0_f64];
        // Wrong obs_dim (3 instead of 2)
        let obs = vec![array![1.0_f64, 2.0, 3.0]];
        let result = model.encode(&times, &obs);
        assert!(result.is_err());
    }
}