#!/usr/bin/env python3

import glob
from datetime import datetime

import matplotlib.pyplot as plt

def load_data(path):
    all_data = {}
    for file in glob.glob(path):
        with open(file, 'r') as f:
            lines = list(map(lambda l: l.strip(), f))
            ts = file.split('_')[-1][:-4]
            bucket = datetime.fromtimestamp(float(ts))
            data = {}
            data['avg_out'] = float(lines[1].split(': ')[1])
            data['min_out'] = float(lines[2].split(': ')[1].split(', ')[3][0:-1])
            data['max_out'] = float(lines[3].split(': ')[1].split(', ')[3][0:-1])
            data['known_out'] = float(lines[4].split(': ')[1].split(' ')[-1])
            pr_sr_line = lines[11].split(' & ')
            data['known_pr'] = float(pr_sr_line[2])
            data['avg_pr'] = float(pr_sr_line[1])
            data['known_sr'] = float(pr_sr_line[4].split(' ')[0])
            data['avg_sr'] = float(pr_sr_line[3])
            # pr_percentiles = filter(lambda l: l.startswith('PR4('), lines)
            # sr_percentiles = filter(lambda l: l.startswith('SR4('), lines)
            percentiles_to_plot = [50, 60, 70, 80, 90, 95]
            for rank in ['PR', 'SR']:
                percentiles = list(filter(lambda l: l.startswith(rank), lines))
                for to_plot in percentiles_to_plot:
                    line = list(filter(lambda l: l.endswith(f': {to_plot}'), percentiles))[0]
                    percentile = line.split(': ')[2]
                    data[f'{rank.lower()}_perc_{to_plot}'] = float(percentile)

            all_data[bucket] = data

    return all_data


def main():
    # for edges in [0, 20, 30, 40, 50, 60, 70, 80, 90, 100, 130, 150, 200]:
    for edges in [3, 5, 10]:
        print(f'plotting for {edges} edges')
        all_data = load_data(f'./churn_rank/{edges}_*')
        x = []
        avg_out = []
        min_out = []
        max_out = []
        known_out = []
        known_pr = []
        known_sr = []
        avg_pr = []
        avg_sr = []
        pr_perc_50 = []
        pr_perc_60 = []
        pr_perc_70 = []
        pr_perc_80 = []
        pr_perc_90 = []
        pr_perc_95 = []
        sr_perc_50 = []
        sr_perc_60 = []
        sr_perc_70 = []
        sr_perc_80 = []
        sr_perc_90 = []
        sr_perc_95 = []
        for bucket, data in sorted(all_data.items(), key=lambda kv: kv[0]):
            x.append(bucket)
            avg_out.append(data['avg_out'])
            min_out.append(data['min_out'])
            max_out.append(data['max_out'])
            known_out.append(data['known_out'])
            known_pr.append(data['known_pr'])
            known_sr.append(data['known_sr'])
            avg_sr.append(data['avg_sr'])
            avg_pr.append(data['avg_pr'])

            pr_perc_50.append(data['pr_perc_50'])
            pr_perc_60.append(data['pr_perc_60'])
            pr_perc_70.append(data['pr_perc_70'])
            pr_perc_80.append(data['pr_perc_80'])
            pr_perc_90.append(data['pr_perc_90'])
            pr_perc_95.append(data['pr_perc_95'])

            sr_perc_50.append(data['sr_perc_50'])
            sr_perc_60.append(data['sr_perc_60'])
            sr_perc_70.append(data['sr_perc_70'])
            sr_perc_80.append(data['sr_perc_80'])
            sr_perc_90.append(data['sr_perc_90'])
            sr_perc_95.append(data['sr_perc_95'])

        # plt.plot(x, avg_out, label='avg_out')
        # plt.plot(x, max_out, label='max_out')
        # plt.plot(x, min_out, label='min_out')
        # plt.plot(x, known_out, label='known_out')
        # plt.plot(x, known_pr, label='known_pr')
        fig, ax = plt.subplots()
        ax.plot(x, known_sr, label='known_sr')

        # plt.plot(x, pr_perc_50, label='pr_perc_50')
        # plt.plot(x, pr_perc_80, label='pr_perc_80')
        # plt.plot(x, pr_perc_90, label='pr_perc_90')
        # plt.plot(x, pr_perc_95, label='pr_perc_95')

        ax.plot(x, sr_perc_50, label='sr_perc_50')
        ax.plot(x, sr_perc_60, label='sr_perc_60')
        ax.plot(x, sr_perc_70, label='sr_perc_70')
        ax.plot(x, sr_perc_80, label='sr_perc_80')
        ax.plot(x, sr_perc_90, label='sr_perc_90')
        ax.plot(x, sr_perc_95, label='sr_perc_95')
        ax.plot(x, avg_sr, label='avg_sr')

        ax.set_title(f'SensorRank after adding {edges} edges')
        fig.autofmt_xdate()
        fig.legend()
        plt.savefig(f'./{edges}_sr_percentiles.png')
        print('created sr plot')
        # plt.show()

        fig, ax = plt.subplots()
        ax.plot(x, known_pr, label='known_pr')

        ax.plot(x, pr_perc_50, label='pr_perc_50')
        ax.plot(x, pr_perc_60, label='pr_perc_60')
        ax.plot(x, pr_perc_70, label='pr_perc_70')
        ax.plot(x, pr_perc_80, label='pr_perc_80')
        ax.plot(x, pr_perc_90, label='pr_perc_90')
        ax.plot(x, pr_perc_95, label='pr_perc_95')
        ax.plot(x, avg_pr, label='avg_pr')

        ax.set_title(f'PageRank after adding {edges} edges')
        fig.autofmt_xdate()
        fig.legend()
        plt.savefig(f'./{edges}_pr_percentiles.png')
        print('created pr plot')
        # plt.show()


if __name__ == '__main__':
    main()