import logging
import scipy.fftpack as sfft
from plotter import Plotter
from scipy import signal
from cache import timeit
from extractors import Rect_Evaluator
import tqdm
from lattices import SCC_Lattice, VO2_Lattice, VO2_New
import sys
from spin_image import SpinImage
import numpy as np
import matplotlib.pyplot as plt
plt.style.use(["style", "colors", "two_column"])
logger = logging.getLogger('fft')
# logger.setLevel(logging.DEBUG)
ch = logging.StreamHandler()
ch.setLevel(logging.DEBUG)
formatter = logging.Formatter(
    '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
ch.setFormatter(formatter)
logger.addHandler(ch)


def random(seed):
    np.random.seed(seed)
    LEN = 40
    #lat = VO2_New(LEN, LEN)
    lat = VO2_Lattice(LEN, LEN)
    maske = np.zeros((LEN, LEN))
    ind = np.arange(maske.size)
    np.random.shuffle(ind)
    rect = Rect_Evaluator(lat.get_spots())

    out_rect = [[] for x in range(4)]
    percentage = []
    weighted_percentage = []
    counter = 0
    si = SpinImage(lat.get_phases())
    already_inited = False
    for i in tqdm.tqdm(ind):
        maske[np.unravel_index(i, maske.shape)] = True
        counter += 1
        if np.mod(counter, 100) != 0 and i != ind[-1] and i != ind[0]:
            continue

        print(counter, i)

        si.apply_mask(lat.parse_mask(maske))
        si.gaussian(20)

        intens = si.fft()
        if not already_inited:
            rect.generate_mask(intens, merge=True)
            already_inited = True

        ir, vr = rect.extract(intens)
        for lis, val in zip(out_rect, vr):
            lis.append(val)
        percentage.append(np.sum(maske))
        [p1, p2] = si.get_intens(lat.parse_mask(maske))
        weighted_percentage.append(p1/(p1+p2))

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

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

def runner():
    np.random.seed(1234)
    seeds = np.random.randint(0, 10000, 200)
    idx = int(sys.argv[1])
    if idx < 0:
        return
    if idx >= seeds.size:
        return
    seed = seeds[idx]
    random(seed)


if __name__ == "__main__":
    np.random.seed(1234)
    runner()