from ditact_pic import plot
from spin_image import SpinImage, FFT
import sys
import numpy as np
import matplotlib.pyplot as plt
import glob
import scipy.interpolate as ip
from lattices import VO2_Lattice
plt.style.use(["style", "colors", "one_column"])


def check_percentage(p1, p2):
    plt.figure()
    plt.plot(p1, p2)


def merge(files):
    merge = []
    plt.figure()
    for file in files:
        print(file)
        data = np.load(file, allow_pickle=True)
        old_percentage = data["percentage"]
        w_percentage = data["w_percentage"]
        # check_percentage(old_percentage, w_percentage)
        percentage = old_percentage
        out = []
        for o in ["out_1", "out_2", "out_3", "out_4"]:
            out.append(np.array(data[o]))
        print(out)
        out = np.array(out)[:, :, 0]

        summe = np.max(np.sum(out, axis=0))
        out = out / summe
        merge.append(out)

        plt.plot(w_percentage,out[0, :], "r")
        plt.plot(w_percentage,out[3, :], "b")
        plt.plot(w_percentage,out[2, :], "g")

    all = sum(merge)
    summe = np.max(np.sum(all, axis=0))
    all = all / summe

    plt.plot(all[0, :], "k")
    plt.plot(all[3, :], "k")
    plt.plot(all[2, :], "k")
    percentage = 1-percentage
    return percentage, all


def debug(percentage, out):
    plt.figure()
    for o in out:
        plt.plot(percentage, o)

    plt.plot(percentage, out[0, :], "k")
    plt.plot(percentage, out[3, :], "k")
    plt.plot(percentage, out[2, :], "k")


def stacked_plot(ax, percentage, out, title=""):
    stacks = ax.stackplot(percentage, out[[0, 3, 2]], colors=[
        "w"], ls=(0, (0, 1)), ec="w")
    hatches = ["//", "|", "\\\\"]
    for stack, hatch, color in zip(stacks, hatches, ["C1", "C0", "C2"]):
        stack.set_hatch(hatch)
        stack.set_edgecolor(color)
    ax.set_xlabel("Metallic Phase (%)")
    ax.set_ylabel("normalized Intensity ")
    ax.set_ylim([0.4, 1])
    ax.set_xlim([0., 1])
    ax.text(0.1, 0.9, "monoclinic", backgroundcolor="w",
            bbox=dict(boxstyle='square,pad=0.0', ec="None", fc="w"))
    ax.text(0.6, 0.5, "rutile", backgroundcolor="w",
            bbox=dict(boxstyle='square,pad=0.0', ec="None", fc="w"))
    ax.text(0.35, 0.73, "diffusive", backgroundcolor="w",
            bbox=dict(boxstyle='square,pad=0.0', ec="None", fc="w"))
    ax.stackplot(percentage, out[[0, 3, 2]], colors=["None"], ec="k")


def time_scale(ax, p, o):
    rut_perc = o[0]
    rut_perc = rut_perc - np.min(rut_perc)
    rut_perc /= np.max(rut_perc)

    mono_perc = -o[2]
    mono_perc = mono_perc - np.min(mono_perc)
    mono_perc /= np.max(mono_perc)

    # cs_rut = ip.CubicSpline(p[::-1], rut_perc[::-1])
    # cs_mono = ip.CubicSpline(p[::-1], mono_perc[::-1])
    cs_rut = ip.interp1d(p[::-1], rut_perc[::-1])
    cs_mono = ip.interp1d(p[::-1], mono_perc[::-1])

    # plt.figure()
    # ph = np.linspace(0.01, 0.99, 100)
    # plt.plot(ph, cs_rut(ph))
    # plt.plot(ph, cs_mono(ph))

    time = np.linspace(0.01, 3, 1000)
    phy_phase = 1-np.exp(-time)
    rut_phase = cs_rut(phy_phase)
    mono_phase = cs_mono(phy_phase)

    ax.plot(time, phy_phase, "k:", label="physical")
    ax.plot(time, rut_phase, label="rutile", color="C1")
    ax.plot(time, mono_phase, label="monoclinic", color="C2")
    ax.set_xlabel("time (a.u.)")
    ax.set_ylabel("Metallic Phase (%)")
    ax.set_xlim([0, 3])
    ax.set_ylim([0, 1])
    ax.legend()


def read_file(file):
    files = np.load("./merged.npz")
    p = files["p"]
    o = files["o"]
    return p, o


def intens(ax, file, p, o):
    intens = FFT()
    intens.load(file)
    plot(intens, ax)
    ax.set_xlim([-1, 1])
    ax.set_ylim([-.9, .9])
    ax.axis("off")

    #rect = plt.Rectangle((-1, -.8), 2, 1.6, facecolor="None", hatch="//")
    # ax.add_patch(rect)
    lat = VO2_Lattice(20, 20)
    reci = lat.get_spots()
    print(reci)
    size = (intens.freqx[1] - intens.freqx[0]) * 20
    size2 = size/2
    # big_rect = plt.Rectangle((-10, -10), 20, 20, fc="None", ec="k", hatch="//")
    # ax.add_patch(big_rect)
    for x, y in zip(reci[0][0], reci[0][1]):
        if x < 1 and x > -1:
            if y < 1 and y > -1:
                print(x, y)
                rect = plt.Rectangle((-y-size2, x-size2),
                                     size, size, fc="C1", ec="k", alpha=0.5)
                # big_rect.set_clip_path(rect)
                ax.add_patch(rect)
    for x, y in zip(reci[1][0], reci[1][1]):
        if x < 1 and x > -1:
            if y < 1 and y > -1:
                print(x, y)
                rect = plt.Rectangle((-y-size2, x-size2),
                                     size, size, fc="C2", ec="k", alpha=0.5)
                ax.add_patch(rect)

    axins = ax.inset_axes([0.0, 0.0, 0.5, 0.5])
    axins.plot(p, o[0], label="rut.", color="C1")
    axins.plot(p, o[2], label="mono.", color="C2")
    axins.plot(p, o[3], label="diff.", color="C0")
    axins.legend(loc='center left', bbox_to_anchor=(1, 0.5))
    axins.set_xlim([0, 1])
    axins.set_ylim([0, 1])
    axins.set_xlabel("phase (%)")
    axins.set_ylabel("signal",labelpad=-5)
    
    # axins.get_yaxis().set_visible(False)
    # axins.yaxis.tick_right()
    axins.set_yticks([0, 1])


if __name__ == "__main__":
    p, o = merge(sys.argv[2:])
    np.savez("merged.npz", p=p, o=o)
    # eval_data_print(f)

    fig, axs = plt.subplots(1, 3)
    fig.set_figheight(2)
    stacked_plot(axs[1], p, o)
    time_scale(axs[2], p, o)
    intens(axs[0], sys.argv[1], p, o)
    plt.tight_layout()
    plt.savefig("analysis.pdf")
    plt.savefig("analysis.png")
    plt.show()