tools/avm_analyzer/avm_stats/src/stats.rs

use std::collections::{HashMap, HashSet};

use itertools::Itertools;

use ordered_float::OrderedFloat;
use serde::{Deserialize, Serialize};

use crate::{CodingUnitKind, Frame, Plane, PlaneType, ProtoEnumMapping, Spatial};

#[derive(Copy, Clone, PartialEq, Eq, Debug, Deserialize, Serialize)]
pub enum FrameStatistic {
    LumaModes,
    ChromaModes,
    BlockSizes,
    Symbols,
    PartitionSplit,
}

#[derive(Default, Debug, Deserialize, Serialize, PartialEq)]
pub struct StatsFilter {
    pub include: Vec<String>,
    pub exclude: Vec<String>,
}

impl StatsFilter {
    fn from_comma_separated(include: &str, exclude: &str) -> Self {
        Self {
            include: include
                .split(',')
                .map(|s| s.trim().to_string())
                .filter(|s| !s.is_empty())
                .collect(),
            exclude: exclude
                .split(',')
                .map(|s| s.trim().to_string())
                .filter(|s| !s.is_empty())
                .collect(),
        }
    }
}

#[derive(Copy, Clone, PartialEq, Eq, Debug, Deserialize, Serialize)]
pub enum StatSortMode {
    Unsorted,
    ByName,
    ByValue,
}
impl StatSortMode {
    pub fn name(&self) -> &'static str {
        match self {
            Self::Unsorted => "Unsorted",
            Self::ByName => "By name",
            Self::ByValue => "By value",
        }
    }
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
pub struct StatsSettings {
    pub sort_by: StatSortMode,
    // Using separate bool + value fields rather than an Option to make the UI design a bit more intuitive (e.g. checkbox + disabled number input).
    pub apply_limit_count: bool,
    pub limit_count: usize,
    pub apply_limit_frac: bool,
    pub limit_frac: f32,
    pub include_filter: String,
    pub exclude_filter: String,
    pub include_filter_exact_match: bool,
    pub exclude_filter_exact_match: bool,
    pub show_relative_total: bool,
    // Comma separated list of block sizes to include for partition split stats.
    pub partition_split_block_sizes: String,
}

impl Default for StatsSettings {
    fn default() -> Self {
        Self {
            sort_by: StatSortMode::ByValue,
            apply_limit_count: false,
            limit_count: 20,
            apply_limit_frac: false,
            limit_frac: 0.01,
            include_filter: "".into(),
            exclude_filter: "".into(),
            include_filter_exact_match: false,
            exclude_filter_exact_match: false,
            show_relative_total: false,
            partition_split_block_sizes: "".into(),
        }
    }
}

#[derive(Debug)]
pub struct Sample {
    pub name: String,
    pub value: f64,
}
impl Sample {
    pub fn new(name: String, value: f64) -> Self {
        Self { name, value }
    }
}

impl FrameStatistic {
    fn luma_modes(&self, frame: &Frame) -> HashMap<String, f64> {
        let modes = frame.iter_coding_units(CodingUnitKind::Shared).map(|ctx| {
            let cu = ctx.coding_unit;
            let prediction_mode = cu.prediction_mode.as_ref().unwrap();
            frame
                .enum_lookup(ProtoEnumMapping::PredictionMode, prediction_mode.mode)
                .unwrap_or("UNKNOWN".into())
        });
        let mut modes_map: HashMap<String, f64> = HashMap::new();

        for mode in modes {
            *modes_map.entry(mode).or_default() += 1.0;
        }

        modes_map
    }

    fn chroma_modes(&self, frame: &Frame) -> HashMap<String, f64> {
        let kind = frame.coding_unit_kind(PlaneType::Planar(Plane::U));
        let modes = frame.iter_coding_units(kind).map(|ctx| {
            let cu = ctx.coding_unit;
            let prediction_mode = cu.prediction_mode.as_ref().unwrap();
            frame
                .enum_lookup(ProtoEnumMapping::UvPredictionMode, prediction_mode.uv_mode)
                .unwrap_or("UNKNOWN".into())
        });
        let mut modes_map: HashMap<String, f64> = HashMap::new();

        for mode in modes {
            *modes_map.entry(mode).or_default() += 1.0;
        }

        modes_map
    }

    fn block_sizes(&self, frame: &Frame) -> HashMap<String, f64> {
        let sizes = frame.iter_coding_units(CodingUnitKind::Shared).map(|ctx| {
            let cu = ctx.coding_unit;
            let w = cu.width();
            let h = cu.height();
            format!("{w}x{h}")
        });
        let mut sizes_map: HashMap<String, f64> = HashMap::new();

        for size in sizes {
            *sizes_map.entry(size).or_default() += 1.0;
        }
        sizes_map
    }

    fn partition_split(&self, frame: &Frame, settings: &StatsSettings) -> HashMap<String, f64> {
        // TODO(comc): Add settings option for partition kind.
        let filter = StatsFilter::from_comma_separated(&settings.partition_split_block_sizes, "");
        let splits = frame.iter_partitions(CodingUnitKind::Shared).filter_map(|ctx| {
            let partition = ctx.partition;
            let size = partition.size_name();
            if !filter.include.is_empty() && !filter.include.iter().any(|incl| &size == incl) {
                return None;
            }
            let partition_type = frame
                .enum_lookup(ProtoEnumMapping::PartitionType, partition.partition_type)
                .unwrap_or("UNKNOWN".into());
            Some(partition_type)
        });
        let mut splits_map: HashMap<String, f64> = HashMap::new();

        for split in splits {
            *splits_map.entry(split).or_default() += 1.0;
        }
        splits_map
    }

    fn symbols(&self, frame: &Frame) -> HashMap<String, f64> {
        let mut symbols: HashMap<String, f64> = HashMap::new();
        // TODO(comc): Use iter_symbols. Add iter_symbols method for partition blocks as well.
        let sbs = frame.iter_superblocks().map(|sb_ctx| {
            let sb = sb_ctx.superblock;
            let mut symbols_sb: HashMap<String, f64> = HashMap::new();
            for symbol in sb.symbols.iter() {
                let info = symbol.info_id;
                let info = &frame.symbol_info[&info];
                let name = info.source_function.clone();
                let bits = symbol.bits;
                *symbols_sb.entry(name.clone()).or_default() += bits as f64;
            }
            symbols_sb
        });
        for symbols_sb in sbs {
            for (name, bits) in symbols_sb {
                *symbols.entry(name.clone()).or_default() += bits;
            }
        }
        symbols
    }

    fn apply_settings(&self, mapping: HashMap<String, f64>, settings: &StatsSettings) -> Vec<Sample> {
        let mut samples: Vec<_> = mapping
            .into_iter()
            .map(|(name, value)| Sample::new(name, value))
            .collect();
        let filter: StatsFilter = StatsFilter::from_comma_separated(&settings.include_filter, &settings.exclude_filter);
        let total: f64 = samples.iter().map(|sample| sample.value).sum();
        let mut other = 0.0;
        samples.retain(|Sample { name, value }| {
            let mut keep = true;
            // TODO(comc): Make this a method of StatsFilter.
            if !filter.include.is_empty() {
                if settings.include_filter_exact_match {
                    if !filter.include.contains(name) {
                        keep = false;
                    }
                } else if !filter.include.iter().any(|incl| name.contains(incl)) {
                    keep = false;
                }
            }
            if settings.exclude_filter_exact_match {
                if filter.exclude.contains(name) {
                    keep = false
                }
            } else if filter.exclude.iter().any(|excl| name.contains(excl)) {
                keep = false;
            }
            if !keep {
                other += value;
            }
            keep
        });

        let filtered_total: f64 = samples.iter().map(|sample| sample.value).sum();

        let top_n = if settings.apply_limit_count {
            let top_n: HashSet<_> = samples
                .iter()
                .sorted_by_key(|sample| (OrderedFloat(sample.value), &sample.name)) // name used as a tie-breaker.
                .rev()
                .map(|sample| sample.name.clone())
                .take(settings.limit_count)
                .collect();
            Some(top_n)
        } else {
            None
        };

        samples.retain(|Sample { name, value }| {
            let mut keep = true;
            if settings.apply_limit_frac {
                let frac = if settings.show_relative_total {
                    *value / filtered_total
                } else {
                    *value / total
                };
                if frac < settings.limit_frac as f64 {
                    keep = false;
                }
            }

            if let Some(top_n) = &top_n {
                if !top_n.contains(name) {
                    keep = false;
                }
            }

            if !keep {
                other += value;
            }
            keep
        });

        if !settings.show_relative_total && other > 0.0 {
            samples.push(Sample::new("Other".into(), other));
        }
        match settings.sort_by {
            StatSortMode::ByName => samples
                .into_iter()
                .sorted_by_key(|sample| sample.name.clone())
                .collect(),
            StatSortMode::ByValue => samples
                .into_iter()
                .sorted_by_key(|sample| (OrderedFloat(sample.value), sample.name.clone())) // name used as a tie-breaker.
                .collect(),
            StatSortMode::Unsorted => samples,
        }
    }

    pub fn calculate(&self, frame: &Frame, settings: &StatsSettings) -> Vec<Sample> {
        let mapping = match self {
            FrameStatistic::LumaModes => self.luma_modes(frame),
            FrameStatistic::ChromaModes => self.chroma_modes(frame),
            FrameStatistic::BlockSizes => self.block_sizes(frame),
            FrameStatistic::Symbols => self.symbols(frame),
            FrameStatistic::PartitionSplit => self.partition_split(frame, settings),
        };
        self.apply_settings(mapping, settings)
    }

    pub fn name(&self) -> &'static str {
        match self {
            FrameStatistic::LumaModes => "Luma modes",
            FrameStatistic::ChromaModes => "Chroma modes",
            FrameStatistic::BlockSizes => "Block sizes",
            FrameStatistic::Symbols => "Symbols",
            FrameStatistic::PartitionSplit => "Partition split",
        }
    }
}

#[cfg(test)]
mod tests {

    use super::*;

    #[test]
    fn test_exclude_filter() {
        let mapping: HashMap<String, f64> = [("ABC", 1.0), ("AAA", 2.0)]
            .iter()
            .map(|(k, v)| (k.to_string(), *v))
            .collect();
        let settings = StatsSettings {
            exclude_filter: "A".into(),
            exclude_filter_exact_match: true,
            show_relative_total: true,
            ..Default::default()
        };
        let samples = FrameStatistic::Symbols.apply_settings(mapping, &settings);
        assert_eq!(samples.len(), 2);
    }
}