# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*-
# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4
#
# 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
#
import numbers
from multiprocessing.sharedctypes import SynchronizedArray
from multiprocessing import Process, Manager
from joblib import cpu_count
import numpy as np
import sys

import MDAnalysis as mda
from ...coordinates.memory import MemoryReader


class TriangularMatrix(object):
    """Triangular matrix class. This class is designed to provide a
    memory-efficient representation of a triangular matrix that still behaves
    as a square symmetric one. The class wraps a numpy.array object,
    in which data are memorized in row-major order. It also has few additional
    facilities to conveniently load/write a matrix from/to file. It can be
    accessed using the [] and () operators, similarly to a normal numpy array.

    """

    def __init__(self, size, metadata=None, loadfile=None):
        """Class constructor.

        Parameters
        ----------

        size : int / array_like
            Size of the matrix (number of rows or columns). If an
            array is provided instead, the size of the triangular matrix
            will be calculated and the array copied as the matrix
            elements. Otherwise, the matrix is just initialized to zero.
        metadata : dict or None
            Metadata dictionary. Used to generate the metadata attribute.
        loadfile : str or None
            Load the matrix from this file. All the attributes and data will
            be determined by the matrix file itself (i.e. metadata will be
            ignored); size has to be provided though.

        """
        if isinstance(metadata, dict):
            self.metadata = np.array(metadata.items(), dtype=object)
        else:
            self.metadata = metadata

        self.size = size
        if loadfile:
            self.loadz(loadfile)
        elif isinstance(size, numbers.Integral):
            self.size = size
            self._elements = np.zeros((size + 1) * size // 2, dtype=np.float64)
        elif isinstance(size, SynchronizedArray):
            self._elements = np.array(size.get_obj(), dtype=np.float64)
            self.size = int((np.sqrt(1 + 8 * len(size)) - 1) / 2)
        elif isinstance(size, np.ndarray):
            self._elements = size
            self.size = int((np.sqrt(1 + 8 * len(size)) - 1) / 2)
        else:
            raise TypeError

    def __getitem__(self, args):
        x, y = args
        if x < y:
            x, y = y, x
        return self._elements[x * (x + 1) // 2 + y]

    def __setitem__(self, args, val):
        x, y = args
        if x < y:
            x, y = y, x
        self._elements[x * (x + 1) // 2 + y] = val

    def as_array(self):
        """Return standard numpy array equivalent"""
        # pylint: disable=unsubscriptable-object
        a = np.zeros((self.size, self.size))
        a[np.tril_indices(self.size)] = self._elements
        a[np.triu_indices(self.size)] = a.T[np.triu_indices(self.size)]
        return a

    def savez(self, fname):
        """Save matrix in the npz compressed numpy format. Save metadata and
        data as well.

        Parameters
        ----------

        fname : str
            Name of the file to be saved.
        """
        np.savez(fname, elements=self._elements, metadata=self.metadata)

    def loadz(self, fname):
        """Load matrix from the npz compressed numpy format.

        Parameters
        ----------

        fname : str
            Name of the file to be loaded.
        """
        loaded = np.load(fname, allow_pickle=True)

        if loaded["metadata"].shape != ():
            if loaded["metadata"]["number of frames"] != self.size:
                raise TypeError
            self.metadata = loaded["metadata"]
        else:
            if self.size * (self.size - 1) / 2 + self.size != len(
                loaded["elements"]
            ):
                raise TypeError
        self._elements = loaded["elements"]

    def __add__(self, scalar):
        """Add scalar to matrix elements.

        Parameters
        ----------

        scalar : float
            Scalar to be added.
        """
        newMatrix = self.__class__(self.size)
        newMatrix._elements = self._elements + scalar
        return newMatrix

    def __iadd__(self, scalar):
        """Add scalar to matrix elements.

        Parameters
        ----------

        scalar : float
            Scalar to be added.
        """
        self._elements += scalar
        return self

    def __mul__(self, scalar):
        """Multiply with scalar.

        Parameters
        ----------

        scalar : float
            Scalar to multiply with.
        """
        newMatrix = self.__class__(self.size)
        newMatrix._elements = self._elements * scalar
        return newMatrix

    def __imul__(self, scalar):
        """Multiply with scalar.

        Parameters
        ----------

        scalar : float
            Scalar to multiply with.
        """
        self._elements *= scalar
        return self

    __rmul__ = __mul__

    def __str__(self):
        return str(self.as_array())


class ParallelCalculation(object):
    r"""
    Generic parallel calculation class. Can use arbitrary functions,
    arguments to functions and kwargs to functions.

    Attributes
    ----------
    n_jobs : int
        Number of cores to be used for parallel calculation. If -1 use all
        available cores.
    function : callable object
        Function to be run in parallel.
    args : list of tuples
        Each tuple contains the arguments that will be passed to
        function(). This means that a call to function() is performed for
        each tuple. function is called as function(\*args, \*\*kwargs). Runs
        are distributed on the requested numbers of cores.
    kwargs : list of dicts
        Each tuple contains the named arguments that will be passed to
        function, similarly as described for the args attribute.
    nruns : int
        Number of runs to be performed. Must be equal to len(args) and
        len(kwargs).
    """

    def __init__(self, n_jobs, function, args=None, kwargs=None):
        """
        Parameters
        ----------
        n_jobs : int
            Number of cores to be used for parallel calculation. If -1 use all
            available cores.
        function : object that supports __call__, as functions
            function to be run in parallel.
        args : list of tuples
            Arguments for function; see the ParallelCalculation class
            description.
        kwargs : list of dicts or None
            kwargs for function; see the ParallelCalculation
            class description.
        """

        # args[i] should be a list of args, one for each run
        self.n_jobs = n_jobs
        if self.n_jobs == -1:
            self.n_jobs = cpu_count()

        self.functions = function
        if not hasattr(self.functions, "__iter__"):
            self.functions = [self.functions] * len(args)
        if len(self.functions) != len(args):
            self.functions = self.functions[:] * (
                len(args) // len(self.functions)
            )

        # Arguments should be present
        if args is None:
            args = []
        self.args = args

        # If kwargs are not present, use empty dicts
        if kwargs:
            self.kwargs = kwargs
        else:
            self.kwargs = [{} for i in self.args]

        self.nruns = len(args)

    def worker(self, q, results):
        """
        Generic worker. Will run function with the prescribed args and kwargs.

        Parameters
        ----------

        q : multiprocessing.Manager.Queue object
                work queue, from which the worker fetches arguments and
                messages

        results : multiprocessing.Manager.Queue object
                results queue, where results are put after each calculation is
                finished

        """
        while True:
            i = q.get()
            if i == "STOP":
                return

            results.put(
                (i, self.functions[i](*self.args[i], **self.kwargs[i]))
            )

    def run(self):
        r"""
        Run parallel calculation.

        Returns
        -------

        results : tuple of ordered tuples (int, object)
                int is the number of the calculation corresponding to a
                certain argument in the args list, and object is the result of
                corresponding calculation. For instance, in (3, output), output
                is the return of function(\*args[3], \*\*kwargs[3]).
        """
        results_list = []
        if self.n_jobs == 1:
            for i in range(self.nruns):
                results_list.append(
                    (i, self.functions[i](*self.args[i], **self.kwargs[i]))
                )
        else:
            manager = Manager()
            q = manager.Queue()
            results = manager.Queue()

            workers = [
                Process(target=self.worker, args=(q, results))
                for i in range(self.n_jobs)
            ]

            for i in range(self.nruns):
                q.put(i)
            for w in workers:
                q.put("STOP")

            for w in workers:
                w.start()

            for w in workers:
                w.join()

            results.put("STOP")
            for i in iter(results.get, "STOP"):
                results_list.append(i)

        return tuple(sorted(results_list, key=lambda x: x[0]))


def trm_indices(a, b):
    """
    Generate (i,j) indeces of a triangular matrix, between elements a and b.
    The matrix size is automatically determined from the number of elements.
    For instance: trm_indices((0,0),(2,1)) yields (0,0) (1,0) (1,1) (2,0)
    (2,1).

    Parameters
    ----------

    a : (int i, int j) tuple
        starting matrix element.

    b : (int i, int j) tuple
        final matrix element.
    """
    i, j = a
    while i < b[0]:
        if i == j:
            yield (i, j)
            j = 0
            i += 1
        else:
            yield (i, j)
            j += 1
    while j <= b[1]:
        yield (i, j)
        j += 1


def trm_indices_nodiag(n):
    """generate (i,j) indeces of a triangular matrix of n rows (or columns),
    without diagonal (e.g. no elements (0,0),(1,1),...,(n,n))

    Parameters
    ----------

    n : int
        Matrix size
    """

    for i in range(1, n):
        for j in range(i):
            yield (i, j)


def trm_indices_diag(n):
    """generate (i,j) indeces of a triangular matrix of n rows (or columns),
    with diagonal

    Parameters
    ----------

    n : int
        Matrix size
    """

    for i in range(0, n):
        for j in range(i + 1):
            yield (i, j)


def merge_universes(universes):
    """
    Merge list of universes into one

    Parameters
    ----------
    universes : list of Universe objects


    Returns
    ----------
    Universe object
    """

    for universe in universes:
        universe.transfer_to_memory()

    return mda.Universe(
        universes[0].filename,
        np.concatenate(
            tuple([e.trajectory.timeseries(order="fac") for e in universes]),
            axis=0,
        ),
        format=MemoryReader,
    )