# -*- coding: utf-8 -*-

import pandas as pd
import numpy as np
from itertools import repeat
from typing import Type, Optional, List, Tuple
from numba import *
import logging

from .rnkdb import RankingDatabase
from .genesig import GeneSignature, Regulon


__all__ = ["recovery", "aucs", "enrichment4features", "enrichment4cells", "leading_edge4row"]


LOGGER = logging.getLogger(__name__)


def derive_rank_cutoff(auc_threshold, total_genes, rank_threshold: Optional[int] = None):
    """

    """
    if not rank_threshold:
        rank_threshold = total_genes - 1

    assert 0 < rank_threshold < total_genes, \
        "Rank threshold must be an integer between 1 and {0:d}".format(total_genes)
    assert 0.0 < auc_threshold <= 1.0, "AUC threshold must be a fraction between 0.0 and 1.0"

    # In the R implementation the cutoff is rounded.
    rank_cutoff = int(round(auc_threshold * total_genes))
    assert 0 < rank_cutoff <= rank_threshold, \
        "An AUC threshold of {0:f} corresponds to {1:d} top ranked genes/regions in the database. " \
        "Please increase the rank threshold or decrease the AUC threshold.".format(auc_threshold, rank_cutoff)
    # Make sure we have exacly the same AUC values as the R-SCENIC pipeline.
    # In the latter the rank threshold is not included in AUC calculation.
    rank_cutoff -= 1
    return rank_cutoff


# Do not use numba as it dwarfs the performance.
def rcc2d(rankings: np.ndarray, weights: np.ndarray, rank_threshold: int) -> np.ndarray:
    """
    Calculate recovery curves.

    :param rankings: The features rankings for a gene signature (n_features, n_genes).
    :param weights: The weights of these genes.
    :return: Recovery curves (n_features, rank_threshold).
    """
    n_features = rankings.shape[0]
    rccs = np.empty(shape=(n_features, rank_threshold)) # Pre-allocation.
    for row_idx in range(n_features):
        curranking = rankings[row_idx, :]
        rccs[row_idx, :] = np.cumsum(np.bincount(curranking, weights=weights)[:rank_threshold])
    return rccs


def recovery(rnk: pd.DataFrame, total_genes: int, weights: np.ndarray, rank_threshold: int, auc_threshold: float,
             no_auc=False) -> (np.ndarray, np.ndarray):
    """
    Calculate recovery curves and AUCs. This is the workhorse of the recovery algorithm.

    :param rnk: A dataframe containing the rank number of genes of interest. Columns correspond to genes.
    :param total_genes: The total number of genes ranked.
    :param weights: the weights associated with the selected genes.
    :param rank_threshold: The total number of ranked genes to take into account when creating a recovery curve.
    :param auc_threshold: The fraction of the ranked genome to take into account for the calculation of the
        Area Under the recovery Curve.
    :return: A tuple of numpy arrays. The first array contains the recovery curves (n_features/n_cells x rank_threshold),
        the second array the AUC values (n_features/n_cells).
    """
    rank_cutoff = derive_rank_cutoff(auc_threshold, total_genes, rank_threshold)
    features, genes, rankings = rnk.index.values, rnk.columns.values, rnk.values
    weights = np.insert(weights, len(weights), 0.0)
    n_features = len(features)
    rankings = np.append(rankings, np.full(shape=(n_features, 1), fill_value=total_genes), axis=1)

    # Calculate recovery curves.
    rccs = rcc2d(rankings, weights, rank_threshold)
    if no_auc:
        return rccs, np.array([])

    # Calculate AUC.
    # For reason of generating the same results as in R we introduce an error by adding one to the rank_cutoff
    # for calculationg the maximum AUC.
    maxauc = float((rank_cutoff+1) * weights.sum())
    assert maxauc > 0
    # The rankings are 0-based. The position at the rank threshold is included in the calculation.
    # The maximum AUC takes this into account.
    aucs = rccs[:, :rank_cutoff].sum(axis=1) / maxauc

    return rccs, aucs


def enrichment4cells(rnk_mtx: pd.DataFrame, regulon: Type[GeneSignature], auc_threshold: float = 0.05) -> pd.DataFrame:
    """
    Calculate the enrichment of the regulon for the cells in the ranking dataframe.

    :param rnk_mtx: The ranked expression matrix (n_cells, n_genes).
    :param regulon: The regulon the assess for enrichment
    :param auc_threshold: The fraction of the ranked genome to take into account for the calculation of the
        Area Under the recovery Curve.
    :return:
    """
    total_genes = len(rnk_mtx.columns)
    index = pd.MultiIndex.from_tuples(list(zip(rnk_mtx.index.values, repeat(regulon.name))),
                                      names=["Cell", "Regulon"])

    rnk = rnk_mtx.iloc[:,rnk_mtx.columns.isin(regulon.genes)]
    if rnk.empty or float(len(rnk))/len(regulon) < 0.80:
        LOGGER.warning("Less than 80% of the genes in {} are present in the expression matrix.".format(regulon.name))
        return pd.DataFrame(index=index, data={"AUC": np.zeros(shape=(rnk_mtx.shape[0]), dtype=np.float64)})
    else:
        weights = np.asarray([regulon[gene] if gene in regulon.genes else 1.0 for gene in rnk.columns.values])
        return pd.DataFrame(index=index, data={"AUC": aucs(rnk, total_genes, weights, auc_threshold)})


def enrichment4features(rnkdb: Type[RankingDatabase], gs: Type[GeneSignature], rank_threshold: int = 5000, auc_threshold: float = 0.05) -> pd.DataFrame:
    """
    Calculate AUC and NES for all regulatory features in the supplied database using the genes of the give signature.

    :param rnkdb: The database.
    :param gs: The gene signature the assess for enrichment
    :param rank_threshold: The total number of ranked genes to take into account when creating a recovery curve.
    :param auc_threshold: The fraction of the ranked genome to take into account for the calculation of the
        Area Under the recovery Curve.
    :return: A dataframe containing all information.
    """
    assert rnkdb, "A database must be supplied"
    assert gs, "A gene signature must be supplied"

    # Load rank of genes from database.
    df = rnkdb.load(gs)
    features = df.index.values
    genes = df.columns.values
    rankings = df.values
    weights = np.asarray([gs[gene] for gene in genes])

    rccs, aucs = recovery(df, rnkdb.total_genes, weights, rank_threshold, auc_threshold)
    ness = (aucs - aucs.mean()) / aucs.std()

    # The creation of a dataframe is a severe performance penalty.
    df_nes = pd.DataFrame(index=features,
                          data={("Enrichment", "AUC"): aucs, ("Enrichment", "NES"): ness})
    df_rnks = pd.DataFrame(index=features,
                           columns=pd.MultiIndex.from_tuples(list(zip(repeat("Ranking"), genes))),
                           data=rankings)
    df_rccs = pd.DataFrame(index=features,
                           columns=pd.MultiIndex.from_tuples(list(zip(repeat("Recovery"), np.arange(rank_threshold)))),
                           data=rccs)
    return pd.concat([df_nes, df_rccs, df_rnks], axis=1)


def leading_edge(rcc: np.ndarray, avg2stdrcc: np.ndarray,
                 ranking: np.ndarray, genes: np.ndarray,
                 weights: Optional[np.array] = None) -> Tuple[List[Tuple[str,float]], int]:
    """
    Calculate the leading edge for a given recovery curve.

    :param rcc: The recovery curve.
    :param avg2stdrcc: The average + 2 standard deviation recovery curve.
    :param ranking: The rank numbers of the gene signature for a given regulatory feature.
    :param genes: The genes corresponding to the ranking available in the aforementioned parameter.
    :param weights: The weights for these genes.
    :return: The leading edge returned as a list of tuple. Each tuple associates a gene part of the leading edge with
        its rank or with its importance (if gene signature supplied). In addition, the rank at maximum difference is
        returned
    """
    def critical_point():
        """ Returns (rank_at_max, max_recovery). """
        rank_at_max = np.argmax(rcc - avg2stdrcc)
        return rank_at_max, rcc[rank_at_max]

    def get_genes(rank_at_max):
        sorted_idx = np.argsort(ranking)
        sranking = ranking[sorted_idx]
        gene_ids = genes[sorted_idx]
        # Make sure to include the gene at the leading edge itself. This is different from the i-cisTarget implementation
        # but is inline with the RcisTarget implementation.
        filtered_idx = sranking <= rank_at_max
        filtered_gene_ids = gene_ids[filtered_idx]
        return list(zip(filtered_gene_ids, weights[filtered_idx] if weights is not None else sranking[filtered_idx]))

    rank_at_max, n_recovered_genes = critical_point()
    # noinspection PyTypeChecker
    return get_genes(rank_at_max), rank_at_max


def leading_edge4row(row: pd.Series, avg2stdrcc: np.ndarray, genes: np.ndarray,
                     weights: Optional[np.array] = None) -> pd.Series:
    """
    Calculate the leading edge for a row of a dataframe. Should be used with partial function application to make this
    function amenable to the apply idiom common for dataframes.

    :param row: The row of the dataframe to calculate the leading edge for.
    :param avg2stdrcc: The average + 2 standard deviation recovery curve.
    :param genes: The genes corresponding to the ranking available in the supplied row.
    :param weights: The weights for these genes.
    :return: The leading edge returned as a list of tuple. Each tuple associates a gene part of the leading edge with
        its rank or with its importance (if gene signature supplied).
    """
    return pd.Series(data=leading_edge(row['Recovery'].values, avg2stdrcc, row['Ranking'].values, genes, weights))


# Giving numba a signature makes the code marginally faster but with losing flexibility (only being able to use one
# type of integers used in rankings).
#@jit(signature_or_function=float64(int16[:], int_, float64), nopython=True)
@jit(nopython=True)
def auc1d(ranking, rank_cutoff, max_auc):
    """
    Calculate the AUC of the recovery curve of a single ranking. [DEPRECATED]

    :param ranking: The rank numbers of the genes.
    :param rank_cutoff: The maximum rank to take into account when calculating the AUC.
    :param max_auc: The maximum AUC.
    :return: The normalized AUC.
    """
    # Using concatenate and full constructs required by numba.
    # The rankings are 0-based. The position at the rank threshold is included in the calculation.
    x = np.concatenate((np.sort(ranking[ranking < rank_cutoff]), np.full((1,), rank_cutoff, dtype=np.int_)))
    y = np.arange(x.size - 1) + 1.0
    return np.sum(np.diff(x)*y)/max_auc


@jit(nopython=True)
def weighted_auc1d(ranking, weights, rank_cutoff, max_auc):
    """
    Calculate the AUC of the weighted recovery curve of a single ranking.

    :param ranking: The rank numbers of the genes.
    :param weights: The associated weights.
    :param rank_cutoff: The maximum rank to take into account when calculating the AUC.
    :param max_auc: The maximum AUC.
    :return: The normalized AUC.
    """
    # Using concatenate and full constructs required by numba.
    # The rankings are 0-based. The position at the rank threshold is included in the calculation.
    filter_idx = ranking < rank_cutoff
    x = ranking[filter_idx]
    y = weights[filter_idx]
    sort_idx = np.argsort(x)
    x = np.concatenate((x[sort_idx], np.full((1,), rank_cutoff, dtype=np.int_)))
    y = y[sort_idx].cumsum()
    return np.sum(np.diff(x)*y)/max_auc


def auc2d(rankings, weights, rank_cutoff, max_auc):
    """
    Calculate the AUCs of multiple rankings.

    :param ranking: The rankings.
    :param auc_func: The 1d AUC function to use.
    :param rank_cutoff: The maximum rank to take into account when calculating the AUC.
    :param max_auc: The maximum AUC.
    :return: The normalized AUCs.
    """
    n_features = rankings.shape[0]
    aucs = np.empty(shape=(n_features,), dtype=np.float64) # Pre-allocation.
    for row_idx in range(n_features):
        aucs[row_idx] = weighted_auc1d(rankings[row_idx, :], weights, rank_cutoff, max_auc)
    return aucs


def aucs(rnk: pd.DataFrame, total_genes: int, weights: np.ndarray, auc_threshold: float) -> np.ndarray:
    """
    Calculate AUCs (implementation without calculating recovery curves first).

    :param rnk: A dataframe containing the rank number of genes of interest. Columns correspond to genes.
    :param total_genes: The total number of genes ranked.
    :param weights: the weights associated with the selected genes.
    :param auc_threshold: The fraction of the ranked genome to take into account for the calculation of the
        Area Under the recovery Curve.
    :return: An array with the AUCs.
    """
    rank_cutoff = derive_rank_cutoff(auc_threshold, total_genes)
    features, genes, rankings = rnk.index.values, rnk.columns.values, rnk.values
    y_max = weights.sum()
    # The rankings are 0-based. The position at the rank threshold is included in the calculation.
    # The maximum AUC takes this into account.
    # For reason of generating the same results as in R we introduce an error by adding one to the rank_cutoff
    # for calculationg the maximum AUC.
    maxauc = float((rank_cutoff+1) * y_max)
    assert maxauc > 0
    return auc2d(rankings, weights, rank_cutoff, maxauc)
