#!/usr/bin/env /opt/exp_soft/enmr/CIRMMP/xplor-nih/2.21/64/bin/pyXplor


import sys
import psfGen
import protocol
from math import sqrt
from atomSelAction import RMSD, Fit
nFit=0
usage = """usage: %s [options] <target structure file> <structure files to compare>

compute rmsd to a reference structure

option are zero or more of
  -selection <XPLOR selection> - specify atoms used in comparison.
  -fitSelection <selection>    - specify atoms used to best-fit structures.
  -psf          <psf file>     - psf file. If omitted, it is auto-generated.
  -diffSeq                     - target and comparison structures have
                                 different sequences - generate them
                                 separately. All comparison structure files
                                 must have the same sequence, though.
  -noFit                       - omit the fitting procedure. If not
                                 specified, the structures are best-fitted
                                 before pairwise RMSD calculation.

Structures are first best-fit, and then the rmsd is calculated. This is done
for each pair of structures. The average, deviation, min and max rmsd are
printed.

""" % sys.argv[0]


(opts,args) = xplor.parseArguments(('selection:1','fitSelection:1',
                                        'psf:1',
                                        'diffSeq:0',
                                        'help-script','noFit:0'))

sel="not hydro and not resname ANI"
fitSel=None
psf=None
diffSeq=False

for opt in opts:
    if opt[0]=='help-script':
        print
        print usage
        sys.exit(0)
    if opt[0]=='selection':
        sel = opt[1]
    elif opt[0]=='fitSelection':
        fitSel = opt[1]
    elif opt[0]=='psf':
        psf = opt[1]
        pass
    elif opt[0]=='diffSeq':
        diffSeq=True
    elif opt[0]=='noFit':
        nFit=1    
        pass
    pass


if len(args)<2:
    raise Exception("at least one target and one comparison file must be specified")

target=args[0]
pdbFiles=args[1:]

if not fitSel:
    fitSel = sel
    pass

if psf:
    protocol.initStruct(psf)
    protocol.initCoords(target)
else:
    protocol.loadPDB(target)
    
#unknown=AtomSel("not known")
#if len(unknown)>0:
#    print "deleting %d atoms with unknown coordinates" % len(unknown)
#    pass
xplor.command("delete sele=(not known) end")

sim=xplor.simulation

#
# to do:
#  allow structures to have different PSFs
# place each in a sub- XplorSimulation
# RMSD and Fit will need to work on specified sets of coords only
# - should make sure that selections are of the same size.
#


cnt=0
sum=0.
variance=0.
mini=1e20
maxi=0.
comp = sim.atomPosArr()
compSel=AtomSel(sel)
first=True
from atomSel import intersection
for i in range(0,len(pdbFiles)):

    if diffSeq:
        if first:
            # create new XplorSimulation for first comparison structure
            # assume all other comparison structures have same sequence
            # and have same atom coordinates defined.
            import xplorSimulation, simulation
            xSim = xplorSimulation.XplorSimulation()
            simulation.makeCurrent( xSim )
        
            protocol.loadPDB(pdbFiles[i])
            first=False
        else:
            protocol.initCoords(pdbFiles[i])
            pass

        measure = RMSD(AtomSel(sel))
        if nFit==0:
            AtomSel('all').apply( Fit(compSel,fitSel) )
    else:
        protocol.initCoords(pdbFiles[i],erase=True,
                            strictResNames=True,
                            maxUnreadEntries=(400, 20.))

        unknown=intersection(AtomSel("not known"),fitSel)
        if unknown:
            print "these fit atoms have undefined coordinates:"
            for atom in unknown:
                print "\t%s" % atom.string()
                pass
            print "try using the -diffSeq flag"
            raise Exception("unknown fit atoms in %s:" % pdbFiles[i])
        
            

        measure = RMSD(comp)
        if nFit==0:
            AtomSel('all').apply( Fit(comp,fitSel) )
        pass

    compSel.apply( measure )


    sum += measure.rmsd()
    variance += measure.rmsd()**2
    mini=min(mini,measure.rmsd())
    maxi=max(maxi,measure.rmsd())
    pass

cnt=len(pdbFiles)

ave = sum / cnt
sdev = sqrt( variance/cnt  - ave**2 )

print 'rmsd:',
print 'ave: %.4f   std dev: %.4f   min: %.4f   max: %.4f  # atoms: %d' % \
      (ave,sdev,mini,maxi,len(compSel))