#!/usr/bin/env python

import numpy as np

import pandas as pd

pd.set_option('display.max_colwidth', -1)

from functools import partial

from typing import List

from sklearn.metrics import confusion_matrix, roc_auc_score

from scipy import stats

def _mean_confidence_interval(data, conf=0.95, decimal=3):

  assert(conf > 0 and conf < 1), f"Confidence interval must be within (0, 1). It is {conf}"

  a = 1.0 * np.array(data)

  n = len(a)

  m, se = np.mean(a), stats.sem(a)

  h = se * stats.t.ppf((1 + conf) / 2., n-1)

  return np.round(m, decimal), np.round(m-h, decimal), np.round(m+h, decimal)

class BinaryAvgMetrics(object):

  def __init__(self, targets: List[int], predictions: List[int], probs: List[float], decimal=3) -> None:

    assert (len(targets) == len(predictions) == len(probs)), f"Target list (length = {len(targets)}), predictions list (length = {len(predictions)}) and probabilities list (length = {len(probs)}) must all be of the same length!))"

    self.targs = targets

    self.n_runs = len(self.targs)

    self.preds = predictions

    self.probs = probs

    self.decimal = 3

    self.cms = np.zeros((len(self.targs), 2, 2), dtype=np.int64)

    for i, (targ, pred) in enumerate(zip(self.targs, self.preds)):

      self.cms[i] = confusion_matrix(targ, pred)  

  @property

  def tns(self):

    return self.cms[:, 0, 0]

  @property

  def fps(self):

    return self.cms[:, 0, 1]

  @property

  def fns(self):

    return self.cms[:, 1, 0]

  @property

  def tps(self):

    return self.cms[:, 1, 1]

  @property

  def cm_avg(self):

    return np.ceil(np.array([[self.tns.mean(), self.fps.mean()], [self.fns.mean(), self.tps.mean()]])).astype(np.int64)

  @property

  def prevalence_avg(self):

    return np.round(((self.fns + self.tps) / (self.tns + self.fps + self.fns + self.tps)).mean(), self.decimal)

  def sensitivities(self):

    return self.tps / (self.tps + self.fns)

  def sensitivity_avg(self, conf=None):

    se = (self.tps / (self.tps + self.fns))

    if conf is not None:

      return _mean_confidence_interval(se, conf)

    return np.round(se.mean(), self.decimal,)

  def specificities(self):

    return self.tns / (self.tns + self.fps)

  def specificity_avg(self, conf=None):

    sp = (self.tns / (self.tns + self.fps))

    if conf is not None:

      return _mean_confidence_interval(sp, conf)

    return np.round(sp.mean(), self.decimal)

  def ppvs(self):

    return self.tps / (self.tps + self.fps)

  def ppv_avg(self, conf=None):

    ppv = (self.tps / (self.tps + self.fps))

    if conf is not None:

      return _mean_confidence_interval(ppv, conf)

    return np.round(ppv.mean(), self.decimal)  

  def npvs(self):

    return self.tns / (self.tns + self.fns)

  def npv_avg(self, conf=None):

    npv = (self.tns / (self.tns + self.fns))

    if conf is not None:

      return _mean_confidence_interval(npv, conf)

    return np.round(npv.mean(), self.decimal)

  def f1s(self):

    return (2 * self.sensitivities() * self.ppvs()) / (self.sensitivities() + self.ppvs())

  def f1_avg(self, conf=None):

    se = (self.tps / (self.tps + self.fns))

    ppv = (self.tps / (self.tps + self.fps))

    f1 = (2 * se * ppv) / (se + ppv)

    if conf is not None:

      return _mean_confidence_interval(f1, conf)

    return np.round(f1.mean(), self.decimal)

  def aurocs(self):

    return np.array([roc_auc_score(targ, prob) for targ, prob in zip(self.targs, self.probs)])

  def auroc_avg(self, conf=None):

    auroc = np.array([roc_auc_score(targ, prob) for targ, prob in zip(self.targs, self.probs)])

    if conf is not None:

      return _mean_confidence_interval(auroc, conf)

    return np.round(auroc.mean(), self.decimal)

  def get_avg_metrics(self, conf=None, defn=False):

    definitions = {

      'sensitivity': "When it's ACTUALLY YES, how often does it PREDICT YES?",

      'specificity': "When it's ACTUALLY NO, how often does it PREDICT NO?",

      'ppv': "When it PREDICTS YES, how often is it correct?",

      'auroc': "Indicates how well the model is capable of distinguishing between classes",

      'npv': "When it PREDICTS NO, how often is it correct?",

      'f1': "Harmonic mean of sensitivity and ppv",

    }

    if conf is None:

      metrics = {

        'sensitivity': [self.sensitivity_avg() * 100],

        'specificity': [self.specificity_avg() * 100],

        'ppv': [self.ppv_avg() * 100],

        'auroc': [self.auroc_avg() * 100],

        'npv': [self.npv_avg() * 100],

        'f1': [self.f1_avg() * 100],

      }

      if defn:

        for metric, value in metrics.items():

          value.append(definitions[metric])

        d = pd.DataFrame(metrics.values(), index=metrics.keys(), columns=['Value', 'Definition'])

      else:

        d = pd.DataFrame(metrics.values(), index=metrics.keys(), columns=['Value'])

      return d

    else:

      metrics = {

        'sensitivity': [*[value * 100 for value in self.sensitivity_avg(conf)]],        

        'specificity': [*[value * 100 for value in self.specificity_avg(conf)]],

        'ppv': [*[value * 100 for value in self.ppv_avg(conf)]],

        'auroc': [*[value * 100 for value in self.auroc_avg(conf)]],   

        'npv': [*[value * 100 for value in self.npv_avg(conf)]],

        'f1': [*[value * 100 for value in self.f1_avg(conf)]],        

      }

      if defn:

        for metric, value in metrics.items():

          value.append(definitions[metric])

        d = pd.DataFrame(metrics.values(), index=metrics.keys(), columns=['Mean', 'Lower', 'Upper', 'Definition'])

      else:

        d = pd.DataFrame(metrics.values(), index=metrics.keys(), columns=['Mean', 'Lower', 'Upper'])

      return d

  def __repr__(self):

    s = f"Number of Runs: {self.n_runs}\n"

    return s

  def __len__(self):

    return len(self.targs)

def get_best_model(bam: BinaryAvgMetrics, fnames: List[str]):

  best_se, best_se_model = 0, None

  best_sp, best_sp_model = 0, None

  best_ppv, best_ppv_model = 0, None

  best_auroc, best_auroc_model = 0, None

  best_npv, best_npv_model = 0, None

  best_f1, best_f1_model = 0, None

  for i in range(bam.n_runs):

    se = bam.tps[i] / (bam.tps[i] + bam.fns[i])

    sp = bam.tns[i] / (bam.tns[i] + bam.fps[i])

    ppv = bam.tps[i] / (bam.tps[i] + bam.fps[i])

    npv = bam.tns[i] / (bam.tns[i] + bam.fns[i])

    f1 = (2 * se * ppv) / (se + ppv)

    if best_se < se:

      best_se = se

      best_se_model = fnames[i]    

    if best_sp < sp:

      best_sp = sp

      best_sp_model = fnames[i]          

    if best_ppv < ppv:

      best_ppv = ppv

      best_ppv_model = fnames[i]    

    if best_npv < npv:

      best_npv = npv

      best_npv_model = fnames[i]  

    if best_f1 < f1:

      best_f1 = f1

      best_f1_model = fnames[i]    

  for i, (targ, prob) in enumerate(zip(bam.targs, bam.probs)):

    auroc = roc_auc_score(targ, prob)

    if best_auroc < auroc:

      best_auroc = auroc

      best_auroc_model = fnames[i]

  d = {

    'sensitivity': [best_se, best_se_model],

    'specificity': [best_sp, best_sp_model],

    'ppv': [best_ppv, best_ppv_model],

    'auroc': [best_auroc, best_auroc_model],

    'npv': [best_npv, best_npv_model],

    'f1': [best_f1, best_f1_model],

  }

  return pd.DataFrame(d.values(), index=d.keys(), columns=['Value', 'Model File'])