'''
Created on Jan 17, 2014

@author: oliwa
'''
import sys as sys
import numpy as np
from prody.dynamics.anm import calcANM, ANM
from prody.dynamics.editing import extendModel, sliceModel
from prody.dynamics.functions import saveModel, loadModel, writeArray
from prody.proteins.pdbfile import writePDB, parsePDB
from prody.dynamics.mode import Vector
from prody.measure.measure import calcCenter, calcDistance
from prody.dynamics.compare import calcOverlap, calcCumulOverlap,\
	calcSubspaceOverlap, calcCovOverlap, printOverlapTable, getOverlapTable
from prody.apps.prody_apps.prody_contacts import prody_contacts
import traceback
from prody.dynamics.nmdfile import writeNMD
import scipy as
 sp

class ANMs(object):
	"""
	This class holds all the ANMs for an encounter.
	"""

	def __init__(self, utils):
		"""
		Constructor
		"""
		self.utils = utils
		
	def createSlcS
electionString(self, reference, isBoundComplex, ref_chain, referenceTitle):
		""" Under the assumption that is reflected in the Benchmark 4.0 that the receptor atoms are set before the
		ligand atoms (spacially in the PDB file), if the current protein under investigation is a ligand, 
		an offset is added to the selection string to match the atoms of the ligand from the complex. """
		if isBoundComplex and not self.utils.isReceptor(referenceTitle):
			print "adding offset"
			return self.utils.addOffset(ref_chain.getSelstr(), reference.select('segment "R."').numAtoms())
		else:
			print "using original selstr"
			return ref_chain.getSelstr()
		
	def calcANMs(self, reference, ref_chain, numberOfModes, encounter, selstr='calpha', whatAtomsToMatch='calpha', modified="", forceRebuild=False, isBoundComplex=False):
		# if the base model does not exist, it needs to be created along with the
		# extended and slicedback models
		if forceRebuild or not self.doesANMExist(reference, numberOfModes, selstr, whatAtomsToMatch, modified):
			# Create the anm
			anm = calcANM(reference, n_modes=numberOfModes, selstr=selstr)
			# First extend the anm on all atoms
			anm_extend = extendModel(anm[0], anm[1], reference, norm=True)
			# Then slice it back to matched
			selectionAtoms = self.createSlcSelectionString(reference, isBoundComplex, ref_chain, encounter.getReference().getTitle())
			anm_slc = sliceModel(anm_extend[0], anm_extend[1], selectionAtoms)
			# If isBoundComplex, slice one anm back to its overall matched chains
			if isBoundComplex:
				selectionAtomsCounterpart = self.createSlcSelectionString(reference, isBoundComplex, encounter.getBoundCounterpartChain(), encounter.getUnboundCounterpart().getTitle())
				anm_slc_counterpart= sliceModel(anm_extend[0], anm_extend[1], selectionAtomsCounterpart)
				
			# Save the models
#             saveModel(anm[0], 
#                       filename=self.getANMPath(reference, numberOfModes, selstr, whatAtomsToMatch), 
#                       matrices=True)
#             saveModel(anm_extend[0],
#                       filename=self.getANMPath(reference, numberOfModes, selstr, whatAtomsToMatch, modified="extended"),
#                       matrices=True
#                       )
#             saveModel(anm_slc[0],
#                       filename=self.getANMPath(reference, numberOfModes, selstr, whatAtomsToMatch, modified="slicedback"),
#                       matrices=True
#                       )
			print "created and saved models"
#             print "reference, it is the complex: ", reference.select('calpha and segment "R."').numAtoms()
#             print "to slice on, it is the mob_chain: ", ref_chain.numAtoms()
			print "anm hessian               : " + str(anm[0].getHessian().shape)
			print "number of calpha          : " + str(reference.select('calpha').numAtoms())
			print "anm size                  : " + str(anm[0].getArray().shape)
			print "anm_ext size              : " + str(anm_extend[0].getArray().shape)
			print "anm_slice size            : " + str(anm_slc[0].getArray().shape)
			print "selectionAtoms            : " + selectionAtoms
			if isBoundComplex:
				print "anm slice counterpart size: " + str(anm_slc_counterpart[0].getArray().shape)
				print "selectionAtoms counterpart: " + selectionAtomsCounterpart
			# Save the models"
			self._anm = anm
			self._anm_extend = anm_extend
			self._anm_slc = anm_slc
			if isBoundComplex:
				self._anm_slc_counterpart = anm_slc_counterpart
		else:
			#raise Exception("Problem with capturing the selection of saved models, do not use load models from files now.")
			try:
				# load models
				anmModel = loadModel(self.getANMPath(reference, numberOfModes, selstr, whatAtomsToMatch)+".anm.npz")
				anm_extendModel = loadModel(self.getANMPath(reference, numberOfModes, selstr, whatAtomsToMatch, modified="extended")+".nma.npz")
				anm_slcModel = loadModel(self.getANMPath(reference, numberOfModes, selstr, whatAtomsToMatch, modified="slicedback")+".nma.npz")
				
				# store models selections
				anmModelSelection = reference.select(selstr)
				anm_extendModelSelection = reference
				selectionAtoms = self.createSlcSelectionString(reference, isBoundComplex, ref_chain)
				anm_slcModelSelection = reference.select(selectionAtoms)
				
				# recombine models and selections as tuples
				anm = (anmModel, anmModelSelection)
				anm_extend = (anm_extendModel, anm_extendModelSelection)
				anm_slc = (anm_slcModel, anm_slcModelSelection)
				
				print "loaded models"
				print "anm size      : " + str(anm[0].getArray().shape)
				print "anm_ext size  : " + str(anm_extend[0].getArray().shape)
				print "anm_slice size: " + str(anm_slc[0].getArray().shape)
				print "selectionAtoms: " + selectionAtoms
				self._anm = anm
				self._anm_extend = anm_extend
				self._anm_slc = anm_slc
			except IOError as e:
				print "Error loading ANM models from disc: "+str(e)
				
	def calcANMsForPart2a2k(self, reference, counterpart, proteinComplex, ref_chain, counterpart_chain, chain_complex, numberOfModes, selstr='calpha', whatAtomsToMatch='calpha'):
			# Create the anm of reference, counterpart and proteinComplex)
#             print "reference, counterpart, proteinComplex, chain_complex (calphas, calphas*3-6)          : ", (reference.select('calpha').numAtoms(), reference.select('calpha').numAtoms()*3 -6), (counterpart.select('calpha').numAtoms(), counterpart.select('calpha').numAtoms()*3-6), (proteinComplex.select('calpha').numAtoms(), proteinComplex.select('calpha').numAtoms()*3-6), (chain_complex.select('calpha').numAtoms(), chain_complex.select('calpha').numAtoms()*3 -6)
#             print "anm_reference, anm_counterpart, anm_complex hessian shapes                     : ", anm_reference[0].getHessian().shape, anm_counterpart[0].getHessian().shape, anm_complex[0].getHessian().shape
#             print "anm_reference, anm_counterpart, anm_complex, anm_complex_slc getArray() shapes : ", anm_reference[0].getArray().shape, anm_counterpart[0].getArray().shape, anm_complex[0].getArray().shape, anm_complex_slc[0].getArray().shape
			self._anm_reference, self._anm_reference_slc = self._calcANMsUnified(reference, ref_chain, numberOfModes/2, selstr, whatAtomsToMatch)
			self._anm_counterpart, self._anm_counterpart_slc = self._calcANMsUnified(counterpart, counterpart_chain, numberOfModes/2, selstr, whatAtomsToMatch)
			
#             print "15 ang contact before moving atoms:", proteinComplex.select('same residue as exwithin 15 of segment "L." ').numAtoms()
#             self._moveSegment(proteinComplex, "L", 30)
#             if proteinComplex.select('same residue as exwithin 15 of segment "L." ') != None:
#                 print "15 ang contact after moving atoms: ", proteinComplex.select('same residue as exwithin 15 of segment "L." ').numAtoms()
#             else:
#                 print "15 ang contact after moving atoms: 0"
			self._anm_complex, self._anm_complex_slc = self._calcANMsUnified(proteinComplex, chain_complex, numberOfModes, selstr, whatAtomsToMatch)
			
			#self.utils.testHessianSubMatrices(self._anm_reference, self._anm_counterpart, self._anm_complex)
			# check blockmatrix differences and pymol output
#             useRelError = True
			#significantDifferences = self.utils.testBlockMatrixMembership(self.