j2k-native 0.13.1

// SPDX-License-Identifier: MIT OR Apache-2.0

use alloc::vec::Vec;

use super::super::build::Decomposition;
use super::super::codestream::WaveletTransform;
use super::super::decode::DecompositionStorage;
use super::super::rect::IntRect;
use super::banded::{filter_2d_banded, filter_2d_banded_from_bands, packed_band_rects, use_banded};
use super::horizontal::{filter_horizontal, filter_horizontal_codestream};
use super::interleave::{interleave_samples, interleave_samples_into};
use super::model::{CoefficientSource, IDWTInput, IDWTTempOutput};
use super::roi::interleave_samples_roi;
use super::vertical::{filter_vertical, filter_vertical_codestream};
use crate::error::{bail, DecodingError};
use crate::{
    checked_decode_usize_product2, try_resize_decode_elements, HtCodeBlockDecoder, J2kIdwtBand,
    J2kIdwtNormalization, J2kRect, J2kSingleDecompositionIdwtJob, J2kWaveletTransform, Result,
};

#[expect(
    clippy::too_many_arguments,
    reason = "the level driver keeps both IDWT buffers, the band scratch, and the backend hook explicit"
)]
pub(super) fn apply_level(
    input: IDWTInput<'_>,
    target: &mut Vec<f32>,
    band_scratch: &mut Vec<f32>,
    decomposition: &Decomposition,
    transform: WaveletTransform,
    storage: &DecompositionStorage<'_>,
    backend: &mut Option<&mut dyn HtCodeBlockDecoder>,
    parallel: bool,
) -> Result<IDWTTempOutput> {
    let handled = if let Some(backend) = backend.as_deref_mut() {
        let required_len = checked_decode_usize_product2(
            decomposition.rect.width() as usize,
            decomposition.rect.height() as usize,
        )?;
        try_resize_decode_elements(target, required_len, 0.0)?;
        let job = single_decomposition_job(input, decomposition, storage, transform);
        let normalization = match transform {
            WaveletTransform::Reversible53 => J2kIdwtNormalization::Standard,
            WaveletTransform::Irreversible97 => J2kIdwtNormalization::OpenJpegCodestream,
        };
        backend
            .decode_single_decomposition_idwt_with_normalization(job, normalization, target)
            .map_err(|_| DecodingError::CodeBlockDecodeFailure)?
    } else {
        false
    };

    if handled {
        Ok(IDWTTempOutput {
            rect: decomposition.rect,
        })
    } else if use_banded(
        decomposition.rect.width() as usize,
        decomposition.rect.height() as usize,
    ) {
        filter_2d_banded(
            input,
            decomposition,
            target,
            band_scratch,
            transform,
            storage,
            parallel,
        )
    } else {
        filter_2d(input, target, decomposition, transform, storage)
    }
}

/// CPU reconstruction into a previous component's initialized dead prefix.
/// Keep this separate from the ordinary vector and accelerator path.
pub(super) fn apply_level_into(
    input: IDWTInput<'_>,
    target: &mut [f32],
    band_scratch: &mut Vec<f32>,
    decomposition: &Decomposition,
    transform: WaveletTransform,
    storage: &DecompositionStorage<'_>,
    parallel: bool,
) -> Result<IDWTTempOutput> {
    let rect = decomposition.rect;
    let len = checked_decode_usize_product2(rect.width() as usize, rect.height() as usize)?;
    let target = &mut target[..len];
    if use_banded(rect.width() as usize, rect.height() as usize) {
        let band = |index: usize| {
            &storage.coefficients[storage.sub_bands[decomposition.sub_bands[index]]
                .coefficients
                .clone()]
        };
        filter_2d_banded_from_bands(
            rect,
            [input.coefficients, band(0), band(1), band(2)],
            target,
            band_scratch,
            transform,
            super::OPENJPEG_NORMALIZED_HIGH_PASS_F32,
            parallel,
        )?;
    } else {
        interleave_samples_into(input, decomposition, target, storage);
        if rect.width() > 0 && rect.height() > 0 {
            filter_horizontal_codestream(target, rect, transform);
            filter_vertical_codestream(target, rect, transform);
        }
    }
    Ok(IDWTTempOutput { rect })
}

pub(crate) fn apply_single_decomposition_idwt_job(
    job: J2kSingleDecompositionIdwtJob<'_>,
    target: &mut Vec<f32>,
) -> Result<()> {
    apply_single_decomposition_idwt_job_with_normalization(job, target, false, false)
}

pub(crate) fn apply_codestream_single_decomposition_idwt_job(
    job: J2kSingleDecompositionIdwtJob<'_>,
    target: &mut Vec<f32>,
    parallel: bool,
) -> Result<()> {
    apply_single_decomposition_idwt_job_with_normalization(job, target, true, parallel)
}

fn apply_single_decomposition_idwt_job_with_normalization(
    job: J2kSingleDecompositionIdwtJob<'_>,
    target: &mut Vec<f32>,
    openjpeg_codestream: bool,
    parallel: bool,
) -> Result<()> {
    let rect = int_rect_from_public(job.rect);
    validate_direct_band(job.ll)?;
    validate_direct_band(job.hl)?;
    validate_direct_band(job.lh)?;
    validate_direct_band(job.hh)?;

    target.clear();
    let required_len =
        checked_decode_usize_product2(rect.width() as usize, rect.height() as usize)?;
    try_resize_decode_elements(target, required_len, 0.0)?;

    let transform = wavelet_transform_from_public(job.transform);
    if use_banded(rect.width() as usize, rect.height() as usize) {
        let bands = [job.ll, job.hl, job.lh, job.hh];
        let packed = packed_band_rects(rect);
        if bands
            .iter()
            .zip(packed)
            .all(|(band, packed)| int_rect_from_public(band.rect) == packed)
        {
            // Every band is exactly the packed sub-band of this level, so
            // the windowed interleave below reads the same samples the
            // banded reconstruction does.
            let high_pass = if openjpeg_codestream {
                super::OPENJPEG_NORMALIZED_HIGH_PASS_F32
            } else {
                j2k_codec_math::dwt::DWT97_INV_KAPPA_F32
            };
            let mut band_scratch = Vec::new();
            return filter_2d_banded_from_bands(
                rect,
                bands.map(|band| band.coefficients),
                target,
                &mut band_scratch,
                transform,
                high_pass,
                parallel,
            );
        }
    }
    interleave_samples_roi(
        direct_coefficient_source(job.ll),
        direct_coefficient_source(job.hl),
        direct_coefficient_source(job.lh),
        direct_coefficient_source(job.hh),
        target,
        rect,
    );
    if rect.width() > 0 && rect.height() > 0 {
        if openjpeg_codestream && transform == WaveletTransform::Irreversible97 {
            filter_horizontal_codestream(target, rect, transform);
            filter_vertical_codestream(target, rect, transform);
        } else {
            filter_horizontal(target, rect, transform);
            filter_vertical(target, rect, transform);
        }
    }
    Ok(())
}

fn validate_direct_band(band: J2kIdwtBand<'_>) -> Result<()> {
    let rect = int_rect_from_public(band.rect);
    let required_len = rect
        .width()
        .checked_mul(rect.height())
        .and_then(|len| usize::try_from(len).ok())
        .ok_or(DecodingError::CodeBlockDecodeFailure)?;
    if band.coefficients.len() < required_len {
        bail!(DecodingError::CodeBlockDecodeFailure);
    }
    Ok(())
}

fn direct_coefficient_source(band: J2kIdwtBand<'_>) -> CoefficientSource<'_> {
    let rect = int_rect_from_public(band.rect);
    CoefficientSource::new(band.coefficients, rect, rect.width())
}

fn int_rect_from_public(rect: J2kRect) -> IntRect {
    IntRect::from_ltrb(rect.x0, rect.y0, rect.x1, rect.y1)
}

fn wavelet_transform_from_public(transform: J2kWaveletTransform) -> WaveletTransform {
    match transform {
        J2kWaveletTransform::Reversible53 => WaveletTransform::Reversible53,
        J2kWaveletTransform::Irreversible97 => WaveletTransform::Irreversible97,
    }
}

fn single_decomposition_job<'a>(
    input: IDWTInput<'a>,
    decomposition: &'a Decomposition,
    storage: &'a DecompositionStorage<'a>,
    transform: WaveletTransform,
) -> J2kSingleDecompositionIdwtJob<'a> {
    let hl = &storage.sub_bands[decomposition.sub_bands[0]];
    let lh = &storage.sub_bands[decomposition.sub_bands[1]];
    let hh = &storage.sub_bands[decomposition.sub_bands[2]];
    J2kSingleDecompositionIdwtJob {
        rect: J2kRect::from(decomposition.rect),
        transform: J2kWaveletTransform::from(transform),
        ll: J2kIdwtBand {
            rect: J2kRect::from(input.rect),
            coefficients: input.coefficients,
        },
        hl: J2kIdwtBand {
            rect: J2kRect::from(hl.rect),
            coefficients: &storage.coefficients[hl.coefficients.clone()],
        },
        lh: J2kIdwtBand {
            rect: J2kRect::from(lh.rect),
            coefficients: &storage.coefficients[lh.coefficients.clone()],
        },
        hh: J2kIdwtBand {
            rect: J2kRect::from(hh.rect),
            coefficients: &storage.coefficients[hh.coefficients.clone()],
        },
    }
}

/// The `2D_SR` procedure illustrated in Figure F.6.
pub(super) fn filter_2d(
    // The LL sub band of the given decomposition level.
    input: IDWTInput<'_>,
    coefficients: &mut Vec<f32>,
    decomposition: &Decomposition,
    transform: WaveletTransform,
    storage: &DecompositionStorage<'_>,
) -> Result<IDWTTempOutput> {
    // First interleave all sub-bands into a single buffer.
    interleave_samples(input, decomposition, coefficients, storage)?;

    if decomposition.rect.width() > 0 && decomposition.rect.height() > 0 {
        // The buffer may hold more samples than this level; filter only its own.
        let level_len = checked_decode_usize_product2(
            decomposition.rect.width() as usize,
            decomposition.rect.height() as usize,
        )?;
        let coefficients = &mut coefficients[..level_len];
        filter_horizontal_codestream(coefficients, decomposition.rect, transform);
        filter_vertical_codestream(coefficients, decomposition.rect, transform);
    }

    Ok(IDWTTempOutput {
        rect: decomposition.rect,
    })
}