#!/usr/bin/env python3
import matplotlib
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1.inset_locator import zoomed_inset_axes
from mpl_toolkits.axes_grid1.inset_locator import mark_inset
import numpy as np
from math import cos, pi

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

c = 3E8  # m/s
L = 1  # m
# λ = 1E-6  # m
# ω = 2*pi*λ/c  # s^-1
# t0 = λ*(3759398496+0.5)
FSR = 2*pi*c/(2*L)

ω = FSR


def T_e(T, δω):
    R = 1-T
    return T**2 / (1 + R**2 - 2*R*cos(2*L*(ω+δω*FSR)/c))


def R_e(T, δω):
    return 1 - T_e(T, δω)


def main():
    T = 0.15
    δωs = np.linspace(-0.5, 1.5, 5000)
    R_es = [R_e(T, δω) for δω in δωs]
    ax1 = plt.gca()
    ax1.plot(δωs, R_es)

    ax1.set_ylabel("$R_e$")
    ax1.set_xlabel("$\Delta \omega$ (FSR)")
    ax1.set_ylim((0, 1))

    ax1.text(1.15, 0.85, "$\mathcal{F}_c=19.3$",
             bbox={'facecolor': 'grey', 'alpha': 0.3})

    # fmt = '%.0f%%'
    # xticks = matplotlib.ticker.FormatStrFormatter(fmt)
    # ax1.yaxis.set_major_formatter(xticks)
    ax1.yaxis.set_major_formatter(matplotlib.ticker.FuncFormatter(
        lambda x, pos: '{:d}%'.format(int(x*100))))

    axins = zoomed_inset_axes(ax1, 8, bbox_to_anchor=(0.7, 0.5),
                              bbox_transform=ax1.transAxes)
    axins.plot(δωs, R_es)
    axins.set_xlim(-0.03, 0.03)
    axins.set_ylim(0, 0.05)
    axins.set_yticks([])
    axins.set_xticks([])

    mark_inset(ax1, axins, loc1=2, loc2=4, fc="none", ec="0.5")

    plt.tight_layout()
    plt.savefig("pdh_demo.png")
    # plt.show()


if __name__ == '__main__':
    main()
