import math


def keep_top_genes(counts, num_genes_discard, criteria="Variance"):
    """
    This function takes a matrix of counts (Genes as columns and spots as rows)
    and returns a new matrix of counts where a number of genesa re kept
    using the variance or the total count as filtering criterias.
    :param counts: a matrix of counts
    :param num_genes_discard: the % (1-100) of genes to keep
    :param criteria: the criteria used to select ("Variance or "TopRanked")
    :return: a new matrix of counts with only the top ranked genes.
    """
    if num_genes_discard <= 0:
        return counts
    # Spots as columns and genes as rows
    counts = counts.transpose()
    # Keep only the genes with higher over-all variance
    num_genes = len(counts.index)
    print("Removing {}% of genes based on the {}".format(num_genes_discard * 100, criteria))
    if criteria == "Variance":
        var = counts.var(axis=1)
        min_genes_spot_var = var.quantile(num_genes_discard)
        if math.isnan(min_genes_spot_var):
            print("Computed variance is NaN! Check your input data.")
        else:
            print("Min normalized variance a gene must have over all spots "
                  "to be kept ({0}% of total) {1}".format(num_genes_discard, min_genes_spot_var))
            counts = counts[var >= min_genes_spot_var]
    elif criteria == "TopRanked":
        sum = counts.sum(axis=1)
        min_genes_spot_sum = sum.quantile(num_genes_discard)
        if math.isnan(min_genes_spot_sum):
            print("Computed sum is NaN! Check your input data.")
        else:
            print("Min normalized total count a gene must have over all spots "
                  "to be kept ({0}% of total) {1}".format(num_genes_discard, min_genes_spot_sum))
            counts = counts[sum >= min_genes_spot_sum]
    else:
        raise RuntimeError("Error, incorrect criteria method\n")
    print("Dropped {} genes".format(num_genes - len(counts.index)))
    return counts.transpose()