# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Threshold setter for binary classifiers."""

from collections.abc import Mapping, Sequence
import math
from typing import Optional

import attr
from lit_nlp.api import components as lit_components
from lit_nlp.api import dataset as lit_dataset
from lit_nlp.api import model as lit_model
from lit_nlp.api import types
from lit_nlp.components import metrics
from lit_nlp.lib import utils
import numpy as np


JsonDict = types.JsonDict
IndexedInput = types.IndexedInput
Spec = types.Spec


@attr.s(auto_attribs=True, kw_only=True)
class ThresholderConfig(object):
  """Config options for Thresholder component."""
  # Ratio of cost of a false negative to a false positive.
  cost_ratio: Optional[float] = 1
  # Facets of datapoints to calculate individual thresholds for.
  facets: Optional[Mapping[str, JsonDict]] = {'': {}}


class Thresholder(lit_components.Interpreter):
  """Determines optimal thresholds for classifiers given constraints."""

  def __init__(self):
    self.metrics_gen = metrics.BinaryConfusionMetrics()

    # Set up the fairness measure calculators, which take confusion matrix
    # metrics and calculates a score from them.
    def demo_parity_metric(stats):
      return ((stats['TP'] + stats['FP']) /
              (stats['TP'] + stats['FP'] + stats['TN'] + stats['FN']))

    def equal_acc_metric(stats):
      return ((stats['TP'] + stats['TN']) /
              (stats['TP'] + stats['FP'] + stats['TN'] + stats['FN']))

    def equal_opp_metric(stats):
      return stats['TP'] / (stats['TP'] + stats['FN'])

    self.fairness_measures = {
        'Demographic parity': demo_parity_metric,
        'Equal accuracy': equal_acc_metric,
        'Equal opporitunity': equal_opp_metric,
    }

  def is_compatible(self, model: lit_model.Model,
                    dataset: lit_dataset.Dataset) -> bool:
    del dataset  # Unused by Thresholder
    return utils.spec_contains(model.output_spec(), types.MulticlassPreds)

  def threshold_to_margin(self, thresh):
    # Convert between margin and classification threshold when displaying
    # margin as a threshold, as is done for binary classifiers.
    # Threshold is between 0 and 1 and represents the minimum score of the
    # positive (non-null) class before a datapoint is classified as positive.
    # A margin of 0 is the same as a threshold of .5 - meaning we take the
    # argmax class. A negative margin is a threshold below .5. Margin ranges
    # from -5 to 5, and can be converted the threshold through the equation
    # margin = ln(threshold / (1 - threshold)).
    if thresh == 0:
      return -5
    if thresh == 1:
      return 5
    return math.log(thresh / (1 - thresh))

  def get_cost(self, metrics_output, cost_ratio):
    return metrics_output['FP'] * cost_ratio + metrics_output['FN']

  def get_thresholds_for_pred_key(self, pred_key, i, margins_to_try,
                                  dataset_results, faceted_results, config):
    # Find the optimal threshold for the entire dataset
    metrics_list = [result[i]['metrics'] for result in dataset_results]
    costs = [self.get_cost(metrics_for_threshold, config.cost_ratio)
             for metrics_for_threshold in metrics_list]
    single_threshold = np.argmin(costs) / 100

    if not (facets := config.facets):
      return {'pred_key': pred_key, 'thresholds': {}}

    faceted_thresholds = {}
    faceted_costs = {}
    faceted_measures = {}

    # If there is only a single facet, return the single best threshold for
    # this prediction key.
    facets_keys = list(facets.keys())
    if len(facets_keys) == 1:
      faceted_thresholds[facets_keys[0]] = {
          'Single': single_threshold,
      }
      return {'pred_key': pred_key, 'thresholds': faceted_thresholds}

    # Loop through all facets of the dataset.
    for facet_key in facets_keys:
      # Find the optimal threshold for each facet individually.
      faceted_metrics_list = [
          result[i]['metrics'] for result in faceted_results[facet_key]]
      costs = [self.get_cost(metrics_for_threshold, config.cost_ratio)
               for metrics_for_threshold in faceted_metrics_list]
      ind_threshold = np.argmin(costs) / 100
      faceted_thresholds[facet_key] = {
          'Single': single_threshold,
          'Individual': ind_threshold
      }

      # Store the error costs and fairness measure scores for each measure for
      # the current facet across all possible margin values.
      measures = {}
      for measure_key in self.fairness_measures:
        measures[measure_key] = [
            self.fairness_measures[measure_key](metrics_for_threshold)
            for metrics_for_threshold in faceted_metrics_list]
      faceted_costs[facet_key] = costs
      faceted_measures[facet_key] = measures

    # Follows the same logic used in https://github.com/PAIR-code/what-if-tool
    # for calculating these thresholds.
    #
    # For all fairness measures:
    #   For all margins for first facet:
    #     For all other facets:
    #        Find margin with closest fairness measure of first facet at
    #        current margin
    #     Calculate overall cost for these margin settings settings across the
    #     facets.
    #   Save the margin settings that correspond to the lowest overall cost
    #   across faceted results.
    first_facet = facets_keys[0]
    for measure_key in self.fairness_measures:
      first_facet_costs = []
      first_facet_thresholds = []
      for threshold_idx in range(len(margins_to_try)):
        first_facet_measure = faceted_measures[
            first_facet][measure_key][threshold_idx]
        cost = faceted_costs[first_facet][threshold_idx]
        cur_thresholds = [threshold_idx / 100]
        for facet_to_check in facets_keys[1:]:
          distances_to_first_facet_measure = [
              abs(measure - first_facet_measure)
              for measure in faceted_measures[facet_to_check][measure_key]]
          threhold_idx_for_facet = np.argmin(distances_to_first_facet_measure)
          cost += faceted_costs[facet_to_check][threhold_idx_for_facet]
          cur_thresholds.append(threhold_idx_for_facet / 100)
        first_facet_costs.append(cost)
        first_facet_thresholds.append(cur_thresholds)
      min_measure_thresholds_idx = np.argmin(first_facet_costs)
      measure_thresholds = first_facet_thresholds[min_measure_thresholds_idx]
      for t_idx, facet_key in enumerate(facets_keys):
        faceted_thresholds[facet_key][measure_key] = measure_thresholds[t_idx]

    return {'pred_key': pred_key, 'thresholds': faceted_thresholds}

  def run(
      self,
      inputs: Sequence[JsonDict],
      model: lit_model.Model,
      dataset: lit_dataset.Dataset,
      model_outputs: Optional[list[JsonDict]] = None,
      config: Optional[JsonDict] = None) -> Optional[list[JsonDict]]:
    """Calculates optimal thresholds on the provided data.

    Args:
      inputs: all examples in the dataset, in the indexed input format.
      model: the model being explained.
      dataset: the dataset which the current examples belong to.
      model_outputs: optional model outputs from calling model.predict(inputs).
      config: a config which should specify TresholderConfig options.

    Returns:
      A JsonDict containing the calcuated thresholds
    """
    # Convert any facets from using IndexedInputs to using JsonDicts prior to
    # passing to self.run().
    if config and (facets := config.get('facets')):
      for facet_dict in facets.values():
        if (facet_data := facet_dict.get('data')):
          facet_dict['data'] = [ex.get('data', ex) for ex in facet_data]

    config = ThresholderConfig(**(config or {}))

    pred_keys = []
    for pred_key, pred_spec in model.output_spec().items():
      if not isinstance(pred_spec, types.MulticlassPreds):
        continue

      if not (parent_key := pred_spec.parent):
        continue

      parent_spec: Optional[types.LitType] = dataset.spec().get(parent_key)
      if self.metrics_gen.is_field_compatible(pred_spec, parent_spec):
        pred_keys.append(pred_key)

    indexed_outputs = {
        ex['_id']: output for (ex, output) in zip(inputs, model_outputs)
    }

    # Try all margins for thresholds from 0 to 1, by hundreths.
    margins_to_try = [
        self.threshold_to_margin(t) for t in np.linspace(0, 1, 101)]

    # Get binary classification metrics for all margins, for the entire
    # dataset, and also for each facet specified in the config.
    dataset_results = []
    faceted_results = {}
    # Loop over each margin/threshold to check.
    for margin in margins_to_try:
      # Set up an empty config to pass to the metrics generator.
      metrics_config = {}
      for pred_key in pred_keys:
        metrics_config[pred_key] = {'': {'margin': margin}}

      # Get and store the metrics for the entire dataset for this margin.
      dataset_results.append(self.metrics_gen.run(
          list(inputs), model, dataset, model_outputs, metrics_config))

      # Get and store the metrics for each facet of the dataset for this margin.
      if not (facets := config.facets):
        continue

      for facet_key, facet_dict in facets.items():
        if not (facet_data := facet_dict.get('data')):
          continue

        if facet_key not in faceted_results:
          faceted_results[facet_key] = []
        faceted_outputs = [indexed_outputs[ex['_id']] for ex in facet_data]
        faceted_results[facet_key].append(
            self.metrics_gen.run(
                facet_data, model, dataset, faceted_outputs, metrics_config
            )
        )

    pred_keys = [result['pred_key'] for result in dataset_results[0]]
    ret = []

    # Find threshold information for each prediction key.
    for i, pred_key in enumerate(pred_keys):
      ret.append(self.get_thresholds_for_pred_key(
          pred_key, i, margins_to_try, dataset_results, faceted_results,
          config))
    return ret
