import numpy as np
import matplotlib.pyplot as plt
import tqdm

import multiprocessing as mp

RESOLUTION = 0.01
LENGTH = 500


def generate_image_from_mask(mask: np.array):
    pos_mono = np.arange(0, mask.size * 2.89 * 2, 2.89)
    pos_mono[::2][mask] -= 0.27
    pos_mono[1::2][mask] += 0.27
    return pos_mono


def pad_zero(img, length):
    pad = np.zeros(length)
    img = np.append(img, pad)
    img = np.append(pad, img)
    return img


def image_from_pos(pos):
    length = np.max(pos) + RESOLUTION
    x = np.arange(0, length, RESOLUTION)  # angstrom
    y = np.zeros_like(x)
    ind = np.searchsorted(x, pos)
    y[ind] = 1
    return y


def beugung(y, resolution):
    fft = np.fft.fft(y)
    fft_clean = np.fft.fftshift(fft)
    fft_freq = np.fft.fftfreq(y.size, resolution)
    fft_freq_clean = np.fft.fftshift(fft_freq)
    return fft_freq_clean, np.abs(fft_clean) ** 2


def gaussian_convol(img):
    sigma = 100 / RESOLUTION
    mu = img.size/2
    x = np.arange(0, img.size)
    gauss = 1/(sigma * np.sqrt(2 * np.pi)) * \
        np.exp(- (x - mu)**2 / (2 * sigma**2))

    return img*gauss


def analyisis(mask):
    pos_h = generate_image_from_mask(mask)
    img = image_from_pos(pos_h)
    img = gaussian_convol(img)
    padded = pad_zero(img, int(100 / RESOLUTION))
    freq, intens = beugung(padded, RESOLUTION)
    return freq, intens


def get_peaks():
    orders = np.arange(1, 2, 1)
    orders = orders / 5.78
    return np.array(orders)


def eval_peaks(freq, fft):
    orders = get_peaks()
    ind = np.searchsorted(freq, orders)
    return fft[ind]


def basic_test():
    mask_h = np.zeros(LENGTH).astype(bool)
    mask_l = np.ones(LENGTH).astype(bool)

    mask_mixed = np.zeros(LENGTH).astype(bool)
    ind = (np.random.rand(30) * (LENGTH - 1)).astype(int)
    mask_mixed[ind] = True

    mask_ner = np.zeros(LENGTH).astype(bool)
    ind = (np.random.rand(1) * (LENGTH - 31)).astype(int)
    ind = np.arange(ind, ind+30).astype(int)
    mask_ner[ind] = True

    fig, axs = plt.subplots(4, 1)
    for mask, ax in zip([mask_h, mask_l, mask_mixed, mask_ner], axs):
        freq, ffty = analyisis(mask)
        ax.plot(freq, ffty)
    for ax in axs:
        ax.plot([1.0 / 5.78, 1.0 / 2.62, 1.0 / 3.16], [0, 0, 0], "kx")
        ax.plot([2.0 / 5.78, 2.0 / 2.62, 2.0 / 3.16], [0, 0, 0], "rx")
        ax.plot([3.0 / 5.78, 3.0 / 2.62, 3.0 / 3.16], [0, 0, 0], "bx")
        ax.set_xlim(0, 3)
    plt.show()


def norm(arr):
    return arr/np.sum(arr)


def norm2(arr):
    return arr
    # return arr/np.max(arr)


def next_mask(mask):
    prob = np.exp((np.roll(mask, 1)*1.0 + np.roll(mask, -1)) / .1)
    prob[mask] = 0.0
    prob = norm(prob)

    ind = np.random.choice(LENGTH, p=prob)
    mask[ind] = True
    return mask


def random_loop():
    mask = np.zeros(LENGTH).astype(bool)
    ind = np.arange(0, LENGTH)
    np.random.shuffle(ind)
    percentage = []
    peaks = []

    masks = []
    for i in ind:
        mask[i] = True
        freq, fft = analyisis(mask)
        peak = eval_peaks(freq, fft)
        percentage.append(np.mean(mask))
        peaks.append(peak)
        masks.append(mask.copy())
    masks = np.array(masks)
    plt.figure()
    plt.imshow(masks)
    plt.plot([0, 500], [406, 406])
    print()
    percentage = np.array(percentage)
    peaks = np.array(peaks)
    return percentage, peaks


def nearest_loop():
    mask = np.zeros(LENGTH).astype(bool)
    percentage = []
    peaks = []
    for i in range(LENGTH):
        mask = next_mask(mask)
        freq, fft = analyisis(mask)
        peak = eval_peaks(freq, fft)
        percentage.append(np.mean(mask))
        peaks.append(peak)
    percentage = np.array(percentage)
    peaks = np.array(peaks)
    return percentage, peaks


def random_helper(seed):
    np.random.seed(seed)
    #percentage_near, peaks_near = nearest_loop()
    percentage_rand, peaks_rand = random_loop()
    print("done")
    return percentage_rand, peaks_rand
    # for i in range(peaks_near.shape[1]):
    #    axs[2].plot(percentage_near, norm2(
    #        peaks_near[:, i]), "-", label="near")

    # for i in range(peaks_rand.shape[1]):
    #    axs[2].plot(percentage_rand, norm2(
    #        peaks_rand[:, i]), ":", label="rand")


def random_increase():
    fig, axs = plt.subplots(3, 1)

    results = []
    for i in np.arange(10):
        results.append(random_helper(i))

    for percentage_rand, peaks_rand in results:
        for i in range(peaks_rand.shape[1]):
            axs[2].plot(percentage_rand, norm2(
                peaks_rand[:, i]), ":", label="rand")

    for ax in [axs[0], axs[1]]:
        orders = get_peaks()
        ax.plot(orders, np.zeros_like(orders), "kx")
        ax.set_xlim(0, 3)

    mask_l = np.ones(LENGTH).astype(bool)
    mask_h = np.zeros(LENGTH).astype(bool)
    freq, ffty = analyisis(mask_l)
    axs[0].plot(freq, ffty)
    freq, ffty = analyisis(mask_h)
    axs[1].plot(freq, ffty)
    plt.xlabel("percentage")
    plt.ylabel("peak intensity")
    plt.show()
    plt.legend()


if __name__ == "__main__":
    random_increase()