"""
======================COPYRIGHT/LICENSE START==========================

CloudHomologueAssign.py: Part of the CcpNmr Clouds program

Copyright (C) 2003-2010 Wayne Boucher and Tim Stevens (University of Cambridge)

=======================================================================

The CCPN license can be found in ../../../license/CCPN.license.

======================COPYRIGHT/LICENSE END============================

for further information, please contact :

- CCPN website (http://www.ccpn.ac.uk/)

- email: ccpn@bioc.cam.ac.uk

- contact the authors: wb104@bioc.cam.ac.uk, tjs23@cam.ac.uk
=======================================================================

If you are using this software for academic purposes, we suggest
quoting the following references:

===========================REFERENCE START=============================
A. Grishaev and M. Llinas (2002).
CLOUDS, a protocol for deriving a molecular proton density via NMR.
Proc Natl Acad Sci USA. 99, 6707-6712.

Wim F. Vranken, Wayne Boucher, Tim J. Stevens, Rasmus
H. Fogh, Anne Pajon, Miguel Llinas, Eldon L. Ulrich, John L. Markley, John
Ionides and Ernest D. Laue (2005). The CCPN Data Model for NMR Spectroscopy:
Development of a Software Pipeline. Proteins 59, 687 - 696.

===========================REFERENCE END===============================

"""
from ccpnmr.analysis.core.AssignmentBasic  import assignSpinSystemResidue
from ccpnmr.analysis.core.MoleculeBasic    import newMolSystem
from ccp.lib.MoleculeAlign import getChainResidueMapping
from ccp.lib.StructureIo      import getStructureFromFile
from ccp.lib.StructureLib     import getAtomSetCoords

from ccpnmr.clouds.CloudThreader import cloudThreader
from ccpnmr.clouds.CloudBasic    import getMeanCloud

def cloudHomologueAssign(clouds, chain, pdbFileName, threshold=3.0, nSteps=3000, graph=None, progressBar=None):

  project   = chain.root
  molSystem = newMolSystem(project)
  structure = getStructureFromFile(molSystem, pdbFileName) # Done
  chainH    = structure.coordChains[0].chain

  # Match the amide protons of the homologue structure to the coordinates
  # of the amide protons in the mean of the family of clouds
  # - returns a mapping of spin system to residue
  focMapping  = matchCloudsToHomologue(clouds, chainH, structure, threshold)
  spinSystems = focMapping.values()
 
  # Use a sequence bases alignment to determine which residues
  # of the homologue chain are equivalent to thos ein the subject chain
  seqMapping = getChainResidueMapping(chainH, chain) # Done
  
  # Generate an assignment mapping of spin system to residue for the subject
  # by comparing the mapping obtained by matching the clouds to the homologue structure
  # with the mapping generated by a sequence alignment
  assignment = linkCloudsToSequence(focMapping, seqMapping) # Done
  
  # Assign spin systems to residues according to the cloud-via-homologue mapping
  assignSpinSystems(assignment) # Done
  
  # Use CloudThreader to fill in rmaining and mismatched sequence
  # TBD: Will eventually fill in the side chain atoms
  s1, s2, assignment = cloudThreader(chain, clouds, spinSystems, nSteps, graph, 1, progressBar) # Done
  
  # Do a final assignment of spin systems to residues after CLoudsThreader
  assignSpinSystems(assignment) # Done
  
  return s1, s2, assignment
  
def matchCloudsToHomologue(clouds, chainH, structure, threshold=3.0):

  T          = threshold * float(threshold)
  mapping    = {}
  meanCloud  = getMeanCloud(clouds)
  amideCloud = {}
  
  for resonance in meanCloud.keys():
    if resonance.name == 'H':
      amideCloud[resonance] = tuple(meanCloud[resonance])
  
  amideCoords = []
  for residue in chain.residues:
    for atom in residue.atoms:
      if (atom.name == 'H') and atom.atomSet:
        coords = getAtomSetCoords(atomSet, structure)
        amideCoords.append(coords[0], residue)
        break

  for resonance in amideCloud.keys():
    spinSystem = resonance.resonanceSet
    if not spinSystem:
      print 'Warning: Missing spin system for cloud resonance', resonance
      continue
      
    x1,y1,z1 = amideCloud[resonance]

    close    = []
    for (x2,y2,z2), residue in amideCoords:
      dx = x1-x2
      dy = y1-y2
      dz = z1-z2
      d2 = (dx*dx) + (dy*dy) + (dz*dz)
      if d2 <= T:
        close.append( (d2, (x2,y2,z2), residue) )
    
    if not close:
      print 'Warning: No homologue match for resonance', resonance
    
    close.sort()
    d2, bestCoord, residue = close[0]
    
    if mapping.get(residue):
      print 'Warning: Multiple matches for homologue residue', residue
    else:
      mapping[residue] = spinSystem  
  
  return mapping

def linkCloudsToSequence(focMapping, seqMapping):

  assignment = {}
  
  for residueH, residue in seqMapping:
    spinSystem = focMapping.get(residueH)
    if spinSystem:
      assignment[spinSystem] = residue
  
  return assignment

def assignSpinSystems(assignment):

  for spinSystem in assignment.keys():
    residue = assignment[spinSystem]
    assignSpinSystemResidue(spinSystem, residue, warnMerge=1)


