from plotter import Plotter
import matplotlib
from lattices import VO2_New
from spin_image import SpinImage, FFT
import numpy as np
import matplotlib.pyplot as plt

if __name__ == "__main__":
    plt.style.use(["style", "colors", "two_column"])


def simulate():
    LEN = 50
    lat = VO2_New(LEN, LEN)
    plot = Plotter(lat)
    si = SpinImage(lat.get_phases())
    mask_misk = np.ones((2*LEN, 2*LEN))
    ind = np.arange(mask_misk.size)
    np.random.shuffle(ind)
    mask_misk[np.unravel_index(
        ind[:int(mask_misk.size/2)], mask_misk.shape)] = 0
    print(mask_misk.shape)

    si.apply_mask(lat.parse_mask(np.ones((2*LEN, 2*LEN))))
    si.gaussian(20)
    intens_mono = si.fft()
    intens_mono.clean()

    si.apply_mask(lat.parse_mask(np.zeros((2*LEN, 2*LEN))))
    si.gaussian(20)
    intens_rutile = si.fft()
    intens_rutile.clean()

    si.apply_mask(lat.parse_mask(mask_misk))
    # si.apply_mask(mask_misk)
    si.gaussian(20)
    intens_mixed = si.fft()
    intens_mixed.clean()

    intens_rutile.save("intens_rut.npz")
    intens_mono.save("intens_mono.npz")
    intens_mixed.save("intens_mixed.npz")


def plot(fft, ax):
    return ax.imshow(
        fft.intens,
        extent=fft.extents(),
        norm=matplotlib.colors.LogNorm(vmin=1e-10, vmax=1),
        # norm=matplotlib.colors.Normalize(vmax=1, vmin=1e-10),
        cmap="magma",
        origin="lower"
    )


def norm(*intenss):
    max = 1e-10
    for intens in intenss:
        m = np.max(intens.intens)
        max = np.maximum(max, m)
    return max


def plot_all(intens_rutile, intens_mono, intens_mixed):
    fig, axs = plt.subplots(3, 2)
    fig.set_figheight(5.2)
    for ax in axs.flatten():
        ax.axis("off")
    axs = axs[:, 1]

    ax = axs[0]
    plot(intens_rutile, ax)

    ax = axs[1]
    plot(intens_mono, ax)
    y_shift = 0.175
    h_shift = 2*y_shift
    l_shift = 0.108
    big_shift = l_shift + h_shift
    # c = plt.Circle((-l_shift, y_shift), radius=0.07,
    #               label='patch', fill=False, ec="w", ls=":")
    # ax.add_patch(c)
    # c = plt.Circle((l_shift, -y_shift), radius=0.07,
    #               label='patch', fill=False, ec="w", ls=":")
    # ax.add_patch(c)

    # c = plt.Circle((-h_shift, -y_shift), radius=0.07,
    #               label='patch', fill=False, ec="w", ls=":")
    # ax.add_patch(c)
    # c = plt.Circle((h_shift, y_shift), radius=0.07,
    #               label='patch', fill=False, ec="w", ls=":")
    # ax.add_patch(c)

    ax.annotate("", (-l_shift, y_shift), (0, 0),
                arrowprops=dict(ec='w', facecolor='white', width=.8, headwidth=3, headlength=5))

    ax = axs[2]
    cmap = plot(intens_mixed, ax)

    cut_off = 0.8
    for ax in axs:
        ax.axis("off")
        ax.set_xlim(-cut_off, cut_off)
        ax.set_ylim(-cut_off, cut_off)
    plt.tight_layout()
    fig.subplots_adjust(bottom=0.1, right=0.95, left=0.15, wspace=0.)
    cbar_ax = fig.add_axes([0.55, 0.07, 0.4, 0.015])
    cbar = fig.colorbar(cmap, cax=cbar_ax,
                        orientation="horizontal", ticks=[1e-10, 1e-5, 1e0])
    #cbar.ax.set_xticklabels(['Low', 'Medium', 'High'])

    fig.savefig("erklaerbaer.pdf")
    fig.savefig("erklaerbaer.png")
    # Plotting cuts


def load():
    r = FFT()
    r.load("intens_rut.npz")
    mo = FFT()
    mo.load("intens_mono.npz")
    mi = FFT()
    mi.load("intens_mixed.npz")

    return r, mo, mi


if __name__ == "__main__":
    np.random.seed(1234)
    simulate()
    np.savez("intens.npz", r=r, mo=mo, mi=mi)
    r, mo, mi = load()
    max = norm(r, mo, mi)
    r.intens = r.intens/max
    mo.intens = mo.intens/max
    mi.intens = mi.intens/max
    plot_all(r, mo, mi)

    plt.show()
