import numpy as np
import pandas as pd
from sklearn.metrics import (
    roc_auc_score, 
    average_precision_score
)

class Performancer:

    def __init__(self):
        self.control = {}       # id -> y_true
        self.predictions = []   # [(id, score)]
        self.y_true = []
        self.y_score = []
        self.factors = [] # e.g. "algorithm1", "configset2",...
        self.metric_table = []

    def load_controls(self, ref_array):
        for _id, y in ref_array:
            self.control[_id] = float(y)

    def load_predictions(self, pred_array):
        self.predictions = pred_array

    def _align_y(self):
        y_true, y_score = [], []
        for _id, s in self.predictions:
            if _id in self.control:
                y_true.append(int(self.control[_id]))
                y_score.append(float(s))
        self.y_true = np.asarray(y_true, dtype=int)
        self.y_score = np.asarray(y_score, dtype=float)

    def metrics_table_by_score( self, metrics=("tpr", "fnr", "tnr", "fpr"), descending=True, 
                               score_cut=np.inf, add_factors=False, add_counts=False):
        if len(self.y_true) == 0: self._align_y()

        y_true = np.asarray(self.y_true, dtype=int)
        y_score = np.asarray(self.y_score, dtype=float)

        n = len(y_true)

        order = np.argsort(-y_score) if descending else np.argsort(y_score)
        s = y_score[order]
        y = y_true[order]

        P = int((y == 1).sum())
        N = n - P
        tp_cum = np.cumsum(y == 1)
        fp_cum = np.cumsum(y == 0)

        change_idx = np.r_[np.where(np.diff(s) != 0)[0], n - 1]
        thr = s[change_idx]
        tp = tp_cum[change_idx].astype(float)
        fp = fp_cum[change_idx].astype(float)
        fn = float(P) - tp
        tn = float(N) - fp

        def safe_div(num, den):
            return np.divide(num, den, out=np.zeros_like(num, dtype=float), where=den > 0)

        P_f = float(P)
        N_f = float(N)
        T_f = float(n)

        # Primary metrics
        tpr = safe_div(tp, P_f)
        fnr = safe_div(fn, P_f)
        tnr = safe_div(tn, N_f)
        fpr = safe_div(fp, N_f)
        # Secondary metrics
        prec = safe_div(tp, tp + fp)
        f1 = safe_div(2 * tp, 2 * tp + fp + fn)
        acc = safe_div(tp + tn, T_f)
        bal_acc = 0.5 * (tpr + tnr)
        denom = np.sqrt((tp + fp) * (tp + fn) * (tn + fp) * (tn + fn))
        mcc = safe_div(tp * tn - fp * fn, denom)
        coverage = safe_div(tp + fp, T_f)
        cols = {
            "score": thr.astype(float),
            "tpr": tpr, "fnr": fnr, "tnr": tnr, "fpr": fpr,
            "prec": prec, "f1": f1,
            "acc": acc, "bal_acc": bal_acc,
            "mcc": mcc,
            "coverage": coverage,
        }
        if add_counts:
            cols.update({"tp": tp, "fp": fp, "fn": fn, "tn": tn})

        df = pd.DataFrame(cols)

        keep = ["score"]
        for m in metrics:
            if m not in df.columns:
                raise KeyError(f"Metric '{m}' not available. Availables: {list(df.columns)}")
            keep.append(m)
        if add_counts:
            keep += ["tp", "fp", "fn", "tn"]

        df = df[keep]

        if score_cut is None or score_cut is np.inf:
            pass 
        else:
            k = int(score_cut)
            if k > 0 and k < len(df):
                idx = np.linspace(0, len(df) - 1, num=k, dtype=int)
                idx = np.unique(idx)
                df = df.iloc[idx].reset_index(drop=True)

        df = df.sort_values("score", ascending=not descending, ignore_index=True)
        if add_factors and self.factors:
            for i, f in enumerate(self.factors, start=1):
                df[f"factor{i}"] = f

        self.metric_table = df

    def get_auc_metrics(self, descending=True):
        if len(self.y_true) == 0: self._align_y()
        if len(np.unique(self.y_true)) != 2:
            return {"roc_auc": np.nan, "pr_auc": np.nan}
        c = 1 if descending else -1
        
        return {
            "roc_auc": float(roc_auc_score(self.y_true, c * self.y_score)),
            "pr_auc": float(average_precision_score(self.y_true, c * self.y_score)),
        }

    def get_best_thresholds_from_table(self, best_for=("f1", "mcc")):
        out = {}
        for m in best_for:
            col = self.metric_table[m].to_numpy(dtype=float)
            idx = int(np.nanargmax(col))
            out[f"{m}_best"] = float(self.metric_table.iloc[idx][m])
            out[f"{m}_best_threshold"] = float(self.metric_table.iloc[idx]["score"])
        return out

    def get_summary_metrics(self, best_for=("f1", "mcc"), descending=True, add_factors=False):
        out = {}
        out.update(self.get_auc_metrics(descending))
        out.update(self.get_best_thresholds_from_table(best_for=best_for))
        if add_factors and self.factors:
            for i, f in enumerate(self.factors, start=1):
                out[f"factor{i}"] = f
        return out