import os
import numpy

from iplotter import C3Plotter

result = {}
methods = []
directory = "testfiles"
videofile = "5_face.mov"

def kpis(path):
    with open(path) as f:
        lines = f.readlines()
    
    # collect all values (skipping first and last line)
    vals = [float(line.split(",")[1]) for line in lines[1:-1]]        
    
    return {"max": max(vals), "min": min(vals), "avg": numpy.average(vals)}

for _, _, files in os.walk(directory):
    for file in files:
        if file.endswith(".csv") and file.startswith(videofile):
            method_res = kpis(os.path.join(directory, file))
            
            method = file.split(".")[-2].split("_")[1]
            if method not in methods:
                methods.append(method)
                
            bitrate = int(file.split(".")[-3][:-1])

            if bitrate not in result:
                result[bitrate] = {}
                
            result[bitrate][method] = method_res

bitrates = sorted(list(result.keys()))
results_by_method = {}
for method in methods:
    results_by_method[method] = [result[bitrate][method] for bitrate in bitrates]

plotter = C3Plotter()
charts = {}

for method, results in results_by_method.items():
    mins = [x["min"] for x in results]
    maxs = [x["max"] for x in results]
    avgs = [x["avg"] for x in results]
    charts[method] = {
        "data": {
            "x": "bitrate",
            "columns": [
                ["bitrate", *bitrates],
                [method + " min", *mins],
                [method + " max", *maxs],
                [method + " avg", *avgs]
            ]
        },
        "padding": {
            "right": 15
        }
    }

plotter.plot(charts["psnr"])

plotter.plot(charts["msssim"])

plotter.plot(charts["vifp"])

plotter.plot(charts["psnrhvsm"])