#!/usr/bin/env python
import sys, commands
from os import path, getcwd
from os import chdir, system, popen
import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
from mpl_toolkits.axes_grid1.inset_locator import inset_axes

Nrows = 3
Ncols = 2
mpl.rc("text", usetex=True)
fig, ax = plt.subplots(nrows=Nrows, ncols=Ncols, figsize=(6.69,6.3))

legend_idx = 1
folders = [ '/home/nick/AIMD_CHD3-Ni111_SRP0.32-vdW_Ts550K_ET/LaserOff_101.1kJmol+Shift', '/home/nick/AIMD_CHD3-Ni111_SRP0.32-vdW_Ts550K_ET/101.1v1', '/home/nick/AIMDdatapd111/CHD3_Pd111_LO/102-5kJmol', '/home/nick/AIMDdatapd111/CHD3_Pd111_v1/102.4v1', '/home/nick/AIMD_CHD3-Pt111_SRP32vdW_Ts500K/Pt111-SRP-LO/102.5kJmol', '/home/nick/AIMD_CHD3-Pt111_SRP32vdW_Ts500K/Pt111-SRP-v1/104.6kJmol' ]
titles = ['Ni(111)', 'Ni(111)', 'Pd(111)', 'Pd(111)', 'Pt(111)', 'Pt(111)']

def stattop(ax, Cell):
	for idx1 in range(2):
		for idx2 in range(2):
			if idx1 == 0 and idx2 == 0:
				top = np.array(
					[[2./3. / 2., 1./3. / 2.],
					[1./3. / 2., 2./3. / 2.],
					[0., 1./2. / 2.],
					[0., 0.],
					[1./2. / 2., 0.]]
					) 
			elif idx1 == 0 and idx2 == 1:
				top = np.array(
					[[1./2. / 2., 0.],
					[0., 0.],
					[0., -1./2. / 2.],
					[1./3. / 2., -1./3. / 2.]]
					)
			elif idx1 == 1 and idx2 == 0:
				top = np.array(
					[[0., 1./2. / 2.],
					[-1./3. / 2., 1./3. / 2.],
					[-1./2. / 2., 0.],
					[0., 0.]]
					)
			elif idx1 == 1 and idx2 == 1:
				top = np.array(
					[[0., 0.],
					[-1./2. / 2., 0.],
					[-2./3. / 2., -1./3. / 2.],
					[-1./3. / 2., -2./3. / 2.],
					[0., -1./2. / 2.]]
					)

			top[:,0] += idx1
			top[:,1] += idx2
			top = np.dot( top, Cell )

			ax.fill( top[:,0], top[:,1], color='b', edgecolor='k')

def statfcc(ax, Cell):
	fcc = np.array(
	[[2./3. / 2., 1./3. / 2.],
	[5./6., 2./3.],
	[5./6., 1./3. / 2.]]
	)

	fcc = np.dot( fcc, Cell )

	ax.fill( fcc[:,0], fcc[:,1], color='r')

def stathcp(ax, Cell):
	hcp = np.array(
	[[1./3. / 2., 2./3. / 2.],
	[1./3. / 2., 5./6.],
	[2./3., 5./6.]]
	)

	hcp = np.dot( hcp, Cell )

	ax.fill( hcp[:,0], hcp[:,1], color='g')

def statbridge(ax, Cell):
	for idx in range(2):
		if idx == 0:
			bridge = np.array(
			[[1./2. / 2., 0.],
			[2./3. / 2., 1./3. / 2.],
			[5./6., 1./3. / 2.],
			[3./4., 0.]]
			)
		elif idx == 1:
			bridge = np.array(
			[[1./3. / 2., -1./3. / 2.],
			[1./2. / 2., 0.],
			[3./4., 0.],
			[2./3., -1./3. / 2.]]
			)

		bridge[:,1] += idx
		bridge = np.dot( bridge, Cell )

		ax.fill( bridge[:,0], bridge[:,1], color='k')

		bridge = np.array(
		[[2./3. / 2., 1./3. / 2.],
		[1./3. / 2., 2./3. / 2.],
		[2./3., 5./6.],
		[5./6., 2./3.]]
		)

		bridge = np.dot( bridge, Cell )

		ax.fill( bridge[:,0], bridge[:,1], color='k')

	for idx in range(2):
		if idx == 0:
			bridge = np.array(
			[[1./3. / 2., 2./3. / 2.],
			[0., 1./2. / 2.],
			[0., 3./4.],
			[1./3. / 2., 5./6.]]
			)
		elif idx == 1:
			bridge = np.array(
			[[0., 1./2. / 2.],
			[-1./3. / 2., 1./3. / 2.],
			[-1./3. / 2., 2./3.],
			[0., 3./4.]]
			)

		bridge[:,0] += idx
		bridge = np.dot( bridge, Cell )

		ax.fill( bridge[:,0], bridge[:,1], color='k')

###########################################################

here=getcwd()

def plot(idx, legend_idx, folder, title):
	chdir(folder)

	subax = plt.subplot(Nrows,Ncols,idx)

	if idx == 1 or idx == 2: #Ni111
		Cell = [[ 7.6011662803200002,    0.0000000000000000],
			[-3.8005831401600001,    6.5828030971400002]]
	if idx == 3 or idx == 4: #Pd111
		Cell = [[ 8.5137224254107995,    0.0000000000000000],
			[-4.2568612127053997,    7.3730999011750002]]
	if idx == 5 or idx == 6: #Pt111
		Cell = [[ 8.5609344674269003,    0.0000000000000000],
			[-4.2804672337134004,    7.4139867289254999]]

	Cell = np.array( Cell )
	Cell2 = Cell / 3.
	Cell_= np.linalg.inv(Cell)
	Cell2_= np.linalg.inv(Cell2)

	statbridge(subax, Cell2)
	statfcc(subax, Cell2)
	stathcp(subax, Cell2)
	stattop(subax, Cell2)

	Surface_d = np.array(	[[ 0.0000000000000,  0.0000000000000],
				[  0.0000000000000,  1.0000000000000],
				[  1.0000000000000,  0.0000000000000],
				[  1.0000000000000,  1.0000000000000]] )

	Surface_c = np.dot(Surface_d, Cell2)
	Surface = np.zeros((len(Surface_d),2))
	for i in range(len(Surface_d)):
		for j in range(2):
			Surface[i][j] = Surface_c[i][j]

	Pd_key = plt.scatter(* Surface.T, c='grey', s=50, marker='o', facecolors='none', linewidths=2, edgecolor='grey', zorder=5)

	outcome_dat = []
	dist2top    = []
	dist2fcc    = []
	dist2hcp    = []
	dist2bridge = []
	Ntop = 0
	Nfcc = 0
	Nhcp = 0
	Nbridge = 0

	here_tmp=getcwd()				#workdir
	Ntraj = int( popen('ls | egrep "^[0-9]...$" | wc -l').read() )
	for i in range(1,(Ntraj+1)):
		traj= '%.4d' % i		# define working directory
		newdir=path.join(here_tmp, traj)	# updating mypath

		if not path.exists(newdir): 
			print "folder", newdir, "not found. exiting..."
			quit()

		chdir(newdir)

		if path.exists('PostAnalysis.dat') and path.exists('outcome'):
			outcome = open('outcome', 'r').readline()
			outcome_dat.append(str(outcome))

			if 'UNCLEAR' in outcome:
				continue

			post_dat = open('PostAnalysis.dat', 'r').readlines()
			if "SCATTERING" in outcome or "TRAPPING" in outcome:
				continue

               	        dist2top.append(float(post_dat[14].split()[2]))
                       	dist2fcc.append(float(post_dat[16].split()[2]))
                        dist2hcp.append(float(post_dat[18].split()[2]))
       	                dist2bridge.append(float(post_dat[20].split()[2]))

			x = float(post_dat[31].split()[2])
			y = float(post_dat[31].split()[3])
			pos = np.dot( [x,y], Cell_ )
			pos *= 3.
			pos %= 1.
			pos = np.dot( pos, Cell2 )
			if "top" in post_dat[21]:
				Ntop += 1
				top_key = plt.scatter(pos[0], pos[1], c='b', s=50, zorder=10)
			if "fcc" in post_dat[21]:
				Nfcc += 1
				fcc_key = plt.scatter(pos[0], pos[1], c='g', s=50, zorder=10)
			if "hcp" in post_dat[21]:
				Nhcp += 1
				hcp_key = plt.scatter(pos[0], pos[1], c='r', s=50, zorder=10)
			if "bridge" in post_dat[21]:
				Nbridge += 1
				bridge_key = plt.scatter(pos[0], pos[1], c='k', s=50, edgecolor='grey', zorder=10)
		else:
			outcome_dat.append('UNCLEAR')

	if idx > (Nrows-1)*Ncols:
		plt.xlabel(r'x (\r{A})')
	else:
		plt.tick_params(labelbottom='off')
	if idx % Ncols == 0:
		plt.tick_params(labelleft='off')
	if idx % 2 == 1:
		plt.ylabel(r'y (\r{A})')
	plt.tick_params(length=6, width=1)

	plt.annotate(title, xy=(2.25, 2.25), size=12)

	plt.axis('scaled')
	plt.axis([-2., 3.5, -0.4, 2.9])
#	plt.xlim(-5., 9.)

#	if idx == 2:
#		plt.legend((top_key), ("top"), loc='upper right')
#		plt.legend((top_key, hcp_key), ('top', 'hcp'), loc='upper right', handletextpad=0, borderaxespad=0)
#	elif idx == 4:
#		plt.legend(hcp_key, "hcp", loc='upper right')
#	elif idx == 6:
#		plt.legend(fcc_key, "fcc", loc='upper right')
#	elif idx == 8:
#		plt.legend(bridge_key, "bridge", loc='upper right')
#	if idx % 2 == 0:
#	plt.legend(handletextpad=0, borderaxespad=0)

	print folder, "Amount of closest site (top, hollow, bridge, total):"
	print Ntop, Nfcc + Nhcp, Nbridge, Ntop + Nfcc + Nhcp + Nbridge

	chdir(here)

for i in range(len(folders)):
	plot(i+1, legend_idx, folders[i], titles[i])

plt.tight_layout()
plt.subplots_adjust(wspace=0, hspace=-0.13)

plt.text(0.25, 0.97, 'Laser off', fontsize=12, transform=plt.gcf().transFigure)
plt.text(0.7, 0.97, r'$\nu_1=1$', fontsize=12, transform=plt.gcf().transFigure)

plt.savefig('site_analytical.pdf')