import numpy as np

from inferelator.postprocessing.precision_recall import RankSummaryPR
from inferelator.postprocessing import (
    TARGET_COLUMN,
    REGULATOR_COLUMN,
    CONFIDENCE_COLUMN,
    F1_COLUMN,
    PRECISION_COLUMN,
    RECALL_COLUMN
)

from ._plot_fix import plt


class RankSummaryF1(RankSummaryPR):
    """
    This class extends RankSumming and calculates F1 score
    """

    name = "F1"
    curve_file_name = "f1_curve.pdf"

    # F1 values
    f1 = None

    # Plotter function

    @property
    def optconff1(self):
        return RankSummaryF1.calculate_opt_conf_f1(self.filtered_data)

    @property
    def maxf1(self):
        return self.calculate_opt_f1(self.filtered_data)

    def __init__(
        self,
        rankable_data,
        gold_standard,
        filter_method='keep_all_gold_standard'
    ):
        super(RankSummaryPR, self).__init__(
            rankable_data,
            gold_standard,
            filter_method=filter_method
        )

        # Calculate the precision and recall and store them with confidence
        # data
        self.filtered_data = self.calculate_precision_recall(
            self.filtered_data.copy(),
            transform_ties='mean'
        )
        self.filtered_data = self.calculate_f1(self.filtered_data.copy())

        # Join the filtered F1 score onto the full confidences
        join_data = self.filtered_data.loc[
            :,
            [TARGET_COLUMN, REGULATOR_COLUMN, F1_COLUMN]
        ].set_index(
            [TARGET_COLUMN, REGULATOR_COLUMN]
        )
        self.confidence_data = self.confidence_data.join(
            join_data,
            on=[TARGET_COLUMN, REGULATOR_COLUMN]
        )

    def score(self):
        return self.name, self.maxf1

    def curve_dataframe(self):
        return self.filtered_data[[CONFIDENCE_COLUMN, F1_COLUMN]]

    def output_curve(self, ax=None, figsize=(6, 4)):
        if ax is None:
            fig, ax = plt.subplots(figsize=figsize, constrained_layout=True)
            ax = fig.add_subplot(1, 1, 1)

        # Extract the recall and precision data
        curve = self.curve_dataframe()
        self.plot_f1_conf(
            curve[F1_COLUMN].values,
            curve[CONFIDENCE_COLUMN].values,
            self.maxf1,
            self.optconff1,
            ax,
            num_edges=(
                self.confidence_data[CONFIDENCE_COLUMN] >= self.optconff1
            ).sum())

        return ax

    @staticmethod
    def plot_f1_conf(f1, conf, optf1, optconf, ax, num_edges=None):

        num_edges = np.sum(conf >= optconf) if num_edges is None else num_edges

        # Generate a plot
        ax.plot(conf, f1)
        ax.set_xlabel('Confidence')
        ax.set_xlim(1, 0)
        ax.set_ylim(0, 1)
        ax.set_ylabel('F1')
        ax.vlines(
            float(optconf),
            0,
            1,
            transform=ax.get_xaxis_transform(),
            colors='r',
            linestyles='dashed'
        )

        ax.annotate(
            f"max F1 = {optf1:.4f}\n"
            f"optimal conf = {optconf:.4f}\n"
            f"num_edges = {num_edges}",
            xy=(0.4, 0.075),
            xycoords='axes fraction'
        )

        return ax

    @staticmethod
    def calculate_opt_f1(data):

        return np.nanmax(data[F1_COLUMN])

    @staticmethod
    def calculate_opt_conf_f1(data):

        return data.loc[
            data[F1_COLUMN] >= np.max(data[F1_COLUMN]),
            CONFIDENCE_COLUMN
        ].min()

    @staticmethod
    def calculate_f1(data):

        data[F1_COLUMN] = RankSummaryF1.pr_to_f1(
            data[PRECISION_COLUMN],
            data[RECALL_COLUMN]
        )
        return data

    @staticmethod
    def pr_to_f1(prec, recall):
        f1 = 2 * (prec * recall) / (prec + recall)
        f1[np.isnan(f1)] = 0.
        return f1
