import os
import re
import sys
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

mydir = os.path.dirname( __file__ )
statusdir = os.path.join(mydir, '..', 'status')
sys.path.append(statusdir)
from status import status

def echart(files=[],loc='./',energy_pattern='Energy',hartree=False,absolute=False,fig_pattern='line',figdata=True,savefig=True,figname='Result',savedata=True,dataname='Result',xx=False):
    """
        A method to make a chart with calculated Electronic Energy datas or Gibbs Free Energy data.
        echart(files=[],loc='./',energy_pattern='Energy',hartree=False,absolute=False,figname='Result',fig_pattern='line',savefig=False,dataname='Result',savedata=True)
        files: The name of ORCA output files already calculated.
        loc: File Location. The default is your current location.
        energy_pattern: Energy you need. There are 2 options, 'Energy' for normal Electronic Energy and 'Gibbs' for Gibbs Free Energy. If you want to use 'Gibbs', please make sure your output file including frequency and thermodynamics calculation.
        hartree: Energy unit is normally kcal/mol, if you need unit 'Hartree', please input 'hartree=True'
                 TIPS: 1 Hartree ≈ 627.509 kcal/mol
        absolute: Normally we'll automatical use Relative Energy to show energy change, if you want Absolute Energy, please input 'absolute=True'
        fig_pattern: Pattern of the figure. There are 2 options, 'line' means Line Chart and 'bar' means Bar Chart.
        figdata: Show the File Name and Energy on the chart. The defaulr is True.
        savefig / savedata: To show if you want to save the chart and .txt result. The default is True.
        figname / savename: To set the name of the chart and .txt result. The default is 'Result'. 
        TIPS: If you use it to collect Gibbs Free Energy, please make sure you have frequency calculation. If the programme can't find out thermodynamics calculation, it will return None(NaN).
        
        Example:
            Input:
                from MCPoly.status import echart
                a=echart(files['Molecule1','Molecule2','Molecule3','Molecule4','Molecule5','Molecule6'],fig_pattern='line',\
                    figdata=True)
                print(a)
            Output:
                <<matplotlib.plot>>
                [ 0.          0.79379889  3.3396029   8.29064891 16.99670877 33.95325697]
            In Result.csv:
                Step,File,∆E(kcal/mol)
                0,Molecule1,0.0
                1,Molecule2,0.7937988850646964
                2,Molecule3,3.3396028980577155
                3,Molecule4,8.290648908041941
                4,Molecule5,16.996708774033824
                5,Molecule6,33.95325697202637
        
        TIPS: For some limitation, the chart created on this function might not very good.
    """
    datas=[]
    opath=os.getcwd()
    if files==[]:
        for path in os.listdir(loc):
            if os.path.isfile(os.path.join(loc, path)):
                a=re.search('_Result.txt', path)
                if a:
                    files.append(path[:-11])
        files.sort()
    #print(files)
    if energy_pattern=='Energy':
        for file in files:
            #try:
            data=status(file,loc).converged_energy()
            datas.append(data)
            #except:
            #    datas.append(None)
    elif energy_pattern=='Gibbs':
        for file in files:
            #try:
            data=status(file,loc).gibbs()
            datas.append(data)
            #except:
            #    datas.append(None)
    os.chdir(opath)
    os.chdir(loc)
    Ebar=np.array(datas)
    if absolute==False:
        Ebar=Ebar-datas[0]
    if hartree==False:
        Ebar=Ebar*627.509
    Ebar.tolist()
    Emax=-99999
    for E in Ebar:
        try:
            if E>Emax:
                Emax=E
        except:
            pass
    Emin=min(Ebar)
    gap=Emax-Emin
    if xx==False:
        if figdata==True:
            if fig_pattern=='line':
                for i in range(len(Ebar)):
                    if files[i]=='':
                        continue
                    if Emax-Ebar[i]<=5.0:
                        plt.text(i-0.5, Ebar[i]-0.05*gap, '{0:.1f}'.format(Ebar[i]))
                    elif Emax-Ebar[i]>5.0:
                        plt.text(i-0.2, Ebar[i]+0.02*gap, '{0:.1f}'.format(Ebar[i]))
            elif fig_pattern=='bar':
                for i in range(len(Ebar)):
                    plt.text(i, Ebar[i], '{0:.1f}'.format(Ebar[i]))
        if fig_pattern=='line':
            ax=plt.plot(files,Ebar,'*-')
            if energy_pattern=='Energy' and hartree==True:
                plt.ylabel('∆E(Eh)')
            elif energy_pattern=='Energy' and hartree==False:
                plt.ylabel('∆E(kcal/mol)')
            elif energy_pattern=='Gibbs' and hartree==True:
                plt.ylabel('∆G(Eh)')
            elif energy_pattern=='Gibbs' and hartree==False:
                plt.ylabel('∆G(kcal/mol)')
        elif fig_pattern=='bar':
            ax=plt.bar(files,Ebar)
            if energy_pattern=='Energy' and hartree==True:
                plt.ylabel('∆E(Eh)')
            elif energy_pattern=='Energy' and hartree==False:
                plt.ylabel('∆E(kcal/mol)')
            elif energy_pattern=='Gibbs' and hartree==True:
                plt.ylabel('∆G(Eh)')
            elif energy_pattern=='Gibbs' and hartree==False:
                plt.ylabel('∆G(kcal/mol)')
    if savefig==True:
        plt.savefig('{0}.png'.format(figname))
    if xx==False:
        plt.show()
    if savedata==True:
        try:
            f=open('{0}.csv'.format(dataname),'x')
        except:
            f=open('{0}.csv'.format(dataname),'w')
        f.write('Step,File,')
        if energy_pattern=='Energy':
            f.write('∆E')
        elif energy_pattern=='Gibbs':
            f.write('∆G')
        if hartree==True:
            f.write('(Eh)')
        else:
            f.write('(kcal/mol)')
        f.write('\n')
        for i in range(len(files)):
            f.write('{0},{1},{2}\n'.format(i,files[i],Ebar[i]))
        f.close()
    os.chdir(opath)
    return Ebar.tolist()