gspx 0.1.2

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

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

/// Configuration for vector fitting.
#[derive(Debug, Clone)]
pub struct VfFitOptions {
    /// Optional initial poles. If `None`, poles are initialized automatically.
    pub poles: Option<Array1<Complex64>>,
    /// Maximum relocation iterations.
    pub iters: usize,
    /// Relative convergence tolerance on pole movement.
    pub tol: f64,
    /// Fit the direct term `d`.
    pub fit_d: bool,
    /// Fit the proportional term `e`.
    pub fit_e: bool,
    /// Restrict pole initialization/relocation to the real axis.
    pub real_only: bool,
}

impl Default for VfFitOptions {
    fn default() -> Self {
        Self {
            poles: None,
            iters: 30,
            tol: 1e-9,
            fit_d: true,
            fit_e: false,
            real_only: false,
        }
    }
}

/// Vector-fitting model result.
#[derive(Debug, Clone)]
pub struct VfModel {
    /// Fitted poles.
    pub poles: Array1<Complex64>,
    /// Residues with shape `(n_poles, n_channels)`.
    pub residues: Array2<Complex64>,
    /// Direct term.
    pub d: Array1<Complex64>,
    /// Proportional term.
    pub e: Array1<Complex64>,
    /// Root-mean-square error on the fitting samples.
    pub rmse: f64,
    /// Number of iterations used.
    pub iters: usize,
}

impl VfModel {
    /// Evaluates the fitted rational model on complex samples `s`.
    pub fn evaluate(&self, s: ArrayView1<Complex64>) -> Array2<Complex64> {
        let c = cauchy(s, self.poles.view());
        let mut out = c.dot(&self.residues);
        for mut row in out.rows_mut() {
            row += &self.d;
        }
        for (i, mut row) in out.rows_mut().into_iter().enumerate() {
            let si = s[i];
            for j in 0..row.len() {
                row[j] += si * self.e[j];
            }
        }
        out
    }

    /// Fits a rational model to sampled frequency response data.
    ///
    /// # Errors
    /// Returns an error when fitting fails or dimensions are inconsistent.
    pub fn fit(
        s: ArrayView1<Complex64>,
        f: ArrayView2<Complex64>,
        n_poles: usize,
        opts: VfFitOptions,
    ) -> Result<Self, GspError> {
        if n_poles == 0 {
            return Err(GspError::InvalidKernel(
                "n_poles must be at least 1".to_string(),
            ));
        }
        let s = s.to_owned();
        let mut fmat = f.to_owned();
        if fmat.nrows() != s.len() {
            if fmat.ncols() == s.len() {
                fmat = fmat.t().to_owned();
            } else {
                return Err(GspError::Dimensions(format!(
                    "response shape {:?} is incompatible with sample length {}",
                    fmat.raw_dim(),
                    s.len()
                )));
            }
        }
        let mut poles = if let Some(p) = opts.poles.clone() {
            p
        } else {
            initial_poles(s.view(), n_poles, opts.real_only)
        };

        let mut it_used = 0usize;
        for it in 0..opts.iters {
            let prev = poles.clone();
            poles = relocate(
                s.view(),
                fmat.view(),
                poles.view(),
                opts.fit_d,
                opts.fit_e,
                opts.real_only,
            )?;
            it_used = it + 1;
            let rel = poles
                .iter()
                .zip(prev.iter())
                .map(|(a, b)| (*a - *b).norm() / (b.norm() + 1e-15))
                .fold(0.0f64, |acc, v| acc.max(v));
            if rel < opts.tol {
                break;
            }
        }

        let x = complex_lstsq(
            basis(s.view(), poles.view(), opts.fit_d, opts.fit_e).view(),
            fmat.view(),
        )?;
        let n = poles.len();
        let ch = fmat.ncols();
        let d = if opts.fit_d {
            x.row(n).to_owned()
        } else {
            Array1::from_elem(ch, Complex64::new(0.0, 0.0))
        };
        let e = if opts.fit_e {
            x.row(n + (opts.fit_d as usize)).to_owned()
        } else {
            Array1::from_elem(ch, Complex64::new(0.0, 0.0))
        };
        let residues = x.slice(s![0..n, ..]).to_owned();

        let mut model = Self {
            poles,
            residues,
            d,
            e,
            rmse: 0.0,
            iters: it_used,
        };
        let y = model.evaluate(s.view());
        let rmse = (&fmat - &y)
            .mapv(|v| v.norm_sqr())
            .mean()
            .ok_or_else(|| GspError::Dimensions("empty fit matrix".to_string()))?
            .sqrt();
        model.rmse = rmse;
        Ok(model)
    }

    /// Fits a real-valued [`VfKernel`] over Laplacian eigenvalue samples.
    ///
    /// # Errors
    /// Returns an error when fitting fails.
    pub fn kernel(
        lam: ArrayView1<f64>,
        g: ArrayView2<f64>,
        n_poles: usize,
        mut opts: VfFitOptions,
    ) -> Result<VfKernel, GspError> {
        opts.real_only = true;
        opts.fit_e = false;
        let s = lam.mapv(|v| Complex64::new(v, 0.0));
        let g_complex = g.mapv(|v| Complex64::new(v, 0.0));
        let model = Self::fit(s.view(), g_complex.view(), n_poles, opts)?;
        Ok(VfKernel {
            residues: model.residues.mapv(|v| v.re),
            poles: model.poles.mapv(|v| v.re),
            direct: model.d.mapv(|v| v.re),
        })
    }
}

fn cauchy(s: ArrayView1<Complex64>, poles: ArrayView1<Complex64>) -> Array2<Complex64> {
    let mut out = Array2::<Complex64>::zeros((s.len(), poles.len()));
    for i in 0..s.len() {
        for j in 0..poles.len() {
            out[[i, j]] = Complex64::new(1.0, 0.0) / (s[i] - poles[j]);
        }
    }
    out
}

fn basis(
    s: ArrayView1<Complex64>,
    poles: ArrayView1<Complex64>,
    fit_d: bool,
    fit_e: bool,
) -> Array2<Complex64> {
    let c = cauchy(s, poles);
    let mut n_cols = poles.len();
    if fit_d {
        n_cols += 1;
    }
    if fit_e {
        n_cols += 1;
    }
    let mut out = Array2::<Complex64>::zeros((s.len(), n_cols));
    out.slice_mut(s![.., 0..poles.len()]).assign(&c);
    let mut col = poles.len();
    if fit_d {
        out.column_mut(col).fill(Complex64::new(1.0, 0.0));
        col += 1;
    }
    if fit_e {
        out.column_mut(col).assign(&s);
    }
    out
}

fn relocate(
    s: ArrayView1<Complex64>,
    f: ArrayView2<Complex64>,
    poles: ArrayView1<Complex64>,
    fit_d: bool,
    fit_e: bool,
    real_only: bool,
) -> Result<Array1<Complex64>, GspError> {
    let k = f.nrows();
    let a_ch = f.ncols();
    let n = poles.len();

    let phi_sig = basis(s, poles, true, false);
    let phi = basis(s, poles, fit_d, fit_e);

    let mut flat = Array2::<Complex64>::zeros((k, a_ch * (n + 1)));
    for i in 0..k {
        for a in 0..a_ch {
            for j in 0..(n + 1) {
                flat[[i, a * (n + 1) + j]] = f[[i, a]] * phi_sig[[i, j]];
            }
        }
    }

    let proj_coef = complex_lstsq(phi.view(), flat.view())?;
    let proj = phi.dot(&proj_coef);
    let residual = &flat - &proj;

    let mut rows = Array2::<Complex64>::zeros((k * a_ch, n + 1));
    for i in 0..k {
        for a in 0..a_ch {
            let ridx = i * a_ch + a;
            for j in 0..(n + 1) {
                rows[[ridx, j]] = residual[[i, a * (n + 1) + j]];
            }
        }
    }

    let w = ((k * a_ch) as f64).sqrt();
    let mut m = Array2::<Complex64>::zeros((rows.nrows() + 1, rows.ncols()));
    m.slice_mut(s![0..rows.nrows(), ..]).assign(&rows);
    let mut sum_phi = phi_sig.sum_axis(Axis(0));
    sum_phi.mapv_inplace(|v| v * w);
    m.row_mut(rows.nrows()).assign(&sum_phi);

    let mut rhs = Array2::<Complex64>::zeros((m.nrows(), 1));
    rhs[[m.nrows() - 1, 0]] = Complex64::new((k as f64) * w, 0.0);
    let z = complex_lstsq(m.view(), rhs.view())?;
    let c_tilde = z.slice(s![0..n, 0]).to_owned();
    let d_tilde = z[[n, 0]] + Complex64::new(1e-30, 0.0);

    let mut h = Array2::<Complex64>::zeros((n, n));
    for i in 0..n {
        h[[i, i]] = poles[i];
    }
    for i in 0..n {
        for j in 0..n {
            h[[i, j]] -= c_tilde[j] / d_tilde;
        }
    }
    let mut new_poles = complex_eigvals(h.view())?;
    new_poles.mapv_inplace(|p| Complex64::new(-p.re.abs(), p.im));
    if real_only {
        let mut vals = new_poles
            .iter()
            .map(|p| Complex64::new(p.re, 0.0))
            .collect::<Vec<_>>();
        vals.sort_by(|a, b| a.re.partial_cmp(&b.re).unwrap_or(std::cmp::Ordering::Equal));
        return Ok(Array1::from(vals));
    }
    Ok(enforce_conjugates(new_poles.view(), n))
}

fn enforce_conjugates(poles: ArrayView1<Complex64>, n: usize) -> Array1<Complex64> {
    let maxabs = poles.iter().map(|p| p.norm()).fold(0.0f64, |a, b| a.max(b));
    let tol = 1e-6 * (maxabs + 1.0);
    let mut real = poles
        .iter()
        .filter(|p| p.im.abs() < tol)
        .map(|p| Complex64::new(p.re, 0.0))
        .collect::<Vec<_>>();
    let upper = poles
        .iter()
        .filter(|p| p.im >= tol)
        .copied()
        .collect::<Vec<_>>();
    real.extend(upper.iter().copied());
    real.extend(upper.iter().map(|p| p.conj()));
    real.sort_by(|a, b| {
        a.re.partial_cmp(&b.re)
            .unwrap_or(std::cmp::Ordering::Equal)
            .then(a.im.partial_cmp(&b.im).unwrap_or(std::cmp::Ordering::Equal))
    });
    if real.len() > n {
        real.truncate(n);
    } else if real.len() < n {
        let mut fallback = poles.to_vec();
        fallback.sort_by(|a, b| {
            a.re.partial_cmp(&b.re)
                .unwrap_or(std::cmp::Ordering::Equal)
                .then(a.im.partial_cmp(&b.im).unwrap_or(std::cmp::Ordering::Equal))
        });
        real.extend(fallback.into_iter().take(n - real.len()));
    }
    Array1::from(real)
}

fn initial_poles(s: ArrayView1<Complex64>, n_poles: usize, real_only: bool) -> Array1<Complex64> {
    if real_only || n_poles < 2 {
        let mut r: Vec<f64> = s
            .iter()
            .filter(|z| z.norm() > 0.0)
            .map(|z| z.norm())
            .collect();
        if r.is_empty() {
            r = s.iter().map(|z| z.norm() + 1.0).collect();
        }
        let lo = (r.iter().fold(f64::INFINITY, |a, &b| a.min(b)) * 0.1).max(1e-6);
        let hi = r.iter().fold(0.0f64, |a, &b| a.max(b)) * 2.0;
        let vals = geomspace(lo, hi, n_poles);
        return Array1::from(
            vals.into_iter()
                .map(|v| Complex64::new(-v, 0.0))
                .collect::<Vec<_>>(),
        );
    }
    let mut w: Vec<f64> = s.iter().map(|z| z.im.abs()).filter(|x| *x > 0.0).collect();
    if w.is_empty() {
        w = vec![1e-2, 1e2];
    }
    let lo = (w.iter().fold(f64::INFINITY, |a, &b| a.min(b)) * 0.5).max(1e-3);
    let hi = w.iter().fold(0.0f64, |a, &b| a.max(b)) * 1.5;
    let pts = geomspace(lo, hi, n_poles / 2);
    let mut poles = Vec::new();
    for &p in &pts {
        poles.push(Complex64::new(-0.01 * p, p));
    }
    let mut cc = poles.clone();
    cc.extend(poles.iter().map(|p| p.conj()));
    if n_poles % 2 == 1 {
        cc.push(Complex64::new(-pts.last().copied().unwrap_or(1.0), 0.0));
    }
    cc.truncate(n_poles);
    Array1::from(cc)
}

fn geomspace(lo: f64, hi: f64, n: usize) -> Vec<f64> {
    if n == 0 {
        return Vec::new();
    }
    if n == 1 {
        return vec![lo];
    }
    let log_lo = lo.log10();
    let log_hi = hi.log10();
    (0..n)
        .map(|i| {
            let t = i as f64 / ((n - 1) as f64);
            10f64.powf(log_lo + t * (log_hi - log_lo))
        })
        .collect()
}

fn complex_lstsq(
    a: ArrayView2<Complex64>,
    b: ArrayView2<Complex64>,
) -> Result<Array2<Complex64>, GspError> {
    let (m, n) = a.dim();
    if b.nrows() != m {
        return Err(GspError::Dimensions(format!(
            "least squares row mismatch: A is {m}x{n}, B has {} rows",
            b.nrows()
        )));
    }
    let p = b.ncols();
    let a_data = a.iter().copied().collect::<Vec<_>>();
    let b_data = b.iter().copied().collect::<Vec<_>>();
    let adm = nalgebra::DMatrix::<Complex64>::from_row_slice(m, n, &a_data);
    let bdm = nalgebra::DMatrix::<Complex64>::from_row_slice(m, p, &b_data);
    let svd = adm.svd(true, true);
    let x = svd
        .solve(&bdm, 1e-12)
        .map_err(|_| GspError::Factorization("complex least-squares solve failed".to_string()))?;
    let mut out = Array2::<Complex64>::zeros((n, p));
    for r in 0..n {
        for c in 0..p {
            out[[r, c]] = x[(r, c)];
        }
    }
    Ok(out)
}

fn complex_eigvals(a: ArrayView2<Complex64>) -> Result<Array1<Complex64>, GspError> {
    let n = a.nrows();
    if n != a.ncols() {
        return Err(GspError::Dimensions(format!(
            "eigendecomposition requires square matrix, got {}x{}",
            a.nrows(),
            a.ncols()
        )));
    }
    let data = a
        .as_slice()
        .ok_or_else(|| GspError::Factorization("matrix is not contiguous".to_string()))?;
    let dm = nalgebra::DMatrix::<Complex64>::from_row_slice(n, n, data);
    let schur = dm.schur();
    let (_q, t) = schur.unpack();
    let mut vals = Vec::with_capacity(n);
    for i in 0..n {
        vals.push(t[(i, i)]);
    }
    Ok(Array1::from(vals))
}