#!/usr/bin/env python
import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
from scipy.optimize import curve_fit

mpl.rc("text", usetex=True)
#fig = plt.figure(figsize=(4.69,4.))
Nrows = 3
Ncols = 2
fig, ax = plt.subplots(nrows=Nrows, ncols=Ncols, figsize=(3.69,4.))
#plt.rcParams.update({'font.size': 10})

#Ei, S0, standard error
# Serious trapping problem, should increase simulation time
NN_111_D2 = np.array(
[[0.0073, 0.211500, 0.009131],
[0.0084, 0.196500, 0.008885],
[0.0121, 0.132500, 0.007581],
[0.0212, 0.100500, 0.006723],
[0.0304, 0.105500, 0.006869],
[0.0437, 0.105500, 0.006869],
[0.0530, 0.122500, 0.007331],
[0.0692, 0.148500, 0.007951],
[0.0889, 0.170500, 0.008409],
[0.1078, 0.197500, 0.008902],
[0.1249, 0.226500, 0.009359],
[0.1405, 0.233500, 0.009460],
[0.1436, 0.251000, 0.009695]]
)

NN_533_D2 = np.array(
[[0.0073, 0.696500, 0.004598],
[0.0084, 0.675000, 0.010473],
[0.0121, 0.626000, 0.010820],
[0.0212, 0.477000, 0.011169],
[0.0304, 0.392900, 0.004884],
[0.0437, 0.309500, 0.010337],
[0.0530, 0.277000, 0.010007],
[0.0692, 0.262000, 0.009832],
[0.0889, 0.265000, 0.009869],
[0.1078, 0.279000, 0.010029],
[0.1249, 0.271500, 0.009945],
[0.1405, 0.272000, 0.009950],
[0.1436, 0.277000, 0.004475]]
)

NN_322_D2 = np.array(
[[0.0073, 0.596900, 0.004905],
[0.0084, 0.605000, 0.010931],
[0.0121, 0.503000, 0.011180],
[0.0212, 0.380500, 0.010856],
[0.0304, 0.296000, 0.004565],
[0.0437, 0.271500, 0.009945],
[0.0530, 0.228500, 0.009388],
[0.0692, 0.215000, 0.009186],
[0.0889, 0.216500, 0.009209],
[0.1078, 0.223500, 0.223500],
[0.1249, 0.223500, 0.009315],
[0.1405, 0.254000, 0.009734],
[0.1436, 0.254500, 0.004356]]
)

NN_755_D2 = np.array(
[[0.0073, 0.544100, 0.004981],
[0.0084, 0.535500, 0.011152],
[0.0121, 0.464000, 0.011151],
[0.0212, 0.325000, 0.010473],
[0.0304, 0.247800, 0.004317],
[0.0437, 0.215000, 0.009186],
[0.0530, 0.215500, 0.009194],
[0.0692, 0.192000, 0.008807],
[0.0889, 0.205500, 0.009035],
[0.1078, 0.208000, 0.009076],
[0.1249, 0.232000, 0.009439],
[0.1405, 0.229500, 0.009403],
[0.1436, 0.241600, 0.004281]]
)

NN_433_D2 = np.array(
[[0.0073, 0.500070, 0.001581],
[0.0084, 0.457000, 0.011139],
[0.0121, 0.270500, 0.009933],
[0.0212, 0.287500, 0.010120],
[0.0304, 0.231000, 0.002980],
[0.0437, 0.195500, 0.008868],
[0.0530, 0.195000, 0.008859],
[0.0692, 0.186500, 0.008710],
[0.0889, 0.191000, 0.008790],
[0.1078, 0.208500, 0.009084],
[0.1249, 0.205000, 0.009027],
[0.1405, 0.220500, 0.009270],
[0.1436, 0.224800, 0.002952]]
)

NN_977_D2 = np.array(
[[0.0073, 0.478600, 0.004995],
[0.0084, 0.433000, 0.011080],
[0.0121, 0.362500, 0.010749],
[0.0212, 0.257500, 0.009777],
[0.0304, 0.211200, 0.004082],
[0.0437, 0.185500, 0.008692],
[0.0530, 0.177500, 0.008544],
[0.0692, 0.181000, 0.008609],
[0.0889, 0.188000, 0.008737],
[0.1078, 0.206000, 0.009043],
[0.1249, 0.206500, 0.009051],
[0.1405, 0.217000, 0.009217],
[0.1436, 0.242600, 0.004287]]
)

NN_544_D2 = np.array(
[[0.0073, 0.421000, 0.011040],
[0.0084, 0.398500, 0.010948],
[0.0121, 0.341500, 0.010604],
[0.0212, 0.248000, 0.009657],
[0.0304, 0.209500, 0.009100],
[0.0437, 0.189000, 0.008754],
[0.0530, 0.175000, 0.008496],
[0.0692, 0.188500, 0.008746],
[0.0889, 0.199500, 0.008936],
[0.1078, 0.201000, 0.008961],
[0.1249, 0.221000, 0.009278],
[0.1405, 0.221000, 0.009278],
[0.1436, 0.242500, 0.009584]]
)

NN_1199_D2 = np.array(
[[0.0073, 0.427500, 0.004947],
[0.0084, 0.404000, 0.010972],
[0.0121, 0.317000, 0.010405],
[0.0212, 0.226500, 0.009359],
[0.0304, 0.186500, 0.003895],
[0.0437, 0.170000, 0.008399],
[0.0530, 0.161500, 0.008229],
[0.0692, 0.181000, 0.008609],
[0.0889, 0.204000, 0.009011],
[0.1078, 0.221500, 0.009285],
[0.1249, 0.206500, 0.009051],
[0.1405, 0.231000, 0.009424],
[0.1436, 0.240400, 0.004273]]
)

NN_655_D2 = np.array(
[[0.0073, 0.424400, 0.004943],
[0.0084, 0.398500, 0.004896],
[0.0121, 0.316500, 0.004651],
[0.0212, 0.220000, 0.004142],
[0.0304, 0.177900, 0.003824],
[0.0437, 0.170200, 0.003758],
[0.0530, 0.168500, 0.003743],
[0.0692, 0.169300, 0.003750],
[0.0889, 0.189700, 0.003921],
[0.1078, 0.206600, 0.004049],
[0.1249, 0.225000, 0.004176],
[0.1405, 0.243300, 0.004291],
[0.1436, 0.244700, 0.004299]]
)

NN_131111_D2 = np.array(
[[0.0073, 0.416000, 0.011021],
[0.0084, 0.371000, 0.010802],
[0.0121, 0.313500, 0.010373],
[0.0212, 0.192000, 0.008807],
[0.0304, 0.168500, 0.008370],
[0.0437, 0.154500, 0.008082],
[0.0530, 0.161500, 0.008229],
[0.0692, 0.188000, 0.008737],
[0.0889, 0.183500, 0.008655],
[0.1078, 0.191500, 0.008799],
[0.1249, 0.226000, 0.009352],
[0.1405, 0.242000, 0.009577],
[0.1436, 0.253000, 0.009721]]
)

# NN_XXX_D2 = np.array(
# [[0.0073, , ],
# [0.0084, , ],
# [0.0121, , ],
# [0.0212, , ],
# [0.0304, , ],
# [0.0437, , ],
# [0.0530, , ],
# [0.0692, , ],
# [0.0889, , ],
# [0.1078, , ],
# [0.1249, , ],
# [0.1405, , ],
# [0.1436, , ]]
# )

NN_877_D2 = np.array(
[[0.0073, 0.361189, 0.003532],
[0.0084, 0.352000, 0.004776],
[0.0121, 0.275200, 0.004466],
[0.0212, 0.192700, 0.003944],
[0.0304, 0.159400, 0.003660],
[0.0437, 0.154000, 0.003609],
[0.0530, 0.150800, 0.003579],
[0.0692, 0.163100, 0.003695],
[0.0889, 0.183100, 0.003867],
[0.1078, 0.202600, 0.004019],
[0.1249, 0.222200, 0.004157],
[0.1405, 0.249400, 0.004327],
[0.1436, 0.250000, 0.004330]]
)

NN_171515_D2 = np.array(
[[0.0073, 0.367706, 0.010814],
[0.0084, 0.349024, 0.010666],
[0.0121, 0.288500, 0.010131],
[0.0212, 0.199500, 0.008936],
[0.0304, 0.163000, 0.008259],
[0.0437, 0.123500, 0.007357],
[0.0530, 0.141000, 0.007782],
[0.0692, 0.154000, 0.008071],
[0.0889, 0.188000, 0.008737],
[0.1078, 0.209000, 0.009092],
[0.1249, 0.236000, 0.009495],
[0.1405, 0.247500, 0.009650],
[0.1436, 0.247000, 0.009643]]
)

#Ei, S0, cross section
Exp_433_D2 = np.array(
[[0.0073, 3.123712389561522640e-01, 1.464460754155847699e-01],
[0.0084, 3.240269159548841982e-01, 1.512013867184430205e-01],
[0.0121, 2.411905823089349754e-01, 1.078497405297281914e-01],
[0.0212, 1.972592989011457965e-01, 7.912843815126073543e-02],
[0.0304, 1.550099885235217689e-01, 5.098276325280999555e-02],
[0.0437, 1.469015221906095070e-01, 3.607253662921112991e-02],
[0.0530, 1.559784549342670967e-01, 3.294469661002290273e-02],
[0.0692, 1.621111172662298427e-01, 2.268013607117089214e-02],
[0.0889, 1.711319016052875730e-01, 1.124813089545564394e-02],
[0.1078, 2.102532778741268427e-01, 1.535152157636498343e-02],
[0.1249, 2.474944172906987094e-01, 1.999770284208108456e-02],
[0.1405, 2.499631544856877885e-01, 8.367615943038652498e-03],
[0.1436, 2.978204103664884861e-01, 3.005281645279138381e-02]]
)

Exp_977_D2 = np.array(
[[0.0073, 3.008774290191905676e-01, 1.578008123690265441e-01],
[0.0084, 2.866483244714388423e-01, 1.490519036950181120e-01],
[0.0121, 2.136752846247129412e-01, 1.061059031011866438e-01],
[0.0212, 1.652458076713756441e-01, 7.191446731680038729e-02],
[0.0304, 1.439192230593519528e-01, 5.224121551289109128e-02],
[0.0437, 1.314539405168047292e-01, 3.376097792521811497e-02],
[0.0530, 1.407488031763507730e-01, 3.077430899497045347e-02],
[0.0692, 1.406653035345863112e-01, 1.659454953231510529e-02],
[0.0889, 1.724971506382059205e-01, 1.707130579764278011e-02],
[0.1078, 2.065907511537382613e-01, 1.969572934564971928e-02],
[0.1249, 2.294851973738967488e-01, 1.776982445511855746e-02],
[0.1405, 2.530301854358751323e-01, 1.692350201421983016e-02],
[0.1436, 2.800238193809891918e-01, 2.988196486494650722e-02]]
)

Exp_544_D2 = np.array(
[[0.0073, 0.285032323017987,	0.169327181323935],
[0.0084, 0.274355848690535,	0.161649132663104],
[0.0121, 0.207048230592713,	0.116610355577717],
[0.0212, 0.152714511775869,	0.074541598902488],
[0.0304, 0.137609006966197,	0.056432825241528],
[0.0437, 0.126527252795045,	0.036842116597675],
[0.0530, 0.130997570354976,	0.030819869102323],
[0.0692, 0.150094344582123,	0.027207199137055],
[0.0889, 0.171700774978555,	0.022150044322864],
[0.1078, 0.205869813919058,	0.025854231248374],
[0.1249, 0.22218655569569,	0.020222604157579],
[0.1405, 0.245015970660793,	0.019406570387122],
[0.1436, 0.282482904500684,	0.040641928391885]]
)

Exp_655_D2 = np.array(
[[0.0073,	0.260373176997955,	0.176351917263227],
[0.0084,	0.254989061983479,	0.171057197554525],
[0.0121,	0.193156042012056,	0.122343226685087],
[0.0212,	0.153085559864757,	0.082325583250056],
[0.0304,	0.129346825735136,	0.053564108598744],
[0.0437,	0.118879482020317,	0.02890431191972],
[0.0530,	0.124851395760717,	0.021302198166042],
[0.0692,	0.139979401902363,	0.011276307248063],
[0.0889,	0.177113799942644,	0.012891446840027],
[0.1078,	0.211915900443956,	0.014190383063207],
[0.1249,	0.232139609435594,	0.007390300109167],
[0.1405,	0.247505518274276,	-0.001804964647905],
[0.1436,	0.292701191336413,	0.027814819771207]]
)

Exp_766_D2 = np.array(
[[0.0073,	0.244803089638541,	0.19610221983666],
[0.0084,	0.233221331099433,	0.184809676656973],
[0.0121,	0.176157728686029,	0.131869663506526],
[0.0212,	0.127440838533376,	0.078699957328175],
[0.0304,	0.121765692336696,	0.061326219902316],
[0.0437,	0.11220177097052,	0.034981417292398],
[0.0530,	0.12243639332014,	0.030972891808197],
[0.0692,	0.142975901270236,	0.026109987560749],
[0.0889,	0.176050626619163,	0.027488767613286],
[0.1078,	0.20360972651643,	0.025660326401239],
[0.1249,	0.227071856908255,	0.022823758229027],
[0.1405,	0.251117145083097,	0.021613103075357],
[0.1436,	0.2763785852809,	0.040103417339483]]
)

Exp_877_D2 = np.array(
[[0.0073, 2.268677033865130477e-01, 2.201142424075755999e-01],
[0.0084, 2.125730668461016659e-01,  2.034817799259183713e-01],
[0.0121, 1.431993509475461557e-01, 1.260545699016710486e-01],
[0.0212, 1.255073375628939780e-01, 9.238105597489912335e-02],
[0.0304, 1.194183809029822074e-01, 7.008422289232152380e-02],
[0.0437, 1.044692111029262016e-01, 3.144032794449071189e-02],
[0.0530, 1.158765695715728850e-01, 2.706123685649320515e-02],
[0.0692, 1.375259327914914453e-01, 2.105150974691582730e-02],
[0.0889, 1.836682259473485512e-01, 3.467396492853316858e-02],
[0.1078, 2.203040385672185297e-01, 4.037380353209309158e-02],
[0.1249, 2.328171077070683437e-01, 2.444581941677504561e-02],
[0.1405, 2.451922230602875952e-01, 9.822967520053588891e-03],
[0.1436, 2.881658557907037510e-01, 5.037138474048460934e-02]]
)

# In terms of unit cell vector parallel to the terrace and in nm^-1
# Please check the non-rectangular unit cells, there might be an error (or not)
NN_step_density = {
'211': 1.33925980477299,
'533': 1.10245499657604,
'322': 0.867222846526904,
'755': 0.717738278424414,
'433': 0.609594900345529,
'977': 0.531226638707765,
'544': 0.469525140800388,
'1199': 0.42152023770247,
'655': 0.381732224349042,
'131111': 0.349345915692593,
'766': 0.321582789240287,
'151513': 0.298239412749505,
'877': 0.277790623773609,
'171515': 0.260184400328028,
'111': 0.
}

# Note that the step densities do not correspond exactly miller index surfaces due to experimental reasons
# Hence, I've chosen a surface that is close to the experimental step density
Exp_step_density = {
'433': 5.690392986153920418e-01,
'977': 5.088661521814783484e-01,
'544': 4.489717382902294607e-01,
'655': 3.893218217006779724e-01,
'766': 3.298828424342739041e-01,
'877': 2.706218182028406471e-01
#2.115062505692402728e-01
#1.525040342030422047e-01
#9.358336863015850882e-02
#3.471267190296287164e-02
}

def sticking_function( E, a, b, c, d ):
	S = a*np.exp( -E/b ) + c + d*E

	return S

def cross_section( E, S, sigma, w_step, d_step ):
	parameters, pcov = curve_fit(sticking_function, E, S, sigma=sigma, method='dogbox')

	AS = (S - parameters[3]*E) / d_step * w_step

	return AS

def plot_NN( NN, exp=False ):
	parameters, pcov = curve_fit(sticking_function, NN[:,0], NN[:,1], sigma=NN[:,2], method='dogbox')

	E = np.linspace( min(NN[:,0]), max(NN[:,0]), 100 )
	S_step = parameters[0]*np.exp( -E/parameters[1] ) + parameters[2]
	S_terrace = parameters[3]*E
	
	plt.plot( E*96.485307, S_step, color='C00', label='Step' )
	plt.plot( E*96.485307, S_terrace, color='C01', label='Terrace' )
	plt.plot( E*96.485307, S_step+S_terrace, color='C02', label='Step + Terrace' )
	plt.scatter( NN[:,0]*96.485307, NN[:,1], edgecolor='C02', facecolor='None', marker='o' )
	#plt.errorbar( NN[:,0]*96.485307, NN[:,1], yerr=NN[:,2], color='C02', marker='o', capsize=4, mfc='None', ls='' )

	return
	
plt.subplot(Nrows, Ncols, 1)

plot_NN( NN_977_D2 )

Step_ratio = np.array(
[[0.9741, 0.0024],
[0.9704, 0.0060],
[0.9668, 0.0068],
[0.9115, 0.0127],
[0.8621, 0.0076],
[0.7973, 0.0210],
[0.6838, 0.0248],
[0.6537, 0.0250],
[0.6176, 0.0251],
[0.5583, 0.0245],
[0.5377, 0.0246],
[0.4596, 0.0239],
[0.5072, 0.0102]]
)
plt.scatter( NN_977_D2[:,0]*96.485307, Step_ratio[:,0]*NN_977_D2[:,1], fc='None', edgecolor='C00' )
plt.scatter( NN_977_D2[:,0]*96.485307, (1.-Step_ratio[:,0])*NN_977_D2[:,1], fc='None', edgecolor='C01' )
#plt.errorbar( NN_977_D2[:,0]*96.485307, Step_ratio[:,0]*NN_977_D2[:,1], yerr=Step_ratio[:,1], mfc='None', ls='', marker='o', capsize=4 )
#plt.errorbar( NN_977_D2[:,0]*96.485307, (1.-Step_ratio[:,0])*NN_977_D2[:,1], yerr=Step_ratio[:,1], mfc='None', ls='', marker='o', capsize=4 )

#plt.annotate('(977)', xy=(0.9, 0.1), xycoords='axes fraction', size=12)

plt.title('Theory')
plt.legend(loc='best', numpoints=1, frameon=True, fontsize=7)
plt.xlim(0., 14.9)
plt.ylim( 0., 0.55 )
#plt.ylim(1e-7, 1.)
#plt.yscale('log')
plt.tick_params(length=6, width=1, direction='in', top=True, right=True, labelbottom=False)
#plt.ylabel('Sticking probability')
#plt.xlabel('Incidence energy (kJ/mol)')

plt.subplot(Nrows, Ncols, 2)

plot_NN( Exp_977_D2, True )

plt.annotate('Pt(977)', xy=(0.68, 0.85), xycoords='axes fraction', size=10)

plt.title('Experiment')
#plt.legend(loc='best', numpoints=1, frameon=True, fontsize=8)
plt.xlim(0., 15.)
plt.ylim( 0., 0.55 )
#plt.ylim(0.0, 0.8)
#plt.ylim(1e-7, 1.)
#plt.yscale('log')
plt.tick_params(length=6, width=1, direction='in', top=True, right=True, labelbottom=False, labelleft=False)
#plt.ylabel(r'Sticking cross section step (nm$^2$)')
#plt.xlabel(r'Step density (nm$^{-1}$)')

plt.subplot(Nrows, Ncols, 3)

plot_NN( NN_655_D2 )

Step_ratio = np.array([ 0.9371, 0.9340, 0.9037, 0.8234, 0.7529, 0.5975, 0.5621, 0.4655, 0.4389, 0.3866, 0.3643, 0.3306, 0.3413 ])
plt.scatter( NN_655_D2[:,0]*96.485307, Step_ratio[:]*NN_655_D2[:,1], fc='None', edgecolor='C00' )
plt.scatter( NN_655_D2[:,0]*96.485307, (1.-Step_ratio[:])*NN_655_D2[:,1], fc='None', edgecolor='C01' )

plt.annotate('Pt(655)', xy=(0.68, 0.85), xycoords='axes fraction', size=10)

#plt.legend(loc='best', numpoints=1, frameon=True, fontsize=8)
plt.xlim(0., 14.9)
plt.ylim( 0., 0.55 )
#plt.ylim(0.0, 0.8)
#plt.ylim(1e-7, 1.)
#plt.yscale('log')
plt.tick_params(length=6, width=1, direction='in', top=True, right=True, labelbottom=False)
plt.ylabel(r'Sticking probability')
#plt.xlabel(r'Step density (nm$^{-1}$)')

plt.subplot(Nrows, Ncols, 4)

plot_NN( Exp_655_D2, True )

plt.annotate('Pt(655)', xy=(0.68, 0.85), xycoords='axes fraction', size=10)

#plt.legend(loc='best', numpoints=1, frameon=True, fontsize=8)
plt.xlim(0., 15.)
plt.ylim( 0., 0.55 )
#plt.ylim(0.0, 0.8)
#plt.ylim(1e-7, 1.)
#plt.yscale('log')
plt.tick_params(length=6, width=1, direction='in', top=True, right=True, labelbottom=False, labelleft=False)
#plt.ylabel(r'Sticking cross section step (nm$^2$)')
#plt.xlabel(r'Step density (nm$^{-1}$)')

plt.subplot(Nrows, Ncols, 5)

plot_NN( NN_877_D2 )

Step_ratio = np.array([ 0.8958, 0.8902, 0.8505, 0.7400, 0.6231, 0.4751, 0.4202, 0.3526, 0.3055, 0.2758, 0.2428, 0.2336, 0.2285])
plt.scatter( NN_877_D2[:,0]*96.485307, Step_ratio[:]*NN_877_D2[:,1], fc='None', edgecolor='C00' )
plt.scatter( NN_877_D2[:,0]*96.485307, (1.-Step_ratio[:])*NN_877_D2[:,1], fc='None', edgecolor='C01' )

plt.annotate('Pt(877)', xy=(0.68, 0.85), xycoords='axes fraction', size=10)

#plt.legend(loc='best', numpoints=1, frameon=True, fontsize=8)
plt.xlim(0., 14.9)
plt.ylim( 0., 0.55 )
#plt.ylim(0.0, 0.8)
#plt.ylim(1e-7, 1.)
#plt.yscale('log')
plt.tick_params(length=6, width=1, direction='in', top=True, right=True)
#plt.ylabel(r'Sticking cross section step (nm$^2$)')
plt.xlabel(r'$E_\textrm{i}$ (kJ/mol)')

plt.subplot(Nrows, Ncols, 6)

plot_NN( Exp_877_D2, True )

plt.annotate('Pt(877)', xy=(0.68, 0.85), xycoords='axes fraction', size=10)

#plt.legend(loc='best', numpoints=1, frameon=True, fontsize=8)
plt.xlim(0., 15.)
plt.ylim( 0., 0.55 )
#plt.ylim(0.0, 0.8)
#plt.ylim(1e-7, 1.)
#plt.yscale('log')
plt.tick_params(length=6, width=1, direction='in', top=True, right=True, labelleft=False)
#plt.ylabel(r'Sticking cross section step (nm$^2$)')
plt.xlabel(r'$E_\textrm{i}$ (kJ/mol)')

plt.tight_layout()
plt.subplots_adjust(wspace=0, hspace=0)
plt.savefig('crosssection_fit.pdf')
#plt.show()