#!/usr/bin/env python3
import matplotlib
import matplotlib.pyplot as plt
from matplotlib import animation
import numpy as np
from numpy import cos, pi

matplotlib.rcParams.update({'font.size': 15})

c = 3E8  # m/s
L = 1.0  # m
FSR = 2*pi*c/(2*L)  # Hz

A0 = 1.0
w0 = 1E9*FSR
f0 = w0 / 2 / pi
a_mod = 200
w_mod = 0.0002 * w0
f_mod = w_mod / 2 / pi


def R_e(w):
    T = 0.05
    R = 1-T
    return 1 - T**2 / (1 + R**2 - 2*R*cos(w*2*L/c))


def get_ref_power(freqs, powers):
    return sum(power*R_e(freq) for freq, power in zip(freqs, powers))


def main():
    sample_n = 2000000
    ts = np.linspace(0, 10000*(1/f0), sample_n)
    dt = ts[1]-ts[0]
    mod = a_mod*np.sin(w_mod*ts)
    Es = A0*np.cos(w0*ts + mod)
    print("FSR", FSR)
    print(1/f0, dt)

    n_groups = 100
    win_size = int((1/f_mod) / dt / 12)
    win_spacing = (sample_n - win_size) // (n_groups+2)
    fig, (ax0, ax1, ax2) = plt.subplots(nrows=3, ncols=1)

    def do_plot(i):
        print(i)
        # ax0.clear()
        ax1.clear()
        n = ts.size // n_groups
        win_start = (win_size//2) + i*win_spacing
        win_end = win_start + win_size
        print(f"{win_start},{win_end}")

        # sec_ts = ts[win_start:win_end]
        sec = Es[win_start:win_end]
        # mod_sec = mod[win_start:win_end]
        sec_ft = np.fft.rfftn(sec)
        freq = np.fft.fftfreq(sec.size, d=dt)[:n//2]
        power = (sec_ft[:n//2].real**2 + sec_ft[:n//2].imag**2) / win_size

        # ax0.plot(sec_ts, mod_sec)
        ax0.plot(ts[(win_start+win_end) // 2], mod[(win_start+win_end) // 2], 'k.')
        ax0.set_ylim(-1.1*a_mod, 1.1*a_mod)
        ax0.set_xlim(0, ts[-1])

        ax1.plot(2*pi*freq/w0, power)
        ax1.set_xlim(0.80, 1.20)

        ref_power = get_ref_power(freq, power)
        ax2.plot(ts[(win_start+win_end) // 2], ref_power, 'r.')
        ax2.set_xlim(0, ts[-1])

        plt.tight_layout()

    anim = animation.FuncAnimation(fig, do_plot, n_groups)
    anim.save("output.mp4")


if __name__ == '__main__':
    main()