import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from scipy.optimize import curve_fit

A = 1e-10   # Angstroms
um = 1e-6
mm = 1e-3

def normalize(arr):
    return arr/np.max(arr)

def get_layer_index(x, layerBounds):
    ''' Function to assign a layer/panel index to each ion based on X position '''
    for idx, bounds in enumerate(layerBounds):
        if bounds[0] <= x*A/um < bounds[1]:
            return idx
    return len(layerBounds)


# read file
path = '../Simulation Data/'
fname = 'TRIM_data.csv'

df = pd.read_csv(path+fname)

# create column in dataset which contains the initial energy of the current proton
first_instance = df.drop_duplicates(subset='Ion no.', keep='first')
df['Energy init'] = first_instance['Energy (keV)']
df['Energy init'] = df['Energy init'].fillna(method = 'ffill')

# calculate maximum proton range for each initial proton energy
ebins = np.unique(df['Energy init'])
pRange = np.zeros_like(ebins)
for i, e in enumerate(ebins):
    pRange[i] = np.max(df['X'].loc[df['Energy init']==e])

Range = (pRange*A/mm)   # convert from Angstroms to mm
EBins = (ebins/1e3)   # convert from keV to MeV

# fit proton ranges with a power law function
maxProtonEnergy = lambda x,A,b: ((1/A)*x)**(1/b)
params, pcov = curve_fit(maxProtonEnergy, Range, EBins, p0 = (10, 0.5))


fitLayers = np.linspace(0, 5000, 50000)

nPanels = 8
layers = np.arange(200)
fiberThicknesses = [0.5, 0.25 , 0.1]

colors = ['r', 'purple', 'blue', 'g']
handles = []

fig, ax = plt.subplots(1,1, figsize = (3.5, 2.5))
fig.set_dpi(150)

for i, fThick in enumerate(fiberThicknesses):

    # calculate energy resolution for each fiber thickness 
    midLayerDepthsSmooth = (fitLayers+0.5)*fThick*nPanels
    dE = np.array(maxProtonEnergy((fitLayers+1)*nPanels*fThick, *params)) - np.array(maxProtonEnergy((fitLayers)*nPanels*fThick, *params))
    midLayerEnergiesSmooth = np.array(maxProtonEnergy(midLayerDepthsSmooth, *params))

    # define low and high resolution regions
    lowRes = dE/midLayerEnergiesSmooth > 0.1
    highRes = dE/midLayerEnergiesSmooth <= 0.1

    # add plot to axis
    ax.plot(midLayerEnergiesSmooth[lowRes], dE[lowRes]/midLayerEnergiesSmooth[lowRes], '-', color = colors[i], lw = 1.5, alpha = 0.3)
    fplot, = ax.plot(midLayerEnergiesSmooth[highRes], dE[highRes]/midLayerEnergiesSmooth[highRes], '-', color = colors[i], lw = 1.5, alpha = 1, label = f'$d_f={int(fThick*1e3)}\ \mu m$')
    handles.append(fplot)
    

# RCF HDV2 plot
HDV2filterThick = 0.095
HDV2actLayerThick = 0.012
HDV2LayerThick = HDV2actLayerThick + HDV2filterThick 

HDV2Emin = maxProtonEnergy(layers*HDV2LayerThick, *params)
HDV2Emax = maxProtonEnergy((layers+1)*HDV2actLayerThick + layers*HDV2filterThick, *params)
dEHDV2 = np.array(HDV2Emax) - np.array(HDV2Emin)

# add resolution
midLayerEnergiesHDV2 = np.array(HDV2Emin)+dEHDV2/2

HDV2EminSmooth = maxProtonEnergy(fitLayers*HDV2LayerThick, *params)
HDV2EmaxSmooth = maxProtonEnergy((fitLayers+1)*HDV2actLayerThick + fitLayers*HDV2filterThick, *params)
dEHDV2_smooth = np.array(HDV2EmaxSmooth) - np.array(HDV2EminSmooth)
midLayerEnergiesHDV2_smooth = np.array(HDV2EminSmooth)+dEHDV2_smooth/2

lowRes = dEHDV2_smooth/midLayerEnergiesHDV2_smooth > 0.1
highRes = dEHDV2_smooth/midLayerEnergiesHDV2_smooth <= 0.1

ax.plot(midLayerEnergiesHDV2_smooth[lowRes], dEHDV2_smooth[lowRes]/midLayerEnergiesHDV2_smooth[lowRes], '-', color = 'k', lw = 1.5, alpha = 0.3)
hdv2, = ax.plot(midLayerEnergiesHDV2_smooth[highRes], dEHDV2_smooth[highRes]/midLayerEnergiesHDV2_smooth[highRes], '-', color = 'k', lw = 1.5, alpha = 1, label = 'HDV2')

# RCF EBT3 plot
EBT3filterThick = 0.125
EBT3actLayerThick = 0.028
EBT3LayerThick = EBT3actLayerThick + EBT3filterThick 

EBT3Emin = maxProtonEnergy(layers*EBT3LayerThick + EBT3filterThick/2, *params)
EBT3Emax = maxProtonEnergy(layers*EBT3LayerThick + EBT3filterThick/2 + EBT3actLayerThick, *params)
dEEBT3 = np.array(EBT3Emax) - np.array(EBT3Emin)

# add resolution
midLayerEnergiesEBT3 = np.array(EBT3Emax)+dEEBT3/2

EBT3EminSmooth = maxProtonEnergy(fitLayers*EBT3LayerThick, *params)
EBT3EmaxSmooth = maxProtonEnergy((fitLayers+1)*EBT3actLayerThick + fitLayers*EBT3filterThick, *params)
dEEBT3_smooth = np.array(EBT3EmaxSmooth) - np.array(EBT3EminSmooth)
midLayerEnergiesEBT3_smooth = np.array(EBT3EminSmooth)+dEEBT3_smooth/2

lowRes = dEEBT3_smooth/midLayerEnergiesEBT3_smooth > 0.1
highRes = dEEBT3_smooth/midLayerEnergiesEBT3_smooth <= 0.1

ax.plot(midLayerEnergiesEBT3_smooth[lowRes], dEEBT3_smooth[lowRes]/midLayerEnergiesEBT3_smooth[lowRes], '-', color = 'grey', lw = 1.5, alpha = 0.3)
ebt3, = ax.plot(midLayerEnergiesEBT3_smooth[highRes], dEEBT3_smooth[highRes]/midLayerEnergiesEBT3_smooth[highRes], '-', color = 'grey', lw = 1.5, alpha = 1, label = 'EBT3')

ax.hlines(1e-1, xmin=-10, xmax=200, alpha = 1, ls = '--', color = 'k', lw = 1)

ax.set_yscale('log')
ax.set_yticks([1e-4, 1e-3, 1e-2, 1e-1, 1e0])
ax.set_yticklabels([0.01, 0.1, 1, 10, 100])
ax.set_ylim(1e-4,1e0)
ax.set_xlim(-9,150)
ax.set_xlabel('Layer center energy, $E_{mid}$ (MeV)')
ax.set_ylabel(r'$dE/E_{mid}$ (%)')

handles.append(hdv2)
handles.append(ebt3)

handles = [handles[0], handles[3], handles[1], handles[4], handles[2]]
fig.legend(handles = handles, #title = 'Fiber thickness / RCF type', title_fontsize=9, 
            fontsize = 8, loc = 'center', bbox_to_anchor = (0.5, 1.04,), 
            columnspacing=1, ncol = 3)

ax.annotate('2b)', xy = (0.99, 0.97), xycoords = 'axes fraction', ha = 'right', va = 'top', color = 'k', fontsize = 15)

plt.savefig('../Figures/TRIM_Figure_2b).png', bbox_inches = 'tight')
