# Copyright (c) 2013, GlaxoSmithKline Research & Development Ltd.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#
#     * Redistributions of source code must retain the above copyright
#       notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above
#       copyright notice, this list of conditions and the following
#       disclaimer in the documentation and/or other materials provided
#       with the distribution.
#     * Neither the name of GlaxoSmithKline Research & Development Ltd.
#       nor the names of its contributors may be used to endorse or promote
#       products derived from this software without specific prior written
#       permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# Created by Jameed Hussain, July 2013

import sys
import re
import os
import subprocess
import sqlite3
from optparse import OptionParser
from rdkit import Chem


def heavy_atom_count(smi):
  m = Chem.MolFromSmiles(smi)
  return m.GetNumAtoms()


def add_to_db(id, core, context, context_size):

  cursor.execute("INSERT OR IGNORE INTO context_table(context_smi) VALUES(?)", (context, ))

  the_id_of_the_row = None
  cursor.execute("SELECT context_id FROM context_table WHERE context_smi = ?", (context, ))
  the_id_of_the_row = cursor.fetchone()[0]

  core_ni = re.sub(r'\[\*\:1\]', '[*]', core)
  core_ni = re.sub(r'\[\*\:2\]', '[*]', core_ni)
  core_ni = re.sub(r'\[\*\:3\]', '[*]', core_ni)

  cursor.execute(
    "INSERT INTO core_table(context_id, cmpd_id, core_smi, core_smi_ni) VALUES(?,?,?,?);",
    (the_id_of_the_row, id, core, core_ni))
  #print "the_id_of_the_row:%s, id:%s, core:%s" % (the_id_of_the_row,id,core)

  #check if heavy atom count for context_id has been calculated
  cursor.execute("select context_size from context_table where context_id = ?",
                 (the_id_of_the_row, ))
  heavy_calculated = cursor.fetchone()
  #add heavy atom count if needed
  if (heavy_calculated[0] is None):
    cursor.execute("update context_table set context_size = ? where context_id = ? ",
                   (context_size, the_id_of_the_row))


def index_hydrogen_change():
  #Algorithm details
  #1) Loop through context smiles in the db
  #2) If context smiles is the result of a single cut (so contains only 1 *) replace the * with H, and cansmi
  #3) If full smiles matches this new context smi, add *-H to that core_smi index.

  cursor.execute("SELECT context_smi FROM context_table")
  res = cursor.fetchall()

  for row in res:
    key = row[0]
    attachments = key.count('*')
    #print attachments

    if (attachments == 1):
      smi = str(key)
      #simple method
      smi = re.sub(r'\[\*\:1\]', '[H]', smi)

      #now cansmi it
      temp = Chem.MolFromSmiles(smi)

      if (temp == None):
        sys.stderr.write('Error with key: %s, Added H: %s\n' % (key, smi))
      else:
        c_smi = Chem.MolToSmiles(temp, isomericSmiles=True)

        cursor.execute("SELECT cmpd_id FROM cmpd_smisp WHERE smiles = ?", (c_smi, ))
        cmpd_id = cursor.fetchone()
        if (cmpd_id):
          core = "[*:1][H]"
          cmpd_id = cmpd_id[0]
          key_size = temp.GetNumAtoms() - 1
          add_to_db(cmpd_id, core, key, key_size)
          #print "Added: id:%s, core:%s, context:%s" % (str(cmpd_id),core,key)

        #set up the command line options


parser = OptionParser(description="Program to create an MMP db. ")
parser.add_option(
  '-p', '--prefix', action='store', dest='prefix', type='string',
  help='Prefix to use for the db file (and directory for SMARTS index). DEFAULT=mmp')
parser.add_option(
  '-m', '--maxsize', action='store', dest='maxsize', type='int',
  help="Maximum size of change (in heavy atoms) that is stored in the database. DEFAULT=15. \t\
                  Note: Any MMPs that involve a change greater than this value will not be stored in the database and hence not be identified in the searching.")
parser.add_option(
  '-s', '--smarts', default=False, action='store_true', dest='sma',
  help='Build SMARTS db so can perform SMARTS searching against db. Note: Will make the build process somewhat slower.')
#parse the command line options
(options, args) = parser.parse_args()

#this is the heavy atom limit
#of the size of change that goes in the db
max_size = 15
db_name = "mmp.db"
pre = "mmp"

#print options
if (options.maxsize != None):
  max_size = options.maxsize

if (options.prefix != None):
  pre = options.prefix
  db_name = "%s.db" % (pre)

con = sqlite3.connect(db_name)
cursor = con.cursor()

#these setting increase performance
cursor.execute('PRAGMA main.page_size = 4096;')
cursor.execute('PRAGMA main.cache_size=10000;')
cursor.execute('PRAGMA main.locking_mode=EXCLUSIVE;')
cursor.execute('PRAGMA main.synchronous=NORMAL;')
cursor.execute('PRAGMA main.journal_mode=WAL;')
cursor.execute('PRAGMA main.cache_size=5000;')
cursor.execute('PRAGMA main.temp_store = MEMORY;')

#tables
cursor.execute("DROP TABLE IF EXISTS context_table")
cursor.execute("""CREATE TABLE context_table
                  (
                    context_id INTEGER PRIMARY KEY AUTOINCREMENT,
                    context_smi VARCHAR(1000) NOT NULL UNIQUE,
                    context_size INTEGER
                  )""")

cursor.execute("DROP TABLE IF EXISTS core_table")
cursor.execute("""CREATE TABLE core_table
                  (
                    context_id INTEGER NOT NULL,
                    cmpd_id VARCHAR(20) NOT NULL,
                    core_smi VARCHAR(1000) NOT NULL,
                    core_smi_ni VARCHAR(1000) NOT NULL
                  )""")

cursor.execute("DROP TABLE IF EXISTS cmpd_smisp")
cursor.execute("""CREATE TABLE cmpd_smisp
                  (
                    cmpd_id VARCHAR(20) NOT NULL UNIQUE,
                    smiles VARCHAR(1000),
                    cmpd_size INTEGER
                  )""")

for line in sys.stdin:

  line = line.rstrip()
  smi, id, core, context = line.split(',')

  cursor.execute("INSERT OR IGNORE INTO cmpd_smisp(smiles,cmpd_id) VALUES(?,?)", (smi, id))

  #check if heavy atom count for compound has been calculated
  cursor.execute("select cmpd_size from cmpd_smisp where cmpd_id = ?", (id, ))
  heavy_calculated = cursor.fetchone()
  #add heavy atom count if needed
  if (heavy_calculated[0] is None):
    cmpd_heavy_atoms = heavy_atom_count(smi)
    cursor.execute("update cmpd_smisp set cmpd_size = ? where cmpd_id = ? ", (cmpd_heavy_atoms, id))
  else:
    cmpd_heavy_atoms = heavy_calculated[0]

  #deal with cmpds that have not been core_smied
  if (len(core) == 0) and (len(context) == 0):
    continue

  #deal with single cuts
  if (len(core) == 0):
    side_chains = context.split('.')

    context = side_chains[0]
    core = side_chains[1]
    context_size = heavy_atom_count(context) - context.count("*")

    #limit what goes in db based on max_size
    if (cmpd_heavy_atoms - context_size <= max_size):
      #put in db
      add_to_db(id, core, context, context_size)

    context = side_chains[1]
    core = side_chains[0]
    context_size = heavy_atom_count(context) - context.count("*")

    #limit what goes in db based on max_size
    if (cmpd_heavy_atoms - context_size <= max_size):
      #put in db
      add_to_db(id, core, context, context_size)

  #double or triple cut
  else:
    context_size = heavy_atom_count(context) - context.count("*")
    if (cmpd_heavy_atoms - context_size <= max_size):
      #put in db
      add_to_db(id, core, context, context_size)

#index the H change
index_hydrogen_change()

#Note: Indices automatically built for primary key and unique columns
#so need to build the remaining indices that are needed
cursor.execute("DROP INDEX IF EXISTS context_smi_idx")
cursor.execute("CREATE INDEX context_smi_idx ON context_table (context_smi)")
cursor.execute("DROP INDEX IF EXISTS context_id_idx")
cursor.execute("CREATE INDEX context_id_idx ON core_table (context_id)")
cursor.execute("DROP INDEX IF EXISTS cmpd_id_idx")
cursor.execute("CREATE INDEX cmpd_id_idx ON core_table (cmpd_id)")
cursor.execute("DROP INDEX IF EXISTS core_smi_idx")
cursor.execute("CREATE INDEX core_smi_ni_idx ON core_table (core_smi_ni)")
cursor.execute("DROP INDEX IF EXISTS smiles_idx")
cursor.execute("CREATE INDEX smiles_idx ON cmpd_smisp (smiles)")

#add smarts searching capability if required
if (options.sma):
  #first dump out the distinct core_smi_ni from core_table
  #and write to a file
  temp_core_ni_file = 'temp_core_ni_file_%s' % (os.getpid())
  outfile = open(temp_core_ni_file, 'w')

  query_sql = """select distinct(core_smi_ni) from core_table"""

  cursor.execute(query_sql)
  results = cursor.fetchall()

  for r in results:
    #print "%s %s" % (r[0],r[0])
    outfile.write("%s %s\n" % (r[0], r[0]))

  outfile.close()

  #set os enviroment for rdkit to use sqllite
  os.environ['RD_USESQLLITE'] = '1'
  #use the DbCli utility in RDKit: http://code.google.com/p/rdkit/wiki/UsingTheDbCLI
  cmd = 'python $RDBASE/Projects/DbCLI/CreateDb.py --dbDir=%s_smarts --molFormat=smiles %s --noPairs --noFingerprints --noDescriptors --noProps --noMorganFps --noSmiles --silent' % (
    pre, temp_core_ni_file)
  subprocess.Popen(cmd, shell=True).wait()
  #remove temporary file
  os.unlink(temp_core_ni_file)