/*
 * // 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::filter_zpk::filter_zpk;
use crate::zpk::{zpk2sos, zpk2tf};
use crate::{BaFilter, FilterType, ImpliirError, SosFilter, ZpkFilter};
use num_complex::Complex64;
use pxfm::f_sincospi;

/// Builder for a Butterworth IIR filter.
#[derive(Copy, Clone)]
pub struct ButterworthFilter {
    order: usize,
    filter_type: FilterType,
    freqs: [f64; 2],
    /// Sample rate in Hz; defaults to 2.0 so that frequencies are normalised
    /// fractions of Nyquist, matching scipy's convention when `fs` is omitted.
    fs: f64,
}

impl ButterworthFilter {
    pub fn new(order: usize) -> Self {
        Self {
            order,
            filter_type: FilterType::Lowpass,
            freqs: [0.0; 2],
            fs: 2.0,
        }
    }

    pub fn lowpass(mut self, cutoff: f64) -> Self {
        self.filter_type = FilterType::Lowpass;
        self.freqs[0] = cutoff;
        self
    }

    pub fn highpass(mut self, cutoff: f64) -> Self {
        self.filter_type = FilterType::Highpass;
        self.freqs[0] = cutoff;
        self
    }

    pub fn bandpass(mut self, low: f64, high: f64) -> Self {
        self.filter_type = FilterType::Bandpass;
        self.freqs[0] = low;
        self.freqs[1] = high;
        self
    }

    /// Set the sample rate in Hz (default: `2.0`).
    pub fn sample_rate(mut self, fs: f64) -> Self {
        self.fs = fs;
        self
    }

    /// Design the filter and return second-order sections.
    pub fn sos(self) -> Result<Vec<SosFilter>, ImpliirError> {
        let zpk = filter_zpk(
            self.order,
            0.,
            self.filter_type,
            self.freqs,
            self.fs,
            buttap,
        )?;
        Ok(zpk2sos(zpk))
    }

    /// Design the filter and return transfer function coefficients `(b, a)`.
    pub fn ba(self) -> Result<BaFilter, ImpliirError> {
        let zpk = filter_zpk(
            self.order,
            0.,
            self.filter_type,
            self.freqs,
            self.fs,
            buttap,
        )?;
        Ok(zpk2tf(&zpk))
    }

    /// Design the filter and return zeros, poles, and gain.
    pub fn zpk(self) -> Result<ZpkFilter, ImpliirError> {
        let zpk = filter_zpk(
            self.order,
            0.,
            self.filter_type,
            self.freqs,
            self.fs,
            buttap,
        )?;
        Ok(zpk)
    }
}

fn buttap(order: usize, _: f64) -> ZpkFilter {
    let n = order as f64;
    // scipy: m = arange(-N+1, N, 2)  →  -N+1, -N+3, ..., N-3, N-1
    // pole k: -exp(i * pi * m[k] / (2*N))
    let poles = (0..order)
        .map(|k| {
            let m = -(n as i32) + 1 + 2 * k as i32;
            let angle = m as f64 / (2.0 * n);
            let v_sincos = f_sincospi(angle);
            -Complex64::new(v_sincos.1, v_sincos.0)
        })
        .collect();
    ZpkFilter {
        z: vec![],
        p: poles,
        k: 1.0,
    }
}

#[cfg(test)]
mod butter_tests {
    use super::*;

    const ATOL: f64 = 1e-8;

    fn assert_sos_close(got: &[SosFilter], expected: &[[f64; 6]], label: &str) {
        assert_eq!(
            got.len(),
            expected.len(),
            "{label}: got {} sections, expected {}",
            got.len(),
            expected.len()
        );
        for (i, (s, e)) in got.iter().zip(expected).enumerate() {
            for (j, (&g, &ex)) in s.b.iter().chain(s.a.iter()).zip(e).enumerate() {
                let err = (g - ex).abs();
                assert!(
                    err <= ATOL + ATOL * ex.abs(),
                    "{label} section {i} coeff {j}: got {g:.10e}, expected {ex:.10e}, err={err:.3e}"
                );
            }
        }
    }

    fn assert_vec_close(got: &[f64], expected: &[f64], label: &str) {
        assert_eq!(got.len(), expected.len(), "{label}: length mismatch");
        for (i, (&g, &e)) in got.iter().zip(expected).enumerate() {
            let err = (g - e).abs();
            assert!(
                err <= ATOL + ATOL * e.abs(),
                "{label}[{i}]: got {g:.10e}, expected {e:.10e}, err={err:.3e}"
            );
        }
    }
    // -----------------------------------------------------------------------
    // Reproduce in Python:
    //   from scipy.signal import butter
    //   sos = butter(N, freq, btype='...', fs=fs, output='sos')
    //   for row in sos: print(list(row))
    // -----------------------------------------------------------------------

    /// butter(2, 100.0, btype='lowpass', fs=1000.0)
    #[test]
    fn test_lowpass_order2() {
        let sos = ButterworthFilter::new(2)
            .lowpass(100.0)
            .sample_rate(1000.0)
            .sos()
            .unwrap();
        let expected = [[
            6.74552738890718956e-02,
            1.34910547778143791e-01,
            6.74552738890718956e-02,
            1.00000000000000000e+00,
            -1.14298050253990113e+00,
            4.12801598096188771e-01,
        ]];
        assert_sos_close(&sos, &expected, "lowpass order=2");
    }

    /// butter(3, 200.0, btype='highpass', fs=1000.0)
    #[test]
    fn test_highpass_order3() {
        let sos = ButterworthFilter::new(3)
            .highpass(200.0)
            .sample_rate(1000.0)
            .sos()
            .unwrap();
        let expected = [
            [
                2.56915601248463299e-01,
                -5.13831202496926598e-01,
                2.56915601248463299e-01,
                1.00000000000000000e+00,
                -1.58384440324536274e-01,
                0.00000000000000000e+00,
            ],
            [
                1.00000000000000000e+00,
                -1.00000000000000000e+00,
                0.00000000000000000e+00,
                1.00000000000000000e+00,
                -4.18856084481766311e-01,
                3.55446762172390407e-01,
            ],
        ];
        assert_sos_close(&sos, &expected, "highpass order=3");
    }

    /// butter(4, [100.0, 300.0], btype='bandpass', fs=1000.0)
    #[test]
    fn test_bandpass_order4() {
        let sos = ButterworthFilter::new(4)
            .bandpass(100.0, 300.0)
            .sample_rate(1000.0)
            .sos()
            .unwrap();
        let expected = [
            [
                4.65829066364436620e-02,
                9.31658132728873239e-02,
                4.65829066364436620e-02,
                1.00000000000000000e+00,
                -7.36083629467364307e-03,
                1.95604139443682934e-01,
            ],
            [
                1.00000000000000000e+00,
                -2.00000000000000000e+00,
                1.00000000000000000e+00,
                1.00000000000000000e+00,
                -8.82228713489221494e-01,
                3.30195746880084184e-01,
            ],
            [
                1.00000000000000000e+00,
                2.00000000000000000e+00,
                1.00000000000000000e+00,
                1.00000000000000000e+00,
                4.57532645411652839e-01,
                6.24538359972242496e-01,
            ],
            [
                1.00000000000000000e+00,
                -2.00000000000000000e+00,
                1.00000000000000000e+00,
                1.00000000000000000e+00,
                -1.39454092383501460e+00,
                7.46672423426420817e-01,
            ],
        ];
        assert_sos_close(&sos, &expected, "bandpass order=4");
    }

    /// butter(2, 0.1) — normalised (default fs=2.0)
    #[test]
    fn test_normalised_lowpass_order2() {
        let sos = ButterworthFilter::new(2).lowpass(0.1).sos().unwrap();
        let expected = [[
            2.00833655642112321e-02,
            4.01667311284224643e-02,
            2.00833655642112321e-02,
            1.00000000000000000e+00,
            -1.56101807580071816e+00,
            6.41351538057563064e-01,
        ]];
        assert_sos_close(&sos, &expected, "normalised lowpass order=2");
    }

    /// butter(1, 0.2, btype='highpass') — normalised, order 1
    #[test]
    fn test_normalised_highpass_order1() {
        let sos = ButterworthFilter::new(1).highpass(0.2).sos().unwrap();
        let expected = [[
            7.54762724747214397e-01,
            -7.54762724747214397e-01,
            0.00000000000000000e+00,
            1.00000000000000000e+00,
            -5.09525449494428795e-01,
            0.00000000000000000e+00,
        ]];
        assert_sos_close(&sos, &expected, "normalised highpass order=1");
    }

    /// butter(6, 0.15, btype='lowpass') — normalised, higher order
    #[test]
    fn test_normalised_lowpass_order6() {
        let sos = ButterworthFilter::new(6).lowpass(0.15).sos().unwrap();
        let expected = [
            [
                7.62145490093100178e-05,
                1.52429098018620036e-04,
                7.62145490093100178e-05,
                1.00000000000000000e+00,
                -1.23878126513851550e+00,
                3.90316716554843735e-01,
            ],
            [
                1.00000000000000000e+00,
                2.00000000000000000e+00,
                1.00000000000000000e+00,
                1.00000000000000000e+00,
                -1.34896774525279484e+00,
                5.13981894219675883e-01,
            ],
            [
                1.00000000000000000e+00,
                2.00000000000000000e+00,
                1.00000000000000000e+00,
                1.00000000000000000e+00,
                -1.59464056877718496e+00,
                7.89706949934815161e-01,
            ],
        ];
        assert_sos_close(&sos, &expected, "normalised lowpass order=6");
    }

    /// butter(2, 100.0, btype='lowpass', fs=1000.0, output='ba')
    #[test]
    fn test_ba_lowpass_order2() {
        let ba = ButterworthFilter::new(2)
            .lowpass(100.0)
            .sample_rate(1000.0)
            .ba()
            .unwrap();
        let expected_b = vec![
            6.74552738890718956e-02,
            1.34910547778143791e-01,
            6.74552738890718956e-02,
        ];
        let expected_a = vec![
            1.00000000000000000e+00,
            -1.14298050253990113e+00,
            4.12801598096188771e-01,
        ];
        assert_vec_close(&ba.b, &expected_b, "ba lowpass order=2 b");
        assert_vec_close(&ba.a, &expected_a, "ba lowpass order=2 a");
    }

    /// butter(3, 200.0, btype='highpass', fs=1000.0, output='ba')
    #[test]
    fn test_ba_highpass_order3() {
        let ba = ButterworthFilter::new(3)
            .highpass(200.0)
            .sample_rate(1000.0)
            .ba()
            .unwrap();
        let expected_b = vec![
            2.56915601248463299e-01,
            -7.70746803745389952e-01,
            7.70746803745389952e-01,
            -2.56915601248463299e-01,
        ];
        let expected_a = vec![
            1.00000000000000000e+00,
            -5.77240524806302613e-01,
            4.21787048689561628e-01,
            -5.62972364918425949e-02,
        ];
        assert_vec_close(&ba.b, &expected_b, "ba highpass order=3 b");
        assert_vec_close(&ba.a, &expected_a, "ba highpass order=3 a");
    }

    /// butter(2, 0.1, output='ba') — normalised
    #[test]
    fn test_ba_normalised_lowpass_order2() {
        let ba = ButterworthFilter::new(2).lowpass(0.1).ba().unwrap();
        let expected_b = vec![
            2.00833655642112321e-02,
            4.01667311284224643e-02,
            2.00833655642112321e-02,
        ];
        let expected_a = vec![
            1.00000000000000000e+00,
            -1.56101807580071816e+00,
            6.41351538057563064e-01,
        ];
        assert_vec_close(&ba.b, &expected_b, "ba normalised lowpass order=2 b");
        assert_vec_close(&ba.a, &expected_a, "ba normalised lowpass order=2 a");
    }

    // ── ZPK output tests ─────────────────────────────────────────────────────

    /// butter(2, 100.0, output='zpk', fs=1000.0)
    #[test]
    fn test_zpk_lowpass_order2() {
        let zpk = ButterworthFilter::new(2)
            .lowpass(100.0)
            .sample_rate(1000.0)
            .zpk()
            .unwrap();
        assert_eq!(zpk.z.len(), 2);
        assert_eq!(zpk.p.len(), 2);
        // All zeros at z = -1 for a lowpass
        for z in &zpk.z {
            assert!(
                (z.re + 1.0).abs() < 1e-10,
                "zero re should be -1, got {}",
                z.re
            );
            assert!(z.im.abs() < 1e-10, "zero im should be 0, got {}", z.im);
        }
        let tol = 1e-8;
        assert!(
            (zpk.k - 6.74552738890718956e-02).abs() < tol,
            "gain mismatch: {}",
            zpk.k
        );
    }

    /// ZPK → BA round-trip: zpk2tf(butter_zpk) == butter_ba
    #[test]
    fn test_zpk_ba_roundtrip() {
        let ba = ButterworthFilter::new(3)
            .highpass(200.0)
            .sample_rate(1000.0)
            .ba()
            .unwrap();
        let zpk = ButterworthFilter::new(3)
            .highpass(200.0)
            .sample_rate(1000.0)
            .zpk()
            .unwrap();
        let ba_from_zpk = zpk2tf(&zpk);
        assert_vec_close(&ba_from_zpk.b, &ba.b, "roundtrip b");
        assert_vec_close(&ba_from_zpk.a, &ba.a, "roundtrip a");
    }
}
