'''
    SASSIE: Copyright (C) 2011 Joseph E. Curtis, Ph.D. 

    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 3 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.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
'''
import os,sys,locale,time,subprocess
from write_mdx_input import *
import sassie.sasmol.sasmol as sasmol

#       OPEN_MINIMIZE
#
#       12/05/2004       --      initial coding                	:       jc
#       01/02/2011       --      added sasmol support 		:       jc
#       08/26/2011       --      adapted for mdx 			:       jc
#
#LC      1         2         3         4         5         6         7
#LC4567890123456789012345678901234567890123456789012345678901234567890123456789
#                                                                      *      **
'''
	OPEN_MINIMIZE is the module that contains the functions
	that are used to run a series of energy minimization calculations
	on a set of structures in a supplied pdb/dcd file.

	This module is called from Structure Minization in the main 
	GUI through the graphical_minimize.py script.

      REFERENCE:

	http://www.ks.uiuc.edu/Development/MDTools/namdlite/, Universty of Illinois at Urbana-Champaign, 2007.

	The code in this script calls the open-source mdx (mdsim) engine.  The
	license for this code is as follows:

	University of Illinois/NCSA
	Open Source License

	Copyright (C) 2003-2005, David J. Hardy.  All rights reserved.

	Developed by:

	David J. Hardy
	dhardy@ks.uiuc.edu

	Theoretical and Computational Biophysics Group
	Beckman Institute
	University of Illinois at Urbana - Champaign

	http://www.ks.uiuc.edu/~dhardy/
	http://www.ks.uiuc.edu/

	Permission is hereby granted, free of charge, to any person obtaining a
	copy of this software and associated documentation files (the "Software"),
	to deal with the Software without restriction, including without limitation
	the rights to use, copy, modify, merge, publish, distribute, sublicense,
	and/or sell copies of the Software, and to permit persons to whom the
	Software is furnished to do so, subject to the following conditions:
	
  	* Redistributions of source code must retain the above copyright notice,
    	this list of conditions and the following disclaimers.
	
  	* Redistributions in binary form must reproduce the above copyright
    	notice, this list of conditions and the following disclaimers in the
    	documentation and/or other materials provided with the distribution.
	
  	* Neither the names of David J. Hardy, Theoretical and Computational
    	Biophysics Group, Beckman Institute, University of Illinois, nor the
    	names of its contributors may be used to endorse or promote products
    	derived from this Software without specific prior written permission.
	
	THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
	IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
	FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
	THE CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR
	OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
	ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
	OTHER DEALINGS WITH THE SOFTWARE.

'''

def unpack_variables(variables):

        runname         = variables['runname'][0]
        infile          = variables['infile'][0]
        pdbfile         = variables['pdbfile'][0]
      #  convergence     = variables['convergence'][0]
        outfile         = variables['outfile'][0]
        nsteps          = variables['nsteps'][0]
        parmfile        = variables['parmfile'][0]
        #ncpu            = variables['ncpu'][0]
        psffile         = variables['psffile'][0]
        #energyfile      = variables['energyfile'][0]
        #keepout         = variables['keepout'][0]
        #dcdfreq         = variables['dcdfreq'][0]
        infiletype      = variables['infiletype'][0]
	
	return runname,infile,pdbfile,outfile,nsteps,parmfile,psffile,infiletype

def print_failure(message,txtOutput):

        txtOutput.put("\n\n>>>> RUN FAILURE <<<<\n")
        txtOutput.put(">>>> RUN FAILURE <<<<\n")
        txtOutput.put(">>>> RUN FAILURE <<<<\n\n")
        txtOutput.put(message)

        return

def minimize(variables,txtOutput):

	'''
        MINIMIZE is the function to read in variables from GUI input and 
        used to run a series of energy minimization calculations
	on a set of structures in a supplied pdb/dcd file.

        INPUT:  variable descriptions:
              
		runname:              	run_name 
                infile:               	input pdb or dcd filename
                pdbfile:               	input pdb file (reference)
                nsteps:               	number of steps
                parmfile:              	path nand name of topology file
			psffile: 		name of psf file

        OUTPUT:
                
                txtOutput:              TK handler for output to GUI textbox
                
		files stored in ~/run_name/minimize directory:
                outfile:                   output filename (dcd usually)

        '''
	
	runname,infile,pdbfile,outfile,nsteps,parmfile,psffile,infiletype=unpack_variables(variables)

	keepout = 1
	dcdfreq = 1

	path=runname+'/minimization/' 
	print 'path = ',path
	print 'infile = ',infile

	vers='version 0.6 : 07/16/08 : jc'
	direxist=os.path.exists(path)
	if(direxist==0):
		try:
                	result=os.system('mkdir -p '+ path)
			os.system('cp '+psffile+' '+path)	
			print 'result = ',result		
		except:
			message='can not create project directory: '+path
			message+='\nstopping here\n'
			print_failure(message,txtOutput)
		if(result!=0):
			message='can not create project directory: '+path
			message+='\nstopping here\n'
			print_failure(message,txtOutput)

	m1=sasmol.SasMol(0)
	m1.read_pdb(pdbfile)

	print 'infiletype = ',infiletype

	if(infiletype=='dcd'):
		dcdfile = m1.open_dcd_read(infile)
		nf = dcdfile[2]       
	else:
		m1.read_pdb(infile)
        	nf=m1.coor()[:,0,0].shape[0]

	print 'number of frames = ',nf

	ttxt=time.ctime()
  	st=''.join(['=' for x in xrange(60)])

	txtOutput.put("\n%s \n" %(st))
	txtOutput.put("DATA FROM RUN: %s \n\n" %(ttxt))

	final_energy=[] ; dcdlist=[] ; coorlist=[]
	for i in range(nf):
		print 'minimizing frame ',i+1,' of ',nf
		print 'minimizing frame ',i+1,' of ',nf
		print 'minimizing frame ',i+1,' of ',nf
		print 'writing temporary PDB file'

		if(infiletype == 'dcd'):
			m1.read_dcd_step(dcdfile,i)
			m1.write_pdb(path+'junk.pdb',0,'w')
		else:
			m1.write_pdb(path+'junk.pdb',i,'w')
			
		print 'writing temporary NAMD input file'
		if(i<9):
			istr='0000'+str(i+1)
		elif(i<99):
			istr='000'+str(i+1)
		elif(i<999):
			istr='00'+str(i+1)
		elif(i<9999):
			istr='0'+str(i+1)
		elif(i<99999):
			istr=str(i+1)
		else:
			print 'wow, man!'
			istr=str(i+1)

		thisdcd=path+'min_'+istr+'.dcd'
		dcdlist.append(thisdcd)

		write_mdx_input('temp.inp',str(nsteps),str(dcdfreq),path+'junk.pdb',path+psffile,thisdcd,parmfile) 
		print 'starting minimization ( nfiles = ',nf,')'
		ttime=time.ctime()	
		runstring=vers+' : '+outfile+' run stated at : '+ttime
		print runstring
		ncpu=1
		if(ncpu==1):
			print 'starting mdx minimization'
			nst='/usr/local/bin/mdsim temp.inp >& junk.out &'

			p=subprocess.Popen(nst,shell=True,executable='/bin/bash')
			sts=os.waitpid(p.pid,0)[1]
			print 'p.pid = ',p.pid
			thisjob=str(int(p.pid)+1)	

		print 'starting mdsim : thisjob = ',thisjob

		run=1
		esteps=0
		while(run==1):
			#time.sleep(5)
			lsst='ls junk.out | grep -c "junk.out" '
			lsfile=os.popen(lsst,'r').readlines()
			stls=string.split(lsfile[0])
			nstls=locale.atoi(stls[0])
			if(nstls>0):			
				tout2=os.popen('tail -15 junk.out | grep Success! ','r').readlines()

			if(len(tout2)>0):
				print 'finished minimization'
				run=0	

		fraction_done = (float(i+1)/float(nf))
		progress_string='COMPLETED '+str(i+1)+' of '+str(nf)+' : '+str(fraction_done*100.0)+' % done'
		print('%s\n' % progress_string)
		print('%s\n' % progress_string)
		report_string='STATUS\t'+str(fraction_done)
		txtOutput.put(report_string)

		print 'finished run'
		
		os.system('mv junk.out '+path+'min_'+istr+'.out') 
		os.system('mv energy_results.out '+path+'energy_results_'+istr+'.out') 
		os.system('rm -f junk.out junk.coor junk.xs* junk.vel '+path+'junk.pdb')
		os.system('rm -f '+path+' *.BAK')
			
		print 'thisdcd = ',thisdcd	

		temp_mol = sasmol.SasMol(0)
		temp_mol.read_pdb(pdbfile)
			
		header = temp_mol.open_dcd_read(thisdcd)
		temp_mol.close_dcd_read(header[0])

		ndcdfiles = header[2]

		# get the last configuration from the dcd file
		if ndcdfiles>1:
			temp_mol.read_single_dcd_step(thisdcd,ndcdfiles)
			temp_mol.write_dcd(thisdcd)

		if ndcdfiles<1:
			print 'ndcdfiles = ',ndcdfiles
			message='Did not save any dcd files.  Decrease dcd write frequency?'		
			message+' :  stopping here'
			print_failure(message,txtOutput)	
	
	if(infiletype == 'dcd'):
		m1.close_dcd_read(dcdfile[0])

	print 'dcdlist = ',dcdlist

	final_mol = sasmol.SasMol(0)
	final_mol.read_pdb(pdbfile)

	finaldcdfile = final_mol.open_dcd_write(outfile)

	for i in range(len(dcdlist)):
		final_mol.read_dcd(dcdlist[i])
		final_mol.write_dcd_step(finaldcdfile,0,i+1)

	final_mol.close_dcd_write(finaldcdfile)

	rmcmd='rm -f ' 
	for i in range(len(dcdlist)):
		rmcmd=rmcmd+dcdlist[i]+' '
	print 'rmcmd = ',rmcmd
	os.system(rmcmd)
	os.system('mv temp.inp '+path)

 	txtOutput.put("Total number of frames = %d\n\n" % (nf))
	txtOutput.put("Minimized structures saved to : %s\n" % ('./'+path))
	txtOutput.put("\n%s \n" %(st))
	time.sleep(0.5)
	print 'OPEN MINIMIZE IS DONE'

	return()