/*
 * // Copyright (c) Radzivon Bartoshyk 3/2026. All rights reserved.
 * //
 * // Redistribution and use in source and binary forms, with or without modification,
 * // are permitted provided that the following conditions are met:
 * //
 * // 1.  Redistributions of source code must retain the above copyright notice, this
 * // list of conditions and the following disclaimer.
 * //
 * // 2.  Redistributions in binary form must reproduce the above copyright notice,
 * // this list of conditions and the following disclaimer in the documentation
 * // and/or other materials provided with the distribution.
 * //
 * // 3.  Neither the name of the copyright holder nor the names of its
 * // contributors may be used to endorse or promote products derived from
 * // this software without specific prior written permission.
 * //
 * // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * // AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * // DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
 * // FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * // DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * // SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 * // CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 * // OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */
use crate::{BaFilter, SosFilter, ZpkFilter};
use num_complex::Complex64;

/// Split into (one-per-conjugate-pair complex, real) lists.
/// Each entry in the complex list is the positive-imaginary element.
pub(crate) fn split_cplx_real(v: Vec<Complex64>) -> (Vec<Complex64>, Vec<f64>) {
    const EPS: f64 = 1e-10;
    let mut used = vec![false; v.len()];
    let mut cplx: Vec<Complex64> = Vec::new();
    let mut real: Vec<f64> = Vec::new();

    for i in 0..v.len() {
        if used[i] {
            continue;
        }
        used[i] = true;
        if v[i].im.abs() < EPS {
            real.push(v[i].re);
        } else {
            let c = if v[i].im > 0.0 { v[i] } else { v[i].conj() };
            cplx.push(c);
            // mark conjugate partner as consumed
            for j in (i + 1)..v.len() {
                if !used[j] && (v[j] - v[i].conj()).norm() < EPS {
                    used[j] = true;
                    break;
                }
            }
        }
    }
    (cplx, real)
}

/// Pop the pole closest to the unit circle. Returns `(pole, is_complex)`.
pub(crate) fn pop_worst_pole(
    p_cplx: &mut Vec<Complex64>,
    p_real: &mut Vec<f64>,
) -> (Complex64, bool) {
    let unit_dist = |c: Complex64| (1.0 - c.norm()).abs();

    let best_c = p_cplx
        .iter()
        .enumerate()
        .min_by(|&(_, &a), &(_, &b)| unit_dist(a).partial_cmp(&unit_dist(b)).unwrap());
    let best_r = p_real.iter().enumerate().min_by(|&(_, &a), &(_, &b)| {
        unit_dist(Complex64::from(a))
            .partial_cmp(&unit_dist(Complex64::from(b)))
            .unwrap()
    });

    match (best_c, best_r) {
        (Some((ci, &c)), Some((ri, &r))) => {
            if unit_dist(c) <= unit_dist(Complex64::from(r)) {
                p_cplx.remove(ci);
                (c, true)
            } else {
                p_real.remove(ri);
                (Complex64::from(r), false)
            }
        }
        (Some((ci, &c)), None) => {
            p_cplx.remove(ci);
            (c, true)
        }
        (None, Some((ri, &r))) => {
            p_real.remove(ri);
            (Complex64::from(r), false)
        }
        (None, None) => (Complex64::new(0.0, 0.0), false),
    }
}

pub(crate) enum ZeroKind {
    Cplx(Complex64),
    Real(f64),
}

/// Pop the zero nearest to `p` from either list.
pub(crate) fn pop_nearest_zero_any(
    z_cplx: &mut Vec<Complex64>,
    z_real: &mut Vec<f64>,
    p: Complex64,
) -> Option<ZeroKind> {
    let best_c = z_cplx
        .iter()
        .enumerate()
        .min_by(|&(_, &a), &(_, &b)| (p - a).norm().partial_cmp(&(p - b).norm()).unwrap());
    let best_r = z_real
        .iter()
        .enumerate()
        .min_by(|&(_, &a), &(_, &b)| (p - a).norm().partial_cmp(&(p - b).norm()).unwrap());

    match (best_c, best_r) {
        (Some((ci, &c)), Some((ri, &r))) => {
            if (p - c).norm() <= (p - r).norm() {
                z_cplx.remove(ci);
                Some(ZeroKind::Cplx(c))
            } else {
                z_real.remove(ri);
                Some(ZeroKind::Real(r))
            }
        }
        (Some((ci, &c)), None) => {
            z_cplx.remove(ci);
            Some(ZeroKind::Cplx(c))
        }
        (None, Some((ri, &r))) => {
            z_real.remove(ri);
            Some(ZeroKind::Real(r))
        }
        (None, None) => None,
    }
}

/// Pop the real zero nearest to `p`.
pub(crate) fn pop_nearest_real_zero(z_real: &mut Vec<f64>, p: Complex64) -> Option<f64> {
    let idx = z_real
        .iter()
        .enumerate()
        .min_by(|&(_, &a), &(_, &b)| (p - a).norm().partial_cmp(&(p - b).norm()).unwrap())
        .map(|(i, _)| i)?;
    Some(z_real.remove(idx))
}

/// Pop the real pole nearest to `p` (same logic as zero, separate name for clarity).
pub(crate) fn pop_nearest_real_pole(p_real: &mut Vec<f64>, p: Complex64) -> Option<f64> {
    pop_nearest_real_zero(p_real, p)
}

/// Build one `[b0,b1,b2, a0,a1,a2]` section from up to 2 zeros/poles and a gain.
pub(crate) fn sos_from_zpk(z: &[Complex64], p: &[Complex64], k: f64) -> [f64; 6] {
    // (x - z0)(x - z1) = x^2 - (z0+z1)x + z0*z1
    let (b1, b2) = match z.len() {
        0 => (0.0, 0.0),
        1 => (-z[0].re, 0.0),
        _ => (-(z[0] + z[1]).re, (z[0] * z[1]).re),
    };
    let (a1, a2) = match p.len() {
        0 => (0.0, 0.0),
        1 => (-p[0].re, 0.0),
        _ => (-(p[0] + p[1]).re, (p[0] * p[1]).re),
    };
    [k, k * b1, k * b2, 1.0, a1, a2]
}

#[allow(clippy::same_item_push)]
pub(crate) fn zpk2sos(zpk_filter: ZpkFilter) -> Vec<SosFilter> {
    let (mut z_cplx, mut z_real) = split_cplx_real(zpk_filter.z);
    let (mut p_cplx, mut p_real) = split_cplx_real(zpk_filter.p);

    // Pad zeros to equal pole count
    let n_poles = p_cplx.len() * 2 + p_real.len();
    let n_zeros = z_cplx.len() * 2 + z_real.len();
    for _ in 0..(n_poles as i32 - n_zeros as i32).max(0) {
        z_real.push(0.0);
    }

    // If total pole count is odd, add one pole and zero at origin
    if (p_cplx.len() * 2 + p_real.len()) % 2 == 1 {
        p_real.push(0.0);
        z_real.push(0.0);
    }

    // p_cplx holds ONE of each conjugate pair → each entry = 1 section.
    // p_real is always even after the odd-pad above → each pair = 1 section.
    let n_sections = p_cplx.len() + p_real.len() / 2;
    let mut sections: Vec<[f64; 6]> = Vec::with_capacity(n_sections);

    for _ in 0..n_sections {
        let (p1, is_cplx) = pop_worst_pole(&mut p_cplx, &mut p_real);

        let section = if is_cplx {
            match pop_nearest_zero_any(&mut z_cplx, &mut z_real, p1) {
                Some(ZeroKind::Cplx(z1)) => sos_from_zpk(&[z1, z1.conj()], &[p1, p1.conj()], 1.0),
                Some(ZeroKind::Real(z1)) => {
                    let z2 = pop_nearest_real_zero(&mut z_real, p1);
                    let zeros: Vec<Complex64> = std::iter::once(Complex64::from(z1))
                        .chain(z2.map(Complex64::from))
                        .collect();
                    sos_from_zpk(&zeros, &[p1, p1.conj()], 1.0)
                }
                None => sos_from_zpk(&[], &[p1, p1.conj()], 1.0),
            }
        } else {
            // Real pole: pair with another real pole
            let p2 = pop_nearest_real_pole(&mut p_real, p1)
                .map(Complex64::from)
                .unwrap_or(Complex64::new(0.0, 0.0));

            // FIX: use pop_nearest_zero_any (searches both z_cplx AND z_real),
            // not pop_nearest_real_zero (searches only z_real).
            // scipy uses _nearest_real_complex_idx(z, p1, 'any') here too.

            match pop_nearest_zero_any(&mut z_cplx, &mut z_real, p1) {
                Some(ZeroKind::Cplx(z1)) => {
                    // complex zero → add its conjugate
                    sos_from_zpk(&[z1, z1.conj()], &[p1, p2], 1.0)
                }
                Some(ZeroKind::Real(z1)) => {
                    // real zero → try to grab a second real zero
                    let z2 = pop_nearest_real_zero(&mut z_real, p1);
                    let zeros: Vec<Complex64> = std::iter::once(Complex64::from(z1))
                        .chain(z2.map(Complex64::from))
                        .collect();
                    sos_from_zpk(&zeros, &[p1, p2], 1.0)
                }
                None => sos_from_zpk(&[], &[p1, p2], 1.0),
            }
        };

        sections.push(section);
    }

    // scipy processes worst (closest to unit circle) first at index n_sections-1
    // and fills downward; reversing matches that ordering.
    sections.reverse();

    // Gain goes into the first section's numerator
    if let Some(s) = sections.first_mut() {
        s[0] *= zpk_filter.k;
        s[1] *= zpk_filter.k;
        s[2] *= zpk_filter.k;
    }

    sections
        .into_iter()
        .map(|s| SosFilter::new([s[0], s[1], s[2]], [s[3], s[4], s[5]]))
        .collect()
}

/// Expand roots into monic polynomial coefficients (highest power first),
/// matching `numpy.poly(roots)`.
///
/// Uses the recurrence:  coeffs = convolve(coeffs, [1, -root])  for each root,
/// starting from `[1.0]`.  Only the real part is kept at the end (imaginary
/// parts cancel for conjugate-pair inputs).
pub(crate) fn poly(roots: &[Complex64]) -> Vec<f64> {
    let mut coeffs = vec![Complex64::from(1.0)];
    for &r in roots {
        // multiply current polynomial by (x - r)
        let mut next = vec![Complex64::from(0.0); coeffs.len() + 1];
        for (i, &c) in coeffs.iter().enumerate() {
            next[i] += c;
            next[i + 1] -= c * r;
        }
        coeffs = next;
    }
    coeffs.iter().map(|c| c.re).collect()
}

/// matching `scipy.signal.zpk2tf`.
///
/// `b = k * poly(z)` and `a = poly(p)` where `poly(r)` gives the monic
/// polynomial whose roots are `r`, expanding
/// `(x - r[0])(x - r[1])...(x - r[n])` into coefficient form.
pub(crate) fn zpk2tf(zpk: &ZpkFilter) -> BaFilter {
    let b = poly(&zpk.z).into_iter().map(|c| c * zpk.k).collect();
    let a = poly(&zpk.p);
    BaFilter { b, a }
}

/// Return a digital IIR filter from an analog one using a bilinear transform.
///
/// Transforms poles and zeros from the analog s-plane to the digital z-plane
/// using Tustin's method: `s = 2*fs*(z-1)/(z+1)`.
///
/// No pre-warping is done here — pre-warp frequencies before calling
/// (see `warped` computation in the pipeline).
pub(crate) fn bilinear_zpk(
    z: Vec<Complex64>,
    p: Vec<Complex64>,
    k: f64,
    fs: f64,
) -> (Vec<Complex64>, Vec<Complex64>, f64) {
    let fs2 = 2.0 * fs; // real scalar, same as scipy's fs2
    let fs2_c = Complex64::from(fs2);

    // z_z = (fs2 + z) / (fs2 - z)
    let z_d: Vec<Complex64> = z.iter().map(|&zi| (fs2_c + zi) / (fs2_c - zi)).collect();

    // p_z = (fs2 + p) / (fs2 - p)
    let p_d: Vec<Complex64> = p.iter().map(|&pi| (fs2_c + pi) / (fs2_c - pi)).collect();

    // Zeros at infinity → Nyquist frequency (z = -1)
    // degree = len(p) - len(z), clamped to 0
    let degree = (p.len() as i32 - z.len() as i32).max(0) as usize;
    let mut z_d_full = z_d;
    z_d_full.extend(std::iter::repeat_n(Complex64::new(-1.0, 0.0), degree));

    // k_z = k * real(prod(fs2 - z) / prod(fs2 - p))
    let num: Complex64 = z
        .iter()
        .fold(Complex64::from(k), |acc, &zi| acc * (fs2_c - zi));
    let den: Complex64 = p
        .iter()
        .fold(Complex64::from(1.0), |acc, &pi| acc * (fs2_c - pi));
    let k_d = (num / den).re;

    (z_d_full, p_d, k_d)
}

pub(crate) fn lp2lp_zpk(
    z: Vec<Complex64>,
    p: Vec<Complex64>,
    k: f64,
    wo: f64,
) -> (Vec<Complex64>, Vec<Complex64>, f64) {
    let degree = p.len() as i32 - z.len() as i32;
    let z2 = z.iter().map(|&zi| zi * wo).collect();
    let p2 = p.iter().map(|&pi| pi * wo).collect();
    (z2, p2, k * wo.powi(degree))
}

/// Transform a lowpass filter prototype to a highpass filter.
///
/// Returns an analog high-pass filter with cutoff frequency `wo` from an
/// analog low-pass filter prototype with unity cutoff frequency, using
/// zeros, poles, and gain (ZPK) representation.
///
/// Derived from the s-plane substitution `s → ω₀ / s`.
///
/// # Arguments
///
/// * `z`  — zeros of the lowpass prototype
/// * `p`  — poles of the lowpass prototype
/// * `k`  — system gain of the lowpass prototype
/// * `wo` — desired cutoff angular frequency (rad/s)
///
/// # Example
///
/// ```
/// // lp2hp_zpk(z=[-2+3j, -0.5-0.8j], p=[-1, -4], k=10, wo=0.6)
/// // → z=[-0.0923-0.1385j, -0.3371+0.5393j], p=[-0.6, -0.15], k=8.5
/// ```
pub(crate) fn lp2hp_zpk(
    z: Vec<Complex64>,
    p: Vec<Complex64>,
    k: f64,
    wo: f64,
) -> (Vec<Complex64>, Vec<Complex64>, f64) {
    // degree = len(p) - len(z): number of poles with no matching zero.
    // These correspond to zeros at infinity in the LP prototype, which
    // map to zeros at the origin in the HP filter.
    let degree = p.len() as i32 - z.len() as i32;

    // z_hp = wo / z  (invert + scale each zero)
    let mut z_hp: Vec<Complex64> = z.iter().map(|&zi| Complex64::from(wo) / zi).collect();

    // p_hp = wo / p  (invert + scale each pole)
    let p_hp: Vec<Complex64> = p.iter().map(|&pi| Complex64::from(wo) / pi).collect();

    // LP zeros at infinity → HP zeros at origin
    z_hp.extend(std::iter::repeat_n(
        Complex64::new(0.0, 0.0),
        degree.max(0) as usize,
    ));

    // Gain: k * real(prod(-z) / prod(-p))
    // Must be done in complex arithmetic; .re taken only at the end.
    let prod_neg_z: Complex64 = z.iter().fold(Complex64::from(1.0), |acc, &zi| acc * (-zi));
    let prod_neg_p: Complex64 = p.iter().fold(Complex64::from(1.0), |acc, &pi| acc * (-pi));
    let k_hp = k * (prod_neg_z / prod_neg_p).re;

    (z_hp, p_hp, k_hp)
}

pub(crate) fn lp2bp_zpk(
    z: Vec<Complex64>,
    p: Vec<Complex64>,
    k: f64,
    wo: f64,
    bw: f64,
) -> (Vec<Complex64>, Vec<Complex64>, f64) {
    let wo2 = wo * wo;

    let bp_pair = |c: Complex64| -> [Complex64; 2] {
        let t = c * bw;
        let disc = t * t - Complex64::from(4.0 * wo2);
        let sq = disc.sqrt();
        [(t + sq) / 2.0, (t - sq) / 2.0]
    };

    let z2: Vec<Complex64> = z.iter().flat_map(|&zi| bp_pair(zi)).collect();
    let p2: Vec<Complex64> = p.iter().flat_map(|&pi| bp_pair(pi)).collect();

    let degree = p.len() - z.len();
    let mut z2_full = z2;
    z2_full.extend(std::iter::repeat_n(Complex64::new(0.0, 0.0), degree));

    (z2_full, p2, k * bw.powi(degree as i32))
}
