text: # -*- coding: utf-8 -*-
  #
  # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
  # SPDX-License-Identifier: Apache-2.0
  
  import torch
  
  from rdkit import Chem
  from dgllife.utils.featurizers import *
  
  def test_one_hot_encoding():
      x = 1.
      allowable_set = [0., 1., 2.]
      assert one_hot_encoding(x, allowable_set) == [0, 1, 0]
      assert one_hot_encoding(x, allowable_set, encode_unknown=True) == [0, 1, 0, 0]
  
      assert one_hot_encoding(x, allowable_set) == [0, 1, 0, 0]
      assert one_hot_encoding(x, allowable_set, encode_unknown=True) == [0, 1, 0, 0]
      assert one_hot_encoding(-1, allowable_set, encode_unknown=True) == [0, 0, 0, 1]
  
  def test_mol1():
      return Chem.MolFromSmiles('CCO')
  
  def test_mol2():
      return Chem.MolFromSmiles('C1=CC2=CC=CC=CC2=C1')
  
  def test_mol3():
      return Chem.MolFromSmiles('O=C(O)/C=C/C(=O)O')
  
  def test_atom_type_one_hot():
      mol = test_mol1()
      assert atom_type_one_hot(mol.GetAtomWithIdx(0), ['C', 'O']) == [1, 0]
      assert atom_type_one_hot(mol.GetAtomWithIdx(2), ['C', 'O']) == [0, 1]
  
  def test_atomic_number_one_hot():
      mol = test_mol1()
      assert atomic_number_one_hot(mol.GetAtomWithIdx(0), [6, 8]) == [1, 0]
      assert atomic_number_one_hot(mol.GetAtomWithIdx(2), [6, 8]) == [0, 1]
  
  def test_atomic_number():
      mol = test_mol1()
      assert atomic_number(mol.GetAtomWithIdx(0)) == [6]
      assert atomic_number(mol.GetAtomWithIdx(2)) == [8]
  
  def test_atom_degree_one_hot():
      mol = test_mol1()
      assert atom_degree_one_hot(mol.GetAtomWithIdx(0), [0, 1, 2]) == [0, 1, 0]
      assert atom_degree_one_hot(mol.GetAtomWithIdx(1), [0, 1, 2]) == [0, 0, 1]
  
  def test_atom_degree():
      mol = test_mol1()
      assert atom_degree(mol.GetAtomWithIdx(0)) == [1]
      assert atom_degree(mol.GetAtomWithIdx(1)) == [2]
  
  def test_atom_total_degree_one_hot():
      mol = test_mol1()
      assert atom_total_degree_one_hot(mol.GetAtomWithIdx(0), [0, 2, 4]) == [0, 0, 1]
      assert atom_total_degree_one_hot(mol.GetAtomWithIdx(2), [0, 2, 4]) == [0, 1, 0]
  
  def test_atom_total_degree():
      mol = test_mol1()
      assert atom_total_degree(mol.GetAtomWithIdx(0)) == [4]
      assert atom_total_degree(mol.GetAtomWithIdx(2)) == [2]
  
  def test_atom_explicit_valence_one_hot():
      mol = test_mol1()
      assert atom_explicit_valence_one_hot(mol.GetAtomWithIdx(0), [1, 2, 3]) == [1, 0, 0]
      assert atom_explicit_valence_one_hot(mol.GetAtomWithIdx(1), [1, 2, 3]) == [0, 1, 0]
  
  def test_atom_explicit_valence():
      mol = test_mol1()
      assert atom_explicit_valence(mol.GetAtomWithIdx(0)) == [1]
      assert atom_explicit_valence(mol.GetAtomWithIdx(1)) == [2]
  
  def test_atom_implicit_valence_one_hot():
      mol = test_mol1()
      assert atom_implicit_valence_one_hot(mol.GetAtomWithIdx(0), [1, 2, 3]) == [0, 0, 1]
      assert atom_implicit_valence_one_hot(mol.GetAtomWithIdx(1), [1, 2, 3]) == [0, 1, 0]
  
  def test_atom_implicit_valence():
      mol = test_mol1()
      assert atom_implicit_valence(mol.GetAtomWithIdx(0)) == [3]
      assert atom_implicit_valence(mol.GetAtomWithIdx(1)) == [2]
  
  def test_atom_hybridization_one_hot():
      mol = test_mol1()
      assert atom_hybridization_one_hot(mol.GetAtomWithIdx(0)) == [0, 0, 1, 0, 0]
  
  def test_atom_total_num_H_one_hot():
      mol = test_mol1()
      assert atom_total_num_H_one_hot(mol.GetAtomWithIdx(0)) == [0, 0, 0, 1, 0]
      assert atom_total_num_H_one_hot(mol.GetAtomWithIdx(1)) == [0, 0, 1, 0, 0]
  
  def test_atom_total_num_H():
      mol = test_mol1()
      assert atom_total_num_H(mol.GetAtomWithIdx(0)) == [3]
      assert atom_total_num_H(mol.GetAtomWithIdx(1)) == [2]
  
  def test_atom_formal_charge_one_hot():
      mol = test_mol1()
      assert atom_formal_charge_one_hot(mol.GetAtomWithIdx(0)) == [0, 0, 1, 0, 0]
  
  def test_atom_formal_charge():
      mol = test_mol1()
      assert atom_formal_charge(mol.GetAtomWithIdx(0)) == [0]
  
  def test_atom_num_radical_electrons_one_hot():
      mol = test_mol1()
      assert atom_num_radical_electrons_one_hot(mol.GetAtomWithIdx(0)) == [1, 0, 0, 0, 0]
  
  def test_atom_num_radical_electrons():
      mol = test_mol1()
      assert atom_num_radical_electrons(mol.GetAtomWithIdx(0)) == [0]
  
  def test_atom_is_aromatic_one_hot():
      mol = test_mol1()
      assert atom_is_aromatic_one_hot(mol.GetAtomWithIdx(0)) == [1, 0]
      mol = test_mol2()
      assert atom_is_aromatic_one_hot(mol.GetAtomWithIdx(0)) == [0, 1]
  
  def test_atom_is_aromatic():
      mol = test_mol1()
      assert atom_is_aromatic(mol.GetAtomWithIdx(0)) == [0]
      mol = test_mol2()
      assert atom_is_aromatic(mol.GetAtomWithIdx(0)) == [1]
  
  def test_atom_is_in_ring_one_hot():
      mol = test_mol1()
      assert atom_is_in_ring_one_hot(mol.GetAtomWithIdx(0)) == [1, 0]
      mol = test_mol2()
      assert atom_is_in_ring_one_hot(mol.GetAtomWithIdx(0)) == [0, 1]
  
  def test_atom_is_in_ring():
      mol = test_mol1()
      assert atom_is_in_ring(mol.GetAtomWithIdx(0)) == [0]
      mol = test_mol2()
      assert atom_is_in_ring(mol.GetAtomWithIdx(0)) == [1]
  
  def test_atom_chiral_tag_one_hot():
      mol = test_mol1()
      assert atom_chiral_tag_one_hot(mol.GetAtomWithIdx(0)) == [1, 0, 0, 0]
  
  def test_atom_mass():
      mol = test_mol1()
      atom = mol.GetAtomWithIdx(0)
      assert atom_mass(atom) == [atom.GetMass() * 0.01]
      atom = mol.GetAtomWithIdx(1)
      assert atom_mass(atom) == [atom.GetMass() * 0.01]
  
  def test_concat_featurizer():
      test_featurizer = ConcatFeaturizer(
          [atom_is_aromatic_one_hot, atom_chiral_tag_one_hot]
      )
      mol = test_mol1()
      assert test_featurizer(mol.GetAtomWithIdx(0)) == [1, 0, 1, 0, 0, 0]
      mol = test_mol2()
      assert test_featurizer(mol.GetAtomWithIdx(0)) == [0, 1, 1, 0, 0, 0]
  
  class TestAtomFeaturizer(BaseAtomFeaturizer):
      def __init__(self):
          super(TestAtomFeaturizer, self).__init__(
              featurizer_funcs={
                  'h1': ConcatFeaturizer([atom_total_degree_one_hot,
                                          atom_formal_charge_one_hot]),
                  'h2': ConcatFeaturizer([atom_num_radical_electrons_one_hot])
              }
          )
  
  def test_base_atom_featurizer():
      test_featurizer = TestAtomFeaturizer()
      assert test_featurizer.feat_size('h1') == 11
      assert test_featurizer.feat_size('h2') == 5
      mol = test_mol1()
      feats = test_featurizer(mol)
      assert torch.allclose(feats['h1'],
                            torch.tensor([[0., 0., 0., 0., 1., 0., 0., 0., 1., 0., 0.],
                                          [0., 0., 0., 0., 1., 0., 0., 0., 1., 0., 0.],
                                          [0., 0., 1., 0., 0., 0., 0., 0., 1., 0., 0.]]))
      assert torch.allclose(feats['h2'],
                            torch.tensor([[1., 0., 0., 0., 0.],
                                          [1., 0., 0., 0., 0.],
                                          [1., 0., 0., 0., 0.]]))
  
  def test_canonical_atom_featurizer():
      test_featurizer = CanonicalAtomFeaturizer()
      assert test_featurizer.feat_size() == 74
      assert test_featurizer.feat_size('h') == 74
      mol = test_mol1()
      feats = test_featurizer(mol)
      assert list(feats.keys()) == ['h']
      assert torch.allclose(feats['h'],
                            torch.tensor([[1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                                           0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                                           0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                                           0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0.,
                                           0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0.,
                                           0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0.,
                                           1., 0.],
                                          [1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                                           0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                                           0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                                           0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0.,
                                           0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0.,
                                           0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 1.,
                                           0., 0.],
                                          [0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                                           0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                                           0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                                           0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0.,
                                           0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0.,
                                           0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 1., 0.,
                                           0., 0.]]))
  
  def test_weave_atom_featurizer():
      featurizer = WeaveAtomFeaturizer()
      assert featurizer.feat_size() == 27
      mol = test_mol1()
      feats = featurizer(mol)
      assert list(feats.keys()) == ['h']
      assert torch.allclose(feats['h'],
                            torch.tensor([[0.0000, 1.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                                           0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                                           0.0000, 0.0000, -0.0418, 0.0000, 0.0000, 0.0000,
                                           1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                                           0.0000, 0.0000, 0.0000],
                                          [0.0000, 1.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                                           0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                                           0.0000, 0.0000, 0.0402, 0.0000, 0.0000, 0.0000,
                                           1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                                           0.0000, 0.0000, 0.0000],
                                          [0.0000, 0.0000, 0.0000, 1.0000, 0.0000, 0.0000,
                                           0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
                                           0.0000, 0.0000, -0.3967, 0.0000, 0.0000, 0.0000,
                                           1.0000, 1.0000, 1.0000, 0.0000, 0.0000, 0.0000,
                                           0.0000, 0.0000, 0.0000]]), rtol=1e-3)
  
  def test_pretrain_atom_featurizer():
      featurizer = PretrainAtomFeaturizer()
      mol = test_mol1()
      feats = featurizer(mol)
      assert list(feats.keys()) == ['atomic_number', 'chirality_type']
      assert torch.allclose(feats['atomic_number'], torch.tensor([[5, 5, 7]]))
      assert torch.allclose(feats['chirality_type'], torch.tensor([[0, 0, 0]]))
  
  def test_bond_type_one_hot():
      mol = test_mol1()
      assert bond_type_one_hot(mol.GetBondWithIdx(0)) == [1, 0, 0, 0]
      mol = test_mol2()
      assert bond_type_one_hot(mol.GetBondWithIdx(0)) == [0, 0, 0, 1]
  
  def test_bond_is_conjugated_one_hot():
      mol = test_mol1()
      assert bond_is_conjugated_one_hot(mol.GetBondWithIdx(0)) == [1, 0]
      mol = test_mol2()
      assert bond_is_conjugated_one_hot(mol.GetBondWithIdx(0)) == [0, 1]
  
  def test_bond_is_conjugated():
      mol = test_mol1()
      assert bond_is_conjugated(mol.GetBondWithIdx(0)) == [0]
      mol = test_mol2()
      assert bond_is_conjugated(mol.GetBondWithIdx(0)) == [1]
  
  def test_bond_is_in_ring_one_hot():
      mol = test_mol1()
      assert bond_is_in_ring_one_hot(mol.GetBondWithIdx(0)) == [1, 0]
      mol = test_mol2()
      assert bond_is_in_ring_one_hot(mol.GetBondWithIdx(0)) == [0, 1]
  
  def test_bond_is_in_ring():
      mol = test_mol1()
      assert bond_is_in_ring(mol.GetBondWithIdx(0)) == [0]
      mol = test_mol2()
      assert bond_is_in_ring(mol.GetBondWithIdx(0)) == [1]
  
  def test_bond_stereo_one_hot():
      mol = test_mol1()
      assert bond_stereo_one_hot(mol.GetBondWithIdx(0)) == [1, 0, 0, 0, 0, 0]
  
  def test_bond_direction_one_hot():
      mol = test_mol3()
      assert bond_direction_one_hot(mol.GetBondWithIdx(0)) == [1, 0, 0]
      assert bond_direction_one_hot(mol.GetBondWithIdx(2)) == [0, 1, 0]
  
  class TestBondFeaturizer(BaseBondFeaturizer):
      def __init__(self):
          super(TestBondFeaturizer, self).__init__(
              featurizer_funcs={
                  'h1': ConcatFeaturizer([bond_is_in_ring, bond_is_conjugated]),
                  'h2': ConcatFeaturizer([bond_stereo_one_hot])
              }
          )
  
  def test_base_bond_featurizer():
      test_featurizer = TestBondFeaturizer()
      assert test_featurizer.feat_size('h1') == 2
      assert test_featurizer.feat_size('h2') == 6
      mol = test_mol1()
      feats = test_featurizer(mol)
      assert torch.allclose(feats['h1'], torch.tensor([[0., 0.], [0., 0.], [0., 0.], [0., 0.]]))
      assert torch.allclose(feats['h2'], torch.tensor([[1., 0., 0., 0., 0., 0.],
                                                       [1., 0., 0., 0., 0., 0.],
                                                       [1., 0., 0., 0., 0., 0.],
                                                       [1., 0., 0., 0., 0., 0.]]))
  
  def test_canonical_bond_featurizer():
      test_featurizer = CanonicalBondFeaturizer()
      assert test_featurizer.feat_size() == 12
      assert test_featurizer.feat_size('e') == 12
      mol = test_mol1()
      feats = test_featurizer(mol)
      assert torch.allclose(feats['e'], torch.tensor(
          [[1., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],
           [1., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],
           [1., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],
           [1., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0.]]))
  
  def test_weave_edge_featurizer():
      test_featurizer = WeaveEdgeFeaturizer()
      assert test_featurizer.feat_size() == 12
      mol = test_mol1()
      feats = test_featurizer(mol)
      assert torch.allclose(feats['e'],
                            torch.tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
                                          [1., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0.],
                                          [1., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
                                          [1., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0.],
                                          [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
                                          [1., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0.],
                                          [1., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
                                          [1., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0.],
                                          [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]]))
  
  def test_pretrain_bond_featurizer():
      mol = test_mol3()
      test_featurizer = PretrainBondFeaturizer()
      feats = test_featurizer(mol)
      assert torch.allclose(feats['bond_type'].nonzero(),
                            torch.tensor([[0], [1], [6], [7], [10], [11], [14], [15],
                                          [16], [17], [18], [19], [20], [21]]))
      assert torch.allclose(feats['bond_direction_type'].nonzero(),
                            torch.tensor([[4], [5], [8], [9]]))
  
      test_featurizer = PretrainBondFeaturizer(self_loop=False)
      feats = test_featurizer(mol)
      assert torch.allclose(feats['bond_type'].nonzero(),
                            torch.tensor([[0], [1], [6], [7], [10], [11]]))
      assert torch.allclose(feats['bond_direction_type'].nonzero(),
                            torch.tensor([[4], [5], [8], [9]]))
  
  if __name__ == '__main__':
      test_one_hot_encoding()
      test_atom_type_one_hot()
      test_atomic_number_one_hot()
      test_atomic_number()
      test_atom_degree_one_hot()
      test_atom_degree()
      test_atom_total_degree_one_hot()
      test_atom_total_degree()
      test_atom_explicit_valence()
      test_atom_implicit_valence_one_hot()
      test_atom_implicit_valence()
      test_atom_hybridization_one_hot()
      test_atom_total_num_H_one_hot()
      test_atom_total_num_H()
      test_atom_formal_charge_one_hot()
      test_atom_formal_charge()
      test_atom_num_radical_electrons_one_hot()
      test_atom_num_radical_electrons()
      test_atom_is_aromatic_one_hot()
      test_atom_is_aromatic()
      test_atom_is_in_ring_one_hot()
      test_atom_is_in_ring()
      test_atom_chiral_tag_one_hot()
      test_atom_mass()
      test_concat_featurizer()
      test_base_atom_featurizer()
      test_canonical_atom_featurizer()
      test_weave_atom_featurizer()
      test_pretrain_atom_featurizer()
      test_bond_type_one_hot()
      test_bond_is_conjugated_one_hot()
      test_bond_is_conjugated()
      test_bond_is_in_ring_one_hot()
      test_bond_is_in_ring()
      test_bond_stereo_one_hot()
      test_bond_direction_one_hot()
      test_base_bond_featurizer()
      test_canonical_bond_featurizer()
      test_weave_edge_featurizer()
      test_pretrain_bond_featurizer()
  