import argparse
import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics import precision_recall_curve, average_precision_score
from matplotlib_venn import venn2, venn3

##################################### Functions #############################################


def plot_confusion_matrix(cm, target_names, cmap=None, normalize=False):
    """Function that plots the confusion matrix given cm. Mattias Ohlsson's code extended."""

    import itertools

    accuracy = np.trace(cm) / float(np.sum(cm))
    misclass = 1 - accuracy

    if cmap is None:
        cmap = plt.get_cmap("Blues")

    fig = plt.figure(figsize=(4, 3))
    plt.imshow(cm, interpolation="nearest", cmap=cmap)
    plt.colorbar()

    if target_names is not None:
        tick_marks = np.arange(len(target_names))
        plt.xticks(tick_marks, target_names, rotation=0, fontsize=12)
        plt.yticks(tick_marks, target_names, fontsize=12)

    if normalize:
        cm = cm.astype("float") / cm.sum(axis=1)[:, np.newaxis]

    thresh = cm.max() / 1.5 if normalize else cm.max() / 2
    for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):
        if normalize:
            plt.text(
                j,
                i,
                "{:0.4f}".format(cm[i, j]),
                horizontalalignment="center",
                color="white" if cm[i, j] > thresh else "black",
                fontsize=14,
            )
        else:
            plt.text(
                j,
                i,
                "{:,}".format(cm[i, j]),
                horizontalalignment="center",
                color="white" if cm[i, j] > thresh else "black",
                fontsize=14,
            )

    plt.tight_layout()
    plt.ylabel("True label", fontsize=14)
    plt.xlabel(
        "Predicted label\naccuracy={:0.4f}; misclass={:0.4f}".format(
            accuracy, misclass
        ),
        fontsize=14,
    )

    return fig


def classify_associations(target_file, assoc_tuples):
    self_assoc = 0  # Self associations
    found_assoc_dict = {}
    false_assoc_dict = {}
    tp_fp = np.array([[0, 0]])
    with open(target_file, "r") as f:
        for line in f:
            if line[0] != "f":
                splitline = line.strip().split("\t")
                feat_a = splitline[2]
                feat_b = splitline[3]
                score = abs(float(splitline[5]))
                if feat_a == feat_b:  # Self associations will not be counted
                    self_assoc += 1
                else:
                    if (feat_a, feat_b) in assoc_tuples:
                        found_assoc_dict[(feat_a, feat_b)] = score
                        if (
                            feat_b,
                            feat_a,
                        ) not in found_assoc_dict.keys():  # If we had not found it yet
                            tp_fp = np.vstack((tp_fp, tp_fp[-1] + [0, 1]))
                    elif (feat_a, feat_b) not in assoc_tuples:
                        false_assoc_dict[(feat_a, feat_b)] = score
                        if (feat_b, feat_a) not in false_assoc_dict.keys():
                            tp_fp = np.vstack((tp_fp, tp_fp[-1] + [1, 0]))

    # Remove duplicated associations:
    for i, j in list(found_assoc_dict.keys()):
        if (j, i) in found_assoc_dict.keys():
            del found_assoc_dict[
                (j, i)
            ]  # remove the weakest direction for the association

    for i, j in list(false_assoc_dict.keys()):
        if (j, i) in false_assoc_dict.keys():
            del false_assoc_dict[(i, j)]

    return self_assoc, found_assoc_dict, false_assoc_dict, tp_fp


def create_confusion_matrix(n_feat, associations, real_assoc, false_assoc):
    cm = np.empty((2, 2))
    # TN: only counting the upper half matrix (non doubled associations)
    cm[0, 0] = (n_feat * n_feat - n_feat) / 2 - (
        associations + false_assoc
    )  # Diagonal is discarded
    cm[0, 1] = false_assoc
    cm[1, 0] = associations - real_assoc
    cm[1, 1] = real_assoc

    return cm


def get_precision_recall(found_assoc_dict, false_assoc_dict, associations):
    y_true = []
    y_pred = []

    # True Positives
    for score in found_assoc_dict.values():
        y_true.append(1)
        y_pred.append(score)
    # False Positives
    for score in false_assoc_dict.values():
        y_true.append(0)
        y_pred.append(score)
    # False negatives
    for _ in range(associations - len(found_assoc_dict)):
        y_true.append(1)
        y_pred.append(0)

    precision, recall, thr = precision_recall_curve(
        y_true, y_pred
    )  # thr will tell us score values
    avg_prec = average_precision_score(y_true, y_pred)

    return precision, recall, thr, avg_prec


def plot_precision_recall(precision, recall, avg_prec, label, ax):
    ax.scatter(
        recall,
        precision,
        lw=0,
        marker=".",
        s=5,
        edgecolors="none",
        label=f"{label} - APS:{avg_prec:.2f}",
    )
    ax.legend()
    return ax


def plot_thr_recall(thr, recall, label, ax):
    ax.scatter(recall[:-1], thr, lw=0, marker=".", s=5, edgecolors="none", label=label)
    ax.legend()
    return ax


def plot_TP_vs_FP(tp_fp, label, ax):
    ax.scatter(tp_fp[:, 0], tp_fp[:, 1], s=2, label=label, edgecolors="none")
    ax.legend()
    return ax


def plot_filling_order(order_list, last_rank=None):

    if last_rank is None:
        last_rank = len(order_list)
    fig = plt.figure()
    order_img = np.zeros((np.max(order_list), len(order_list)))
    for i, element in enumerate(order_list):
        order_img[element - 1, i:] = 1

    plt.imshow(order_img[:last_rank, :], cmap="binary")
    plt.xlabel("Correct prediction number")
    plt.ylabel("Association ranking")
    plt.plot(np.arange(last_rank), np.arange(last_rank))
    return fig


def plot_effect_size_matching(
    assoc_tuples_dict, found_assoc_dict, label, ALGORITHM, ax
):

    ground_truth_effects = [
        assoc_tuples_dict[key] for key in list(found_assoc_dict.keys())
    ]
    predicted_effects = np.array(list(found_assoc_dict.values()))

    if ALGORITHM == "ttest":
        # Eq 15 on https://doi.org/10.1146/annurev-statistics-031017-100307
        predicted_effects = [-np.log10(p) if p != 0 else -1 for p in predicted_effects]
        predicted_effects[predicted_effects == -1] = np.max(
            predicted_effects
        )  # Change zeros for max likelihood, -1 as dummy value
        predicted_effects = np.array(predicted_effects)

    max, min = np.max(predicted_effects), np.min(predicted_effects)
    standarized_pred_effects = (predicted_effects - min) / (max - min)
    ax.scatter(
        ground_truth_effects,
        standarized_pred_effects,
        s=12,
        edgecolors="none",
        label=label,
    )
    ax.legend()
    return ax


def plot_venn_diagram(venn, ax, mode="all", scale="log"):
    sets = [set(venn[key][mode]) for key in list(venn.keys())]
    labels = (key for key in list(venn.keys()))

    if len(venn) == 2:
        venn2(sets, labels, ax=ax)
    elif len(venn) == 3:
        venn3(sets, labels, ax=ax)
    else:
        raise ValueError("Unsupported number of input files.")


def plot_upsetplot(venn, assoc_tuples):
    from upsetplot import UpSet
    import pandas as pd
    from matplotlib import cm

    all_assoc = set(
        [
            association
            for ALGORITHM in venn.keys()
            for association in venn[ALGORITHM]["all"]
        ]
    )
    columns = ["ground truth"]
    columns.extend([ALGORITHM for ALGORITHM in list(venn.keys())])

    df = {}
    for association in all_assoc:
        df[association] = []

        if association in assoc_tuples:
            df[association].append("TP")
        else:
            df[association].append("FP")

        for ALGORITHM in list(venn.keys()):
            if association in venn[ALGORITHM]["all"]:
                df[association].append(1)
            else:
                df[association].append(0)

    df = pd.DataFrame.from_dict(df, orient="index", columns=columns)
    df = df.set_index([pd.Index(df[ALGORITHM] == 1) for ALGORITHM in list(venn.keys())])
    upset = UpSet(df, intersection_plot_elements=0, show_counts=True)

    upset.add_stacked_bars(
        by="ground truth",
        colors=cm.Pastel1,
        title="Count by ground truth value",
        elements=10,
    )

    return upset


###################################### Main code ################################################

parser = argparse.ArgumentParser(
    description="Read two files with ground truth associations and predicted associations."
)
parser.add_argument(
    "-p",
    "--perturbed",
    metavar="pert",
    type=str,
    required=True,
    help="perturbed feature names",
)
parser.add_argument(
    "-n",
    "--features",
    metavar="feat",
    type=int,
    required=True,
    help=" total number of features",
)
parser.add_argument(
    "-r",
    "--reference",
    metavar="ref",
    type=str,
    required=True,
    help="path to the ground truth associations file",
)
parser.add_argument(
    "-t",
    "--targets",
    metavar="tar",
    type=str,
    required=True,
    nargs="+",
    help="path to the predicted associations files",
)
args = parser.parse_args()


# Defining main performance evaluation figures:
fig_0, ax_0 = plt.subplots(figsize=(7, 7))
fig_1, ax_1 = plt.subplots(figsize=(7, 7))
fig_2, ax_2 = plt.subplots()
fig_3, ax_3 = plt.subplots()

assoc_tuples_dict = {}

# Reading the file with the ground truth changes:
with open(args.reference, "r") as f:
    for line in f:
        if line[0] != "f" and line[0] != "n":
            splitline = line.strip().split("\t")
            feat_a = splitline[2]
            feat_b = splitline[3]
            assoc_strength = abs(float(splitline[4]))
            # Only can detect associations with perturbed features
            if args.perturbed in feat_a or args.perturbed in feat_b:
                assoc_tuples_dict[(feat_a, feat_b)] = assoc_strength
                assoc_tuples_dict[(feat_b, feat_a)] = assoc_strength

associations = int(len(assoc_tuples_dict) / 2)
venn = {}
# Count and save found associations
for target_file in args.targets:

    ALGORITHM = target_file.split("/")[-1].split("_")[3][:-4]
    self_assoc, found_assoc_dict, false_assoc_dict, tp_fp = classify_associations(
        target_file, list(assoc_tuples_dict.keys())
    )
    real_assoc = len(found_assoc_dict)  # True predicted associations
    false_assoc = len(false_assoc_dict)  # False predicted associations
    total_assoc = real_assoc + false_assoc

    venn[ALGORITHM] = {}
    venn[ALGORITHM]["correct"] = list(found_assoc_dict.keys())
    venn[ALGORITHM]["all"] = list(found_assoc_dict.keys()) + list(
        false_assoc_dict.keys()
    )

    # Assess ranking of associations (they are doubled in assoc_tuples):
    order_list = [
        list(assoc_tuples_dict.keys()).index((feat_a, feat_b)) // 2
        for (feat_a, feat_b) in list(found_assoc_dict.keys())
    ]
    fig = plot_filling_order(order_list)
    fig.savefig(f"Order_image_{ALGORITHM}.png", dpi=200)

    ax_0 = plot_effect_size_matching(
        assoc_tuples_dict, found_assoc_dict, ALGORITHM, ALGORITHM, ax_0
    )

    # Plot confusion matrix:
    cm = create_confusion_matrix(args.features, associations, real_assoc, false_assoc)
    fig = plot_confusion_matrix(
        cm, ["No assoc", "Association"], cmap=None, normalize=False
    )

    fig.savefig(f"Confusion_matrix_{ALGORITHM}.png", dpi=100, bbox_inches="tight")

    # Plot precision-recall and TP-FP curves
    precision, recall, thr, avg_prec = get_precision_recall(
        found_assoc_dict, false_assoc_dict, associations
    )

    ax_1 = plot_precision_recall(precision, recall, avg_prec, ALGORITHM, ax_1)
    ax_2 = plot_TP_vs_FP(tp_fp, ALGORITHM, ax_2)
    ax_3 = plot_thr_recall(thr, recall, ALGORITHM, ax_3)

    # Write results:
    with open("Performance_evaluation_summary_results.txt", "a") as f:
        f.write(f" File:  {target_file}\n")
        f.write(
            f"Ground truth detectable associations (i.e. involving perturbed feature,{args.perturbed}):{associations}\n"
        )
        f.write(
            f"{total_assoc} unique associations found\n{self_assoc} self-associations were found before filtering\n{real_assoc} were real associations\n{false_assoc} were either false or below the significance threshold\n"
        )
        # print("Correct associations:\n", found_assoc_tuples, "\n")
        f.write(
            f"Sensitivity:{real_assoc}/{associations} = {real_assoc/associations}\n"
        )
        f.write(f"Precision:{real_assoc}/{total_assoc} = {(real_assoc)/total_assoc}\n")
        f.write(f"Order list:{order_list}\n\n")
        f.write("______________________________________________________\n")


# Edit figures: layout
ax_0.set_xlabel("Real effect")
ax_0.set_ylabel("Predicted effect")
ax_0.set_ylim((-0.02, 1.02))
ax_0.set_xlim((0, 1.02))
ax_0.legend(
    loc="upper center", bbox_to_anchor=(0.5, 1.1), ncol=3, fancybox=True, shadow=True
)

ax_1.set_xlabel("Recall")
ax_1.set_ylabel("Precision")
ax_1.legend()
ax_1.set_ylim((0, 1.05))
ax_1.set_xlim((0, 1.05))

ax_2.set_xlabel("False Positives")
ax_2.set_ylabel("True Positives")
ax_2.set_aspect("equal")

ax_3.set_ylabel("Threshold")
ax_3.set_xlabel("Recall")


# Save main figures:
fig_0.savefig("Effect_size_matchin.png", dpi=200)
fig_1.savefig("Precision_recall.png", dpi=200)
fig_2.savefig("TP_vs_FP.png", dpi=200)
fig_3.savefig("thr_vs_recall.png", dpi=200)

# Plotting venn diagram:
if len(venn) == 2 or len(venn) == 3:
    fig_v, ax_v = plt.subplots()
    ax_v = plot_venn_diagram(venn, ax_v, mode="correct")
    fig_v.savefig("Venn_diagram.png", dpi=200)

# Plotting UpSet plot
upset = plot_upsetplot(venn, list(assoc_tuples_dict.keys()))
upset.plot()
plt.savefig("UpSet.png", dpi=200)