"""
Homology and Smith Normal Form
==============================
The purpose of computing the Homology groups for data generated
hypergraphs is to identify data sources that correspond to interesting
features in the topology of the hypergraph.

The elements of one of these Homology groups are generated by $k$
dimensional cycles of relationships in the original data that are not
bound together by higher order relationships. Ideally, we want the
briefest description of these cycles; we want a minimal set of
relationships exhibiting interesting cyclic behavior. This minimal set
will be a bases for the Homology group.

The cyclic relationships in the data are discovered using a **boundary
map** represented as a matrix. To discover the bases we compute the
**Smith Normal Form** of the boundary map.

Homology Mod2
-------------
This module computes the homology groups for data represented as an
abstract simplicial complex with chain groups $\{C_k\}$ and $Z_2$ additions.
The boundary matrices are represented as rectangular matrices over $Z_2$.
These matrices are diagonalized and represented in Smith
Normal Form. The kernel and image bases are computed and the Betti
numbers and homology bases are returned.

Methods for obtaining SNF for Z/2Z are based on Ferrario's work:
http://www.dlfer.xyz/post/2016-10-27-smith-normal-form/

"""

import numpy as np
import hypernetx as hnx
import warnings
import copy
from hypernetx import HyperNetXError
from collections import defaultdict
import itertools as it
import pickle
from scipy.sparse import csr_matrix

__all__ = [
    "kchainbasis",
    "bkMatrix",
    "swap_rows",
    "swap_columns",
    "add_to_row",
    "add_to_column",
    "logical_dot",
    "logical_matmul",
    "matmulreduce",
    "logical_matadd",
    "smith_normal_form_mod2",
    "reduced_row_echelon_form_mod2",
    "boundary_group",
    "chain_complex",
    "betti",
    "betti_numbers",
    "homology_basis",
    "hypergraph_homology_basis",
    "interpret",
]


def kchainbasis(h, k):
    """
    Compute the set of k dimensional cells in the abstract simplicial
    complex associated with the hypergraph.

    Parameters
    ----------
    h : hnx.Hypergraph
    k : int
        dimension of cell

    Returns
    -------
     : list
        an ordered list of kchains represented as tuples of length k+1

    See also
    --------
    hnx.hypergraph.toplexes

    Notes
    -----
    - Method works best if h is simple [Berge], i.e. no edge contains another and there are no duplicate edges (toplexes).
    - Hypergraph node uids must be sortable.

    """

    import itertools as it

    kchains = set()
    for e in h.edges():
        en = sorted(h.edges[e])
        if len(en) == k + 1:
            kchains.add(tuple(en))
        elif len(en) > k + 1:
            kchains.update(set(it.combinations(en, k + 1)))
    return sorted(list(kchains))


def bkMatrix(km1basis, kbasis):
    """
    Compute the boundary map from $C_{k-1}$-basis to $C_k$ basis with
    respect to $Z_2$

    Parameters
    ----------
    km1basis : indexable iterable
        Ordered list of $k-1$ dimensional cell
    kbasis : indexable iterable
        Ordered list of $k$ dimensional cells

    Returns
    -------
    bk : np.array
        boundary matrix in $Z_2$ stored as boolean

    """
    bk = np.zeros((len(km1basis), len(kbasis)), dtype=int)
    for cell in kbasis:
        for idx in range(len(cell)):
            face = cell[:idx] + cell[idx + 1 :]
            row = km1basis.index(face)
            col = kbasis.index(cell)
            bk[row, col] = 1
    return bk


def _rswap(i, j, S):
    """
    Swaps ith and jth row of copy of S

    Parameters
    ----------
    i : int
    j : int
    S : np.array

    Returns
    -------
    N : np.array
    """
    N = copy.deepcopy(S)
    row = copy.deepcopy(N[i])
    N[i] = copy.deepcopy(N[j])
    N[j] = row
    return N


def _cswap(i, j, S):
    """
    Swaps ith and jth column of copy of S

    Parameters
    ----------
    i : int
    j : int
    S : np.array
        matrix

    Returns
    -------
    N : np.array
    """
    N = _rswap(i, j, S.transpose()).transpose()
    return N


def swap_rows(i, j, *args):
    """
    Swaps ith and jth row of each matrix in args
    Returns a list of new matrices

    Parameters
    ----------
    i : int
    j : int
    args : np.arrays

    Returns
    -------
    list
        list of copies of args with ith and jth row swapped
    """
    output = list()
    for M in args:
        output.append(_rswap(i, j, M))
    return output


def swap_columns(i, j, *args):
    """
    Swaps ith and jth column of each matrix in args
    Returns a list of new matrices

    Parameters
    ----------
    i : int
    j : int
    args : np.arrays

    Returns
    -------
    list
        list of copies of args with ith and jth row swapped
    """
    output = list()
    for M in args:
        output.append(_cswap(i, j, M))
    return output


def add_to_row(M, i, j):
    """
    Replaces row i with logical xor between row i and j

    Parameters
    ----------
    M : np.array
    i : int
        index of row being altered
    j : int
        index of row being added to altered

    Returns
    -------
    N : np.array
    """
    N = copy.deepcopy(M)
    N[i] = 1 * np.logical_xor(N[i], N[j])
    return N


def add_to_column(M, i, j):
    """
    Replaces column i (of M) with logical xor between column i and j

    Parameters
    ----------
    M : np.array
        matrix
    i : int
        index of column being altered
    j : int
        index of column being added to altered

    Returns
    -------
    N : np.array
    """
    N = M.transpose()
    return add_to_row(N, i, j).transpose()


def logical_dot(ar1, ar2):
    """
    Returns the boolean equivalent of the dot product mod 2 on two 1-d arrays of
    the same length.

    Parameters
    ----------
    ar1 : numpy.ndarray
        1-d array
    ar2 : numpy.ndarray
        1-d array

    Returns
    -------
    : bool
        boolean value associated with dot product mod 2

    Raises
    ------
    HyperNetXError
        If arrays are not of the same length an error will be raised.
    """
    if len(ar1) != len(ar2):
        raise HyperNetXError("logical_dot requires two 1-d arrays of the same length")
    else:
        return 1 * np.logical_xor.reduce(np.logical_and(ar1, ar2))


def logical_matmul(mat1, mat2):
    """
    Returns the boolean equivalent of matrix multiplication mod 2 on two
    binary arrays stored as type boolean

    Parameters
    ----------
    mat1 : np.ndarray
        2-d array of boolean values
    mat2 : np.ndarray
        2-d array of boolean values

    Returns
    -------
    mat : np.ndarray
        boolean matrix equivalent to the mod 2 matrix multiplication of the
        matrices as matrices over Z/2Z

    Raises
    ------
    HyperNetXError
        If inner dimensions are not equal an error will be raised.

    """
    L1, R1 = mat1.shape
    L2, R2 = mat2.shape
    if R1 != L2:
        raise HyperNetXError(
            "logical_matmul called for matrices with inner dimensions mismatched"
        )

    mat = np.zeros((L1, R2), dtype=int)
    mat2T = mat2.transpose()
    for i in range(L1):
        if np.any(mat1[i]):
            for j in range(R2):
                mat[i, j] = logical_dot(mat1[i], mat2T[j])
        else:
            mat[i] = np.zeros((1, R2), dtype=int)
    return mat


def matmulreduce(arr, reverse=False):
    """
    Recursively applies a 'logical multiplication' to a list of boolean arrays.

    For arr = [arr[0],arr[1],arr[2]...arr[n]] returns product arr[0]arr[1]...arr[n]
    If reverse = True, returns product arr[n]arr[n-1]...arr[0]

    Parameters
    ----------
    arr : list of np.array
        list of nxm matrices represented as np.array
    reverse : bool, optional
        order to multiply the matrices

    Returns
    -------
    P : np.array
        Product of matrices in the list
    """
    if reverse:
        items = range(len(arr) - 1, -1, -1)
    else:
        items = range(len(arr))
    P = arr[items[0]]
    for i in items[1:]:
        P = logical_matmul(P, arr[i]) * 1
    return P


def logical_matadd(mat1, mat2):
    """
    Returns the boolean equivalent of matrix addition mod 2 on two
    binary arrays stored as type boolean

    Parameters
    ----------
    mat1 : np.ndarray
        2-d array of boolean values
    mat2 : np.ndarray
        2-d array of boolean values

    Returns
    -------
    mat : np.ndarray
        boolean matrix equivalent to the mod 2 matrix addition of the
        matrices as matrices over Z/2Z

    Raises
    ------
    HyperNetXError
        If dimensions are not equal an error will be raised.

    """
    S1 = mat1.shape
    S2 = mat2.shape
    mat = np.zeros(S1, dtype=int)
    if S1 != S2:
        raise HyperNetXError(
            "logical_matadd called for matrices with different dimensions"
        )
    if len(S1) == 1:
        for idx in range(S1[0]):
            mat[idx] = 1 * np.logical_xor(mat1[idx], mat2[idx])
    else:
        for idx in range(S1[0]):
            for jdx in range(S1[1]):
                mat[idx, jdx] = 1 * np.logical_xor(mat1[idx, jdx], mat2[idx, jdx])
    return mat


# Convenience methods for computing Smith Normal Form
# All of these operations have themselves as inverses


def _sr(i, j, M, L):
    return swap_rows(i, j, M, L)


def _sc(i, j, M, R):
    return swap_columns(i, j, M, R)


def _ar(i, j, M, L):
    return add_to_row(M, i, j), add_to_row(L, i, j)


def _ac(i, j, M, R):
    return add_to_column(M, i, j), add_to_column(R, i, j)


def _get_next_pivot(M, s1, s2=None):
    """
    Determines the first r,c indices in the submatrix of M starting
    with row s1 and column s2 index (row,col) that is nonzero,
    if it exists.

    Search starts with the s2th column and looks for the first nonzero
    s1 row. If none is found, search continues to the next column and so
    on.

    Parameters
    ----------
    M : np.array
        matrix represented as np.array
    s1 : int
        index of row position to start submatrix of M
    s2 : int, optional, default = s1
        index of column position to start submatrix of M

    Returns
    -------
    (r,c) : tuple of int or None

    """
    # find the next nonzero pivot to put in s,s spot for Smith Normal Form
    m, n = M.shape
    if not s2:
        s2 = s1
    for c in range(s2, n):
        for r in range(s1, m):
            if M[r, c]:
                return (r, c)
    return None


def smith_normal_form_mod2(M):
    """
    Computes the invertible transformation matrices needed to compute the
    Smith Normal Form of M modulo 2

    Parameters
    ----------
    M : np.array
        a rectangular matrix with data type bool
    track : bool
        if track=True will print out the transformation as Z/2Z matrix as it
        discovers L[i] and R[j]

    Returns
    -------
    L, R, S, Linv : np.arrays
        LMR = S is the Smith Normal Form of the matrix M.

    Note
    ----
    Given a mxn matrix $M$ with
    entries in $Z_2$ we start with the equation: $L M R = S$, where
    $L = I_m$, and $R=I_n$ are identity matrices and $S = M$. We
    repeatedly apply actions to the left and right side of the equation
    to transform S into a diagonal matrix.
    For each action applied to the left side we apply its inverse
    action to the right side of I_m to generate $L^{-1}$.
    Finally we verify:
    $L M R = S$ and  $LLinv = I_m$.
    """

    S = copy.copy(M)
    dimL, dimR = M.shape

    # initialize left and right transformations with identity matrices
    L = np.eye(dimL, dtype=int)
    R = np.eye(dimR, dtype=int)
    Linv = np.eye(dimL, dtype=int)
    for s in range(min(dimL, dimR)):
        # Find index pair (rdx,cdx) with value 1 in submatrix M[s:,s:]
        pivot = _get_next_pivot(S, s)
        if not pivot:
            break
        else:
            rdx, cdx = pivot
        # Swap rows and columns as needed so that 1 is in the s,s position
        if rdx > s:
            S, L = _sr(s, rdx, S, L)
            Linv = swap_columns(rdx, s, Linv)[0]
        if cdx > s:
            S, R = _sc(s, cdx, S, R)
        # add sth row to every row with 1 in sth column & sth column to every column with 1 in sth row
        row_indices = [idx for idx in range(s + 1, dimL) if S[idx, s] == 1]
        for rdx in row_indices:
            S, L = _ar(rdx, s, S, L)
            Linv = add_to_column(Linv, s, rdx)
        column_indices = [jdx for jdx in range(s + 1, dimR) if S[s, jdx] == 1]
        for cdx in column_indices:
            S, R = _ac(cdx, s, S, R)
    return L, R, S, Linv


def reduced_row_echelon_form_mod2(M):
    """
    Computes the invertible transformation matrices needed to compute
    the reduced row echelon form of M modulo 2

    Parameters
    ----------
    M : np.array
        a rectangular matrix with elements in $Z_2$

    Returns
    -------
    L, S, Linv : np.arrays
        LM = S where S is the reduced echelon form of M
        and M = LinvS
    """
    S = copy.deepcopy(M)
    dimL, dimR = M.shape

    # method with numpy
    Linv = np.eye(dimL, dtype=int)
    L = np.eye(dimL, dtype=int)

    s2 = 0
    s1 = 0
    while s2 <= dimR and s1 <= dimL:
        # Find index pair (rdx,cdx) with value 1 in submatrix M[s1:,s2:]
        # look for the first 1 in the s2 column
        pivot = _get_next_pivot(S, s1, s2)

        if not pivot:
            return L, S, Linv
        else:
            rdx, cdx = pivot
            if rdx > s1:
                # Swap rows as needed so that 1 leads the row
                S, L = _sr(s1, rdx, S, L)
                Linv = swap_columns(rdx, s1, Linv)[0]
            # add s1th row to every nonzero row
            row_indices = [
                idx for idx in range(0, dimL) if idx != s1 and S[idx, cdx] == 1
            ]
            for idx in row_indices:
                S, L = _ar(idx, s1, S, L)
                Linv = add_to_column(Linv, s1, idx)
            s1, s2 = s1 + 1, cdx + 1

    return L, S, Linv


def boundary_group(image_basis):
    """
    Returns a csr_matrix with rows corresponding to the elements of the
    group  generated by image basis over $\mathbb{Z}_2$

    Parameters
    ----------
    image_basis : numpy.ndarray or scipy.sparse.csr_matrix
        2d-array of basis elements

    Returns
    -------
     : scipy.sparse.csr_matrix
    """
    if len(image_basis) > 10:
        msg = """
        This method is inefficient for large image bases.
        """
        warnings.warn(msg, stacklevel=2)
    if np.sum(image_basis) == 0:
        return None
    dim = image_basis.shape[0]
    itm = csr_matrix(list(it.product([0, 1], repeat=dim)))
    return csr_matrix(np.mod(itm * image_basis, 2))


def _compute_matrices_for_snf(bd):
    """
    Helper method for smith normal form decomposition for boundary maps
    associated to chain complex

    Parameters
    ----------
    bd : dict
        dict of k-boundary matrices keyed on dimension of domain

    Returns
    -------
    L,R,S,Linv : dict
        dict of matrices ranging over krange

    """
    L, R, S, Linv = [dict() for i in range(4)]

    for kdx in bd:
        L[kdx], R[kdx], S[kdx], Linv[kdx] = smith_normal_form_mod2(bd[kdx])
    return L, R, S, Linv


def _get_krange(max_dim, k=None):
    """
    Helper method to compute range of appropriate k dimensions for homology
    computations given k and the max dimension of a simplicial complex
    """
    if k is None:
        krange = [1, max_dim]
    elif isinstance(k, int):
        if k == 0:
            msg = (
                "Only kth simplicial homology groups for k>0 may be computed."
                "If you are interested in k=0, compute the number connected components."
            )
            print(msg)
            return
        if k > max_dim:
            msg = f"No simplices of dim {k} exist. k adjusted to max dim."
            print(msg)
        krange = [min([k, max_dim])] * 2
    elif not len(k) == 2:
        msg = f"Please enter krange as a positive integer or list of integers: [<min k>,<max k>] inclusive."
        print(msg)
        return None
    elif not k[0] <= k[1]:
        msg = f"k must be an integer or a list of two integers [min,max] with min <=max"
        print(msg)
        return None
    else:
        krange = k

    if krange[1] > max_dim:
        msg = f"No simplices of dim {krange[1]} exist. Range adjusted to max dim."
        print(msg)
        krange[1] = max_dim
    if krange[0] < 1:
        msg = (
            "Only kth simplicial homology groups for k>0 may be computed."
            "If you are interested in k=0, compute the number of connected components."
        )
        print(msg)
        krange[0] = 1
    return krange


def chain_complex(h, k=None):
    """
    Compute the k-chains and k-boundary maps required to compute homology
    for all values in k

    Parameters
    ----------
    h : hnx.Hypergraph
    k : int or list of length 2, optional, default=None
        k must be an integer greater than 0 or a list of
        length 2 indicating min and max dimensions to be
        computed. eg. if k = [1,2] then 0,1,2,3-chains
        and boundary maps for k=1,2,3 will be returned,
        if None than k = [1,max dimension of edge in h]

    Returns
    -------
    C, bd : dict
        C is a dictionary of lists
        bd is a dictionary of numpy arrays
    """
    max_dim = np.max([len(h.edges[e]) for e in h.edges()]) - 1
    krange = _get_krange(max_dim, k)
    if not krange:
        return
    # Compute chain complex

    C = defaultdict(list)
    C[krange[0] - 1] = kchainbasis(h, krange[0] - 1)
    bd = dict()
    for kdx in range(krange[0], krange[1] + 2):
        C[kdx] = kchainbasis(h, kdx)
        bd[kdx] = bkMatrix(C[kdx - 1], C[kdx])
    return C, bd


def betti(bd, k=None):
    """
    Generate the kth-betti numbers for a chain complex with boundary
    matrices given by bd

    Parameters
    ----------
    bd: dict of k-boundary matrices keyed on dimension of domain
    k : int, list or tuple, optional, default=None
        list must be min value and max value of k values inclusive
        if None, then all betti numbers for dimensions of existing cells will be
        computed.

    Returns
    -------
    betti : dict
        Description
    """
    rank = defaultdict(int)
    if k:
        max_dim = max(bd.keys())
        krange = _get_krange(max_dim, k)
        if not krange:
            return
        kvals = sorted(set(range(krange[0], krange[1] + 2)).intersection(bd.keys()))
    else:
        kvals = bd.keys()
    for kdx in kvals:
        _, S, _ = hnx.reduced_row_echelon_form_mod2(bd[kdx])
        rank[kdx] = np.sum(np.sum(S, axis=1).astype(bool))

    betti = dict()
    for kdx in kvals:
        if kdx + 1 in kvals:
            betti[kdx] = bd[kdx].shape[1] - rank[kdx] - rank[kdx + 1]
        else:
            continue

    return betti


def betti_numbers(h, k=None):
    """
    Return the kth betti numbers for the simplicial homology of the ASC
    associated to h

    Parameters
    ----------
    h : hnx.Hypergraph
        Hypergraph to compute the betti numbers from
    k : int or list, optional, default=None
        list must be min value and max value of k values inclusive
        if None, then all betti numbers for dimensions of existing cells will be
        computed.

    Returns
    -------
    betti : dict
        A dictionary of betti numbers keyed by dimension
    """
    _, bd = chain_complex(h, k)
    return betti(bd)


def homology_basis(bd, k=None, boundary=False, **kwargs):
    """
    Compute a basis for the kth-simplicial homology group, $H_k$, defined by a
    chain complex $C$ with boundary maps given by bd $= \{k:\partial_k$\}$

    Parameters
    ----------
    bd : dict
        dict of boundary matrices on k-chains to k-1 chains keyed on k
        if krange is a tuple then all boundary matrices k \in [krange[0],..,krange[1]]
        inclusive must be in the dictionary
    k : int or list of ints, optional, default=None
        k must be a positive integer or a list of
        2 integers indicating min and max dimensions to be
        computed, if none given all homology groups will be computed from
        available boundary matrices in bd
    boundary : bool
        option to return a basis for the boundary group from each dimension.
        Needed to compute the shortest generators in the homology group.

    Returns
    -------
    basis : dict
        dict of generators as 0-1 tuples keyed by dim
        basis for dimension k will be returned only if bd[k] and bd[k+1] have
        been provided.
    im : dict
        dict of boundary group generators keyed by dim
    """
    max_dim = max(bd.keys())
    if k:
        krange = _get_krange(max_dim, k)
        kvals = sorted(
            set(bd.keys()).intersection(range(krange[0], krange[1] + 2))
        )  # to get kth dim need k+1 bdry matrix
    else:
        kvals = bd.keys()

    L, R, S, Linv = _compute_matrices_for_snf(
        {k: v for k, v in bd.items() if k in kvals}
    )

    rank = dict()
    for kdx in kvals:
        rank[kdx] = np.sum(S[kdx])

    basis = dict()
    im = dict()
    for kdx in kvals:
        if kdx + 1 not in kvals:
            continue
        rank1 = rank[kdx]
        rank2 = rank[kdx + 1]
        ker1 = R[kdx][:, rank1:]
        im2 = Linv[kdx + 1][:, :rank2]
        cokernel2 = Linv[kdx + 1][:, rank2:]
        cokproj2 = L[kdx + 1][rank2:, :]

        proj = matmulreduce([cokernel2, cokproj2, ker1]).transpose()
        _, proj, _ = reduced_row_echelon_form_mod2(proj)
        # proj = np.array(proj)
        basis[kdx] = np.array([row for row in proj if np.any(row)])
        if boundary:
            im[kdx] = im2.transpose()
    if boundary:
        return basis, im
    else:
        return basis


def hypergraph_homology_basis(h, k=None, shortest=False, interpreted=True):
    """
    Computes the kth-homology groups mod 2 for the ASC
    associated with the hypergraph h for k in krange inclusive

    Parameters
    ----------
    h : hnx.Hypergraph
    k : int or list of length 2, optional, default = None
        k must be an integer greater than 0 or a list of
        length 2 indicating min and max dimensions to be
        computed
    shortest : bool, optional, default=False
        option to look for shortest representative for each coset in the
        homology group, only good for relatively small examples
    interpreted : bool, optional, default = True
        if True will return an explicit basis in terms of the k-chains

    Returns
    -------
    basis : list
        list of generators as k-chains as boolean vectors
    interpreted_basis :
        lists of kchains in basis

    """
    C, bd = chain_complex(h, k)
    if shortest:
        basis = defaultdict(list)
        tbasis, im = homology_basis(bd, boundary=True)
        for kdx in tbasis:
            imgrp = boundary_group(im[kdx])
            if imgrp is None:
                basis[kdx] = tbasis[kdx]
            else:
                for b in tbasis[kdx]:
                    coset = np.array(
                        np.mod(imgrp + b, 2)
                    )  # dimensions appear to be wrong. See tests2 cell 5
                    idx = np.argmin(np.sum(coset, axis=1))
                    basis[kdx].append(coset[idx])
        basis = dict(basis)

    else:
        basis = homology_basis(bd)

    if interpreted:
        interpreted_basis = dict()
        for kdx in basis:
            interpreted_basis[kdx] = interpret(C[kdx], basis[kdx], labels=None)
        return basis, interpreted_basis
    else:
        return basis


def interpret(Ck, arr, labels=None):
    """
    Returns the data as represented in Ck associated with the arr

    Parameters
    ----------
    Ck : list
        a list of k-cells being referenced by arr
    arr : np.array
        array of 0-1 vectors
    labels : dict, optional
        dictionary of labels to associate to the nodes in the cells

    Returns
    ----
    : list
        list of k-cells referenced by data in Ck

    """

    def translate(cell, labels=labels):
        if not labels:
            return cell
        else:
            temp = list()
            for node in cell:
                temp.append(labels[node])
            return tuple(temp)

    output = list()
    for vec in arr:
        if len(Ck) != len(vec):
            raise HyperNetXError("elements of arr must have the same length as Ck")
        output.append([translate(Ck[idx]) for idx in range(len(vec)) if vec[idx] == 1])
    return output
