use num_complex::Complex;

/// Return relative degree of transfer function from zeros and poles
fn relative_degree(z: &[Complex<f64>], p: &[Complex<f64>]) -> usize {
    let degree = p.len() as isize - z.len() as isize;
    if degree < 0 {
        panic!("Improper transfer function. Must have at least as many poles as zeros.");
    }
    degree as usize
}

/// Lowpass to Lowpass transformation in ZPK format
pub fn lp2lp_zpk(
    z: &[Complex<f64>],
    p: &[Complex<f64>],
    k: f64,
    wo: f64,
) -> (Vec<Complex<f64>>, Vec<Complex<f64>>, f64) {
    let degree = relative_degree(z, p);

    let z_lp: Vec<Complex<f64>> = z.iter().map(|&zi| wo * zi).collect();
    let p_lp: Vec<Complex<f64>> = p.iter().map(|&pi| pi * wo).collect();

    let k_lp = k * wo.powi(degree as i32);

    (z_lp, p_lp, k_lp)
}

/// Lowpass to Highpass transformation in ZPK format
pub fn lp2hp_zpk(
    z: &[Complex<f64>],
    p: &[Complex<f64>],
    k: f64,
    wo: f64,
) -> (Vec<Complex<f64>>, Vec<Complex<f64>>, f64) {
    let degree = relative_degree(z, p);

    let z_hp: Vec<Complex<f64>> = z.iter().map(|&zi| wo / zi).collect();
    let p_hp: Vec<Complex<f64>> = p.iter().map(|&pi| Complex::new(wo, 0.0) / pi).collect();

    let mut z_hp_full = z_hp;
    z_hp_full.extend(vec![Complex::new(0.0, 0.0); degree]);

    let num: Complex<f64> = z.iter().map(|&zi| -zi).product();
    let den: Complex<f64> = p.iter().map(|&pi| -pi).product();
    let k_hp = k * (num / den).re;

    (z_hp_full, p_hp, k_hp)
}
/// Lowpass to Bandpass transformation in ZPK format
pub fn lp2bp_zpk(
    z: &[Complex<f64>],
    p: &[Complex<f64>],
    k: f64,
    wo: f64,
    bw: f64,
) -> (Vec<Complex<f64>>, Vec<Complex<f64>>, f64) {
    let degree = relative_degree(z, p);

    // Scale poles and zeros by bandwidth
    let z_lp: Vec<Complex<f64>> = z.iter().map(|&zi| zi * (bw / 2.0)).collect();
    let p_lp: Vec<Complex<f64>> = p.iter().map(|&pi| pi * (bw / 2.0)).collect();

    // Duplicate poles and zeros and shift to +wo and -wo
    let mut z_bp = Vec::new();
    for &zi in &z_lp {
        let sqrt_term = (zi * zi - Complex::new(wo * wo, 0.0)).sqrt();
        z_bp.push(zi + sqrt_term);
        z_bp.push(zi - sqrt_term);
    }

    let mut p_bp = Vec::new();
    for &pi in &p_lp {
        let sqrt_term = (pi * pi - Complex::new(wo * wo, 0.0)).sqrt();
        p_bp.push(pi + sqrt_term);
        p_bp.push(pi - sqrt_term);
    }

    // Add zeros at origin to match filter degree
    z_bp.extend(vec![Complex::new(0.0, 0.0); degree]);

    // Adjust gain
    let k_bp = k * bw.powi(degree as i32);

    (z_bp, p_bp, k_bp)
}

/// Lowpass to Bandstop (Notch) transformation in ZPK format
pub fn lp2bs_zpk(
    z: &[Complex<f64>],
    p: &[Complex<f64>],
    k: f64,
    wo: f64,
    bw: f64,
) -> (Vec<Complex<f64>>, Vec<Complex<f64>>, f64) {
    let degree = relative_degree(z, p);

    // Invert to highpass
    let z_hp: Vec<Complex<f64>> = z.iter().map(|&zi| Complex::new(bw/2.0,0.0) / zi).collect();
    let p_hp: Vec<Complex<f64>> = p.iter().map(|&pi| Complex::new(bw/2.0,0.0) / pi).collect();

    // Duplicate poles and zeros and shift to +wo and -wo
    let mut z_bs = Vec::new();
    for &zi in &z_hp {
        let sqrt_term = (zi*zi - Complex::new(wo*wo, 0.0)).sqrt();
        z_bs.push(zi + sqrt_term);
        z_bs.push(zi - sqrt_term);
    }

    let mut p_bs = Vec::new();
    for &pi in &p_hp {
        let sqrt_term = (pi*pi - Complex::new(wo*wo, 0.0)).sqrt();
        p_bs.push(pi + sqrt_term);
        p_bs.push(pi - sqrt_term);
    }

    // Add zeros at +j*wo and -j*wo
    z_bs.extend(vec![Complex::new(0.0, wo); degree]);
    z_bs.extend(vec![Complex::new(0.0, -wo); degree]);

    // Adjust gain
    let num: Complex<f64> = z.iter().map(|&zi| -zi).product();
    let den: Complex<f64> = p.iter().map(|&pi| -pi).product();
    let k_bs = k * num.re / den.re;

    (z_bs, p_bs, k_bs)
}

/// Bilinear transform in ZPK format
pub fn bilinear_zpk(
    z: &[Complex<f64>],
    p: &[Complex<f64>],
    k: f64,
    fs: Option<f64>,
) -> (Vec<Complex<f64>>, Vec<Complex<f64>>, f64) {
    let degree = relative_degree(z, p);
    let fs2 = if let Some(f) = fs {
        2.0 * f
    }else{
        0.0
    };

    let z_z: Vec<Complex<f64>> = z.iter().map(|&zi| (fs2 + zi) / (fs2 - zi)).collect();
    let p_z: Vec<Complex<f64>> = p.iter().map(|&pi| (Complex::new(fs2, 0.0) + pi) / (Complex::new(fs2, 0.0) - pi)).collect();
    let mut z_z_full = z_z;
    z_z_full.extend(vec![Complex::new(-1.0, 0.0); degree]);

    let num: Complex<f64> = z.iter().map(|&zi| Complex::new(fs2, 0.0)- zi).product();
    let den: Complex<f64> = p.iter().map(|&pi| Complex::new(fs2, 0.0) - pi).product();
    let k_z = k * (num / den).re;

    (z_z_full, p_z, k_z)
}