#!/usr/bin/env python3

import matplotlib.pyplot as plt
from datetime import datetime
import glob
import json
import reduce_edges

def load_json(path):
    with open(path, 'r') as f:
        return json.load(f)


def plot_in_out(percentage, added_percentage, data):
    times = []
    avg_in = []
    known_in = []
    known_out = []
    # same value, independent of algo
    algo = 'sr'
    for when, d in sorted(data.items(), key=lambda kv: kv[0]):
        times.append(when)
        avg_in.append(d[algo]['avg_in'])
        known_in.append(d[algo]['known_in'])
        known_out.append(d[algo]['known_out'])

    fig, ax = plt.subplots()
    a = 'SensorRank' if algo == 'sr' else 'RageRank'
    ax.set_ylabel(f'{a}')
    ax.plot(times, avg_in, label='Avg. In')
    # ax.plot(times, known_in, label='Known In') # TODO
    ax.plot(times, known_in, label='Known In')
    ax.plot(times, known_out, label='Known out')
    title = f'In And Out after removing {percentage * 100}% edges and adding {added_percentage * 100}%'
    ax.set_title(title)

    fig.autofmt_xdate()
    fig.legend()
    plt.savefig(f'./plot_reduced/{percentage:.02f}/{added_percentage:.02f}/in_out.png')


def plot(percentage, added_percentage, algo, data):
    times = []
    avg_rank = []
    a_avg_in = []
    known_rank = []
    known_in = []
    known_out = []
    for when, d in sorted(data.items(), key=lambda kv: kv[0]):
        times.append(when)
        avg_rank.append(d[algo]['avg_rank'])
        a_avg_in.append(d[algo]['avg_in'])
        known_rank.append(d[algo]['known_rank'])
        known_in.append(d[algo]['known_in'])
        known_out.append(d[algo]['known_out'])

    # avg_out = sum(known_out) / len(known_out)
    # avg_in = sum(known_in) / len(known_in)

    fig, ax = plt.subplots()
    a = 'SensorRank' if algo == 'sr' else 'RageRank'
    ax.set_ylabel(f'{a}')
    ax.plot(times, avg_rank, label='Avg. Rank')
    # ax.plot(times, known_in, label='Known In') # TODO
    ax.plot(times, known_rank, label='Known Rank')
    # title = f'{a} after removing {percentage * 100}% edges and adding {added_percentage * 100}%\nin = {avg_in:.02f} out = {avg_out:.02f}'
    title = f'{a} after removing {percentage * 100}% edges and adding {added_percentage * 100}%'
    ax.set_title(title)

    # ax2 = ax.twinx()
    # ax2.set_ylabel('Edges')
    # ax2.plot(times, known_in, label='Known In', color='red')
    # ax2.plot(times, known_out, label='Known Out', color='green')
    # # print(f'{len(times)=}, {len(a_avg_in)=}')
    # # assert len(times) == len(a_avg_in)
    # ax2.plot(times, a_avg_in, label='Avg. In', color='violet')

    fig.autofmt_xdate()
    fig.legend()
    plt.savefig(f'./plot_reduced/{percentage:.02f}/{added_percentage:.02f}/{algo}.png')
    # print('created sr plot')
    # plt.show()


def main2():
    for reduced_percentage in reduce_edges.percentages:
        perc = reduce_edges.percentages.copy()
        perc.append(1.0)
        perc.append(1.2)
        for added_percentage in perc:
            data = {}
            for file in glob.glob(f'./data_reduced/{reduced_percentage:.02f}/{added_percentage:.02f}/*.json'):
                when = datetime.fromtimestamp(float(file.split('/')[-1][:-5]))
                print(f'{reduced_percentage=:.02f}, {added_percentage=:.02f}, {when=}')
                data[when] = load_json(file)

            plot(reduced_percentage, added_percentage, 'sr', data)
            plot(reduced_percentage, added_percentage, 'pr', data)
            plot_in_out(reduced_percentage, added_percentage, data)




def plot2(percentage, algo, data):
    times = []
    mean = []
    stdev = []
    a_avg_in = []
    known_rank = []
    known_in = []
    known_out = []
    for when, d in sorted(data.items(), key=lambda kv: kv[0]):
        times.append(when)
        mean.append(d[algo]['mean'])
        stdev.append(d[algo]['stdev'])
        a_avg_in.append(d[algo]['avg_in'])
        known_rank.append(d[algo]['known_rank'])
        known_in.append(d[algo]['known_in'])
        known_out.append(d[algo]['known_out'])

    # avg_out = sum(known_out) / len(known_out)
    # avg_in = sum(known_in) / len(known_in)

    fig, ax = plt.subplots()
    a = 'SensorRank' if algo == 'sr' else 'RageRank'
    ax.set_ylabel(f'{a}')
    ax.plot(times, mean, label='Avg. Rank')
    # ax.errorbar(times, mean, stdev, label='Avg. Rank')
    ax.plot(times, mean, label='Avg. Rank')
    # ax.plot(times, known_in, label='Known In') # TODO
    ax.plot(times, known_rank, label='Known Rank')
    # title = f'{a} after removing {percentage * 100}% edges and adding {added_percentage * 100}%\nin = {avg_in:.02f} out = {avg_out:.02f}'
    title = f'{a} after removing {percentage * 100}% edges'
    ax.set_title(title)

    fig.autofmt_xdate()
    fig.legend()
    path = f'./plot_reduced/{percentage:.02f}/{algo}.png'
    with reduce_edges.open_mkdir(path, 'w'):
        print('created')
    plt.savefig(path)
    plt.close(fig)


def plot_in_out2(percentage, data):
    times = []
    avg_in = []
    known_in = []
    known_out = []
    # same value, independent of algo
    algo = 'sr'
    for when, d in sorted(data.items(), key=lambda kv: kv[0]):
        times.append(when)
        avg_in.append(d[algo]['avg_in'])
        known_in.append(d[algo]['known_in'])
        known_out.append(d[algo]['known_out'])

    fig, ax = plt.subplots()
    a = 'SensorRank' if algo == 'sr' else 'RageRank'
    ax.set_ylabel(f'{a}')
    ax.plot(times, avg_in, label='Avg. In')
    # ax.plot(times, known_in, label='Known In') # TODO
    ax.plot(times, known_in, label='Known In')
    ax.plot(times, known_out, label='Known out')
    title = f'In And Out after removing {percentage * 100}% edges'
    ax.set_title(title)

    fig.autofmt_xdate()
    fig.legend()
    path = f'./plot_reduced/{percentage:.02f}/in_out.png'
    with reduce_edges.open_mkdir(path, 'w'):
        print('created')
    plt.savefig(path)
    plt.close(fig)


def main():
    for perc in reduce_edges.percentages:
        data = {}
        for file in glob.glob(f'./data_reduced/{perc:.02f}/*.json'):
            when = datetime.fromtimestamp(float(file.split('/')[-1][:-5]))
            print(f'{perc=:.02f}, {when=}')
            data[when] = load_json(file)

        plot2(perc, 'sr', data)
        plot2(perc, 'pr', data)
        plot_in_out2(perc, data)

if __name__ == '__main__':
    main()