#!/usr/bin/env python3
## Copyright (c) 2023, Alliance for Open Media. All rights reserved
##
## This source code is subject to the terms of the BSD 3-Clause Clear License and the
## Alliance for Open Media Patent License 1.0. If the BSD 3-Clause Clear License was
## not distributed with this source code in the LICENSE file, you can obtain it
## at aomedia.org/license/software-license/bsd-3-c-c/.  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 aomedia.org/license/patent-license/.
##
from __future__ import annotations

import abc
from collections import OrderedDict, defaultdict
from functools import partial
import pathlib
import tempfile
from typing import Sequence

from absl import app
from absl import flags
from absl import logging
from avm_stats.extract_proto import *
from avm_stats.frame_visualizations import *
from avm_stats.proto_helpers import *
from avm_stats.yuv_tools import *
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import PIL

_STREAM = flags.DEFINE_string("stream", None, "Path to AVM stream.")
flags.mark_flag_as_required("stream")

_SOURCE = flags.DEFINE_string(
    "source", None, "Path to source YUV/Y4M before encoding (optional)."
)

_EXTRACT_PROTO_BIN = flags.DEFINE_string(
    "extract_proto_bin", None, "Path to extract_proto binary."
)
flags.mark_flag_as_required("extract_proto_bin")

_FRAME = flags.DEFINE_integer(
    "frame", 0, "Frame number to visualize (defaults to first frame)."
)
_PLANE = flags.DEFINE_enum(
    "plane", "y", ["y", "u", "v"], "Plane to visualize (defaults to luma: 'y')"
)


def main(argv: Sequence[str]) -> None:
  if len(argv) > 1:
    raise app.UsageError("Too many command-line arguments.")
  with tempfile.TemporaryDirectory() as tmp_dir:
    tmp_path = pathlib.Path(tmp_dir)
    stream_path = pathlib.Path(_STREAM.value)
    extract_proto_path = pathlib.Path(_EXTRACT_PROTO_BIN.value)
    yuv_path = _SOURCE.value or None
    frames = extract_and_load_protos(
        extract_proto_path=extract_proto_path,
        stream_path=stream_path,
        output_path=tmp_path,
        skip_if_output_already_exists=False,
        yuv_path=yuv_path,
    )
    seq = list(frames)
    num_frames = len(seq)
    logging.info(f"Loaded {num_frames} frame protos.")
  plane = Plane[_PLANE.value.upper()]
  mode_counts = defaultdict(int)
  for sb in seq[_FRAME.value].superblocks:
    for cu in sb.get_coding_units(use_chroma=plane.is_chroma()):
      mode_counts[cu.get_prediction_mode()] += 1
  total = sum(mode_counts.values())
  keys = mode_counts.keys()
  labels = [f"{k} ({100.0*mode_counts[k]/total:.2f}%)" for k in keys]
  sizes = [mode_counts[k] for k in keys]

  fig, ax = plt.subplots()
  ax.pie(sizes, labels=labels)
  plt.show()


if __name__ == "__main__":
  app.run(main)