import os
import re  
from math import pow  
from math import sqrt  
import numpy as np  
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt  
#from matplotlib import rcParams
from matplotlib.font_manager import FontProperties
from xlrd import *
from scipy.interpolate import griddata
from mpldatacursor import datacursor
import xlwt
from xlwt import*  
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.cm as cmap
import matplotlib.pyplot as plt

#must have (untrimmed) J Arrays in working directory for this code to run
#if phi is listed first in the J file names then order should be x,y,z... if psi is listed first order is y,x,z
def open_sesame(file_name,x,y,z,shift1,shift2,wd):
    inpt = open_workbook(str(wd)+'Array '+str(file_name)+'.xls', 'r')
    array = inpt.sheet_by_name('Array')
    for i in xrange(array.nrows):
        x.append(float(array.cell_value(i, 0))+shift1)
        y.append(float(array.cell_value(i, 1))+shift2)
        z.append(float(array.cell_value(i, 2)))
    for i in xrange(array.nrows):
        if array.cell_value(i, 0)+shift1 < 15:
            x.append(float(array.cell_value(i, 0))+360.0+shift1)
            y.append(float(array.cell_value(i, 1))+shift2)
            z.append(float(array.cell_value(i, 2)))
        elif array.cell_value(i,0)+shift1 > 345:
            x.append(float(array.cell_value(i, 0))-360.0+shift1)
            y.append(float(array.cell_value(i, 1))+shift2)
            z.append(float(array.cell_value(i, 2)))
    for i in xrange(array.nrows):
        if array.cell_value(i, 1)+shift2 < 15:
            x.append(float(array.cell_value(i, 0))+shift1)
            y.append(float(array.cell_value(i, 1))+360.0+shift2)
            z.append(float(array.cell_value(i, 2)))
        elif array.cell_value(i,1)+shift2 > 345:
            x.append(float(array.cell_value(i, 0))+shift1)
            y.append(float(array.cell_value(i, 1))-360.0+shift2)
            z.append(float(array.cell_value(i, 2)))
    for i in xrange(array.nrows):
        if array.cell_value(i, 0)+shift1 < 15 and array.cell_value(i, 1)+shift2 < 15:
            x.append(float(array.cell_value(i, 0))+360.0+shift1)
            y.append(float(array.cell_value(i, 1))+360.0+shift2)
            z.append(float(array.cell_value(i, 2)))
        elif array.cell_value(i, 0)+shift1 < 15 and array.cell_value(i,1)+shift2 > 345:
            x.append(float(array.cell_value(i, 0))+360.0+shift1)
            y.append(float(array.cell_value(i, 1))-360.0+shift2)
            z.append(float(array.cell_value(i, 2)))
        elif array.cell_value(i, 0)+shift1 > 345 and array.cell_value(i,1)+shift2 < 15:
            x.append(float(array.cell_value(i, 0))-360.0+shift1)
            y.append(float(array.cell_value(i, 1))+360.0+shift2)
            z.append(float(array.cell_value(i, 2)))
        elif array.cell_value(i, 0)+shift1 > 345 and array.cell_value(i,1)+shift2 > 345:
            x.append(float(array.cell_value(i, 0))-360.0+shift1)
            y.append(float(array.cell_value(i, 1))-360.0+shift2)
            z.append(float(array.cell_value(i, 2)))


#Smoothing = prevents sharp peaks at data points
#power = prevents  
#power=[5]*2700
#smoothing=[25]*2700
xv=[]
yv=[]
values=[]


def plotter (x,y,z,i):
        
    
 
    grid_x, grid_y = np.mgrid[0:360:360j, 0:360:360j]    

        # Plotting the result  
    grid_z2 = griddata(np.column_stack((xv,yv)), values,(grid_x, grid_y), method='cubic')
    plt.subplot(111)
            
    
    plt.figure(i+1)  

    plt.imshow(grid_z2.T)   

    fontP = FontProperties()
    fontP.set_size('x-small')
    fontP.set_style('italic')
    fontP.set_weight('bold')
    fontL = fontP.copy()
    fontL.set_size('medium')
    fontT = fontL.copy()
    fontT.set_size('xx-large')
    plt.xlim(0, 360)  
    plt.ylim(0, 360)
    plt.yticks(np.arange(0,361,60))
    plt.xticks(np.arange(0,361,60))     
    

    #This is where you specify your axes labels and plot title. Change between molecule/linkage accordingly. 

    plt.text(180, -25, r'$\mathregular{\phi:C2\'-C1\'-O1\'-C3}$',horizontalalignment='center',verticalalignment='top',fontproperties=fontL)
    plt.text(-30, 180, r'$\mathregular{\psi:C1\'-O1\'-C3-C2}$',horizontalalignment='right',verticalalignment='center',fontproperties=fontL,rotation='vertical')

    #Color bar (z-axis)

    plt.text(440, 170, "Calculated\n    J(Hz)",horizontalalignment='left',verticalalignment='center',fontproperties=fontL)
   
    #Plot Title
    plt.text(180, 370, f_n,horizontalalignment='center',verticalalignment='bottom',fontproperties=fontT)    

    plt.colorbar() 
    plt.tick_params(axis='both', which='major',length=10,width=2)

    datacursor(display='multiple',draggable=True,formatter="phi: {x:0.3g}\npsi: {y:0.3g}\nE: {z:0.3g}".format,bbox=dict(fc='white',boxstyle='round,pad=0.5',alpha=1),keybindings=dict(hide='h', toggle='t'))
    plt.show()
    #eps works, svg works, png works, jpg works, jpeg doesn't...
    
    plt.savefig(str(wd)+f_n+'.jpeg', format='jpg', dpi=1000)

#This is where you specify your working directory!
wd = "/Users/Desktop/Workbench/"

#The remaining lines automate the code to search for all J Array.xls files

allFiles = os.listdir(wd)
results = []
for f in os.listdir(wd):
        if f.endswith('.xls'):
            results.append(f)
m = []
new=[]
for i in xrange(len(results)):
    m= re.sub(r'(\w+)\s(\w+)\s(\w+.*)(\.xls)',('\g<2> \g<3>'), results[i])
    new.insert(i, m)
    f_n = new[i]
    open_sesame(f_n,xv,yv,values,0,0,wd)
    plotter(xv,yv,values,i)
    xv = []
    yv = []
    values = []