import matplotlib.pyplot as plt
import os
import numpy as np

colors = [
    'orange',
    'green',
    'royalblue',
    'purple',
    'red',
    'gold'
]


def get_data(metric, run_id=None, run_path=None):
    if run_path is not None:
        metric_file_path = run_path + '/metric-' + metric + '.txt'
    else:
        runs = os.listdir('../results')
        prefix = str(run_id).zfill(5)
        run = None
        for cur_run in runs:
            if cur_run.startswith(prefix):
                run = cur_run
                break
        assert run is not None
        metric_file_path = '../results/' + run + '/metric-' + metric + '.txt'
    with open(metric_file_path) as file:
        line = file.readline()
        kimg = []
        score = []
        while line:
            kimg.append(int(line[17:23]) / 1000)
            score.append(np.float32(line[50 + len(metric):]))
            line = file.readline()
    return kimg, score


def plot_multiple_runs(
        descriptions,
        xlabel,
        ylabel,
        title,
        filename,
        run_ids=None,
        run_paths=None,
        x_min=0,
        x_max=15,
        y_min=0,
        y_max=10,
        xticks=1,
        yticks=1,
        metric='fid50k',
        arrows=False,
        show_lowest_value=True):
    fig, ax = plt.subplots()
    offset = 0.62
    for run_index in range(len(descriptions)):
        if run_paths is None:
            x, y = get_data(metric, run_id=run_ids[run_index])
        else:
            x, y = get_data(metric, run_path=run_paths[run_index])
        xmin = x[np.argmin(y)]
        ymin = min(y)
        text = "lowest: {:.4f}".format(ymin)
        if len(descriptions) <= 1:
            ax.plot(x, y, colors[run_index])
        else:
            ax.plot(x, y, colors[run_index], label=descriptions[run_index])

        bbox_props = dict(boxstyle="square,pad=0.3", fc="w", ec="k", lw=0.52)

        arrowprops = None
        if arrows:
            arrowprops = dict(arrowstyle="->", connectionstyle="angle,angleA=0,angleB=80")
        kw = dict(xycoords='data', textcoords="axes fraction", annotation_clip=False,
                  arrowprops=arrowprops, bbox=bbox_props, color=colors[run_index])
        if show_lowest_value:
            ax.annotate(text, xy=(xmin, ymin), xytext=(0.8, offset), **kw)
        offset -= 0.1
    ax.xaxis.set_ticks(np.arange(x_min, x_max, xticks))
    ax.yaxis.set_ticks(np.arange(y_min, y_max, yticks))
    ax.set_xlim(x_min, x_max)
    ax.set_ylim([y_min, y_max])
    ax.set(xlabel=xlabel, ylabel=ylabel, title=title)
    ax.grid()
    if len(descriptions) <= 1:
        ax.legend()
    fig.tight_layout()
    fig.savefig(filename, dpi=400)
    plt.show()

import matplotlib.pyplot as plt
import os
import numpy as np

colors = [
    'orange',
    'green',
    'royalblue',
    'purple',
    'red',
    'gold'
]


def get_data(metric, run_id=None, run_path=None):
    if run_path is not None:
        metric_file_path = run_path + '/metric-' + metric + '.txt'
    else:
        runs = os.listdir('../results')
        prefix = str(run_id).zfill(5)
        run = None
        for cur_run in runs:
            if cur_run.startswith(prefix):
                run = cur_run
                break
        assert run is not None
        metric_file_path = '../results/' + run + '/metric-' + metric + '.txt'
    with open(metric_file_path) as file:
        line = file.readline()
        kimg = []
        score = []
        while line:
            kimg.append(int(line[17:23]) / 1000)
            score.append(np.float32(line[50 + len(metric):]))
            line = file.readline()
    return kimg, score


def plot_multiple_runs(
        descriptions,
        xlabel,
        ylabel,
        title,
        filename,
        run_ids=None,
        run_paths=None,
        x_min=0,
        x_max=15,
        y_min=0,
        y_max=10,
        xticks=1,
        yticks=1,
        metric='fid50k',
        arrows=False,
        show_lowest_value=True):
    fig, ax = plt.subplots()
    offset = 0.62
    for run_index in range(len(descriptions)):
        if run_paths is None:
            x, y = get_data(metric, run_id=run_ids[run_index])
        else:
            x, y = get_data(metric, run_path=run_paths[run_index])
        xmin = x[np.argmin(y)]
        ymin = min(y)
        text = "lowest: {:.4f}".format(ymin)
        if len(descriptions) <= 1:
            ax.plot(x, y, colors[run_index])
        else:
            ax.plot(x, y, colors[run_index], label=descriptions[run_index])

        bbox_props = dict(boxstyle="square,pad=0.3", fc="w", ec="k", lw=0.52)

        arrowprops = None
        if arrows:
            arrowprops = dict(arrowstyle="->", connectionstyle="angle,angleA=0,angleB=80")
        kw = dict(xycoords='data', textcoords="axes fraction", annotation_clip=False,
                  arrowprops=arrowprops, bbox=bbox_props, color=colors[run_index])
        if show_lowest_value:
            ax.annotate(text, xy=(xmin, ymin), xytext=(0.8, offset), **kw)
        offset -= 0.1
    ax.xaxis.set_ticks(np.arange(x_min, x_max, xticks))
    ax.yaxis.set_ticks(np.arange(y_min, y_max, yticks))
    ax.set_xlim(x_min, x_max)
    ax.set_ylim([y_min, y_max])
    ax.set(xlabel=xlabel, ylabel=ylabel, title=title)
    ax.grid()
    if len(descriptions) <= 1:
        ax.legend()
    fig.tight_layout()
    fig.savefig(filename, dpi=400)
    plt.show()
