#!/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 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 *
import matplotlib.pyplot as plt

_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 the first frame)."
)
_PLANE = flags.DEFINE_enum(
    "plane", "y", ["y", "u", "v"], "Plane to visualize (defaults to luma: 'y')"
)
_FILTER = flags.DEFINE_string("filter", None, "Symbol filter to apply.")


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.")

  frame = seq[_FRAME.value]
  c_functions = set()
  c_functions |= set(
      i.source_function for i in frame.proto.symbol_info.values()
  )
  print("Available symbol functions:")
  for c_function in sorted(c_functions):
    print(c_function)
  symbol_filter = _FILTER.value

  _, axes = plt.subplots(2, 1)

  def my_filter(sym: Symbol):
    return sym.source_function == symbol_filter or symbol_filter is None

  plane = Plane[_PLANE.value.upper()]
  use_chroma = plane.is_chroma()
  ReconstructionYuvLayer(plane=plane).show(frame, axes[0])
  BitsHeatmapLayer(use_chroma=use_chroma, filt=my_filter).add(
      PartitionTreeLayer(use_chroma=use_chroma)
  ).add(SuperblockLayer()).show(frame, axes[1])
  total_bits = 0
  bits_for_symbol = 0
  for sb in frame.superblocks:
    total_bits += sb.get_total_bits()
    bits_for_symbol += sb.get_total_bits(use_chroma=use_chroma, filt=my_filter)
  if symbol_filter is not None:
    percent = 100.0 * bits_for_symbol / total_bits
    print(f"Bits for {symbol_filter}: {bits_for_symbol} ({percent:.2f}%)")
  plt.show()


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