#!/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.8))

legend_idx = 1
folders = [ '74.0LO', '87.8LO', '102.9LO', '102.9LO_475K', '119.5LO', '119.5LO_475K' ]
titles = [r'$\left<E_i\right>=74$ kJ/mol ($1100$ K)', r'$\left<E_i\right>=88$ kJ/mol ($1100$ K)', r'$\left<E_i\right>=103$ kJ/mol ($1100$ K)', r'$\left<E_i\right>=103$ kJ/mol ($475$ K)', r'$\left<E_i\right>=120$ kJ/mol ($1100$ K)', r'$\left<E_i\right>=120$ kJ/mol ($475$ K)']

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

	fcc = np.dot( fcc, Cell )

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

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

	hcp = np.dot( hcp, Cell )

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

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

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

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

	bridge = np.array(
	[[2./3., 1./3. / 2.],
	[5./6., 1./3.],
	[1./3., 5./6.],
	[1./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., 1./3. / 2.],
			[0., 1./2. / 2.],
			[0., 3./4.],
			[1./3. / 2., 2./3.]]
			)
		elif idx == 1:
			bridge = np.array(
			[[0., 1./2. / 2.],
			[-1./3. / 2., 1./3.],
			[-1./3. / 2., 5./6.],
			[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)

	Cell = [[8.3136945076669999,    0.0000000000000000],
		[4.1568472538335000,    7.1998706429428001]]

	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)

	dist2top    = []
	dist2fcc    = []
	dist2hcp    = []
	dist2bridge = []
	dist2top2    = []
	dist2fcc2    = []
	dist2hcp2    = []
	dist2bridge2 = []
	Ntop = 0
	Nfcc = 0
	Nhcp = 0
	Nbridge = 0
	Ntop2 = 0
	Nfcc2 = 0
	Nhcp2 = 0
	Nbridge2 = 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()

			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()[0]))
                       	dist2fcc.append(float(post_dat[16].split()[0]))
                        dist2hcp.append(float(post_dat[18].split()[0]))
       	                dist2bridge.append(float(post_dat[20].split()[0]))

			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 )

               	        dist2top2.append(float(post_dat[14].split()[2]))
                       	dist2fcc2.append(float(post_dat[16].split()[2]))
                        dist2hcp2.append(float(post_dat[18].split()[2]))
       	                dist2bridge2.append(float(post_dat[20].split()[2]))

			x = float(post_dat[32].split()[2])
			y = float(post_dat[32].split()[3])
			pos2 = np.dot( [x,y], Cell_ )
			pos2 *= 3.
			pos2 %= 1.
			pos2 = np.dot( pos2, Cell2 )

			if "top" in post_dat[21]:	Ntop += 1
			if "fcc" in post_dat[21]:	Nfcc += 1
			if "hcp" in post_dat[21]:	Nhcp += 1
			if "bridge" in post_dat[21]:	Nbridge += 1

			if "top" in post_dat[22]:
				Ntop2 += 1
				top_key = plt.scatter(pos[0], pos[1], c='b', s=50, zorder=9, marker='o')
#				top_key = plt.scatter(pos2[0], pos2[1], c='b', s=50, zorder=10)
			if "fcc" in post_dat[22]:
				Nfcc2 += 1
				fcc_key = plt.scatter(pos[0], pos[1], c='g', s=50, zorder=9, marker='o')
#				fcc_key = plt.scatter(pos2[0], pos2[1], c='g', s=50, zorder=10)
			if "hcp" in post_dat[22]:
				Nhcp2 += 1
				hcp_key = plt.scatter(pos[0], pos[1], c='r', s=50, zorder=9, marker='o')
#				hcp_key = plt.scatter(pos2[0], pos2[1], c='r', s=50, zorder=10)
			if "bridge" in post_dat[22]:
				Nbridge2 += 1
				bridge_key = plt.scatter(pos[0], pos[1], c='k', s=50, edgecolor='grey', zorder=9, marker='o')
#				bridge_key = plt.scatter(pos2[0], pos2[1], c='k', s=50, edgecolor='grey', zorder=10)

	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=(0., 2.6), size=12)
#	plt.annotate(title2, xy=(-1, 2.6), size=10)

	plt.axis('scaled')
	plt.axis([-0.5, 4.5, -0.4, 3.1])
#	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 at initial time step(top, hollow, bridge, total):"
	print Ntop, Nfcc + Nhcp, Nbridge, Ntop + Nfcc + Nhcp + Nbridge

	print folder, "Amount of closest site at the moment of dissociation (top, hollow, bridge, total):"
	print Ntop2, Nfcc2 + Nhcp2, Nbridge2, Ntop2 + Nfcc2 + Nhcp2 + Nbridge2

	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.06, hspace=-0.0)

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

plt.savefig('site.pdf')