#!/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)  # (rad/s)

A0 = 1.0
w0 = 1E1*FSR
f0 = w0 / 2 / pi
a_mod = 2
w_mod = w0 / 3000
f_mod = w_mod / 2 / pi


def R_e(w):
    T = 0.25
    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 main3():
    N = 2000000
    ts = np.linspace(0, 10000*(1/f0), N)
    mod = a_mod*np.sin(w_mod*ts)
    Es = A0*np.cos((w0 + (-0.012)*FSR)*ts + mod)
    maxima = ts[np.r_[True, Es[1:] < Es[:-1]] &
                np.r_[Es[:-1] < Es[1:], True]]
    maxima = maxima[1:-1]

    n_periods = 100
    periods = (maxima[n_periods:] - maxima[:-n_periods])/n_periods
    centers = maxima[:-n_periods] + (periods/2)
    powers = R_e(2*pi/periods)
    freqs = (2*pi/periods - w0)/FSR


    def moving_average(a, n=3):
        ret = np.cumsum(a, dtype=float)
        ret[n:] = ret[n:] - ret[:-n]
        return ret[n - 1:] / n

    smooth = 25
    freqs = moving_average(freqs, n=smooth)
    powers = moving_average(powers, n=smooth)
    centers = centers[smooth-1:]

    Rs_ws = np.linspace(w0-0.5*FSR, w0+0.5*FSR, 300)
    Rs_Rs = R_e(Rs_ws)
    Rs_ws = (Rs_ws - w0)/FSR

    fig, (ax1, ax2, ax3) = plt.subplots(nrows=3, ncols=1)

    n_frames = 300
    p_max = 0.16

    def do_plot(i):
        print(i)
        idx = int(i/n_frames * len(centers))
        t = centers[idx]

        f = freqs[idx]
        p = powers[idx]

        ax1.clear()
        ax1.plot(Rs_ws, Rs_Rs, 'b')
        ax1.plot([f], [p], 'r.')
        ax1.set_xlim((-0.1, 0.1))
        ax1.set_xlabel(r"$\Delta\omega(\mathrm{FSR})$")
        ax1.set_ylim((0, p_max))

        ax2.clear()
        ax2.plot(centers*f_mod, powers, 'b')
        ax2.set_ylabel('Refl. Power')
        ax2.set_ylim((0, p_max))
        ax2.set_xlim((0, ts[-1]*f_mod))
        ax2.set_xticks([])
        rect = plt.Rectangle((0, 0), t*f_mod, 1, facecolor='#aaaaaa')
        ax2.add_patch(rect)

        ax3.clear()
        ax3.plot(ts*f_mod, mod, 'r')
        ax3.set_ylabel('Modulation')
        ax3.set_ylim((-a_mod, a_mod))
        ax3.set_xlim((0, ts[-1]*f_mod))
        ax3.set_xlabel(r"$t(\tau_{\mathrm{mod}})$")
        rect = plt.Rectangle((0, -a_mod),
                             t*f_mod, 2*a_mod, facecolor='#aaaaaa')
        ax3.add_patch(rect)

    # do_plot(0)
    # plt.show()
    anim = animation.FuncAnimation(fig, do_plot, n_frames)
    anim.save('output.mp4', fps=30, dpi=270)


def main():
    N = 2000000
    ts = np.linspace(0, 10000*(1/f0), 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 = (N - win_size) // (n_groups-1)
    fig, (ax0, ax1, ax2) = plt.subplots(nrows=3, ncols=1)
    fig.set_tight_layout(True)

    Rs_ws = np.linspace(w0-0.5*FSR, w0+0.5*FSR, 300)
    Rs_Rs = R_e(Rs_ws)
    Rs_ws = (Rs_ws - w0)/FSR

    def do_plot(i):
        print(i)

        n = ts.size // n_groups
        win_start = 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 = 2*pi*np.fft.fftfreq(sec.size, d=dt)[:n//2]
        power = (sec_ft[:n//2]*sec_ft[:n//2].conj()).real / win_size
        # power = (sec_ft[:n//2].real**2 + sec_ft[:n//2].imag**2) / win_size

        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.clear()
        ax1.plot(freq/w0, power)
        ax1.set_xlim(0.80, 1.20)

        mean = np.average(freq, weights=power)
        print('{:2.2f} : {:2.2E}'.format((mean-w0)/FSR, R_e(mean*w0)))

        ax2.clear()
        ax2.plot(Rs_ws, Rs_Rs, 'b')
        ax2.plot((mean-w0)/FSR, R_e(mean), 'r.')
        ax2.set_xlim(-.5, .5)
        # ref_power = get_ref_power(freq, power)
        # ax2.plot(ts[(win_start+win_end) // 2], ref_power, 'r.')
        # ax2.set_xlim(0, ts[-1])

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


if __name__ == '__main__':
    # main()
    # main2()
    main3()