import matplotlib.pyplot as plt
import numpy as np
import matplotlib.colors as mcolors 
from scipy.constants import c

### supporting functions
def gamma(e, m0=0.511):
    return e/m0 + 1


def beta(g):
    return np.sqrt(1-(1/g)**2)


### plotting
mn = 939  # MeV of a neutron
mk = np.linspace(1, 100, 501)  # kinetic energy
r = np.geomspace(0.2, 100, 501)
R, MK = np.meshgrid(r, mk)

Arrival = R/(beta(gamma(MK, m0=mn))*c)
C_Arrival = R/c

cmap = mcolors.Colormap('Blues', 24)

plt.contourf(mk, r, (Arrival-C_Arrival).T*1e9, norm=mcolors.LogNorm(),
             cmap='Blues', levels=[*np.geomspace(1, 10000, 13)])
plt.ylabel('Detector Distance (m)')
plt.xlabel('Neutron Energy (MeV)')
plt.colorbar(label='Arrival Delta, $t_n-t_g$ (ns)')


# # limits
# limit=[]
# # gate_time = 0.45
# # scintillator = 0.7
# # limit.append(np.sqrt(gate_time**2 + scintillator**2))

detector = 17.9
print(detector)
MHzLimits = 1e9/np.array([1e6, 10e6, 100e6])

CS = plt.contour(mk, r, (Arrival-C_Arrival).T*1e9,
                 levels=MHzLimits[::-1], colors=['k'])
CS1 = plt.contour(mk, r, (Arrival-C_Arrival).T*1e9,
                  levels=[detector], colors=['r'], linestyles=['--'])
fmt = {}
strs = ['100 MHz \n 10 ns', '10 MHz \n 100 ns', '1 MHz \n 1000 ns']
for l, s in zip(CS.levels, strs):
    fmt[l] = s
# plt.axhline(20,color='b',alpha=0.4)
# plt.gca().clabel(CS, CS.levels[:-1], inline=True,
#                  inline_spacing=1, fmt=fmt, fontsize=8)
labs2 = plt.gca().clabel(CS, [CS.levels[-2]], fmt=fmt,
                         inline_spacing=-1, inline=True, fontsize=9, manual=[(40, 40)])
labs2[-1].set_rotation(10)

labs1 = plt.gca().clabel(CS, [CS.levels[-3]], fmt=fmt,
                         inline_spacing=-1, inline=True, fontsize=9, manual=[(40, 40)])
labs1[-1].set_rotation(10)

labs = plt.gca().clabel(CS, [CS.levels[-1]], fmt=fmt,
                        inline_spacing=10, inline=True, fontsize=9, manual=[(40, 40)])
labs[-1].set_rotation(30)
plt.yscale('log')

plt.gca().clabel(CS1, CS1.levels, fmt={
    CS1.levels[0]: '56 MHz \n 18 ns'}, fontsize=9, inline_spacing=-2)
plt.tight_layout()
plt.show()
# plt.savefig('./figures/figure1_timeofflightdetectors_contourf_withtiming_delta.pdf',dpi=800)
