# coding=utf-8
# Copyright 2021 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python2, python3
"""Generates molecules that satisfy two targets.

Target1: SAS
Target2: QED
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import functools
import json
import os
from absl import app
from absl import flags
from rdkit import Chem
from rdkit.Chem import QED
from rdkit.Contrib import SA_Score
from tensorflow.compat.v1 import gfile
from mol_dqn.chemgraph.mcts import deep_q_networks
from mol_dqn.chemgraph.mcts import molecules as molecules_mdp
from mol_dqn.chemgraph.mcts import run_dqn
from mol_dqn.chemgraph.tensorflow import core


flags.DEFINE_float('target_sas', 1, 'The target SAS of the molecule.')
flags.DEFINE_float('target_qed', 0.5, 'The target QED of the molecule.')
flags.DEFINE_float('gamma', 0.999, 'discount')
FLAGS = flags.FLAGS


class MultiObjectiveRewardMolecule(molecules_mdp.Molecule):
  """Defines the subclass of generating a molecule with a specific reward.

  The reward is defined as a 1-D vector with 2 entries: similarity and QED
    reward = (similarity_score, qed_score)
  """

  def _reward(self):
    """Calculates the reward of the current state.

    The reward is defined as a tuple of the similarity and QED value.

    Returns:
      A tuple of the similarity and qed value
    """
    # calculate similarity.
    # if the current molecule does not contain the scaffold of the target,
    # similarity is zero.
    if self._state is None:
      return 0.0, 0.0
    mol = Chem.MolFromSmiles(self._state)
    if mol is None:
      return 0.0, 0.0

    qed_value = QED.qed(mol)
    sas = SA_Score.sascorer.calculateScore(mol)
    return -abs(sas - FLAGS.target_sas), -abs(qed_value - FLAGS.target_qed)


def soft_cst(v, l, r):
  if l <= v <= r:
    return 1
  return -min(abs(l - v), abs(r - v))


class Molecule(molecules_mdp.Molecule):
  """SAS and QED reward molecule."""

  def _reward(self):
    """Calculates the reward of the current state.

    The reward is defined as a tuple of the similarity and QED value.

    Returns:
      A tuple of the similarity and qed value
    """
    # calculate similarity.
    # if the current molecule does not contain the scaffold of the target,
    # similarity is zero.
    if self._state is None:
      return 0.0, 0.0
    mol = Chem.MolFromSmiles(self._state)
    if mol is None:
      return 0.0, 0.0

    qed_value = QED.qed(mol)
    sas = SA_Score.sascorer.calculateScore(mol)

    # c1 = soft_cst(sas, FLAGS.target_sas - 0.2, FLAGS.target_sas + 0.2)
    # c2 = soft_cst(qed_value, FLAGS.target_qed - 0.1, FLAGS.target_qed + 0.1)
    # # if c1 < 0 and c2 < 0:
    # #   return - c1 * c2
    # # else:
    # #   return c1 * c2
    return (soft_cst(sas, FLAGS.target_sas - 0.2, FLAGS.target_sas + 0.2) +
            soft_cst(qed_value, FLAGS.target_qed - 0.1,
                     FLAGS.target_qed + 0.1)) * FLAGS.gamma**(
                         self.max_steps - self._counter)


def main(argv):
  del argv
  if FLAGS.hparams is not None:
    with gfile.Open(FLAGS.hparams, 'r') as f:
      hparams = deep_q_networks.get_hparams(**json.load(f))
  else:
    hparams = deep_q_networks.get_hparams()

  hparams.add_hparam('target_qed', FLAGS.target_qed)
  hparams.add_hparam('target_sas', FLAGS.target_sas)

  environment = Molecule(
      atom_types=set(hparams.atom_types),
      init_mol='CCc1c(C)[nH]c2CCC(CN3CCOCC3)C(=O)c12',
      allow_removal=hparams.allow_removal,
      allow_no_modification=hparams.allow_no_modification,
      allow_bonds_between_rings=False,
      allowed_ring_sizes={3, 4, 5, 6},
      max_steps=hparams.max_steps_per_episode)

  dqn = deep_q_networks.DeepQNetwork(
      input_shape=(hparams.batch_size, hparams.fingerprint_length + 1),
      q_fn=functools.partial(
          deep_q_networks.multi_layer_model, hparams=hparams),
      optimizer=hparams.optimizer,
      grad_clipping=hparams.grad_clipping,
      num_bootstrap_heads=hparams.num_bootstrap_heads,
      gamma=hparams.gamma,
      epsilon=1.0)

  run_dqn.run_training(
      hparams=hparams,
      environment=environment,
      dqn=dqn,
  )

  core.write_hparams(hparams, os.path.join(FLAGS.model_dir, 'config.json'))


if __name__ == '__main__':
  app.run(main)

[((42, 0, 42, 70), 'absl.flags.DEFINE_float', 'flags.DEFINE_float', ({(42, 19, 42, 31): '"""target_sas"""', (42, 33, 42, 34): '(1)', (42, 36, 42, 69): '"""The target SAS of the molecule."""'}, {}), "('target_sas', 1, 'The target SAS of the molecule.')", False, 'from absl import flags\n'), ((43, 0, 43, 72), 'absl.flags.DEFINE_float', 'flags.DEFINE_float', ({(43, 19, 43, 31): '"""target_qed"""', (43, 33, 43, 36): '(0.5)', (43, 38, 43, 71): '"""The target QED of the molecule."""'}, {}), "('target_qed', 0.5, 'The target QED of the molecule.')", False, 'from absl import flags\n'), ((44, 0, 44, 46), 'absl.flags.DEFINE_float', 'flags.DEFINE_float', ({(44, 19, 44, 26): '"""gamma"""', (44, 28, 44, 33): '(0.999)', (44, 35, 44, 45): '"""discount"""'}, {}), "('gamma', 0.999, 'discount')", False, 'from absl import flags\n'), ((148, 2, 152, 3), 'mol_dqn.chemgraph.mcts.run_dqn.run_training', 'run_dqn.run_training', (), '', False, 'from mol_dqn.chemgraph.mcts import run_dqn\n'), ((158, 2, 158, 15), 'absl.app.run', 'app.run', ({(158, 10, 158, 14): 'main'}, {}), '(main)', False, 'from absl import app\n'), ((68, 10, 68, 41), 'rdkit.Chem.MolFromSmiles', 'Chem.MolFromSmiles', ({(68, 29, 68, 40): 'self._state'}, {}), '(self._state)', False, 'from rdkit import Chem\n'), ((72, 16, 72, 28), 'rdkit.Chem.QED.qed', 'QED.qed', ({(72, 24, 72, 27): 'mol'}, {}), '(mol)', False, 'from rdkit.Chem import QED\n'), ((73, 10, 73, 47), 'rdkit.Contrib.SA_Score.sascorer.calculateScore', 'SA_Score.sascorer.calculateScore', ({(73, 43, 73, 46): 'mol'}, {}), '(mol)', False, 'from rdkit.Contrib import SA_Score\n'), ((99, 10, 99, 41), 'rdkit.Chem.MolFromSmiles', 'Chem.MolFromSmiles', ({(99, 29, 99, 40): 'self._state'}, {}), '(self._state)', False, 'from rdkit import Chem\n'), ((103, 16, 103, 28), 'rdkit.Chem.QED.qed', 'QED.qed', ({(103, 24, 103, 27): 'mol'}, {}), '(mol)', False, 'from rdkit.Chem import QED\n'), ((104, 10, 104, 47), 'rdkit.Contrib.SA_Score.sascorer.calculateScore', 'SA_Score.sascorer.calculateScore', ({(104, 43, 104, 46): 'mol'}, {}), '(mol)', False, 'from rdkit.Contrib import SA_Score\n'), ((124, 14, 124, 43), 'mol_dqn.chemgraph.mcts.deep_q_networks.get_hparams', 'deep_q_networks.get_hparams', ({}, {}), '()', False, 'from mol_dqn.chemgraph.mcts import deep_q_networks\n'), ((154, 30, 154, 74), 'os.path.join', 'os.path.join', ({(154, 43, 154, 58): 'FLAGS.model_dir', (154, 60, 154, 73): '"""config.json"""'}, {}), "(FLAGS.model_dir, 'config.json')", False, 'import os\n'), ((121, 9, 121, 39), 'tensorflow.compat.v1.gfile.Open', 'gfile.Open', ({(121, 20, 121, 33): 'FLAGS.hparams', (121, 35, 121, 38): '"""r"""'}, {}), "(FLAGS.hparams, 'r')", False, 'from tensorflow.compat.v1 import gfile\n'), ((140, 11, 141, 61), 'functools.partial', 'functools.partial', (), '', False, 'import functools\n'), ((122, 46, 122, 58), 'json.load', 'json.load', ({(122, 56, 122, 57): 'f'}, {}), '(f)', False, 'import json\n')]