import numpy as np
import scipy.sparse as sparse

from hic3defdr.util.clusters import find_clusters


def classify(row, col, value, clusters):
    """
    Classifies pixels (specified by row, col coordinates) into classes
    (corresponding to the columns of value) where their value is highest,
    ignoring those pixels that don't belong to the specified clusters.

    Parameters
    ----------
    row, col : np.ndarray
        Row and column indices, respectively, identifying pixels which
        correspond to rows of ``value``.
    value : np.ndarray
        Value to use to assign classes. Rows correspond to pixels, columns to
        conditions.
    clusters : list of set of tuple
        Only pixels in these clusters will be assigned classes.

    Returns
    -------
    list of list of set of tuple
        The outer list represents classes (one per condition). Its elements
        represent clusters of points with that class/condition.
    """
    # convert clusters to idx, parallel to row/col
    pixels = set().union(*(c for c in clusters))
    idx = np.zeros_like(row, dtype=bool)
    for i, p in enumerate(zip(row, col)):
        if p in pixels:
            idx[i] = True

    # assign classes
    classes = np.argmax(value[idx, :], axis=1)

    # re-cluster each class
    n = max(row.max(), col.max()) + 1
    class_clusters = []
    for c in range(value.shape[1]):
        coo = sparse.coo_matrix(
            (np.ones((classes == c).sum()),
             (row[idx][classes == c], col[idx][classes == c])), shape=(n, n))
        class_clusters.append(find_clusters(coo))

    return class_clusters
