from typing import List, Dict, Tuple, Optional
from .utils import to_scipy
import numpy as np
import scipy.sparse
import scipy.sparse.linalg
import scipy.linalg
import scipy.stats
import sklearn.linear_model
import sklearn.isotonic


def minmax(arr: np.array) -> np.array:
    data = np.array(arr)
    xmin = data.min()
    xmax = data.max()
    return (data - xmin) / (xmax - xmin)


def calibrate(scores: np.array,
              labels: np.array,
              method: Optional[str] = None) -> np.array:
    """Wrapper function to calibrate scores with its binary labels

    Parameters:
    -----------
    scores: np.array
        The scores generated by a model. It's assumed that these scores
          are probabilities with values between [0,1]. For example, apply
          min-max-scaling for ratio-scale data types (i.e. score>0.0).

    labels: np.array
        The binary labels that are supposed to be classified by the scores.

    method: str (Default: None)
        The calibration algorithm:
            - 'platt' for Platt-Scaling (Platt, 1999)
            - 'isotonic' for Isotonic Regression (Zadrozny and Elkan, 2002)

    Return:
    -------
    calibrated_scores : np.array
        The predicted probabilities

    References:
    -----------
    Platt, J., 1999. Probabilistic outputs for support vector machines and
        comparisons to regularized likelihood methods.

    Zadrozny, B., Elkan, C., 2002. Transforming classifier scores into
        accurate multiclass probability estimates, in: Proceedings of the
        Eighth ACM SIGKDD International Conference on Knowledge Discovery
        and Data Mining, KDD ’02. Association for Computing Machinery,
        New York, NY, USA, pp. 694–699. https://doi.org/10.1145/775047.775151
    """
    scores = np.array(scores)
    labels = np.array(labels)
    if method == 'platt':
        cls = sklearn.linear_model.LogisticRegression()
        cls.fit(X=scores.reshape(-1, 1), y=labels)
        return cls.predict_proba(scores.reshape(-1, 1))[:, 1]
    elif method == 'isotonic':
        cls = sklearn.isotonic.IsotonicRegression(out_of_bounds='raise')
        cls.fit(X=scores, y=labels)
        return cls.transform(scores)
    else:
        return scores


def rank(dok: Dict[Tuple[str, str], int],
         method: Optional[str] = 'pvalue', **kwargs):
    """Rank and score items based on pairwise comparison frequencies

    Parameters:
    -----------
    dok : Dict[Tuple[str, str], int]
        Count/Frequency data as Dictionary of Keys (DoK)

    method : Optional[str]
        The procedure to compute ranks and scores.
        - 'ratios'
        - 'btl'
        - 'pvalue'
        - 'orme'
        - 'eigen'
        - 'transition'

    Returns:
    --------
    ranked : List[int]
        The array positions to order/sort the original data by indexing.

    ordids : List[int]
        The item IDs in the new order.

    scores : List[float]
        The scores for each item ID. Also sorted in descending order.

    info
        Output depends on the selected method

    Example:
    --------
        import bwsample as bws
        data = (
            ([1, 0, 0, 2], ['A', 'B', 'C', 'D']),
            ([1, 0, 0, 2], ['A', 'B', 'C', 'D']),
            ([2, 0, 0, 1], ['A', 'B', 'C', 'D']),
            ([0, 1, 2, 0], ['A', 'B', 'C', 'D']),
            ([0, 1, 0, 2], ['A', 'B', 'C', 'D']),
        )
        dok, _, _, _ = bws.extract_pairs_batch2(data)
        ranked, ordids, scores, info = bws.ranking(dok, method='pvalue')
    """
    cnt, indices = to_scipy(dok)
    if method in ('ratios'):
        return ranking_maximize_ratios(cnt, indices, **kwargs)
    elif method in ('pvalue'):
        return ranking_minus_pvalues(cnt, indices, **kwargs)
    elif method in ('btl', 'bradley', 'hunter'):
        return ranking_btl(cnt, indices, **kwargs)
    elif method in ('orme'):
        return scoring_orme(cnt, indices, **kwargs)
    elif method in ('eigen'):
        return scoring_eigenvector(cnt, indices, **kwargs)
    elif method in ('transition'):
        return transition_simulation(cnt, indices, **kwargs)
    else:
        raise Exception(f"method='{method}' not available.")


def ranking_maximize_ratios(cnt: scipy.sparse.csr_matrix,
                            indices: List[str],
                            avg: Optional[str] = 'exist',
                            calibration: Optional[str] = 'platt'):
    """Rank items based simple ratios, and calibrate row sums as scores

    Parameters:
    -----------
    cnt : scipy.sparse.dok.dok_matrix
        Quadratic sparse matrix with frequency data

    indices : List[str]
        Identifiers, e.g. UUID4, of each row/column of the `cnt` matrix.

    avg : Optional[str]
        How to compute denominator for averaging.
        - 'all': divide the sum of ratios by the row length
        - 'exist': divide the sum of ratios by the number of ratios in the row

    calibration: str (Default: None)
        The calibrated scores. For 'platt' and 'isotonic' we assume
          `label[i]=rowsum[i]>mean(rowsum)`.

    Returns:
    --------
    ranked : List[int]
        The array positions to order/sort the original data by indexing.

    ordids : List[int]
        The item IDs in the new order.

    scores : List[float]
        The scores for each item ID. Also sorted in descending order.

    info
        The matrix with the ratios

    Example:
    --------
        import bwsample as bws
        data = (
            ([1, 0, 0, 2], ['A', 'B', 'C', 'D']),
            ([1, 0, 0, 2], ['A', 'B', 'C', 'D']),
            ([2, 0, 0, 1], ['A', 'B', 'C', 'D']),
            ([0, 1, 2, 0], ['A', 'B', 'C', 'D']),
            ([0, 1, 0, 2], ['A', 'B', 'C', 'D']),
        )
        dok, _, _, _ = bws.extract_pairs_batch2(data)
        ranked, ordids, scores, ratios = bws.rank(
            dok, method='ratio', avg='exist', calibration='platt')
    """
    # compute ratios
    cnt = cnt.tocsr()
    ratios = cnt + cnt.T
    ratios.data = 1.0 / ratios.data
    ratios = ratios.multiply(cnt)

    # sum rows in DoK matrix
    metric = np.array(ratios.sum(axis=1).flatten())[0]
    # averaging
    if avg == 'all':
        metric /= len(metric)
    elif avg == 'exist':
        ridx, _ = (ratios + ratios.T).nonzero()  # ensure actual 0s are counted
        for i, c in zip(*np.unique(ridx, return_counts=True)):
            metric[i] /= c

    # sort, larger row sums are better
    ranked = np.argsort(-metric)  # maximize
    ordids = np.array(indices)[ranked].tolist()
    scores = metric[ranked]

    # calibrate scores
    if calibration in ('platt', 'isotonic'):
        labels = scores > np.mean(scores)  # TRUE: s>mean(s)
        scores = calibrate(scores, labels, method=calibration)
    elif calibration == 'minmax':
        scores = minmax(scores)

    # done
    return ranked.tolist(), ordids, scores.tolist(), ratios


def ranking_minus_pvalues(cnt: scipy.sparse.csr_matrix,
                          indices: List[str],
                          avg: Optional[str] = 'exist',
                          calibration: Optional[str] = 'platt'):
    """Rank based on p-values of a Chi-Squard tests between reciprocal pairs,
        and calibrate row sums as scores

    Parameters:
    -----------
    cnt : scipy.sparse.dok.dok_matrix
        Quadratic sparse matrix with frequency data

    indices : List[str]
        Identifiers, e.g. UUID4, of each row/column of the `cnt` matrix.

    avg : Optional[str]
        How to compute denominator for averaging.
        - 'all': divide the sum of ratios by the row length
        - 'exist': divide the sum of ratios by the number of ratios in the row

    calibration: str (Default: None)
        The calibrated scores. For 'platt' and 'isotonic' we assume
          `label[i]=rowsum[i]>mean(rowsum)`.

    Returns:
    --------
    ranked : List[int]
        The array positions to order/sort the original data by indexing.

    ordids : List[int]
        The item IDs in the new order.

    scores : List[float]
        The scores for each item ID. Also sorted in descending order.

    info
        The matrix with the `1-p`-values

    Example:
    --------
        import bwsample as bws
        data = (
            ([1, 0, 0, 2], ['A', 'B', 'C', 'D']),
            ([1, 0, 0, 2], ['A', 'B', 'C', 'D']),
            ([2, 0, 0, 1], ['A', 'B', 'C', 'D']),
            ([0, 1, 2, 0], ['A', 'B', 'C', 'D']),
            ([0, 1, 0, 2], ['A', 'B', 'C', 'D']),
        )
        dok, _, _, _ = bws.extract_pairs_batch2(data)
        ranked, ordids, scores, (eigval, eigenvec) = bws.rank(
            dok, method='pvalue', avg='exist', calibration='platt')
    """
    # compute p-values for Nij>Nji or 1
    n, _ = cnt.shape
    P = scipy.sparse.dok_matrix((n, n), dtype=np.float64)
    for i in range(n):
        for j in range(n):
            if i > j:
                f1, f2 = cnt[i, j], cnt[j, i]
                fe = (f1 + f2) / 2  # E: Expected Frequency
                if fe > 0 and f1 != f2:
                    _, pval = scipy.stats.chisquare(
                        f_obs=[f1, f2], f_exp=[fe, fe], ddof=0)
                    if f1 > f2:
                        P[i, j] = 1 - pval
                    else:
                        P[j, i] = 1 - pval

    # sum rows in DoK matrix
    metric = np.array(P.sum(axis=1).flatten())[0]
    # averaging
    if avg == 'all':
        metric /= len(metric)
    elif avg == 'exist':
        ridx, _ = (P + P.T).nonzero()  # ensure actual 0s are counted
        for i, c in zip(*np.unique(ridx, return_counts=True)):
            metric[i] /= c

    # sort, larger row sums are better
    ranked = np.argsort(-metric)  # minimize
    ordids = np.array(indices)[ranked].tolist()
    scores = metric[ranked]

    # calibrate scores
    if calibration in ('platt', 'isotonic'):
        labels = scores > np.mean(scores)  # TRUE: s>mean(s)
        scores = calibrate(scores, labels, method=calibration)
    elif calibration == 'minmax':
        scores = minmax(scores)

    # done
    return ranked.tolist(), ordids, scores.tolist(), P


def scoring_orme(cnt: scipy.sparse.csr_matrix,
                 indices: List[str],
                 calibration: Optional[str] = 'platt'):
    """Scoring based on Orme (2009)

    Parameters:
    -----------
    cnt : scipy.sparse.dok.dok_matrix
        Quadratic sparse matrix with frequency data

    indices : List[str]
        Identifiers, e.g. UUID4, of each row/column of the `cnt` matrix.

    calibration: str (Default: None)
        The calibrated scores. For 'platt' and 'isotonic' we assume
          `label[i]=rowsum[i]>mean(rowsum)`.

    Returns:
    --------
    ranked : List[int]
        The array positions to order/sort the original data by indexing.

    ordids : List[int]
        The item IDs in the new order.

    scores : List[float]
        The scores for each item ID. Also sorted in descending order.

    info
        the ratios from [-1, +1]

    Example:
    --------
        import bwsample as bws
        data = (
            ([1, 0, 0, 2], ['A', 'B', 'C', 'D']),
            ([1, 0, 0, 2], ['A', 'B', 'C', 'D']),
            ([2, 0, 0, 1], ['A', 'B', 'C', 'D']),
            ([0, 1, 2, 0], ['A', 'B', 'C', 'D']),
            ([0, 1, 0, 2], ['A', 'B', 'C', 'D']),
        )
        _, dok_direct, _, _ = bws.extract_pairs_batch2(data)
        ranked, ordids, scores, ratios = bws.rank(
            dok_direct, method='omre', calibration='platt')
    References:
    -----------
    Orme, B., 2009. MaxDiff Analysis: Simple Counting, Individual-Level
      Logit, and HB. https://api.semanticscholar.org/CorpusID:202605777
    """
    # compute ratios
    cnt = cnt.tocsr()
    metric = (cnt - cnt.T).sum(axis=1) / (cnt + cnt.T).sum(axis=1)
    metric = np.array(metric.flatten())[0]

    # sort, larger row sums are better
    ranked = np.argsort(-metric)  # maximize
    ordids = np.array(indices)[ranked].tolist()
    scores = metric[ranked]

    # calibrate scores
    if calibration in ('platt', 'isotonic'):
        labels = scores > 0  # TRUE: s>0
        scores = calibrate(scores, labels, method=calibration)
    elif calibration == 'minmax':
        scores = minmax(scores)

    # done
    return ranked.tolist(), ordids, scores.tolist(), metric


def scoring_eigenvector(cnt: scipy.sparse.csr_matrix,
                        indices: List[str],
                        calibration: Optional[str] = None):
    """Compute the eigenvectors of the pairwise comparison matrix, and
        calibrate eigenvectors as scores.

    Parameters:
    -----------
    cnt : scipy.sparse.dok.dok_matrix
        Quadratic sparse matrix with frequency data

    indices : List[str]
        Identifiers, e.g. UUID4, of each row/column of the `cnt` matrix.

    calibration: str (Default: None)
        The calibrated scores. For 'platt' and 'isotonic' we assume
          `label[i]=eigenvector[i]>0.5`. There is also the option to run
          Min-Max-Scaling (`'minmax'`) but won't recommend using it.

    Returns:
    --------
    ranked : List[int]
        The array positions to order/sort the original data by indexing.

    ordids : List[int]
        The item IDs in the new order.

    scores : List[float]
        The scores for each item ID. Also sorted in descending order.

    info
        (eigval, eigenvec)

    Example:
    --------
        import bwsample as bws
        data = (
            ([1, 0, 0, 2], ['A', 'B', 'C', 'D']),
            ([1, 0, 0, 2], ['A', 'B', 'C', 'D']),
            ([2, 0, 0, 1], ['A', 'B', 'C', 'D']),
            ([0, 1, 2, 0], ['A', 'B', 'C', 'D']),
            ([0, 1, 0, 2], ['A', 'B', 'C', 'D']),
        )
        dok, _, _, _ = bws.extract_pairs_batch2(data)
        ranked, ordids, scores, (eigval, eigenvec) = bws.rank(
            dok, method='eigen', calibration=None)

    References:
    -----------
    Saaty, T.L., 2003. Decision-making with the AHP: Why is the principal
        eigenvector necessary. European Journal of Operational Research 145,
        85–91. https://doi.org/10.1016/S0377-2217(02)00227-8
    """
    # set diagonals to 1
    n = cnt.shape[0]
    cnt = cnt.tolil()
    for i in range(n):
        cnt[i, i] = 1

    # compute "positive reciprocal near consistent pairwise comparison matrix"
    cnt = cnt.tocsr()
    cntT = cnt.T
    cntT.data = 1.0 / cntT.data
    ratios = cnt.multiply(cntT)

    # compute eigenvectors as scores
    eigval, eigenvec = scipy.sparse.linalg.eigs(ratios, k=1)
    metric = np.abs(np.real(eigenvec[:, 0]))

    # sort, larger row sums are better
    ranked = np.argsort(-metric)  # maximize
    ordids = np.array(indices)[ranked].tolist()
    scores = metric[ranked]

    # calibrate scores
    if calibration in ('platt', 'isotonic'):
        labels = scores > .5  # TRUE: s>.5
        scores = calibrate(scores, labels, method=calibration)
    elif calibration == 'minmax':
        scores = minmax(scores)

    # done
    return ranked.tolist(), ordids, scores.tolist(), (eigval, eigenvec)


def transition_simulation(cnt: scipy.sparse.dok.dok_matrix,
                          indices: List[str],
                          n_rounds: Optional[int] = 3,
                          calibration: Optional[str] = 'platt'):
    """Estimate transition matrix of item_i>item_j, simulate the item
        probabilities that are calibrated to scores.

    Parameters:
    -----------
    cnt : scipy.sparse.dok.dok_matrix
        Quadratic sparse matrix with frequency data

    indices : List[str]
        Identifiers, e.g. UUID4, of each row/column of the `cnt` matrix.

    calibration: Optional[str]  (Default: 'platt')
        The calibrated scores. We are predicting transition probabilities
          here, i.e. `SUM[transprob]=1`. Thus, we interpret `transprob[i]>1/N`
          as our true binary label for Platt Scaling (`'platt'`) and Isotonic
          Regression (`'isotonic'`). We don't recommend using Min-Max-Scaling
          (`'minmax'`).

    Returns:
    --------
    ranked : List[int]
        The array positions to order/sort the original data by indexing.

    ordids : List[int]
        The item IDs in the new order.

    scores : List[float]
        The scores for each item ID. Also sorted in descending order.

    info : Tuple
        (x, transmat) `x` is is the predicted/simulated item probability,
          and `transmat` the estimated transition probability matrix.

    Example:
    --------
        import bwsample as bws
        data = (
            ([1, 0, 0, 2], ['A', 'B', 'C', 'D']),
            ([1, 0, 0, 2], ['A', 'B', 'C', 'D']),
            ([2, 0, 0, 1], ['A', 'B', 'C', 'D']),
            ([0, 1, 2, 0], ['A', 'B', 'C', 'D']),
            ([0, 1, 0, 2], ['A', 'B', 'C', 'D']),
        )
        dok, _, _, _ = bws.extract_pairs_batch2(data)
        ranked, ordids, scores, (x, transmat) = bws.rank(
            dok, method='transition', n_rounds=3, calibration='platt')
    """
    n = cnt.shape[0]

    # create generator matrix
    genmat = cnt.T.tolil()
    rowsum = genmat.sum(axis=1)
    for i in range(n):
        genmat[i, i] = -rowsum[i]  # set diagonals to -sum(row)
        genmat[i, :] /= max(1, rowsum[i])

    # compute transition matrix
    transmat = scipy.sparse.linalg.expm(genmat.tocsc())

    # simulation: transition from an item1 to the next item2
    #   that is most likely "item2 > item1"
    x = np.ones(n) / n
    for i in range(n_rounds):
        x = x * transmat

    # sort, larger state probabilities are better
    ranked = np.argsort(-x)  # maximize
    ordids = np.array(indices)[ranked].tolist()
    scores = x[ranked]

    # calibrate scores
    if calibration in ('platt', 'isotonic'):
        labels = scores > 1.0 / len(scores)  # TRUE: s>1/N
        scores = calibrate(scores, labels, method=calibration)
    elif calibration == 'minmax':
        scores = minmax(scores)

    # done
    return ranked.tolist(), ordids, scores.tolist(), (x, transmat)


def mle_btl_sparse(cnt: scipy.sparse.csr_matrix,
                   x0: Optional[np.array] = None,
                   max_iter: Optional[int] = 50,
                   tol: Optional[float] = 1e-5) -> (np.array, bool):
    """MLE by Hunter (2004, p.386-387)

    Parameters:
    -----------
    cnt : scipy.sparse.dok.dok_matrix
        Quadratic sparse matrix with frequency data

    x0 : np.array
        Initial values.

    max_iter : int
        maximum number of iterations

    tol : float
        termination criteria

    Returns:
    --------
    gamma : np.array
        Estimated gamma parameters. SUM[gammas]=1. The estimated
          parameters can be used as scores

    flag : bool
        True if solution was found within `max_iter` optimization steps.
        False if not.

    References:
    -----------
    Hunter, D.R., 2004. MM algorithms for generalized Bradley-Terry models.
      The Annals of Statistics 32, 384–406.
      https://doi.org/10.1214/aos/1079120141
    """
    # ensure CSR format
    cnt = cnt.tocsr()
    m = cnt.shape[0]

    # set initial values
    if x0 is None:
        x = np.ones(m) / m
    else:
        x = np.array(x0)
        x = x / x.sum()

    # rowsum, i.e. the number of wins `W_i`
    rowsum = cnt.sum(axis=1)

    # copy sparse structure for `yi + yj`
    gamij = cnt.copy()
    ridx, cidx = gamij.nonzero()

    for k in range(max_iter):
        # assign new weights
        gamij[(ridx, cidx)] = x[ridx] + x[cidx]

        # don't forget to use the inverse: [(Nij + Nji) / (yi + yj)]^{-1}
        cntij = (cnt + cnt.T)
        cntij.data = 1 / cntij.data
        tmp = gamij.multiply(cntij)

        # sum up rows, and elementwise multiply
        gamk = np.multiply(rowsum, tmp.sum(axis=1))

        # normalize to `gam_i^(k)`
        x1 = gamk / gamk.sum(axis=0)

        # abort
        if np.linalg.norm(x1 - x, ord=np.inf) < tol:
            return np.array(x1.flatten())[0], True

        # update
        x = x1

    # last result
    return np.array(x1.flatten())[0], False


def ranking_btl(cnt: scipy.sparse.csr_matrix,
                indices: List[str],
                calibration: Optional[str] = 'platt',
                prefit: Optional[bool] = True,
                max_iter: Optional[int] = 50,
                tol: Optional[float] = 1e-5):
    """Bradley-Terry-Luce (BTL) probability model for pairwise comparisons

    Parameters:
    -----------
    cnt : scipy.sparse.dok.dok_matrix
        Quadratic sparse matrix with frequency data

    indices : List[str]
        Identifiers, e.g. UUID4, of each row/column of the `cnt` matrix.

    calibration: str (Default: None)
        The calibrated scores. For 'platt' and 'isotonic' we assume
          `label[i]=mleparams[i]>1/N`. There is also the option to run
          Min-Max-Scaling (`'minmax'`) but won't recommend using it.

    prefit : bool
        flag to prefit parameters with 'ratio' method
        (see `ranking_maximize_ratios`)

    max_iter : int  (see `mle_btl_sparse`)
        maximum number of iterations

    tol : float  (see `mle_btl_sparse`)
        termination criteria

    Returns:
    --------
    ranked : List[int]
        The array positions to order/sort the original data by indexing.

    ordids : List[int]
        The item IDs in the new order.

    scores : List[float]
        The scores for each item ID. Also sorted in descending order.

    info
        x are the estimated MLE parameters that can be used for scoring

    Example:
    --------
        import bwsample as bws
        data = (
            ([1, 0, 0, 2], ['A', 'B', 'C', 'D']),
            ([1, 0, 0, 2], ['A', 'B', 'C', 'D']),
            ([2, 0, 0, 1], ['A', 'B', 'C', 'D']),
            ([0, 1, 2, 0], ['A', 'B', 'C', 'D']),
            ([0, 1, 0, 2], ['A', 'B', 'C', 'D']),
        )
        dok, _, _, _ = bws.extract_pairs_batch2(data)
        ranked, ordids, scores, x = bws.rank(
            dok, method='btl', calibration='platt')
    """
    cnt = cnt.tocsr()
    x0 = None
    if prefit:
        ratios = cnt + cnt.T
        ratios.data = 1.0 / ratios.data
        ratios = ratios.multiply(cnt)
        x0 = np.array(ratios.sum(axis=1).flatten())[0]
        x0 = minmax(x0)

    # estimate Bradley-Terry-Luce model parameters as metric
    x, flag = mle_btl_sparse(cnt, x0=x0, max_iter=max_iter, tol=tol)

    # sort, larger state probabilities are better
    ranked = np.argsort(-x)  # maximize
    ordids = np.array(indices)[ranked].tolist()
    scores = x[ranked]

    # calibrate scores
    if calibration in ('platt', 'isotonic'):
        labels = scores > 1.0 / len(scores)  # TRUE: s>1/N
        scores = calibrate(scores, labels, method=calibration)
    elif calibration == 'minmax':
        scores = minmax(scores)

    # done
    return ranked.tolist(), ordids, scores.tolist(), x
