dsmc 0.3.1

//! Partial-fraction expansion of a rational function by the Heaviside expansion theorem.

use super::TransferFunction;
use crate::{dka_method, vieta_formula, Continuous, Polynomial};
use num_complex::Complex;
use num_traits::{Float, Zero};
use serde::Serialize;
use std::ops::AddAssign;

/// Terms of one (possibly repeated) pole.
#[derive(Clone, Debug, Serialize)]
pub struct PoleTerm<T> {
    /// The pole `p`.
    pub pole: Complex<T>,
    /// `residues[k]` is the coefficient of `1 / (s - pole)^(k + 1)`; its length is the multiplicity.
    pub residues: Vec<Complex<T>>,
}

impl<T> PoleTerm<T> {
    /// Multiplicity of the pole.
    pub fn multiplicity(&self) -> usize {
        self.residues.len()
    }
}

/// `X(s) = direct(s) + Σ terms`.
#[derive(Clone, Debug, Serialize)]
pub struct PartialFraction<T> {
    /// Polynomial part `Q(s)` (descending order). Empty when `X(s)` is strictly proper.
    pub direct: Polynomial<T>,
    /// Terms of each distinct pole.
    pub terms: Vec<PoleTerm<T>>,
}

impl<T: Float + AddAssign> TransferFunction<T, Continuous> {
    /// Partial-fraction expansion by the Heaviside expansion theorem (repeated poles supported).
    /// Poles closer than a relative distance of 1e-4 are treated as one repeated pole.
    ///
    /// ```text
    /// X(s) = Q(s) + Σ_p Σ_{k=1..m_p} r_{p,k} / (s - p)^k
    /// ```
    /// For a pole `p` of multiplicity `m`, `r_{p,k} = G^(m-k)(p) / (m-k)!` with `G(s) = (s - p)^m X(s)`.
    /// The derivatives are obtained exactly as Taylor coefficients of `G` around `p`
    /// (synthetic division + power-series division), not by numerical differentiation.
    /// The time-domain counterpart (inverse Laplace transform) is in `laplace_transform`.
    pub fn partial_fraction(&self) -> PartialFraction<T> {
        self.partial_fraction_with_tolerance(T::from(1e-4).unwrap())
    }

    /// Same as `partial_fraction`, but with an explicit relative tolerance for grouping
    /// numerically found roots into repeated poles.
    pub fn partial_fraction_with_tolerance(&self, rel_tol: T) -> PartialFraction<T> {
        let numer = trim_leading_zeros(&self.numerator);
        let denom = trim_leading_zeros(&self.denominator);
        assert!(!denom.is_empty(), "partial_fraction: denominator is zero");

        let (direct, remainder) = polynomial_division(&numer, &denom);
        let gain = Complex::from(denom[0]);
        let remainder: Vec<Complex<T>> = remainder.iter().map(|&c| Complex::from(c)).collect();

        let denom_complex = Polynomial(denom.iter().map(|&c| Complex::from(c)).collect());
        let roots = dka_method(&denom_complex).unwrap_or_default();
        let poles = group_roots(&denom_complex, &roots, rel_tol, rel_tol);

        let terms = poles
            .iter()
            .enumerate()
            .map(|(i, &(p, m))| {
                // B(s) = D(s) / (s - p)^m, so that G(s) = (s - p)^m X(s) = R(s) / B(s).
                let others: Vec<Complex<T>> = poles
                    .iter()
                    .enumerate()
                    .filter(|&(j, _)| j != i)
                    .flat_map(|(_, &(q, mq))| std::iter::repeat_n(q, mq))
                    .collect();
                let b: Vec<Complex<T>> = vieta_formula(&others).0.iter().map(|&c| c * gain).collect();
                PoleTerm { pole: p, residues: principal_part(&remainder, &b, p, m) }
            })
            .collect();

        PartialFraction { direct, terms }
    }
}

fn trim_leading_zeros<T: Float>(p: &Polynomial<T>) -> Vec<T> {
    p.iter().copied().skip_while(|c| c.is_zero()).collect()
}

/// `numer = quotient * denom + remainder` for descending-order coefficients
/// (`denom` must have a nonzero leading coefficient). `remainder` has `denom.len() - 1` entries.
fn polynomial_division<T: Float>(numer: &[T], denom: &[T]) -> (Polynomial<T>, Vec<T>) {
    let rem_len = denom.len() - 1;
    if numer.len() < denom.len() {
        let mut remainder = vec![T::zero(); rem_len - numer.len()];
        remainder.extend_from_slice(numer);
        return (Polynomial(vec![]), remainder);
    }
    let mut work = numer.to_vec();
    let quotient_len = numer.len() - rem_len;
    let mut quotient = vec![T::zero(); quotient_len];
    for i in 0..quotient_len {
        let c = work[i] / denom[0];
        quotient[i] = c;
        for (j, &d) in denom.iter().enumerate() {
            work[i + j] = work[i + j] - c * d;
        }
    }
    (Polynomial(quotient), work.split_off(quotient_len))
}

/// Group numerically found roots into `(pole, multiplicity)`. A repeated root of multiplicity m
/// is found as a ring-like cluster whose radius grows like eps^(1/m), so a fixed tolerance cannot
/// separate clusters reliably for large m. Instead, starting from each group as a seed, the
/// nearest groups are added one by one, and the largest union that is a consistent cluster
/// (`is_cluster`) is merged; this repeats until no merge is possible. Partial arcs of a ring are
/// usually not consistent clusters, so growing past them (rather than merging pairs) is needed.
/// The cluster mean is then refined by Newton's method on `D^(m-1)(s)`, for which the pole is
/// a simple root. Real / imaginary parts within `snap_tol` (relative to `max(|p|, 1)`) are
/// snapped to zero, so that poles nearly on the real / imaginary axis lie exactly on it.
pub(super) fn group_roots<T: Float>(
    denom: &Polynomial<Complex<T>>,
    roots: &[Complex<T>],
    rel_tol: T,
    snap_tol: T,
) -> Vec<(Complex<T>, usize)> {
    let mut groups: Vec<Vec<Complex<T>>> = roots.iter().map(|&r| vec![r]).collect();
    loop {
        // (member count, spread, group indices) of the best merge found so far
        let mut best: Option<(usize, T, Vec<usize>)> = None;
        for seed in 0..groups.len() {
            let mut chosen = vec![seed];
            let mut members = groups[seed].clone();
            while chosen.len() < groups.len() {
                let center = mean(&members);
                let next = (0..groups.len())
                    .filter(|k| !chosen.contains(k))
                    .min_by(|&a, &b| {
                        let da = (mean(&groups[a]) - center).norm();
                        let db = (mean(&groups[b]) - center).norm();
                        da.partial_cmp(&db).unwrap_or(std::cmp::Ordering::Equal)
                    })
                    .unwrap();
                chosen.push(next);
                members.extend_from_slice(&groups[next]);
                let others = (0..groups.len()).filter(|k| !chosen.contains(k)).flat_map(|k| &groups[k]);
                if let Some(spread) = is_cluster(denom, &members, others, rel_tol) {
                    let better = match &best {
                        None => true,
                        Some((n, s, _)) => members.len() > *n || (members.len() == *n && spread < *s),
                    };
                    if better {
                        best = Some((members.len(), spread, chosen.clone()));
                    }
                }
            }
        }
        let Some((_, _, mut chosen)) = best else { break };
        chosen.sort_unstable();
        let merged: Vec<Complex<T>> = chosen.iter().flat_map(|&k| groups[k].clone()).collect();
        for &k in chosen.iter().rev() {
            groups.swap_remove(k);
        }
        groups.push(merged);
    }
    groups
        .into_iter()
        .map(|members| {
            let m = members.len();
            let mut p = mean(&members);
            if m > 1 {
                p = refine_repeated_root(denom, p, m);
            }
            let scale = snap_tol * p.norm().max(T::one());
            if p.im.abs() <= scale {
                p.im = T::zero();
            }
            if p.re.abs() <= scale {
                p.re = T::zero();
            }
            (p, m)
        })
        .collect()
}

/// `Some(spread)` if `members` form one repeated root: they lie within `cluster_radius` of their
/// mean, and no other root is closer to the mean than a member (otherwise e.g. a conjugate pair
/// centered on another root would be merged).
fn is_cluster<'a, T: Float + 'a>(
    denom: &Polynomial<Complex<T>>,
    members: &[Complex<T>],
    mut others: impl Iterator<Item = &'a Complex<T>>,
    rel_tol: T,
) -> Option<T> {
    let center = mean(members);
    let spread = members.iter().fold(T::zero(), |acc, &r| acc.max((r - center).norm()));
    let isolated = others.all(|&r| (r - center).norm() > spread);
    (isolated && spread <= cluster_radius(denom, center, members.len(), rel_tol)).then_some(spread)
}

fn mean<T: Float>(members: &[Complex<T>]) -> Complex<T> {
    members.iter().fold(Complex::zero(), |acc, &r| acc + r) / T::from(members.len()).unwrap()
}

/// Admissible spread of a cluster of `m` roots around `center`: the larger of
/// `rel_tol * max(|center|, 1)` and the round-off spread of an m-fold root,
/// `(K eps Σ|a_i||center|^i / |D^(m)(center) / m!|)^(1/m)` (K: safety factor, above the
/// threshold at which `dka_method` stops refining a root).
fn cluster_radius<T: Float>(denom: &Polynomial<Complex<T>>, center: Complex<T>, m: usize, rel_tol: T) -> T {
    let base = rel_tol * center.norm().max(T::one());
    let c_m = taylor_coefficients(&denom.0, center, m + 1)[m].norm();
    if c_m.is_zero() {
        return T::infinity();
    }
    let x = center.norm();
    let bound = denom.0.iter().fold(T::zero(), |acc, a| acc * x + a.norm());
    let safety = T::from(100.max(8 * denom.0.len())).unwrap();
    let spread = (safety * T::epsilon() * bound / c_m).powf(T::one() / T::from(m).unwrap());
    base.max(spread)
}

/// Newton's method on `D^(m-1)(s)`, whose root at a pole of multiplicity m is simple.
/// With Taylor coefficients `c_j = D^(j)(p) / j!`, the step is `c_{m-1} / (m c_m)`.
fn refine_repeated_root<T: Float>(denom: &Polynomial<Complex<T>>, mut p: Complex<T>, m: usize) -> Complex<T> {
    let eps = T::epsilon();
    for _ in 0..20 {
        let c = taylor_coefficients(&denom.0, p, m + 1);
        if c[m].is_zero() {
            break;
        }
        let step = c[m - 1] / (c[m] * T::from(m).unwrap());
        p = p - step;
        if step.norm() <= eps * p.norm().max(T::one()) {
            break;
        }
    }
    p
}

/// Residues of `R(s) / ((s - p)^m B(s))` at `p` (`B(p) != 0`, descending `numer` / `rest`):
/// element k is the coefficient of `1 / (s - p)^(k + 1)`. With `G = R / B`, these are
/// `G^(m-1-k)(p) / (m-1-k)!`, obtained by power-series division of Taylor coefficients.
pub(crate) fn principal_part<T: Float>(
    numer: &[Complex<T>],
    rest: &[Complex<T>],
    p: Complex<T>,
    m: usize,
) -> Vec<Complex<T>> {
    let r_taylor = taylor_coefficients(numer, p, m);
    let b_taylor = taylor_coefficients(rest, p, m);

    // Power-series division: g_j = G^(j)(p) / j!
    let mut g = vec![Complex::zero(); m];
    for j in 0..m {
        let mut acc = r_taylor[j];
        for l in 1..=j {
            acc = acc - b_taylor[l] * g[j - l];
        }
        g[j] = acc / b_taylor[0];
    }
    // g_j is the coefficient of 1 / (s - p)^(m - j)
    g.reverse();
    g
}

/// First `count` ascending Taylor coefficients of `P(x0 + h)` (descending `coeffs`),
/// i.e. `P^(j)(x0) / j!` for j = 0..count, by repeated synthetic division by `(s - x0)`.
pub(super) fn taylor_coefficients<T: Float>(coeffs: &[Complex<T>], x0: Complex<T>, count: usize) -> Vec<Complex<T>> {
    let mut work = coeffs.to_vec();
    let mut out = Vec::with_capacity(count);
    for _ in 0..count {
        let mut quotient = Vec::with_capacity(work.len());
        let mut acc = Complex::zero();
        for &a in &work {
            acc = acc * x0 + a;
            quotient.push(acc);
        }
        out.push(quotient.pop().unwrap_or_else(Complex::zero));
        work = quotient;
    }
    out
}

pub(crate) fn fmt_num<T: Float + std::fmt::Display>(c: T, precision: Option<usize>) -> String {
    match precision {
        Some(p) => format!("{:.*}", p, c),
        None => c.to_string(),
    }
}

fn fmt_complex<T: Float + std::fmt::Display>(c: Complex<T>, precision: Option<usize>) -> String {
    if c.im.is_zero() {
        return fmt_num(c.re, precision);
    }
    let sign = if c.im.is_sign_negative() { '-' } else { '+' };
    format!("({} {} j{})", fmt_num(c.re, precision), sign, fmt_num(c.im.abs(), precision))
}

/// Append `coef * body` to a sum, folding the sign of `coef` into the separator.
pub(crate) fn push_term<T: Float + std::fmt::Display>(out: &mut String, coef: T, body: &str, precision: Option<usize>) {
    push_term_with(out, coef, " * ", body, precision);
}

/// Append `coef{op}body` (`op` = `" * "` or `" / "`) to a sum, folding the sign of `coef` into the separator.
/// A unit coefficient is omitted for `" * "`.
fn push_term_with<T: Float + std::fmt::Display>(
    out: &mut String,
    coef: T,
    op: &str,
    body: &str,
    precision: Option<usize>,
) {
    if coef.is_zero() {
        return;
    }
    let negative = coef.is_sign_negative();
    match (out.is_empty(), negative) {
        (true, true) => out.push('-'),
        (true, false) => {}
        (false, true) => out.push_str(" - "),
        (false, false) => out.push_str(" + "),
    }
    let abs = coef.abs();
    if body.is_empty() {
        out.push_str(&fmt_num(abs, precision));
    } else if abs == T::one() && op == " * " {
        out.push_str(body);
    } else {
        out.push_str(&format!("{}{}{}", fmt_num(abs, precision), op, body));
    }
}

impl<T: Float + std::fmt::Display> std::fmt::Display for PartialFraction<T> {
    /// e.g. `s + 1 + 2 / (s + 1) + 3 / (s + 1)^2` (complex poles are printed as `(a + jb)`).
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let prec = f.precision();
        let mut out = String::new();
        let degree = self.direct.len().saturating_sub(1);
        for (i, &c) in self.direct.iter().enumerate() {
            let body = match degree - i {
                0 => String::new(),
                1 => "s".to_string(),
                n => format!("s^{}", n),
            };
            push_term(&mut out, c, &body, prec);
        }
        for term in &self.terms {
            let p = term.pole;
            let factor = if p.is_zero() {
                "s".to_string()
            } else if p.im.is_zero() {
                let sign = if p.re.is_sign_negative() { '+' } else { '-' };
                format!("(s {} {})", sign, fmt_num(p.re.abs(), prec))
            } else {
                format!("(s - {})", fmt_complex(p, prec))
            };
            for (k, &r) in term.residues.iter().enumerate() {
                let power = if k == 0 { factor.clone() } else { format!("{}^{}", factor, k + 1) };
                if r.im.is_zero() {
                    push_term_with(&mut out, r.re, " / ", &power, prec);
                } else if !r.is_zero() {
                    if !out.is_empty() {
                        out.push_str(" + ");
                    }
                    out.push_str(&format!("{} / {}", fmt_complex(r, prec), power));
                }
            }
        }
        if out.is_empty() {
            out = fmt_num(T::zero(), prec);
        }
        write!(f, "{}", out)
    }
}