# $Id$
#
#  Copyright (C) 2003-2008 Greg Landrum and Rational Discovery LLC
#
#   @@ All Rights Reserved @@
#  This file is part of the RDKit.
#  The contents are covered by the terms of the BSD license
#  which is included in the file license.txt, found at the root
#  of the RDKit source tree.
#
"""  command line utility for working with FragmentCatalogs (CASE-type analysis)

**Usage**

  BuildFragmentCatalog [optional args] <filename>

 filename, the name of a delimited text file containing InData, is required
 for some modes of operation (see below)

**Command Line Arguments**

 - -n *maxNumMols*:  specify the maximum number of molecules to be processed

 - -b: build the catalog and OnBitLists
    *requires InData*

 - -s: score compounds
    *requires InData and a Catalog, can use OnBitLists*

 - -g: calculate info gains
    *requires Scores*

 - -d: show details about high-ranking fragments
    *requires a Catalog and Gains*

 - --catalog=*filename*: filename with the pickled catalog.
    If -b is provided, this file will be overwritten.

 - --onbits=*filename*: filename to hold the pickled OnBitLists.
   If -b is provided, this file will be overwritten
  
 - --scores=*filename*: filename to hold the text score data.
   If -s is provided, this file will be overwritten

 - --gains=*filename*: filename to hold the text gains data.
   If -g is provided, this file will be overwritten

 - --details=*filename*: filename to hold the text details data.
   If -d is provided, this file will be overwritten.

 - --minPath=2: specify the minimum length for a path

 - --maxPath=6: specify the maximum length for a path

 - --smiCol=1: specify which column in the input data file contains
     SMILES

 - --actCol=-1: specify which column in the input data file contains
     activities

 - --nActs=2: specify the number of possible activity values

 - --nBits=-1: specify the maximum number of bits to show details for

"""
import sys,os,cPickle
from rdkit import Chem
from rdkit import RDConfig
from rdkit.Chem import FragmentCatalog
from rdkit.Dbase.DbConnection import DbConnect
import numpy
from rdkit.ML import InfoTheory
import types

_cvsVersion="$Revision$"
idx1 = _cvsVersion.find(':')+1
idx2 = _cvsVersion.rfind('$')
__VERSION_STRING="%s"%(_cvsVersion[idx1:idx2])

def message(msg,dest=sys.stdout):
  dest.write(msg)

def BuildCatalog(suppl,maxPts=-1,groupFileName=None,
                 minPath=2,maxPath=6,reportFreq=10):
  """ builds a fragment catalog from a set of molecules in a delimited text block

    **Arguments**

      - suppl: a mol supplier

      - maxPts: (optional) if provided, this will set an upper bound on the
        number of points to be considered

      - groupFileName: (optional) name of the file containing functional group
        information

      - minPath, maxPath: (optional) names of the minimum and maximum path lengths
        to be considered

      - reportFreq: (optional) how often to display status information  

    **Returns**

      a FragmentCatalog
      
  """
  if groupFileName is None:
    groupFileName = os.path.join(RDConfig.RDDataDir,"FunctionalGroups.txt")

  fpParams = FragmentCatalog.FragCatParams(minPath,maxPath,groupFileName)
  catalog = FragmentCatalog.FragCatalog(fpParams)
  fgen = FragmentCatalog.FragCatGenerator()
  if maxPts >0:
    nPts = maxPts
  else:
    if hasattr(suppl,'__len__'):
      nPts = len(suppl)
    else:
      nPts = -1
  for i,mol in enumerate(suppl):
    if i == nPts:
      break
    if i and not i%reportFreq:
      if nPts>-1:
        message('Done %d of %d, %d paths\n'%(i,nPts,catalog.GetFPLength()))
      else:
        message('Done %d, %d paths\n'%(i,catalog.GetFPLength()))
    fgen.AddFragsFromMol(mol,catalog)
  return catalog

def ScoreMolecules(suppl,catalog,maxPts=-1,actName='',acts=None,
                   nActs=2,reportFreq=10):
  """ scores the compounds in a supplier using a catalog

    **Arguments**

      - suppl: a mol supplier

      - catalog: the FragmentCatalog

      - maxPts: (optional) the maximum number of molecules to be
        considered

      - actName: (optional) the name of the molecule's activity property.
        If this is not provided, the molecule's last property will be used.

      - acts: (optional) a sequence of activity values (integers).
        If not provided, the activities will be read from the molecules.

      - nActs: (optional) number of possible activity values

      - reportFreq: (optional) how often to display status information  

    **Returns**

      a 2-tuple:

        1) the results table (a 3D array of ints nBits x 2 x nActs)

        2) a list containing the on bit lists for each molecule

  """
  nBits = catalog.GetFPLength()
  resTbl = numpy.zeros((nBits,2,nActs),numpy.int)
  obls = []

  if not actName and not acts:
    actName = suppl[0].GetPropNames()[-1]

  
  fpgen = FragmentCatalog.FragFPGenerator()
  suppl.reset()
  i = 1
  for mol in suppl:
    if i and not i%reportFreq:
      message('Done %d.\n'%(i))
    if mol:
      if not acts:
        act = int(mol.GetProp(actName))
      else:
        act = acts[i-1]
      fp = fpgen.GetFPForMol(mol,catalog)
      obls.append([x for x in fp.GetOnBits()])
      for j in range(nBits):
        resTbl[j,0,act] += 1
      for id in obls[i-1]:
        resTbl[id-1,0,act] -= 1
        resTbl[id-1,1,act] += 1
    else:
      obls.append([])
    i+=1  
  return resTbl,obls

def ScoreFromLists(bitLists,suppl,catalog,maxPts=-1,actName='',acts=None,
                   nActs=2,reportFreq=10):
  """  similar to _ScoreMolecules()_, but uses pre-calculated bit lists
    for the molecules (this speeds things up a lot)

    
    **Arguments**

      - bitLists: sequence of on bit sequences for the input molecules

      - suppl: the input supplier (we read activities from here)

      - catalog: the FragmentCatalog

      - maxPts: (optional) the maximum number of molecules to be
        considered

      - actName: (optional) the name of the molecule's activity property.
        If this is not provided, the molecule's last property will be used.

      - nActs: (optional) number of possible activity values

      - reportFreq: (optional) how often to display status information  

    **Returns**

       the results table (a 3D array of ints nBits x 2 x nActs)

  """
  nBits = catalog.GetFPLength()
  if maxPts >0:
    nPts = maxPts
  else:
    nPts = len(bitLists)
  resTbl = numpy.zeros((nBits,2,nActs),numpy.int)
  if not actName and not acts:
    actName = suppl[0].GetPropNames()[-1]
  suppl.reset()
  for i in range(1,nPts+1):
    mol = suppl.next()
    if not acts:
      act = int(mol.GetProp(actName))
    else:
      act = acts[i-1]
    if i and not i%reportFreq:
      message('Done %d of %d\n'%(i,nPts))
    ids = set()
    for id in bitLists[i-1]:
      ids.add(id-1)
    for j in range(nBits):
      resTbl[j,0,act] += 1
    for id in ids:
      resTbl[id,0,act] -= 1
      resTbl[id,1,act] += 1
  return resTbl

def CalcGains(suppl,catalog,topN=-1,actName='',acts=None,
              nActs=2,reportFreq=10,biasList=None,collectFps=0):
  """ calculates info gains by constructing fingerprints
    *DOC*

    Returns a 2-tuple:
       1) gains matrix
       2) list of fingerprints
    
  """
  nBits = catalog.GetFPLength()
  if topN < 0:
    topN = nBits
  if not actName and not acts:
    actName = suppl[0].GetPropNames()[-1]

  gains = [0]*nBits
  if hasattr(suppl,'__len__'):
    nMols = len(suppl)
  else:
    nMols = -1
  fpgen = FragmentCatalog.FragFPGenerator()
  #ranker = InfoTheory.InfoBitRanker(nBits,nActs,InfoTheory.InfoType.ENTROPY)
  if biasList:
    ranker = InfoTheory.InfoBitRanker(nBits,nActs,InfoTheory.InfoType.BIASENTROPY)
    ranker.SetBiasList(biasList)
  else:
    ranker = InfoTheory.InfoBitRanker(nBits,nActs,InfoTheory.InfoType.ENTROPY)
  i = 0
  fps = []
  for mol in suppl:
    if not acts:
      try:
        act = int(mol.GetProp(actName))
      except KeyError:
        message('ERROR: Molecule has no property: %s\n'%(actName))
        message('\tAvailable properties are: %s\n'%(str(mol.GetPropNames())))
        raise KeyError,actName
    else:
      act = acts[i]
    if i and not i%reportFreq:
      if nMols>0:
        message('Done %d of %d.\n'%(i,nMols))
      else:
        message('Done %d.\n'%(i))
    fp = fpgen.GetFPForMol(mol,catalog)
    ranker.AccumulateVotes(fp,act)
    i+=1;
    if collectFps:
      fps.append(fp)
  gains = ranker.GetTopN(topN)
  return gains,fps

def CalcGainsFromFps(suppl,fps,topN=-1,actName='',acts=None,
              nActs=2,reportFreq=10,biasList=None):
  """ calculates info gains from a set of fingerprints

    *DOC*
    
  """
  nBits = len(fps[0])
  if topN < 0:
    topN = nBits
  if not actName and not acts:
    actName = suppl[0].GetPropNames()[-1]

  gains = [0]*nBits
  if hasattr(suppl,'__len__'):
    nMols = len(suppl)
  else:
    nMols = -1
  if biasList:
    ranker = InfoTheory.InfoBitRanker(nBits,nActs,InfoTheory.InfoType.BIASENTROPY)
    ranker.SetBiasList(biasList)
  else:
    ranker = InfoTheory.InfoBitRanker(nBits,nActs,InfoTheory.InfoType.ENTROPY)
  for i,mol in enumerate(suppl):
    if not acts:
      try:
        act = int(mol.GetProp(actName))
      except KeyError:
        message('ERROR: Molecule has no property: %s\n'%(actName))
        message('\tAvailable properties are: %s\n'%(str(mol.GetPropNames())))
        raise KeyError,actName
    else:
      act = acts[i]
    if i and not i%reportFreq:
      if nMols>0:
        message('Done %d of %d.\n'%(i,nMols))
      else:
        message('Done %d.\n'%(i))
    fp = fps[i]
    ranker.AccumulateVotes(fp,act)
  gains = ranker.GetTopN(topN)
  return gains

def OutputGainsData(outF,gains,cat,nActs=2):
  actHeaders = ['Act-%d'%(x) for x in range(nActs)]
  if cat:
    outF.write('id,Description,Gain,%s\n'%(','.join(actHeaders)))
  else:
    outF.write('id,Gain,%s\n'%(','.join(actHeaders)))
  for entry in gains:
    id = int(entry[0])
    outL = [str(id)]
    if cat:
      descr = cat.GetBitDescription(id)
      outL.append(descr)
    outL.append('%.6f'%entry[1])  
    outL += ['%d'%x for x in entry[2:]]
    outF.write(','.join(outL))
    outF.write('\n')
      

def ProcessGainsData(inF,delim=',',idCol=0,gainCol=1):
  """ reads a list of ids and info gains out of an input file

  """
  res = []
  inL = inF.readline()
  for line in inF.xreadlines():
    splitL = line.strip().split(delim)
    res.append((splitL[idCol],float(splitL[gainCol])))
  return res  
    

def ShowDetails(catalog,gains,nToDo=-1,outF=sys.stdout,idCol=0,gainCol=1,
                outDelim=','):
  """
   gains should be a sequence of sequences.  The idCol entry of each
   sub-sequence should be a catalog ID.  _ProcessGainsData()_ provides
   suitable input.

  """
  if nToDo < 0:
    nToDo = len(gains)
  for i in range(nToDo):
    id = int(gains[i][idCol])
    gain = float(gains[i][gainCol])
    descr = catalog.GetFragDescription(id)
    if descr:
      outF.write('%s\n'%(outDelim.join((str(id),descr,str(gain)))))

def SupplierFromDetails(details):
  from rdkit.VLib.NodeLib.DbMolSupply import DbMolSupplyNode
  from rdkit.VLib.NodeLib.SmilesSupply import SmilesSupplyNode

  if details.dbName:
    conn = DbConnect(details.dbName,details.tableName)
    suppl = DbMolSupplyNode(conn.GetData())
  else:
    suppl = SmilesSupplyNode(details.inFileName,delim=details.delim,
                             nameColumn=details.nameCol,
                             smilesColumn=details.smiCol,
                             titleLine=details.hasTitle)
    if type(details.actCol)==types.IntType:
      suppl.reset()
      m = suppl.next()
      actName = m.GetPropNames()[details.actCol]
      details.actCol = actName
    if type(details.nameCol)==types.IntType:
      suppl.reset()
      m = suppl.next()
      nameName = m.GetPropNames()[details.nameCol]
      details.nameCol = nameName
      suppl.reset()
  if type(details.actCol)==types.IntType:
    suppl.reset()
    m = suppl.next()
    actName = m.GetPropNames()[details.actCol]
    details.actCol = actName
  if type(details.nameCol)==types.IntType:
    suppl.reset()
    m = suppl.next()
    nameName = m.GetPropNames()[details.nameCol]
    details.nameCol = nameName
    suppl.reset()
  return suppl


def Usage():
  print "This is BuildFragmentCatalog version %s"%(__VERSION_STRING)
  print 'usage error'
  #print __doc__
  sys.exit(-1)

class RunDetails(object):
  numMols=-1
  doBuild=0
  doSigs=0
  doScore=0
  doGains=0
  doDetails=0
  catalogName=None
  onBitsName=None
  scoresName=None
  gainsName=None
  dbName=''
  tableName=None
  detailsName=None
  inFileName=None
  fpName=None
  minPath=2
  maxPath=6
  smiCol=1
  actCol=-1
  nameCol=-1
  hasTitle=1
  nActs = 2
  nBits=-1
  delim=','
  biasList=None
  topN=-1
  
def ParseArgs(details):
  import getopt
  try:
    args,extras = getopt.getopt(sys.argv[1:],'n:d:cst',
                                ['catalog=','onbits=',
                                 'scoresFile=','gainsFile=','detailsFile=','fpFile=',
                                 'minPath=','maxPath=','smiCol=','actCol=','nameCol=','nActs=',
                                 'nBits=','biasList=','topN=',
                                 'build','sigs','gains','details','score','noTitle'])
  except:
    sys.stderr.write('Error parsing command line:\n')
    import traceback
    traceback.print_exc()
    Usage()
  for arg,val in args:
    if arg=='-n':
      details.numMols=int(val)
    elif arg=='-c':
      details.delim=','
    elif arg=='-s':
      details.delim=' '
    elif arg=='-t':
      details.delim='\t'
    elif arg=='-d':
      details.dbName=val
    elif arg=='--build':
      details.doBuild=1
    elif arg=='--score':
      details.doScore=1
    elif arg=='--gains':
      details.doGains=1
    elif arg=='--sigs':
      details.doSigs=1
    elif arg=='-details':
      details.doDetails=1
    elif arg=='--catalog':
      details.catalogName=val
    elif arg=='--onbits':
      details.onBitsName=val
    elif arg=='--scoresFile':
      details.scoresName=val
    elif arg=='--gainsFile':
      details.gainsName=val
    elif arg=='--detailsFile':
      details.detailsName=val
    elif arg=='--fpFile':
      details.fpName=val
    elif arg=='--minPath':
      details.minPath=int(val)
    elif arg=='--maxPath':
      details.maxPath=int(val)
    elif arg=='--smiCol':
      try:
        details.smiCol=int(val)
      except ValueError:
        details.smiCol=val
    elif arg=='--actCol':
      try:
        details.actCol=int(val)
      except ValueError:
        details.actCol=val
    elif arg=='--nameCol':
      try:
        details.nameCol=int(val)
      except ValueError:
        details.nameCol=val
    elif arg=='--nActs':
      details.nActs=int(val)
    elif arg=='--nBits':
      details.nBits=int(val)
    elif arg=='--noTitle':
      details.hasTitle=0
    elif arg=='--biasList':
      details.biasList=tuple(eval(val))
    elif arg=='--topN':
      details.topN=int(val)
    elif arg=='-h':
      Usage()
      sys.exit(0)
    else:
      Usage()
  if len(extras):
    if details.dbName:
      details.tableName=extras[0]
    else:
      details.inFileName = extras[0]
  else:
    Usage()
if __name__=='__main__':
  import time
  details = RunDetails()
  ParseArgs(details)
  from cStringIO import StringIO
  suppl = SupplierFromDetails(details)

  cat = None
  obls = None
  if details.doBuild:
    if not suppl:
      message("We require inData to generate a catalog\n")
      sys.exit(-2)
    message("Building catalog\n")
    t1 = time.time()
    cat = BuildCatalog(suppl,maxPts=details.numMols,
                       minPath=details.minPath,maxPath=details.maxPath)
    t2 = time.time()
    message("\tThat took %.2f seconds.\n"%(t2-t1))
    if details.catalogName:
      message("Dumping catalog data\n")
      cPickle.dump(cat,open(details.catalogName,'wb+'))
  elif details.catalogName:
    message("Loading catalog\n")
    cat = cPickle.load(open(details.catalogName,'rb'))
    if details.onBitsName:
      try:
        obls = cPickle.load(open(details.onBitsName,'rb'))
      except:
        obls = None
      else:
        if len(obls)<(inD.count('\n')-1):
          obls = None
  scores = None
  if details.doScore:
    if not suppl:
      message("We require inData to score molecules\n")
      sys.exit(-2)
    if not cat:
      message("We require a catalog to score molecules\n")
      sys.exit(-2)
    message("Scoring compounds\n")
    if not obls or len(obls)<details.numMols:
      scores,obls = ScoreMolecules(suppl,cat,maxPts=details.numMols,
                                actName=details.actCol,
                                nActs=details.nActs)
      if details.scoresName:
        cPickle.dump(scores,open(details.scoresName,'wb+'))
      if details.onBitsName:
        cPickle.dump(obls,open(details.onBitsName,'wb+'))
    else:
      scores = ScoreFromLists(obls,suppl,cat,maxPts=details.numMols,
                              actName=details.actCol,
                              nActs=details.nActs)
  elif details.scoresName:
    scores = cPickle.load(open(details.scoresName,'rb'))
    
  if details.fpName and os.path.exists(details.fpName) and not details.doSigs:
    message("Reading fingerprints from file.\n")
    fps = cPickle.load(open(details.fpName,'rb'))
  else:
    fps = []
  gains = None
  if details.doGains:
    if not suppl:
      message("We require inData to calculate gains\n")
      sys.exit(-2)
    if not (cat or fps):
      message("We require either a catalog or fingerprints to calculate gains\n")
      sys.exit(-2)
    message("Calculating Gains\n")
    t1 = time.time()
    if details.fpName:
      collectFps=1
    else:
      collectFps=0
    if not fps:
      gains,fps = CalcGains(suppl,cat,topN=details.topN,actName=details.actCol,
                            nActs=details.nActs,biasList=details.biasList,
                            collectFps=collectFps)
      if details.fpName:
        message("Writing fingerprint file.\n")
        tmpF=open(details.fpName,'wb+')
        cPickle.dump(fps,tmpF,1)
        tmpF.close()
    else:
      gains = CalcGainsFromFps(suppl,fps,topN=details.topN,actName=details.actCol,
                               nActs=details.nActs,biasList=details.biasList)
    t2=time.time()
    message("\tThat took %.2f seconds.\n"%(t2-t1))
    if details.gainsName:
      outF = open(details.gainsName,'w+')
      OutputGainsData(outF,gains,cat,nActs=details.nActs)
  else:
    if details.gainsName:
      inF = open(details.gainsName,'r')
      gains = ProcessGainsData(inF)

  
  if details.doDetails:
    if not cat:
      message("We require a catalog to get details\n")
      sys.exit(-2)
    if not gains:
      message("We require gains data to get details\n")
      sys.exit(-2)
    io = StringIO()
    io.write('id,SMILES,gain\n')
    ShowDetails(cat,gains,nToDo=details.nBits,outF=io)
    if details.detailsName:
      open(details.detailsName,'w+').write(io.getvalue())
    else:
      sys.stderr.write(io.getvalue())