import numpy as np
import picos as pc


def irev_blocked_lp(S, R):
    """
    Linear Program to detect all blocked irreversible reactions.
    Code adapted from the blocked.m method from Mojtaba Tefagh, Stephen P. Boyd, 2019 (Stanford University)
    Original code: https://github.com/mtefagh/swiftcore/blob/master/src/blocked.m

    :param S: Stoichiometric matrix
    :param R: Matrix of size Num. R x 3 with the reactions
    :return: tuple LP, V (LP = picos problem, V = optimized variables)
    """
    m, n = S.shape
    # Find irreversible reactions
    irev_rxn = R['lb'] >= 0
    irev = m + np.where(irev_rxn)[0]
    A = np.concatenate((S.T, -np.eye(n)), axis=1)
    lb = np.concatenate(([-np.inf] * m, [0] * n))
    lb[irev] = -1
    ub = np.concatenate(([np.inf] * m, [0] * n))
    V = pc.RealVariable("V", (n + m, 1), lower=lb, upper=ub)
    LP = pc.Problem()
    LP.add_constraint(A[irev_rxn, :] * V <= 0)
    LP.add_constraint(A[np.logical_not(irev_rxn), :] * V == 0)
    obj = np.zeros((m + n, 1))
    obj[irev] = 1
    C = pc.Constant("C", value=obj)
    LP.set_objective('min', C.T * V)
    return LP, V


def swiftcc(S, R, solver='glpk', verbosity=0):
    """
    SWIFTCC method for finding the largest flux consistent subnetwork of
    the original metabolic network. The implementation is based on the
    matlab code from Mojtaba Tefagh, Stephen P. Boyd, 2019 (Stanford University).
    https://github.com/mtefagh/swiftcore/blob/master/src/swiftcc.m

    :param S: Stoichiometric matrix
    :param R: Matrix of size Num. R x 3 with the reactions
    :param solver: picos-compatible linear solver (gplk, cplex, gurobi...)
    :param verbosity: 1 to show solver specific information, 0 to disable output
    :return: indexes of the consistent reactions
    """
    from scipy.linalg import qr
    m, n = S.shape
    rev = R['lb'] < 0
    LP, V = irev_blocked_lp(S, R)
    LP.options.verbosity = verbosity
    LP.options.solver = solver
    lp_sol = LP.solve()
    consistent = np.array([True] * n)
    consistent[np.array(V[m:].value).squeeze() < -0.5] = False
    # Estimate the rank of the consistent matrix using qr
    # (equivalent to np.linalg.rank_matrix)
    tol = np.linalg.norm(S[:, consistent]) * np.finfo(S.dtype).eps
    Q, R, _ = qr(S[:, consistent].T, pivoting=True)
    srank = np.sum(np.abs(np.diag(R)) > tol)
    Z = Q[rev[consistent] == True, srank:]
    consistent[np.logical_and(consistent, rev == True)] = np.diag(np.dot(Z, Z.T)) > tol ** 2
    return np.where(consistent == True)[0], lp_sol

def consistent_subnetwork(network, solver='glpk', verbosity=0):
    """Finds the largest consistent subnetwork of the original network using SwiftCC.
    Note that SwiftCC does not take into account the bounds of the reactions.

    Args:
        network (MiomNetwork): A MiomNetwork instance.

    Returns:
        MiomNetwork: Flux consistent subnetwork.
        picos.modeling.solution: The solution of the LP problem.
    """
    c, lp = swiftcc(network.S, network.R, solver=solver, verbosity=verbosity)
    return network.subnet(c), lp
