// Copyright (c) 2018-2023, The rav1e contributors. All rights reserved

//

// This source code is subject to the terms of the BSD 2 Clause License and

// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License

// was not distributed with this source code in the LICENSE file, you can

// obtain it at www.aomedia.org/license/software. If the Alliance for Open

// Media Patent License 1.0 was not distributed with this source code in the

// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

use crate::context::*;

use crate::header::PRIMARY_REF_NONE;

use crate::partition::BlockSize;

use crate::rdo::spatiotemporal_scale;

use crate::rdo::DistortionScale;

use crate::tiling::TileStateMut;

use crate::util::Pixel;

use crate::FrameInvariants;

use crate::FrameState;

pub const MAX_SEGMENTS: usize = 8;

#[profiling::function]

Unexecuted instantiation: rav1e::segmentation::segmentation_optimize::<u16>

Unexecuted instantiation: rav1e::segmentation::segmentation_optimize::<u8>

pub fn segmentation_optimize<T: Pixel>(

  fi: &FrameInvariants<T>, fs: &mut FrameState<T>,

) {

  assert!(fi.enable_segmentation);

  fs.segmentation.enabled = true;

  if fs.segmentation.enabled {

    fs.segmentation.update_map = true;

    // We don't change the values between frames.

    fs.segmentation.update_data = fi.primary_ref_frame == PRIMARY_REF_NONE;

    // Avoid going into lossless mode by never bringing qidx below 1.

    // Because base_q_idx changes more frequently than the segmentation

    // data, it is still possible for a segment to enter lossless, so

    // enforcement elsewhere is needed.

    let offset_lower_limit = 1 - fi.base_q_idx as i16;

    if !fs.segmentation.update_data {

      let mut min_segment = MAX_SEGMENTS;

      for i in 0..MAX_SEGMENTS {

        if fs.segmentation.features[i][SegLvl::SEG_LVL_ALT_Q as usize]

          && fs.segmentation.data[i][SegLvl::SEG_LVL_ALT_Q as usize]

            >= offset_lower_limit

        {

          min_segment = i;

          break;

        }

      }

      assert_ne!(min_segment, MAX_SEGMENTS);

      fs.segmentation.min_segment = min_segment as u8;

      fs.segmentation.update_threshold(fi.base_q_idx, fi.config.bit_depth);

      return;

    }

    segmentation_optimize_inner(fi, fs, offset_lower_limit);

    /* Figure out parameters */

    fs.segmentation.preskip = false;

    fs.segmentation.last_active_segid = 0;

    for i in 0..MAX_SEGMENTS {

      for j in 0..SegLvl::SEG_LVL_MAX as usize {

        if fs.segmentation.features[i][j] {

          fs.segmentation.last_active_segid = i as u8;

          if j >= SegLvl::SEG_LVL_REF_FRAME as usize {

            fs.segmentation.preskip = true;

          }

        }

      }

    }

  }

}

// Select target quantizers for each segment by fitting to log(scale).

fn segmentation_optimize_inner<T: Pixel>(

  fi: &FrameInvariants<T>, fs: &mut FrameState<T>, offset_lower_limit: i16,

) {

  use crate::quantize::{ac_q, select_ac_qi};

  use crate::util::kmeans;

  use arrayvec::ArrayVec;

  // Minimize the total distance from a small set of values to all scales.

  // Find k-means of log(spatiotemporal scale), k in 3..=8

  let c: ([_; 8], [_; 7], [_; 6], [_; 5], [_; 4], [_; 3]) = {

    let spatiotemporal_scores =

      &fi.coded_frame_data.as_ref().unwrap().spatiotemporal_scores;

    let mut log2_scale_q11 = Vec::with_capacity(spatiotemporal_scores.len());

    log2_scale_q11.extend(spatiotemporal_scores.iter().map(|&s| s.blog16()));

Unexecuted instantiation: rav1e::segmentation::segmentation_optimize_inner::<u16>::{closure#0}

Unexecuted instantiation: rav1e::segmentation::segmentation_optimize_inner::<u8>::{closure#0}

    log2_scale_q11.sort_unstable();

    let l = &log2_scale_q11;

    (kmeans(l), kmeans(l), kmeans(l), kmeans(l), kmeans(l), kmeans(l))

  };

  // Find variance in spacing between successive log(scale)

  let var = |c: &[i16]| {

    let delta = ArrayVec::<_, MAX_SEGMENTS>::from_iter(

      c.iter().skip(1).zip(c).map(|(&a, &b)| b as i64 - a as i64),

Unexecuted instantiation: rav1e::segmentation::segmentation_optimize_inner::<u16>::{closure#1}::{closure#0}

Unexecuted instantiation: rav1e::segmentation::segmentation_optimize_inner::<u8>::{closure#1}::{closure#0}

    );

    let mean = delta.iter().sum::<i64>() / delta.len() as i64;

    delta.iter().map(|&d| (d - mean).pow(2)).sum::<i64>() as u64

Unexecuted instantiation: rav1e::segmentation::segmentation_optimize_inner::<u16>::{closure#1}::{closure#1}

Unexecuted instantiation: rav1e::segmentation::segmentation_optimize_inner::<u8>::{closure#1}::{closure#1}

  };

Unexecuted instantiation: rav1e::segmentation::segmentation_optimize_inner::<u16>::{closure#1}

Unexecuted instantiation: rav1e::segmentation::segmentation_optimize_inner::<u8>::{closure#1}

  let variance =

    [var(&c.0), var(&c.1), var(&c.2), var(&c.3), var(&c.4), var(&c.5)];

  // Choose the k value with minimal variance in spacing

  let min_variance = *variance.iter().min().unwrap();

  let position = variance.iter().rposition(|&v| v == min_variance).unwrap();

Unexecuted instantiation: rav1e::segmentation::segmentation_optimize_inner::<u16>::{closure#2}

Unexecuted instantiation: rav1e::segmentation::segmentation_optimize_inner::<u8>::{closure#2}

  // For the selected centroids, derive a target quantizer:

  //   scale Q'^2 = Q^2

  // See `distortion_scale_for` for more information.

  let compute_delta = |centroids: &[i16]| {

    use crate::util::{bexp64, blog64};

    let log2_base_ac_q_q57 =

      blog64(ac_q(fi.base_q_idx, 0, fi.config.bit_depth).get().into());

    centroids

      .iter()

      .rev()

      // Rewrite in log form and exponentiate:

      //   scale Q'^2 = Q^2

      //           Q' = Q / sqrt(scale)

      //      log(Q') = log(Q) - 0.5 log(scale)

      .map(|&log2_scale_q11| {

        bexp64(log2_base_ac_q_q57 - ((log2_scale_q11 as i64) << (57 - 11 - 1)))

      })

Unexecuted instantiation: rav1e::segmentation::segmentation_optimize_inner::<u16>::{closure#3}::{closure#0}

Unexecuted instantiation: rav1e::segmentation::segmentation_optimize_inner::<u8>::{closure#3}::{closure#0}

      // Find the index of the nearest quantizer to the target,

      // and take the delta from the base quantizer index.

      .map(|q| {

        // Avoid going into lossless mode by never bringing qidx below 1.

        select_ac_qi(q, fi.config.bit_depth).max(1) as i16

          - fi.base_q_idx as i16

      })

Unexecuted instantiation: rav1e::segmentation::segmentation_optimize_inner::<u16>::{closure#3}::{closure#1}

Unexecuted instantiation: rav1e::segmentation::segmentation_optimize_inner::<u8>::{closure#3}::{closure#1}

      .collect::<ArrayVec<_, MAX_SEGMENTS>>()

  };

Unexecuted instantiation: rav1e::segmentation::segmentation_optimize_inner::<u16>::{closure#3}

Unexecuted instantiation: rav1e::segmentation::segmentation_optimize_inner::<u8>::{closure#3}

  // Compute segment deltas for best value of k

  let seg_delta = match position {

    0 => compute_delta(&c.0),

    1 => compute_delta(&c.1),

    2 => compute_delta(&c.2),

    3 => compute_delta(&c.3),

    4 => compute_delta(&c.4),

    _ => compute_delta(&c.5),

  };

  // Update the segmentation data

  fs.segmentation.min_segment = 0;

  fs.segmentation.max_segment = seg_delta.len() as u8 - 1;

  for (&delta, (features, data)) in seg_delta

    .iter()

    .zip(fs.segmentation.features.iter_mut().zip(&mut fs.segmentation.data))

  {

    features[SegLvl::SEG_LVL_ALT_Q as usize] = true;

    data[SegLvl::SEG_LVL_ALT_Q as usize] = delta.max(offset_lower_limit);

  }

  fs.segmentation.update_threshold(fi.base_q_idx, fi.config.bit_depth);

}

Unexecuted instantiation: rav1e::segmentation::segmentation_optimize_inner::<u16>

Unexecuted instantiation: rav1e::segmentation::segmentation_optimize_inner::<u8>

#[profiling::function]

Unexecuted instantiation: rav1e::segmentation::select_segment::<u16>

Unexecuted instantiation: rav1e::segmentation::select_segment::<u8>

pub fn select_segment<T: Pixel>(

  fi: &FrameInvariants<T>, ts: &TileStateMut<'_, T>, tile_bo: TileBlockOffset,

  bsize: BlockSize, skip: bool,

) -> std::ops::RangeInclusive<u8> {

  // If skip is true or segmentation is turned off, sidx is not coded.

  if skip || !fi.enable_segmentation {

    return 0..=0;

  }

  use crate::api::SegmentationLevel;

  if fi.config.speed_settings.segmentation == SegmentationLevel::Full {

    return ts.segmentation.min_segment..=ts.segmentation.max_segment;

  }

  let frame_bo = ts.to_frame_block_offset(tile_bo);

  let scale = spatiotemporal_scale(fi, frame_bo, bsize);

  let sidx = segment_idx_from_distortion(&ts.segmentation.threshold, scale);

  // Avoid going into lossless mode by never bringing qidx below 1.

  let sidx = sidx.max(ts.segmentation.min_segment);

  if fi.config.speed_settings.segmentation == SegmentationLevel::Complex {

    return sidx..=ts.segmentation.max_segment.min(sidx.saturating_add(1));

  }

  sidx..=sidx

}

fn segment_idx_from_distortion(

  threshold: &[DistortionScale; MAX_SEGMENTS - 1], s: DistortionScale,

) -> u8 {

  threshold.partition_point(|&t| s.0 < t.0) as u8

}