import itertools as it

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

from . import perturbation


def update(model, attributions, X, y, masker, sort_order, perturbation_method, scores):
    metric = perturbation_method + ' ' + sort_order
    sp = perturbation.SequentialPerturbation(model, masker, sort_order, perturbation_method)
    xs, ys, auc = sp.model_score(attributions, X, y=y)
    scores['metrics'].append(metric)
    scores['values'][metric] = [xs, ys, auc]

def get_benchmark(model, attributions, X, y, masker, metrics):
    # convert dataframes
    if isinstance(X, (pd.Series, pd.DataFrame)):
        X = X.values
    if isinstance(masker, (pd.Series, pd.DataFrame)):
        masker = masker.values

    # record scores per metric
    scores = {'metrics': list(), 'values': dict()}
    for sort_order, perturbation_method in list(it.product(metrics['sort_order'], metrics['perturbation'])):
        update(model, attributions, X, y, masker, sort_order, perturbation_method, scores)

    return scores

def get_metrics(benchmarks, selection):
    # select metrics to plot using selection function
    explainer_metrics = set()
    for explainer in benchmarks:
        scores = benchmarks[explainer]
        if len(explainer_metrics) == 0:
            explainer_metrics = set(scores['metrics'])
        else:
            explainer_metrics = selection(explainer_metrics, set(scores['metrics']))

    return list(explainer_metrics)

def trend_plot(benchmarks):
    explainer_metrics = get_metrics(benchmarks, lambda x, y: x.union(y))

    # plot all curves if metric exists
    for metric in explainer_metrics:
        plt.clf()

        for explainer in benchmarks:
            scores = benchmarks[explainer]
            if metric in scores['values']:
                x, y, auc = scores['values'][metric]
                plt.plot(x, y, label=f'{round(auc, 3)} - {explainer}')

        if 'keep' in metric:
            xlabel = 'Percent Unmasked'
        if 'remove' in metric:
            xlabel = 'Percent Masked'

        plt.ylabel('Model Output')
        plt.xlabel(xlabel)
        plt.title(metric)
        plt.legend()
        plt.show()

def compare_plot(benchmarks):
    explainer_metrics = get_metrics(benchmarks, lambda x, y: x.intersection(y))
    explainers = list(benchmarks.keys())
    num_explainers = len(explainers)
    num_metrics = len(explainer_metrics)

    # dummy start to evenly distribute explainers on the left
    # can later be replaced by boolean metrics
    aucs = dict()
    for i in range(num_explainers):
        explainer = explainers[i]
        aucs[explainer] = [i/(num_explainers-1)]

    # normalize per metric
    for metric in explainer_metrics:
        max_auc, min_auc = -float('inf'), float('inf')

        for explainer in explainers:
            scores = benchmarks[explainer]
            _, _, auc = scores['values'][metric]
            min_auc = min(auc, min_auc)
            max_auc = max(auc, max_auc)

        for explainer in explainers:
            scores = benchmarks[explainer]
            _, _, auc = scores['values'][metric]
            aucs[explainer].append((auc-min_auc)/(max_auc-min_auc))

    # plot common curves
    ax = plt.gca()
    for explainer in explainers:
        plt.plot(np.linspace(0, 1, len(explainer_metrics)+1), aucs[explainer], '--o')

    ax.tick_params(which='major', axis='both', labelsize=8)

    ax.set_yticks([i/(num_explainers-1) for i in range(0, num_explainers)])
    ax.set_yticklabels(explainers, rotation=0)

    ax.set_xticks(np.linspace(0, 1, num_metrics+1))
    ax.set_xticklabels([' '] + explainer_metrics, rotation=45, ha='right')

    plt.grid(which='major', axis='x', linestyle='--')
    plt.tight_layout()
    plt.ylabel('Relative Performance of Each Explanation Method')
    plt.xlabel('Evaluation Metrics')
    plt.title('Explanation Method Performance Across Metrics')
    plt.show()
