from lattices import SCC_Lattice, VO2_Lattice
from spin_image import SpinImage
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import matplotlib
import tqdm
from extractors import Rect_Evaluator, Voronoi_Evaluator, Image_Wrapper


class Plotter:
    def __init__(self, lat):
        self.lattice = lat
        self.length_2 = 0.05

    def reduce(self, arr):
        arr = np.array(arr)
        arr = arr.flatten()
        return np.mean(arr)
        # return np.sum(arr[np.argpartition(arr, -8)[-8:]])

    def rect_at_point(self, x, y, color):
        length_2 = self.length_2
        rect = patches.Rectangle(
            (x - length_2, y - length_2),
            2 * length_2,
            2 * length_2,
            linewidth=1,
            edgecolor=color,
            facecolor="none",
        )
        return rect

    def plot(self, freqx, freqy, intens, ax_log=None, ax_lin=None, vmax=None):
        if ax_log:
            t = ax_log.imshow(
                intens,
                extent=(np.min(freqx), np.max(freqx),
                        np.min(freqy), np.max(freqy)),
                norm=matplotlib.colors.LogNorm(vmin=10),
                cmap="viridis",
                origin="lower"
            )
            plt.colorbar(t, ax=ax_log)
        if ax_lin:
            t = ax_lin.imshow(
                intens,
                extent=(np.min(freqx), np.max(freqx),
                        np.min(freqy), np.max(freqy)),
                # vmax=vmax,
                cmap="viridis",
                origin="lower"
            )
            plt.colorbar(t, ax=ax_lin)


def rotate(x, y, angle):
    radian = angle / 180 * 2 * np.pi
    return np.cos(radian) * x - np.sin(radian) * y,\
        np.sin(radian) * x + np.cos(radian) * y


def test_square():
    LEN = 40
    # lat = SCC_Lattice(LEN, LEN)
    lat = VO2_Lattice(LEN, LEN)
    plot = Plotter(lat)
    pos_x, pos_y = lat.get_from_mask(np.zeros((40, 40)))
    # pos_x, pos_y = rotate(pos_x, pos_y, 30)
    si = SpinImage(pos_x, pos_y)
    fig, axs = plt.subplots(2, 2)
    si.pad_it_square(10)
    si.plot(axs[0, 0], 2)
    # si.gaussian(LEN)
    # si.blur(3)
    si.plot(axs[0, 1], 2)

    plt.pause(0.1)
    fx, fy, intens = si.fft()
    plot.plot(fx, fy, intens, axs[1, 0], axs[1, 1])
    plt.savefig("test.png")
    plt.show()


def helper(intens, fx, fy, voro, rect):
    print(np.mean(intens))
    v_idx, v_val = voro.extract(Image_Wrapper(intens, fx, fy))
    r_idx, r_val = rect.extract(Image_Wrapper(intens, fx, fy))

    for iv, av, ir, ar in zip(v_idx, v_val, r_idx, r_val):
        av = np.array(av)
        ar = np.array(ar)
        mask = np.isin(ir, iv)
        print("Test: ", np.all(np.isclose(
            ar[mask], av)), np.isclose(ar[mask], av))
        print(iv, "\n",  ar[mask], "\n", av, "\n\n\n")


def test_mixed():
    LEN = 40
    lat = VO2_Lattice(LEN, LEN)
    plot = Plotter(lat)
    all_rutile = np.stack(lat.reci()[0]).T
    all_mono = np.stack(lat.reci()[1]).T
    all_mono2 = np.stack(lat.reci()[2]).T
    voro = Voronoi_Evaluator([all_rutile, all_mono, all_mono2])
    rect = Rect_Evaluator([all_rutile, all_mono, all_mono2])
    #voro = Voronoi_Evaluator([all_mono])
    #rect = Rect_Evaluator([all_mono])

    pos_x, pos_y = lat.get_from_mask(np.zeros((40, 40)))
    si = SpinImage(pos_x, pos_y)
    si.pad_it_square(10, size=2300)
    fx, fy, intens_mono = si.fft()

    pos_x, pos_y = lat.get_from_mask(np.ones((40, 40)))
    si = SpinImage(pos_x, pos_y)
    si.pad_it_square(10, size=2300)
    fx, fy, intens_rutile = si.fft()

    mask_misk = np.ones((40, 40))
    ind = np.arange(mask_misk.size)
    np.random.shuffle(ind)
    mask_misk[np.unravel_index(ind[:800], (40, 40))] = False
    pos_x, pos_y = lat.get_from_mask(mask_misk)
    si = SpinImage(pos_x, pos_y)
    si.pad_it_square(10, size=2300)
    fx, fy, intens_mixed = si.fft()
    img = Image_Wrapper(intens_mono, fx, fy)
    voro.generate_mask(img)
    rect.generate_mask(Image_Wrapper(intens_mono, fx, fy))

    fig, axs = plt.subplots(1, 2)
    axs[0].imshow(rect.get_mask(), extent=img.ext(), origin="lower")
    axs[0].plot(all_rutile[:, 0], all_rutile[:, 1], ".")
    axs[1].plot(all_rutile[:, 0], all_rutile[:, 1], ".")
    axs[0].plot(all_mono[:, 0], all_mono[:, 1], ".")
    axs[1].plot(all_mono[:, 0], all_mono[:, 1], ".")
    axs[0].plot(all_mono2[:, 0], all_mono2[:, 1], ".")
    axs[1].plot(all_mono2[:, 0], all_mono2[:, 1], ".")
    axs[1].imshow(voro.get_mask(), extent=img.ext(), origin="lower")

    print("mono")
    helper(intens=intens_mono, fx=fx, fy=fy, voro=voro, rect=rect)
    print("mixed")
    helper(intens=intens_mixed, fx=fx, fy=fy, voro=voro, rect=rect)
    print("rutile")
    helper(intens=intens_rutile, fx=fx, fy=fy, voro=voro, rect=rect)

    new_intens_mono = rect.debug(Image_Wrapper(intens_mono, fx, fy))
    new_intens_mixed = rect.debug(Image_Wrapper(intens_mixed, fx, fy))
    new_intens_rutile = rect.debug(Image_Wrapper(intens_rutile, fx, fy))

    fig, axs = plt.subplots(2, 3)
    plot.plot(freqx=fx, freqy=fy, intens=new_intens_rutile,
              ax_log=axs[0, 0], ax_lin=axs[1, 0], vmax=10e7)
    plot.plot(freqx=fx, freqy=fy, intens=new_intens_mono,
              ax_log=axs[0, 2], ax_lin=axs[1, 2], vmax=10e7)
    plot.plot(freqx=fx, freqy=fy, intens=new_intens_mixed,
              ax_log=axs[0, 1], ax_lin=axs[1, 1], vmax=10e7)

    for ax in axs.flatten():
        ax.set_xlim(-1, 1)
        ax.set_ylim(-1, 1)

    new_intens_mono = voro.debug(Image_Wrapper(new_intens_mono, fx, fy))
    new_intens_mixed = voro.debug(Image_Wrapper(intens_mixed, fx, fy))
    new_intens_rutile = voro.debug(Image_Wrapper(intens_rutile, fx, fy))

    fig, axs = plt.subplots(2, 3)
    plot.plot(freqx=fx, freqy=fy, intens=new_intens_rutile,
              ax_log=axs[0, 0], ax_lin=axs[1, 0], vmax=10e7)
    plot.plot(freqx=fx, freqy=fy, intens=new_intens_mono,
              ax_log=axs[0, 2], ax_lin=axs[1, 2], vmax=10e7)
    plot.plot(freqx=fx, freqy=fy, intens=new_intens_mixed,
              ax_log=axs[0, 1], ax_lin=axs[1, 1], vmax=10e7)

    for ax in axs.flatten():
        ax.set_xlim(-1, 1)
        ax.set_ylim(-1, 1)


def random(seed):
    np.random.seed(seed)
    LEN = 40
    lat = VO2_Lattice(LEN, LEN)
    maske = np.zeros((LEN, LEN))
    ind = np.arange(LEN * LEN)
    np.random.shuffle(ind)
    all_rutile = np.stack(lat.reci()[0]).T
    all_mono = np.stack(lat.reci()[1]).T
    all_mono2 = np.stack(lat.reci()[2]).T
    voro = Voronoi_Evaluator([all_rutile, all_mono, all_mono2])
    rect = Rect_Evaluator([all_rutile, all_mono, all_mono2])

    out_rect = [[] for x in range(len(lat.reci())+1)]
    out_voro = [[] for x in range(len(lat.reci())+1)]
    percentage = []
    counter = 0
    already_inited = False
    for i in tqdm.tqdm(ind):
        maske[np.unravel_index(i, (LEN, LEN))] = True
        counter += 1
        if np.mod(counter, 100) != 0:
            continue

        pos_x, pos_y = lat.get_from_mask(maske)
        si = SpinImage(pos_x, pos_y)
        si.pad_it_square(10, size=2300)
        fx, fy, intens = si.fft()
        img = Image_Wrapper(intens, fx, fy)
        if not already_inited:
            voro.generate_mask(img)
            voro.merge_mask()
            rect.generate_mask(img)
            rect.merge_mask()
            already_inited = True

        iv, vv = voro.extract(img)
        ir, vr = rect.extract(img)
        for lis, val in zip(out_rect, vr):
            lis.append(val)
        for lis, val in zip(out_voro, vv):
            lis.append(val)
        percentage.append(np.sum(maske))

    percentage = np.array(percentage)
    percentage /= np.max(percentage)

    np.savez(f"random_rect_{seed}.npz",
             percentage=percentage, out_1=out_rect[0],
             out_2=out_rect[1], out_3=out_rect[2], out_4=out_rect[3])
    np.savez(f"random_voro_{seed}.npz",
             percentage=percentage, out_1=out_voro[0],
             out_2=out_voro[1], out_3=out_voro[2], out_4=out_voro[3])


def sample_index(p):
    i = np.random.choice(np.arange(p.size), p=p.ravel())
    return np.unravel_index(i, p.shape)


def ising(seed):
    np.random.seed(seed)
    LEN = 40
    temp = 0.1
    maske = np.zeros((LEN, LEN), dtype=bool)

    lat = VO2_Lattice(LEN, LEN)
    all_rutile = np.stack(lat.reci()[0]).T
    all_mono = np.stack(lat.reci()[1]).T
    all_mono2 = np.stack(lat.reci()[2]).T
    voro = Voronoi_Evaluator([all_rutile, all_mono, all_mono2])
    rect = Rect_Evaluator([all_rutile, all_mono, all_mono2])

    out_rect = [[] for x in range(len(lat.reci())+1)]
    out_voro = [[] for x in range(len(lat.reci())+1)]
    percentage = []
    counter = 0
    already_inited = False
    for i in tqdm.tqdm(range(LEN*LEN)):
        probability = np.roll(maske, 1, axis=0).astype(float)
        probability += np.roll(maske, -1, axis=0).astype(float)
        probability += np.roll(maske, 1, axis=1).astype(float)
        probability += np.roll(maske, -1, axis=1).astype(float)

        probability = np.exp(probability/temp)
        probability[maske] = 0
        probability /= np.sum(probability)
        maske[sample_index(probability)] = True

        counter += 1
        if np.mod(counter, 100) != 0:
            continue

        pos_x, pos_y = lat.get_from_mask(maske)
        si = SpinImage(pos_x, pos_y)
        si.pad_it_square(10, size=2300)
        fx, fy, intens = si.fft()
        img = Image_Wrapper(intens, fx, fy)
        if not already_inited:
            voro.generate_mask(img)
            voro.merge_mask()
            rect.generate_mask(img)
            rect.merge_mask()
            already_inited = True

        iv, vv = voro.extract(img)
        ir, vr = rect.extract(img)
        for lis, val in zip(out_rect, vr):
            lis.append(val)
        for lis, val in zip(out_voro, vv):
            lis.append(val)
        percentage.append(np.sum(maske))

    percentage = np.array(percentage, dtype=np.float64)
    percentage /= np.max(percentage)

    np.savez(f"ising_rect_{seed}.npz",
             percentage=percentage, out_1=out_rect[0],
             out_2=out_rect[1], out_3=out_rect[2], out_4=out_rect[3])
    np.savez(f"ising_voro_{seed}.npz",
             percentage=percentage, out_1=out_voro[0],
             out_2=out_voro[1], out_3=out_voro[2], out_4=out_voro[3])


if __name__ == "__main__":
    # test_square()
    #test_mixed()
    #plt.show()
    # random()
    np.random.seed(1234)
    for i in np.random.randint(0, 10000, 1):
        random(i)
        ising(i)

    plt.show()