# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*-
# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4 fileencoding=utf-8
#
# MDAnalysis --- https://www.mdanalysis.org
# Copyright (c) 2006-2017 The MDAnalysis Development Team and contributors
# (see the file AUTHORS for the full list of names)
#
# Released under the Lesser GNU Public Licence, v2.1 or any higher version
#
# Please cite your use of MDAnalysis in published work:
#
# R. J. Gowers, M. Linke, J. Barnoud, T. J. E. Reddy, M. N. Melo, S. L. Seyler,
# D. L. Dotson, J. Domanski, S. Buchoux, I. M. Kenney, and O. Beckstein.
# MDAnalysis: A Python package for the rapid analysis of molecular dynamics
# simulations. In S. Benthall and S. Rostrup editors, Proceedings of the 15th
# Python in Science Conference, pages 102-109, Austin, TX, 2016. SciPy.
# doi: 10.25080/majora-629e541a-00e
#
# N. Michaud-Agrawal, E. J. Denning, T. B. Woolf, and O. Beckstein.
# MDAnalysis: A Toolkit for the Analysis of Molecular Dynamics Simulations.
# J. Comput. Chem. 32 (2011), 2319--2327, doi:10.1002/jcc.21787
#
from collections import OrderedDict

import numpy as np
import pytest
from MDAnalysis.coordinates.base import (
    ReaderBase,
    SingleFrameReaderBase,
    StreamReaderBase,
    Timestep,
)
from numpy.testing import assert_allclose, assert_equal

"""
Isolate the API definitions of Readers independent of implementations
"""


class AmazingMultiFrameReader(ReaderBase):
    format = "AmazingMulti"

    def __init__(self, filename, **kwargs):
        self.filename = filename
        self.n_frames = 10
        self.n_atoms = 10
        self._auxs = {}
        self._transformations = []
        # ts isn't a real timestep, but just an integer
        # whose value represents the frame number (0 based)
        self.ts = Timestep(self.n_atoms)
        self.ts.frame = -1
        self._read_next_timestep()

    def _read_next_timestep(self):
        self.ts.frame += 1
        if (self.ts.frame + 1) > self.n_frames:
            raise IOError
        else:
            return self.ts

    def _read_frame(self, frame):
        if frame < 0:
            frame = self.n_frames + frame
        if not (0 <= frame < self.n_frames):
            raise IOError
        self.ts.frame = frame

        return self.ts

    def _reopen(self):
        self.ts.frame = -1


class AmazingReader(SingleFrameReaderBase):
    format = "Amazing"

    # have to hack this in to get the base class to "work"
    def _read_first_frame(self):
        self.n_atoms = 10
        self.ts = Timestep(self.n_atoms)
        self.ts.frame = 0


class AmazingStreamReader(StreamReaderBase):
    format = "AmazingStream"

    def __init__(self, filename, n_atoms):
        self.n_atoms = n_atoms
        self._mocked_frames = [Timestep(n_atoms) for _ in range(3)]
        super().__init__(filename)

    def _read_frame(self, frame):
        self._frame = frame
        if self._frame >= len(self._mocked_frames):
            raise EOFError("End of stream")
        ts = self._mocked_frames[self._frame]
        ts.frame = self._frame
        self.ts = ts
        return ts


class _TestReader(object):
    __test__ = False
    """Basic API readers"""

    @pytest.fixture()
    def reader(self):
        return self.readerclass("test.txt")

    @pytest.fixture()
    def ts(self, reader):
        return reader.ts

    def test_required_attributes(self, reader):
        """Test that Reader has the required attributes"""
        for attr in [
            "filename",
            "n_atoms",
            "n_frames",
            "ts",
            "units",
            "format",
        ]:
            assert_equal(
                hasattr(reader, attr), True, "Missing attr: {0}".format(attr)
            )

    def test_iter(self, reader):
        l = [ts for ts in reader]

        assert_equal(len(l), self.n_frames)

    def test_close(self):
        sfr = self.readerclass("text.txt")

        ret = sfr.close()
        # Check that method works?
        assert_equal(ret, None)

    def test_rewind(self, reader):
        ret = reader.rewind()

        assert_equal(ret, None)
        assert_equal(reader.ts.frame, 0)

    def test_context(self):
        with self.readerclass("text.txt") as sfr:
            l = sfr.ts.frame

        assert_equal(l, 0)

    def test_len(self, reader):
        l = len(reader)

        assert_equal(l, self.n_frames)

    def test_raises_StopIteration(self, reader):
        reader[-1]
        with pytest.raises(StopIteration):
            next(reader)

    @pytest.mark.parametrize("order", ["turnip", "abc"])
    def test_timeseries_raises_unknown_order_key(self, reader, order):
        with pytest.raises(ValueError, match="Unrecognized order key"):
            reader.timeseries(order=order)

    @pytest.mark.parametrize("order", ["faac", "affc", "afcc", ""])
    def test_timeseries_raises_incorrect_order_key(self, reader, order):
        with pytest.raises(ValueError, match="Repeated or missing keys"):
            reader.timeseries(order=order)


class _Multi(_TestReader):
    n_frames = 10
    n_atoms = 10
    readerclass = AmazingMultiFrameReader
    reference = np.arange(10)


class TestMultiFrameReader(_Multi):

    __test__ = True

    @pytest.mark.parametrize(
        "start, stop, step",
        [
            (None, None, None),  # blank slice
            (None, 5, None),  # set end point
            (2, None, None),  # set start point
            (2, 5, None),  # start & end
            (None, None, 2),  # set skip
            (None, None, -1),  # backwards skip
            (None, -1, -1),
            (10, 0, -1),
            (0, 10, 1),
            (0, 10, 2),
            (None, 20, None),  # end beyond real end
            (None, 20, 2),  # with skip
            (0, 5, 2),
            (5, None, -1),
            (None, 5, -1),
            (100, 10, 1),
            (-10, None, 1),
            (100, None, -1),  # beyond real end
            (100, 5, -20),
            (5, 1, 1),  # Stop less than start
            (1, 5, -1),  # Stop less than start
            (-100, None, None),
            (100, None, None),  # Outside of range of trajectory
            (-2, 10, -2),
            (0, 0, 1),  # empty
            (10, 1, 2),  # empty
        ],
    )
    def test_slice(self, start, stop, step, reader):
        """Compare the slice applied to trajectory, to slice of list"""
        res = [ts.frame for ts in reader[start:stop:step]]
        ref = self.reference[start:stop:step]

        assert_equal(res, ref)

    def test_slice_VE_1(self, reader):
        def sl():
            return list(reader[::0])

        with pytest.raises(ValueError):
            sl()

    def test_slice_TE_1(self, reader):
        def sl():
            return list(reader[1.2:2.5:0.1])

        with pytest.raises(TypeError):
            sl()

    @pytest.mark.parametrize("slice_cls", [list, np.array])
    @pytest.mark.parametrize(
        "sl",
        [
            [0, 1, 4, 5],
            [5, 1, 6, 2, 7, 3, 8],
            [0, 1, 1, 1, 0, 0, 2, 3, 4],
            [True, False, True, False, True, False, True, False, True, False],
            [True, True, False, False, True, True, False, True, False, True],
            [True, True, True, True, True, True, True, True, True, True],
            [
                False,
                False,
                False,
                False,
                False,
                False,
                False,
                False,
                False,
                False,
            ],
        ],
    )
    def test_getitem(self, slice_cls, sl, reader):
        sl = slice_cls(sl)
        res = [ts.frame for ts in reader[sl]]

        sl = np.asarray(sl)
        ref = self.reference[sl]

        assert_equal(res, ref)

    @pytest.mark.parametrize(
        "sl",
        [
            [0, 1, 2, 3],  # ordered list of indices without duplicates
            [1, 3, 4, 2, 9],  # disordered list of indices without duplicates
            [0, 1, 1, 2, 2, 2],  # ordered list with duplicates
            [-1, -2, 3, -1, 0],  # disordered list with duplicates
            [True] * 10,
            [False] * 10,
            [True, False] * 5,
            slice(None, None, None),
            slice(0, 10, 1),
            slice(None, None, -1),
            slice(10, 0, -1),
            slice(2, 7, 2),
            slice(7, 2, -2),
            slice(7, 2, 1),  # empty
            slice(0, 0, 1),  # empty
        ],
    )
    def test_getitem_len(self, sl, reader):
        traj_iterable = reader[sl]
        if not isinstance(sl, slice):
            sl = np.array(sl)
        ref = self.reference[sl]
        assert len(traj_iterable) == len(ref)

    @pytest.mark.parametrize("iter_type", (list, np.array))
    def test_getitem_len_empty(self, reader, iter_type):
        # Indexing a numpy array with an empty array tends to break.
        traj_iterable = reader[iter_type([])]
        assert len(traj_iterable) == 0

    # All the sl1 slice must be 5 frames long so that the sl2 can be a mask
    @pytest.mark.parametrize(
        "sl1",
        [
            [0, 1, 2, 3, 4],
            [1, 1, 1, 1, 1],
            [True, False] * 5,
            slice(None, None, 2),
            slice(None, None, -2),
        ],
    )
    @pytest.mark.parametrize(
        "sl2",
        [
            [0, -1, 2],
            [-1, -1, -1],
            [True, False, True, True, False],
            np.array([True, False, True, True, False]),
            slice(None, None, None),
            slice(None, 3, None),
            slice(4, 0, -1),
            slice(None, None, -1),
            slice(None, None, 2),
        ],
    )
    def test_double_getitem(self, sl1, sl2, reader):
        traj_iterable = reader[sl1][sl2]
        # Old versions of numpy do not behave the same when indexing with a
        # list or with an array.
        if not isinstance(sl1, slice):
            sl1 = np.asarray(sl1)
        if not isinstance(sl2, slice):
            sl2 = np.asarray(sl2)
        print(sl1, sl2, type(sl1), type(sl2))
        ref = self.reference[sl1][sl2]
        res = [ts.frame for ts in traj_iterable]
        assert_equal(res, ref)
        assert len(traj_iterable) == len(ref)

    @pytest.mark.parametrize(
        "sl1",
        [
            [0, 1, 2, 3, 4],
            [1, 1, 1, 1, 1],
            [True, False] * 5,
            slice(None, None, 2),
            slice(None, None, -2),
            slice(None, None, None),
        ],
    )
    @pytest.mark.parametrize("idx2", [0, 2, 4, -1, -2, -4])
    def test_double_getitem_int(self, sl1, idx2, reader):
        ts = reader[sl1][idx2]
        # Old versions of numpy do not behave the same when indexing with a
        # list or with an array.
        if not isinstance(sl1, slice):
            sl1 = np.asarray(sl1)
        ref = self.reference[sl1][idx2]
        assert ts.frame == ref

    def test_list_TE(self, reader):
        def sl():
            return list(reader[[0, "a", 5, 6]])

        with pytest.raises(TypeError):
            sl()

    def test_array_TE(self, reader):
        def sl():
            return list(reader[np.array([1.2, 3.4, 5.6])])

        with pytest.raises(TypeError):
            sl()

    @pytest.mark.parametrize(
        "sl1",
        [
            [0, 1, 2, 3, 4],
            [1, 1, 1, 1, 1],
            [True, False] * 5,
            slice(None, None, 2),
            slice(None, None, -2),
        ],
    )
    @pytest.mark.parametrize("idx2", [5, -6])
    def test_getitem_IE(self, sl1, idx2, reader):
        partial_reader = reader[sl1]
        with pytest.raises(IndexError):
            partial_reader[idx2]


class _Single(_TestReader):
    n_frames = 1
    n_atoms = 10
    readerclass = AmazingReader


class TestSingleFrameReader(_Single):
    __test__ = True

    def test_next(self, reader):
        with pytest.raises(StopIteration):
            reader.next()

    # Getitem tests
    # only 0 & -1 should work
    # others should get IndexError
    def _check_get_results(self, l, ts):
        assert_equal(len(l), 1)
        assert_equal(ts in l, True)

    def test_getitem(self, reader, ts):
        fr = [reader[0]]

        self._check_get_results(fr, ts)

    def test_getitem_2(self, reader, ts):
        fr = [reader[-1]]

        self._check_get_results(fr, ts)

    def test_getitem_IE(self, reader):
        with pytest.raises(IndexError):
            reader.__getitem__(1)

    def test_getitem_IE_2(self, reader):
        with pytest.raises(IndexError):
            reader.__getitem__(-2)

    # Slicing should still work!
    def test_slice_1(self, reader, ts):
        l = list(reader[::])
        self._check_get_results(l, ts)

    def test_slice_2(self, reader, ts):
        l = list(reader[::-1])
        self._check_get_results(l, ts)

    def test_reopen(self, reader, ts):
        reader._reopen()
        assert_equal(ts.frame, 0)

    def test_rewind(self, reader, ts):
        reader.rewind()
        assert_equal(ts.frame, 0)

    def test_read_frame(self, reader):
        with pytest.raises(IndexError):
            reader._read_frame(1)

    def test_iter_rewind(self, reader):
        # Issue #3423
        # positions should be zero to start with
        assert_allclose(reader.ts.positions, np.zeros((10, 3)))

        # modify positions in place
        reader.ts.positions = np.ones((10, 3))
        assert_allclose(reader.ts.positions, np.ones((10, 3)))

        # iterate the reader
        for ts in reader:
            assert_allclose(ts.positions, np.zeros((10, 3)))

        assert_allclose(reader.ts.positions, np.zeros((10, 3)))


class _Stream:
    n_atoms = 3
    readerclass = AmazingStreamReader


class TestStreamReader(_Stream):
    @pytest.fixture
    def reader(self):
        return self.readerclass("dummy", n_atoms=self.n_atoms)

    def test_repr(self, reader):
        rep = repr(reader)
        assert "AmazingStreamReader" in rep
        assert "continuous stream" in rep
        assert "3 atoms" in rep

    def test_read_and_exhaust_stream(self, reader):
        ts0 = reader.next()
        ts1 = reader.next()
        ts2 = reader.next()
        assert ts0.frame == 0
        assert ts1.frame == 1
        assert ts2.frame == 2

        with pytest.raises(StopIteration):
            reader.next()

    def test_len_and_n_frames_raise(self, reader):
        with pytest.raises(RuntimeError):
            _ = len(reader)
        with pytest.raises(RuntimeError):
            _ = reader.n_frames

    def test_rewind_raises(self, reader):
        with pytest.raises(RuntimeError, match="can't be rewound"):
            reader.rewind()

    def test_copy_raises(self, reader):
        with pytest.raises(RuntimeError, match="does not support copying"):
            reader.copy()

    def test_timeseries_raises(self, reader):
        with pytest.raises(RuntimeError, match="cannot access timeseries"):
            reader.timeseries()

    def test_reopen_only_once(self, reader):
        reader._reopen()
        with pytest.raises(RuntimeError, match="Cannot reopen stream"):
            reader._reopen()

    def test_slice_reader(self, reader):
        sliced = reader[slice(None, None, 2)]
        with pytest.raises(RuntimeError, match="has unknown length"):
            len(sliced)
        with pytest.raises(RuntimeError, match="does not support indexing"):
            sliced[0]

        for i, ts in enumerate(sliced):
            assert ts.frame == i * 2

    def test_check_slice_index_errors(self, reader):
        with pytest.raises(ValueError, match="start.*must be None"):
            reader.check_slice_indices(0, None, 1)
        with pytest.raises(ValueError, match="stop.*must be None"):
            reader.check_slice_indices(None, 1, 1)
        with pytest.raises(ValueError, match="must be > 0"):
            reader.check_slice_indices(None, None, 0)
        with pytest.raises(ValueError, match="must be an integer"):
            reader.check_slice_indices(None, None, 1.5)

    def test_pickle_methods(self, reader):
        with pytest.raises(NotImplementedError):
            reader.__getstate__()
        with pytest.raises(NotImplementedError):
            reader.__setstate__({})