import numpy as np 
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from matplotlib import cm 

## define colours for plot
CB_color_cycle = ['#377eb8', '#ff7f00', '#4daf4a',
                  '#f781bf', '#a65628', '#984ea3',
                  '#999999', '#e41a1c', '#dede00']

def NSr_func(Ny, theta):
    return Ny / (2*np.pi*(1-np.cos(theta)))

## Set Scale and intial datapoint for graph axis
Ny = 7e11 
theta = 0.005 
NSr = Ny / (2*np.pi*(1-np.cos(theta)))
N_range = np.geomspace(NSr/20000,NSr*10000,1000)

## Set Scale and intial datapoint for graph axis
dx = 6.5e-6
z = 1
k = 1
dxzk = (dx**2 / z**2) * k
d_range = np.geomspace(dxzk/20000000,dxzk*100,1000)

## Construct Mesh Grid
N,D = np.meshgrid(N_range,d_range)

## Colormap for Plotting
Blues2 = cm.get_cmap('viridis', 255)
Blues2.set_over('w')
Blues2.set_under('w')

## Plot Main Figure
plt.figure(figsize=(10,5))
im = plt.pcolormesh(N, D , N*D,norm=mcolors.LogNorm(vmin=0.01,vmax=2**18),cmap=Blues2,shading='auto',rasterized='True')
plt.colorbar(im,label='Counts / px',extend='both')

CS = plt.contour(N, D, N*D,
            levels=[1], colors='k', linestyles='dashed')

CS = plt.contour(N, D, N*D,
            levels=[0.01,2**18], colors='k', linestyles='solid')

plt.axvline(NSr_func(7e11,5.08/1000),color=CB_color_cycle[0],ls='-')
plt.text(NSr_func(7e11,5.08/1000)*0.9,1e-10,'100 MeV',rotation=90, fontsize=9 ,color=CB_color_cycle[0], 
         bbox = {'facecolor':'white', 'alpha':0.8, 'edgecolor':'white'})
plt.axvline(NSr_func(1e12,0.51/1000),color=CB_color_cycle[1],ls='-')
plt.text(NSr_func(1e12,0.51/1000)*0.9,1e-10,'1 GeV',rotation=90, fontsize=9,color=CB_color_cycle[1],
         bbox = {'facecolor':'white', 'alpha':0.8, 'edgecolor':'white'})
plt.axvline(NSr_func(3e12,0.13/1000),color=CB_color_cycle[2],ls='-')
plt.text(NSr_func(3e12,0.13/1000)*0.9,1e-10,'10 GeV',rotation=90, fontsize=9,color=CB_color_cycle[2],
         bbox = {'facecolor':'white', 'alpha':0.8, 'edgecolor':'white'})

## horizontal bar - for Andor Zyla
dx = 6.5e-6
z = 2
k = 1
y = (dx**2 / z**2) * k
c = np.array([10,2**16])
x =  c/y 

plt.plot(x,[y,y],'ko-')
plt.text(x[0],y*1.4,'Andor Zyla @ 2m', fontsize=9)

## horizontal bar - for Hexitec
dx = 250e-6
z = 20
k = 1
y = (dx**2 / z**2) * k
c = np.array([0.1,2**8])
x =  c/y 

plt.plot(x,[y,y],'ko-')
plt.text(x[1],y*1.4,'Hexitec @ 20m', fontsize=9)

## horizontal bar - for Xineos 3030
dx = 100e-6
z = 100
k = 1
y = (dx**2 / z**2) * k
c = np.array([100,2**14])
x =  c/y 

plt.plot(x,[y,y],'ko-')
plt.text(x[0],y*1.4,'Xineos 3030 @ 100m', fontsize=9)

## horizontal bar - for MMPAD
dx = 150e-6
z = 2000
k = 1
y = (dx**2 / z**2) * k
c = np.array([0.1,1.6e6 / 3.56])
x =  c/y 

plt.plot(x,[y,y],'ko-')
plt.text(x[0]*10,y*1.4,'MMPAD @ 2km',fontsize=9)

## horizontal bar - for idealised detector
dx = 10e-6
z = 20
k = 100e-6
y = (dx**2 / z**2) * k
c = np.array([0.05,1.6e6 / 3.56])
x =  c/y 

plt.plot(x,[y,y],'wo-')
plt.text(x[0]*1000,y*1.4,'Idealised @ 20m',fontsize=9,color='k')

## figure formatting
plt.yscale('log')
plt.xscale('log')

plt.ylabel('$\kappa \cdot dx^2 / z^2$ ')
plt.xlabel('$N_\gamma$/Sr ')

plt.xlim([N_range[0],N_range[-1]])
plt.show()