from io import StringIO
from collections import defaultdict

from numpy.testing import (
    assert_equal, assert_array_equal,)
import pytest

import MDAnalysis
from MDAnalysis.analysis.hydrogenbonds.wbridge_analysis import (
    WaterBridgeAnalysis, )


class TestWaterBridgeAnalysis(object):
    @staticmethod
    @pytest.fixture(scope='class')
    def universe_empty():
        '''A universe with no hydrogen bonds'''
        grofile = '''Test gro file
5
    1ALA      N    1   0.000   0.000   0.000
    1ALA      H    2   0.100   0.000   0.000
    2SOL     OW    3   3.000   0.000   0.000
    4ALA      H    4   0.500   0.000   0.000
    4ALA      N    5   0.600   0.000   0.000
  1.0   1.0   1.0'''
        u = MDAnalysis.Universe(StringIO(grofile), format='gro')
        return u

    @staticmethod
    @pytest.fixture(scope='class')
    def universe_DA():
        '''A universe with one hydrogen bond acceptor bonding to a hydrogen bond
        donor'''
        grofile = '''Test gro file
3
    1ALA      N    1   0.000   0.000   0.000
    1ALA      H    2   0.100   0.000   0.000
    4ALA      O    3   0.300   0.000   0.000
  1.0   1.0   1.0'''
        u = MDAnalysis.Universe(StringIO(grofile), format='gro')
        return u

    @staticmethod
    @pytest.fixture(scope='class')
    def universe_DA_PBC():
        '''A universe with one hydrogen bond acceptor bonding to a hydrogen bond
        donor but in a PBC condition'''
        grofile = '''Test gro file
3
    1ALA      N    1   0.800   0.000   0.000
    1ALA      H    2   0.900   0.000   0.000
    4ALA      O    3   0.100   0.000   0.000
  1.0   1.0   1.0'''
        u = MDAnalysis.Universe(StringIO(grofile), format='gro')
        return u

    @staticmethod
    @pytest.fixture(scope='class')
    def universe_AD():
        '''A universe with one hydrogen bond donor bonding to a hydrogen bond
        acceptor'''
        grofile = '''Test gro file
3
    1ALA      O    1   0.000   0.000   0.000
    4ALA      H    2   0.200   0.000   0.000
    4ALA      N    3   0.300   0.000   0.000
  1.0   1.0   1.0'''
        u = MDAnalysis.Universe(StringIO(grofile), format='gro')
        return u

    @staticmethod
    @pytest.fixture(scope='class')
    def universe_loop():
        '''A universe with one hydrogen bond acceptor bonding to a water which
        bonds back to the first hydrogen bond  acceptor and thus form a loop'''
        grofile = '''Test gro file
5
    1ALA      O    1   0.000   0.001   0.000
    2SOL     OW    2   0.300   0.001   0.000
    2SOL    HW1    3   0.200   0.002   0.000
    2SOL    HW2    4   0.200   0.000   0.000
    4ALA      O    5   0.600   0.000   0.000
  1.0   1.0   1.0'''
        u = MDAnalysis.Universe(StringIO(grofile), format='gro')
        return u

    @staticmethod
    @pytest.fixture(scope='class')
    def universe_DWA():
        '''A universe with one hydrogen bond donor bonding to a hydrogen bond
        acceptor through a water'''
        grofile = '''Test gro file
5
    1ALA      N    1   0.000   0.000   0.000
    1ALA      H    2   0.100   0.000   0.000
    2SOL     OW    3   0.300   0.000   0.000
    2SOL    HW2    4   0.400   0.000   0.000
    4ALA      O    5   0.600   0.000   0.000
  1.0   1.0   1.0'''
        u = MDAnalysis.Universe(StringIO(grofile), format='gro')
        return u

    @staticmethod
    @pytest.fixture(scope='class')
    def universe_DWD():
        '''A universe with one hydrogen bond donor bonding to a hydrogen bond
        donor through a water'''
        grofile = '''Test gro file
5
    1ALA      N    1   0.000   0.000   0.000
    1ALA      H    2   0.100   0.000   0.000
    2SOL     OW    3   0.300   0.000   0.000
    4ALA      H    4   0.500   0.000   0.000
    4ALA      N    5   0.600   0.000   0.000
  1.0   1.0   1.0'''
        u = MDAnalysis.Universe(StringIO(grofile), format='gro')
        return u

    @staticmethod
    @pytest.fixture(scope='class')
    def universe_AWA():
        '''A universe with two hydrogen bond acceptor are joined by a water'''
        grofile = '''Test gro file
5
    1ALA      O    1   0.000   0.000   0.000
    2SOL     OW    2   0.300   0.000   0.000
    2SOL    HW1    3   0.200   0.000   0.000
    2SOL    HW2    4   0.400   0.000   0.000
    4ALA      O    5   0.600   0.000   0.000
  1.0   1.0   1.0'''
        u = MDAnalysis.Universe(StringIO(grofile), format='gro')
        return u

    @staticmethod
    @pytest.fixture(scope='class')
    def universe_AWD():
        '''A universe with one hydrogen bond acceptor bonding to a hydrogen
        bond donor through a water'''
        grofile = '''Test gro file
5
    1ALA      O    1   0.000   0.000   0.000
    2SOL     OW    2   0.300   0.000   0.000
    2SOL    HW1    3   0.200   0.000   0.000
    4ALA      H    4   0.500   0.000   0.000
    4ALA      N    5   0.600   0.000   0.000
  1.0   1.0   1.0'''
        u = MDAnalysis.Universe(StringIO(grofile), format='gro')
        return u

    @staticmethod
    @pytest.fixture(scope='class')
    def universe_AWWA():
        '''A universe with one hydrogen bond acceptor bonding to a hydrogen bond
        acceptor through two waters'''
        grofile = '''Test gro file
7
    1ALA      O    1   0.000   0.000   0.000
    2SOL     OW    2   0.300   0.000   0.000
    2SOL    HW1    3   0.200   0.000   0.000
    2SOL    HW2    4   0.400   0.000   0.000
    3SOL     OW    5   0.600   0.000   0.000
    3SOL    HW1    6   0.700   0.000   0.000
    4ALA      O    7   0.900   0.000   0.000
  1.0   1.0   1.0'''
        u = MDAnalysis.Universe(StringIO(grofile), format='gro')
        return u

    @staticmethod
    @pytest.fixture(scope='class')
    def universe_AWWWA():
        '''A universe with one hydrogen bond acceptor bonding to a hydrogen bond
        acceptor through three waters'''
        grofile = '''Test gro file
9
    1ALA      O    1   0.000   0.000   0.000
    2SOL     OW    2   0.300   0.000   0.000
    2SOL    HW1    3   0.200   0.000   0.000
    2SOL    HW2    4   0.400   0.000   0.000
    3SOL     OW    5   0.600   0.000   0.000
    3SOL    HW1    6   0.700   0.000   0.000
    4SOL     OW    7   0.900   0.000   0.000
    4SOL    HW1    8   1.000   0.000   0.000
    5ALA      O    9   1.200   0.000   0.000
  10.0   10.0   10.0'''
        u = MDAnalysis.Universe(StringIO(grofile), format='gro')
        return u

    @staticmethod
    @pytest.fixture(scope='class')
    def universe_AWWWWA():
        '''A universe with one hydrogen bond acceptor bonding to a hydrogen bond
        acceptor through three waters'''
        grofile = '''Test gro file
11
    1ALA      O    1   0.000   0.000   0.000
    2SOL     OW    2   0.300   0.000   0.000
    2SOL    HW1    3   0.200   0.000   0.000
    2SOL    HW2    4   0.400   0.000   0.000
    3SOL     OW    5   0.600   0.000   0.000
    3SOL    HW1    6   0.700   0.000   0.000
    4SOL     OW    7   0.900   0.000   0.000
    4SOL    HW1    8   1.000   0.000   0.000
    5SOL     OW    9   1.200   0.000   0.000
    5SOL    HW1   10   1.300   0.000   0.000
    6ALA      O   11   1.400   0.000   0.000
  10.0   10.0   10.0'''
        u = MDAnalysis.Universe(StringIO(grofile), format='gro')
        return u

    @staticmethod
    @pytest.fixture(scope='class')
    def universe_branch():
        '''A universe with one hydrogen bond acceptor bonding to two hydrogen
        bond acceptor in selection 2'''
        grofile = '''Test gro file
9
    1ALA      O    1   0.000   0.000   0.000
    2SOL     OW    2   0.300   0.000   0.000
    2SOL    HW1    3   0.200   0.000   0.000
    2SOL    HW2    4   0.400   0.000   0.000
    3SOL     OW    5   0.600   0.000   0.000
    3SOL    HW1    6   0.700   0.000   0.000
    3SOL    HW2    7   0.600   0.100   0.000
    4ALA      O    8   0.900   0.000   0.000
    5ALA      O    9   0.600   0.300   0.000
  1.0   1.0   1.0'''
        u = MDAnalysis.Universe(StringIO(grofile), format='gro')
        return u

    @staticmethod
    @pytest.fixture(scope='class')
    def universe_AWA_AWWA():
        '''A universe with one hydrogen bond acceptors are bonded through one or
        two water'''
        grofile = '''Test gro file
12
    1ALA      O    1   0.000   0.000   0.000
    2SOL     OW    2   0.300   0.000   0.000
    2SOL    HW1    3   0.200   0.000   0.000
    2SOL    HW2    4   0.400   0.000   0.000
    4ALA      O    5   0.600   0.000   0.000
    5ALA      O    6   0.000   1.000   0.000
    6SOL     OW    7   0.300   1.000   0.000
    6SOL    HW1    8   0.200   1.000   0.000
    6SOL    HW2    9   0.400   1.000   0.000
    7SOL     OW   10   0.600   1.000   0.000
    7SOL    HW1   11   0.700   1.000   0.000
    8ALA      O   12   0.900   1.000   0.000
  1.0   1.0   1.0'''
        u = MDAnalysis.Universe(StringIO(grofile), format='gro')
        return u

    @staticmethod
    @pytest.fixture(scope='class')
    def wb_multiframe():
        '''A water bridge object with multipley frames'''
        grofile = '''Test gro file
 13
    1ALA      O    1   0.000   0.000   0.000
    1ALA      H    2   0.000   0.000   0.000
    2SOL     OW    3   0.300   0.000   0.000
    2SOL    HW1    4   0.200   0.000   0.000
    2SOL    HW2    5   0.400   0.000   0.000
    3SOL     OW    6   0.600   0.000   0.000
    3SOL    HW1    7   0.700   0.000   0.000
    4SOL     OW    8   0.900   0.000   0.000
    4SOL    HW1    9   1.000   0.000   0.000
    5SOL     OW   10   1.200   0.000   0.000
    5SOL    HW1   11   1.300   0.000   0.000
    6ALA      H   12   1.400   0.000   0.000
    6ALA      O   13   1.400   0.000   0.000
   10.0   10.0   10.0'''
        u = MDAnalysis.Universe(StringIO(grofile), format='gro')
        wb = WaterBridgeAnalysis(u, 'protein and (resid 1)', 'protein and (resid 4)',
                                 order=4)
        # Build an dummy WaterBridgeAnalysis object for testing
        wb.results.network = []
        wb.results.network.append({(1, 0, 12, None, 2.0, 180.0): None})
        wb.results.network.append({(0, None, 12, 13, 2.0, 180.0): None})
        wb.results.network.append({(1, 0, 3, None, 2.0, 180.0):
                            {(4, 2, 12, None, 2.0, 180.0): None}})
        wb.results.network.append({(0, None, 3, 2, 2.0, 180.0):
                            {(4, 2, 5, None, 2.0, 180.0):
                             {(5, None, 11, 12, 2.0, 180.0): None}}})
        wb.timesteps = range(len(wb.results.network))
        return wb

    def test_nodata(self, universe_DA):
        '''Test if the funtions can run when there is no data.
        This is achieved by not runing the run() first.'''
        wb = WaterBridgeAnalysis(universe_DA, 'protein and (resid 1)',
                                 'protein and (resid 4)', order=0)
        wb.generate_table()
        assert_equal(wb.timesteps_by_type(), None)
        assert_equal(wb.count_by_time(), None)
        assert_equal(wb.count_by_type(), None)

    def test_selection_type_error(self, universe_DA):
        '''Test the case when the wrong selection1_type is given'''
        try:
            wb = WaterBridgeAnalysis(universe_DA, 'protein and (resid 1)',
                'protein and (resid 4)', order=0, selection1_type='aaa')
        except ValueError:
            pass
        else:
            raise pytest.fail("selection_type aaa should rasie error")

    def test_empty_selection(self, universe_DA):
        '''Test the case when selection yields empty result'''
        wb = WaterBridgeAnalysis(universe_DA, 'protein and (resid 9)',
                                 'protein and (resid 10)', order=0)
        wb.run()
        assert wb.results.network == [{}]

    def test_loop(self, universe_loop):
        '''Test if loop can be handled correctly'''
        wb = WaterBridgeAnalysis(universe_loop, 'protein and (resid 1)',
                                 'protein and (resid 1 or resid 4)')
        wb.run()
        assert_equal(len(wb.results.network[0].keys()), 2)

    def test_donor_accepter(self, universe_DA):
        '''Test zeroth order donor to acceptor hydrogen bonding'''
        wb = WaterBridgeAnalysis(universe_DA, 'protein and (resid 1)',
        'protein and (resid 4)', order=0, update_selection=True, debug=True)
        wb.run(verbose=False)
        network = wb.results.network[0]
        assert_equal(list(network.keys())[0][:4], (1, 0, 2, None))

    def test_donor_accepter_heavy(self, universe_DA):
        '''Test zeroth order donor to acceptor hydrogen bonding'''
        wb = WaterBridgeAnalysis(universe_DA, 'protein and (resid 1)',
        'protein and (resid 4)', order=0, update_selection=True, debug=True, distance_type='heavy')
        wb.run(verbose=False)
        network = wb.results.network[0]
        assert_equal(list(network.keys())[0][:4], (1, 0, 2, None))

    def test_donor_accepter_pbc(self, universe_DA_PBC):
        '''Test zeroth order donor to acceptor hydrogen bonding in PBC conditions'''
        wb = WaterBridgeAnalysis(universe_DA_PBC, 'protein and (resid 1)',
            'protein and (resid 4)', order=0, pbc=True)
        wb.run(verbose=False)
        network = wb.results.network[0]
        assert_equal(list(network.keys())[0][:4], (1, 0, 2, None))

    def test_accepter_donor(self, universe_AD):
        '''Test zeroth order acceptor to donor hydrogen bonding'''
        wb = WaterBridgeAnalysis(universe_AD, 'protein and (resid 1)',
                                 'protein and (resid 4)', order=0)
        wb.run(verbose=False)
        network = wb.results.network[0]
        assert_equal(list(network.keys())[0][:4], (0, None, 1, 2))

    def test_acceptor_water_accepter(self, universe_AWA):
        '''Test case where the hydrogen bond acceptor from selection 1 form
        water bridge with hydrogen bond acceptor from selection 2'''
        wb = WaterBridgeAnalysis(universe_AWA, 'protein and (resid 1)',
                                 'protein and (resid 4)')
        wb.run(verbose=False)
        network = wb.results.network[0]
        assert_equal(list(network.keys())[0][:4], (0, None, 2, 1))
        second = network[list(network.keys())[0]]
        assert_equal(list(second.keys())[0][:4], (3, 1, 4, None))
        assert_equal(second[list(second.keys())[0]], None)

    def test_donor_water_accepter(self, universe_DWA):
        '''Test case where the hydrogen bond donor from selection 1 form
        water bridge with hydrogen bond acceptor from selection 2'''
        wb = WaterBridgeAnalysis(universe_DWA, 'protein and (resid 1)',
                                 'protein and (resid 4)')
        wb.run(verbose=False)
        network = wb.results.network[0]
        assert_equal(list(network.keys())[0][:4], (1, 0, 2, None))
        second = network[list(network.keys())[0]]
        assert_equal(list(second.keys())[0][:4], (3, 2, 4, None))
        assert_equal(second[list(second.keys())[0]], None)

    def test_acceptor_water_donor(self, universe_AWD):
        '''Test case where the hydrogen bond acceptor from selection 1 form
        water bridge with hydrogen bond donor from selection 2'''
        wb = WaterBridgeAnalysis(universe_AWD, 'protein and (resid 1)',
                                 'protein and (resid 4)')
        wb.run(verbose=False)
        network = wb.results.network[0]
        assert_equal(list(network.keys())[0][:4], (0, None, 2, 1))
        second = network[list(network.keys())[0]]
        assert_equal(list(second.keys())[0][:4], (1, None, 3, 4))
        assert_equal(second[list(second.keys())[0]], None)

    def test_donor_water_donor(self, universe_DWD):
        '''Test case where the hydrogen bond donor from selection 1 form
        water bridge with hydrogen bond donor from selection 2'''
        wb = WaterBridgeAnalysis(universe_DWD, 'protein and (resid 1)',
                                 'protein and (resid 4)')
        wb.run(verbose=False)
        network = wb.results.network[0]
        assert_equal(list(network.keys())[0][:4], (1, 0, 2, None))
        second = network[list(network.keys())[0]]
        assert_equal(list(second.keys())[0][:4], (2, None, 3, 4))
        assert_equal(second[list(second.keys())[0]], None)

    def test_empty(self, universe_empty):
        '''Test case where no water bridge exists'''
        wb = WaterBridgeAnalysis(universe_empty, 'protein', 'protein')
        wb.run(verbose=False)
        assert_equal(wb.results.network[0], defaultdict(dict))

    def test_same_selection(self, universe_DWA):
        '''
        This test tests that if the selection 1 and selection 2 are both protein.
        However, the protein only forms one hydrogen bond with the water.
        This entry won't be included.
        '''
        wb = WaterBridgeAnalysis(universe_DWA, 'protein and resid 1',
                                 'protein and resid 1')
        wb.run(verbose=False)
        assert_equal(wb.results.network[0], defaultdict(dict))

    def test_acceptor_2water_accepter(self, universe_AWWA):
        '''Test case where the hydrogen bond acceptor from selection 1 form second order
        water bridge with hydrogen bond acceptor from selection 2'''
        # test first order
        wb = WaterBridgeAnalysis(universe_AWWA, 'protein and (resid 1)',
                                 'protein and (resid 4)')
        wb.run(verbose=False)
        assert_equal(wb.results.network[0], defaultdict(dict))
        # test second order
        wb = WaterBridgeAnalysis(universe_AWWA, 'protein and (resid 1)',
                                'protein and (resid 4)', order=2)
        wb.run(verbose=False)
        network = wb.results.network[0]
        assert_equal(list(network.keys())[0][:4], (0, None, 2, 1))
        second = network[list(network.keys())[0]]
        assert_equal(list(second.keys())[0][:4], (3, 1, 4, None))
        third = second[list(second.keys())[0]]
        assert_equal(list(third.keys())[0][:4], (5, 4, 6, None))
        assert_equal(third[list(third.keys())[0]], None)
        # test third order
        wb = WaterBridgeAnalysis(universe_AWWA, 'protein and (resid 1)',
                                 'protein and (resid 4)', order=3)
        wb.run(verbose=False)
        network = wb.results.network[0]
        assert_equal(list(network.keys())[0][:4], (0, None, 2, 1))
        second = network[list(network.keys())[0]]
        assert_equal(list(second.keys())[0][:4], (3, 1, 4, None))
        third = second[list(second.keys())[0]]
        assert_equal(list(third.keys())[0][:4], (5, 4, 6, None))
        assert_equal(third[list(third.keys())[0]], None)

    def test_acceptor_3water_accepter(self, universe_AWWWA):
        '''Test case where the hydrogen bond acceptor from selection 1 form third order
        water bridge with hydrogen bond acceptor from selection 2'''
        wb = WaterBridgeAnalysis(universe_AWWWA, 'protein and (resid 1)',
                                 'protein and (resid 5)', order=2)
        wb.run(verbose=False)
        assert_equal(wb.results.network[0], defaultdict(dict))

        wb = WaterBridgeAnalysis(universe_AWWWA, 'protein and (resid 1)',
                                 'protein and (resid 5)', order=3)
        wb.run(verbose=False)
        network = wb.results.network[0]
        assert_equal(list(network.keys())[0][:4], (0, None, 2, 1))
        second = network[list(network.keys())[0]]
        assert_equal(list(second.keys())[0][:4], (3, 1, 4, None))
        third = second[list(second.keys())[0]]
        assert_equal(list(third.keys())[0][:4], (5, 4, 6, None))
        fourth = third[list(third.keys())[0]]
        assert_equal(list(fourth.keys())[0][:4], (7, 6, 8, None))
        assert_equal(fourth[list(fourth.keys())[0]], None)

        wb = WaterBridgeAnalysis(universe_AWWWA, 'protein and (resid 1)',
                                 'protein and (resid 5)', order=4)
        wb.run(verbose=False)
        network = wb.results.network[0]
        assert_equal(list(network.keys())[0][:4], (0, None, 2, 1))
        second = network[list(network.keys())[0]]
        assert_equal(list(second.keys())[0][:4], (3, 1, 4, None))
        third = second[list(second.keys())[0]]
        assert_equal(list(third.keys())[0][:4], (5, 4, 6, None))
        fourth = third[list(third.keys())[0]]
        assert_equal(list(fourth.keys())[0][:4], (7, 6, 8, None))
        assert_equal(fourth[list(fourth.keys())[0]], None)

    def test_acceptor_4water_accepter(self, universe_AWWWWA):
        '''Test case where the hydrogen bond acceptor from selection 1 form fourth order
        water bridge with hydrogen bond acceptor from selection 2'''
        wb = WaterBridgeAnalysis(universe_AWWWWA, 'protein and (resid 1)',
                                 'protein and (resid 6)', order=3)
        wb.run(verbose=False)
        assert_equal(wb.results.network[0], defaultdict(dict))

        wb = WaterBridgeAnalysis(universe_AWWWWA, 'protein and (resid 1)',
                                 'protein and (resid 6)', order=4)
        wb.run(verbose=False)
        network = wb.results.network[0]
        assert_equal(list(network.keys())[0][:4], (0, None, 2, 1))
        second = network[list(network.keys())[0]]
        assert_equal(list(second.keys())[0][:4], (3, 1, 4, None))
        third = second[list(second.keys())[0]]
        assert_equal(list(third.keys())[0][:4], (5, 4, 6, None))
        fourth = third[list(third.keys())[0]]
        assert_equal(list(fourth.keys())[0][:4], (7, 6, 8, None))
        fifth = fourth[list(fourth.keys())[0]]
        assert_equal(list(fifth.keys())[0][:4], (9, 8, 10, None))
        assert_equal(fifth[list(fifth.keys())[0]], None)

        wb = WaterBridgeAnalysis(universe_AWWWWA, 'protein and (resid 1)',
                                 'protein and (resid 6)', order=5)
        wb.run(verbose=False)
        network = wb.results.network[0]
        assert_equal(list(network.keys())[0][:4], (0, None, 2, 1))
        second = network[list(network.keys())[0]]
        assert_equal(list(second.keys())[0][:4], (3, 1, 4, None))
        third = second[list(second.keys())[0]]
        assert_equal(list(third.keys())[0][:4], (5, 4, 6, None))
        fourth = third[list(third.keys())[0]]
        assert_equal(list(fourth.keys())[0][:4], (7, 6, 8, None))
        fifth = fourth[list(fourth.keys())[0]]
        assert_equal(list(fifth.keys())[0][:4], (9, 8, 10, None))
        assert_equal(fifth[list(fifth.keys())[0]], None)

    def test_acceptor_22water_accepter(self, universe_branch):
        '''Test case where the hydrogen bond acceptor from selection 1 form a second order
        water bridge with hydrogen bond acceptor from selection 2
        and the last water is linked to two residues in selection 2'''
        wb = WaterBridgeAnalysis(universe_branch, 'protein and (resid 1)',
                                 'protein and (resid 4 or resid 5)', order=2)
        wb.run(verbose=False)
        network = wb.results.network[0]
        assert_equal(list(network.keys())[0][:4], (0, None, 2, 1))
        second = network[list(network.keys())[0]]
        assert_equal(list(second.keys())[0][:4], (3, 1, 4, None))
        third = second[list(second.keys())[0]]
        assert_equal([(5, 4, 7, None), (6, 4, 8, None)],
                     sorted([key[:4] for key in list(third.keys())]))

    def test_timeseries_wba(self, universe_branch):
        '''Test if the time series data is correctly generated in water bridge analysis format'''
        wb = WaterBridgeAnalysis(universe_branch, 'protein and (resid 1)',
                                 'protein and (resid 4 or resid 5)', order=2)
        wb.output_format = 'sele1_sele2'
        wb.run(verbose=False)
        timeseries = sorted(wb.results.timeseries[0])

        assert_equal(timeseries[0][:4], (0, 2, ('ALA', 1, 'O'), ('SOL', 2, 'HW1')))
        assert_equal(timeseries[1][:4], (3, 4, ('SOL', 2, 'HW2'), ('SOL', 3, 'OW')))
        assert_equal(timeseries[2][:4], (5, 7, ('SOL', 3, 'HW1'), ('ALA', 4, 'O')))
        assert_equal(timeseries[3][:4], (6, 8, ('SOL', 3, 'HW2'), ('ALA', 5, 'O')))

    def test_timeseries_hba(self, universe_branch):
        '''Test if the time series data is correctly generated in hydrogen bond analysis format'''
        wb = WaterBridgeAnalysis(universe_branch, 'protein and (resid 1)',
                                 'protein and (resid 4 or resid 5)', order=2)
        wb.output_format = 'donor_acceptor'
        wb.run(verbose=False)
        timeseries = sorted(wb.results.timeseries[0])

        assert_equal(timeseries[0][:4], (2, 0, ('SOL', 2, 'HW1'), ('ALA', 1, 'O')))
        assert_equal(timeseries[1][:4], (3, 4, ('SOL', 2, 'HW2'), ('SOL', 3, 'OW')))
        assert_equal(timeseries[2][:4], (5, 7, ('SOL', 3, 'HW1'), ('ALA', 4, 'O')))
        assert_equal(timeseries[3][:4], (6, 8, ('SOL', 3, 'HW2'), ('ALA', 5, 'O')))

    def test_acceptor_12water_accepter(self, universe_AWA_AWWA):
        '''Test of independent first order and second can be recognised correctely'''
        wb = WaterBridgeAnalysis(universe_AWA_AWWA, 'protein and (resid 1 or resid 5)',
                                 'protein and (resid 4 or resid 8)', order=1)
        wb.run(verbose=False)
        network = wb.results.network[0]
        assert_equal(list(network.keys())[0][:4], (0, None, 2, 1))
        second = network[list(network.keys())[0]]
        assert_equal(list(second.keys())[0][:4], (3, 1, 4, None))
        assert_equal(second[list(second.keys())[0]], None)
        network = wb.results.network[0]
        wb = WaterBridgeAnalysis(universe_AWA_AWWA, 'protein and (resid 1 or resid 5)',
                                 'protein and (resid 4 or resid 8)', order=2)
        wb.run(verbose=False)
        network = wb.results.network[0]
        assert_equal([(0, None, 2, 1), (5, None, 7, 6)],
                     sorted([key[:4] for key in list(network.keys())]))

    def test_count_by_type_single_link(self, universe_DWA):
        '''
        This test tests the simplest water bridge to see if count_by_type() works.
        '''
        wb = WaterBridgeAnalysis(universe_DWA, 'protein and (resid 1)',
                                 'protein and (resid 4)')
        wb.run(verbose=False)
        assert_equal(wb.count_by_type(), [(1, 4, 'ALA', 1, 'H', 'ALA', 4, 'O',  1.)])

    def test_count_by_type_multiple_link(self, universe_AWA_AWWA):
        '''
        This test tests if count_by_type() can give the correct result for more than 1 links.
        '''
        wb = WaterBridgeAnalysis(universe_AWA_AWWA, 'protein and (resid 1 or resid 5)',
                                 'protein and (resid 4 or resid 8)', order=2)
        wb.run(verbose=False)
        assert_equal(sorted(wb.count_by_type()),
        [[0, 4, 'ALA', 1, 'O', 'ALA', 4, 'O', 1.0],
         [5, 11, 'ALA', 5, 'O', 'ALA', 8, 'O', 1.0]])


    def test_count_by_type_multiple_frame(self, wb_multiframe):
        '''
        This test tests if count_by_type() works in multiply situations.
        :return:
        '''
        result = [[0, 11, 'ALA', 1, 'O', 'ALA', 6, 'H', 0.25],
                  [0, 12, 'ALA', 1, 'O', 'ALA', 6, 'O', 0.25],
                  [1, 12, 'ALA', 1, 'H', 'ALA', 6, 'O', 0.5]]

        assert_equal(sorted(wb_multiframe.count_by_type()), result)

    def test_count_by_type_filter(self, wb_multiframe):
        '''
        This test tests if modifying analysis_func
        allows some results to be filtered out in count_by_type().
        :return:
        '''
        def analysis(current, output, u):
            sele1_index, sele1_heavy_index, atom2, heavy_atom2, dist, angle = current[0]
            atom1, heavy_atom1, sele2_index, sele2_heavy_index, dist, angle = current[-1]
            sele1 = u.atoms[sele1_index]
            sele2 = u.atoms[sele2_index]
            (s1_resname, s1_resid, s1_name) = (sele1.resname, sele1.resid, sele1.name)
            (s2_resname, s2_resid, s2_name) = (sele2.resname, sele2.resid, sele2.name)

            key = (sele1_index, sele2_index, s1_resname, s1_resid, s1_name, s2_resname, s2_resid, s2_name)
            if s2_name == 'H':
                output[key] += 1
        result = [((0, 11, 'ALA', 1, 'O', 'ALA', 6, 'H'), 0.25)]
        assert_equal(sorted(wb_multiframe.count_by_type(analysis_func=analysis)), result)

    def test_count_by_type_merge(self, wb_multiframe):
        '''
        This test tests if modifying analysis_func
        allows some same residue to be merged in count_by_type().
        '''
        def analysis(current, output, u):
            sele1_index, sele1_heavy_index, atom2, heavy_atom2, dist, angle = current[0]
            atom1, heavy_atom1, sele2_index, sele2_heavy_index, dist, angle = current[-1]
            sele1 = u.atoms[sele1_index]
            sele2 = u.atoms[sele2_index]
            (s1_resname, s1_resid, s1_name) = (sele1.resname, sele1.resid, sele1.name)
            (s2_resname, s2_resid, s2_name) = (sele2.resname, sele2.resid, sele2.name)

            key = (s1_resname, s1_resid, s2_resname, s2_resid)
            output[key] = 1
        result = [(('ALA', 1, 'ALA', 6), 1.0)]
        assert_equal(sorted(wb_multiframe.count_by_type(analysis_func=analysis)), result)

    def test_count_by_type_order(self, wb_multiframe):
        '''
        This test tests if modifying analysis_func
        allows the order of water bridge to be separated in count_by_type().
        :return:
        '''
        def analysis(current, output, u):
            sele1_index, sele1_heavy_index, atom2, heavy_atom2, dist, angle = current[0]
            atom1, heavy_atom1, sele2_index, sele2_heavy_index, dist, angle = current[-1]
            sele1 = u.atoms[sele1_index]
            sele2 = u.atoms[sele2_index]
            (s1_resname, s1_resid, s1_name) = (sele1.resname, sele1.resid, sele1.name)
            (s2_resname, s2_resid, s2_name) = (sele2.resname, sele2.resid, sele2.name)
            key = (s1_resname, s1_resid, s2_resname, s2_resid, len(current)-1)
            output[key] = 1
        result = [(('ALA', 1, 'ALA', 6, 0), 0.5),
                  (('ALA', 1, 'ALA', 6, 1), 0.25),
                  (('ALA', 1, 'ALA', 6, 2), 0.25)]
        assert_equal(sorted(wb_multiframe.count_by_type(analysis_func=analysis)), result)

    def test_count_by_time(self, wb_multiframe):
        '''
        This test tests if count_by_times() works.
        :return:
        '''
        assert_equal(wb_multiframe.count_by_time(), [(0, 1), (1, 1), (2, 1), (3, 1)])


    def test_count_by_time_weight(self, universe_AWA_AWWA):
        '''
        This test tests if modyfing the analysis_func allows the weight to be changed
        in count_by_type().
        :return:
        '''
        wb = WaterBridgeAnalysis(universe_AWA_AWWA, 'protein and (resid 1 or resid 5)',
                                 'protein and (resid 4 or resid 8)', order=2)
        wb.run(verbose=False)
        def analysis(current, output, u):
            sele1_index, sele1_heavy_index, atom2, heavy_atom2, dist, angle = current[0]
            atom1, heavy_atom1, sele2_index, sele2_heavy_index, dist, angle = current[-1]
            sele1 = u.atoms[sele1_index]
            sele2 = u.atoms[sele2_index]
            (s1_resname, s1_resid, s1_name) = (sele1.resname, sele1.resid, sele1.name)
            (s2_resname, s2_resid, s2_name) = (sele2.resname, sele2.resid, sele2.name)
            key = (s1_resname, s1_resid, s2_resname, s2_resid)
            output[key] += len(current)-1
        assert_equal(wb.count_by_time(analysis_func=analysis), [(0,3), ])

    def test_count_by_time_empty(self, universe_AWA_AWWA):
        '''
        See if count_by_time() can handle zero well.
        :return:
        '''
        wb = WaterBridgeAnalysis(universe_AWA_AWWA, 'protein and (resid 1 or resid 5)',
                                 'protein and (resid 4 or resid 8)', order=2)
        wb.run(verbose=False)
        def analysis(current, output, u):
            pass
        assert_equal(wb.count_by_time(analysis_func=analysis), [(0,0), ])

    def test_generate_table_hba(self, wb_multiframe):
        '''Test generate table using hydrogen bond analysis format'''
        table = wb_multiframe.generate_table(output_format='donor_acceptor')
        assert_array_equal(
            sorted(table.donor_resid),
            [1, 1, 2, 2, 2, 6, 6],
        )

    def test_generate_table_s1s2(self, wb_multiframe):
        '''Test generate table using hydrogen bond analysis format'''
        table = wb_multiframe.generate_table(output_format='sele1_sele2')
        assert_array_equal(
            sorted(table.sele1_resid),
            [1, 1, 1, 1, 2, 2, 3],
        )

    def test_timesteps_by_type(self, wb_multiframe):
        '''Test the timesteps_by_type function'''

        timesteps = sorted(wb_multiframe.timesteps_by_type())
        assert_array_equal(timesteps[3], [1, 12, 'ALA', 1, 'H', 'ALA', 6, 'O', 0, 2])

    def test_duplicate_water(self):
        '''A case #3119 where
        Acceptor···H−O···H-Donor
                     |
                     H···O-H
        will be recognised as 3rd order water bridge.
        '''
        grofile = '''Test gro file
    7
    1LEU      O    1   1.876   0.810   1.354
  117SOL    HW1    2   1.853   0.831   1.162
  117SOL     OW    3   1.877   0.890   1.081
  117SOL    HW2    4   1.908   0.828   1.007
  135SOL     OW    5   1.924   0.713   0.845
    1LEU      H    6   1.997   0.991   1.194
    1LEU      N    7   2.041   1.030   1.274
   2.22092   2.22092   2.22092'''
        u = MDAnalysis.Universe(StringIO(grofile), format='gro')
        wb = WaterBridgeAnalysis(u, 'resname LEU and name O',
                                 'resname LEU and name N H', order=4)
        wb.run()
        assert len(wb.results.timeseries[0]) == 2

    def test_warn_results_deprecated(self, universe_DA):
        wb = WaterBridgeAnalysis(universe_DA, 'protein and (resid 9)',
                                 'protein and (resid 10)', order=0)
        wb.run()

        wmsg = "The `network` attribute was deprecated in MDAnalysis 2.0.0"
        with pytest.warns(DeprecationWarning, match=wmsg):
            assert_equal(wb.network, wb.results.network)

        wmsg = "The `timeseries` attribute was deprecated in MDAnalysis 2.0.0"
        with pytest.warns(DeprecationWarning, match=wmsg):
            assert_equal(wb.timeseries, wb.results.timeseries)
