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

#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.440800, 0.004965],
[0.0084, 0.398500, 0.010948],
[0.0121, 0.341500, 0.010604],
[0.0212, 0.248000, 0.009657],
[0.0304, 0.211834, 0.006286],
[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.250700, 0.004334]]
)

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.404600, 0.004908],
[0.0084, 0.371000, 0.010802],
[0.0121, 0.313500, 0.010373],
[0.0212, 0.192000, 0.008807],
[0.0304, 0.169200, 0.003749],
[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.248900, 0.004324]]
)

NN_766_D2 = np.array(
[[0.0073, 0.392500, 0.004883],
[0.0084, 0.362000, 0.010746],
[0.0121, 0.310500, 0.010346],
[0.0212, 0.206500, 0.009051],
[0.0304, 0.166600, 0.003726],
[0.0437, 0.154000, 0.008071],
[0.0530, 0.164000, 0.008280],
[0.0692, 0.166500, 0.008330],
[0.0889, 0.168500, 0.008370],
[0.1078, 0.205000, 0.009027],
[0.1249, 0.230500, 0.009417],
[0.1405, 0.232000, 0.009439],
[0.1436, 0.249500, 0.004327]]
)

NN_151313_D2 = np.array(
[[0.0073, 0.388600, 0.004874],
[0.0084, 0.361500, 0.010743],
[0.0121, 0.290500, 0.010152],
[0.0212, 0.203500, 0.009002],
[0.0304, 0.172600, 0.003779],
[0.0437, 0.143000, 0.007828],
[0.0530, 0.158500, 0.008166],
[0.0692, 0.182500, 0.008637],
[0.0889, 0.174500, 0.008487],
[0.1078, 0.203500, 0.009002],
[0.1249, 0.221500, 0.009285],
[0.1405, 0.243500, 0.009597],
[0.1436, 0.246800, 0.004311]]
)

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.375490, 0.004855],
[0.0084, 0.349024, 0.010666],
[0.0121, 0.288500, 0.010131],
[0.0212, 0.199500, 0.008936],
[0.0304, 0.151375, 0.004531],
[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.256126, 0.004365]]
)

# 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, , ]]
# )

#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,
'151313': 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

AS_533_D2 = cross_section( NN_533_D2[:,0], NN_533_D2[:,1], NN_533_D2[:,2], 0.2832, NN_step_density['533'] )
AS_322_D2 = cross_section( NN_322_D2[:,0], NN_322_D2[:,1], NN_322_D2[:,2], 0.2832, NN_step_density['322'] )
AS_755_D2 = cross_section( NN_755_D2[:,0], NN_755_D2[:,1], NN_755_D2[:,2], 0.2832, NN_step_density['755'] )
AS_433_D2 = cross_section( NN_433_D2[:,0], NN_433_D2[:,1], NN_433_D2[:,2], 0.2832, NN_step_density['433'] )
AS_977_D2 = cross_section( NN_977_D2[:,0], NN_977_D2[:,1], NN_977_D2[:,2], 0.2832, NN_step_density['977'] )
AS_544_D2 = cross_section( NN_544_D2[:,0], NN_544_D2[:,1], NN_544_D2[:,2], 0.2832, NN_step_density['544'] )
AS_1199_D2 = cross_section( NN_1199_D2[:,0], NN_1199_D2[:,1], NN_1199_D2[:,2], 0.2832, NN_step_density['1199'] )
AS_655_D2 = cross_section( NN_655_D2[:,0], NN_655_D2[:,1], NN_655_D2[:,2], 0.2832, NN_step_density['655'] )
AS_131111_D2 = cross_section( NN_131111_D2[:,0], NN_131111_D2[:,1], NN_131111_D2[:,2], 0.2832, NN_step_density['131111'] )
AS_766_D2 = cross_section( NN_766_D2[:,0], NN_766_D2[:,1], NN_766_D2[:,2], 0.2832, NN_step_density['766'] )
AS_877_D2 = cross_section( NN_877_D2[:,0], NN_877_D2[:,1], NN_877_D2[:,2], 0.2832, NN_step_density['877'] )
AS_171515_D2 = cross_section( NN_171515_D2[:,0], NN_171515_D2[:,1], NN_171515_D2[:,2], 0.2832, NN_step_density['171515'] )

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

def plot_sticking( idx, nn, exp, title ):
	plt.subplot(Nrows, Ncols, idx)
	
	plt.errorbar( nn[:,0]*96.485307, nn[:,1], yerr=nn[:,2], capsize=4, marker='o' )
	plt.plot( exp[:,0]*96.485307, exp[:,1], marker='o' )
	
	plt.annotate('({:s})'.format( title ), xy=(0.5, 0.7), xycoords='axes fraction', horizontalalignment='center', size=12)
	
	plt.xlim(0., 15.)
	plt.xticks([0.,5.,10.])
	plt.ylim(0., plt.ylim()[1])
	plt.tick_params(length=6, width=1, direction='in', top=True, right=True, labelbottom=False)
	if idx == 9:
		plt.tick_params( labelbottom=True )
		#plt.xlabel('Incidence energy (kJ/mol)')
	if idx == 5:
		plt.ylabel('Sticking probability')
	
	return

def plot_crosssection( idx, nn, nn2, exp, title ):
	ax = plt.subplot(Nrows, Ncols, idx)
	
	plt.errorbar( nn[:,0]*96.485307, nn2[:], capsize=4, marker='o' )
	plt.plot( exp[:,0]*96.485307, exp[:,2], marker='o' )
	
	plt.annotate('({:s})'.format( title ), xy=(0.5, 0.7), xycoords='axes fraction', horizontalalignment='center', size=12)
	
	plt.xlim(0., 15.)
	plt.xticks([0.,5.,10.])
	plt.ylim(0., plt.ylim()[1])
	plt.tick_params(length=6, width=1, direction='in', top=True, right=True, labelbottom=False, labelleft=False, labelright=True)
	ax.yaxis.set_label_position("right")
	if idx == 10:
		plt.tick_params( labelbottom=True )
		plt.xlabel('Incidence energy (kJ/mol)')
		ax.xaxis.set_label_coords(0.,-0.3)
	if idx == 6:
		plt.ylabel(r'Step sticking cross section (nm$^2$)')
	
	return

plot_sticking( 1, NN_433_D2, Exp_433_D2, '433' )
plot_sticking( 3, NN_544_D2, Exp_544_D2, '544' )
plot_sticking( 5, NN_655_D2, Exp_655_D2, '655' )
plot_sticking( 7, NN_766_D2, Exp_766_D2, '766' )
plot_sticking( 9, NN_877_D2, Exp_877_D2, '877' )

plot_crosssection( 2, NN_433_D2, AS_433_D2, Exp_433_D2, '433' )
plot_crosssection( 4, NN_544_D2, AS_544_D2, Exp_544_D2, '544' )
plot_crosssection( 6, NN_655_D2, AS_655_D2, Exp_655_D2, '655' )
plot_crosssection( 8, NN_766_D2, AS_766_D2, Exp_766_D2, '766' )
plot_crosssection( 10, NN_877_D2, AS_877_D2, Exp_877_D2, '877' )

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