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


def eval_data(file):
    data = np.load(file)
    percentage = data["percentage"]
    out = data["out"]
    out = np.array(out)
    print(out.shape)
    fig, all_axs = plt.subplots(2, out.shape[0])
    axs = all_axs[0, :]
    axs2 = all_axs[1, :]
    for o, ax, ax2,  lab in zip(out, axs, axs2, ["rutile", "mono_twin", "mono"]):
        # ax.plot(percentage, o/np.max(o, axis=0))
        ax.plot(percentage, o/o[0])
        # ax.plot(percentage, o)
        o = np.mean(o, axis=1)
        o = o/o[0]
        ax2.plot(percentage, o)
        ax2.plot([0, 1], [o[0], o[-1]], "k:")
        ax.set_title(lab)

    title = ""
    if "ising" in file:
        title += "Ising "
    else:
        title += "Random "
    if "rect" in file:
        title += "rect "
    if "voro" in file:
        title += "voro "
    fig.suptitle(title)
    fig.savefig(f"{file}.png")


def parse_lists(out):
    lists = []
    for o in out:
        lists.append(np.stack(o))

    max = 0
    for l in lists:
        print(l.shape)
        if max < l.shape[1]:
            max = l.shape[1]
    lists = [np.pad(l, ((0, 0), (0, max-l.shape[1]))) for l in lists]
    for l in lists:
        print(l.shape)
    return np.stack(lists)


def eval_data_print(file):
    data = np.load(file, allow_pickle=True)
    percentage = data["percentage"]
    #out = parse_lists(data["out"])
    out = []
    for o in ["out_1","out_2","out_3","out_4"]:
        out.append(np.array(data[o]))
    fig, all_axs = plt.subplots(2, len(out))
    axs = all_axs[0, :]
    axs2 = all_axs[1, :]
    for o, ax, ax2,  lab in zip(out, axs, axs2, ["rutile", "monoclinic", "mono", "rest"]):
        # ax.plot(percentage, o/np.max(o, axis=0))
        ax.plot(percentage, o/o[0])
        # ax.plot(percentage, o)
        o = np.mean(o, axis=1)
        #o = o/o[0]
        ax2.plot(percentage, o)
        ax2.plot([0, 1], [o[0], o[-1]], "k:")
        ax.set_title(lab)
    for ax in all_axs.flatten():
        ax.set_xlabel("Rutile Phase")
        ax.set_ylabel("Normalized Intensity")

    title = ""
    if "ising" in file:
        title += "Ising "
    else:
        title += "Random "
    if "rect" in file:
        title += "rect "
    if "voro" in file:
        title += "voro "
    fig.suptitle(title)
    plt.tight_layout()


if __name__ == "__main__":
    for f in glob.glob("*.npz"):
        eval_data_print(f)
    plt.show()