//! Single-token recurrent step for Mamba SSM.
//!
//! `MambaRecurrence` extracts all CPU-side weight arrays from a trained
//! [`MambaBlock`] once (via [`MambaRecurrence::from_block`]) and then exposes
//! a pure-Rust, allocation-light [`MambaRecurrence::step`] that advances the
//! hidden state one token at a time — the same S6 selective-scan arithmetic
//! used in [`MambaBlock::forward`], but without the batch / sequence loop.

use super::mamba::MambaBlock;
use crate::error::TransformerError;

/// Convenience alias so callers inside rustytorch can write `ThinkError`.
pub type ThinkError = TransformerError;

// ─── state ───────────────────────────────────────────────────────────────────

/// Recurrent hidden state for a single-channel Mamba SSM.
///
/// `h` carries the SSM hidden state `[d_inner * d_state]`; `conv_buf` carries
/// the last `d_conv − 1` x_in values per channel `[d_inner * (d_conv − 1)]`
/// so the causal conv1d can be computed correctly in a single-step setting.
#[derive(Debug, Clone, PartialEq)]
pub struct MambaState {
    /// Flattened SSM hidden state `[d_inner * d_state]`.
    pub h: Vec<f32>,
    /// Causal-conv history buffer `[d_inner * (d_conv - 1)]`, oldest-first per channel.
    pub conv_buf: Vec<f32>,
}

impl MambaState {
    /// Create a zero-initialised state for the given dimensions.
    #[must_use]
    pub fn zeros(d_inner: usize, d_state: usize, d_conv: usize) -> Self {
        Self {
            h: vec![0.0f32; d_inner * d_state],
            conv_buf: vec![0.0f32; d_inner * d_conv.saturating_sub(1)],
        }
    }

    /// The SSM hidden-state slice `[d_inner * d_state]`.
    #[must_use]
    pub fn hidden(&self) -> &[f32] {
        &self.h
    }
}

// ─── recurrence cell ─────────────────────────────────────────────────────────

/// CPU weight snapshot of a [`MambaBlock`] for single-step recurrent inference.
///
/// All weights are extracted once from the block's `Tensor` fields as flat
/// `Vec<f32>` and kept on the CPU.  After construction the block is no longer
/// needed.
pub struct MambaRecurrence {
    // ── dims ──────────────────────────────────────────────────────────────
    /// Model input dimension.
    pub d_model: usize,
    /// Inner (expanded) dimension `= d_model * expand`.
    pub d_inner: usize,
    /// SSM state dimension.
    pub d_state: usize,
    /// Causal conv1d kernel width.
    pub d_conv: usize,
    /// Rank of the Δ (dt) projection.
    pub dt_rank: usize,

    // ── weights (flat, row-major) ──────────────────────────────────────────
    /// `[d_model, 2*d_inner]` — input projection (SSM branch ++ gate branch).
    in_proj: Vec<f32>,
    /// `[d_inner, d_conv]` — depthwise conv1d weights (d_inner channels, kernel d_conv).
    conv1d_weight: Vec<f32>,
    /// `[d_inner]` — conv1d bias (may be all-zero when `conv_bias=false`).
    conv1d_bias: Vec<f32>,
    /// `[d_inner, dt_rank + 2*d_state]` — x_proj.
    x_proj: Vec<f32>,
    /// `[dt_rank, d_inner]` — dt_proj.
    dt_proj: Vec<f32>,
    /// `[d_inner]` — dt_bias (added before softplus).
    dt_bias: Vec<f32>,
    /// `[d_inner, d_state]` — A_log (log of positive eigenvalues; A = -exp(A_log)).
    a_log: Vec<f32>,
    /// `[d_inner]` — D skip coefficient.
    d_skip: Vec<f32>,
    /// `[d_inner, d_model]` — output projection.
    out_proj: Vec<f32>,
}

impl MambaRecurrence {
    /// Return the model input dimension.
    #[must_use]
    pub fn d_model(&self) -> usize {
        self.d_model
    }

    /// Return `d_inner * d_state` — the length of the hidden-state vector.
    #[must_use]
    pub fn state_size(&self) -> usize {
        self.d_inner * self.d_state
    }

    /// Create a zero-initialised [`MambaState`] compatible with this recurrence cell.
    #[must_use]
    pub fn init_state(&self) -> MambaState {
        MambaState::zeros(self.d_inner, self.d_state, self.d_conv)
    }

    /// Extract all weights from `block` into CPU `Vec<f32>` buffers.
    ///
    /// # Errors
    /// Returns `TransformerError` if any tensor cannot be read as f32.
    pub fn from_block(block: &MambaBlock) -> Result<Self, TransformerError> {
        // Pull every tensor via the public persistence API so we don't rely on
        // private fields.
        let pts = block.persistence_tensors();
        let get = |name: &str| -> Result<Vec<f32>, TransformerError> {
            let t = pts
                .iter()
                .find(|(n, _)| *n == name)
                .map(|(_, t)| *t)
                .ok_or_else(|| {
                    TransformerError::Generic(format!(
                        "MambaRecurrence::from_block: missing tensor `{name}`"
                    ))
                })?;
            t.to_vec().map_err(|e| {
                TransformerError::Generic(format!(
                    "MambaRecurrence::from_block: to_vec failed for `{name}`: {e}"
                ))
            })
        };

        let in_proj = get("in_proj")?; // [d_model, 2*d_inner]
        let conv1d_weight_raw = get("conv1d_weight")?; // [d_inner, 1, d_conv]
        let conv1d_bias = get("conv1d_bias").unwrap_or_else(|_| Vec::new());
        let a_log = get("A_log")?; // [d_inner, d_state]
        let x_proj = get("x_proj")?; // [d_inner, dt_rank + 2*d_state]
        let dt_proj = get("dt_proj")?; // [dt_rank, d_inner]
        let dt_bias = get("dt_bias")?; // [d_inner]
        let d_skip = get("D")?; // [d_inner]
        let out_proj = get("out_proj")?; // [d_inner, d_model]

        // Derive dims from tensor sizes.
        let d_skip_len = d_skip.len(); // d_inner
        let d_inner = d_skip_len;
        let a_log_len = a_log.len(); // d_inner * d_state
        let d_state = a_log_len / d_inner;
        let d_model = in_proj.len() / (2 * d_inner); // in_proj: [d_model, 2*d_inner]
        let dt_proj_len = dt_proj.len(); // dt_rank * d_inner
        let dt_rank = dt_proj_len / d_inner;
        let conv1d_weight_len = conv1d_weight_raw.len(); // d_inner * 1 * d_conv
        let d_conv = conv1d_weight_len / d_inner;

        // conv1d is stored as [d_inner, 1, d_conv]; flatten to [d_inner, d_conv].
        let conv1d_weight = conv1d_weight_raw;

        // Ensure conv1d_bias is exactly d_inner long (pad with zeros if absent).
        let conv1d_bias = if conv1d_bias.len() == d_inner {
            conv1d_bias
        } else {
            vec![0.0f32; d_inner]
        };

        Ok(Self {
            d_model,
            d_inner,
            d_state,
            d_conv,
            dt_rank,
            in_proj,
            conv1d_weight,
            conv1d_bias,
            x_proj,
            dt_proj,
            dt_bias,
            a_log,
            d_skip,
            out_proj,
        })
    }

    /// Perform one recurrent step of the S6 selective scan.
    ///
    /// `x` is the current input token `[d_model]`.  `state` is updated in
    /// place; initialise with [`MambaRecurrence::init_state`] for a fresh
    /// zero-state run.
    ///
    /// Returns the output token `[d_model]`.
    ///
    /// # Panics
    /// Panics if `x.len() != d_model` or `state.h.len() != d_inner * d_state`
    /// (programmer error — callers should always pass matching dimensions).
    pub fn step(&self, x: &[f32], state: &mut MambaState) -> Vec<f32> {
        let d = self.d_inner;
        let n = self.d_state;
        let dm = self.d_model;
        let kc = self.d_conv;
        let dt_rank = self.dt_rank;

        assert_eq!(
            x.len(),
            dm,
            "MambaRecurrence::step: x.len()={} != d_model={dm}",
            x.len()
        );
        assert_eq!(
            state.h.len(),
            d * n,
            "MambaRecurrence::step: state.h.len()={} != d_inner*d_state={}",
            state.h.len(),
            d * n
        );

        // ── 1. in_proj → x_in (SSM branch) + z (gate branch) ──────────────
        // in_proj: [d_model, 2*d_inner]
        let mut x_in = vec![0.0f32; d]; // inner activation
        let mut z = vec![0.0f32; d]; // gate
        for j in 0..d {
            let mut sx = 0.0f32;
            let mut sz = 0.0f32;
            for m in 0..dm {
                let w = x[m];
                sx += w * self.in_proj[m * (2 * d) + j];
                sz += w * self.in_proj[m * (2 * d) + d + j];
            }
            x_in[j] = sx;
            z[j] = sz;
        }

        // ── 2. causal depthwise conv1d (with history buffer) → SiLU → u ──
        // conv_buf holds the last (kc-1) x_in values per channel, oldest-first.
        // conv1d_weight layout: [d_inner, 1, d_conv] = [d_inner * d_conv] flat.
        let cb = kc.saturating_sub(1); // history buffer length per channel
        let mut u = vec![0.0f32; d];
        for j in 0..d {
            let mut acc = self.conv1d_bias[j];
            for kk in 0..cb {
                acc += self.conv1d_weight[j * kc + kk] * state.conv_buf[j * cb + kk];
            }
            acc += self.conv1d_weight[j * kc + (kc - 1)] * x_in[j];
            u[j] = silu_f32(acc);
        }
        // Shift buffer left and append current x_in as newest entry.
        for j in 0..d {
            let base = j * cb;
            for kk in 0..cb.saturating_sub(1) {
                state.conv_buf[base + kk] = state.conv_buf[base + kk + 1];
            }
            if cb > 0 {
                state.conv_buf[base + cb - 1] = x_in[j];
            }
        }

        // ── 3. x_proj → (dt_raw [dt_rank], B [d_state], C [d_state]) ───────
        let dbc = dt_rank + 2 * n;
        let mut xdbl = vec![0.0f32; dbc];
        for q in 0..dbc {
            let mut s = 0.0f32;
            for j in 0..d {
                s += u[j] * self.x_proj[j * dbc + q];
            }
            xdbl[q] = s;
        }
        let b_vec = &xdbl[dt_rank..dt_rank + n];
        let c_vec = &xdbl[dt_rank + n..dt_rank + 2 * n];

        // ── 4. discretize Δ = softplus(dt_raw @ dt_proj^T + dt_bias) ────────
        // dt_proj: [dt_rank, d_inner]
        let mut delta = vec![0.0f32; d];
        for j in 0..d {
            let mut s = self.dt_bias[j];
            for r in 0..dt_rank {
                s += xdbl[r] * self.dt_proj[r * d + j];
            }
            delta[j] = softplus_f32(s);
        }

        // ── 5. scan: h_new = dA * h + dB*u ;  y = C @ h_new + D*u ──────────
        let mut h_new = vec![0.0f32; d * n];
        let mut y_raw = vec![0.0f32; d];
        for j in 0..d {
            let dj = delta[j];
            let uj = u[j];
            let mut yj = self.d_skip[j] * uj;
            for nn in 0..n {
                let a = -self.a_log[j * n + nn].exp();
                let da = (dj * a).exp();
                let dbu = dj * b_vec[nn] * uj;
                let hv = da * state.h[j * n + nn] + dbu;
                h_new[j * n + nn] = hv;
                yj += c_vec[nn] * hv;
            }
            y_raw[j] = yj;
        }
        state.h = h_new;

        // ── 6. gate: y_gated = y_raw * silu(z) ──────────────────────────────
        let mut y_gated = vec![0.0f32; d];
        for j in 0..d {
            y_gated[j] = y_raw[j] * silu_f32(z[j]);
        }

        // ── 7. out_proj: output = y_gated @ out_proj^T ───────────────────────
        // out_proj: [d_inner, d_model]
        let mut output = vec![0.0f32; dm];
        for j in 0..d {
            let yg = y_gated[j];
            for m in 0..dm {
                output[m] += yg * self.out_proj[j * dm + m];
            }
        }

        output
    }
}

// ─── helpers ─────────────────────────────────────────────────────────────────

#[inline]
fn silu_f32(x: f32) -> f32 {
    let s = if x >= 0.0 {
        1.0 / (1.0 + (-x).exp())
    } else {
        let e = x.exp();
        e / (1.0 + e)
    };
    x * s
}

#[inline]
fn softplus_f32(x: f32) -> f32 {
    x.max(0.0) + (1.0 + (-x.abs()).exp()).ln()
}