text: # coding: utf-8
  
  
  import os
  import unittest
  
  import pytest
  
  from pymatgen.analysis.fragmenter import Fragmenter, metal_edge_extender
  from pymatgen.analysis.graphs import MoleculeGraph
  from pymatgen.analysis.local_env import OpenBabelNN
  from pymatgen.core.structure import Molecule
  from pymatgen.util.testing import PymatgenTest
  
  __author__ = "Samuel Blau"
  __email__ = "samblau1@gmail.com"
  
  
  test_dir = os.path.join(PymatgenTest.TEST_FILES_DIR, "fragmenter_files")
  
  
  class TestFragmentMolecule(PymatgenTest):
      @classmethod
      def setUpClass(cls):
          cls.pc = Molecule.from_file(os.path.join(test_dir, "PC.xyz"))
          cls.ec = Molecule.from_file(os.path.join(test_dir, "EC.xyz"))
          cls.pos_pc = Molecule.from_file(os.path.join(test_dir, "PC.xyz"))
          cls.pos_pc.set_charge_and_spin(charge=1)
          cls.pc_edges = [
              [5, 10],
              [5, 12],
              [5, 11],
              [5, 3],
              [3, 7],
              [3, 4],
              [3, 0],
              [4, 8],
              [4, 9],
              [4, 1],
              [6, 1],
              [6, 0],
              [6, 2],
          ]
          cls.pc_frag1 = Molecule.from_file(os.path.join(test_dir, "PC_frag1.xyz"))
          cls.pc_frag1_edges = [[0, 2], [4, 2], [2, 1], [1, 3]]
          cls.tfsi = Molecule.from_file(os.path.join(test_dir, "TFSI.xyz"))
          cls.tfsi_edges = (
              [14, 1],
              [1, 4],
              [1, 5],
              [1, 7],
              [7, 11],
              [7, 12],
              [7, 13],
              [14, 0],
              [0, 2],
              [0, 3],
              [0, 6],
              [6, 8],
              [6, 9],
              [6, 10],
          )
          cls.LiEC = Molecule.from_file(os.path.join(test_dir, "LiEC.xyz"))
  
      def test_edges_given_PC_frag1(self):
          fragmenter = Fragmenter(molecule=self.pc_frag1, edges=self.pc_frag1_edges, depth=0)
          self.assertEqual(fragmenter.total_unique_fragments, 12)
  
      def test_babel_PC_frag1(self):
          pytest.importorskip("openbabel", reason="OpenBabel not installed")
          fragmenter = Fragmenter(molecule=self.pc_frag1, depth=0)
          self.assertEqual(fragmenter.total_unique_fragments, 12)
  
      def test_babel_PC_old_defaults(self):
          pytest.importorskip("openbabel", reason="OpenBabel not installed")
          fragmenter = Fragmenter(molecule=self.pc, open_rings=True)
          self.assertEqual(fragmenter.open_rings, True)
          self.assertEqual(fragmenter.opt_steps, 10000)
          default_mol_graph = MoleculeGraph.with_local_env_strategy(self.pc, OpenBabelNN())
          self.assertEqual(fragmenter.mol_graph, default_mol_graph)
          self.assertEqual(fragmenter.total_unique_fragments, 13)
  
      def test_babel_PC_defaults(self):
          pytest.importorskip("openbabel", reason="OpenBabel not installed")
          fragmenter = Fragmenter(molecule=self.pc)
          self.assertEqual(fragmenter.open_rings, False)
          self.assertEqual(fragmenter.opt_steps, 10000)
          default_mol_graph = MoleculeGraph.with_local_env_strategy(self.pc, OpenBabelNN())
          self.assertEqual(fragmenter.mol_graph, default_mol_graph)
          self.assertEqual(fragmenter.total_unique_fragments, 8)
  
      def test_edges_given_PC_not_defaults(self):
          fragmenter = Fragmenter(
              molecule=self.pc,
              edges=self.pc_edges,
              depth=2,
              open_rings=False,
              opt_steps=0,
          )
          self.assertEqual(fragmenter.open_rings, False)
          self.assertEqual(fragmenter.opt_steps, 0)
          edges = {(e[0], e[1]): None for e in self.pc_edges}
          default_mol_graph = MoleculeGraph.with_edges(self.pc, edges=edges)
          self.assertEqual(fragmenter.mol_graph, default_mol_graph)
          self.assertEqual(fragmenter.total_unique_fragments, 20)
  
      def test_edges_given_TFSI(self):
          fragmenter = Fragmenter(molecule=self.tfsi, edges=self.tfsi_edges, depth=0)
          self.assertEqual(fragmenter.total_unique_fragments, 156)
  
      def test_babel_TFSI(self):
          pytest.importorskip("openbabel", reason="OpenBabel not installed")
          fragmenter = Fragmenter(molecule=self.tfsi, depth=0)
          self.assertEqual(fragmenter.total_unique_fragments, 156)
  
      def test_babel_PC_with_RO_depth_0_vs_depth_10(self):
          pytest.importorskip("openbabel", reason="OpenBabel not installed")
          fragmenter0 = Fragmenter(molecule=self.pc, depth=0, open_rings=True, opt_steps=1000)
          self.assertEqual(fragmenter0.total_unique_fragments, 509)
  
          fragmenter10 = Fragmenter(molecule=self.pc, depth=10, open_rings=True, opt_steps=1000)
          self.assertEqual(fragmenter10.total_unique_fragments, 509)
  
          fragments_by_level = fragmenter10.fragments_by_level
          num_frags_by_level = [13, 51, 95, 115, 105, 75, 39, 14, 2, 0]
          for ii in range(10):
              num_frags = 0
              for key in fragments_by_level[str(ii)]:
                  num_frags += len(fragments_by_level[str(ii)][key])
              self.assertEqual(num_frags, num_frags_by_level[ii])
  
      def test_PC_depth_0_vs_depth_10(self):
          fragmenter0 = Fragmenter(molecule=self.pc, edges=self.pc_edges, depth=0, open_rings=False)
          self.assertEqual(fragmenter0.total_unique_fragments, 295)
  
          fragmenter10 = Fragmenter(molecule=self.pc, edges=self.pc_edges, depth=10, open_rings=False)
          self.assertEqual(fragmenter10.total_unique_fragments, 63)
  
          fragments_by_level = fragmenter10.fragments_by_level
          num_frags_by_level = [8, 12, 15, 14, 9, 4, 1]
          for ii in range(7):
              num_frags = 0
              for key in fragments_by_level[str(ii)]:
                  num_frags += len(fragments_by_level[str(ii)][key])
              self.assertEqual(num_frags, num_frags_by_level[ii])
  
      def test_PC_frag1_then_PC(self):
          frag1 = Fragmenter(molecule=self.pc_frag1, edges=self.pc_frag1_edges, depth=0)
          self.assertEqual(frag1.new_unique_fragments, frag1.total_unique_fragments)
          frag2 = Fragmenter(
              molecule=self.pc,
              edges=self.pc_edges,
              depth=0,
              open_rings=False,
              prev_unique_frag_dict=frag1.unique_frag_dict,
          )
          self.assertEqual(frag2.new_unique_fragments, 295 - 12)
  
      def test_PC_then_EC_depth_10(self):
          pytest.importorskip("openbabel", reason="OpenBabel not installed")
          fragPC = Fragmenter(molecule=self.pc, depth=10, open_rings=True)
          fragEC = Fragmenter(
              molecule=self.ec,
              depth=10,
              open_rings=True,
              prev_unique_frag_dict=fragPC.unique_frag_dict,
          )
          self.assertEqual(fragEC.new_unique_fragments, 11)
          self.assertEqual(fragEC.total_unique_fragments, 509 + 11)
  
      def test_metal_edge_extender(self):
          mol_graph = MoleculeGraph.with_edges(
              molecule=self.LiEC,
              edges={
                  (0, 2): None,
                  (0, 1): None,
                  (1, 3): None,
                  (1, 4): None,
                  (2, 7): None,
                  (2, 5): None,
                  (2, 8): None,
                  (3, 6): None,
                  (4, 5): None,
                  (5, 9): None,
                  (5, 10): None,
              },
          )
          self.assertEqual(len(mol_graph.graph.edges), 11)
          extended_mol_graph = metal_edge_extender(mol_graph)
          self.assertEqual(len(mol_graph.graph.edges), 12)
  
  
  if __name__ == "__main__":
      unittest.main()
  