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


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 = 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()