#!/usr/bin/env python3
# Copyright 2016 Google Inc. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import ast
from pathlib import Path
from visual_metrics import HandleFiles
from collections import OrderedDict
import matplotlib.pyplot as plt
import os
import sys
import re

layer_regex_pattern = re.compile(r"^(\d)sl(\d)tl$")


def writable_dir(directory):
    if not os.path.isdir(directory) or not os.access(directory, os.W_OK):
        raise argparse.ArgumentTypeError(
            "'%s' is either not a directory or cannot be opened for writing.\n"
            % directory)
    return directory


def formats(formats_list):
    formats = formats_list.split(',')
    for extension in formats:
        if extension not in ['png', 'svg']:
            raise argparse.ArgumentTypeError(
                "'%s' is not a valid file format.\n" % extension)
    return formats


parser = argparse.ArgumentParser(description='Generate graphs from data files.')
parser.add_argument('graph_files',
                    nargs='+',
                    metavar='graph_file.txt',
                    type=argparse.FileType('r'))
parser.add_argument('--out-dir', required=True, type=writable_dir)
parser.add_argument('--formats',
                    type=formats,
                    metavar='png,svg',
                    help='comma-separated list of output formats',
                    default=['png', 'svg'])


def split_data(graph_data, attribute):
  groups = {}
  for element in graph_data:
    value = element[attribute]
    if value not in groups:
      groups[value] = []
    groups[value].append(element)
  return groups.values()


def normalize_bitrate_config_string(config):
    return ":".join([str(int(x * 100.0 / config[-1])) for x in config])

def generate_stt(data, output_dir=''):

  metrics = [
          'vpx-ssim',
          'ssim',
          'ssim-y',
          'ssim-u',
          'ssim-v',
          'avg-psnr',
          'avg-psnr-y',
          'avg-psnr-u',
          'avg-psnr-v',
          'glb-psnr',
          'glb-psnr-y',
          'glb-psnr-u',
          'glb-psnr-v',
          'encode-time-utilization',
          'actual-encode-time-ms',
          'vmaf',
          'psnr-dmos'
        ]
  encoder_codecs = set([(item['encoder'], item['codec']) for item in data])
  videos = set([item['input-file'] for item in data])

## For target bitrate


  for encoder, codec in encoder_codecs:
      ## Create a directory for every encoder-codec tool
    Path(f"./{output_dir}/{encoder}:{codec}").mkdir(parents=True, exist_ok=True)
    for video in videos:
        ## Create a file for the metrics of this certain video
      sb = ""

      encoder_metrics = list(filter(lambda item: item['encoder'] == encoder and item['codec'] == codec and item['input-file'] == video, data))

      ##  For Target bitrate
      #   bitrates = sorted(list(set([[item['actual-bitrate-bps'] // 1000 ][0] for item in encoder_metrics])))

      bitrates = sorted(list(set([item['actual-bitrate-bps'] for item in encoder_metrics])))
      for bitrate in bitrates:
        ## Append this to the table

        ## For target bitrate
        # filtered_item = list(filter(lambda item: item['encoder'] == encoder and item['codec'] == codec and item['input-file'] == video and bitrate in [item['actual-bitrate-bps'] // 1000 ], data))


        filtered_item = list(filter(lambda item: item['encoder'] == encoder and item['codec'] == codec and item['input-file'] == video and bitrate == item['actual-bitrate-bps'], data))
           
        if len(filtered_item) != 1:
            print("WARNING: %s was encoded with %s:%s using different target bitrates but has the same actual bitrate" 
                  % (item['input-file'], item['encoder'], item['codec']),
                  file=sys.stderr)
            
        for item in filtered_item:

            ## Filter my dictionary to contain the required key values
            required_data = OrderedDict()
            required_data['bitrate'] = str(bitrate)
            for key in metrics:
                if key in item:
                    required_data[key] = "{:.2f}".format(item[key])

        header = '\t'.join(required_data.keys()) + '\n'
        values = '\t'.join(required_data.values()) + '\n'

        sb += header + values
      with open(f'./{output_dir}/{encoder}:{codec}/{video}.stt', 'w') as file:
        file.write(sb)

  enc_cod_dirs = [f"./{output_dir}/{encoder}:{codec}" for encoder, codec in encoder_codecs]
  html = HandleFiles(['', 'metrics_template.html', '*stt'] + enc_cod_dirs)

  with open(f"{output_dir}/results.html", "w") as file:
    file.write(html)



def generate_graphs(output_dict, graph_data, target_metric,
                    bitrate_config_string):
    lines = {}
    for encoder in split_data(graph_data, 'encoder'):
        for codec in split_data(encoder, 'codec'):
            for layer in split_data(codec, 'temporal-layer'):
                metric_data = []
                for data in layer:
                    if target_metric not in data:
                        return
                    metric_data.append(
                        (data['actual-bitrate-bps'] / 1000, data[target_metric],
                         1))
                line_name = '%s:%s (tl%d)' % (layer[0]['encoder'],
                                              layer[0]['codec'],
                                              layer[0]['temporal-layer'])
                # Sort points on target bitrate.
                lines[line_name] = sorted(metric_data,
                                          key=lambda point: point[0])

    graph_name = "%s-%s-%s:%s" % (graph_data[0]['input-file'],
                                  graph_data[0]['layer-pattern'],
                                  bitrate_config_string, target_metric)
    output_dict[('', graph_name)] = lines


def main():
    args = parser.parse_args()
    graph_data = []
    generate_images = False
    for f in args.graph_files:
        graph_data += ast.literal_eval(f.read())
    generate_stt(graph_data, args.out_dir)
    #Generate images defined above, constant (change to true if images of graphs are wanted.)
    if not generate_images:
        return
    graph_dict = {}
    for input_files in split_data(graph_data, 'input-file'):
        for layer_pattern in split_data(input_files, 'layer-pattern'):
            for data in layer_pattern:
                metrics = [
                    'vpx-ssim', 'ssim', 'ssim-y', 'ssim-u', 'ssim-v',
                    'avg-psnr', 'avg-psnr-y', 'avg-psnr-u', 'avg-psnr-v',
                    'glb-psnr', 'glb-psnr-y', 'glb-psnr-u', 'glb-psnr-v',
                    'encode-time-utilization', 'actual-encode-time-ms','vmaf'
                ]
                for metric in metrics:
                    generate_graphs(
                        graph_dict, layer_pattern, metric,
                        normalize_bitrate_config_string(
                            [data['actual-bitrate-bps'] // 1000]))

    for point in graph_data:
        pattern_match = layer_regex_pattern.match(point['layer-pattern'])
        num_temporal_layers = int(pattern_match.group(2))
        temporal_divide = 2**(num_temporal_layers - 1 - point['temporal-layer'])
        frame_metrics = [
            'frame-ssim', 'frame-ssim-y', 'frame-ssim-u', 'frame-ssim-v',
            'frame-psnr', 'frame-psnr-y', 'frame-psnr-u', 'frame-psnr-v',
            'frame-qp', 'frame-bytes', 'frame-vmaf'
        ]
        for target_metric in frame_metrics:
            if target_metric not in point:
                continue

            split_on_codecs = target_metric == 'frame-qp'

            if split_on_codecs:
                graph_name = "%s-%s-%s-%dkbps-tl%d-%s:%s" % (
                    point['input-file'], point['layer-pattern'],
                    normalize_bitrate_config_string(
                        [point['actual-bitrate-bps'] // 1000 ]),
                    [point['actual-bitrate-bps'] // 1000 ][-1], point['temporal-layer'],
                    point['codec'], target_metric)
                line_name = '%s' % point['encoder']
            else:
                graph_name = "%s-%s-%s-%dkbps-tl%d:%s" % (
                    point['input-file'], point['layer-pattern'],
                    normalize_bitrate_config_string(
                        [point['actual-bitrate-bps'] // 1000 ]),
                    [point['actual-bitrate-bps'] // 1000 ][-1], point['temporal-layer'],
                    target_metric)
                line_name = '%s:%s' % (point['encoder'], point['codec'])
            graph_info = ('frame-data-%s/' % point['input-file'], graph_name)
            if not graph_info in graph_dict:
                graph_dict[graph_info] = {}
            line = []
            for idx, val in enumerate(point[target_metric]):
                frame_size = point['frame-bytes'][
                    idx] if 'frame-bytes' in point else -1
                line.append((point['frame-offset'] + temporal_divide * idx + 1,
                             val, frame_size))
            graph_dict[graph_info][line_name] = line

    current_graph = 1
    total_graphs = len(graph_dict)
    for (subdir, graph_name), lines in graph_dict.items():
        print("[%d/%d] %s" % (current_graph, total_graphs, graph_name))
        current_graph += 1
        metric = graph_name.split(':')[-1]
        fig, ax = plt.subplots()
        ax.set_title(graph_name)
        frame_data = 'frame-' in metric
        ax2 = None
        ax2_bitrate_utilization = False
        linestyle = 'o--'
        ax2_linestyle = 'x-'

        if frame_data:
            ax.set_xlabel('Frame')
            linestyle = '-'
            if metric == 'frame-bytes':
                ax.set_ylabel('Frame Size (bytes / frame)')
            else:
                ax.set_ylabel(metric.replace('frame-', '').upper())
                ax2 = ax.twinx()
                ax2.set_ylabel('Frame Size (bytes / frame)')
                ax2_linestyle = '-'
        elif metric == 'encode-time-utilization':
            ax.set_xlabel('Layer Target Bitrate (kbps)')
            ax.set_ylabel('Encode Time (fraction)')
            # Draw a reference line for realtime.
            ax.axhline(1.0, color='k', alpha=0.2, linestyle='--')
        else:
            ax.set_xlabel('Layer Target Bitrate (kbps)')
            ax.set_ylabel(metric.upper())
            ax2 = ax.twinx()
            ax2.set_ylabel('Bitrate Utilization (actual / target)')
            ax2_bitrate_utilization = True

        for title in sorted(lines.keys()):
            points = lines[title]
            x = []
            y = []
            y2 = []
            for bitrate_kbps, value, utilization in points:
                x.append(bitrate_kbps)
                y.append(value)
                y2.append(utilization)
            ax.plot(x, y, linestyle, linewidth=1, label=title)
            if ax2:
                ax2.plot(x, y2, ax2_linestyle, alpha=0.2)
            ax.legend(loc='best', fancybox=True, framealpha=0.5)

        if metric == 'encode-time-utilization':
            # Make sure the horizontal reference line at 1.0 can be seen.
            (lower, upper) = ax.get_ylim()
            if upper < 1.10:
                ax.set_ylim(top=1.10)

        # TODO(pbos): Read 'input-total-frames' from input and set as graph xlim.
        if frame_data:
            ax.set_xlim(left=0)

        if ax2_bitrate_utilization:
            # Set bitrate limit axes to +/- 20%.
            ax2.set_ylim(bottom=0.80, top=1.20)

        for extension in args.formats:
            graph_dir = os.path.join(args.out_dir, extension, subdir)
            if not os.path.exists(graph_dir):
                os.makedirs(graph_dir)
            plt.savefig(
                os.path.join(
                    graph_dir,
                    "%s.%s" % (graph_name.replace(":", "-"), extension)))
        plt.close()


if __name__ == '__main__':
    main()
