gspx 0.1.2

use std::collections::HashMap;

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

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

use super::shared::{
    SparseFactor, accumulate, apply_direct_term, check_scale, check_scales, check_signal_rows,
    combine_complex_matrix, combine_complex_vec, factorize_shifted_matrix, shifted_matrix,
    solve_multi_rhs,
};

/// Options for lowpass solves.
#[derive(Debug, Clone, Copy, Default)]
pub struct LowpassOptions {
    /// When true, invalidate and rebuild the shifted-factor cache before solving.
    pub refactor: bool,
}

impl LowpassOptions {
    /// Builds options with explicit factor refactor behavior.
    pub const fn with_refactor(refactor: bool) -> Self {
        Self { refactor }
    }
}

/// Convolver for fixed (static) graph topology.
pub struct StaticConvolver {
    laplacian: Laplacian,
    n_vertices: usize,
    shift_factors: HashMap<u64, SparseFactor>,
}

impl StaticConvolver {
    /// Creates a static convolver for a fixed Laplacian.
    ///
    /// # 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(),
            laplacian: l,
            shift_factors: HashMap::new(),
        })
    }

    /// Convolves a real-valued signal matrix with a VF kernel.
    ///
    /// # 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_shifted(q, b)?;
            accumulate(&mut w, &z, k.residues.row(idx));
        }
        Ok(w)
    }

    /// 1D wrapper around [`StaticConvolver::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 [`StaticConvolver::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)
    }

    /// Runs lowpass filtering with explicit cache-control options.
    ///
    /// # Errors
    /// Returns an error when dimensions/scales are invalid or solves fail.
    pub fn lowpass_with_options(
        &mut self,
        b: ArrayView2<f64>,
        scales: &[f64],
        options: LowpassOptions,
        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_with_options(b, s, options, 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.lowpass_one_with_options(b, scale, LowpassOptions::default(), order)
    }

    /// Applies one real low-pass solve with cache-control options.
    ///
    /// # Errors
    /// Returns an error when dimensions/scales are invalid or solves fail.
    pub fn lowpass_one_with_options(
        &mut self,
        b: ArrayView2<f64>,
        scale: f64,
        options: LowpassOptions,
        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;
        if options.refactor {
            self.shift_factors.remove(&q.to_bits());
        }
        let mut x = b.to_owned();
        for _ in 0..order {
            let mut next = self.solve_shifted(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 = b.to_owned();
        for _ in 0..order {
            let x2 = self.solve_shifted(q, x.view())?;
            let x1 = self.solve_shifted(q, x2.view())?;
            let mut next = x2;
            next.scaled_add(-q, &x1);
            next *= 4.0 * q;
            x = next;
        }
        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 x1 = self.solve_shifted(q, b)?;
        let mut y = b.to_owned();
        y.scaled_add(-q, &x1);
        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))
    }

    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 factor_for_shift(&mut self, q: f64) -> Result<&SparseFactor, GspError> {
        let key = q.to_bits();
        if !self.shift_factors.contains_key(&key) {
            let shifted = shifted_matrix(&self.laplacian, q);
            let factor = factorize_shifted_matrix(&shifted)?;
            self.shift_factors.insert(key, factor);
        }
        self.shift_factors
            .get(&key)
            .ok_or_else(|| GspError::Factorization("failed to cache factor".to_string()))
    }

    fn solve_shifted(&mut self, q: f64, b: ArrayView2<f64>) -> Result<Array2<f64>, GspError> {
        let factor = self.factor_for_shift(q)?;
        solve_multi_rhs(factor, b)
    }
}