/*
 * // 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::WaveletSample;
use crate::err::{OscletError, try_vec};
use crate::{
    Dwt, DwtExecutor, DwtForwardExecutor, DwtInverseExecutor, DwtSize, IncompleteDwtExecutor,
    MultiDwt,
};
use num_traits::AsPrimitive;
use std::marker::PhantomData;

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

impl<T: WaveletSample> DwtForwardExecutor<T> for Cdf53<T>
where
    f64: AsPrimitive<T>,
{
    /// Perform the forward CDF 5/3 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 < 3 {
            return Err(OscletError::MinFilterSize(n, 3));
        }
        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));
        }

        details[0] = fmla(input[0] + input[1], (-0.5f64).as_(), input[0]);

        // Split into even (approx) and odd (detail)
        // Predict step: d[i] = odd - floor((even_left + even_right)/2)
        let signal_n = &input[1..];
        let signal_n2 = &input[2..];

        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, (-0.5f64).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, (-0.5f64).as_(), input[i]);
        }

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

        // Update step: s[i] = even + floor((d_left + d_right + 2)/4)
        approx[0] = fmla(details[0] + details[0], 0.25f64.as_(), 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, 0.25f64.as_(), *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, 0.25f64.as_(), approx[i]);
        }
        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
    }

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

impl<T: WaveletSample> DwtInverseExecutor<T> for Cdf53<T>
where
    f64: AsPrimitive<T>,
{
    /// Perform the inverse CDF 5/3 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 < 3 {
            return Err(OscletError::MinFilterSize(n, 3));
        }

        // Inverse update step: s[i] = even - floor((d_left + d_right + 2)/4)
        let mut approx_inv = approx.to_vec();
        approx_inv[0] = fmla(details[0] + details[0], (-0.25f64).as_(), approx_inv[0]);
        let approx_k = &mut approx_inv[1..];
        let detail_next = &details[1..];

        let mut iters = 1usize;

        for ((dst, detail), detail_next) in approx_k
            .iter_mut()
            .zip(details.iter())
            .zip(detail_next.iter())
        {
            let d_left = *detail;
            let d_right = *detail_next;
            *dst = fmla(d_left + d_right, (-0.25f64).as_(), *dst);
            iters += 1;
        }

        if iters < approx.len() {
            let i = approx.len() - 1;
            let d_left = details[i - 1];
            let d_right = if i < details.len() {
                details[i]
            } else {
                details[details.len() - 1]
            };
            approx_inv[i] = fmla(d_left + d_right, (-0.25f64).as_(), approx_inv[i]);
        }

        // Inverse predict step: odd = detail + floor((even_left + even_right)/2)
        let approx_next = &approx_inv[1..];
        let mut odd_values: Vec<T> = try_vec![T::default(); details.len()];

        for (((dst, &current), &next), &detail) in odd_values
            .iter_mut()
            .zip(approx_inv.iter())
            .zip(approx_next.iter())
            .zip(details.iter())
        {
            let left = current;
            let right = next;
            *dst = fmla(left + right, 0.5f64.as_(), detail);
        }

        if !details.is_empty() {
            let i = details.len() - 1;
            let left = approx_inv[i];
            let right = if i + 1 < approx_inv.len() {
                approx_inv[i + 1]
            } else {
                left
            };
            odd_values[i] = fmla(left + right, 0.5f64.as_(), details[i]);
        }

        // Interleave even and odd
        for ((dst, &even), &odd) in output
            .chunks_exact_mut(2)
            .zip(approx_inv.iter())
            .zip(odd_values.iter())
        {
            dst[0] = even;
            dst[1] = odd;
        }
        let chunks = output.chunks_exact_mut(2).into_remainder();
        if let Some(dst) = chunks.get_mut(0) {
            // Only one even value left at the end
            *dst = *approx_inv.last().unwrap();
        }

        Ok(())
    }

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

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

macro_rules! define_dwt_cdf_float {
    ($for_clazz: ident, $min_size: expr) => {
        use crate::DwtRef;
        use crate::MultiLevelDwtRef;

        impl<
            T: WaveletSample,
        > DwtExecutor<T> for $for_clazz<T>
        where
            f64: AsPrimitive<T>,
        {
            fn dwt(&self, signal: &[T], level: usize) -> Result<Dwt<T>, OscletError> {
                if signal.len() < $min_size {
                    return Err(OscletError::BufferWasTooSmallForLevel);
                }
                if level == 0 || level == 1 {
                    let mut approx = try_vec![T::default(); signal.len().div_ceil(2)];
                    let mut details = try_vec![T::default(); signal.len() / 2];

                    self.execute_forward(signal, &mut approx, &mut details)?;

                    Ok(Dwt {
                        approximations: approx,
                        details,
                    })
                } else {
                    let mut current_signal = signal.to_vec();
                    let mut approx = vec![];
                    let mut details = vec![];

                    for _ in 0..level {
                        if current_signal.len() < $min_size {
                            return Err(OscletError::BufferWasTooSmallForLevel);
                        }

                        approx = try_vec![T::default(); current_signal.len().div_ceil(2)];
                        details = try_vec![T::default(); current_signal.len() / 2];

                        // Forward DWT on current signal
                        self.execute_forward(&current_signal, &mut approx, &mut details)?;

                        // Next level uses only the approximation
                        current_signal = approx.to_vec();
                    }

                    Ok(Dwt {
                        approximations: approx,
                        details,
                    })
                }
            }

            fn multi_dwt(&self, signal: &[T], levels: usize) -> Result<MultiDwt<T>, OscletError> {
                if signal.len() < $min_size {
                    return Err(OscletError::BufferWasTooSmallForLevel);
                }
                if levels == 0 || levels == 1 {
                    let mut approx = try_vec![T::default(); signal.len().div_ceil(2)];
                    let mut details = try_vec![T::default(); signal.len() / 2];

                    self.execute_forward(signal, &mut approx, &mut details)?;

                    return Ok(MultiDwt {
                        levels: vec![Dwt {
                            approximations: approx,
                            details,
                        }],
                    });
                }

                let mut current_signal = signal.to_vec();
                let mut approx = try_vec![T::default(); signal.len().div_ceil(2)];
                let mut levels_store = Vec::with_capacity(levels);

                for _ in 0..levels {
                    if current_signal.len() < $min_size {
                        return Err(OscletError::BufferWasTooSmallForLevel);
                    }

                    let approx_len = current_signal.len().div_ceil(2);
                    let details_len = current_signal.len() / 2;

                    approx.resize(approx_len, T::default());
                    let mut details = try_vec![T::default(); details_len];

                    self.execute_forward(&current_signal, &mut approx, &mut details)?;

                    std::mem::swap(&mut current_signal, &mut approx);

                    levels_store.push(Dwt {
                        approximations: current_signal.clone(),
                        details,
                    });
                }

                Ok(MultiDwt {
                    levels: levels_store,
                })
            }

            fn idwt(&self, dwt: &DwtRef<'_, T>) -> Result<Vec<T>, OscletError> {
                let mut output = try_vec![T::default(); dwt.details.len() + dwt.approximations.len()];
                self.execute_inverse(&dwt.approximations, &dwt.details, &mut output)?;
                Ok(output)
            }

            fn multi_idwt(&self, dwt: &MultiLevelDwtRef<'_, T>) -> Result<Vec<T>, OscletError> {
                  if dwt.details.is_empty() || dwt.approximations.is_empty() {
                    return Err(OscletError::ZeroedBaseSize);
                }
                let mut current_approximations = dwt.approximations;
                let mut output = self.idwt(&DwtRef {
                    approximations: current_approximations,
                    details: dwt.details.last().unwrap(),
                })?;
                if dwt.details.len() == 1 {
                    return Ok(output);
                }
                let details_remainder = &dwt.details[..dwt.details.len() - 1];
                current_approximations = &output;
                for details in details_remainder.iter().rev() {
                    output = self.idwt(&DwtRef {
                        approximations: current_approximations,
                        details,
                    })?;
                    current_approximations = &output;
                }

                Ok(output)
            }
        }
    };
}

use crate::mla::fmla;
pub(crate) use define_dwt_cdf_float;

define_dwt_cdf_float!(Cdf53, 3);

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

    #[test]
    fn test_cdf53() {
        let i16_cdf53 = Cdf53::<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()];

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

        i16_cdf53
            .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_cdf53_1() {
        let i16_cdf53 = Cdf53::<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()];

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

        i16_cdf53
            .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_cdf53_2() {
        let i16_cdf53 = Cdf53::<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()];

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

        i16_cdf53
            .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_db4_even_big_mdwt() {
        let data_length = 86;
        let mut input = vec![0.; data_length];
        for i in 0..data_length {
            input[i] = i as f32 / data_length as f32;
        }
        let i16_cdf53 = Cdf53::<f32> {
            phantom0: Default::default(),
        };
        _ = i16_cdf53.multi_dwt(&input, 4).unwrap();
    }

    #[test]
    fn test_cdf97_multilevel() {
        let m_cdf97 = Cdf53::<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<_>>();

        for levels in 1..=4 {
            let multi = m_cdf97.multi_dwt(&o_signal, levels).unwrap();

            let refs: Vec<&[f32]> = multi.levels.iter().map(|d| d.details.as_slice()).collect();
            let restored = m_cdf97
                .multi_idwt(&MultiLevelDwtRef {
                    approximations: multi.levels.last().unwrap().approximations.as_slice(),
                    details: refs,
                })
                .unwrap();

            assert_eq!(
                restored.len(),
                o_signal.len(),
                "levels={levels}: restored length mismatch"
            );
            o_signal
                .iter()
                .zip(restored.iter())
                .enumerate()
                .for_each(|(idx, (o, re))| {
                    assert!(
                        (o - re).abs() < 1e-3,
                        "levels={levels} idx={idx}: original={o}, restored={re}, diff={}",
                        (o - re).abs()
                    );
                });
        }
    }
}
