#!/usr/bin/env python
import sys, commands
from os import path, getcwd
from os import chdir, system
import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
from scipy.interpolate import spline

fig, ax = plt.subplots(nrows=3, ncols=1, figsize=(3.37,4))
mpl.rc("font", size=10)

xlabels = [r'$\theta$', r'$\beta$', r'$\gamma$']

def plot(angles, label, color, TS):
	for plot_idx in range(1,4):
		plt.subplot(3,1,plot_idx)

		if len(angles) != 0:
			n, bins = np.histogram( angles[:,plot_idx-1], bins=20, density=True )
			x = []
			x.append( bins[0] - (bins[1]-bins[0])/2 )
			for i in range(len(bins)-1):
				x.append( (bins[i]+bins[i+1])/2. )
			x.append( bins[len(bins)-1] + (bins[len(bins)-1]-bins[len(bins)-2])/2 )
			n = np.insert(n, 0, 0)
			n = np.append(n,0)

			xnew = np.linspace(min(x),max(x),200)
			n_smooth = spline(x,n,xnew)

			plt.plot(xnew,n_smooth, linestyle='--', color=color)

			n, bins = np.histogram( angles[:,plot_idx+2], bins=20, density=True )
			x = []
			x.append( bins[0] - (bins[1]-bins[0])/2 )
			for i in range(len(bins)-1):
				x.append( (bins[i]+bins[i+1])/2. )
			x.append( bins[len(bins)-1] + (bins[len(bins)-1]-bins[len(bins)-2])/2 )
			n = np.insert(n, 0, 0)
			n = np.append(n,0)

			xnew = np.linspace(min(x),max(x),200)
			n_smooth = spline(x,n,xnew)

			plt.plot(xnew,n_smooth, label=label, color=color)

		if plot_idx == 1: TS_angle = TS[0]
		if plot_idx == 2: TS_angle = TS[1]
		if plot_idx == 3: TS_angle = TS[2]
		y_limits = plt.ylim()
		plt.plot( [TS_angle, TS_angle], [0., y_limits[1] * 1.5], linestyle='dotted', c=color )
		plt.ylim(0., y_limits[1]*1.05)

		plt.tick_params(labelleft='off')
		plt.xlim(0,180)
		if plot_idx == 1:
			plt.annotate(xlabels[plot_idx-1], xy=(40, 0.03), size=12)
		if plot_idx == 2:
			plt.annotate(xlabels[plot_idx-1], xy=(40, 0.02), size=12)
		if plot_idx == 3:
			plt.annotate(xlabels[plot_idx-1], xy=(40, 0.04), size=12)
		if plot_idx == 3:
			plt.legend(loc='best', fontsize=9, frameon=False)
		if plot_idx == 1 or plot_idx == 2:
			plt.tick_params(labelbottom='off')
		if plot_idx == 2:
			plt.ylabel('Distribution (arb.)')
		if plot_idx == 3:
			plt.xlabel('Angle (deg)')

angles = []
dat = open('ni_angles_reac.dat', 'r').readlines()
for i in range(1, len(dat)):
	angles.append( [float(dat[i].split()[0]), float(dat[i].split()[1]), float(dat[i].split()[2]),
			float(dat[i].split()[3]), float(dat[i].split()[4]), float(dat[i].split()[5])] )
angles = np.array(angles)
plot(angles, 'Ni(111)', 'b', [135.7, 164.7, 29.1])

angles = []
dat = open('pd_angles_reac.dat', 'r').readlines()
for i in range(1, len(dat)):
	angles.append( [float(dat[i].split()[0]), float(dat[i].split()[1]), float(dat[i].split()[2]),
			float(dat[i].split()[3]), float(dat[i].split()[4]), float(dat[i].split()[5])] )
angles = np.array(angles)
plot(angles, 'Pd(111)', 'k', [135.9, 165.0, 29.1])

angles = []
dat = open('pt_angles_reac.dat', 'r').readlines()
for i in range(1, len(dat)):
	angles.append( [float(dat[i].split()[0]), float(dat[i].split()[1]), float(dat[i].split()[2]),
			float(dat[i].split()[3]), float(dat[i].split()[4]), float(dat[i].split()[5])] )
angles = np.array(angles)
plot(angles, 'Pt(111)', 'r', [133.4, 168.3, 34.8])

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