#!/usr/bin/env python

import numpy as np
from math import pi
import matplotlib.pyplot as plt 
from matplotlib import rc

rc('mathtext', fontset ='stix')

params = np.array([20, 0.591864648903, 0.0979828313012, 3.64682747229, 7.6466378363]) # AIMD S0 0.06, theta terr = -20deg, theta step = 20deg

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def ang_dep(theta_inc, theta_100, A100, S0_100, n100, n111):
         theta_ter = (theta_inc+20)*pi/180.
         theta_step = (theta_inc-20)*pi/180.
         return (1-A100)*0.06*(np.cos(theta_ter))**n111+A100*S0_100*(np.cos(theta_step))**n100, (1-A100)*0.06*(np.cos(theta_ter))**n111, A100*S0_100*(np.cos(theta_step))**n100
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

# Atoms definition
#			   Theta  Tot    Err    Terr	 Err	  Step	   Err
AIMD_rot_surf = np.array([[-50,  0.008,  0.003, 0.002,   0.001,   0.006,   0.002],
                          [-40,  0.021,  0.005, 0.01,    0.003,   0.011,   0.003],
                          [-30,  0.037,  0.006, 0.005,   0.002,   0.032,   0.006],
                          [-20,  0.039,  0.006, 0.003,   0.002,   0.036,   0.006],
                          [-10,  0.060,  0.008, 0.007,   0.003,   0.053,   0.007],
                          [0,    0.069,  0.008, 0.004,   0.002,   0.065,   0.008],
                          [10,   0.069,  0.008, 0.002,   0.001,   0.067,   0.008],
                          [20,   0.053,  0.007, 0.   ,   0.001,   0.053,   0.007],
                          [30,   0.054,  0.007, 0.   ,   0.001,   0.054,   0.007],
                          [40,   0.050,  0.007, 0.   ,   0.001,   0.05 ,   0.007],
                          [50,   0.033,  0.006, 0.   ,   0.001,   0.033,   0.006]])


AIMD_rot_surf_fac = np.array([[ 0.007,	0.003,	0.001,	0.001],  # POTCAR 2005 SURFACE
			      [	0.014,	0.004,	0.007,	0.003],
			      [	0.023,	0.005,	0.014,	0.004],
			      [	0.022,	0.005,	0.017,	0.004],
			      [	0.038,	0.006,	0.022,	0.005],
			      [	0.038,	0.006,	0.031,	0.006],
			      [	0.033,	0.005,	0.036,	0.006],
			      [	0.013,	0.004,	0.04,	0.006],
			      [	0.008,	0.003,	0.046,	0.007],
			      [	0.004,	0.002,	0.046,	0.007],
 			      [	0.,	0.001,	0.033,	0.006]])

xfit = np.array([i for i in range(-60,60)])

# Fitting facets individually

params_terr = [-7.65365286,  0.03552937,  6.21320268]
params_step = [ 31.25179864,   0.04172631,   1.81187216]

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def ang_fac(theta_inc, theta, A, n):
         theta_inc = (theta_inc-theta)*pi/180.
         return A* (np.cos(theta_inc)) ** n

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

terr_fit = ang_fac(xfit, params_terr[0], params_terr[1], params_terr[2])
step_fit = ang_fac(xfit, params_step[0], params_step[1], params_step[2])
ried_fit = ang_dep(xfit, 20, 0.591864648903, 0.0979828313012, 3.64682747229, 7.6466378363)

ax = plt.subplot(221)
plt.errorbar(AIMD_rot_surf[:,0],AIMD_rot_surf[:,1],yerr = AIMD_rot_surf[:,2], color='k', marker='o', mec = 'k', mew = 1.5, markersize = 5, linestyle='None', label = 'Total')
plt.errorbar(AIMD_rot_surf[:,0],AIMD_rot_surf_fac[:,0],yerr = AIMD_rot_surf_fac[:,1], color='white', marker='o', mec = 'black', mew = 1.5, markersize = 5, linestyle='None', ecolor = 'k', label = '(111)')
plt.errorbar(AIMD_rot_surf[:,0],AIMD_rot_surf_fac[:,2],yerr = AIMD_rot_surf_fac[:,3], color='gray', marker='o', mec = 'gray', mew = 1.5, markersize = 5, linestyle='None', label = '(100)')
plt.plot(xfit, ried_fit[0], color = 'orange')
plt.plot(xfit, terr_fit+step_fit, 'k')
plt.plot(xfit, terr_fit, 'k--')
plt.plot(xfit, step_fit, color = 'gray', linestyle = '--')
#plt.plot(xfit, ried_fit[0], color = 'orange')
plt.xlim(-55,55)
plt.ylim(0,0.1)
plt.xticks(np.arange(-40., 50., 20.0))
plt.xlabel('$\\rm{\\theta_i}$ ('u'\xb0)', fontsize = 12)
plt.ylabel('S$_0$', fontsize = 12)
plt.annotate('$\\rm{\\phi_i} = 0$'u'\xb0', xy = (25,0.0885), xytext=(25,0.0885), fontsize = 12)
handles, labels = ax.get_legend_handles_labels()
handles = [h[0] for h in handles]
lg = ax.legend(handles, labels, loc = 'upper left', numpoints = 1, prop={'size': 11})
lg.draw_frame(False)
plt.savefig('Fig_S7_Riedmuller_facet.png',bbox_inches='tight', dpi=300)
plt.show()