# Source code for MDAnalysis.visualization.streamlines

    # -*- 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
    #

    """Streamplots (2D)  --- :mod:`MDAnalysis.visualization.streamlines`
    =================================================================

    :Authors: Tyler Reddy and Matthieu Chavent
    :Year: 2014
    :Copyright: Lesser GNU Public License v2.1+

    The :func:`generate_streamlines` function can generate a 2D flow field from a
    MD trajectory, for instance, lipid molecules in a flat membrane. It can make
    use of multiple cores to perform the analyis in parallel (using
    :mod:`multiprocessing`).

    See Also
    --------
    MDAnalysis.visualization.streamlines_3D : streamplots in 3D

    .. autofunction:: generate_streamlines

    """
    import multiprocessing

    import numpy as np
    import scipy

    try:
        import matplotlib
        import matplotlib.path
    except ImportError:
        raise ImportError(
            "2d streamplot module requires: matplotlib.path for its "
            "path.Path.contains_points method. The installation "
            "instructions for the matplotlib module can be found here: "
            "http://matplotlib.org/faq/installing_faq.html?highlight=install"
        ) from None

    import MDAnalysis

    def produce_grid(tuple_of_limits, grid_spacing):
        """Produce a 2D grid for the simulation system.

        The grid is based on the tuple of Cartesian Coordinate limits calculated in
        an earlier step.

        Parameters
        ----------
        tuple_of_limits : tuple
            ``x_min, x_max, y_min, y_max``
        grid_spacing : float
            grid size in all directions in ångström

        Returns
        -------
        grid : array
           ``numpy.mgrid[x_min:x_max:grid_spacing, y_min:y_max:grid_spacing]``
        """
        x_min, x_max, y_min, y_max = tuple_of_limits
        grid = np.mgrid[x_min:x_max:grid_spacing, y_min:y_max:grid_spacing]
        return grid

    def split_grid(grid, num_cores):
        """Split the grid into blocks of vertices.

        Take the overall `grid` for the system and split it into lists of
        square vertices that can be distributed to each core.

        Parameters
        ----------
        grid : numpy.array
            2D array
        num_cores : int
            number of partitions to generate

        Returns
        -------
        list_square_vertex_arrays_per_core : array of arrays
             split the list of square vertices
             ``[[v1,v2,v3,v4],[v1,v2,v3,v4],...,...]`` into roughly equally-sized
             sublists to be distributed over the available cores on the system
        list_parent_index_values : array of arrays
             arrays of `[[row, column], [row, column], ...]`` for each core
        current_row : int
             last row + 1
        current_column : int
             last column + 1

        Note
        ----
        Limited to 2D for now.

        """

        # produce an array containing the cartesian coordinates of all vertices in the grid:
        x_array, y_array = grid
        grid_vertex_cartesian_array = np.dstack((x_array, y_array))
        # the grid_vertex_cartesian_array has N_rows, with each row corresponding to a column of coordinates in the grid (
        # so a given row has shape N_rows, 2); overall shape (N_columns_in_grid, N_rows_in_a_column, 2)
        # although I'll eventually want a pure numpy/scipy/vector-based solution, for now I'll allow loops to simplify the
        #  division of the cartesian coordinates into a list of the squares in the grid
        list_all_squares_in_grid = (
            []
        )  # should eventually be a nested list of all the square vertices in the grid/system
        list_parent_index_values = (
            []
        )  # want an ordered list of assignment indices for reconstructing the grid positions
        # in the parent process
        current_column = 0
        while current_column < grid_vertex_cartesian_array.shape[0] - 1:
            # go through all the columns except the last one and account for the square vertices (the last column
            #  has no 'right neighbour')
            current_row = 0
            while current_row < grid_vertex_cartesian_array.shape[1] - 1:
                # all rows except the top row, which doesn't have a row above it for forming squares
                bottom_left_vertex_current_square = grid_vertex_cartesian_array[
                    current_column, current_row
                ]
                bottom_right_vertex_current_square = grid_vertex_cartesian_array[
                    current_column + 1, current_row
                ]
                top_right_vertex_current_square = grid_vertex_cartesian_array[
                    current_column + 1, current_row + 1
                ]
                top_left_vertex_current_square = grid_vertex_cartesian_array[
                    current_column, current_row + 1
                ]
                # append the vertices of this square to the overall list of square vertices:
                list_all_squares_in_grid.append(
                    [
                        bottom_left_vertex_current_square,
                        bottom_right_vertex_current_square,
                        top_right_vertex_current_square,
                        top_left_vertex_current_square,
                    ]
                )
                list_parent_index_values.append([current_row, current_column])
                current_row += 1
            current_column += 1
        # split the list of square vertices [[v1,v2,v3,v4],[v1,v2,v3,v4],...,...] into roughly equally-sized sublists to
        # be distributed over the available cores on the system:
        list_square_vertex_arrays_per_core = np.array_split(
            list_all_squares_in_grid, num_cores
        )
        list_parent_index_values = np.array_split(
            list_parent_index_values, num_cores
        )
        return [
            list_square_vertex_arrays_per_core,
            list_parent_index_values,
            current_row,
            current_column,
        ]

    # some private utility functions for trajectory iteration
    def _produce_list_indices_point_in_polygon_this_frame(
        vertex_coord_list, relevant_particle_coordinate_array_xy
    ):
        """
        Produce a list of indices for particles of interest in the current frame.
        The list is ordered according to the order of the squares in the grid.
        Each entry in the list is a tuple of indices for the particles in that square.
        The list is the same length as the number of squares in the grid.
        """
        list_indices_point_in_polygon = []
        for square_vertices in vertex_coord_list:
            path_object = matplotlib.path.Path(square_vertices)
            index_list_in_polygon = np.where(
                path_object.contains_points(relevant_particle_coordinate_array_xy)
            )
            list_indices_point_in_polygon.append(index_list_in_polygon)
        return list_indices_point_in_polygon

    def _produce_list_centroids_this_frame(
        list_indices_in_polygon, relevant_particle_coordinate_array_xy
    ):
        """
        Produce a list of centroids for the particles in the current frame.
        The list is ordered according to the order of the squares in the grid.
        Each entry in the list is a numpy array of the centroid coordinates for the particles in that square.
        The list is the same length as the number of squares in the grid.
        If there are no particles in a square, the entry is None.
        """
        list_centroids_this_frame = []
        for indices in list_indices_in_polygon:
            if (
                not indices[0].size > 0
            ):  # if there are no particles of interest in this particular square
                list_centroids_this_frame.append(None)
            else:
                current_coordinate_array_in_square = (
                    relevant_particle_coordinate_array_xy[indices]
                )
                current_square_indices_centroid = np.average(
                    current_coordinate_array_in_square, axis=0
                )
                list_centroids_this_frame.append(current_square_indices_centroid)
        return list_centroids_this_frame  # a list of numpy xy centroid arrays for this frame

    def per_core_work(
        topology_file_path,
        trajectory_file_path,
        list_square_vertex_arrays_this_core,
        MDA_selection,
        start_frame,
        end_frame,
        reconstruction_index_list,
        maximum_delta_magnitude,
    ):
        """Run the analysis on one core.

        The code to perform on a given core given the list of square vertices assigned to it.
        """
        # obtain the relevant coordinates for particles of interest
        universe_object = MDAnalysis.Universe(
            topology_file_path, trajectory_file_path
        )
        list_previous_frame_centroids = []
        list_previous_frame_indices = []
        for ts in universe_object.trajectory:
            if ts.frame < start_frame:  # don't start until first specified frame
                continue
            relevant_particle_coordinate_array_xy = universe_object.select_atoms(
                MDA_selection
            ).positions[..., :-1]
            # only 2D / xy coords for now
            # I will need a list of indices for relevant particles falling within each square in THIS frame:
            list_indices_in_squares_this_frame = (
                _produce_list_indices_point_in_polygon_this_frame(
                    list_square_vertex_arrays_this_core,
                    relevant_particle_coordinate_array_xy,
                )
            )
            # likewise, I will need a list of centroids of particles in each square (same order as above list):
            list_centroids_in_squares_this_frame = (
                _produce_list_centroids_this_frame(
                    list_indices_in_squares_this_frame,
                    relevant_particle_coordinate_array_xy,
                )
            )
            if (
                list_previous_frame_indices
            ):  # if the previous frame had indices in at least one square I will need to use
                #  those indices to generate the updates to the corresponding centroids in this frame:
                list_centroids_this_frame_using_indices_from_last_frame = (
                    _produce_list_centroids_this_frame(
                        list_previous_frame_indices,
                        relevant_particle_coordinate_array_xy,
                    )
                )
                # I need to write a velocity of zero if there are any 'empty' squares in either frame:
                xy_deltas_to_write = []
                for square_1_centroid, square_2_centroid in zip(
                    list_centroids_this_frame_using_indices_from_last_frame,
                    list_previous_frame_centroids,
                ):
                    if square_1_centroid is None or square_2_centroid is None:
                        xy_deltas_to_write.append([0, 0])
                    else:
                        xy_deltas_to_write.append(
                            np.subtract(
                                square_1_centroid, square_2_centroid
                            ).tolist()
                        )

                # xy_deltas_to_write = np.subtract(np.array(
                # list_centroids_this_frame_using_indices_from_last_frame),np.array(list_previous_frame_centroids))
                xy_deltas_to_write = np.array(xy_deltas_to_write)
                # now filter the array to only contain distances in the range [-8,8] as a placeholder for dealing with PBC
                #  issues (Matthieu seemed to use a limit of 8 as well);
                xy_deltas_to_write = np.clip(
                    xy_deltas_to_write,
                    -maximum_delta_magnitude,
                    maximum_delta_magnitude,
                )

                # with the xy and dx,dy values calculated I need to set the values from this frame to previous frame
                # values in anticipation of the next frame:
                list_previous_frame_centroids = (
                    list_centroids_in_squares_this_frame[:]
                )
                list_previous_frame_indices = list_indices_in_squares_this_frame[:]
            else:  # either no points in squares or after the first frame I'll just reset the 'previous' values so they
                # can be used when consecutive frames have proper values
                list_previous_frame_centroids = (
                    list_centroids_in_squares_this_frame[:]
                )
                list_previous_frame_indices = list_indices_in_squares_this_frame[:]
            if ts.frame > end_frame:
                break  # stop here
        return list(zip(reconstruction_index_list, xy_deltas_to_write.tolist()))

    def generate_streamlines(
        topology_file_path,
        trajectory_file_path,
        grid_spacing,
        MDA_selection,
        start_frame,
        end_frame,
        xmin,
        xmax,
        ymin,
        ymax,
        maximum_delta_magnitude,
        num_cores="maximum",
    ):
        r"""Produce the x and y components of a 2D streamplot data set.

        Parameters
        ----------
        topology_file_path : str
                Absolute path to the topology file
        trajectory_file_path : str
                Absolute path to the trajectory file. It will normally be desirable
                to filter the trajectory with a tool such as GROMACS
                :program:`g_filter` (see :footcite:p:`Chavent2014`)
        grid_spacing : float
                The spacing between grid lines (angstroms)
        MDA_selection : str
                MDAnalysis selection string
        start_frame : int
                First frame number to parse
        end_frame : int
                Last frame number to parse
        xmin : float
                Minimum coordinate boundary for x-axis (angstroms)
        xmax : float
                Maximum coordinate boundary for x-axis (angstroms)
        ymin : float
                Minimum coordinate boundary for y-axis (angstroms)
        ymax : float
                Maximum coordinate boundary for y-axis (angstroms)
        maximum_delta_magnitude : float
                Absolute value of the largest displacement tolerated for the
                centroid of a group of particles ( angstroms). Values above this
                displacement will not count in the streamplot (treated as
                excessively large displacements crossing the periodic boundary)
        num_cores : int or 'maximum' (optional)
                The number of cores to use. (Default 'maximum' uses all available
                cores)

        Returns
        -------
        dx_array : array of floats
                An array object containing the displacements in the x direction
        dy_array : array of floats
                An array object containing the displacements in the y direction
        average_displacement : float
                :math:`\frac{\sum\sqrt[]{dx^2 + dy^2}}{N}`
        standard_deviation_of_displacement : float
                standard deviation of :math:`\sqrt[]{dx^2 + dy^2}`

        Examples
        --------
        Generate 2D streamlines and plot::

            import matplotlib, matplotlib.pyplot, np
            import MDAnalysis, MDAnalysis.visualization.streamlines

            u1, v1, average_displacement, standard_deviation_of_displacement =
                MDAnalysis.visualization.streamlines.generate_streamlines('testing.gro', 'testing_filtered.xtc',
                        grid_spacing=20, MDA_selection='name PO4', start_frame=2, end_frame=3,
                        xmin=-8.73000049591, xmax= 1225.96008301,
                        ymin= -12.5799999237, ymax=1224.34008789,
                        maximum_delta_magnitude=1.0, num_cores=16)
            x = np.linspace(0, 1200, 61)
            y = np.linspace(0, 1200, 61)
            speed = np.sqrt(u1*u1 + v1*v1)
            fig = matplotlib.pyplot.figure()
            ax = fig.add_subplot(111, aspect='equal')
            ax.set_xlabel('x ($\AA$)')
            ax.set_ylabel('y ($\AA$)')
            ax.streamplot(x, y, u1, v1, density=(10,10), color=speed, linewidth=3*speed/speed.max())
            fig.savefig('testing_streamline.png',dpi=300)

        .. image:: testing_streamline.png

        References
        ----------
        .. footbibliography::

        See Also
        --------
        MDAnalysis.visualization.streamlines_3D.generate_streamlines_3d

        """
        # work out the number of cores to use:
        if num_cores == "maximum":
            num_cores = multiprocessing.cpu_count()  # use all available cores
        else:
            num_cores = num_cores  # use the value specified by the user
            # assert isinstance(num_cores,(int,long)), "The number of specified cores must (of course) be an integer."
        np.seterr(all="warn", over="raise")
        parent_list_deltas = []  # collect all data from child processes here

        def log_result_to_parent(delta_array):
            parent_list_deltas.extend(delta_array)

        tuple_of_limits = (xmin, xmax, ymin, ymax)
        grid = produce_grid(
            tuple_of_limits=tuple_of_limits, grid_spacing=grid_spacing
        )
        (
            list_square_vertex_arrays_per_core,
            list_parent_index_values,
            total_rows,
            total_columns,
        ) = split_grid(grid=grid, num_cores=num_cores)
        pool = multiprocessing.Pool(num_cores)
        for vertex_sublist, index_sublist in zip(
            list_square_vertex_arrays_per_core, list_parent_index_values
        ):
            pool.apply_async(
                per_core_work,
                args=(
                    topology_file_path,
                    trajectory_file_path,
                    vertex_sublist,
                    MDA_selection,
                    start_frame,
                    end_frame,
                    index_sublist,
                    maximum_delta_magnitude,
                ),
                callback=log_result_to_parent,
            )
        pool.close()
        pool.join()
        dx_array = np.zeros((total_rows, total_columns))
        dy_array = np.zeros((total_rows, total_columns))
        # the parent_list_deltas is shaped like this: [ ([row_index,column_index],[dx,dy]), ... (...),...,]
        for (
            index_array,
            delta_array,
        ) in (
            parent_list_deltas
        ):  # go through the list in the parent process and assign to the
            #  appropriate positions in the dx and dy matrices:
            # build in a filter to replace all values at the cap (currently between -8,8) with 0 to match Matthieu's code
            # (I think eventually we'll reduce the cap to a narrower boundary though)
            index_1 = index_array.tolist()[0]
            index_2 = index_array.tolist()[1]
            if abs(delta_array[0]) == maximum_delta_magnitude:
                dx_array[index_1, index_2] = 0
            else:
                dx_array[index_1, index_2] = delta_array[0]
            if abs(delta_array[1]) == maximum_delta_magnitude:
                dy_array[index_1, index_2] = 0
            else:
                dy_array[index_1, index_2] = delta_array[1]

        # at Matthieu's request, we now want to calculate the average and standard deviation of the displacement values:
        displacement_array = np.sqrt(dx_array**2 + dy_array**2)
        average_displacement = np.average(displacement_array)
        standard_deviation_of_displacement = np.std(displacement_array)

        return (
            dx_array,
            dy_array,
            average_displacement,
            standard_deviation_of_displacement,
        )