/*
 * // Copyright (c) Radzivon Bartoshyk 11/2025. 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::cdf53f::define_dwt_cdf_float;
use crate::err::{OscletError, try_vec};
use crate::mla::fmla;
use crate::{
    Dwt, DwtExecutor, DwtForwardExecutor, DwtInverseExecutor, DwtSize, IncompleteDwtExecutor,
    MultiDwt, WaveletSample,
};
use num_traits::AsPrimitive;
use std::marker::PhantomData;

const ALPHA: f64 = -1.5861343420693648;
const BETA: f64 = -0.0529801185718856;
const GAMMA: f64 = 0.8829110755411875;
const DELTA: f64 = 0.4435068520511142;
const K: f64 = 1.1496043988602418;
const INV_K: f64 = 0.86986445162478079;

#[derive(Default)]
pub(crate) struct Cdf97<T> {
    phantom0: PhantomData<T>,
}

fn dwt97_forward_update_even<T: WaveletSample>(approx: &mut [T], details: &mut [T], c: T)
where
    f64: AsPrimitive<T>,
{
    approx[0] = fmla(details[0] + details[0], c, approx[0]);

    let details_next = &details[1..];
    let approx_len = approx.len();
    let approx_k = &mut approx[1..approx_len - 1];

    for ((dst, detail), detail_next) in approx_k
        .iter_mut()
        .zip(details.iter())
        .zip(details_next.iter())
    {
        let d_left = *detail;
        let d_right = *detail_next;
        *dst = fmla(d_left + d_right, c, *dst);
    }

    if approx.len() > 1 {
        let i = approx.len() - 1;
        let d_left = details[i - 1];
        let d_right = *details.last().unwrap();
        approx[i] = fmla(d_left + d_right, c, approx[i]);
    }
}

fn dwt97_forward_update_odd<T: WaveletSample>(approx: &mut [T], details: &mut [T], c: T)
where
    f64: AsPrimitive<T>,
{
    let approx_next = &approx[1..];
    for ((dst, &src_left), &src_right) in details
        .iter_mut()
        .zip(approx.iter())
        .zip(approx_next.iter())
    {
        *dst = fmla(src_left + src_right, c, *dst);
    }

    if approx.len() == details.len() {
        let i = details.len();
        let src_left = approx[i - 1];
        // at boundary, mirror last sample
        let src_right = src_left;
        details[i - 1] = fmla(src_left + src_right, c, details[i - 1]);
    }
}

fn dwt97_scale<T: WaveletSample>(approx: &mut [T], details: &mut [T], k: T, inv_k: T) {
    for (a_dst, d_dst) in approx.iter_mut().zip(details.iter_mut()) {
        *a_dst = *a_dst * k;
        *d_dst = *d_dst * inv_k;
    }
    if approx.len() < details.len() {
        let det_left = &mut details[approx.len()..];
        for x in det_left.iter_mut() {
            *x = *x * inv_k;
        }
    } else if approx.len() > details.len() {
        let app_left = &mut approx[details.len()..];
        for x in app_left.iter_mut() {
            *x = *x * k;
        }
    }
}

impl<T: WaveletSample> DwtForwardExecutor<T> for Cdf97<T>
where
    f64: AsPrimitive<T>,
{
    /// Perform the forward CDF 9/7 DWT (lifting scheme) on a 1D signal.
    /// Splits the signal into approximation (low-pass) and detail (high-pass) coefficients.
    fn execute_forward(
        &self,
        input: &[T],
        approx: &mut [T],
        details: &mut [T],
    ) -> Result<(), OscletError> {
        let n = input.len();
        if n < 4 {
            return Err(OscletError::MinFilterSize(n, 4));
        }
        if approx.len() != n.div_ceil(2) {
            return Err(OscletError::ApproxSizeNotMatches(approx.len(), n));
        }
        if details.len() != n / 2 {
            return Err(OscletError::ApproxSizeNotMatches(details.len(), n));
        }

        let signal_n = &input[1..];
        let signal_n2 = &input[2..];

        // First lifting step
        details[0] = fmla(input[0] + input[1], ALPHA.as_(), input[0]);

        for (((dst, &s_previous), &s_current), &s_next) in details
            .iter_mut()
            .zip(input.iter().step_by(2))
            .zip(signal_n.iter().step_by(2))
            .zip(signal_n2.iter().step_by(2))
        {
            let left = s_previous;
            let right = s_next;
            *dst = fmla(left + right, ALPHA.as_(), s_current);
        }

        // Handle last odd index if n is even (boundary)
        if n.is_multiple_of(2) {
            let i = n - 1;
            let left = input[i - 1];
            *details.last_mut().unwrap() = fmla(left + left, ALPHA.as_(), input[i]);
        }

        for (dst, src) in approx.iter_mut().zip(input.iter().step_by(2)) {
            *dst = *src;
        }

        // Update 1: even += beta * (odd_left + odd_right)
        dwt97_forward_update_even(approx, details, BETA.as_());

        // Predict 2: odd += gamma * (even_left + even_right)
        dwt97_forward_update_odd(approx, details, GAMMA.as_());

        // Update 2: even += delta * (odd_left + odd_right)
        dwt97_forward_update_even(approx, details, DELTA.as_());

        dwt97_scale(approx, details, K.as_(), INV_K.as_());

        Ok(())
    }

    fn execute_forward_with_scratch(
        &self,
        input: &[T],
        approx: &mut [T],
        details: &mut [T],
        _: &mut [T],
    ) -> Result<(), OscletError> {
        self.execute_forward(input, approx, details)
    }

    fn required_scratch_size(&self, _: usize) -> usize {
        0
    }

    fn dwt_size(&self, input_length: usize) -> DwtSize {
        DwtSize {
            approx_length: input_length.div_ceil(2).max(2),
            details_length: (input_length / 2).max(2),
        }
    }
}

impl<T: WaveletSample> DwtInverseExecutor<T> for Cdf97<T>
where
    f64: AsPrimitive<T>,
{
    /// Perform the inverse CDF 9/7 DWT (lifting scheme) to reconstruct the original signal.
    fn execute_inverse(
        &self,
        approx: &[T],
        details: &[T],
        output: &mut [T],
    ) -> Result<(), OscletError> {
        let n = approx.len() + details.len();
        if n != output.len() {
            return Err(OscletError::OutputSizeIsNotValid(output.len(), n));
        }
        let inputs_size = self.dwt_size(output.len());
        if inputs_size.approx_length != approx.len() {
            return Err(OscletError::ApproxSizeNotMatches(
                inputs_size.approx_length,
                approx.len(),
            ));
        }
        if inputs_size.details_length != details.len() {
            return Err(OscletError::ApproxSizeNotMatches(
                inputs_size.approx_length,
                approx.len(),
            ));
        }
        if n < 4 {
            return Err(OscletError::MinFilterSize(n, 4));
        }

        let mut approx_inv = approx.to_vec();
        let mut detail_inv = details.to_vec();

        dwt97_scale(&mut approx_inv, &mut detail_inv, INV_K.as_(), K.as_());

        // Inverse update 2: even -= delta * (odd_left + odd_right)
        dwt97_forward_update_even(&mut approx_inv, &mut detail_inv, (-DELTA).as_());

        // Inverse predict 2: odd -= gamma * (even_left + even_right)
        dwt97_forward_update_odd(&mut approx_inv, &mut detail_inv, (-GAMMA).as_());

        // Inverse update 1: even -= beta * (odd_left + odd_right)
        dwt97_forward_update_even(&mut approx_inv, &mut detail_inv, (-BETA).as_());

        // Inverse predict 1: odd -= alpha * (even_left + even_right)
        dwt97_forward_update_odd(&mut approx_inv, &mut detail_inv, (-ALPHA).as_());

        // Interleave approx and detail to reconstruct signal
        for ((dst, &src_even), &src_odd) in output
            .as_chunks_mut::<2>()
            .0
            .iter_mut()
            .zip(approx_inv.iter())
            .zip(detail_inv.iter())
        {
            dst[0] = src_even;
            dst[1] = src_odd;
        }
        if !output.len().is_multiple_of(2) {
            *output.last_mut().unwrap() = *approx_inv.last().unwrap();
        }

        Ok(())
    }

    fn idwt_size(&self, input_length: DwtSize) -> usize {
        (input_length.approx_length + input_length.details_length).max(4)
    }
}

impl<T: WaveletSample> IncompleteDwtExecutor<T> for Cdf97<T>
where
    f64: AsPrimitive<T>,
{
    fn filter_length(&self) -> usize {
        10
    }
}

define_dwt_cdf_float!(Cdf97, 4);

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

    #[test]
    fn test_cdf97() {
        let m_cdf97 = Cdf97::<f32> {
            phantom0: Default::default(),
        };
        let o_signal = vec![
            10, 20, 30, 40, 50, 60, 70, 80, 52, 63, 13, 255, 63, 42, 32, 12, 52, 54, 23, 125, 23,
            255, 43, 23, 123, 54, 34, 255, 255, 23, 255, 32, 13, 15, 65, 23, 5, 7, 7, 3, 9,
        ]
        .iter()
        .map(|&x| x as f32)
        .collect::<Vec<_>>();

        let mut approx: Vec<f32> = vec![0.; (o_signal.len() + 1) / 2];
        let mut details: Vec<f32> = vec![0.; o_signal.len() / 2];

        let mut restored = vec![0.; o_signal.len()];

        m_cdf97
            .execute_forward(&o_signal, &mut approx, &mut details)
            .unwrap();

        m_cdf97
            .execute_inverse(&approx, &details, &mut restored)
            .unwrap();
        o_signal.iter().zip(restored.iter()).enumerate().for_each(|(idx, (o, re))| {
            assert!((o - re).abs() < 1e-3, "Reconstruction difference should be less than 1e-3, but it's not for original o {o}, restored {re} at idx {idx}");
        });
    }

    #[test]
    fn test_cdf97_1() {
        let m_cdf97 = Cdf97::<f32> {
            phantom0: Default::default(),
        };
        let o_signal = vec![
            10i16, 20, 30, 40, 50, 60, 70, 80, 52, 63, 13, 255, 63, 42, 32, 12, 52, 54, 23, 125,
            23, 255, 43, 23, 123, 54, 34, 255, 255, 23, 255, 32, 13, 15, 65, 23, 5, 7, 7, 3, 9, 1,
        ]
        .iter()
        .map(|&x| x as f32)
        .collect::<Vec<_>>();

        let mut approx: Vec<f32> = vec![0.; (o_signal.len() + 1) / 2];
        let mut details: Vec<f32> = vec![0.; o_signal.len() / 2];

        let mut restored = vec![0.; o_signal.len()];

        m_cdf97
            .execute_forward(&o_signal, &mut approx, &mut details)
            .unwrap();

        m_cdf97
            .execute_inverse(&approx, &details, &mut restored)
            .unwrap();
        o_signal.iter().zip(restored.iter()).enumerate().for_each(|(idx, (o, re))| {
            assert!((o - re).abs() < 1e-3, "Reconstruction difference should be less than 1e-3, but it's not for original o {o}, restored {re} at idx {idx}");
        });
    }

    #[test]
    fn test_cdf97_2() {
        let m_cdf97 = Cdf97::<f32> {
            phantom0: Default::default(),
        };
        let o_signal = vec![
            1, 55, 523, 40, 8, 32, 45, 166, 52, 63, 13, 255, 63, 42, 32, 12, 52, 54, 23, 125, 23,
            255, 43, 23, 123, 54, 34, 255, 255, 23, 255, 32, 13, 15, 65, 23, 5, 7, 7, 3, 9, 1,
        ]
        .iter()
        .map(|&x| x as f32)
        .collect::<Vec<_>>();

        let mut approx: Vec<f32> = vec![0.; (o_signal.len() + 1) / 2];
        let mut details: Vec<f32> = vec![0.; o_signal.len() / 2];

        let mut restored = vec![0.; o_signal.len()];

        m_cdf97
            .execute_forward(&o_signal, &mut approx, &mut details)
            .unwrap();

        m_cdf97
            .execute_inverse(&approx, &details, &mut restored)
            .unwrap();
        o_signal.iter().zip(restored.iter()).enumerate().for_each(|(idx, (o, re))| {
            assert!((o - re).abs() < 1e-3, "Reconstruction difference should be less than 1e-3, but it's not for original o {o}, restored {re} at idx {idx}");
        });
    }
}
