import numpy as np
import matplotlib.pyplot as plt


def analysis(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 play_1d():
    RESOLUTION = 0.001
    LENGTH = 10000

    x = np.arange(0, LENGTH, RESOLUTION)  # angstrom
    y = np.zeros_like(x)

    pos_mono = np.arange(0, x.size, 2890)
    pos_mono = (
        pos_mono + np.random.normal(size=pos_mono.shape, loc=0, scale=10)
    ).astype(int)

    pos_rut = np.arange(0, x.size, 5780)

    pos_rut = np.append(pos_rut, pos_rut - 3160)

    # pos_rut = (pos_rut + np.random.normal(size=pos_rut.shape, loc=0, scale=10)).astype(int)

    y[pos_rut] = 1
    y[pos_rut + 1] = 1
    y[pos_rut + 2] = 1

    # y = np.sin(x)

    fig, axs = plt.subplots(3, 1)
    ax = axs[0]
    ax.plot(x, y)

    ax = axs[1]
    fft_x, fft_y = analysis(y, RESOLUTION)
    ax.plot(fft_x, fft_y)
    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)


def from_mask(mask):
    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 image_from_pos(pos):
    RESOLUTION = 0.001
    LENGTH = 1000000

    x = np.arange(0, LENGTH, RESOLUTION)  # angstrom
    y = np.zeros_like(x)
    ind = np.searchsorted(x, pos)
    if np.any(ind > LENGTH):
        print("overflow")
    ind = ind[ind < LENGTH]
    y[ind] = 1

    sigma = 500
    mu = int(LENGTH / 2)
    gaussian = (
        1
        / (sigma * np.sqrt(2 * np.pi))
        * np.exp(-((x - mu) ** 2) / (2 * sigma * sigma))
    )
    # y = np.multiply(y, gaussian)

    return x, y


def plot_img(x, y, ax):
    ax.plot(x, y)


if __name__ == "__main__":
    RESOLUTION = 0.001
    print("Done")

    LENGTH = 1000
    mask_h = np.ones(LENGTH).astype(bool)
    pos_h = from_mask(mask_h)
    x, img_h = image_from_pos(pos_h)
    fftx, ffty_h = analysis(img_h, RESOLUTION)

    mask_l = np.zeros(LENGTH).astype(bool)
    pos_l = from_mask(mask_l)
    x, img_l = image_from_pos(pos_l)
    fftx, ffty_l = analysis(img_l, RESOLUTION)
    print("Done")

    mask_mixed = np.zeros(LENGTH).astype(bool)
    ind = (np.random.rand(400) * (LENGTH - 1)).astype(int)
    mask_mixed[ind] = True
    pos_mixed = from_mask(mask_mixed)
    x, img_mixed = image_from_pos(pos_mixed)
    fftx, ffty_mixed = analysis(img_mixed, RESOLUTION)

    print("Done")
    mask_near = np.zeros(LENGTH).astype(bool)
    ind = (np.random.rand(50) * (LENGTH - 1)).astype(int)
    #for i in range(1, 8):
    #    ind = np.append(ind, ind+i)
    print("Done")
    

    mask_near[ind] = True
    pos_near = from_mask(mask_near)
    x, img_near = image_from_pos(pos_near)
    fftx, ffty_near = analysis(img_near, RESOLUTION)

    fig, axs = plt.subplots(4, 1)
    plot_img(fftx, ffty_h, axs[0])
    plot_img(fftx, ffty_l, axs[1])
    plot_img(fftx, ffty_mixed, axs[2])
    plot_img(fftx, ffty_near, axs[3])
    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)
    # play_1d()
    plt.show()
    print("Done")
    pass