zoomvtools 1.1.1

use std::num::NonZeroUsize;

use anyhow::{Result, anyhow, bail};

use crate::{
    analysis::MVAnalysisData,
    dct::DctHelper,
    fake::group_of_planes::FakeGroupOfPlanes,
    filters::analyse::SuperClipInfo,
    frame::FrameView,
    group_of_planes::GroupOfPlanes,
    mv_gof::MVGroupOfFrames,
    params::{DctMode, DivideMode, MVPlaneSet, MotionFlags, SearchType, Subpel},
    plane_of_blocks::MvsOutput,
    util::Pixel,
    video::{ColorFamily, SampleType, VideoInfo},
};

#[cfg(test)]
mod tests;

/// Options used to construct [`Recalculate`].
///
/// `None` selects the MVTools-compatible default for optional fields.
#[derive(Debug, Clone, Copy)]
pub struct RecalculateOptions {
    /// SAD threshold above which vectors are recalculated.
    pub th_sad: Option<u64>,
    /// Enables temporal smoothing of the refined vectors.
    pub smooth: Option<bool>,
    /// Horizontal block size in pixels.
    pub blk_size_x: Option<usize>,
    /// Vertical block size in pixels.
    pub blk_size_y: Option<usize>,
    /// Search strategy for refinement.
    pub search_type: Option<SearchType>,
    /// Search radius or step count for the selected search strategy.
    pub search_param: Option<i32>,
    /// Motion-vector penalty scale for refinement.
    pub lambda: Option<u32>,
    /// Enables chroma-guided refinement when `true`.
    pub chroma: Option<bool>,
    /// Enables MVTools true-motion defaults.
    pub truemotion: bool,
    /// Penalty applied to new candidate vectors.
    pub penalty_new: Option<u16>,
    /// Horizontal overlap between neighbouring blocks.
    pub overlap_x: Option<usize>,
    /// Vertical overlap between neighbouring blocks.
    pub overlap_y: Option<usize>,
    /// Optional extra block subdivision mode.
    pub divide_extra: Option<DivideMode>,
    /// Enables meander block traversal.
    pub meander: bool,
    /// Treats the clip as field-based.
    pub fields: bool,
    /// DCT or SATD mode used for block matching.
    pub dct_mode: Option<DctMode>,
}

impl Default for RecalculateOptions {
    #[inline]
    fn default() -> Self {
        Self {
            th_sad: None,
            smooth: None,
            blk_size_x: None,
            blk_size_y: None,
            search_type: None,
            search_param: None,
            lambda: None,
            chroma: None,
            truemotion: true,
            penalty_new: None,
            overlap_x: None,
            overlap_y: None,
            divide_extra: None,
            meander: true,
            fields: false,
            dct_mode: None,
        }
    }
}

/// Motion-vector refinement filter state.
#[derive(Debug, Clone)]
pub struct Recalculate {
    th_sad: u64,
    smooth: bool,
    search_type: SearchType,
    search_param: i32,
    lambda: u32,
    penalty_new: u16,
    dct_mode: DctMode,
    divide_extra: DivideMode,
    meander: bool,
    fields: bool,
    super_info: SuperClipInfo,
    vectors_data: MVAnalysisData,
    analysis_data: MVAnalysisData,
    analysis_data_divided: Option<MVAnalysisData>,
    plane_count: usize,
    frame_count: usize,
}

/// Output of [`Recalculate::recalculate_frame`].
#[derive(Debug, Clone)]
pub struct RecalculateResult {
    /// Metadata describing the returned vector stream.
    pub analysis_data: MVAnalysisData,
    /// Packed vectors for the refined frame.
    pub vectors: MvsOutput,
}

impl Recalculate {
    /// Builds a recalculation filter from clip, super, and vector metadata.
    ///
    /// # Errors
    /// Returns an error if the clip format, super metadata, source vectors, or options are
    /// invalid.
    #[inline]
    pub fn new(
        info: VideoInfo,
        super_info: SuperClipInfo,
        vectors_data: MVAnalysisData,
        options: RecalculateOptions,
    ) -> Result<Self> {
        let blk_size_x = options.blk_size_x.unwrap_or(8);
        let blk_size_y = options.blk_size_y.unwrap_or(blk_size_x);
        let search_type = options.search_type.unwrap_or(SearchType::Hex2);
        let mut search_param = options.search_param.unwrap_or(2);
        let chroma = options.chroma.unwrap_or(true);
        let lambda = options.lambda.unwrap_or(if options.truemotion {
            (1000 * blk_size_x * blk_size_y / 64) as u32
        } else {
            0
        });
        let penalty_new = options
            .penalty_new
            .unwrap_or(if options.truemotion { 50 } else { 0 });
        let overlap_x = options.overlap_x.unwrap_or(0);
        let overlap_y = options.overlap_y.unwrap_or(overlap_x);
        let dct_mode = options.dct_mode.unwrap_or(DctMode::Spatial);
        let divide_extra = options.divide_extra.unwrap_or(DivideMode::None);

        if dct_mode.uses_satd() && blk_size_x == 16 && blk_size_y == 2 {
            bail!("Recalculate: dct 5..10 cannot work with 16x2 blocks.");
        }

        match (blk_size_x, blk_size_y) {
            (4, 4)
            | (8, 4)
            | (8, 8)
            | (16, 2)
            | (16, 8)
            | (16, 16)
            | (32, 16)
            | (32, 32)
            | (64, 32)
            | (64, 64)
            | (128, 64)
            | (128, 128) => (),
            _ => {
                bail!(
                    "Recalculate: the block size must be 4x4, 8x4, 8x8, 16x2, 16x8, 16x16, 32x16, \
                     32x32, 64x32, 64x64, 128x64, or 128x128."
                );
            }
        }

        if penalty_new > 256 {
            bail!("Recalculate: pnew must be between 0 and 256 (inclusive).");
        }
        if overlap_x > blk_size_x / 2 || overlap_y > blk_size_y / 2 {
            bail!(
                "Recalculate: overlap must be at most half of blksize, and overlapv must be at \
                 most half of blksizev"
            );
        }
        if divide_extra != DivideMode::None && (blk_size_x < 8 || blk_size_y < 8) {
            bail!("Recalculate: blksize and blksizev must be at least 8 when divide=True.");
        }

        if search_type == SearchType::Nstep {
            if search_param < 0 {
                search_param = 0;
            }
        } else if search_param < 1 {
            search_param = 1;
        }

        let format = info.format;
        if format.bits_per_sample.get() > 16 {
            bail!("Recalculate: input clip must be 8-16 bits");
        }
        if format.sample_type != SampleType::Integer {
            bail!("Recalculate: input clip must be integer super_format");
        }
        if ![ColorFamily::Yuv, ColorFamily::Gray].contains(&format.color_family)
            || format.sub_sampling_w > 1
            || format.sub_sampling_h > 1
        {
            bail!("Recalculate: input clip must be GRAY, 420, 422, 440, or 444");
        }

        if !overlap_x.is_multiple_of(1 << format.sub_sampling_w)
            || !overlap_y.is_multiple_of(1 << format.sub_sampling_h)
        {
            bail!(
                "Recalculate: The requested overlap is incompatible with the super clip's \
                 subsampling."
            );
        }
        if divide_extra != DivideMode::None
            && (!overlap_x.is_multiple_of(2 << format.sub_sampling_w)
                || !overlap_y.is_multiple_of(2 << format.sub_sampling_h))
        {
            bail!(
                "Recalculate: overlap and overlapv must be multiples of 2 or 4 when divide=True, \
                 depending on the super clip's subsampling."
            );
        }

        let mode_yuv = if chroma {
            MVPlaneSet::YUVPLANES
        } else {
            MVPlaneSet::YPLANE
        };
        if mode_yuv & super_info.mode_yuv != mode_yuv {
            bail!("Recalculate: super clip does not contain needed colour data.");
        }

        let width = NonZeroUsize::new(info.resolution.width)
            .ok_or_else(|| anyhow!("Recalculate: variable resolution clips are not supported."))?;
        let expected_width = width
            .get()
            .checked_sub(super_info.hpad * 2)
            .ok_or_else(|| anyhow!("Recalculate: wrong frame size."))?;
        if super_info.height != vectors_data.height || expected_width != vectors_data.width.get() {
            bail!("Recalculate: wrong frame size.");
        }

        let bits_per_sample = format.bits_per_sample;
        let pixel_max = (1u32 << bits_per_sample.get()) - 1;
        let mut th_sad = options.th_sad.unwrap_or(200);
        th_sad = (th_sad as f64 * pixel_max as f64 / 255.0 + 0.5) as u64;
        let lambda = (lambda as f64 * pixel_max as f64 / 255.0 + 0.5) as u32;

        let reference_block_size = 8 * 8;
        th_sad = th_sad * (blk_size_x * blk_size_y) as u64 / reference_block_size;
        if chroma {
            th_sad +=
                th_sad / (vectors_data.x_ratio_uv.get() * vectors_data.y_ratio_uv.get()) as u64 * 2;
        }

        let mut motion_flags = MotionFlags::empty();
        if vectors_data.is_backward {
            motion_flags |= MotionFlags::IS_BACKWARD;
        }
        if chroma {
            motion_flags |= MotionFlags::USE_CHROMA_MOTION;
        }

        let blk_x = (vectors_data.width.get() - overlap_x) / (blk_size_x - overlap_x);
        let blk_y = (vectors_data.height.get() - overlap_y) / (blk_size_y - overlap_y);
        let analysis_data = MVAnalysisData {
            blk_size_x: NonZeroUsize::new(blk_size_x)
                .ok_or_else(|| anyhow!("Recalculate: blksize must be greater than 0"))?,
            blk_size_y: NonZeroUsize::new(blk_size_y)
                .ok_or_else(|| anyhow!("Recalculate: blksizev must be greater than 0"))?,
            pel: super_info.pel,
            level_count: 1,
            delta_frame: vectors_data.delta_frame,
            is_backward: vectors_data.is_backward,
            motion_flags,
            width: vectors_data.width,
            height: vectors_data.height,
            overlap_x,
            overlap_y,
            blk_x: NonZeroUsize::new(blk_x).expect("block count should not be zero"),
            blk_y: NonZeroUsize::new(blk_y).expect("block count should not be zero"),
            bits_per_sample,
            y_ratio_uv: vectors_data.y_ratio_uv,
            x_ratio_uv: vectors_data.x_ratio_uv,
            h_padding: super_info.hpad,
            v_padding: super_info.vpad,
        };
        let analysis_data_divided = (divide_extra != DivideMode::None).then(|| {
            let two = NonZeroUsize::new(2).expect("non-zero const");
            MVAnalysisData {
                blk_x: analysis_data.blk_x.saturating_mul(two),
                blk_y: analysis_data.blk_y.saturating_mul(two),
                blk_size_x: NonZeroUsize::new(analysis_data.blk_size_x.get() / 2)
                    .expect("cannot be zero"),
                blk_size_y: NonZeroUsize::new(analysis_data.blk_size_y.get() / 2)
                    .expect("cannot be zero"),
                overlap_x: analysis_data.overlap_x / 2,
                overlap_y: analysis_data.overlap_y / 2,
                level_count: analysis_data.level_count + 1,
                ..analysis_data
            }
        });

        Ok(Self {
            th_sad,
            smooth: options.smooth.unwrap_or(true),
            search_type,
            search_param,
            lambda,
            penalty_new,
            dct_mode,
            divide_extra,
            meander: options.meander,
            fields: options.fields,
            super_info,
            vectors_data,
            analysis_data,
            analysis_data_divided,
            plane_count: format.plane_count(),
            frame_count: info.num_frames,
        })
    }

    /// Returns metadata for vectors produced by this recalculator.
    #[must_use]
    #[inline]
    pub const fn analysis_data(&self) -> MVAnalysisData {
        self.analysis_data
    }

    /// Returns metadata for the divided vector stream when division is enabled.
    #[must_use]
    #[inline]
    pub const fn analysis_data_divided(&self) -> Option<MVAnalysisData> {
        self.analysis_data_divided
    }

    /// Returns whether the recalculator is configured for field-based search.
    #[must_use]
    #[inline]
    pub const fn fields(&self) -> bool {
        self.fields
    }

    /// Returns the configured search strategy.
    #[must_use]
    #[inline]
    pub const fn search_type(&self) -> SearchType {
        self.search_type
    }

    /// Returns the validated search parameter.
    #[must_use]
    #[inline]
    pub const fn search_param(&self) -> i32 {
        self.search_param
    }

    /// Returns the motion penalty scale.
    #[must_use]
    #[inline]
    pub const fn lambda(&self) -> u32 {
        self.lambda
    }

    /// Returns the penalty applied to new candidate vectors.
    #[must_use]
    #[inline]
    pub const fn penalty_new(&self) -> u16 {
        self.penalty_new
    }

    /// Returns the configured block-matching metric.
    #[must_use]
    #[inline]
    pub const fn dct_mode(&self) -> DctMode {
        self.dct_mode
    }

    /// Returns the extra block-division mode.
    #[must_use]
    #[inline]
    pub const fn divide_extra(&self) -> DivideMode {
        self.divide_extra
    }

    /// Returns whether temporal smoothing is enabled.
    #[must_use]
    #[inline]
    pub const fn smooth(&self) -> bool {
        self.smooth
    }

    /// Returns whether meander traversal is enabled.
    #[must_use]
    #[inline]
    pub const fn meander(&self) -> bool {
        self.meander
    }

    /// Returns the reference frame index paired with source frame `n`.
    ///
    /// Returns `None` when the reference would fall outside the source clip.
    #[must_use]
    #[inline]
    pub fn reference_frame_number(&self, n: usize) -> Option<usize> {
        let nref = if self.analysis_data.delta_frame > 0 {
            let offset = if self.analysis_data.is_backward {
                self.analysis_data.delta_frame
            } else {
                -self.analysis_data.delta_frame
            };
            n as isize + offset
        } else {
            -self.analysis_data.delta_frame
        };

        (nref >= 0 && (nref as usize) < self.frame_count).then_some(nref as usize)
    }

    /// Recalculates one packed vector frame and returns the refined result.
    ///
    /// # Errors
    /// Returns an error if the frames, field flags, or packed vectors do not match the filter
    /// configuration.
    #[inline]
    pub fn recalculate_frame<T: Pixel>(
        &self,
        src: &FrameView<'_, T>,
        reference: Option<&FrameView<'_, T>>,
        vectors_data: &[u8],
        src_top_field: Option<bool>,
        ref_top_field: Option<bool>,
    ) -> Result<RecalculateResult> {
        let src_top_field = self.resolve_top_field(src_top_field, "source")?;
        let mut fake_gop = FakeGroupOfPlanes::new(&self.vectors_data);
        fake_gop.update(vectors_data)?;

        let mut vector_fields = GroupOfPlanes::<T>::new(
            self.analysis_data.blk_size_x,
            self.analysis_data.blk_size_y,
            self.analysis_data.level_count,
            self.analysis_data.pel,
            self.analysis_data.motion_flags,
            self.analysis_data.overlap_x,
            self.analysis_data.overlap_y,
            self.analysis_data.blk_x,
            self.analysis_data.blk_y,
            self.analysis_data.x_ratio_uv,
            self.analysis_data.y_ratio_uv,
            self.divide_extra,
            self.analysis_data.bits_per_sample,
        )?;

        let vectors = if fake_gop.is_valid() {
            if let Some(reference) = reference {
                let ref_top_field = self.resolve_top_field(ref_top_field, "reference")?;
                let field_shift = self.field_shift(src_top_field, ref_top_field);
                let src_gof = MVGroupOfFrames::new(
                    self.super_info.levels,
                    self.analysis_data.width,
                    self.analysis_data.height,
                    self.super_info.pel,
                    self.super_info.hpad,
                    self.super_info.vpad,
                    self.super_info.mode_yuv,
                    self.analysis_data.x_ratio_uv,
                    self.analysis_data.y_ratio_uv,
                    self.analysis_data.bits_per_sample,
                    src.pitch(),
                    self.plane_count,
                )?;
                let ref_gof = MVGroupOfFrames::new(
                    self.super_info.levels,
                    self.analysis_data.width,
                    self.analysis_data.height,
                    self.super_info.pel,
                    self.super_info.hpad,
                    self.super_info.vpad,
                    self.super_info.mode_yuv,
                    self.analysis_data.x_ratio_uv,
                    self.analysis_data.y_ratio_uv,
                    self.analysis_data.bits_per_sample,
                    reference.pitch(),
                    self.plane_count,
                )?;

                let dct_helper = if matches!(
                    self.dct_mode,
                    DctMode::Dct
                        | DctMode::MixedSpatialDct
                        | DctMode::AdaptiveSpatialMixed
                        | DctMode::AdaptiveSpatialDct
                ) {
                    Some(DctHelper::new(
                        self.analysis_data.blk_size_x,
                        self.analysis_data.blk_size_y,
                        self.analysis_data.bits_per_sample,
                    )?)
                } else {
                    None
                };

                let mut vectors = vector_fields.recalculate_mvs(
                    &fake_gop,
                    &src_gof,
                    src.planes(),
                    &ref_gof,
                    reference.planes(),
                    self.search_type,
                    self.search_param,
                    self.lambda,
                    self.penalty_new,
                    field_shift,
                    self.th_sad,
                    dct_helper,
                    self.dct_mode,
                    self.smooth,
                    self.meander,
                )?;
                if self.divide_extra != DivideMode::None {
                    vector_fields.extra_divide(&mut vectors);
                }
                vectors
            } else {
                vector_fields.write_default_to_array()
            }
        } else {
            vector_fields.write_default_to_array()
        };

        Ok(RecalculateResult {
            analysis_data: self.analysis_data_for_output(),
            vectors,
        })
    }

    #[must_use]
    #[inline]
    fn analysis_data_for_output(&self) -> MVAnalysisData {
        self.analysis_data_divided.unwrap_or(self.analysis_data)
    }

    #[inline]
    fn resolve_top_field(&self, field: Option<bool>, frame_name: &str) -> Result<bool> {
        if self.fields {
            field.ok_or_else(|| {
                anyhow!(
                    "Recalculate: {frame_name} field order is required when fields=true and the wrapper does not provide tff."
                )
            })
        } else {
            Ok(field.unwrap_or(false))
        }
    }

    #[must_use]
    #[inline]
    fn field_shift(&self, src_top_field: bool, ref_top_field: bool) -> i32 {
        if self.fields
            && self.analysis_data.pel > Subpel::Full
            && (self.analysis_data.delta_frame % 2) > 0
        {
            if src_top_field && !ref_top_field {
                (u8::from(self.analysis_data.pel) as u32 / 2) as i32
            } else if ref_top_field && !src_top_field {
                -((u8::from(self.analysis_data.pel) as u32 / 2) as i32)
            } else {
                0
            }
        } else {
            0
        }
    }
}