#! /usr/bin/env python

# computeDistributionOfNearestNeighbors.py
# P.Clote




import sys,math, os,copy
from misc import basePairList
from math import exp,sqrt
from vienna import evalViennaEnergy
from misc import RNAfold, RNAeval, RNAsubopt, RNAshapes, DANGLE_STATUS


R         =  0.00198717
T         = 310.15
RT        = (R*T)
THETA     = 3
PRINT     = True
PRINTDist = True #print prob distribution of num nbors
DEBUG     = False
EPSILON   = 1e-12

def Sqrt(x): #Without this can get math error of sqrt(-1e-15)!
  if abs(x)<EPSILON:
    return 0
  return sqrt(x)

def pf(rna):
  #compute ens free energy
  cmd = "echo %s | RNAfold -d%d -p0" % (rna,DANGLE_STATUS)
  file = os.popen(cmd)
  line = file.readline() #sequence
  line = file.readline() #structure
  mfeStr = line.split()[0]
  line = file.readline() #contains ensemble free energy
  ensFreeEnergy = float(line.split()[-2])
  Z             = exp(-ensFreeEnergy/RT)
  return Z


def canBasePair(X,Y):
  if (X,Y) in [('A','U'),('U','A'),('C','G'),('G','C'),('G','U'),('U','G')]:
    return True
  else:
    return False

def generateAllSecStrUsingRNAsubopt(rna,E):
  #generate all sec str having energy in range [MFE,MFE+E]
  cmd  = 'echo %s | %s  -e %f ' % (rna,RNAsubopt,E)
  file = os.popen(cmd)
  line = file.readline()  #skip first line which contains sequence
  SecStrs  = {}
  line = file.readline() 
  while line:
    words = line.split()
    SecStrs[words[0]] = float(words[1]) #ensures overwriting of duplicates
    line = file.readline() 
  Z        = 0.0
  for secStr in SecStrs.keys():
    Z += exp(-SecStrs[secStr]/RT)
  SecStrs = SecStrs.keys()
  PF      = pf(rna)
  return SecStrs,Z,PF

def noCrossingArcs(i,j,L):
  #L is list of base pairs
  for (x,y) in L:
    if i<=x<=j<=y or x<=i<=y<=j:
      return False
  return True

def basePairListToDotBracketList(bplist,rna):
  n      = len(rna)
  secStr = list('.'*n)
  for (i,j) in bplist:
    secStr[i]='('
    secStr[j]=')'
  return secStr

def dotBracketListToString(dblist):
  return "".join(dblist)
  
def computeSumNx(rna,E):
  SecStrs,Z,PF  = generateAllSecStrUsingRNAsubopt(rna,E)
  Nhbors        = []
  Dist          = [] #Dist=[ (numNbors,BoltzmannFactor), ... ]
  sum           = 0.0
  sumSq         = 0.0
  n             = len(rna)
  totalNum      = 0
  sumBoltzmannFactors = 0 #check that Z = sum of boltzmann factors
  for secStr in SecStrs:
    num             = 0.0 #number of nhbors of secStr
    energy          = evalViennaEnergy(rna,secStr)
    boltzmannFactor = exp(-energy/RT) #boltzmann factor of secStr
    prob   = boltzmannFactor/Z #prob of secStr
    bplist = basePairList(secStr)
    num   += len(bplist) #str obtained by omitting bp from secStr
    #Enumerate the nhbors obtained by removing one base pair
    ss = basePairListToDotBracketList(bplist,rna)
    if DEBUG:
      print "Current sec str: %s" % dotBracketListToString(ss)
    for (i,j) in bplist:
      ss0 = copy.deepcopy(ss)
      ss0[i]='.'
      ss0[j]='.'
      ss0dotBracket = dotBracketListToString(ss0)
      Nhbors.append( ss0dotBracket )
      if DEBUG:
        print "neighbor: %s" % dotBracketListToString(ss0)
    #Now enumerate remaining nhbors obtained by adding one base pair
    for i in range(n-THETA-1): #i in [0,...,n-THETA-2]
      for j in range(i+THETA+1,n): #j in [i+THETA+1...n-1]
        if canBasePair(rna[i],rna[j]) and noCrossingArcs(i,j,bplist):
          num += 1
          if DEBUG:
            ss0 = copy.deepcopy(ss)
            ss0[i] = '('
            ss0[j] = ')'
            Nhbors.append("".join(ss0))
            if DEBUG:
              print "neighbor: %s" % dotBracketListToString(ss0)
    totalNum += num
    sumBoltzmannFactors += boltzmannFactor
    val        = num*boltzmannFactor
    Dist.append( (num,prob) )
#    Dist.append( (num,boltzmannFactor) )
    if DEBUG: print  "%d\t%f\t%f\t%f" % (num,prob,boltzmannFactor,val)
    sum      += val
    sumSq += (num**2)*boltzmannFactor
  mean  = sum/Z
  stdev = Sqrt(sumSq/Z-mean**2)
  ratio = Z/PF
  if PRINT: print "Approx partition function: %s" % Z
  if PRINT: print "Real partition function: %s" % PF
  return totalNum,mean,stdev,SecStrs,Nhbors,ratio,Dist,Z

def printDistribution(Dist,Z):
  Dist.sort(); Processed = []
  for num,prob in Dist:
    if num not in Processed:
      if Processed!=[]:
        print "%s\t%s" % (num,sum)
      Processed.append(num)
      sum  = prob
    else:
      sum += prob
      


def mainSingleRNA(rna,E):  
  rna                                 = rna.upper()
  totalNum,mean,stdev,SecStrs,Nhbors,ratio,Dist,Z = computeSumNx(rna,E)
  uniformMean                         = totalNum/float(len(SecStrs))
  print "Exhaustive number of nearest nhbors and str:", totalNum,len(SecStrs)
  print "Exhaustive computation of mean & stdev:", mean,stdev
  print "mean\tstdev\tuniformMean\tratio"
  print "%f\t%f\t%f\t%f" % (mean,stdev,uniformMean,ratio)
  if PRINTDist: printDistribution(Dist,Z)
  if DEBUG:
    print "Nearest neighbor list of size %d" % len(Nhbors)
    print Nhbors
  if DEBUG:
    print "Sec str list of size %d" % len(SecStrs)
    print SecStrs

def mainFASTAfile(filename,E):  
  file = open(filename)
  line = file.readline()
  while line:
    FASTA = line.replace('>','').strip()
    rna   = file.readline().strip().upper()
    totalNum,mean,stdev,SecStrs,Nhbors,ratio,Dist,Z = computeSumNx(rna,E)
    if PRINTDist: printDistribution(Dist,Z)
    numStr = len(SecStrs)
    uniformMean = totalNum/float(len(SecStrs))
    print "mean\tstdev\tuniformMean\tratio\tnumStr"
#    print "%.2f\t%.4f\t%.2f\t%.4f\t%d" % (mean,stdev,uniformMean,ratio,numStr)
    print "%.4f\t%.4f\t%.4f\t\t%.4f\t%d" % (mean,stdev,uniformMean,ratio,numStr)
    line  = file.readline()
  file.close()


if __name__ == '__main__':
  if len(sys.argv)< 3:
    print "Usage: %s -s/-f rna/rnafilename [EnergyBound]" % sys.argv[0]
    sys.exit(1)
  assert( sys.argv[1] in ['-s','-f'] )
  flag = sys.argv[1]
  foo  = sys.argv[2]
  if len(sys.argv)>3:
    E = float(sys.argv[3])
  else:
    E = 2.0
  if flag=='-s':
    mainSingleRNA(foo,E)
  else:
    mainFASTAfile(foo,E)