#!/usr/bin/python
# -*- coding: utf-8 -*-
#
#--------------------------------------------------------------------------------
#anaout.py v1.1, Copyright Bjoern Olausson
#--------------------------------------------------------------------------------
#This program is free software; you can redistribute it and/or modify
#it under the terms of the GNU General Public License as published by
#the Free Software Foundation; either version 2 of the License, or
#(at your option) any later version.
#
#This program is distributed in the hope that it will be useful,
#but WITHOUT ANY WARRANTY; without even the implied warranty of
#MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
#GNU General Public License for more details.
#
#To view the license visit
#http://www.gnu.org/licenses/old-licenses/gpl-2.0.html
#or write to
#Free Software Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA
#--------------------------------------------------------------------------------
#--------------------------------------------------------------------------------
#
#This script is intended to pars CHARMM output files and plot energy, temperature, and
#whatever information CHARMM writes down for DYNA.
#
#For example:
#
#DYNA DYN: Step         Time      TOTEner        TOTKe       ENERgy  TEMPerature
#DYNA PROP:             GRMS      HFCTote        HFCKe       EHFCor        VIRKe
#DYNA INTERN:          BONDs       ANGLes       UREY-b    DIHEdrals    IMPRopers
#DYNA CROSS:           CMAPs
#DYNA EXTERN:        VDWaals         ELEC       HBONds          ASP         USER
#DYNA IMAGES:        IMNBvdw       IMELec       IMHBnd       RXNField    EXTElec
#DYNA EWALD:          EWKSum       EWSElf       EWEXcl       EWQCor       EWUTil
#DYNA CONSTR:       HARMonic    CDIHedral          CIC     RESDistance       NOE
#DYNA MMFP:              GEO        MDIP           SSBP        SHEL       DROFfa
#DYNA PRESS:            VIRE         VIRI       PRESSE       PRESSI       VOLUme
#DYNA XTLE:                       XTLTe         SURFtension  XTLPe        XTLtemp
#----------       ---------    ---------    ---------    ---------    ---------

#DYNA>        0     50.00000-154327.02601  76739.59712-231066.62313    303.04814
#DYNA PROP>         16.53147-154187.74874  77157.31193    139.27726  35759.35840
#DYNA INTERN>     6875.17059  26031.47702   8099.17560  20014.17660    468.25788
#DYNA CROSS>      -671.51578
#DYNA EXTERN>     5323.04913-211972.77063      0.00000      0.00000      0.00000
#DYNA IMAGES>    -1480.14932 -13098.37023      0.00000      0.00000      0.00000
#DYNA EWALD>       4283.1528-1565911.3520 1490022.1635       0.0000       0.0000
#DYNA CONSTR>      410.87665      0.00000      0.00000      0.00000      0.00000
#DYNA MMFP>        540.03511      0.00000      0.00000      0.00000      0.00000
#DYNA PRESS>      -5704.1194  -18135.4529     364.4002    2127.5031 1073329.5749
#DYNA XTLE>                  -14733.71807    -52.29616  -2039.36602    128.28182
#----------       ---------    ---------    ---------    ---------    ---------
import sys, os, shutil, string, glob, subprocess, itertools, pylab, scipy

from optparse import OptionParser, OptionGroup
opa = OptionParser()

opa.add_option("-f", action="append", dest="chmout", metavar="FILE",
	help="CHARMM out file")
opa.add_option("-o", action="store", dest="saveplot", metavar="FILE",
	help="Filename to which we save the plot (.png) and RAW data (.dat)")
	
opa.add_option("-s", action="store_true", dest="singleplot",
	help="Plot to single image")
	
group = OptionGroup(opa, "You can choose freely from the following Values which will occure on either X or Y Axis",
" Step Time TOTEner TOTKe ENERgy TEMPerature GRMS HFCTote HFCKe EHFCor VIRKe BONDs ANGLes UREYb DIHEdrals IMPRopers CMAPs VDWaals ELEC HBONds ASP USER IMNBvdw IMELec IMHBnd RXNField EXTElec EWKSum EWSElf EWEXcl EWQCor EWUTil HARMonic CDIHedral CIC RESDistance NOE GEO MDIP SSBP SHEL DROFfa VIRE VIRI PRESSE PRESSI VOLUme XTLTe SURFtension XTLPe XTLtemp")

group.add_option("-x", action="append", dest="xax",
	help="Choose Value from above you want to appear on the X-axis.")

group.add_option("-y", action="append", dest="yax",
	help="Choose Value from above you want to appear on the Y-axis. If you want to plot multible Y values, you have to asign every Y value a X value and vice versa. For example: anaout.py -x Time -x Time -y PRESSE -y TEMPerature")
opa.add_option_group(group)

# instruct optparse to parse the program's command line
(options, args) = opa.parse_args()

# Checking for required otions
if not options.chmout:
	parser.error("-f must be set")

#chmout = options.chmout.split(" ")
chmout = options.chmout
xax = options.xax
yax = options.yax
saveplot = options.saveplot
singleplot = options.singleplot

#If "*" is found in a parameter, glob it.
for element in chmout :
	if '*' in element :
		chmout.extend(glob.glob(element))
		chmout.remove(element)

#Defining the lists
Step = [];	Time = [];	TOTEner = [];	TOTKe = [];		ENERgy = [];	TEMPerature = []
GRMS = [];	HFCTote = [];	HFCKe = [];	EHFCor = [];		VIRKe = []
BONDs = [];	ANGLes = [];	UREYb = [];	DIHEdrals = [];		IMPRopers = []
CMAPs = []
VDWaals = [];	ELEC = [];	HBONds = [];	ASP = [];		USER = []
IMNBvdw = [];	IMELec = [];	IMHBnd = [];	RXNField = [];		EXTElec = []
EWKSum = [];	EWSElf = [];	EWEXcl = [];	EWQCor = [];		EWUTil = []
HARMonic = [];	CDIHedral = [];	CIC = [];	RESDistance = [];	NOE = []
GEO = [];	MDIP = [];	SSBP = [];	SHEL = [];		DROFfa = []
VIRE  = [];	VIRI = [];	PRESSE = [];	PRESSI = [];		VOLUme = []
XTLTe = [];	SURFtension = [];	XTLPe = [];	XTLtemp = []

#DYNA DYN: Step         Time      TOTEner        TOTKe       ENERgy  TEMPerature
#DYNA>        0     50.00000-154327.02601  76739.59712-231066.62313    303.04814
def Dyna(line) :
	if fnum > 0 :
		CurStep = int(line[5:14].strip())
		NewStep = CurStep + LastStep
		Step.append(str(NewStep))
	else :
		Step.append(line[5:14].strip())

	Time.append(line[14:27].strip())
	TOTEner.append(line[27:40].strip())
	TOTKe.append(line[40:53].strip())
	ENERgy.append(line[53:66].strip())
	TEMPerature.append(line[66:79].strip())

#DYNA PROP:             GRMS      HFCTote        HFCKe       EHFCor        VIRKe
#DYNA PROP>         16.53147-154187.74874  77157.31193    139.27726  35759.35840
def DynaProp(line) :
	GRMS.append(line[10:27].strip())
	HFCTote.append(line[27:40].strip())
	HFCKe.append(line[40:53].strip())
	EHFCor.append(line[53:66].strip())
	VIRKe.append(line[66:79].strip())
	
#DYNA INTERN:          BONDs       ANGLes       UREY-b    DIHEdrals    IMPRopers
#DYNA INTERN>     6875.17059  26031.47702   8099.17560  20014.17660    468.25788
def DynaIntern(line) :
	BONDs.append(line[12:27].strip())
	ANGLes.append(line[27:40].strip())
	UREYb.append(line[40:53].strip())
	DIHEdrals.append(line[53:66].strip())
	IMPRopers.append(line[66:79].strip())
	
#DYNA CROSS:           CMAPs
#DYNA CROSS>      -671.51578
def DynaCross(line) :
	CMAPs.append(line[11:27].strip())

#DYNA EXTERN:        VDWaals         ELEC       HBONds          ASP         USER
#DYNA EXTERN>     5323.04913-211972.77063      0.00000      0.00000      0.00000
def DynaExtern(line) :
	VDWaals.append(line[12:27].strip())
	ELEC.append(line[27:40].strip())
	HBONds.append(line[40:53].strip())
	ASP.append(line[53:66].strip())
	USER.append(line[66:79].strip())
	
#DYNA IMAGES:        IMNBvdw       IMELec       IMHBnd       RXNField    EXTElec
#DYNA IMAGES>    -1480.14932 -13098.37023      0.00000      0.00000      0.00000
def DynaImages(line) :
	IMNBvdw.append(line[12:27].strip())
	IMELec.append(line[27:40].strip())
	IMHBnd.append(line[40:53].strip())
	RXNField.append(line[53:66].strip())
	EXTElec.append(line[66:79].strip())
	
#DYNA EWALD:          EWKSum       EWSElf       EWEXcl       EWQCor       EWUTil
#DYNA EWALD>       4283.1528-1565911.3520 1490022.1635       0.0000       0.0000
def DynaEwald(line) :
	EWKSum.append(line[11:27].strip())
	EWSElf.append(line[27:40].strip())
	EWEXcl.append(line[40:53].strip())
	EWQCor.append(line[53:66].strip())
	EWUTil.append(line[66:79].strip())
	
#DYNA CONSTR:       HARMonic    CDIHedral          CIC     RESDistance       NOE
#DYNA CONSTR>      410.87665      0.00000      0.00000      0.00000      0.00000
def DynaConstr(line) :
	HARMonic.append(line[12:27].strip())
	CDIHedral.append(line[27:40].strip())
	CIC.append(line[40:53].strip())
	RESDistance.append(line[53:66].strip())
	NOE.append(line[66:79].strip())
	
#DYNA MMFP:              GEO        MDIP           SSBP        SHEL       DROFfa
#DYNA MMFP>        540.03511      0.00000      0.00000      0.00000      0.00000
def DynaMmfp(line) :
	GEO.append(line[10:27].strip())
	MDIP.append(line[22:40].strip())
	SSBP.append(line[40:53].strip())
	SHEL.append(line[53:66].strip())
	DROFfa.append(line[66:79].strip())
	
#DYNA PRESS:            VIRE         VIRI       PRESSE       PRESSI       VOLUme
#DYNA PRESS>      -5704.1194  -18135.4529     364.4002    2127.5031 1073329.5749
def DynaPress(line) :
	VIRE.append(line[11:27].strip())
	VIRI.append(line[22:40].strip())
	PRESSE.append(line[40:53].strip())
	PRESSI.append(line[53:66].strip())
	VOLUme.append(line[66:79].strip())

#DYNA XTLE:                       XTLTe         SURFtension  XTLPe        XTLtemp
#DYNA XTLE>                  -14733.71807    -52.29616  -2039.36602    128.28182
def DynaXtle(line) :
	XTLTe.append(line[22:40].strip())
	SURFtension.append(line[40:53].strip())
	XTLPe.append(line[53:66].strip())
	XTLtemp.append(line[66:79].strip())

LineFunc = {
	'DYNA': Dyna,
	'DYNA PROP': DynaProp,
	'DYNA INTERN': DynaIntern,
	'DYNA CROSS': DynaCross,
	'DYNA EXTERN': DynaExtern,
	'DYNA IMAGES': DynaImages,
	'DYNA EWALD': DynaEwald,
	'DYNA CONSTR': DynaConstr,
	'DYNA MMFP': DynaMmfp,
	'DYNA PRESS': DynaPress,
	'DYNA XTLE': DynaXtle
}

fnum = 0
for outfile in chmout :
	for lines in open(outfile, "r") :
		LineStart = lines.split(">")
		if LineStart[0] in LineFunc :
			LineFunc[LineStart[0]](lines)
#	print "Step: ", Step
	LastStep = int(Step[-1])
#	print "LastStep: ", LastStep
	fnum = fnum + 1

#Convert Lists to scipy arrays
Step = scipy.array(Step)
Time = scipy.array(Time)
TOTEner = scipy.array(TOTEner)
TOTKe = scipy.array(TOTKe)
ENERgy = scipy.array(ENERgy)
TEMPerature = scipy.array(TEMPerature)

GRMS = scipy.array(GRMS)
HFCTote = scipy.array(HFCTote)
HFCKe = scipy.array(HFCKe)
EHFCor = scipy.array(EHFCor)
VIRKe = scipy.array(VIRKe)

BONDs = scipy.array(BONDs)
ANGLes = scipy.array(ANGLes)
UREYb = scipy.array(UREYb)
DIHEdrals = scipy.array(DIHEdrals)
IMPRopers = scipy.array(IMPRopers)

CMAPs = scipy.array(CMAPs)

VDWaals = scipy.array(VDWaals)
ELEC = scipy.array(ELEC)
HBONds = scipy.array(HBONds)
ASP = scipy.array(ASP)
USER = scipy.array(USER)

IMNBvdw = scipy.array(IMNBvdw)
IMELec = scipy.array(IMELec)
IMHBnd = scipy.array(IMHBnd)
RXNField = scipy.array(RXNField)
EXTElec = scipy.array(EXTElec)

EWKSum = scipy.array(EWKSum)
EWSElf = scipy.array(EWSElf)
EWEXcl = scipy.array(EWEXcl)
EWQCor = scipy.array(EWQCor)
EWUTil = scipy.array(EWUTil)

HARMonic = scipy.array(HARMonic)
CDIHedral = scipy.array(CDIHedral)
CIC = scipy.array(CIC)
RESDistance = scipy.array(RESDistance)
NOE = scipy.array(NOE)

GEO = scipy.array(GEO)
MDIP = scipy.array(MDIP)
SSBP = scipy.array(SSBP)
SHEL = scipy.array(SHEL)
DROFfa = scipy.array(DROFfa)

VIRE = scipy.array(VIRE)
VIRI = scipy.array(VIRI)
PRESSE = scipy.array(PRESSE)
PRESSI = scipy.array(PRESSI)
VOLUme = scipy.array(VOLUme)

XTLTe = scipy.array(XTLTe)
SURFtension = scipy.array(SURFtension)
XTLPe = scipy.array(XTLPe)
XTLtemp = scipy.array(XTLtemp)

#COMBINED = scipy.array([[Step], [Time], [TOTEner], [TOTKe], [ENERgy], [TEMPerature], [GRMS], [HFCTote], [HFCKe], [EHFCor], [VIRKe], [BONDs], [ANGLes], [UREYb], [DIHEdrals], [IMPRopers], [CMAPs], [VDWaals], [ELEC], [HBONds], [ASP], [USER], [IMNBvdw], [IMELec], [IMHBnd], [RXNField], [EXTElec], [EWKSum], [EWSElf], [EWEXcl], [EWQCor], [EWUTil], [HARMonic], [CDIHedral], [CIC], [RESDistance], [NOE], [GEO], [MDIP], [SSBP], [SHEL], [DROFfa], [VIRE], [VIRI], [PRESSE], [PRESSI], [VOLUme]])
PlotDict = {
	'Step': Step,
	'Time': Time,
	'TOTEner': TOTEner,
	'TOTKe': TOTKe,
	'ENERgy': ENERgy,
	'TEMPerature': TEMPerature,
	'GRMS': GRMS,
	'HFCTote': HFCTote,
	'HFCKe': HFCKe,
	'EHFCor': EHFCor,
	'VIRKe': VIRKe,
	'BONDs': BONDs,
	'ANGLes': ANGLes,
	'UREYb': UREYb,
	'DIHEdrals': DIHEdrals ,
	'CMAPs': CMAPs,
	'IMPRopers': IMPRopers,
	'ASP': ASP,
	'USER': USER,
	'IMNBvdw': IMNBvdw,
	'IMELec': IMELec,
	'IMHBnd': IMHBnd,
	'RXNField': RXNField,
	'EXTElec': EXTElec,
	'EWKSum': EWKSum,
	'EWSElf': EWSElf,
	'EWEXcl': EWEXcl,
	'EWQCor': EWQCor,
	'EWUTil': EWUTil,
	'HARMonic': HARMonic,
	'CDIHedral': CDIHedral,
	'CIC': CIC,
	'RESDistance': RESDistance,
	'NOE': NOE,
	'GEO': GEO,
	'MDIP': MDIP,
	'SSBP': SSBP,
	'SHEL': SHEL,
	'DROFfa': DROFfa,
	'VIRE': VIRE,
	'VIRI': VIRI,
	'PRESSE': PRESSE,
	'PRESSI': PRESSI,
	'VOLUme': VOLUme,
	'XTLTe' : XTLTe,
	'SURFtension' : SURFtension,
	'XTLPe' : XTLPe,
	'XTLtemp' : XTLtemp
}

def SaveData(Xvalues,Yvalues, Xname, Yname, saveplot) :
	if saveplot :
		savedata = Xname+"-"+Yname+"_"+saveplot+".dat"
		fout = open(savedata, "w")
		for i,j in itertools.izip(Xvalues,Yvalues) :
			fout.writelines("%s %s \n" % (i,j))
		fout.close()

def SaveImage(Xname, Yname, saveplot) :
	if saveplot :
		saveimg = Xname+"-"+Yname+"_"+saveplot+".png"
		pylab.savefig(saveimg)

def DynaPltSingle(Xvalues,Yvalues,Xname,Yname) :
	#Do some nice plots later
	SaveData(Xvalues, Yvalues, Xname, Yname, saveplot)
	pylab.plot(Xvalues, Yvalues)
	Xlegend.append(Xname)
	Ylegend.append(Yname)

def DynaPltMulti(Xvalues,Yvalues,Xname,Yname, figure) :
	pylab.figure(figure)
	pylab.plot(Xvalues, Yvalues)
	pylab.xlabel(Xname)
	pylab.ylabel(Yname)
	SaveImage(Xname, Yname, saveplot)
	SaveData(Xvalues,Yvalues, Xname, Yname, saveplot)

fign = 0
Xlegend = []
Xlabel = []
Ylegend = []
Ylabel = []
if singleplot :
	for xAxis,yAxis in itertools.izip(xax,yax) :
		DynaPltSingle(PlotDict[xAxis], PlotDict[yAxis], xAxis, yAxis)
	pylab.legend((Ylegend), shadow = True)
	pylab.ylabel(Ylegend)
	pylab.xlabel(Xlegend)
	SaveImage(xAxis, 'MULTI', saveplot)
else :
	figure = 0
	for xAxis,yAxis in itertools.izip(xax,yax) :
		DynaPltMulti(PlotDict[xAxis], PlotDict[yAxis], xAxis, yAxis, figure)
		figure = figure + 1
		

pylab.show()
