""""

Script to export number of constraints for RECOORD projects.
Does not work with pre-2007 releases.

"""


from memops.general.Io import readXmlProjectFile
import sys

def getResonanceToResidueDict(nmrConstraintStore):

  resonanceToResidueDict = {}
  
  for fixedResonance in nmrConstraintStore.fixedResonances:
    if fixedResonance.resonanceSet:
      refAtom = fixedResonance.resonanceSet.atomSets[0].atoms[0]
      residue = refAtom.residue
      
      resonanceToResidueDict[fixedResonance] = residue
  
  return resonanceToResidueDict
  
def setConstraintsPerResidue(resonance,resonanceToResidueDict,constraintsPerResidueAndList,addItem = 1):

  if resonanceToResidueDict.has_key(resonance):
    residue = resonanceToResidueDict[resonance]
    if not constraintsPerResidueAndList.has_key(residue):
      constraintsPerResidueAndList[residue] = 0
    constraintsPerResidueAndList[residue] += addItem

if __name__ == '__main__':

  #
  # Get script input variables
  #
  
  xmlFile = sys.argv[1]
    
  #
  # Load the CCPN project
  #
  
  ccpnProject = readXmlProjectFile(file = xmlFile)
  
  #
  # Get the correct structure generation
  #

  nmrProject = ccpnProject.nmrProjects[0]
  strucGen = nmrProject.structureGenerations[0]
  nmrConstraintStore = strucGen.nmrConstraintStore
  
  #
  # Set up dictionary to connect CCPN 'resonances' to the residues in the chain
  #
  
  resonanceToResidueDict = getResonanceToResidueDict(nmrConstraintStore)
  
  #
  # Loop over the lists, get results
  #
  
  constraintListTypes = ('Distance','Dihedral','HBond','Rdc')
  constraintsPerResidue = {}
  
#   print constraintHead.constraintLists
  
  for constraintListType in constraintListTypes:
  
    constraintLists = nmrConstraintStore.findAllConstraintLists(className = "%sConstraintList" % constraintListType)
    
    if constraintLists:
      constraintsPerResidue[constraintListType] = {}
      
      for constraintList in constraintLists:
        
        for constraint in constraintList.constraints:
          
          if constraintListType == 'Dihedral':
            for resonance in constraint.resonances:
              setConstraintsPerResidue(resonance,resonanceToResidueDict,constraintsPerResidue[constraintListType],addItem = 0.25)
          
          else:
            constraintItems = constraint.items
            
            contribPerItem = 1 / len(constraintItems)
                        
            for item in constraintItems:
              for resonance in item.resonances:
                setConstraintsPerResidue(resonance,resonanceToResidueDict,constraintsPerResidue[constraintListType], addItem = contribPerItem / 2.0)
          
  #
  # Print results
  #
  
  constraintListTypes = constraintsPerResidue.keys()
  constraintListTypes.sort()
  
  print "%15s" % "",
  for constraintListType in constraintListTypes:
    print "%-10s" % constraintListType,
  print
  
  for chain in ccpnProject.molSystems[0].chains:
    for residue in chain.residues:
      resId = "'%s.%s.%s':" % (chain.code,residue.ccpCode,residue.seqCode)
      print "%-15s" % resId,
      
      for constraintListType in constraintListTypes:
        if constraintsPerResidue[constraintListType].has_key(residue):
          numConstraints = constraintsPerResidue[constraintListType][residue]
        else:
          numConstraints = 0

        print "%6.2f    " % numConstraints,
      
      print