"""
Featurizer implementations used in Smiles2Vec models.
SmilesToSeq featurizer for Smiles2Vec models taken from https://arxiv.org/abs/1712.02734
"""
from typing import Dict, List
import numpy as np
import pandas as pd

from deepchem.utils.typing import RDKitMol
from deepchem.feat.base_classes import MolecularFeaturizer

PAD_TOKEN = "<pad>"
OUT_OF_VOCAB_TOKEN = "<unk>"


def create_char_to_idx(filename: str,
                       max_len: int = 250,
                       smiles_field: str = "smiles") -> Dict[str, int]:
  """Creates a dictionary with character to index mapping.

  Parameters
  ----------
  filename: str
      Name of the file containing the SMILES strings
  max_len: int, default 250
      Maximum allowed length of the SMILES string
  smiles_field: str, default "smiles"
      Field indicating the SMILES strings int the file.

  Returns
  -------
  Dict[str, int]
    A dictionary mapping characters to their integer indexes.
  """
  smiles_df = pd.read_csv(filename)
  char_set = set()
  for smile in smiles_df[smiles_field]:
    if len(smile) <= max_len:
      char_set.update(set(smile))

  unique_char_list = list(char_set)
  unique_char_list += [PAD_TOKEN, OUT_OF_VOCAB_TOKEN]
  char_to_idx = {letter: idx for idx, letter in enumerate(unique_char_list)}
  return char_to_idx


class SmilesToSeq(MolecularFeaturizer):
  """
  SmilesToSeq Featurizer takes a SMILES string, and turns it into a sequence.
  Details taken from [1]_.

  SMILES strings smaller than a specified max length (max_len) are padded using
  the PAD token while those larger than the max length are not considered. Based
  on the paper, there is also the option to add extra padding (pad_len) on both
  sides of the string after length normalization. Using a character to index (char_to_idx)
  mapping, the SMILES characters are turned into indices and the
  resulting sequence of indices serves as the input for an embedding layer.

  References
  ----------
  .. [1] Goh, Garrett B., et al. "Using rule-based labels for weak supervised
     learning: a ChemNet for transferable chemical property prediction."
     Proceedings of the 24th ACM SIGKDD International Conference on Knowledge
     Discovery & Data Mining. 2018.

  Note
  ----
  This class requires RDKit to be installed.
  """

  def __init__(self,
               char_to_idx: Dict[str, int],
               max_len: int = 250,
               pad_len: int = 10):
    """Initialize this class.

    Parameters
    ----------
    char_to_idx: Dict
      Dictionary containing character to index mappings for unique characters
    max_len: int, default 250
      Maximum allowed length of the SMILES string.
    pad_len: int, default 10
      Amount of padding to add on either side of the SMILES seq
    """
    self.max_len = max_len
    self.char_to_idx = char_to_idx
    self.idx_to_char = {idx: letter for letter, idx in self.char_to_idx.items()}
    self.pad_len = pad_len

  def to_seq(self, smile: List[str]) -> np.ndarray:
    """Turns list of smiles characters into array of indices"""
    out_of_vocab_idx = self.char_to_idx[OUT_OF_VOCAB_TOKEN]
    seq = [
        self.char_to_idx.get(character, out_of_vocab_idx) for character in smile
    ]
    return np.array(seq)

  def remove_pad(self, characters: List[str]) -> List[str]:
    """Removes PAD_TOKEN from the character list."""
    characters = characters[self.pad_len:len(characters) - self.pad_len]
    chars = list()

    for char in characters:
      if char != PAD_TOKEN:
        chars.append(char)
    return chars

  def smiles_from_seq(self, seq: List[int]) -> str:
    """Reconstructs SMILES string from sequence."""
    characters = [self.idx_to_char[i] for i in seq]

    characters = self.remove_pad(characters)
    smile = "".join([letter for letter in characters])
    return smile

  def _featurize(self, datapoint: RDKitMol, **kwargs) -> np.ndarray:
    """Featurizes a SMILES sequence.

    Parameters
    ----------
    datapoints: rdkit.Chem.rdchem.Mol
      RDKit Mol object

    Returns
    -------
    np.ndarray
      A 1D array of a SMILES sequence.
      If the length of SMILES is longer than `max_len`, this value is an empty array.
    """
    try:
      from rdkit import Chem
    except ModuleNotFoundError:
      raise ImportError("This class requires RDKit to be installed.")

    if 'mol' in kwargs:
      datapoint = kwargs.get("mol")
      raise DeprecationWarning(
          'Mol is being phased out as a parameter, please pass "datapoint" instead.'
      )
    smile = Chem.MolToSmiles(datapoint)
    if len(smile) > self.max_len:
      return np.array([])

    smile_list = list(smile)
    # Extend shorter strings with padding
    if len(smile) < self.max_len:
      smile_list.extend([PAD_TOKEN] * (self.max_len - len(smile)))

    # Padding before and after
    smile_list += [PAD_TOKEN] * self.pad_len
    smile_list = [PAD_TOKEN] * self.pad_len + smile_list

    smile_seq = self.to_seq(smile_list)
    return smile_seq
