"""
DGL-based AttentiveFP for graph property prediction.
"""
import torch.nn as nn
import torch.nn.functional as F

from deepchem.models.losses import Loss, L2Loss, SparseSoftmaxCrossEntropy
from deepchem.models.torch_models.torch_model import TorchModel


class AttentiveFP(nn.Module):
  """Model for Graph Property Prediction.

  This model proceeds as follows:

  * Combine node features and edge features for initializing node representations,
    which involves a round of message passing
  * Update node representations with multiple rounds of message passing
  * For each graph, compute its representation by combining the representations
    of all nodes in it, which involves a gated recurrent unit (GRU).
  * Perform the final prediction using a linear layer

  Examples
  --------

  >>> import deepchem as dc
  >>> import dgl
  >>> from deepchem.models import AttentiveFP
  >>> smiles = ["C1CCC1", "C1=CC=CN=C1"]
  >>> featurizer = dc.feat.MolGraphConvFeaturizer(use_edges=True)
  >>> graphs = featurizer.featurize(smiles)
  >>> print(type(graphs[0]))
  <class 'deepchem.feat.graph_data.GraphData'>
  >>> dgl_graphs = [graphs[i].to_dgl_graph(self_loop=True) for i in range(len(graphs))]
  >>> # Batch two graphs into a graph of two connected components
  >>> batch_dgl_graph = dgl.batch(dgl_graphs)
  >>> model = AttentiveFP(n_tasks=1, mode='regression')
  >>> preds = model(batch_dgl_graph)
  >>> print(type(preds))
  <class 'torch.Tensor'>
  >>> preds.shape == (2, 1)
  True

  References
  ----------
  .. [1] Zhaoping Xiong, Dingyan Wang, Xiaohong Liu, Feisheng Zhong, Xiaozhe Wan, Xutong Li,
         Zhaojun Li, Xiaomin Luo, Kaixian Chen, Hualiang Jiang, and Mingyue Zheng. "Pushing
         the Boundaries of Molecular Representation for Drug Discovery with the Graph Attention
         Mechanism." Journal of Medicinal Chemistry. 2020, 63, 16, 8749–8760.

  Notes
  -----
  This class requires DGL (https://github.com/dmlc/dgl) and DGL-LifeSci
  (https://github.com/awslabs/dgl-lifesci) to be installed.
  """

  def __init__(self,
               n_tasks: int,
               num_layers: int = 2,
               num_timesteps: int = 2,
               graph_feat_size: int = 200,
               dropout: float = 0.,
               mode: str = 'regression',
               number_atom_features: int = 30,
               number_bond_features: int = 11,
               n_classes: int = 2,
               nfeat_name: str = 'x',
               efeat_name: str = 'edge_attr'):
    """
    Parameters
    ----------
    n_tasks: int
      Number of tasks.
    num_layers: int
      Number of graph neural network layers, i.e. number of rounds of message passing.
      Default to 2.
    num_timesteps: int
      Number of time steps for updating graph representations with a GRU. Default to 2.
    graph_feat_size: int
      Size for graph representations. Default to 200.
    dropout: float
      Dropout probability. Default to 0.
    mode: str
      The model type, 'classification' or 'regression'. Default to 'regression'.
    number_atom_features: int
      The length of the initial atom feature vectors. Default to 30.
    number_bond_features: int
      The length of the initial bond feature vectors. Default to 11.
    n_classes: int
      The number of classes to predict per task
      (only used when ``mode`` is 'classification'). Default to 2.
    nfeat_name: str
      For an input graph ``g``, the model assumes that it stores node features in
      ``g.ndata[nfeat_name]`` and will retrieve input node features from that.
      Default to 'x'.
    efeat_name: str
      For an input graph ``g``, the model assumes that it stores edge features in
      ``g.edata[efeat_name]`` and will retrieve input edge features from that.
      Default to 'edge_attr'.
    """
    try:
      import dgl
    except:
      raise ImportError('This class requires dgl.')
    try:
      import dgllife
    except:
      raise ImportError('This class requires dgllife.')

    if mode not in ['classification', 'regression']:
      raise ValueError("mode must be either 'classification' or 'regression'")

    super(AttentiveFP, self).__init__()

    self.n_tasks = n_tasks
    self.mode = mode
    self.n_classes = n_classes
    self.nfeat_name = nfeat_name
    self.efeat_name = efeat_name
    if mode == 'classification':
      out_size = n_tasks * n_classes
    else:
      out_size = n_tasks

    from dgllife.model import AttentiveFPPredictor as DGLAttentiveFPPredictor

    self.model = DGLAttentiveFPPredictor(
        node_feat_size=number_atom_features,
        edge_feat_size=number_bond_features,
        num_layers=num_layers,
        num_timesteps=num_timesteps,
        graph_feat_size=graph_feat_size,
        n_tasks=out_size,
        dropout=dropout)

  def forward(self, g):
    """Predict graph labels

    Parameters
    ----------
    g: DGLGraph
      A DGLGraph for a batch of graphs. It stores the node features in
      ``dgl_graph.ndata[self.nfeat_name]`` and edge features in
      ``dgl_graph.edata[self.efeat_name]``.

    Returns
    -------
    torch.Tensor
      The model output.

      * When self.mode = 'regression',
        its shape will be ``(dgl_graph.batch_size, self.n_tasks)``.
      * When self.mode = 'classification', the output consists of probabilities
        for classes. Its shape will be
        ``(dgl_graph.batch_size, self.n_tasks, self.n_classes)`` if self.n_tasks > 1;
        its shape will be ``(dgl_graph.batch_size, self.n_classes)`` if self.n_tasks is 1.
    torch.Tensor, optional
      This is only returned when self.mode = 'classification', the output consists of the
      logits for classes before softmax.
    """
    node_feats = g.ndata[self.nfeat_name]
    edge_feats = g.edata[self.efeat_name]
    out = self.model(g, node_feats, edge_feats)

    if self.mode == 'classification':
      if self.n_tasks == 1:
        logits = out.view(-1, self.n_classes)
        softmax_dim = 1
      else:
        logits = out.view(-1, self.n_tasks, self.n_classes)
        softmax_dim = 2
      proba = F.softmax(logits, dim=softmax_dim)
      return proba, logits
    else:
      return out


class AttentiveFPModel(TorchModel):
  """Model for Graph Property Prediction.

  This model proceeds as follows:

  * Combine node features and edge features for initializing node representations,
    which involves a round of message passing
  * Update node representations with multiple rounds of message passing
  * For each graph, compute its representation by combining the representations
    of all nodes in it, which involves a gated recurrent unit (GRU).
  * Perform the final prediction using a linear layer

  Examples
  --------
  >>> import deepchem as dc
  >>> from deepchem.models import AttentiveFPModel
  >>> # preparing dataset
  >>> smiles = ["C1CCC1", "C1=CC=CN=C1"]
  >>> labels = [0., 1.]
  >>> featurizer = dc.feat.MolGraphConvFeaturizer(use_edges=True)
  >>> X = featurizer.featurize(smiles)
  >>> dataset = dc.data.NumpyDataset(X=X, y=labels)
  >>> # training model
  >>> model = AttentiveFPModel(mode='classification', n_tasks=1,
  ...    batch_size=16, learning_rate=0.001)
  >>> loss = model.fit(dataset, nb_epoch=5)

  References
  ----------
  .. [1] Zhaoping Xiong, Dingyan Wang, Xiaohong Liu, Feisheng Zhong, Xiaozhe Wan, Xutong Li,
         Zhaojun Li, Xiaomin Luo, Kaixian Chen, Hualiang Jiang, and Mingyue Zheng. "Pushing
         the Boundaries of Molecular Representation for Drug Discovery with the Graph
         Attention Mechanism." Journal of Medicinal Chemistry. 2020, 63, 16, 8749–8760.

  Notes
  -----
  This class requires DGL (https://github.com/dmlc/dgl) and DGL-LifeSci
  (https://github.com/awslabs/dgl-lifesci) to be installed.
  """

  def __init__(self,
               n_tasks: int,
               num_layers: int = 2,
               num_timesteps: int = 2,
               graph_feat_size: int = 200,
               dropout: float = 0.,
               mode: str = 'regression',
               number_atom_features: int = 30,
               number_bond_features: int = 11,
               n_classes: int = 2,
               self_loop: bool = True,
               **kwargs):
    """
    Parameters
    ----------
    n_tasks: int
      Number of tasks.
    num_layers: int
      Number of graph neural network layers, i.e. number of rounds of message passing.
      Default to 2.
    num_timesteps: int
      Number of time steps for updating graph representations with a GRU. Default to 2.
    graph_feat_size: int
      Size for graph representations. Default to 200.
    dropout: float
      Dropout probability. Default to 0.
    mode: str
      The model type, 'classification' or 'regression'. Default to 'regression'.
    number_atom_features: int
      The length of the initial atom feature vectors. Default to 30.
    number_bond_features: int
      The length of the initial bond feature vectors. Default to 11.
    n_classes: int
      The number of classes to predict per task
      (only used when ``mode`` is 'classification'). Default to 2.
    self_loop: bool
      Whether to add self loops for the nodes, i.e. edges from nodes to themselves.
      When input graphs have isolated nodes, self loops allow preserving the original feature
      of them in message passing. Default to True.
    kwargs
      This can include any keyword argument of TorchModel.
    """
    model = AttentiveFP(
        n_tasks=n_tasks,
        num_layers=num_layers,
        num_timesteps=num_timesteps,
        graph_feat_size=graph_feat_size,
        dropout=dropout,
        mode=mode,
        number_atom_features=number_atom_features,
        number_bond_features=number_bond_features,
        n_classes=n_classes)
    if mode == 'regression':
      loss: Loss = L2Loss()
      output_types = ['prediction']
    else:
      loss = SparseSoftmaxCrossEntropy()
      output_types = ['prediction', 'loss']
    super(AttentiveFPModel, self).__init__(
        model, loss=loss, output_types=output_types, **kwargs)

    self._self_loop = self_loop

  def _prepare_batch(self, batch):
    """Create batch data for AttentiveFP.

    Parameters
    ----------
    batch: tuple
      The tuple is ``(inputs, labels, weights)``.

    Returns
    -------
    inputs: DGLGraph
      DGLGraph for a batch of graphs.
    labels: list of torch.Tensor or None
      The graph labels.
    weights: list of torch.Tensor or None
      The weights for each sample or sample/task pair converted to torch.Tensor.
    """
    try:
      import dgl
    except:
      raise ImportError('This class requires dgl.')

    inputs, labels, weights = batch
    dgl_graphs = [
        graph.to_dgl_graph(self_loop=self._self_loop) for graph in inputs[0]
    ]
    inputs = dgl.batch(dgl_graphs).to(self.device)
    _, labels, weights = super(AttentiveFPModel, self)._prepare_batch(
        ([], labels, weights))
    return inputs, labels, weights
