gspx 0.1.2

use std::collections::HashMap;

use ndarray::{Array1, Array2, Array3, ArrayView1, ArrayView2, Axis, Zip};
use num_complex::Complex64;

use crate::{Laplacian, error::GspError, kernel::VfKernel};

use super::shared::{
    BranchUpdate, PoleState, accumulate, apply_direct_term, check_scale, check_scales,
    check_signal_rows, combine_complex_matrix, combine_complex_vec, dot_update,
    factorize_shifted_matrix, invert_dense, shifted_matrix, solve_multi_rhs, solve_vector,
    update_vector,
};

pub struct DynamicConvolver {
    base_laplacian: Laplacian,
    n_vertices: usize,
    branch_updates: Vec<BranchUpdate>,
    pole_states: HashMap<u64, PoleState>,
}

impl DynamicConvolver {
    /// Creates a dynamic convolver with branch-update support.
    ///
    /// # Errors
    /// Returns an error when `l` is not square.
    pub fn new(l: Laplacian) -> Result<Self, GspError> {
        if l.rows() != l.cols() {
            return Err(GspError::Dimensions(format!(
                "Laplacian must be square, got {}x{}",
                l.rows(),
                l.cols()
            )));
        }
        Ok(Self {
            n_vertices: l.rows(),
            base_laplacian: l,
            branch_updates: Vec::new(),
            pole_states: HashMap::new(),
        })
    }

    /// Creates a dynamic convolver and pre-initializes pole states.
    ///
    /// # Errors
    /// Returns an error for invalid poles or failed factorizations.
    pub fn with_poles(l: Laplacian, poles: &[f64]) -> Result<Self, GspError> {
        let mut conv = Self::new(l)?;
        for &q in poles {
            if !q.is_finite() || q <= 0.0 {
                return Err(GspError::InvalidKernel(format!(
                    "all dynamic poles must be finite and > 0, got {q}"
                )));
            }
            conv.ensure_pole(q)?;
        }
        Ok(conv)
    }

    /// Creates a dynamic convolver using poles from an existing kernel.
    ///
    /// # Errors
    /// Returns an error when kernel validation or pole setup fails.
    pub fn with_kernel(l: Laplacian, k: &VfKernel) -> Result<Self, GspError> {
        k.validate()?;
        let poles = k.poles.to_vec();
        Self::with_poles(l, &poles)
    }

    /// Convolves a real-valued signal matrix with a VF kernel on the updated graph.
    ///
    /// # Errors
    /// Returns an error when kernel validation or solves fail.
    pub fn convolve(&mut self, b: ArrayView2<f64>, k: &VfKernel) -> Result<Array3<f64>, GspError> {
        k.validate()?;
        self.check_signal_2d(b)?;
        let n_dim = k.residues.ncols();
        let mut w = Array3::<f64>::zeros((b.nrows(), b.ncols(), n_dim));
        apply_direct_term(&mut w, b, k.direct.view())?;

        for (idx, &q) in k.poles.iter().enumerate() {
            let z = self.solve_updated(q, b)?;
            accumulate(&mut w, &z, k.residues.row(idx));
        }
        Ok(w)
    }

    /// 1D wrapper around [`DynamicConvolver::convolve`].
    ///
    /// # Errors
    /// Returns an error when dimensions are invalid or solves fail.
    pub fn convolve_1d(
        &mut self,
        b: ArrayView1<f64>,
        k: &VfKernel,
    ) -> Result<Array2<f64>, GspError> {
        self.check_signal_1d(b)?;
        let b2 = b.insert_axis(Axis(1));
        let w = self.convolve(b2, k)?;
        Ok(w.index_axis_move(Axis(1), 0))
    }

    /// Convolves a complex-valued signal matrix by splitting real/imaginary parts.
    ///
    /// # Errors
    /// Returns an error when dimensions are invalid or solves fail.
    pub fn convolve_complex(
        &mut self,
        b: ArrayView2<Complex64>,
        k: &VfKernel,
    ) -> Result<Array3<Complex64>, GspError> {
        self.check_signal_2d_complex(b)?;
        let wr = self.convolve(b.mapv(|v| v.re).view(), k)?;
        let wi = self.convolve(b.mapv(|v| v.im).view(), k)?;
        let mut out = Array3::<Complex64>::zeros(wr.raw_dim());
        Zip::from(&mut out)
            .and(&wr)
            .and(&wi)
            .for_each(|o, &r, &i| *o = Complex64::new(r, i));
        Ok(out)
    }

    /// 1D wrapper around [`DynamicConvolver::convolve_complex`].
    ///
    /// # Errors
    /// Returns an error when dimensions are invalid or solves fail.
    pub fn convolve_complex_1d(
        &mut self,
        b: ArrayView1<Complex64>,
        k: &VfKernel,
    ) -> Result<Array2<Complex64>, GspError> {
        let b2 = b.insert_axis(Axis(1));
        let w = self.convolve_complex(b2, k)?;
        Ok(w.index_axis_move(Axis(1), 0))
    }

    /// Applies the analytical low-pass filter over multiple scales.
    ///
    /// # Errors
    /// Returns an error when dimensions/scales are invalid or solves fail.
    pub fn lowpass(
        &mut self,
        b: ArrayView2<f64>,
        scales: &[f64],
        order: usize,
    ) -> Result<Vec<Array2<f64>>, GspError> {
        self.check_signal_2d(b)?;
        check_scales(scales)?;
        let mut out = Vec::with_capacity(scales.len());
        for &s in scales {
            out.push(self.lowpass_one(b, s, order)?);
        }
        Ok(out)
    }

    /// Applies analytical low-pass filtering to complex signals.
    ///
    /// # Errors
    /// Returns an error when dimensions/scales are invalid or solves fail.
    pub fn lowpass_complex(
        &mut self,
        b: ArrayView2<Complex64>,
        scales: &[f64],
        order: usize,
    ) -> Result<Vec<Array2<Complex64>>, GspError> {
        self.check_signal_2d_complex(b)?;
        let yr = self.lowpass(b.mapv(|z| z.re).view(), scales, order)?;
        let yi = self.lowpass(b.mapv(|z| z.im).view(), scales, order)?;
        combine_complex_vec(yr, yi)
    }

    /// Applies one complex low-pass solve at a single scale.
    ///
    /// # Errors
    /// Returns an error when dimensions/scales are invalid or solves fail.
    pub fn lowpass_complex_one(
        &mut self,
        b: ArrayView2<Complex64>,
        scale: f64,
        order: usize,
    ) -> Result<Array2<Complex64>, GspError> {
        self.check_signal_2d_complex(b)?;
        let yr = self.lowpass_one(b.mapv(|z| z.re).view(), scale, order)?;
        let yi = self.lowpass_one(b.mapv(|z| z.im).view(), scale, order)?;
        Ok(combine_complex_matrix(yr, yi))
    }

    /// Applies one real low-pass solve at a single scale.
    ///
    /// # Errors
    /// Returns an error when dimensions/scales are invalid or solves fail.
    pub fn lowpass_one(
        &mut self,
        b: ArrayView2<f64>,
        scale: f64,
        order: usize,
    ) -> Result<Array2<f64>, GspError> {
        self.check_signal_2d(b)?;
        check_scale(scale)?;
        if order == 0 {
            return Ok(b.to_owned());
        }
        let q = 1.0 / scale;
        let mut x = self.solve_updated(q, b)?;
        x *= q;
        for _ in 1..order {
            let mut next = self.solve_updated(q, x.view())?;
            next *= q;
            x = next;
        }
        Ok(x)
    }

    /// Applies the analytical band-pass filter over multiple scales.
    ///
    /// # Errors
    /// Returns an error when dimensions/scales are invalid or solves fail.
    pub fn bandpass(
        &mut self,
        b: ArrayView2<f64>,
        scales: &[f64],
        order: usize,
    ) -> Result<Vec<Array2<f64>>, GspError> {
        self.check_signal_2d(b)?;
        check_scales(scales)?;
        let mut out = Vec::with_capacity(scales.len());
        for &s in scales {
            out.push(self.bandpass_one(b, s, order)?);
        }
        Ok(out)
    }

    /// Applies one real band-pass solve at a single scale.
    ///
    /// # Errors
    /// Returns an error when dimensions/scales are invalid or solves fail.
    pub fn bandpass_one(
        &mut self,
        b: ArrayView2<f64>,
        scale: f64,
        order: usize,
    ) -> Result<Array2<f64>, GspError> {
        self.check_signal_2d(b)?;
        check_scale(scale)?;
        if order == 0 {
            return Ok(b.to_owned());
        }
        let q = 1.0 / scale;
        let mut x = self.bandpass_step(q, b)?;
        for _ in 1..order {
            x = self.bandpass_step(q, x.view())?;
        }
        Ok(x)
    }

    /// Applies analytical band-pass filtering to complex signals.
    ///
    /// # Errors
    /// Returns an error when dimensions/scales are invalid or solves fail.
    pub fn bandpass_complex(
        &mut self,
        b: ArrayView2<Complex64>,
        scales: &[f64],
        order: usize,
    ) -> Result<Vec<Array2<Complex64>>, GspError> {
        self.check_signal_2d_complex(b)?;
        let yr = self.bandpass(b.mapv(|z| z.re).view(), scales, order)?;
        let yi = self.bandpass(b.mapv(|z| z.im).view(), scales, order)?;
        combine_complex_vec(yr, yi)
    }

    /// Applies one complex band-pass solve at a single scale.
    ///
    /// # Errors
    /// Returns an error when dimensions/scales are invalid or solves fail.
    pub fn bandpass_complex_one(
        &mut self,
        b: ArrayView2<Complex64>,
        scale: f64,
        order: usize,
    ) -> Result<Array2<Complex64>, GspError> {
        self.check_signal_2d_complex(b)?;
        let yr = self.bandpass_one(b.mapv(|z| z.re).view(), scale, order)?;
        let yi = self.bandpass_one(b.mapv(|z| z.im).view(), scale, order)?;
        Ok(combine_complex_matrix(yr, yi))
    }

    /// Applies the analytical high-pass filter over multiple scales.
    ///
    /// # Errors
    /// Returns an error when dimensions/scales are invalid or solves fail.
    pub fn highpass(
        &mut self,
        b: ArrayView2<f64>,
        scales: &[f64],
    ) -> Result<Vec<Array2<f64>>, GspError> {
        self.check_signal_2d(b)?;
        check_scales(scales)?;
        let mut out = Vec::with_capacity(scales.len());
        for &s in scales {
            out.push(self.highpass_one(b, s)?);
        }
        Ok(out)
    }

    /// Applies one real high-pass solve at a single scale.
    ///
    /// # Errors
    /// Returns an error when dimensions/scales are invalid or solves fail.
    pub fn highpass_one(
        &mut self,
        b: ArrayView2<f64>,
        scale: f64,
    ) -> Result<Array2<f64>, GspError> {
        self.check_signal_2d(b)?;
        check_scale(scale)?;
        let q = 1.0 / scale;
        let mut y = self.solve_updated(q, b)?;
        y *= -q;
        y.scaled_add(1.0, &b);
        Ok(y)
    }

    /// Applies analytical high-pass filtering to complex signals.
    ///
    /// # Errors
    /// Returns an error when dimensions/scales are invalid or solves fail.
    pub fn highpass_complex(
        &mut self,
        b: ArrayView2<Complex64>,
        scales: &[f64],
    ) -> Result<Vec<Array2<Complex64>>, GspError> {
        self.check_signal_2d_complex(b)?;
        let yr = self.highpass(b.mapv(|z| z.re).view(), scales)?;
        let yi = self.highpass(b.mapv(|z| z.im).view(), scales)?;
        combine_complex_vec(yr, yi)
    }

    /// Applies one complex high-pass solve at a single scale.
    ///
    /// # Errors
    /// Returns an error when dimensions/scales are invalid or solves fail.
    pub fn highpass_complex_one(
        &mut self,
        b: ArrayView2<Complex64>,
        scale: f64,
    ) -> Result<Array2<Complex64>, GspError> {
        self.check_signal_2d_complex(b)?;
        let yr = self.highpass_one(b.mapv(|z| z.re).view(), scale)?;
        let yi = self.highpass_one(b.mapv(|z| z.im).view(), scale)?;
        Ok(combine_complex_matrix(yr, yi))
    }

    /// Adds a weighted branch `(i, j)` to the dynamic topology.
    ///
    /// # Errors
    /// Returns an error on negative weights or near-singular Woodbury updates.
    pub fn add_branch(&mut self, i: usize, j: usize, w: f64) -> Result<bool, GspError> {
        if i >= self.n_vertices || j >= self.n_vertices {
            return Ok(false);
        }
        if w < 0.0 {
            return Err(GspError::InvalidScales(
                "math domain error: weight must be non-negative".to_string(),
            ));
        }
        let update = BranchUpdate {
            i,
            j,
            sqrt_w: w.sqrt(),
        };
        let previous_updates = &self.branch_updates;
        let mut staged_pole_states = Vec::with_capacity(self.pole_states.len());
        for (&key, state) in &self.pole_states {
            let (solved_update, gram_inv) = Self::stage_update_for_pole_state(
                self.n_vertices,
                previous_updates,
                state,
                &update,
            )?;
            staged_pole_states.push((key, solved_update, gram_inv));
        }
        self.branch_updates.push(update);
        for (key, solved_update, gram_inv) in staged_pole_states {
            let state = self.pole_states.get_mut(&key).ok_or_else(|| {
                GspError::Factorization("pole cache missing during update commit".to_string())
            })?;
            state.solved_cols.push(solved_update);
            state.gram_inv = gram_inv;
        }
        Ok(true)
    }

    fn solve_updated(&mut self, q: f64, b: ArrayView2<f64>) -> Result<Array2<f64>, GspError> {
        self.ensure_pole(q)?;
        let state = self.pole_states.get(&q.to_bits()).ok_or_else(|| {
            GspError::Factorization("pole cache missing after ensure".to_string())
        })?;
        let mut out = solve_multi_rhs(&state.factor, b)?;
        if self.branch_updates.is_empty() {
            return Ok(out);
        }

        let n_updates = self.branch_updates.len();
        let mut t = vec![0.0; n_updates];
        let mut gamma = vec![0.0; n_updates];
        for col in 0..out.ncols() {
            for (idx, upd) in self.branch_updates.iter().enumerate() {
                t[idx] = upd.sqrt_w * (out[[upd.i, col]] - out[[upd.j, col]]);
            }
            for (row, gamma_row) in gamma.iter_mut().enumerate() {
                let mut value = 0.0;
                for (k, &t_k) in t.iter().enumerate() {
                    value += state.gram_inv[[row, k]] * t_k;
                }
                *gamma_row = value;
            }
            let mut out_col = out.column_mut(col);
            for (idx, u) in state.solved_cols.iter().enumerate() {
                let coeff = gamma[idx];
                if coeff != 0.0 {
                    out_col.scaled_add(-coeff, &u.view());
                }
            }
        }
        Ok(out)
    }

    fn bandpass_step(&mut self, q: f64, x: ArrayView2<f64>) -> Result<Array2<f64>, GspError> {
        let x2 = self.solve_updated(q, x)?;
        let x1 = self.solve_updated(q, x2.view())?;
        let mut next = x2;
        next.scaled_add(-q, &x1);
        next *= 4.0 * q;
        Ok(next)
    }

    fn ensure_pole(&mut self, q: f64) -> Result<(), GspError> {
        let key = q.to_bits();
        if self.pole_states.contains_key(&key) {
            return Ok(());
        }
        let shifted = shifted_matrix(&self.base_laplacian, q);
        let factor = factorize_shifted_matrix(&shifted)?;
        let mut state = PoleState {
            factor,
            solved_cols: Vec::new(),
            gram_inv: Array2::zeros((0, 0)),
        };
        if !self.branch_updates.is_empty() {
            for upd in &self.branch_updates {
                let rhs = update_vector(self.n_vertices, upd);
                let u = solve_vector(&state.factor, rhs.view())?;
                state.solved_cols.push(u);
            }
            let n_updates = self.branch_updates.len();
            let mut gram = Array2::<f64>::zeros((n_updates, n_updates));
            for i in 0..n_updates {
                gram[[i, i]] = 1.0;
            }
            for i in 0..n_updates {
                for j in 0..n_updates {
                    gram[[i, j]] +=
                        dot_update(&self.branch_updates[i], state.solved_cols[j].view());
                }
            }
            state.gram_inv = invert_dense(gram)?;
        }
        self.pole_states.insert(key, state);
        Ok(())
    }

    fn check_signal_1d(&self, b: ArrayView1<f64>) -> Result<(), GspError> {
        if b.len() != self.n_vertices {
            return Err(GspError::Dimensions(format!(
                "signal length {} does not match graph size {}",
                b.len(),
                self.n_vertices
            )));
        }
        Ok(())
    }

    fn check_signal_2d(&self, b: ArrayView2<f64>) -> Result<(), GspError> {
        check_signal_rows(b.nrows(), self.n_vertices)
    }

    fn check_signal_2d_complex(&self, b: ArrayView2<Complex64>) -> Result<(), GspError> {
        check_signal_rows(b.nrows(), self.n_vertices)
    }

    fn stage_update_for_pole_state(
        n_vertices: usize,
        previous_updates: &[BranchUpdate],
        state: &PoleState,
        update: &BranchUpdate,
    ) -> Result<(Array1<f64>, Array2<f64>), GspError> {
        let rhs = update_vector(n_vertices, update);
        let solved_update = solve_vector(&state.factor, rhs.view())?;
        if previous_updates.is_empty() {
            let alpha = 1.0 + dot_update(update, solved_update.view());
            if alpha.abs() < 1e-12 {
                return Err(GspError::SingularSolve(
                    "Woodbury update produced near-singular scalar block".to_string(),
                ));
            }
            return Ok((solved_update, Array2::from_elem((1, 1), 1.0 / alpha)));
        }

        let coupling = Array1::from(
            previous_updates
                .iter()
                .map(|previous_update| dot_update(previous_update, solved_update.view()))
                .collect::<Vec<_>>(),
        );
        let old_inverse = &state.gram_inv;
        let schur_rhs = old_inverse.dot(&coupling);
        let alpha = 1.0 + dot_update(update, solved_update.view());
        let schur_complement = alpha - coupling.dot(&schur_rhs);
        if schur_complement.abs() < 1e-12 {
            return Err(GspError::SingularSolve(
                "Woodbury update produced near-singular Schur complement".to_string(),
            ));
        }

        let previous_count = previous_updates.len();
        let mut expanded_inverse = Array2::<f64>::zeros((previous_count + 1, previous_count + 1));
        for row in 0..previous_count {
            for col in 0..previous_count {
                expanded_inverse[[row, col]] =
                    old_inverse[[row, col]] + schur_rhs[row] * schur_rhs[col] / schur_complement;
            }
        }
        for row in 0..previous_count {
            expanded_inverse[[row, previous_count]] = -schur_rhs[row] / schur_complement;
            expanded_inverse[[previous_count, row]] = -schur_rhs[row] / schur_complement;
        }
        expanded_inverse[[previous_count, previous_count]] = 1.0 / schur_complement;
        Ok((solved_update, expanded_inverse))
    }
}