from os import listdir

from os.path import join

import matplotlib.pyplot as plt

import numpy as np

import pandas as pd

def get_predictorwise_distribution(experiment_path):

    mean_ious = []

    for experiment in listdir(experiment_path):

        # mean_ious.append(np.load(join(experiment_path, experiment))) #cool

        mean_ious.append(np.mean(np.load(join(experiment_path, experiment))))

    plt.hist(mean_ious, bins=np.linspace(min(mean_ious), max(mean_ious), 25))

    plt.show()

def get_datawise_distribution(experiment_path):

    mean_ious = []

    for experiment in listdir(experiment_path):

        # mean_ious.append(np.load(join(experiment_path, experiment))) #cool

        mean_ious = np.concatenate((mean_ious, np.mean(np.load(join(experiment_path, experiment)))))

        plt.hist(mean_ious, bins=np.linspace(min(mean_ious), max(mean_ious), 25))

        plt.show()

def plot_training_progression(csv_name):

    df = pd.read_csv(csv_name)

    plt.plot(df["epoch"], df["iid_test_iou"], label="IID")

    plt.plot(df["epoch"], df["ood_iou"], label="OOD")

    plt.plot(df["epoch"], df["iid_test_iou"] - df["ood_iou"], label="Diff")

    plt.legend()

    plt.ylim((0, 1))

    plt.xlim((0, 250))

    plt.show()

if __name__ == '__main__':

    plot_training_progression("logs/consistency/DeepLab/0.csv")

    plot_training_progression("logs/consistency/DeepLab/1.csv")

    # df_iid = df["Augmented" not in df["name"]]

    # get_predictorwise_distribution("experiments/Data/Normal-Pipelines/DeepLab")