import sys

import numpy as np
import scipy.sparse as sp
import torch


from pymde import constraints
from pymde import problem
from pymde import util
from pymde.functions import penalties, losses
from pymde.preprocess import _graph


__this_module = sys.modules[__name__]


LOGGER = problem.LOGGER


def _to_edges(graph):
    if not sp.issparse(graph):
        n = graph.shape[0]
        edges = torch.triu_indices(n, n, 1).T
        distances = graph[edges]
        return edges, distances

    # discard the lower triangular part, which is redundant
    graph = sp.triu(graph, format="coo")
    edges = torch.stack(
        [
            torch.tensor(graph.row, dtype=torch.long),
            torch.tensor(graph.col, dtype=torch.long),
        ],
        dim=1,
    )
    data = torch.tensor(graph.data, dtype=torch.float)
    return edges, data


def _validate_adjacency_matrix(matrix):
    diagonal_nonzero = matrix.diagonal() > 0
    if diagonal_nonzero.any():
        raise ValueError(
            "Adjacency matrices must not contain self edges; "
            "the following nodes were found to have self edges: ",
            np.argwhere(diagonal_nonzero).flatten(),
        )


def _to_graph(edges, distances=None, n_items=None):
    "Distances for repeated edges are summed."
    if isinstance(edges, torch.Tensor):
        edges = edges.cpu().numpy()

    if distances is None:
        distances = np.ones(edges.shape[0], dtype=np.float32)
    elif isinstance(distances, torch.Tensor):
        distances = distances.cpu().float().numpy()

    flip_idx = edges[:, 0] > edges[:, 1]
    edges[flip_idx] = np.stack(
        [edges[flip_idx][:, 1], edges[flip_idx][:, 0]], axis=1
    )

    if n_items is None:
        n_items = edges.max() + 1
    rows = edges[:, 0]
    cols = edges[:, 1]
    graph = sp.coo_matrix((distances, (rows, cols)), shape=(n_items, n_items))
    graph = graph + graph.T
    return Graph(graph.tocsr())


class Graph(object):
    """A weighted graph.

    This class represents a weighted graph. It is backed by a scipy.sparse
    adjacency matrix, and can be constructed from either a dense
    distance matrix, a sparse adjacency matrix, or torch.Tensor's
    of edges and weights (using the `from_edges` static method).

    It is an error for the graph to contain self-edges (i.e., non-zero
    values on the diagonal).

    This class is accepted as an argument to various preprocessing functions.
    """

    def __init__(self, adjacency_matrix):
        if isinstance(adjacency_matrix, np.ndarray):
            adjacency_matrix = sp.csr_matrix(adjacency_matrix)
        elif isinstance(adjacency_matrix, torch.Tensor):
            adjacency_matrix = sp.csr_matrix(adjacency_matrix.cpu().numpy())
        elif sp.issparse(adjacency_matrix) and not isinstance(
            adjacency_matrix, sp.csr_matrix
        ):
            LOGGER.warning(
                "The adjacency matrix for a Graph object should be "
                "a CSR matrix; converting your matrix to a CSR "
                "matrix, which may be costly ..."
            )
            adjacency_matrix = adjacency_matrix.tocsr()

        unreachable = adjacency_matrix.data == np.inf
        adjacency_matrix.data[unreachable] = 0
        adjacency_matrix.eliminate_zeros()
        _validate_adjacency_matrix(adjacency_matrix)

        self._adjacency_matrix = adjacency_matrix
        self._edges = None
        self._distances = None

    @staticmethod
    def from_edges(edges, weights=None, n_items=None):
        """Construct a graph from edges and weights.

        Arguments
        ---------
        edges: torch.Tensor
            Tensor of edges, of shape ``(n_edges, 2)``, with each
            edge represented by a row, i.e. by two integers.
        weights: torch.Tensor, optional
            Tensor of weights, of shape ``(n_edges,)``, with each weight
            a float.
        n_items: int, optional
            The number of items in the graph; if ``None``, this is
            taken to be the maximum value in ``edges`` plus 1.

        Returns
        -------
        pymde.Graph
            A Graph object representing ``edges`` and ``weights``.
        """
        return _to_graph(edges, weights, n_items)

    @property
    def adjacency_matrix(self):
        """The scipy.sparse adjacency matrix backing the graph."""
        return self._adjacency_matrix

    @property
    def A(self):
        return self.adjacency_matrix

    @A.setter
    def A(self, value):
        self.adjacency_matrix = value

    @property
    def edges(self):
        """A torch.Tensor of the edges in the graph."""
        if self._edges is None:
            edges, distances = _to_edges(self.A)
            self._edges = edges
            self._distances = distances
        return self._edges

    @property
    def distances(self):
        """The distances associated with each edge."""
        if self._distances is None:
            edges, distances = _to_edges(self.A)
            self._edges = edges
            self._distances = distances
        return self._distances

    @property
    def weights(self):
        return self.distances

    @property
    def n_items(self):
        """The number of items in the graph."""
        return self.A.shape[0]

    @property
    def n_edges(self):
        """The number of edges in the graph."""
        return self.edges.shape[0]

    @property
    def n_all_edges(self):
        """n_items choose 2."""
        return self.n_items * (self.n_items - 1) // 2

    def neighbors(self, node: int) -> np.ndarray:
        """The indices of the neighbors of ``node``."""
        return self.A.indices[self.A.indptr[node] : self.A.indptr[node + 1]]

    def neighbor_distances(self, node) -> np.ndarray:
        """The distances associated with the edges connected to ``node``."""
        return self.A.data[self.A.indptr[node] : self.A.indptr[node + 1]]

    def __getitem__(self, key):
        return self.A[key]

    def __setitem__(self, key, value):
        raise AttributeError("Graph objects are immutable.")

    def draw(
        self, embedding_dim=2, standardized=False, device="cpu", verbose=False
    ):
        """Draw a graph in the Cartesian plane.

        This method does some basic preprocessing, constructs an MDE problem
        that is often suitable for drawing graphs, and computes/returns an
        embedding by approximately solving the MDE problem.

        Arguments
        ---------
        embedding_dim: int
            The number of dimemsions, 1, 2, or 3.
        standardized: bool
            Whether to impose a standardization constraint.
        device: str
            Device on which to compute/store embedding, 'cpu' or 'cuda'.
        verbose: bool
            Whether to print verbose output.

        Returns
        -------
        torch.Tensor
            The embedding, of shape ``(n_items, embedding_dim)``
        """
        if (self.distances < 0).any():
            raise ValueError(
                "Graphs with negative edge weights cannot be drawn."
            )

        if self.n_edges < self.n_all_edges:
            # copmute shortest paths, but no more than 1e7 of them ...
            retain_fraction = min(1.0, 1e7 / self.n_all_edges)
            distance_graph = shortest_paths(
                self, retain_fraction=retain_fraction, verbose=verbose
            )
        else:
            distance_graph = self

        if not standardized:
            constraint = constraints.Centered()
            f = losses.WeightedQuadratic(distance_graph.distances)
        else:
            constraint = constraints.Standardized()
            # TODO(akshayka) better weights
            f = penalties.Cubic(1 / distance_graph.distances)
        mde = problem.MDE(
            n_items=self.n_items,
            embedding_dim=embedding_dim,
            edges=distance_graph.edges,
            distortion_function=f,
            constraint=constraint,
            device=device,
        )
        X = mde.embed(verbose=verbose)
        mde.plot(edges=self.edges)
        return X


def scale(graph, natural_length):
    """Scale graph weights to have RMS equal to `natural_length`

    Returns a new graph, leaving the old graph unmodified.

    Arguments
    ---------
    graph: pymde.Graph
        The graph whose distances to scale.
    natural_length: float
        Target RMS value.

    Returns
    -------
    pymde.Graph
        A new graph with scaled distances.
    """
    rms = graph.distances.pow(2).mean().sqrt()
    alpha = natural_length / rms
    distances = alpha * graph.distances
    return _to_graph(graph.edges, distances.cpu().numpy())


def _sparsify(graph, retain_fraction):
    raise NotImplementedError


def breadth_first_order(
    csgraph, i_start, directed=True, return_predecessors=True
):
    N = csgraph.shape[0]
    node_list = np.empty(N, dtype=np.int32)
    lengths = np.empty(N, dtype=np.int32)
    predecessors = np.empty(N, dtype=np.int32)
    node_list.fill(-9999)
    lengths.fill(-9999)
    predecessors.fill(-9999)

    if directed:
        _graph._breadth_first_directed(
            i_start,
            csgraph.indices,
            csgraph.indptr,
            node_list,
            lengths,
            predecessors,
        )
    lengths = lengths.astype(np.float32)
    lengths[lengths < 0] = np.inf
    return lengths, predecessors


def _shortest_paths(shape, node, max_length, unweighted, sample_prob):
    data = __this_module.__data
    indptr = __this_module.__indptr
    indices = __this_module.__indices

    data, indptr, indices = map(np.ctypeslib.as_array, [data, indptr, indices])
    A = sp.csr_matrix((data, indices, indptr), shape=shape)
    if unweighted:
        # fast path
        distances, _ = breadth_first_order(
            A,
            node,
        )
        distances[distances > max_length] = np.inf
    else:
        distances = sp.csgraph.dijkstra(A, indices=node, limit=max_length)

    indices = np.argwhere((distances != np.inf) * (distances > 0)).ravel()
    indices = indices[indices > node]
    if sample_prob < 1.0:
        sample_mask = util.np_rng().uniform(size=indices.size) <= sample_prob
        indices = indices[sample_mask]

    distances = distances[indices]
    return distances, indices


def __init_process(data, indptr, indices):
    global __this_module
    __this_module.__data = data
    __this_module.__indptr = indptr
    __this_module.__indices = indices


def shortest_paths(
    graph, max_length=None, retain_fraction=1.0, n_workers=None, verbose=False
):
    """Compute shortest-path distances.

    This function computes the shortest-path distances on a graph.
    It returns a new graph, with an edge for each distance that was
    computed (each edge is weighted by the shortest-path distance between
    the two nodes in the original graph). The new graph can be interpreted
    as a "filled-in" version of the input graph.

    The ``max_length`` and ``retain_fraction`` keyword arguments can be used
    to compute only a subset of the distances. This can be useful
    if the graph has a very large number of nodes, in which storing all graph
    distances may be intractable.

    If the graph is sufficiently large, multiple cores will be used to
    accelerate the computation.

    Arguments
    ---------
    graph: pymde.Graph or scipy.sparse matrix
        Graph instance, or an adjacency matrix
    max_length: float, optional
        The maximum-length path to compute; paths longer than max_length
        are not computed/explored.
    retain_fraction: float
        A float between 0 and 1, specifying the fraction of all (n choose 2)
        to compute. For example, if ``retain_fraction`` is 0.1, only 10
        percent of the edges will be stored.
    n_workers: int, optional
        The number of processes to use. Defaults to the number of available
        cores.
    verbose: bool
        Whether to print verbose output.

    Returns
    -------
    pymde.Graph
        A new graph, with an edge (and weight) for each shortest-path
        distance that was computed/stored.
    """
    if sp.issparse(graph):
        graph = Graph(graph)
    elif not isinstance(graph, Graph):
        raise ValueError(
            "`graph` must be a pymde.Graph instance or "
            "scipy.sparse adjacency matrix."
        )

    if verbose:
        LOGGER.info(
            f"Computing shortest path distances (retaining "
            f"{100*retain_fraction:.2f} percent with "
            f"max_distance={max_length}) ..."
        )

    if n_workers is None and max_length is None and retain_fraction == 1.0:
        # obtain a dense distance matrix
        distance_matrix = sp.csgraph.shortest_path(graph.A, directed=False)
        return Graph(sp.csr_matrix(distance_matrix))
    elif max_length is None:
        max_length = np.inf
    unweighted = (graph.distances == 1.0).all()
    if verbose and not unweighted:
        LOGGER.info("Graph is weighted ... using slow path.")

    # call dijkstra n times and compress the output in each iteration
    # (Calling just once, with indices = [1, ..., n], would return a dense
    # n times n matrix.)
    import multiprocessing
    from multiprocessing import sharedctypes

    def to_shared_memory(array):
        array = np.ctypeslib.as_ctypes(array)
        return sharedctypes.RawArray(array._type_, array)

    global __this_module
    data = __this_module.__data = to_shared_memory(graph.A.data)
    indptr = __this_module.__indptr = to_shared_memory(graph.A.indptr)
    indices = __this_module.__indices = to_shared_memory(graph.A.indices)

    nodes = np.arange(graph.n_items)
    print_every = max(1, nodes.shape[0] // 10)
    n_workers = multiprocessing.cpu_count() if n_workers is None else n_workers
    with multiprocessing.Pool(
        n_workers, initializer=__init_process, initargs=(data, indptr, indices)
    ) as pool:
        try:
            async_ret = [
                pool.apply_async(
                    _shortest_paths,
                    (
                        graph.A.shape,
                        node,
                        max_length,
                        unweighted,
                        retain_fraction,
                    ),
                )
                for node in nodes
            ]
            values = []
            for idx, ret in enumerate(async_ret):
                data, columns = ret.get()
                values.append((data, columns))
                if verbose and idx % print_every == 0:
                    LOGGER.info(f"processed node {idx + 1}/{nodes.shape[0]}")
        except:  # noqa: E722
            del __this_module.__data
            del __this_module.__indptr
            del __this_module.__indices
            raise

    del __this_module.__data
    del __this_module.__indptr
    del __this_module.__indices

    data, columns = zip(*values)
    graph.A.eliminate_zeros()

    indptr = [0]
    for row in range(graph.n_items):
        indptr.append(indptr[-1] + columns[row].size)
    data = np.concatenate(data)
    indices = np.concatenate(columns)
    adj_matrix = sp.csr_matrix((data, indices, indptr), shape=graph.A.shape)

    adj_matrix = adj_matrix + adj_matrix.T.tocsr()
    return Graph(adj_matrix)


def _minimum_spanning_tree(graph):
    return Graph(sp.csgraph.minimum_spanning_tree(graph.A))


def _thresholded_neighborhood_graph(graph, threshold, max_length):
    graph = shortest_paths(graph, max_length)

    data = []
    indices = []
    indptr = [0]
    n = graph.n_items
    for node in range(n):
        distances = graph.neighbor_distances(node)
        columns = graph.neighbors(node)
        thresholded_indices = distances[distances <= threshold]

        data.append(distances[thresholded_indices])
        indices.append(columns[thresholded_indices])
        indptr.append(indptr[-1] + thresholded_indices.size)

    data = np.concatenate(data)
    indices = np.concatenate(indices)
    indptr = np.array(indptr)
    return Graph(sp.csr_matrix((data, indices, indptr), shape=graph.A.shape))


def k_nearest_neighbors(
    graph,
    k,
    graph_distances=False,
    max_distance=None,
    verbose=False,
):
    """Compute k-nearest neighbors for each node in graph.

    This function computes a k-nearest neighbor graph.

    By default, the input graph is interpreted as a distance matrix, with
    ``graph.adjacency_matrix[i, j]`` giving the distance between i and j. If
    `graph_distances` is True, then the shortest-path metric is used to
    compute neighbors.

    In the returned graph, if i is a neighbor of j and j a neighbor of i, then
    the weight w_{ij} will be +2; if only one is a neighbor of the other, then
    w_{ij} will be +1; if neither are neighbors of each other, then (i, j) will
    not be in the graph.

    Arguments
    ---------
    graph: pymde.Graph
        The graph, representing a distance metric on the items.
    k: int
        The number of nearest neighbors per item.
    graph_distances: bool
        If ``True``, computes shortest-path distances on the graph; otherwise,
        interprets the graph as a distance matrix.
    max_distance: float (optional)
        If not ``None``, neighborhoods are restricted to have a radius
        no greater than max_distance.
    verbose: bool
        If ``True``, print verbose output.

    Returns
    -------
    pymde.Graph
        The k-nearest neighbor graph.
    """

    if sp.issparse(graph):
        graph = Graph(graph)
    elif not isinstance(graph, Graph):
        raise ValueError(
            "`graph` must be a pymde.Graph instance or "
            "scipy.sparse adjacency matrix."
        )

    if graph.n_edges == graph.n_all_edges:
        # the graph is already full; no need to compute shortest paths.
        graph_distances = False

    if graph_distances:
        graph = shortest_paths(graph, max_length=max_distance, verbose=verbose)

    if max_distance is None:
        max_distance = np.inf

    coo_data = []
    coo_rows = []
    coo_cols = []
    for node in range(graph.n_items):
        dist = graph.neighbor_distances(node)
        cols = graph.neighbors(node)
        if cols.size == 0:
            continue

        dist[dist > max_distance] = np.inf
        knn_indices = np.argsort(dist)[:k]

        coo_rows.extend([node for _ in range(knn_indices.size)])
        coo_cols.extend([c for c in cols[knn_indices]])
        coo_data.extend([d for d in dist[knn_indices]])
    coo_graph = sp.coo_matrix(
        (coo_data, (coo_rows, coo_cols)), shape=graph.A.shape
    )

    # weights will be 0, 1, or 2
    coo_graph.data = np.ones(coo_graph.data.shape)
    coo_graph = coo_graph + coo_graph.T
    csr_graph = coo_graph.tocsr()
    return Graph(csr_graph)


# TODO(akshayka) including both threshold and max_length is confusing. fix
# this api
def _neighborhood_graph(
    graph, n_neighbors=None, threshold=None, max_length=None
):
    if n_neighbors is None and threshold is None:
        raise ValueError("either n_neighbors or threshold must be non-None")
    if n_neighbors is not None and threshold is not None:
        raise ValueError(
            "only one of n_neighbors and threshold can be non-none"
        )

    if n_neighbors is not None:
        return k_nearest_neighbors(graph, n_neighbors, max_length)
    return _thresholded_neighborhood_graph(graph, threshold, max_length)
