import os
from subprocess import getoutput

import numpy as np
import numbers
from matplotlib import pyplot as plt
from matplotlib.font_manager import FontProperties
from sklearn import metrics

DPI = 100


def mpl_fontprop():
    fc_list = [x.strip() for x in getoutput('fc-list :lang=zh').strip().split('\n')]
    fc = None
    for item in fc_list:
        if 'CN' in item:
            fc = item.strip().split()[0][:-1]
    if fc:
        fontp = FontProperties(fname=fc)
    else:
        fontp = FontProperties()
    return fontp


def smooth_line(x, y):
    newx = np.linspace(np.min(x), np.max(x), 100)
    newy = np.interp(newx, x, y)
    return newx, newy


def optimize_curve(a, b):
    optimal_idxs = np.where(np.r_[True,
                                  np.logical_or(np.diff(a, 2),
                                                np.diff(b, 2)),
                                  True])[0]
    a = a[optimal_idxs]
    b = b[optimal_idxs]
    return a, b, optimal_idxs


def roc_curve_ci(auc, pos_num, neg_num, ci=0.95):
    from scipy import stats as st
    q0 = auc * (1 - auc)
    q1 = auc / (2 - auc) - auc ** 2
    q2 = 2 * (auc ** 2) / (1 + auc) - auc ** 2
    se = np.sqrt((q0 + (pos_num - 1) * q1 + (neg_num - 1) * q2) / (pos_num * neg_num))
    alpha = (1 - ci) / 2
    z_crit = st.norm.ppf(1 - alpha)
    lower = auc - z_crit * se
    upper = auc + z_crit * se
    lower = max(lower, 0.)
    upper = min(upper, 1.)
    return lower, upper


def plot_curve(xs, ys, labels, xlabel='x', ylabel='y', title=None, save_path=None, diagonal=False, smooth=True, fscore=False,
               pos_nums=None, neg_nums=None):
    fig = plt.figure(dpi=DPI)
    fontprop = mpl_fontprop()
    ax = fig.add_subplot(111)
    ax.set_xlabel(xlabel, fontproperties=fontprop)
    ax.set_ylabel(ylabel, fontproperties=fontprop)
    if title is not None:
        ax.set_title(title, fontproperties=fontprop)
    auc_key = 'AUC' if diagonal else 'AP'
    for i, (x, y) in enumerate(zip(xs, ys)):
        y = np.asarray(y)
        x = np.asarray(x)
        # ROC
        auc_l = np.nan_to_num(metrics.auc(x, y))
        if diagonal == True:
            # x = np.concatenate([np.array([0]), x])
            # y = np.concatenate([np.array([0]), y])
            y[x == 0] = 0
        else:
            # sortx
            order = x.argsort()[::-1]
            x = x[order]
            y = y[order]
            y[x == 1] = 0
        x, y, thresh = optimize_curve(x, y)
        # sort end
        if smooth:
            x, y = smooth_line(x, y)
        if pos_nums is not None and neg_nums is not None:
            lower_l, upper_l = roc_curve_ci(auc_l, pos_nums[i], neg_nums[i])
            plt.plot(x, y, label='{}: {} {:0.3f} \n(95%CI:{:0.3f}-{:0.3f})'.format(labels[i], auc_key, auc_l, lower_l,
                                                                                   upper_l),
                     linestyle='-',
                     # c='#03DAC5'
                     )
        else:
            plt.plot(x, y, label='{}={:.03f}'.format(labels[i], auc_l), linestyle='-')
    if diagonal:
        plt.plot([0, 1], [0, 1], color='gray', lw=2, linestyle='--', alpha=.8)
    if fscore:
        f_scores = np.linspace(0.2, 0.8, num=4)
        lines, labels = [], []
        for f_score in f_scores:
            x = np.linspace(0.01, 1)
            y = f_score * x / (2 * x - f_score)
            (l,) = plt.plot(x[y >= 0], y[y >= 0], color="gray", alpha=0.2)
            plt.annotate("f1={0:0.1f}".format(f_score), xy=(0.9, y[45] + 0.02))
            ax.set_xlim([0.0, 1.0])
            ax.set_ylim([0.0, 1.05])
    # plt.legend(prop=fontprop, loc=4)
    plt.legend(prop=fontprop, loc='best')
    if save_path is not None:
        plt.savefig(os.path.join(save_path, '{}.png'.format(title)))


def dca(preds, y_trues, labels, title=None):
    bottom = -0.05
    op_bottom = -0.1
    N = len(y_trues[0])
    pos_num = np.sum(y_trues[0])
    neg_num = N - pos_num
    for p, y, l in zip(preds, y_trues, labels):
        threshs = []
        nbs = []
        lower = 1.
        for th in np.unique(p):
            tp = np.sum((p >= th) * y)
            fp = np.sum((p >= th) * (y != 1))
            net_benefit = (tp - fp * th / (1 - th)) / N
            threshs.append(th)
            # nbs.append(net_benefit)
            lower = min(lower, net_benefit)
            if lower >= bottom:
                nbs.append(lower)
            else:
                break

        # remove null thresh curve
        if len(threshs[:len(nbs)]) < 1:
            print(f"{l} dca is invalid")
            continue

        threshs, nbs = smooth_line(threshs[:len(nbs)], nbs)
        plt.plot(threshs, nbs, label=l)

    # operation curve
    xs = np.linspace(0, 0.5, 20).tolist()
    operations = []
    for idx, NRI_th in enumerate(xs):
        NRI = pos_num / N - neg_num * (NRI_th / (1 - NRI_th)) / N
        if NRI >= op_bottom:
            operations.append(NRI)
        else:
            break
    plt.plot(xs[:len(operations)], operations, label='all with operation')
    plt.plot([0, 1], [0, 0], 'k:', label='none operation')
    if title is not None:
        plt.title(title)
    plt.legend()


def binary_tfpn(match, scores):
    desc_score_indices = np.argsort(scores, kind="mergsort")[::-1]
    match = match[desc_score_indices]
    scores = scores[desc_score_indices]
    # diff means gradient. remove duplicated score values as below
    distinct_value_indices = np.where(np.diff(scores))[0]
    threshold_idxs = np.r_[distinct_value_indices, match.size - 1]
    tps = np.cumsum(match)[threshold_idxs]
    fps = 1 + threshold_idxs - tps
    thresholds = scores[threshold_idxs]
    total_gt_pos = np.sum(match == 1)
    total_gt_neg = np.sum(match == 0)
    fns = total_gt_pos - tps
    tns = total_gt_neg - fps

    return tps, fps, tns, fns, thresholds


def simple_metric(gts, preds, cutoff=0.5):
    """

    :param gts: ground truth [0, 1] array
    :param preds: predictions 0-1 float array
    :param cutoff: could be a number as model threshold or statistics index such as f1-score/g-mean
    :return:
    """
    if not isinstance(preds, np.ndarray):
        preds = np.asarray(preds)
    if not isinstance(gts, np.ndarray):
        gts = np.asarray(gts)
    tps, fps, tns, fns, threshs = binary_tfpn(gts, preds)
    sim_metr = {"tps": tps,
                "fps": fps,
                "tns": tns,
                "fns": fns,
                "sensitivity": tps / (tps + fns),
                "precision": tps / (tps + fps),
                "specificity": tns / (tns + fps),
                "accuracy": (tps+tns)/(tps+tns+fps+fns)}
    sim_metr["f1-score"] = 2 * (sim_metr["precision"] * sim_metr["sensitivity"]) / (
            sim_metr["precision"] + sim_metr["sensitivity"])
    sim_metr["g-mean"] = np.sqrt(sim_metr["sensitivity"] * sim_metr["specificity"])
    sim_metr["youden-index"] = sim_metr["sensitivity"] + sim_metr["specificity"] - 1
    if isinstance(cutoff, numbers.Number):
        target_ind = np.max(np.where(threshs > cutoff)[0])
    elif isinstance(cutoff, str):
        if cutoff in sim_metr.keys():
            target_ind = np.argsort(sim_metr[cutoff], kind="mergsort")[::-1][0]
            # print(f"Set threshold to {threshs[target_ind]}")
            cutoff = threshs[target_ind]
        else:
            raise ValueError(f"{cutoff} is not available. keys are {sim_metr.keys()}")
    else:
        raise ValueError(f"{cutoff} is not available. number of str is required.")
    results = {}
    for k, v in sim_metr.items():
        results[k] = v[target_ind]
    results['youden-index'] = np.max(sim_metr['youden-index'])
    results['auroc'] = metrics.roc_auc_score(gts, preds)
    results['average_precision'] = metrics.average_precision_score(gts, preds)
    return results, cutoff
