text: # -*- coding: utf-8 -*-
  # Copyright 2018 University of Groningen
  #
  # 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.
  
  """
  Unit tests for the GRO reader.
  """
  # Pylint complains about hypothesis strategies not receiving a value for the
  # `draw` parameter. This is because the `draw` parameter is implicitly filled
  # by hypothesis. The pylint warning is disabled for the file instead of being
  # disabled at every call of a strategy.
  # pylint: disable=no-value-for-parameter
  
  from pprint import pprint
  import copy
  import collections
  import itertools
  
  import numpy as np
  import networkx as nx
  
  import pytest
  from hypothesis import given, assume
  import hypothesis.strategies as st
  import hypothesis.extra.numpy as hnp
  
  import vermouth
  from vermouth.file_writer import DeferredFileWriter
  from vermouth.utils import are_different
  from vermouth.molecule import Molecule
  from vermouth.gmx import gro
  
  # pylint: disable=redefined-outer-name
  
  
  # The data comes from residues 3 to 5 of 1BTA.pdb. The atoms from ILE 5 are
  # changed in solvent.
  GRO_CONTENT = """\
      1ALA      N    1
      1ALA     CA    2
      1ALA      C    3
      1ALA      O    4
      1ALA     CB    5
      1ALA      H    6
      1ALA     HA    7
      1ALA    HB1    8
      1ALA    HB2    9
      1ALA    HB3   10
      2VAL      N   11
      2VAL     CA   12
      2VAL      C   13
      2VAL      O   14
      2VAL     CB   15
      2VAL    CG1   16
      2VAL    CG2   17
      2VAL      H   18
      2VAL     HA   19
      2VAL     HB   20
      2VAL   HG11   21
      2VAL   HG12   22
      2VAL   HG13   23
      2VAL   HG21   24
      2VAL   HG22   25
      2VAL   HG23   26
      3SOL     OW   27
      3SOL    HW1   28
      3SOL    HW2   29
      4SOL     OW   30
      4SOL    HW1   31
      4SOL    HW2   32
      5SOL     OW   33
      5SOL    HW1   34
      5SOL    HW2   35
      6SOL     OW   36
      6SOL    HW1   37
      6SOL    HW2   38
      7SOL     OW   39
      7SOL    HW1   40
      7SOL    HW2   41
      8SOL     OW   42
      8SOL    HW1   43
      8SOL    HW2   44
  """.split('\n')
  
  _COORDINATES_STR = """\
    -0.204   0.447   0.725
    -0.109   0.340   0.685
     0.032   0.401   0.675
     0.047   0.521   0.659
    -0.149   0.283   0.549
    -0.199   0.535   0.683
    -0.109   0.262   0.759
    -0.226   0.346   0.505
    -0.063   0.281   0.483
    -0.187   0.183   0.561
     0.134   0.321   0.686
     0.273   0.377   0.678
     0.364   0.277   0.607
     0.384   0.167   0.654
     0.325   0.402   0.819
     0.462   0.471   0.812
     0.226   0.491   0.895
     0.121   0.226   0.700
     0.271   0.470   0.623
     0.335   0.308   0.871
     0.455   0.557   0.747
     0.491   0.503   0.911
     0.535   0.401   0.773
     0.178   0.559   0.825
     0.151   0.430   0.942
     0.279   0.549   0.969
     0.421   0.315   0.496
     0.512   0.221   0.424
     0.657   0.257   0.455
     0.718   0.339   0.389
     0.488   0.230   0.274
     0.336   0.223   0.248
     0.562   0.114   0.206
     0.303   0.125   0.135
     0.404   0.405   0.460
     0.492   0.120   0.458
     0.527   0.324   0.237
     0.286   0.191   0.338
     0.301   0.321   0.221
     0.658   0.101   0.253
     0.504   0.024   0.216
     0.576   0.137   0.101
     0.378   0.133   0.057
     0.304   0.024   0.174\
  """
  COORDINATES = [[float(x) for x in line.split()]
                 for line in _COORDINATES_STR.split('\n')]
  VELOCITIES = [[x + 1 for x in line] for line in COORDINATES]
  
  DIFFERENCE_USE_CASE = (
      '', [], 'a', 2, 3, True, False, None, float('nan'), float('inf'),
      [True], [False], [True, False], [False, False], [None], [None, None],
      [1.1, 1.1], [0.0, 1.1],
      [float('nan'), float('inf')],
      [float('inf'), float('nan')],
      ['a', 'b'], 'ab',
      np.array([1.2, 2.3, 4.5]),
      np.array([[1.2, 2.3, 4.5], [5.6, 7.8, 8.9]]),
      np.array([np.nan, np.inf, 3]),
      [[6, 8], [9, 10, 11]],
      {'a': 3, 2: 'tata'},
      {'a': 3, 2: np.array([2, 3, 4])},
      {'a': 4, 2: 'not tata'},
  )
  
  
  @pytest.fixture(params=[True, False])
  def gro_reference(request, tmpdir_factory):
      """
      Generate a GRO file and the corresponding molecule.
      """
      filename = tmpdir_factory.mktemp("data").join("tmp.gro")
      with open(str(filename), 'w') as outfile:
          write_ref_gro(outfile, velocities=request.param, box='10.0 11.1 12.2')
      molecule = build_ref_molecule(velocities=request.param)
      return filename, molecule
  
  
  @pytest.fixture(params=[43, 45])
  def gro_wrong_length(request, gro_reference, tmpdir_factory):  # pylint: disable=redefined-outer-name
      """
      Generate a GRO file with a wrong number of atoms on line 2.
      """
      path_in, _ = gro_reference
      path_out = tmpdir_factory.mktemp("data").join("wrong.gro")
      with open(str(path_in)) as infile, open(str(path_out), 'w') as outfile:
          outfile.write(next(infile))
          outfile.write('{}\n'.format(request.param))
          for line in infile:
              outfile.write(line)
      return path_out
  
  
  def write_ref_gro(outfile, velocities=False, box='10.0 10.0 10.0'):
      """
      Write a GRO file from the reference data.
  
      Parameters
      ----------
      outfile: file
          The file in which to write.
      velocities: bool
          Set if velocities must be written.
      box: str
          The box string to write at the end of the file.
      """
      velocity_fmt = ''
      if velocities:
          velocity_fmt = '{:8.4f}' * 3
      outfile.write('Just a title\n')
      outfile.write(str(len(COORDINATES)) + '\n')
      for atom, coords, vels in zip(GRO_CONTENT, COORDINATES, VELOCITIES):
          outfile.write(('{}{:8.3f}{:8.3f}{:8.3f}' + velocity_fmt + '\n')
                        .format(atom, *itertools.chain(coords, vels)))
      outfile.write(box)
      outfile.write('\n')
  
  
  def build_ref_molecule(velocities=False):
      """
      Build the molecule graph corresponding to the reference data.
  
      Parameters
      ----------
      velocities: bool
          Set if velocities must be written.
  
      Return
      ------
      Molecule
      """
      nodes = (
          {'resid': 1, 'resname': 'ALA', 'atomname': 'N', 'atomid': 1, 'element': 'N'},
          {'resid': 1, 'resname': 'ALA', 'atomname': 'CA', 'atomid': 2, 'element': 'C'},
          {'resid': 1, 'resname': 'ALA', 'atomname': 'C', 'atomid': 3, 'element': 'C'},
          {'resid': 1, 'resname': 'ALA', 'atomname': 'O', 'atomid': 4, 'element': 'O'},
          {'resid': 1, 'resname': 'ALA', 'atomname': 'CB', 'atomid': 5, 'element': 'C'},
          {'resid': 1, 'resname': 'ALA', 'atomname': 'H', 'atomid': 6, 'element': 'H'},
          {'resid': 1, 'resname': 'ALA', 'atomname': 'HA', 'atomid': 7, 'element': 'H'},
          {'resid': 1, 'resname': 'ALA', 'atomname': 'HB1', 'atomid': 8, 'element': 'H'},
          {'resid': 1, 'resname': 'ALA', 'atomname': 'HB2', 'atomid': 9, 'element': 'H'},
          {'resid': 1, 'resname': 'ALA', 'atomname': 'HB3', 'atomid': 10, 'element': 'H'},
          {'resid': 2, 'resname': 'VAL', 'atomname': 'N', 'atomid': 11, 'element': 'N'},
          {'resid': 2, 'resname': 'VAL', 'atomname': 'CA', 'atomid': 12, 'element': 'C'},
          {'resid': 2, 'resname': 'VAL', 'atomname': 'C', 'atomid': 13, 'element': 'C'},
          {'resid': 2, 'resname': 'VAL', 'atomname': 'O', 'atomid': 14, 'element': 'O'},
          {'resid': 2, 'resname': 'VAL', 'atomname': 'CB', 'atomid': 15, 'element': 'C'},
          {'resid': 2, 'resname': 'VAL', 'atomname': 'CG1', 'atomid': 16, 'element': 'C'},
          {'resid': 2, 'resname': 'VAL', 'atomname': 'CG2', 'atomid': 17, 'element': 'C'},
          {'resid': 2, 'resname': 'VAL', 'atomname': 'H', 'atomid': 18, 'element': 'H'},
          {'resid': 2, 'resname': 'VAL', 'atomname': 'HA', 'atomid': 19, 'element': 'H'},
          {'resid': 2, 'resname': 'VAL', 'atomname': 'HB', 'atomid': 20, 'element': 'H'},
          {'resid': 2, 'resname': 'VAL', 'atomname': 'HG11', 'atomid': 21, 'element': 'H'},
          {'resid': 2, 'resname': 'VAL', 'atomname': 'HG12', 'atomid': 22, 'element': 'H'},
          {'resid': 2, 'resname': 'VAL', 'atomname': 'HG13', 'atomid': 23, 'element': 'H'},
          {'resid': 2, 'resname': 'VAL', 'atomname': 'HG21', 'atomid': 24, 'element': 'H'},
          {'resid': 2, 'resname': 'VAL', 'atomname': 'HG22', 'atomid': 25, 'element': 'H'},
          {'resid': 2, 'resname': 'VAL', 'atomname': 'HG23', 'atomid': 26, 'element': 'H'},
          {'resid': 3, 'resname': 'SOL', 'atomname': 'OW', 'atomid': 27, 'element': 'O'},
          {'resid': 3, 'resname': 'SOL', 'atomname': 'HW1', 'atomid': 28, 'element': 'H'},
          {'resid': 3, 'resname': 'SOL', 'atomname': 'HW2', 'atomid': 29, 'element': 'H'},
          {'resid': 4, 'resname': 'SOL', 'atomname': 'OW', 'atomid': 30, 'element': 'O'},
          {'resid': 4, 'resname': 'SOL', 'atomname': 'HW1', 'atomid': 31, 'element': 'H'},
          {'resid': 4, 'resname': 'SOL', 'atomname': 'HW2', 'atomid': 32, 'element': 'H'},
          {'resid': 5, 'resname': 'SOL', 'atomname': 'OW', 'atomid': 33, 'element': 'O'},
          {'resid': 5, 'resname': 'SOL', 'atomname': 'HW1', 'atomid': 34, 'element': 'H'},
          {'resid': 5, 'resname': 'SOL', 'atomname': 'HW2', 'atomid': 35, 'element': 'H'},
          {'resid': 6, 'resname': 'SOL', 'atomname': 'OW', 'atomid': 36, 'element': 'O'},
          {'resid': 6, 'resname': 'SOL', 'atomname': 'HW1', 'atomid': 37, 'element': 'H'},
          {'resid': 6, 'resname': 'SOL', 'atomname': 'HW2', 'atomid': 38, 'element': 'H'},
          {'resid': 7, 'resname': 'SOL', 'atomname': 'OW', 'atomid': 39, 'element': 'O'},
          {'resid': 7, 'resname': 'SOL', 'atomname': 'HW1', 'atomid': 40, 'element': 'H'},
          {'resid': 7, 'resname': 'SOL', 'atomname': 'HW2', 'atomid': 41, 'element': 'H'},
          {'resid': 8, 'resname': 'SOL', 'atomname': 'OW', 'atomid': 42, 'element': 'O'},
          {'resid': 8, 'resname': 'SOL', 'atomname': 'HW1', 'atomid': 43, 'element': 'H'},
          {'resid': 8, 'resname': 'SOL', 'atomname': 'HW2', 'atomid': 44, 'element': 'H'},
      )
      for node, coords, vels in zip(nodes, COORDINATES, VELOCITIES):
          node['chain'] = ''
          node['position'] = np.array(coords)
          if velocities:
              node['velocity'] = np.array(vels)
      molecule = Molecule()
      molecule.add_nodes_from(enumerate(nodes))
      return molecule
  
  
  def filter_molecule(molecule, exclude, ignh):
      """
      Remove nodes from a graph based on some criteria.
  
      Nodes are removed if their resname matches one from
      the `exclude` argument, or of `ighn` is `True` and the node is a
      hydrogen.
  
      The atoms are renumbered after the filtered atoms are removed. This will
      scramble the edges and the interactions! This function is meant to be used
      on molecules with ONLY nodes.
  
      Parameters
      ----------
      molecule: nx.Graph
          The molecule to filter. The molecule is modified in place.
      exclude: list
          List of residue name to exclude.
      ighn: bool
          If `True`, the hydrogens are excluded.
      """
      to_remove = []
      exclusions = [('resname', value) for value in exclude]
      if ignh:
          exclusions.append(('element', 'H'))
      for node_key, node in molecule.nodes.items():
          for key, value in exclusions:
              if key in node and node[key] == value:
                  to_remove.append(node_key)
                  break
      molecule.remove_nodes_from(to_remove)
      mapping = collections.OrderedDict(
          [(old, new) for new, old in enumerate(molecule.nodes)]
      )
      nx.relabel_nodes(molecule, mapping, copy=False)
      # Networkx mangles the order of the nodes, but we want the nodes sorted.
      new_nodes = sorted(molecule.nodes.items())
      molecule.clear()  # Remove all the nodes
      molecule.add_nodes_from(new_nodes)
  
  
  def compare_dicts(testee, reference):
      """
      Report differences between dictionaries.
      """
      report = []
      dict_equals = testee.keys() == reference.keys()
      diff_keys = []
      if dict_equals:
          for key in testee:
              if are_different(testee[key], reference[key]):
                  diff_keys.append(key)
      if diff_keys:
          report.append('The following keys differ: {}'.format(diff_keys))
      if diff_keys or not dict_equals:
          report.append('+ {}'.format(testee))
          report.append('- {}'.format(reference))
      return report
  
  
  def assert_molecule_equal(molecule, reference):
      """
      Make a pytest test fail with a report if the molecules are not equal.
      """
      report = []
      # Nodes
      if len(molecule.nodes) != len(reference.nodes):
          report.append(
              'Different number of nodes: {} != {}'
              .format(len(molecule.nodes), len(reference.nodes))
          )
      zipped_nodes = zip(molecule.nodes.items(), reference.nodes.items())
      for (key_mol, node_mol), (key_ref, node_ref) in zipped_nodes:
          partial_report = []
          failed = False
          if key_mol != key_ref:
              partial_report.append('Node key {} != {}'.format(key_mol, key_ref))
              failed = True
          partial_report = compare_dicts(node_mol, node_ref)
          if partial_report:
              if not failed:
                  report.append('Node {}'.format(key_mol))
              report.extend(partial_report)
          elif failed:
              report.append('= {}'.format(node_mol))
  
      extra_molecule_nodes = molecule.nodes.keys() - reference.nodes.keys()
      if extra_molecule_nodes:
          report.append('The following nodes are not in the reference: {}.'
                        .format(extra_molecule_nodes))
      extra_reference_nodes = reference.nodes.keys() - molecule.nodes.keys()
      if extra_reference_nodes:
          report.append('The following nodes are not in the molecule: {}.'
                        .format(extra_reference_nodes))
      # Edges
      # The order of the edges does not matter.
      molecule_edges = set(molecule.edges)
      reference_edges = set(reference.edges)
      extra_molecule_edges = molecule_edges - reference_edges
      extra_reference_edges = reference_edges - molecule_edges
      if extra_molecule_edges:
          report.append('The following edges are not in the reference:')
          report.append('{}'.format(extra_molecule_edges))
      if extra_reference_edges:
          report.append('The following edges are not in the molecule:')
          report.append('{}'.format(extra_reference_edges))
      # Meta data
      report.extend(compare_dicts(molecule.meta, reference.meta))
      if molecule._force_field is not reference._force_field:  # pylint: disable=protected-access
          report.append('Force fields do not match: {} != {}'
                        .format(molecule._force_field, reference._force_field))  # pylint: disable=protected-access
      if molecule.nrexcl != reference.nrexcl:
          report.append('nrexcl: {} != {}'.format(molecule.nrexcl, reference.nrexcl))
      # Interactions
      if report:
          pytest.fail('\n'.join(report))
  
  
  @st.composite
  def generate_dict(draw, min_size=0):
      """
      Strategy to generate an arbitrary dictionary.
      """
      keys = st.one_of(st.text(), st.integers())
      values = st.one_of(
          st.text(), st.integers(), st.floats(),
          hnp.arrays(dtype=hnp.scalar_dtypes(), shape=hnp.array_shapes())
      )
      dict_a = draw(st.dictionaries(keys, values, min_size=min_size, max_size=10))
      return dict_a
  
  
  @st.composite
  def generate_equal_dict(draw):
      """
      Strategy to generate two equal dictionaries.
      """
      dict_a = draw(generate_dict())
      dict_b = copy.deepcopy(dict_a)
      return dict_a, dict_b
  
  
  
  @st.composite
  def generate_diff_dict(draw):
      """
      Strategy to generate to similar but different dictionaries.
      """
      dict_a = draw(generate_dict(min_size=1))
      dict_b = copy.deepcopy(dict_a)
      values = st.one_of(
          st.text(), st.integers(), st.floats(),
          hnp.arrays(dtype=hnp.scalar_dtypes(), shape=hnp.array_shapes())
      )
      num_to_change = draw(st.integers(
          min_value=1, max_value=(len(dict_b) // 2) + 1
      ))
      for _ in range(num_to_change):
          key = draw(st.sampled_from(list(dict_b.keys())))
          new_val = draw(values)
          assume(are_different(new_val, dict_a[key]))
          dict_b[key] = new_val
      dict_c = draw(generate_dict())
      for key in set(dict_c.keys()) - set(dict_b.keys()):
          dict_b[key] = dict_c[key]
      return dict_a, dict_b
  
  
  @given(generate_equal_dict())
  def test_compare_dict_equal(dict_a_and_b):
      """
      Test that :func:`compare_dicts` identify equal dictionaries.
      """
      dict_a = dict_a_and_b[0]
      dict_b = dict_a_and_b[1]
      assert not compare_dicts(dict_a, dict_b)  # report is empty
  
  
  @given(generate_diff_dict())
  def test_compare_dict_diff(dict_a_and_b):
      """
      Test that :func:`compare_dicts` identify different dictionaries.
      """
      dict_a = dict_a_and_b[0]
      dict_b = dict_a_and_b[1]
      assert compare_dicts(dict_a, dict_b)  # report is not empty
  
  
  def test_filter_molecule_none(gro_reference):  # pylint: disable=redefined-outer-name
      """
      Test that :func:`filter_molecule` does not remove nodes when not asked for.
      """
      _, molecule = gro_reference
      filter_molecule(molecule, exclude=(), ignh=False)
      assert len(molecule.nodes) == 44
  
  
  def test_filter_molecule_ignh(gro_reference):  # pylint: disable=redefined-outer-name
      """
      Test that :func:`filter_molecule` works with `ighn` argument.
      """
      _, molecule = gro_reference
      filter_molecule(molecule, exclude=(), ignh=True)
      elements = [node['element'] for node in molecule.nodes.values()]
      assert len(molecule.nodes) == (44 - 26)
      assert 'H' not in elements
  
  
  @pytest.mark.parametrize('exclude', (
      ('SOL', ), ('ALA', ), ('ALA', 'VAL'),
      ('ALA', 'VAL', 'SOL'), ('XXX', ), ('ALA', 'XXX'),
  ))
  def test_filter_molecule_exclude(gro_reference, exclude):  # pylint: disable=redefined-outer-name
      """
      Test that :func:`filter_molecule` works with `exclude` argument.
      """
      _, molecule = gro_reference
      filter_molecule(molecule, exclude=exclude, ignh=False)
      resnames = [node['resname'] for node in molecule.nodes.values()]
      should_not_be_there = [exclusion
                             for exclusion in exclude
                             if exclusion in resnames]
      assert not should_not_be_there
  
  
  def test_filter_molecule_identity(gro_reference):  # pylint: disable=redefined-outer-name
      """
      Test that :func:`filter_molecule` modifies the molecule in place.
      """
      _, molecule = gro_reference
      id_before = id(molecule)
      filter_molecule(molecule, exclude=('SOL', ), ignh=True)
      assert id(molecule) == id_before
  
  
  def test_filter_molecule_order(gro_reference):  # pylint: disable=redefined-outer-name
      """
      Test that :func:`filter_molecule` returns the nodes in order.
      """
      _, molecule = gro_reference
      filter_molecule(molecule, exclude=('SOL', ), ignh=True)
      keys = list(molecule.nodes)
      assert  sorted(keys) == keys
  
  
  @pytest.mark.parametrize('exclude', (
      (), ('SOL', ), ('ALA', ), ('ALA', 'VAL'),
      ('ALA', 'VAL', 'SOL'), ('XXX', ), ('ALA', 'XXX'),
  ))
  @pytest.mark.parametrize('ignh', (True, False))
  def test_read_gro(gro_reference, exclude, ignh):  # pylint: disable=redefined-outer-name
      """
      Test the GRO reader.
      """
      filename, reference = gro_reference
      filter_molecule(reference, exclude=exclude, ignh=ignh)
      molecule = gro.read_gro(filename, exclude=exclude, ignh=ignh)
      pprint(list(molecule.nodes.items()))
      assert_molecule_equal(molecule, reference)
  
  
  def test_read_gro_wrong_atom_number(gro_wrong_length):  # pylint: disable=redefined-outer-name
      """
      Test that the GRO reader raises an exception if the number of atoms is not
      consistent.
      """
      with pytest.raises(ValueError):
          gro.read_gro(gro_wrong_length)
  
  
  def test_write_gro(gro_reference, tmpdir):
      """
      Test writing GRO file.
      """
      filename, molecule = gro_reference
      system = vermouth.System()
      system.molecules.append(molecule)
      outname = tmpdir / 'out_test.gro'
      gro.write_gro(
          system,
          outname,
          box=(10.0, 11.1, 12.2),
          title='Just a title',
      )
      DeferredFileWriter().write()
      with open(str(filename)) as ref, open(str(outname)) as out:
          assert out.read() == ref.read()
  