# import torch
# print(torch.__version__)
# from torch.utils.tensorboard import SummaryWriter
# %load_ext tensorboard


import os
import math
import matplotlib.pyplot as plt
import seaborn as sns
# %matplotlib inline
import numpy as np
import pandas as pd
sns.set()


# !rm -R ./runs
# writer = SummaryWriter()

# root_dir = "../../workflow_v3/test_case/big_model/{}/output/mse_analysis".format("Facedetection_SSD_without_BN")

# writer.close()



# plt.plot(threshold_lst, mse_lst)
# plt.title(fn)
# plt.xlabel("threshold")
# plt.ylabel("MSE");


def main():

    # input validation
    if len(sys.argv) != 2:
        print("need mse dump dir")
        return

    root_dir = sys.argv[1]
    num_fig = len(os.listdir(root_dir))
    nrows=math.ceil(num_fig/2.0)
    ncols= 2

    fig = plt.figure(figsize=(15, 100))
    fig.subplots_adjust(hspace=0.4, wspace=0.4)
    for idx, fn in enumerate(os.listdir(root_dir)[:]):
        file_path = os.path.join(root_dir, fn)
        file = open(file_path, "r")

        threshold_lst = []
        mse_lst = []
        for line in file.readlines():
            line = line.rstrip().split(",")
            threshold = float(line[0])
            threshold_lst.append(threshold)
            mse = 10*math.log10(float(line[1]))
            mse_lst.append(mse)
    #         writer.add_scalar(fn, mse, walltime=threshold)
        idx+=1
        ax = fig.add_subplot(nrows, ncols, idx)
        ax.plot(threshold_lst, mse_lst)
        ax.set_title(fn)

        ax.plot(threshold_lst, mse_lst)
        ax.set_xlabel("threshold")
        ax.set_ylabel("MSE")
    plt.savefig("test.pdf",bbox_inches='tight')

if __name__ == "__main__":
    main()