from sklearn.tree import BaseDecisionTree
from sklearn.ensemble import RandomForestRegressor, ExtraTreesRegressor
import time
from operator import itemgetter
from multiprocessing import Pool
from typing import Optional, Union
from numpy import ndarray, array, transpose, std, zeros, ndenumerate, random


def compute_feature_importances(estimator):
    if isinstance(estimator, BaseDecisionTree):
        return estimator.tree_.compute_feature_importances(normalize=False)
    else:
        importances = [
            e.tree_.compute_feature_importances(normalize=False)
            for e in estimator.estimators_
        ]
        importances = array(importances)
        return sum(importances, axis=0) / len(estimator)


def get_link_list(
    VIM: ndarray,
    gene_names: Optional[list[str]] = None,
    regulators: Union[list[str], str] = "all",
    maxcount: Union[int, str] = "all",
    file_name: Optional[str] = None,
):
    """Gets the ranked list of (directed) regulatory links.

    Args
        VIM (np.array): Array as returned by the function GENIE3(), in which the element (i,j) is the score of the edge directed from the i-th gene to the j-th gene.
        gene_names (list[str] optional): List of length p, where p is the number of rows/columns in VIM, containing the names of the genes. The i-th item of gene_names must correspond to the i-th row/column of VIM. When the gene names are not provided, the i-th gene is named Gi. Default is None.
        regulators (list[str], optional): List containing the names of the candidate regulators. When a list of regulators is provided, the names of all the genes must be provided (in gene_names), and the returned list contains only edges directed from the candidate regulators. When regulators is set to 'all', any gene can be a candidate regulator. Default is 'all'.
        maxcount (Union[str, int], optional): Writes only the first maxcount regulatory links of the ranked list. When maxcount is set to 'all', all the regulatory links are written. Default is 'all'.
        file_name (str, optional): Writes the ranked list of regulatory links to the file file_name. Default is None.

    Returns
        pd.Dataframe: The list of regulatory links, ordered according to the edge score. Auto-regulations do not appear in the list. Regulatory links with a score equal to zero are randomly permuted. In the ranked list of edges, each line has format:
            regulator   target gene     score of edge
    """

    # Check input arguments
    if not isinstance(VIM, ndarray):
        raise ValueError("VIM must be a square array")
    elif VIM.shape[0] != VIM.shape[1]:
        raise ValueError("VIM must be a square array")

    ngenes = VIM.shape[0]

    if gene_names is not None:
        if not isinstance(gene_names, (list, tuple)):
            raise ValueError("input argument gene_names must be a list of gene names")
        elif len(gene_names) != ngenes:
            raise ValueError(
                "input argument gene_names must be a list of length p, where p is the number of columns/genes in the expression data"
            )

    if regulators != "all":
        if not isinstance(regulators, (list, tuple)):
            raise ValueError("input argument regulators must be a list of gene names")

        if gene_names is None:
            raise ValueError(
                "the gene names must be specified (in input argument gene_names)"
            )
        else:
            sIntersection = set(gene_names).intersection(set(regulators))
            if not sIntersection:
                raise ValueError(
                    "The genes must contain at least one candidate regulator"
                )

    if maxcount != "all" and not isinstance(maxcount, int):
        raise ValueError('input argument maxcount must be "all" or a positive integer')

    if file_name is not None and not isinstance(file_name, str):
        raise ValueError("input argument file_name must be a string")

    # Get the indices of the candidate regulators
    if regulators == "all":
        input_idx = range(ngenes)
    else:
        input_idx = [i for i, gene in enumerate(gene_names) if gene in regulators]

    # Get the non-ranked list of regulatory links
    vInter = [
        (i, j, score) for (i, j), score in ndenumerate(VIM) if i in input_idx and i != j
    ]

    # Rank the list according to the weights of the edges
    vInter_sort = sorted(vInter, key=itemgetter(2), reverse=True)
    nInter = len(vInter_sort)

    # Random permutation of edges with score equal to 0
    flag = 1
    i = 0
    while flag and i < nInter:
        (TF_idx, target_idx, score) = vInter_sort[i]
        if score == 0:
            flag = 0
        else:
            i += 1

    if not flag:
        items_perm = vInter_sort[i:]
        items_perm = random.permutation(items_perm)
        vInter_sort[i:] = items_perm

    # Write the ranked list of edges
    nToWrite = nInter
    if isinstance(maxcount, int) and maxcount >= 0 and maxcount < nInter:
        nToWrite = maxcount

    if file_name:
        outfile = open(file_name, "w")

        if gene_names is not None:
            for i in range(nToWrite):
                (TF_idx, target_idx, score) = vInter_sort[i]
                TF_idx = int(TF_idx)
                target_idx = int(target_idx)
                outfile.write(
                    "%s\t%s\t%.6f\n"
                    % (gene_names[TF_idx], gene_names[target_idx], score)
                )
        else:
            for i in range(nToWrite):
                (TF_idx, target_idx, score) = vInter_sort[i]
                TF_idx = int(TF_idx)
                target_idx = int(target_idx)
                outfile.write("G%d\tG%d\t%.6f\n" % (TF_idx + 1, target_idx + 1, score))

        outfile.close()

    else:
        if gene_names is not None:
            for i in range(nToWrite):
                (TF_idx, target_idx, score) = vInter_sort[i]
                TF_idx = int(TF_idx)
                target_idx = int(target_idx)
                print(
                    "%s\t%s\t%.6f" % (gene_names[TF_idx], gene_names[target_idx], score)
                )
        else:
            for i in range(nToWrite):
                (TF_idx, target_idx, score) = vInter_sort[i]
                TF_idx = int(TF_idx)
                target_idx = int(target_idx)
                print("G%d\tG%d\t%.6f" % (TF_idx + 1, target_idx + 1, score))


def GENIE3(
    expr_data: ndarray,
    gene_names: Optional[list[str]] = None,
    regulators: Union[list[str], str] = "all",
    tree_method: str = "RF",
    K: Union[str, int] = "sqrt",
    ntrees: int = 1000,
    nthreads: int = 1,
):
    """
    GENIE3 Computation of tree-based scores for all putative regulatory links.

    Args:
        expr_data (numpy.ndarray): Array containing gene expression values. Each row corresponds to a condition and each column corresponds to a gene.
        gene_names (list[str], optional): List of length p, where p is the number of columns in expr_data, containing the names of the genes. The i-th item of gene_names must correspond to the i-th column of expr_data. Defaults to None.
        regulators (list[str], optional): List containing the names of the candidate regulators. When a list of regulators is provided, the names of all the genes must be provided (in gene_names). When regulators is set to 'all', any gene can be a candidate regulator. Defaults to 'all'.
        tree_method (str, optional): Specifies which tree-based procedure is used: either Random Forest ('RF') or Extra-Trees ('ET'). Defaults to 'RF'.
        K (str or int, optional): Specifies the number of selected attributes at each node of one tree: either the square root of the number of candidate regulators ('sqrt'), the total number of candidate regulators ('all'), or any positive integer. Defaults to 'sqrt'.
        ntrees (int, optional): Specifies the number of trees grown in an ensemble. Defaults to 1000.
        nthreads (int, optional): Number of threads used for parallel computing. Defaults to 1.

    Returns:
        numpy.ndarray: An array in which the element (i,j) is the score of the edge directed from the i-th gene to the j-th gene. All diagonal elements are set to zero (auto-regulations are not considered). When a list of candidate regulators is provided, the scores of all the edges directed from a gene that is not a candidate regulator are set to zero.
    """
    time_start = time.time()

    # Check input arguments
    if not isinstance(expr_data, ndarray):
        raise ValueError(
            "expr_data must be an array in which each row corresponds to a condition/sample and each column corresponds to a gene"
        )

    ngenes = expr_data.shape[1]

    if gene_names is not None:
        if not isinstance(gene_names, (list, tuple)):
            raise ValueError("input argument gene_names must be a list of gene names")
        elif len(gene_names) != ngenes:
            raise ValueError(
                "input argument gene_names must be a list of length p, where p is the number of columns/genes in the expr_data"
            )

    if regulators != "all":
        if not isinstance(regulators, (list, tuple)):
            raise ValueError("input argument regulators must be a list of gene names")

        if gene_names is None:
            raise ValueError(
                "the gene names must be specified (in input argument gene_names)"
            )
        else:
            sIntersection = set(gene_names).intersection(set(regulators))
            if not sIntersection:
                raise ValueError(
                    "the genes must contain at least one candidate regulator"
                )

    if tree_method != "RF" and tree_method != "ET":
        raise ValueError(
            'input argument tree_method must be "RF" (Random Forests) or "ET" (Extra-Trees)'
        )

    if K != "sqrt" and K != "all" and not isinstance(K, int):
        raise ValueError(
            'input argument K must be "sqrt", "all" or a stricly positive integer'
        )

    if isinstance(K, int) and K <= 0:
        raise ValueError(
            'input argument K must be "sqrt", "all" or a stricly positive integer'
        )

    if not isinstance(ntrees, int):
        raise ValueError("input argument ntrees must be a stricly positive integer")
    elif ntrees <= 0:
        raise ValueError("input argument ntrees must be a stricly positive integer")

    if not isinstance(nthreads, int):
        raise ValueError("input argument nthreads must be a stricly positive integer")
    elif nthreads <= 0:
        raise ValueError("input argument nthreads must be a stricly positive integer")

    print("Tree method: " + str(tree_method))
    print("K: " + str(K))
    print("Number of trees: " + str(ntrees))
    print("\n")

    # Get the indices of the candidate regulators
    if regulators == "all":
        input_idx = list(range(ngenes))
    else:
        input_idx = [i for i, gene in enumerate(gene_names) if gene in regulators]

    # Learn an ensemble of trees for each target gene, and compute scores for candidate regulators
    VIM = zeros((ngenes, ngenes))

    if nthreads > 1:
        print("running jobs on %d threads" % nthreads)

        input_data = list()
        for i in range(ngenes):
            input_data.append([expr_data, i, input_idx, tree_method, K, ntrees])

        pool = Pool(nthreads)
        alloutput = pool.map(wr_GENIE3_single, input_data)

        for i, vi in alloutput:
            VIM[i, :] = vi

    else:
        print("running single threaded jobs")
        for i in range(ngenes):
            print("Gene %d/%d..." % (i + 1, ngenes))

            vi = GENIE3_single(expr_data, i, input_idx, tree_method, K, ntrees)
            VIM[i, :] = vi

    VIM = transpose(VIM)

    time_end = time.time()
    print("Elapsed time: %.2f seconds" % (time_end - time_start))

    return VIM


def wr_GENIE3_single(args):
    return [
        args[1],
        GENIE3_single(args[0], args[1], args[2], args[3], args[4], args[5]),
    ]


def GENIE3_single(expr_data, output_idx, input_idx, tree_method, K, ntrees):
    ngenes = expr_data.shape[1]

    # Expression of target gene
    output = expr_data[:, output_idx]

    # Normalize output data
    output = output / std(output)

    # Remove target gene from candidate regulators
    input_idx = input_idx[:]
    if output_idx in input_idx:
        input_idx.remove(output_idx)

    expr_data_input = expr_data[:, input_idx]

    # Parameter K of the tree-based method
    if (K == "all") or (isinstance(K, int) and K >= len(input_idx)):
        max_features = "auto"
    else:
        max_features = K

    if tree_method == "RF":
        treeEstimator = RandomForestRegressor(
            n_estimators=ntrees, max_features=max_features
        )
    elif tree_method == "ET":
        treeEstimator = ExtraTreesRegressor(
            n_estimators=ntrees, max_features=max_features
        )

    # Learn ensemble of trees
    treeEstimator.fit(expr_data_input, output)

    # Compute importance scores
    feature_importances = compute_feature_importances(treeEstimator)
    vi = zeros(ngenes)
    vi[input_idx] = feature_importances

    return vi
