import collections
import time
from typing import Any, Dict

from absl import logging
import numpy as np
import robustness_metrics as rm
from sklearn.metrics import accuracy_score
from sklearn.metrics import auc
from sklearn.metrics import log_loss
from sklearn.metrics import precision_recall_curve
from sklearn.metrics import roc_auc_score
from sklearn.metrics import roc_curve
from sklearn.preprocessing import LabelBinarizer
from sklearn.utils import check_array
from sklearn.utils import check_consistent_length
import tensorflow as tf

from . import metric_utils  # local file import
from . import results_storage_utils  # local file import

@tf.function
def eval_step_tf(dataset_iterator, dataset_steps, strategy, estimator,
                 estimator_args, uncertainty_estimator_fn, is_deterministic):
  """Eval step.

  Run TensorFlow model evaluation, using an `uncertainty_estimator_fn`
  to produce predictions and decomposed uncertainty estimates for each example.

  Args:
    dataset_iterator: tf.data.Dataset, dataset on which we will evaluate the
      model.
    dataset_steps: int, number of gradient steps in the dataset.
    strategy: tf.distribute strategy, used to distribute datasets.
    estimator: model wrapped to produce a `tf.Tensor`, predictive mean, with
      shape [B].
    estimator_args: Dict, extra args for the `uncertainty_estimator_fn`, such as
      the number of MC samples, `num_samples`.
    uncertainty_estimator_fn: Callable, method to produce predictive means along
      with various metrics of uncertainty, e.g., predictive_entropy, epistemic
      uncertainty (mutual information).
    is_deterministic: bool, is the model a single deterministic network. In this
      case, we cannot capture epistemic uncertainty.

  Returns:
    Dict, contains `tf.Tensor` predictions, ground truth,
      and uncertainty estimates.
  """
  print('Tracing in `eval_utils.eval_step_tf`.')

  def step_fn(inputs):
    print('Tracing in `eval_utils.eval_step_tf.step_fn`.')
    images = inputs['features']
    labels = inputs['labels']

    # Compute prediction, total, aleatoric, and epistemic uncertainty estimates
    pred_and_uncert = uncertainty_estimator_fn(
        images, estimator, training_setting=False, **estimator_args)

    # Return a tuple
    y_true = labels
    y_pred = pred_and_uncert['prediction']
    y_pred_entropy = pred_and_uncert['predictive_entropy']

    if not is_deterministic:
      y_pred_variance = pred_and_uncert['predictive_variance']
      y_aleatoric_uncert = pred_and_uncert['aleatoric_uncertainty']
      y_epistemic_uncert = pred_and_uncert['epistemic_uncertainty']
    else:
      y_pred_variance = tf.zeros(0)
      y_aleatoric_uncert = tf.zeros(0)
      y_epistemic_uncert = tf.zeros(0)

    return (y_true, y_pred, y_pred_entropy, y_pred_variance, y_aleatoric_uncert,
            y_epistemic_uncert)

  # Containers for storage of
  # predictions, ground truth, uncertainty estimates
  # Construct tf.TensorArrays to store model results
  n_per_core_batches = dataset_steps * strategy.num_replicas_in_sync
  y_true = tf.TensorArray(tf.int32, size=n_per_core_batches)
  y_pred = tf.TensorArray(tf.float32, size=n_per_core_batches)
  y_pred_entropy = tf.TensorArray(tf.float32, size=n_per_core_batches)
  y_pred_variance = tf.TensorArray(tf.float32, size=n_per_core_batches)
  y_aleatoric_uncert = tf.TensorArray(tf.float32, size=n_per_core_batches)
  y_epistemic_uncert = tf.TensorArray(tf.float32, size=n_per_core_batches)

  for i in tf.range(dataset_steps):
    result = strategy.run(step_fn, args=(next(dataset_iterator),))

    # Parse results tuple
    (y_true_, y_pred_, y_pred_entropy_, y_pred_variance_, y_aleatoric_uncert_,
     y_epistemic_uncert_) = result

    # Convert from Per-Replica object to tuple
    if strategy.num_replicas_in_sync > 1:
      y_true_ = y_true_.values
      y_pred_ = y_pred_.values
      y_pred_entropy_ = y_pred_entropy_.values

      if not is_deterministic:
        y_pred_variance_ = y_pred_variance_.values
        y_aleatoric_uncert_ = y_aleatoric_uncert_.values
        y_epistemic_uncert_ = y_epistemic_uncert_.values

      # Iterate through per-batch results
      # This is written in a very un-Pythonic manner to have updates only
      # rely on arguments successfully passed to TPU scope
      for replica_id in tf.range(strategy.num_replicas_in_sync):
        index = (strategy.num_replicas_in_sync * i) + replica_id
        for batch_result in y_true_:
          y_true = y_true.write(index, batch_result)
        for batch_result in y_pred_:
          y_pred = y_pred.write(index, batch_result)
        for batch_result in y_pred_entropy_:
          y_pred_entropy = y_pred_entropy.write(index, batch_result)

        if not is_deterministic:
          for batch_result in y_pred_variance_:
            y_pred_variance = y_pred_variance.write(index, batch_result)
          for batch_result in y_aleatoric_uncert_:
            y_aleatoric_uncert = y_aleatoric_uncert.write(index, batch_result)
          for batch_result in y_epistemic_uncert_:
            y_epistemic_uncert = y_epistemic_uncert.write(index, batch_result)
    else:
      y_true = y_true.write(i, y_true_)
      y_pred = y_pred.write(i, y_pred_)
      y_pred_entropy = y_pred_entropy.write(i, y_pred_entropy_)

      if not is_deterministic:
        y_pred_variance = y_pred_variance.write(i, y_pred_variance_)
        y_aleatoric_uncert = y_aleatoric_uncert.write(i, y_aleatoric_uncert_)
        y_epistemic_uncert = y_epistemic_uncert.write(i, y_epistemic_uncert_)

  results_arrs = {
      'y_true': y_true.stack(),
      'y_pred': y_pred.stack(),
      'y_pred_entropy': y_pred_entropy.stack(),
  }
  if not is_deterministic:
    results_arrs['y_pred_variance'] = y_pred_variance.stack()
    results_arrs['y_aleatoric_uncert'] = y_aleatoric_uncert.stack()
    results_arrs['y_epistemic_uncert'] = y_epistemic_uncert.stack()

  return results_arrs

def evaluate_model_on_datasets(strategy,
                               datasets,
                               steps,
                               estimator,
                               estimator_args,
                               uncertainty_estimator_fn,
                               eval_batch_size,
                               call_dataset_iter,
                               is_deterministic=False,
                               backend='tf',
                               eval_step_jax=None,
                               verbose=False):
  """Eval on dataset.

  Run model evaluation on all provided datasets, using an
  `uncertainty_estimator_fn` to produce predictions and decomposed
  uncertainty estimates for each example.

  Additionally constructs joint dataset predictions, composed of predictions
    on both in-domain and OOD datasets.

  Args:
    strategy: tf.distribute strategy, used to distribute datasets.
    datasets: Dict[str, tf.data.Dataset], datasets on which we evaluate the
      model.
    steps: Dict[str, int], number of gradient steps in each dataset.
    estimator: model wrapped to produce a `tf.Tensor` (if `backend`=='tf') or
      `np.ndarray` (if `backend`=='jax'), predictive mean, with shape [B].
    estimator_args: Dict, extra args for the `uncertainty_estimator_fn`, such as
      the number of MC samples, `num_samples`.
    uncertainty_estimator_fn: Callable, method to produce predictive means along
      with various metrics of uncertainty, e.g., predictive_entropy, epistemic
      uncertainty (mutual information).
    eval_batch_size: int, size of evaluation minibatches.
    call_dataset_iter: bool, if True, should call `iter()` on each dataset. May
      not need if evaluation datasets have been repeated.
    is_deterministic: bool, is the model a single deterministic network. In this
      case, we cannot capture epistemic uncertainty.
    backend: str, in {'tf', 'jax'}, specifies the evaluation method.
    eval_step_jax: Callable, evaluation method used for Jax model.
    verbose: bool, extra logging.

  Returns:
    Dict, for each dataset, contains `np.array` predictions, ground truth,
      and uncertainty estimates.
  """
  # Need to collect these so we can form joint datasets:
  # e.g., joint_test = in_domain_test UNION ood_test
  dataset_split_to_containers = {}

  for dataset_split, dataset in datasets.items():
    # Begin iteration for this dataset split
    start_time = time.time()

    if call_dataset_iter:
      dataset_iterator = iter(dataset)
    else:
      dataset_iterator = dataset

    logging.info(f'Creating iterator took {time.time() - start_time} seconds.')

    dataset_steps = steps[dataset_split]
    logging.info(f'Evaluating split {dataset_split}.')
    if backend == 'jax':
      eval_epoch_arrs = eval_step_jax(dataset_iterator, dataset_steps,
                                      is_deterministic, **estimator_args)
    elif backend == 'tf':
      eval_epoch_arrs = eval_step_tf(dataset_iterator,
                                     tf.convert_to_tensor(dataset_steps),
                                     strategy, estimator, estimator_args,
                                     uncertainty_estimator_fn, is_deterministic)
    else:
      raise NotImplementedError(f'Backend {backend} is not supported yet.')

    # Update metadata
    time_elapsed = time.time() - start_time
    dataset_split_to_containers[dataset_split] = {}
    dataset_split_dict = dataset_split_to_containers[dataset_split]
    dataset_split_dict['total_ms_elapsed'] = time_elapsed * 1e6
    dataset_split_dict['dataset_size'] = dataset_steps * eval_batch_size

    # Use vectorized NumPy containers
    for eval_key, eval_arr in eval_epoch_arrs.items():
      tmp_eval_arr = eval_arr if backend == 'jax' else eval_arr.numpy()

      if tmp_eval_arr.ndim > 1:
        tmp_eval_arr = np.concatenate(tmp_eval_arr).flatten()

      dataset_split_dict[eval_key] = tmp_eval_arr
      if verbose:
        print(f'Concatenated {eval_key} into shape '
              f'{dataset_split_dict[eval_key].shape}')

    dataset_split_dict['y_pred'] = dataset_split_dict['y_pred'].astype(
        'float64')

  # Add Joint Dicts
  dataset_split_to_containers = results_storage_utils.add_joint_dicts(
      dataset_split_to_containers, is_deterministic=is_deterministic)

  return dataset_split_to_containers

def evaluate_model_and_compute_metrics(
    strategy,
    eval_datasets,
    steps,
    metrics,
    eval_estimator,
    uncertainty_estimator_fn,
    eval_batch_size,
    available_splits,
    estimator_args,
    call_dataset_iter,
    is_deterministic=False,
    num_bins=15,
    use_tpu=True,
    return_per_pred_results=False,
    backend='tf',
    eval_step_jax=None):
  """Main.

  Main method for evaluation and computing metrics using TF or Jax
  models. Usable for evaluation during tuning.

  Args:
    strategy: tf.distribute strategy, used to distribute datasets.
    eval_datasets: Dict[str, tf.data.Dataset], datasets on which we evaluate the
      model.
    steps: Dict[str, int], number of gradient steps in each dataset.
    metrics: metrics.
    eval_estimator: model wrapped to produce a `tf.Tensor` (if `backend`=='tf')
      or `np.ndarray` (if `backend`=='jax'), predictive mean, with shape [B].
    uncertainty_estimator_fn: Callable, method to produce predictive means along
      with various metrics of uncertainty, e.g., predictive_entropy, epistemic
      uncertainty (mutual information).
    eval_batch_size: int, size of evaluation minibatches.
    available_splits: List[str], names of the evaluation datasets provided, used
      to log results only for these splits.
    estimator_args: Dict, extra args for the `uncertainty_estimator_fn`, such as
      the number of MC samples, `num_samples`.
    call_dataset_iter: bool, if True, should call `iter()` on each dataset. May
      not need if evaluation datasets have been repeated.
    is_deterministic: bool, is the model a single deterministic network. In this
      case, we cannot capture epistemic uncertainty.
    num_bins: int, number of bins to use with expected calibration error.
    use_tpu: bool, currently exists a bug that disallows collecting ECE during
      training with TPU, this is used to avoid logging that metric.
    return_per_pred_results: bool,
    backend: str, in {'tf', 'jax'}, specifies the evaluation method.
    eval_step_jax: Callable, evaluation method used for Jax model.

  Returns:
    Union[Tuple[Dict, Dict], Dict]
      If return_per_pred_results, return two Dicts. Else, return only the
      second.
      first Dict:
        for each dataset, per-prediction results (e.g., each prediction,
        ground-truth, loss, retention arrays).
      second Dict:
        for each dataset, contains `np.array` predictions, ground truth,
        and uncertainty estimates.
  """
  # Compute predictions on all evaluation datasets
  # When we eval during training we don't need to re-iterate the
  # evaluation datasets
  eval_results = evaluate_model_on_datasets(
      strategy=strategy,
      datasets=eval_datasets,
      steps=steps,
      estimator=eval_estimator,
      estimator_args=estimator_args,
      uncertainty_estimator_fn=uncertainty_estimator_fn,
      eval_batch_size=eval_batch_size,
      call_dataset_iter=call_dataset_iter,
      is_deterministic=is_deterministic,
      backend=backend,
      eval_step_jax=eval_step_jax)

  # For each eval dataset, add NLL and accuracy for each example
  eval_results = compute_loss_and_accuracy_arrs_for_all_datasets(eval_results)

  # Compute all metrics for each dataset --
  # Robustness, Open Set Recognition, Retention AUC
  metrics_results = compute_metrics_for_all_datasets(
      eval_results,
      use_precomputed_arrs=False,
      ece_num_bins=num_bins,
      compute_retention_auc=True,
      verbose=False)

  # Log metrics
  metrics_results = metric_utils.log_epoch_metrics(
      metrics=metrics,
      eval_results=metrics_results,
      use_tpu=use_tpu,
      dataset_splits=available_splits)

  if return_per_pred_results:
    return eval_results, metrics_results
  else:
    return metrics_results

def evaluate_model_on_datasets_np(
    datasets,
    steps,
    estimator,
    estimator_args,
    uncertainty_estimator_fn,
    eval_batch_size,
    is_deterministic,
    np_input=False,
):
  """Main method for evaluation and computing metrics.

  Appropriate for evaluation loops with a single GPU (i.e., no
  distribution strategy), and is framework-agnostic (will work just as
  well with a TF, Jax, or PyTorch model, given an `uncertainty_estimator_fn`
  to cast to NumPy ndarrays.

  Args:
    datasets: Dict[str, tf.data.Dataset], datasets on which we evaluate the
      model.
    steps: Dict[str, int], number of gradient steps in each dataset.
    estimator: model wrapped to produce a `np.ndarray`, predictive mean, with
      shape [B].
    estimator_args: Dict, extra args for the `uncertainty_estimator_fn`, such as
      the number of MC samples, `num_samples`.
    uncertainty_estimator_fn: Callable, method to produce predictive means along
      with various metrics of uncertainty, e.g., predictive_entropy, epistemic
      uncertainty (mutual information).
    eval_batch_size: int, size of evaluation minibatches.
    is_deterministic: bool, is the model a single deterministic network. In this
      case, we cannot capture epistemic uncertainty.
    np_input: bool, True if the model expects a NumPy input; we add a cast.

  Returns:
    Dict, for each dataset, contains `np.array` predictions, ground truth,
      and uncertainty estimates.
  """
  # Need to collect these so we can form joint datasets:
  # e.g., joint_test = in_domain_test UNION ood_test
  dataset_split_to_containers = {}

  for dataset_split, dataset in datasets.items():
    # Containers for numpy storage of
    # image names, predictions, ground truth, uncertainty estimates
    names = list()
    y_true = list()
    y_pred = list()
    y_pred_entropy = list()

    if not is_deterministic:
      y_pred_variance = list()
      y_aleatoric_uncert = list()
      y_epistemic_uncert = list()

    # Begin iteration for this dataset split
    start_time = time.time()
    dataset_iterator = iter(dataset)
    dataset_steps = steps[dataset_split]
    logging.info(f'Evaluating split {dataset_split}.')
    for step in range(dataset_steps):
      if step % 10 == 0:
        logging.info('Evaluated %d/%d batches.', step, dataset_steps)

      inputs = next(dataset_iterator)  # pytype: disable=attribute-error
      images = inputs['features']
      labels = inputs['labels']

      # Compute prediction, total, aleatoric, and epistemic
      # uncertainty estimates
      pred_and_uncert = uncertainty_estimator_fn(
          images._numpy() if np_input else images,  # pylint: disable=protected-access
          estimator,
          training_setting=False,
          **estimator_args)

      # Add this batch of predictions to the containers
      names.append(inputs['name'])
      y_true.append(labels.numpy())
      y_pred.append(pred_and_uncert['prediction'])
      y_pred_entropy.append(pred_and_uncert['predictive_entropy'])
      if not is_deterministic:
        y_pred_variance.append(pred_and_uncert['predictive_variance'])
        y_aleatoric_uncert.append(pred_and_uncert['aleatoric_uncertainty'])
        y_epistemic_uncert.append(pred_and_uncert['epistemic_uncertainty'])

    # Update metadata
    time_elapsed = time.time() - start_time
    dataset_split_to_containers[dataset_split] = {}
    dataset_split_dict = dataset_split_to_containers[dataset_split]
    dataset_split_dict['total_ms_elapsed'] = time_elapsed * 1e6
    dataset_split_dict['dataset_size'] = dataset_steps * eval_batch_size

    dataset_split_dict['names'] = np.concatenate(names).flatten()
    dataset_split_dict['y_true'] = np.concatenate(y_true).flatten()
    dataset_split_dict['y_pred'] = np.concatenate(y_pred).flatten()
    dataset_split_dict['y_pred'] = dataset_split_dict['y_pred'].astype(
        'float64')

    # Use vectorized NumPy containers
    dataset_split_dict['y_pred_entropy'] = (
        np.concatenate(y_pred_entropy).flatten())

    if not is_deterministic:
      dataset_split_dict['y_pred_variance'] = (
          np.concatenate(y_pred_variance).flatten())
      dataset_split_dict['y_aleatoric_uncert'] = (
          np.concatenate(y_aleatoric_uncert).flatten())
      dataset_split_dict['y_epistemic_uncert'] = (
          np.concatenate(y_epistemic_uncert).flatten())

  # Add Joint Dicts
  dataset_split_to_containers = results_storage_utils.add_joint_dicts(
      dataset_split_to_containers, is_deterministic=is_deterministic)

  return dataset_split_to_containers

def eval_model_numpy(datasets,
                     steps,
                     estimator,
                     estimator_args,
                     uncertainty_estimator_fn,
                     eval_batch_size,
                     is_deterministic,
                     distribution_shift,
                     num_bins=15,
                     np_input=False):
  """Main method for evaluation and computing metrics.

  Appropriate for evaluation loops with a single GPU (i.e., no
  distribution strategy), and is framework-agnostic (will work just as
  well with a TF, Jax, or PyTorch model, given an `uncertainty_estimator_fn`
  to cast to NumPy ndarrays.

  Args:
    datasets: Dict[str, tf.data.Dataset], datasets on which we evaluate the
      model.
    steps: Dict[str, int], number of gradient steps in each dataset.
    estimator: model wrapped to produce a `np.ndarray`, predictive mean, with
      shape [B].
    estimator_args: Dict, extra args for the `uncertainty_estimator_fn`, such as
      the number of MC samples, `num_samples`.
    uncertainty_estimator_fn: Callable, method to produce predictive means along
      with various metrics of uncertainty, e.g., predictive_entropy, epistemic
      uncertainty (mutual information).
    eval_batch_size: int, size of evaluation minibatches.
    is_deterministic: bool, is the model a single deterministic network. In this
      case, we cannot capture epistemic uncertainty.
    distribution_shift: which distribution shift to run on.
    num_bins: the number of bins to use for ECE.
    np_input: bool, True if the model expects a NumPy input; we add a cast.

  Returns:
    Dict, for each dataset, contains `np.array` predictions, ground truth,
      and uncertainty estimates.
  """
  eval_results = evaluate_model_on_datasets_np(
      datasets,
      steps,
      estimator,
      estimator_args,
      uncertainty_estimator_fn,
      eval_batch_size,
      is_deterministic=is_deterministic,
      np_input=np_input)

  if distribution_shift == 'aptos':
    # TODO(nband): generalize
    aptos_metadata_path = 'gs://ub-data/aptos/metadata.csv'
    eval_results['ood_test_balanced'] = compute_rebalanced_aptos_dataset(
        aptos_dataset=eval_results['ood_test'],
        aptos_metadata_path=aptos_metadata_path,
        new_aptos_size=10000)

  # For each eval dataset, add NLL and accuracy for each example
  eval_results = compute_loss_and_accuracy_arrs_for_all_datasets(eval_results)

  # Precompute ROC/PR curves, retention and balanced retention curves
  logging.info('Precomputing ROC/PR curves, retention and balanced'
               ' retention curves.')
  eval_results = precompute_arrs_for_all_datasets(eval_results=eval_results)

  logging.info('Computing metrics with precomputed arrs.')
  metrics_results = compute_metrics_for_all_datasets(
      eval_results,
      use_precomputed_arrs=True,
      ece_num_bins=num_bins,
      compute_retention_auc=True,
      verbose=False)

  # Log metrics
  available_splits = [split for split in eval_results.keys()]
  metrics_results = metric_utils.log_epoch_metrics(
      metrics=None,
      eval_results=metrics_results,
      use_tpu=False,
      dataset_splits=available_splits)

  return eval_results, metrics_results

def compute_metrics_for_all_datasets(eval_results,
                                     use_precomputed_arrs,
                                     ece_num_bins=15,
                                     compute_retention_auc=False,
                                     verbose=False):
  """Computes scalar metrics for all datasets.

  Args:
    eval_results: Dict[str, Dict], evaluation results for each dataset.
    use_precomputed_arrs: selects which eval function to use.
    ece_num_bins: int, used to compute expected calibration error metric.
    compute_retention_auc: bool, should compute retention metrics.
    verbose: bool, extra logging.

  Returns:
    Dict[str, Dict], scalar metric results for each dataset.
  """
  dataset_eval_fn = (
      compute_dataset_eval_metrics_with_precomputed_arrs
      if use_precomputed_arrs else compute_dataset_eval_metrics)
  metric_results = {}
  for dataset_key, results_dict in eval_results.items():
    if verbose:
      logging.info(f'Computing metrics for dataset split {dataset_key}.')

    compute_open_set_recognition = 'joint' in dataset_key
    metric_results[dataset_key] = dataset_eval_fn(
        dataset_key,
        results_dict,
        ece_num_bins=ece_num_bins,
        compute_open_set_recognition=compute_open_set_recognition,
        compute_retention_auc=compute_retention_auc)

  return metric_results

def precompute_arrs_for_all_datasets(eval_results, verbose=False):
  """Precompute metric arrays for all datasets, e.g., log loss, retention

  arrays, etc.

  Args:
    eval_results: Dict[str, Dict], evaluation results for each dataset.
    verbose: bool, extra logging.

  Returns:
    Dict[str, Dict], metric arrays for each dataset.
  """
  for dataset_key, results_dict in eval_results.items():
    if verbose:
      logging.info(f'Computing metrics for dataset split {dataset_key}.')

    compute_open_set_recognition = 'joint' in dataset_key
    eval_results[dataset_key] = precompute_metric_arrs(
        results=results_dict,
        compute_open_set_recognition=compute_open_set_recognition)

  return eval_results

def compute_loss_and_accuracy_arrs_for_all_datasets(eval_results):
  """Compute loss and accuracy arrays for each dataset.

  Args:
    eval_results: Dict[str, Dict], evaluation results for each dataset.

  Returns:
    Dict[str, Dict], loss and accuracy arrays for each dataset.
  """
  for dataset_key, results_dict in eval_results.items():
    eval_results[dataset_key] = compute_loss_and_accuracy_arrs(results_dict)

  return eval_results

def compute_loss_and_accuracy_arrs(results):
  """Compute loss and accuracy arrays for a particular dataset results dict.

  Args:
    results: Dict, evaluation results for a single dataset.

  Returns:
    Dict, loss and accuracy arrays for a single dataset.
  """
  results = compute_log_loss_arr(results)
  y_pred, y_true = results['y_pred'], results['y_true']
  results['accuracy_arr'] = y_true == (y_pred > 0.5)
  return results

def compute_log_loss_arr(results, labels=np.asarray([0, 1]), eps=1e-15):
  """Based on sklearn.preprocessing.log_loss, no aggregation.

  Args:
    results: Dict, evaluation results for a single dataset.
    labels: np.ndarray, binary classification task labels.
    eps: float, for numerical stability.

  Returns:
    Dict containing unaggregated log loss `np.ndarray`.
  """
  y_pred, y_true = results['y_pred'], results['y_true']
  y_pred = check_array(y_pred, ensure_2d=False)
  check_consistent_length(y_pred, y_true, None)

  lb = LabelBinarizer()

  if labels is not None:
    lb.fit(labels)
  else:
    lb.fit(y_true)

  if len(lb.classes_) == 1:
    if labels is None:
      raise ValueError('y_true contains only one label ({0}). Please '
                       'provide the true labels explicitly through the '
                       'labels argument.'.format(lb.classes_[0]))
    else:
      raise ValueError('The labels array needs to contain at least two '
                       'labels for log_loss, '
                       'got {0}.'.format(lb.classes_))

  transformed_labels = lb.transform(y_true)

  if transformed_labels.shape[1] == 1:
    transformed_labels = np.append(
        1 - transformed_labels, transformed_labels, axis=1)

  # Clipping
  y_pred = np.clip(y_pred, eps, 1 - eps)

  # If y_pred is of single dimension, assume y_true to be binary
  # and then check.
  if y_pred.ndim == 1:
    y_pred = y_pred[:, np.newaxis]
  if y_pred.shape[1] == 1:
    y_pred = np.append(1 - y_pred, y_pred, axis=1)

  # Check if dimensions are consistent.
  transformed_labels = check_array(transformed_labels)
  if len(lb.classes_) != y_pred.shape[1]:
    if labels is None:
      raise ValueError('y_true and y_pred contain different number of '
                       'classes {0}, {1}. Please provide the true '
                       'labels explicitly through the labels argument. '
                       'Classes found in '
                       'y_true: {2}'.format(transformed_labels.shape[1],
                                            y_pred.shape[1], lb.classes_))
    else:
      raise ValueError('The number of classes in labels is different '
                       'from that in y_pred. Classes found in '
                       'labels: {0}'.format(lb.classes_))

  # Renormalize
  y_pred /= y_pred.sum(axis=1)[:, np.newaxis]
  loss = -(transformed_labels * np.log(y_pred)).sum(axis=1)
  results['nll_arr'] = loss
  return results

def compute_rebalanced_aptos_dataset(aptos_dataset,
                                     aptos_metadata_path,
                                     new_aptos_size=10000):
  """The APTOS dataset has a significantly different underlying distribution

  of clinical severity (e.g., there are far more severe examples) versus
  the EyePACS dataset.

  This might tend to inflate the performance of a model (by providing more
  easily classified "severe" examples) without rebalancing of the dataset
  and metrics.

  Here we match the test empirical distribution of EyePACS by sampling
  from each underlying category with replacement, bringing the total size
  of the dataset to `new_aptos_size`.

  Args:
    aptos_dataset: Dict, predictions and ground truths on the APTOS dataset.
    aptos_metadata_path: str, location of APTOS metadata (names and clinical
      labels of each example).
    new_aptos_size: int, target size of the new rebalanced APTOS dataset.

  Returns:
    Dict, rebalanced APTOS dataset predictions, ground truths, etc.
  """
  # EyePACS Test Data: clinical severity label to proportion of dataset
  label_to_proportion = {
      0: 0.7359503164,
      1: 0.07129130537,
      2: 0.1472228732,
      3: 0.0228966487,
      4: 0.02263885634
  }

  # Load in APTOS metadata
  aptos_metadata_df = results_storage_utils.load_dataframe_gfile(
      file_path=aptos_metadata_path, sep=',')
  name_to_diagnosis = dict(
      zip(aptos_metadata_df['id_code'], aptos_metadata_df['diagnosis']))

  # Determine location of indices corresponding to each diagnosis
  diagnosis_to_indices = collections.defaultdict(list)
  names = aptos_dataset['names']
  for i, name in enumerate(names):
    try:
      name = name.decode('utf-8')
    except UnicodeDecodeError:
      name = str(name)
    diagnosis = int(name_to_diagnosis[name])
    diagnosis_to_indices[diagnosis].append(i)

  # Uniformly sample without replacement to form a new dataset,
  # with approximately same proportions as EyePACS Test Data
  # and total size ~10000
  new_indices = []
  for diagnosis, indices in diagnosis_to_indices.items():
    total_count_in_new_dataset = int(label_to_proportion[diagnosis] *
                                     new_aptos_size)
    new_indices.append(
        np.random.choice(
            indices, size=total_count_in_new_dataset, replace=True))

  new_indices = np.concatenate(new_indices)
  new_aptos_dataset = {}
  for key, value in aptos_dataset.items():
    try:
      new_aptos_dataset[key] = value[new_indices]
    except IndexError:
      new_aptos_dataset[key] = value

  return new_aptos_dataset

def compute_rebalanced_retention_curves(results: Dict[str, Any]):
  """Compute rebalanced retention curves.

  Compute rebalanced retention curves, which are used for joint (ID + OOD)
  tuning. This is done by repeating the OOD indices many times, and then
  bringing the number of OOD indices to the number of ID indices by sampling
  from the OOD indices without replacement.

  Args:
    results: Dict, results for a particular dataset (must be joint to have the
      `is_ood` key, for a binary `np.ndarray`.

  Returns:
    Dict, contains rebalanced retention curves.
  """
  y_pred_entropy, is_ood = results['y_pred_entropy'], results['is_ood']

  # Convert the boolean list to numpy array
  is_ood = np.array(is_ood)

  in_domain_indices = np.where(is_ood == 0)[0]

  n_in_domain = in_domain_indices.shape[0]
  ood_indices = np.where(is_ood == 1)[0]
  n_ood = ood_indices.shape[0]

  # We first tile the OOD indices this many times
  n_ood_repeats = int(n_in_domain / n_ood)

  # To bring the number of OOD indices = the number of ID indices, we sample
  # the necessary indices without replacement
  remaining_n_ood_indices = n_in_domain - (n_ood_repeats * n_ood)

  remaining_ood_indices = np.random.choice(
      ood_indices, size=remaining_n_ood_indices, replace=False)

  # Construct a list of all of the indices to retrieve
  all_indices = (
      in_domain_indices.tolist() + (n_ood_repeats * ood_indices.tolist()) +
      remaining_ood_indices.tolist())

  # Construct predictive entropy, accuracy, NLL arrays with these indices
  rebalanced_y_pred_entropy = y_pred_entropy[all_indices]
  rebalanced_accuracy = results['accuracy_arr'][all_indices]
  rebalanced_nll = results['nll_arr'][all_indices]

  rebalanced_accuracy_retention_curve = compute_retention_curve_on_accuracies(
      accuracies=rebalanced_accuracy, uncertainty=rebalanced_y_pred_entropy)
  rebalanced_nll_retention_curve = compute_retention_curve_on_losses(
      losses=rebalanced_nll, uncertainty=rebalanced_y_pred_entropy)

  y_pred = results['y_pred']
  y_true = results['y_true']
  rebalanced_auroc_retention_curve = compute_auc_retention_curve(
      y_pred=y_pred, y_true=y_true, uncertainty=y_pred_entropy, auc_str='roc')
  rebalanced_auprc_retention_curve = compute_auc_retention_curve(
      y_pred=y_pred, y_true=y_true, uncertainty=y_pred_entropy, auc_str='prc')

  return {
      'accuracy': rebalanced_accuracy_retention_curve,
      'nll': rebalanced_nll_retention_curve,
      'auroc': rebalanced_auroc_retention_curve,
      'auprc': rebalanced_auprc_retention_curve
  }

def compute_rebalanced_retention_scores(results: Dict[str, Any]):
  """Computes rebalanced retention curves, then the mean to get an AUC.

  Args:
    results: Dict, results for a particular dataset (must be joint to have the
      `is_ood` key, for a binary `np.ndarray`.

  Returns:
    Dict, contains the AUCs of the retention curves built on various
      base metrics.
  """
  rebalanced_curves = compute_rebalanced_retention_curves(results)
  return {
      'accuracy': np.mean(rebalanced_curves['accuracy']),
      'nll': np.mean(rebalanced_curves['nll']),
      'auroc': np.mean(rebalanced_curves['auroc']),
      'auprc': np.mean(rebalanced_curves['auprc'])
  }

def precompute_metric_arrs(results, compute_open_set_recognition=False):
  """Compute retention arrays and ROC/PR curves.

  Used for caching to do downstream plots and scalar metrics quickly.

  Args:
    results: dict, results for a particular dataset split.
    compute_open_set_recognition: bool, if True, compute OOD detection PR and
      AUROC metrics.

  Returns:
      Dict, contains retention arrays and ROC/PR curves.
  """
  y_true = results['y_true']
  y_pred = results['y_pred']
  y_pred_entropy = results['y_pred_entropy']

  # Compute ROC curve
  # try:
  results['fpr_arr'], results['tpr_arr'], _ = roc_curve(
      y_true=y_true, y_score=y_pred)
  # except:
  #   pass

  # Compute PR curve
  # try:
  results['precision_arr'], results['recall_arr'], _ = precision_recall_curve(
      y_true=y_true, probas_pred=y_pred)
  # except:
  #   pass

  if compute_open_set_recognition:
    is_ood = results['is_ood']
    results['ood_detection_fpr_arr'], results['ood_detection_tpr_arr'], _ = (
        roc_curve(y_true=is_ood, y_score=y_pred_entropy))
    (results['ood_detection_precision_arr'],
     results['ood_detection_recall_arr'], _) = (
         precision_recall_curve(y_true=is_ood, probas_pred=y_pred_entropy))

    # For the joint datasets, we also compute a rebalanced retention metric,
    # in which we duplicate the OOD dataset to match the size of the in-domain
    # dataset, and then compute the retention metrics.
    ret_curves = compute_rebalanced_retention_curves(results)
    results['balanced_retention_accuracy_arr'] = ret_curves['accuracy']
    results['balanced_retention_nll_arr'] = ret_curves['nll']
    results['balanced_retention_auroc_arr'] = ret_curves['auroc']
    results['balanced_retention_auprc_arr'] = ret_curves['auprc']

  # Retention curves
  assert 'accuracy_arr' in results
  assert 'nll_arr' in results
  results['retention_accuracy_arr'] = compute_retention_curve_on_accuracies(
      accuracies=results['accuracy_arr'], uncertainty=y_pred_entropy)
  results['retention_nll_arr'] = compute_retention_curve_on_losses(
      losses=results['nll_arr'], uncertainty=y_pred_entropy)
  results['retention_auroc_arr'] = compute_auc_retention_curve(
      y_pred=y_pred, y_true=y_true, uncertainty=y_pred_entropy, auc_str='roc')
  results['retention_auprc_arr'] = compute_auc_retention_curve(
      y_pred=y_pred, y_true=y_true, uncertainty=y_pred_entropy, auc_str='prc')

  return results

def compute_dataset_eval_metrics_with_precomputed_arrs(
    dataset_key,
    results,
    ece_num_bins=15,
    compute_open_set_recognition=False,
    compute_retention_auc=False):
  """Compute scalar metrics using cached retention and ROC/PR curves for

  efficiency.

  Args:
    dataset_key: str, name of dataset (prepends each metric in returned Dict).
    results: Dict, results for a particular dataset split including precomputed
      arrays.
    ece_num_bins: int, number of bins used in computing expected calibration
      error.
    compute_open_set_recognition: bool, if True, compute OOD detection PR and
      AUROC metrics.
    compute_retention_auc: bool, if True, compute retention AUC metrics by
      taking the mean of the retention arrays.

  Returns:
      Dict, scalar metrics.
  """
  y_pred = results['y_pred']
  y_true = results['y_true']

  eval_metrics = dict()

  # Standard predictive metrics
  eval_metrics[f'{dataset_key}/negative_log_likelihood'] = log_loss(
      y_pred=y_pred, y_true=y_true, labels=np.asarray([0, 1]))
  if 'fpr_arr' in results.keys() and 'tpr_arr' in results.keys():
    eval_metrics[f'{dataset_key}/auroc'] = auc(results['fpr_arr'],
                                               results['tpr_arr'])
  else:
    eval_metrics[f'{dataset_key}/auroc'] = None

  if 'precision_arr' in results.keys() and 'recall_arr' in results.keys():
    recall = results['recall_arr']
    precision = results['precision_arr']
    eval_metrics[f'{dataset_key}/auprc'] = auc(recall, precision)

  eval_metrics[f'{dataset_key}/accuracy'] = (
      accuracy_score(y_true=y_true, y_pred=(y_pred > 0.5)))

  # Uncertainty metrics
  ece = rm.metrics.ExpectedCalibrationError(num_bins=ece_num_bins)
  ece.add_batch(y_pred, label=y_true)
  eval_metrics[f'{dataset_key}/ece'] = ece.result()['ece']

  if compute_open_set_recognition:
    eval_metrics[f'{dataset_key}/ood_detection_auroc'] = auc(
        results['ood_detection_fpr_arr'], results['ood_detection_tpr_arr'])
    eval_metrics[f'{dataset_key}/ood_detection_auprc'] = auc(
        results['ood_detection_recall_arr'],
        results['ood_detection_precision_arr'])

    eval_metrics[f'{dataset_key}/balanced_retention_accuracy_auc'] = np.mean(
        results['balanced_retention_accuracy_arr'])
    eval_metrics[f'{dataset_key}/balanced_retention_nll_auc'] = np.mean(
        results['balanced_retention_nll_arr'])
    eval_metrics[f'{dataset_key}/balanced_retention_auroc_auc'] = np.mean(
        results['balanced_retention_auroc_arr'])
    eval_metrics[f'{dataset_key}/balanced_retention_auprc_auc'] = np.mean(
        results['balanced_retention_auprc_arr'])
  else:
    # This is added for convenience when logging (so the entry exists
    # in tabular format)
    eval_metrics[f'{dataset_key}/ood_detection_auroc'] = None
    eval_metrics[f'{dataset_key}/ood_detection_auprc'] = None
    eval_metrics[f'{dataset_key}/balanced_retention_accuracy_auc'] = None
    eval_metrics[f'{dataset_key}/balanced_retention_nll_auc'] = None
    eval_metrics[f'{dataset_key}/balanced_retention_auroc_auc'] = None
    eval_metrics[f'{dataset_key}/balanced_retention_auprc_auc'] = None

  if compute_retention_auc:
    assert 'accuracy_arr' in results
    assert 'nll_arr' in results

    eval_metrics[f'{dataset_key}/retention_accuracy_auc'] = np.mean(
        results['retention_accuracy_arr'])
    eval_metrics[f'{dataset_key}/retention_nll_auc'] = np.mean(
        results['retention_nll_arr'])
    eval_metrics[f'{dataset_key}/retention_auroc_auc'] = np.mean(
        results['retention_auroc_arr'])
    eval_metrics[f'{dataset_key}/retention_auprc_auc'] = np.mean(
        results['retention_auprc_arr'])

  return eval_metrics

def compute_dataset_eval_metrics(
    dataset_key,
    results,
    ece_num_bins=15,
    compute_open_set_recognition=False,
    compute_retention_auc=False,
):
  """Compute scalar metrics.

  Args:
    dataset_key: str, name of dataset (prepends each metric in returned Dict).
    results: Dict, results for a particular dataset split.
    ece_num_bins: int, number of bins used in computing expected calibration
      error.
    compute_open_set_recognition: bool, if True, compute OOD detection PR and
      AUROC metrics.
    compute_retention_auc: bool, if True, compute retention AUC metrics by
      taking the mean of the retention arrays.

  Returns:
      Dict, scalar metrics.
  """
  y_pred, y_true, y_pred_entropy = (results['y_pred'], results['y_true'],
                                    results['y_pred_entropy'])

  eval_metrics = dict()

  # Standard predictive metrics
  eval_metrics[f'{dataset_key}/negative_log_likelihood'] = log_loss(
      y_pred=y_pred, y_true=y_true, labels=np.asarray([0, 1]))
  try:
    eval_metrics[f'{dataset_key}/auroc'] = roc_auc_score(
        y_true=y_true, y_score=y_pred, labels=np.asarray([0, 1]))
  except ValueError:
    eval_metrics[f'{dataset_key}/auroc'] = None
  precision, recall, _ = precision_recall_curve(
      y_true=y_true, probas_pred=y_pred)
  eval_metrics[f'{dataset_key}/auprc'] = auc(recall, precision)
  eval_metrics[f'{dataset_key}/accuracy'] = (
      accuracy_score(y_true=y_true, y_pred=(y_pred > 0.5)))

  # Uncertainty metrics
  ece = rm.metrics.ExpectedCalibrationError(num_bins=ece_num_bins)
  ece.add_batch(y_pred, label=y_true)
  eval_metrics[f'{dataset_key}/ece'] = ece.result()['ece']

  if compute_open_set_recognition:
    is_ood = results['is_ood']
    eval_metrics[f'{dataset_key}/ood_detection_auroc'] = roc_auc_score(
        y_true=is_ood, y_score=y_pred_entropy)
    precision, recall, _ = precision_recall_curve(
        y_true=is_ood, probas_pred=y_pred_entropy)
    eval_metrics[f'{dataset_key}/ood_detection_auprc'] = auc(recall, precision)

    # For the joint datasets, we also compute a rebalanced retention metric,
    # in which we duplicate the OOD dataset to match the size of the in-domain
    # dataset, and then compute the retention metrics.
    rebal_ret_scores = compute_rebalanced_retention_scores(results)
    eval_metrics[f'{dataset_key}/balanced_retention_accuracy_auc'] = (
        rebal_ret_scores['accuracy'])
    eval_metrics[f'{dataset_key}/balanced_retention_nll_auc'] = (
        rebal_ret_scores['nll'])
    eval_metrics[f'{dataset_key}/balanced_retention_auroc_auc'] = (
        rebal_ret_scores['auroc'])
    eval_metrics[f'{dataset_key}/balanced_retention_auprc_auc'] = (
        rebal_ret_scores['auprc'])
  else:
    # This is added for convenience when logging (so the entry exists
    # in tabular format)
    eval_metrics[f'{dataset_key}/ood_detection_auroc'] = None
    eval_metrics[f'{dataset_key}/ood_detection_auprc'] = None
    eval_metrics[f'{dataset_key}/balanced_retention_accuracy_auc'] = None
    eval_metrics[f'{dataset_key}/balanced_retention_nll_auc'] = None
    eval_metrics[f'{dataset_key}/balanced_retention_auroc_auc'] = None
    eval_metrics[f'{dataset_key}/balanced_retention_auprc_auc'] = None

  if compute_retention_auc:
    assert 'accuracy_arr' in results
    assert 'nll_arr' in results
    eval_metrics[f'{dataset_key}/retention_accuracy_auc'] = np.mean(
        compute_retention_curve_on_accuracies(
            accuracies=results['accuracy_arr'], uncertainty=y_pred_entropy))
    eval_metrics[f'{dataset_key}/retention_nll_auc'] = np.mean(
        compute_retention_curve_on_losses(
            losses=results['nll_arr'], uncertainty=y_pred_entropy))
    eval_metrics[f'{dataset_key}/retention_auroc_auc'] = np.mean(
        compute_auc_retention_curve(
            y_pred=y_pred,
            y_true=y_true,
            uncertainty=y_pred_entropy,
            auc_str='roc'))
    eval_metrics[f'{dataset_key}/retention_auprc_auc'] = np.mean(
        compute_auc_retention_curve(
            y_pred=y_pred,
            y_true=y_true,
            uncertainty=y_pred_entropy,
            auc_str='prc'))

  return eval_metrics

def compute_roc_curve(y_uncertainty: np.ndarray, is_ood: np.ndarray):
  """Compute OOD detection ROC curve using sklearn methods.

  Args:
    y_uncertainty: np.ndarray, uncertainty scores for each example.
    is_ood: np.ndarray, Boolean array indicating if an example is from the OOD
      dataset (True) or in-domain (False).

  Returns:
      Tuple[np.ndarray, np.ndarray, np.float]:
        FPR curve, TPR curve, and ROC-AUC value.
  """
  fpr, tpr, _ = roc_curve(y_true=is_ood, y_score=y_uncertainty)
  roc_auc = auc(x=fpr, y=tpr)
  return fpr, tpr, roc_auc

def get_retention_curve_normalizer(use_oracle, n_objects):
  """Obtain normalization constants for each entry of the unnormalized
  retention curve.
  When using an oracle, we divide by the total number of objects.
  When not, we divide by the object index (i.e., the number of objects used
    to compute the model metric at each referral rate).
  Args:
    use_oracle: Bool, if True, evaluate the combined predictive performance
      of the model and an oracle that is correct on all referred datapoints.
    n_objects: int, number of objects used to create the retention curve.
  """
  if use_oracle:
    return n_objects
  else:
    # e.g., for 5 objects, returns [5, 4, 3, 2, 1, 1], where the extra
    # element at the end divides the term corresponding to referring
    # all examples.
    normalizer = np.arange(n_objects + 1)
    normalizer[0] = 1
    return normalizer[::-1]

def compute_retention_curve_on_losses(losses, uncertainty, use_oracle=False):
  """Computes a retention curve on a loss (where lower loss is better)
  and corresponding per-example uncertainty values.
  Based on utils by Andrey Malinin, Yandex Research.
  https://github.com/yandex-research/shifts/blob/main/weather/assessment.py
  Args:
    losses: np.ndarray, losses from a particular dataset.
    uncertainty: np.ndarray, per-example uncertainties for the same dataset,
      should follow the order of the losses.
    use_oracle: Bool (default: False), if True, evaluate the combined predictive
      performance of the model and an oracle that is correct on all referred
      datapoints.
  Returns:
    np.ndarray, retention curve at all possible retention thresholds,
      including all examples retained (i.e., no referral) and no examples
      retained (i.e., all points referred to an expert).
  """
  n_objects = losses.shape[0]
  uncertainty_order = uncertainty.argsort()

  # Losses in order of increasing uncertainty
  losses = losses[uncertainty_order]
  error_rates = np.zeros(n_objects + 1)
  error_rates[:-1] = np.cumsum(losses)[::-1]

  # With oracle:
  # * Divide by total number of predictions
  # Without oracle:
  # * Divide by only the number of predictions the model must make at each
  # * referral rate
  normalizer = get_retention_curve_normalizer(use_oracle, n_objects)
  error_rates = error_rates / normalizer

  return error_rates

def compute_retention_curve_on_accuracies(
    accuracies,
    uncertainty,
    use_oracle=False
):
  """Computes a retention curve on an accuracy (where higher accuracy is better)
  and corresponding per-example uncertainty values.
  Based on utils by Andrey Malinin, Yandex Research.
  https://github.com/yandex-research/shifts/blob/main/weather/assessment.py
  Args:
    accuracies: np.ndarray, accuracies from a particular dataset.
    uncertainty: np.ndarray, per-example uncertainties for the same dataset,
      should follow the order of the accuracies.
    use_oracle: Bool (default: False), if True, evaluate the combined predictive
      performance of the model and an oracle that is correct on all referred
      datapoints.
  Returns:
    np.ndarray, retention curve at all possible retention thresholds,
      including all examples retained (i.e., no referral) and no examples
      retained (i.e., all points referred to an expert).
  """
  n_objects = accuracies.shape[0]
  uncertainty_order = uncertainty.argsort()

  # Per-point accuracy (binary) in order of increasing uncertainty
  accuracies = accuracies[uncertainty_order]
  retention_arr = np.zeros(n_objects + 1)

  for i in range(1, n_objects):
    accuracy_i = accuracies[:i].sum()

    if use_oracle:
      j = n_objects - i
      accuracy_i += j

    retention_arr[i] = accuracy_i

  # With oracle:
  # * Divide by total number of predictions
  # Without oracle:
  # * Divide by only the number of predictions the model must make at each
  # * referral rate
  normalizer = get_retention_curve_normalizer(use_oracle, n_objects)

  # Assume perfect performance when all examples have been referred.
  retention_arr[0] = n_objects if use_oracle else 1
  retention_arr[-1] = accuracies.sum()

  acc_rates = retention_arr[::-1] / normalizer
  return acc_rates

def compute_auc_retention_curve(y_pred,
                                y_true,
                                uncertainty,
                                auc_str,
                                n_buckets=100,
                                use_oracle=False):
  """Computes a retention curve for AUC or AUPRC using predictions, ground
  truths, and corresponding per-example uncertainty values.
  Based on utils by Andrey Malinin, Yandex Research.
  https://github.com/yandex-research/shifts/blob/main/weather/assessment.py
  Args:
    y_pred: np.ndarray, predicted sigmoid probabilities.
    y_true: np.ndarray, ground truth values.
    uncertainty: np.ndarray, per-example uncertainties for the same dataset,
      should follow the order of the accuracies.
    auc_str: str, determines if we evaluate the retention AUC or AUPRC.
    n_buckets: int, number of retention thresholds to evaluate (AUC can be
      costly to evaluate for thousands of possible thresholds.
    use_oracle: Bool (default: False), if True, evaluate the combined predictive
      performance of the model and an oracle that is correct on all referred
      datapoints.
  Returns:
    np.ndarray, AUC or AUPRC retention curve at specified number of thresholds.
  """

  def compute_auroc(true, pred):
    return roc_auc_score(y_true=true, y_score=pred)

  def compute_auprc(true, pred):
    precision, recall, _ = precision_recall_curve(y_true=true, probas_pred=pred)
    return auc(recall, precision)

  if auc_str == 'roc':
    auc_fn = compute_auroc
  elif auc_str == 'prc':
    auc_fn = compute_auprc
  else:
    raise NotImplementedError

  n_objects = y_pred.shape[0]
  bucket_indices = (np.arange(n_buckets) / n_buckets) * n_objects
  bucket_indices = bucket_indices.astype(int)

  uncertainty_order = uncertainty.argsort()
  y_pred = y_pred[uncertainty_order]
  y_true = y_true[uncertainty_order]

  retention_arr = np.zeros(n_buckets + 1)

  if use_oracle:
    # We will later divide by n_objects
    perfect_performance = n_objects
  else:
    perfect_performance = 1

  # Assume perfect performance when all examples have been referred.
  retention_arr[0] = perfect_performance

  check_n_unique = True

  for i_buckets in range(1, n_buckets):
    i_objects = bucket_indices[i_buckets]
    j_objects = n_objects - i_objects
    y_pred_curr = y_pred[:i_objects]
    y_true_curr = y_true[:i_objects]

    # For the first few low uncertainty points, we may be predicting the same
    # (correct) label, which would break AUC. We default to assigning
    # perfect AUC until this condition is broken.
    if (check_n_unique and
        len(np.unique(y_true_curr)) == 1 and
        np.array_equal(y_pred_curr > 0.5, y_true_curr)):
      retention_arr[i_buckets] = perfect_performance
    else:
      try:
        auc_val = auc_fn(true=y_true_curr, pred=y_pred_curr)
        if use_oracle:
          # Weight current AUC/AUPRC by number of objects, and
          # add weight of oracle's perfect prediction.
          retention_arr[i_buckets] = (i_objects * auc_val) + j_objects
        else:
          retention_arr[i_buckets] = auc_val
      except ValueError:  # All of the preds are the same
        if use_oracle:
          retention_arr[i_buckets] = j_objects
        else:
          retention_arr[i_buckets] = 0
      check_n_unique = False

  # Handle the case when no examples have been referred.
  try:
    auc_val = auc_fn(true=y_true, pred=y_pred)
    if use_oracle:
      retention_arr[-1] = n_objects * auc_val
    else:
      retention_arr[-1] = auc_val
  except ValueError:
    retention_arr[-1] = 0

  if use_oracle:
    auc_retention = retention_arr[::-1] / n_objects
  else:
    auc_retention = retention_arr[::-1]

  return auc_retention

def compute_ood_calibration_curve(y_pred: np.ndarray):
  """OOO calibration curve.
  Form a curve by sweeping over confidences in [0, 1], and counting
  the number of predictions with confidence over the threshold.
  On an OOD dataset, we should expect low confidence.
  Args:
    y_pred: np.ndarray, predictions on the OOD dataset
  Returns:
    Tuple[np.ndarray, np.ndarray]
      First np.ndarray:
        sorted probabilities for the predicted class
      Second np.ndarray:
        number of predictions with confidence greater than or equal to
        the confidence at a given threshold
  """
  # TODO(nband): Debug OOD calibration curve metric.
  raise NotImplementedError('Implementation in progress.')


#!/usr/bin/env python
# encoding: utf-8
#
# facke-guider.py
#
# Created by José Sánchez-Gallego on 29 Mar 2017.

from __future__ import division
from __future__ import print_function
from __future__ import absolute_import

import opscore
from opscore.protocols.parser import CommandParser
from opscore.utility.qstr import qstr
from opscore.utility.tback import tback

from opscore.utility.sdss3logging import setupRootLogger, setConsoleLevel
import logging

import opscore.protocols.keys as keys
import opscore.protocols.validation as validation

import actorcore.CommandLinkManager as cmdLinkManager
import actorcore.Command as actorCmd
import actorcore.CmdrConnection

import ConfigParser
import Queue

from twisted.internet import reactor

import imp
import inspect
import os
import re
import sys
import threading
import traceback

class FakeGuider(object):

    def __init__(self, name, productName='guiderActor', makeCmdrConnection=True):
        self.ii = 1026
        self.name = name
        self.productName = productName if productName else self.name
        product_dir_name = '$%s_DIR' % (self.productName.upper())
        self.product_dir = os.path.expandvars(product_dir_name)

        self.parser = CommandParser()

        self.config = ConfigParser.ConfigParser()
        self.config.read(os.path.join(os.path.dirname(__file__), 'guider.cfg'))

        self.configureLogs()

        # The list of all connected sources.
        tronInterface = ''
        tronPort = 9994
        self.commandSources = cmdLinkManager.listen(self,
                                                    port=tronPort,
                                                    interface=tronInterface)
        # The Command which we send uncommanded output to.
        self.bcast = actorCmd.Command(self.commandSources,
                                      'self.0', 0, 0, None, immortal=True)

        # IDs to send commands to ourself.
        self.selfCID = self.commandSources.fetchCid()
        self.synthMID = 1

        # commandSets are the command handler packages. Each handles
        # a vocabulary, which it registers when loaded.
        # We gather them in one place mainly so that "meta-commands" (init, status)
        # can find the others.
        self.commandSets = {}

        self.handler = validation.CommandHandler()

        self.attachAllCmdSets()

        self.commandQueue = Queue.Queue()
        self.shuttingDown = False

        if makeCmdrConnection:
            self.cmdr = actorcore.CmdrConnection.Cmdr(name, self)
            self.cmdr.connectionMade = self._connectionMade
            self.cmdr.connect()
        else:
            self.cmdr = None

    def configureLogs(self, cmd=None):
        """ (re-)configure our logs. """

        self.logDir = self.config.get('logging', 'logdir')
        assert self.logDir, "logdir must be set!"

        # Make the root logger go to a rotating file. All others derive from this.
        setupRootLogger(self.logDir)

        # The real stderr/console filtering is actually done through the console Handler.
        try:
            consoleLevel = int(self.config.get('logging', 'consoleLevel'))
        except:
            consoleLevel = int(self.config.get('logging', 'baseLevel'))
        setConsoleLevel(consoleLevel)

        # self.console needs to be renamed ore deleted, I think.
        self.console = logging.getLogger('')
        self.console.setLevel(int(self.config.get('logging', 'baseLevel')))

        self.logger = logging.getLogger('actor')
        self.logger.setLevel(int(self.config.get('logging', 'baseLevel')))
        self.logger.propagate = True
        self.logger.info('(re-)configured root and actor logs')

        self.cmdLog = logging.getLogger('cmds')
        self.cmdLog.setLevel(int(self.config.get('logging', 'cmdLevel')))
        self.cmdLog.propagate = True
        self.cmdLog.info('(re-)configured cmds log')

        if cmd:
            cmd.inform('text="reconfigured logs"')

    def versionString(self, cmd):
        """ Return the version key value.

        If you simply want to generate the keyword, call .sendVersionKey().
        """

        cmd.warn("text='pathetic version string: unknown'")
        return 'unknown'

    def sendVersionKey(self, cmd):
        """ Generate the version keyword in response to cmd. """

        version = self.versionString(cmd)
        cmd.inform('version=%s' % (qstr(version)))

    def triggerHubConnection(self):
        """ Send the hub a command to connect back to us. """

        if not self.cmdr:
            self.bcast.warn('text="CANNOT ask hub to connect to us, '
                            'since we do not have a connection to it yet!"')
            return

        self.bcast.warn('%s is asking the hub to connect back to us' % (self.name))
        self.cmdr.dispatcher.executeCmd(opscore.actor.keyvar.CmdVar
                                        (actor='hub', cmdStr='startNubs %s' % (self.name),
                                         timeLim=5.0))

    def _connectionMade(self):
        """ twisted arranges to call this when self.cmdr has been established. """

        self.bcast.warn('%s is connected to the hub.' % (self.name))

        #
        # Request that tron connect to us.
        #
        self.triggerHubConnection()
        self.connectionMade()

    def connectionMade(self):
        """ For overriding. """
        pass

    def attachCmdSet(self, cname, path=None):
        """ (Re-)load and attach a named set of commands. """

        if path is None:
            path = [os.path.join(self.product_dir, 'python', self.productName, 'Commands')]

        self.logger.info("attaching command set %s from path %s", cname, path)

        file = None
        try:
            file, filename, description = imp.find_module(cname, path)
            self.logger.debug("command set file=%s filename=%s from path %s",
                              file, filename, path)
            mod = imp.load_module(cname, file, filename, description)
        except ImportError as e:
            raise RuntimeError('Import of %s failed: %s' % (cname, e))
        finally:
            if file:
                file.close()

        # Instantiate and save a new command handler.
        cmdSet = getattr(mod, cname)(self)

        # Check any new commands before finishing with the load. This
        # is a bit messy, as the commands might depend on a valid
        # keyword dictionary, which also comes from the module
        # file.
        #
        # BAD problem here: the Keys define a single namespace. We need
        # to check for conflicts and allow unloading. Right now we unilaterally
        # load the Keys and do not unload them if the validation fails.
        if hasattr(cmdSet, 'keys') and cmdSet.keys:
            keys.CmdKey.addKeys(cmdSet.keys)
        valCmds = []
        for v in cmdSet.vocab:
            try:
                verb, args, func = v
            except ValueError as e:
                raise RuntimeError("vocabulary word needs three parts: %s" % (v))

            # Check that the function exists and get its help.
            #
            funcDoc = inspect.getdoc(func)
            valCmd = validation.Cmd(verb, args, help=funcDoc) >> func
            valCmds.append(valCmd)

        # Got this far? Commit. Save the Cmds so that we can delete them later.
        oldCmdSet = self.commandSets.get(cname, None)
        cmdSet.validatedCmds = valCmds
        self.commandSets[cname] = cmdSet

        # Delete previous set of consumers for this named CmdSet, add new ones.
        if oldCmdSet:
            self.handler.removeConsumers(*oldCmdSet.validatedCmds)
        self.handler.addConsumers(*cmdSet.validatedCmds)

        self.logger.debug("handler verbs: %s" % (self.handler.consumers.keys()))

    def attachAllCmdSets(self, path=None):
        """ (Re-)load all command classes -- files in ./Command which end with Cmd.py.
        """

        if path is None:
            self.attachAllCmdSets(path=os.path.join(os.path.expandvars('$ACTORCORE_DIR'),
                                                    'python', 'actorcore', 'Commands'))
            # self.attachAllCmdSets(path=os.path.join(self.product_dir, 'python', self.productName,
            #                                         'Commands'))
            return

        dirlist = os.listdir(path)
        dirlist.sort()
        self.logger.info("loading %s" % (dirlist))

        for f in dirlist:
            if os.path.isdir(f) and not f.startswith('.'):
                self.attachAllCmdSets(path=f)
            if re.match('^[a-zA-Z][a-zA-Z0-9_-]*Cmd\.py$', f):
                self.attachCmdSet(f[:-3], [path])

    def cmdTraceback(self, e):
        eType, eValue, eTraceback = sys.exc_info()
        tbList = traceback.extract_tb(eTraceback)
        where = tbList[-1]

        return "%r at %s:%d" % (eValue, where[0], where[1])

    def runActorCmd(self, cmd):

        self.output_file()

        try:
            cmdStr = cmd.rawCmd
            self.cmdLog.debug('raw cmd: %s' % (cmdStr))

            try:
                validatedCmd, cmdFuncs = self.handler.match(cmdStr)
            except Exception as e:
                cmd.fail('text=%s' % (qstr("Unmatched command: %s (exception: %s)" %
                                           (cmdStr, e))))
                # tback('actor_loop', e)
                return

            if not validatedCmd:
                cmd.fail('text=%s' % (qstr("Unrecognized command: %s" % (cmdStr))))
                return

            self.cmdLog.info('< %s:%d %s' % (cmd.cmdr, cmd.mid, validatedCmd))
            if len(cmdFuncs) > 1:
                cmd.warn('text=%s' % (qstr("command has more than one callback (%s): %s" %
                                           (cmdFuncs, validatedCmd))))
            try:
                cmd.cmd = validatedCmd
                for func in cmdFuncs:
                    func(cmd)
            except Exception as e:
                oneLiner = self.cmdTraceback(e)
                cmd.fail('text=%s' % (qstr("command failed: %s" % (oneLiner))))
                # tback('newCmd', e)
                return

        except Exception as e:
            cmd.fail('text=%s' %
                     (qstr('completely unexpected exception when processing a new command: %s'
                      % (e))))
            try:
                tback('newCmdFail', e)
            except:
                pass

    def actor_loop(self):
        """ Check the command queue and dispatch commands."""

        while True:
            try:
                cmd = self.commandQueue.get(block=True, timeout=3)
            except Queue.Empty:
                if self.shuttingDown:
                    return
                else:
                    continue
            self.runActorCmd(cmd)

    def commandFailed(self, cmd):
        """ Gets called when a command has failed. """
        pass

    def newCmd(self, cmd):
        """ Dispatch a newly received command. """

        self.cmdLog.info('new cmd: %s' % (cmd))

        # Empty cmds are OK; send an empty response...
        if len(cmd.rawCmd) == 0:
            cmd.finish('')
            return None

        if self.runInReactorThread:
            self.runActorCmd(cmd)
        else:
            self.commandQueue.put(cmd)

        return self

    def callCommand(self, cmdStr):
        """ Send ourselves a command. """

        cmd = actorCmd.Command(self.commandSources, 'self.%d' % (self.selfCID),
                               cid=self.selfCID, mid=self.synthMID, rawCmd=cmdStr)
        self.synthMID += 1
        self.newCmd(cmd)

    def _shutdown(self):
        self.shuttingDown = True

    def run(self, doReactor=True):
        """ Actually run the twisted reactor. """
        try:
            self.runInReactorThread = self.config.getboolean(self.name, 'runInReactorThread')
        except:
            self.runInReactorThread = False

        self.logger.info("starting reactor (in own thread=%s)...." % (not self.runInReactorThread))
        try:
            if not self.runInReactorThread:
                threading.Thread(target=self.actor_loop).start()
            if doReactor:
                reactor.run()
        except Exception as e:
            tback('run', e)

        if doReactor:
            self.logger.info("reactor dead, cleaning up...")
            self._shutdown()

    def output_file(self):
        self.bcast.inform('guideState="on"')
        self.bcast.inform('file=/data/gcam/57831/,proc-gimg-{0}.fits.gz'.format(self.ii))
        self.ii += 1
        reactor.callLater(10, self.output_file)

if __name__ == '__main__':
    guider = FakeGuider('guider')
    guider.run()


import util.settings as settings
import time
import sys, getopt
from util.logging import *
from colorama import init
import supermarkets.runner as srunner

def main(argv):
	init()

	# Parsing of arguments
	try:
		#opts, args = getopt.getopt(argv,"hi:o:",["ifile=","ofile="])
		opts, args = getopt.getopt(argv,"hve",["help","verbose","error"])
	except getopt.GetoptError:
		sys.exit(2)
	for opt, arg in opts:
		if opt == '-h' or opt == '--help':
			LogH("SupermarketScraper Help function")
			LogH("Permitted parameters are:")
			LogH("-h, --help (Show this info)")
			LogH("-v, --verbose (Verbose logging)")
			LogH("-e, --error (Error logging)")
			#LogI("# -t, --tor (Send all trafic via Tor")
			sys.exit()
		elif opt == '-e' or opt == '--error':
			settings.debugging = True
		elif opt == '-v' or opt == '--verbose':
			settings.verbose = True
	# End of argument parsing

	start_time = time.time() * 1000

	settings.print_info()

	LogD("Verbose logging enabled!!\n")

	srunner.run()

	seconds = (time.time() * 1000) - start_time
	PrintLine()
	LogI("Scraper finished in {0}ms.\n".format(format(seconds, '.2f')))

if __name__ == '__main__':
	main(sys.argv[1:])


from vegetation_index import psri_index
from osgeo import gdal
import numpy as np
import zipfile
import os

# adjust with band data filename
target_band = ["B04_10m.jp2", "B08_10m.jp2"]
print(target_band)

zip_folder_path = "./data/2020"
files = os.listdir(zip_folder_path)

# make folder if not exits
output_folder = f"./index/psri/2020"
if not os.path.exists(output_folder):
    os.makedirs(output_folder)

# shp file
shp_file = './shp/poly_paddy_13.5768_101.717.shp'

for path in files:
    zip_file_path = f'{zip_folder_path}/{path}'
    file_inside_zip = []

    with zipfile.ZipFile(zip_file_path, 'r') as zip_file:
        for file_name in zip_file.namelist():
            for target_fl in target_band:
                if file_name.endswith(target_fl):
                    file_inside_zip.append(file_name)

    # config to accesing the zip
    gdal.VSICurlClearCache()
    gdal.PushErrorHandler('CPLQuietErrorHandler')
    gdal.SetConfigOption('CPL_VSIL_CURL_ALLOWED_EXTENSIONS', 'jp2')
    gdal.SetConfigOption('CPL_VSIL_CURL_ALLOWED_EXTENSIONS', 'jp2')

    # Open Using Gdal (adjust it with target_band)
    b8 = gdal.Open('/vsizip/' + zip_file_path + '/' + file_inside_zip[0])
    b4 = gdal.Open('/vsizip/' + zip_file_path + '/' + file_inside_zip[1])

    geoTransform = b4.GetGeoTransform()
    geoProjection = b4.GetProjection()

    nir = b8.GetRasterBand(1).ReadAsArray().astype(np.float32)
    red = b4.GetRasterBand(1).ReadAsArray().astype(np.float32)

    psri = psri_index(red, nir)

    # Save ram
    b4 = None
    b8 = None
    nir = None
    red = None

    # save psri
    driver = gdal.GetDriverByName('GTiff')
    output_path = f'{output_folder}/{path.split(".")[0]}.tif'
    cols, rows = psri.shape

    lon_start = geoTransform[0]
    lat_start = geoTransform[3]
    pixel_width = geoTransform[1]
    pixel_height = geoTransform[5]

    lon_values = np.arange(lon_start, lon_start +
                           cols * pixel_width, pixel_width)
    lat_values = np.arange(lat_start, lat_start + rows *
                           pixel_height, pixel_height)

    output_dataset = driver.Create(
        output_path, cols, rows, 1, gdal.GDT_Float32)
    output_dataset.GetRasterBand(1).WriteArray(psri)

    output_dataset.SetGeoTransform(geoTransform)
    output_dataset.SetProjection(geoProjection)

    metadata = {
        'Latitude': ' '.join(map(str, lat_values)),
        'Longitude': ' '.join(map(str, lon_values))
    }

    output_dataset.SetMetadata(metadata)

    print(f"Extract {path} : Completed")


import discord
from discord.ext import commands
from core.classes import Cog_Extension
import json, asyncio, datetime, sqlite3, shutil, os, time

class Task(Cog_Extension):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.time_task_counters = 0
        self.bg_task = self.bot.loop.create_task(self.time_task())

    async def time_task(self):
            await self.bot.wait_until_ready()
            with open("setting.json", 'r', encoding='utf-8') as jFile:
                jdata = json.load(jFile)
            self.channel = self.bot.get_channel(jdata['test_channel'])
            while not self.bot.is_closed():
                now_time = datetime.datetime.now().strftime("%H%M")
                with open('setting.json', 'r', encoding='utf-8') as JJ:
                    Godjj = json.load(JJ)
                if now_time == Godjj['指定排程'] and self.time_task_counters == 0:
                    await self.channel.send('汪汪汪汪!! 瘋狗動出沒')
                    self.time_task_counters = 1
                    await asyncio.sleep(1)
                else:
                    await asyncio.sleep(1) #間隔用
                    pass

    async def record_task(self, num):
        await self.bot.wait_until_ready()
        while not self.bot.is_closed():
            await asyncio.sleep(60)
            with open("setting.json", 'r', encoding='utf-8') as jFile:
                jdata = json.load(jFile)
            now_time = datetime.datetime.now().strftime('%H%M')  # X小時X分才進行備份
            if now_time == jdata['time']:
                #在刪除舊有檔案前 把檔案備份至log資料夾------------------------------------------------------------------
                if not os.path.exists(".\\log"):  # 檢查有沒有log資料夾
                    os.mkdir(".\\log")
                with open('setting.json', 'r', encoding='utf-8-sig') as jFile:
                    jdata = json.load(jFile)
                cur_path = ("{}.db".format(time.strftime("%Y%m%d%H%M%S\
                            ", time.localtime())))
                shutil.copy("log.db", cur_path)
                shutil.move(cur_path, jdata['log_path'])
                #------------------------------------------------------------------
                try:
                    conn = sqlite3.connect('log.db')
                    conn.execute("DROP TABLE member")
                    conn.commit()
                    conn.close()
                except Exception as e:
                    print("讀取失敗 原因: ", e)
                self.channel = self.bot.get_channel(648202897259233355)
                await self.channel.send("清除資料庫完成")

    @commands.command()
    async def cancel_rd_task(self, ctx):
        self.rd_task.cancel()
        await ctx.send(f'關閉背景任務')

    @commands.command()
    async def create_rd_task(self, ctx, num: int):
        self.rd_task = self.bot.loop.create_task(self.record_task(num))

        await ctx.send(f'開始每隔{num}秒自動清除資料庫')

    @commands.command()
    async def set_channel(self, ctx, ch: int):
        self.channel = self.bot.get_channel(ch)
        await ctx.send(f'Set Channel: {self.channel.mention}') #mention 標記功能

    @commands.command() #使用者可以指定排程
    async def set_time(self, ctx, time):
        self.counters = 0
        with open('setting.json', 'r', encoding='utf-8') as JJ:
            Godjj = json.load(JJ)
        Godjj['指定排程'] = time #使用者的資料傳入json裡
        with open('setting.json', 'w', encoding='utf-8') as JJ:
            json.dump(Godjj, JJ, indent=4)
        await ctx.send('排程設定成功')

def setup(bot):
    bot.add_cog(Task(bot))


#!/usr/bin/env python
import sys
import os
import shutil
#os.environ['OPENBLAS_NUM_THREADS'] = '1'

import argparse
import subprocess
import requests
import stat
import json
import pathlib
import zipfile
import pysam
import pandas as pd
import numpy as np
from collections import Counter
from functools import reduce

from sys import platform

from nimble.types import Config
from nimble.parse import parse_fasta, parse_filter_config, parse_csv
from nimble.usage import print_usage_and_exit
from nimble.utils import get_exec_name_from_platform, low_complexity_filter_amount, append_path_string

ALIGN_TRIES = 10
ALIGN_TRIES_THRESHOLD = 0

# Generate and write human-editable config json files to disk. Input data is a CSV, FASTA, or both
def generate(file, opt_file, output_path):

    (data, config, is_csv_req) = process_file(file, opt_file)
    (data_opt, config_opt, is_csv_opt) = process_file(opt_file, file)

    final_config = config
    if data_opt != None and is_csv_opt:
        final_config = config_opt

    final_data = None
    if data_opt != None:
        if is_csv_req:
            final_data = collate_data(data_opt, data)
        elif is_csv_opt:
            final_data = collate_data(data, data_opt)
    else:
        final_data = data

    print("Filtered " + str(low_complexity_filter_amount) + " base pairs from reference library.")

    # Write reference and default config to disk
    with open(output_path, "w") as f:
        json.dump([ final_config.__dict__, final_data.__dict__], f, indent=2)

# Parse a lone FASTA/lone CSV/CSV-FASTA tuple. If there's a lone FASTA, generate a simple library.
# If there's a lone CSV, assume it has sequence information or a genbank link.
# If there is not a lone CSV, assume it contains the metadata and that the sequence information is contained in the FASTA.
def process_file(file, paired_file):
    data = None
    config = None
    is_csv = False

    if file:
        if pathlib.Path(file).suffix == ".fasta":
            (data, config) = parse_fasta(file)
        elif pathlib.Path(file).suffix == ".csv" and paired_file:
            (data, config) = parse_csv(file, False)
            is_csv = True
        elif pathlib.Path(file).suffix == ".csv" and not paired_file:
            (data, config) = parse_csv(file, True)
            is_csv = True

    return (data, config, is_csv)

def collate_data(data, metadata):
    name_idx = data.headers.index("sequence_name")
    sequence_idx = data.headers.index("sequence")
    nt_length_idx = data.headers.index("nt_length")

    meta_name_idx = metadata.headers.index("sequence_name")
    meta_sequence_idx = metadata.headers.index("sequence")
    meta_nt_length_idx = metadata.headers.index("nt_length")

    metadata.columns[meta_sequence_idx] = ["" for _ in range(0, len(data.columns[sequence_idx]))]
    metadata.columns[meta_nt_length_idx] = ["" for _ in range(0, len(data.columns[nt_length_idx]))]

    for (from_idx, name) in enumerate(data.columns[name_idx]):
        if name not in metadata.columns[meta_name_idx]:
            print("Error -- record " + name + " is not found in both input files.")
            sys.exit()

        update_idx = metadata.columns[meta_name_idx].index(name)

        metadata.columns[meta_sequence_idx][update_idx] = data.columns[sequence_idx][from_idx]
        metadata.columns[meta_nt_length_idx][update_idx] = data.columns[nt_length_idx][from_idx]

    return metadata

# Given the name of a release, download the platform-specific executable from that release.
# Given no name, default to the most recent release.
def download(release):
    exec_name = get_exec_name_from_platform()
    print("Downloading " + exec_name)

    url = ""
    if len(release) == 1:
        url = (
            "https://github.com/BimberLab/nimble-aligner/releases/download/"
            + release[0]
            + "/"
            + exec_name
        )
    else:
        url = (
            "https://github.com/BimberLab/nimble-aligner/releases/latest/download/"
            + exec_name
        )

    # Download aligner
    r = requests.get(url)

    # The executable will be placed in the python package directory
    aligner_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "aligner")
    print("Aligner download path: " + aligner_path)

    # If we successfully downloaded it, write the file to disk and give it execution permissions if necessary
    if r.status_code == 200:
        with open(aligner_path, "wb") as f:
            f.write(r.content)
        if sys.platform == "linux" or sys.platform == "darwin": # If we're on a Unix, ensure permission to write
            st = os.stat(aligner_path)
            os.chmod(aligner_path, st.st_mode | stat.S_IEXEC | stat.S_IXOTH)
    else:
        print("Error -- could not download aligner, status code " + str(r.status_code))
        sys.exit()

# Check if the aligner exists -- if it does, call it with the given parameters.
def align(reference, output, input, _alignment_path, log_path, num_cores, strand_filter):
    path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "aligner")
    input_ext = os.path.splitext(input[0])[-1].lower()

    if os.path.exists(path):
        if input_ext == ".bam":
            split = input[0].rsplit("/", 1)
            sort_input_bam(split, num_cores)
            input = [split[0] + "/sorted-" + split[1]]

        print("Aligning input data to the reference libraries")
        sys.stdout.flush()

        library_list = reference.split(",")

        processed_param_list = []
        for input_file in input:
            processed_param_list.extend(["--input", input_file])
        processed_param_list.extend(["-c", str(num_cores), "--strand_filter", strand_filter])

        for library in library_list:
            out_file_append = ""

            if len(library_list) > 1:
                out_file_append = "." + os.path.splitext(os.path.basename(library))[0]

            processed_param_list.extend(["-r", library, "-o", append_path_string(output, out_file_append)])
        print(processed_param_list)

        proc = subprocess.Popen([path] + processed_param_list)
        proc.wait()

        return_code = proc.returncode

        if input_ext == ".bam" and return_code == 0:
            print("Deleting intermediate sorted .bam file")
            os.remove(input[0])

        return return_code
    else:
        print("No aligner found. Attempting to download the latest release.\n")
        download([])

        global ALIGN_TRIES
        ALIGN_TRIES = ALIGN_TRIES + 1
        if ALIGN_TRIES >= ALIGN_THRESHOLD:
            print("Error -- could not find or download aligner.")
            sys.exit()

        return align(reference, output, input, alignment_path, log_path, num_cores, strand_filter)

def intersect_lists(lists):
    # Returns the intersection of all lists in a list
    return list(reduce(set.intersection, map(set, lists)))

def write_empty_df(output):
    print('No data to parse from input file, writing empty output.')
    empty_df = pd.DataFrame(columns=['feature', 'count', 'cell_barcode'])
    empty_df.to_csv(output, sep='\t', index=False, compression='gzip', header=False)

def report(input, output):
    df = None

    # if the file has data, try to read it. write an empty output and return if there is no data.
    if os.path.getsize(input) > 0:
        try:
            df = pd.read_csv(input, sep='\t', compression='gzip')
            df.rename(columns={'r1_cb': 'cb', 'r1_umi': 'umi'}, inplace=True) # Use the r1 version of the cb and umi flags
        except pd.errors.EmptyDataError:
            write_empty_df(output)
            return
    else:
        write_empty_df(output)
        return

    # If the file is not empty but the DataFrame is, write an empty output
    if df.empty:
        write_empty_df(output)
        return

    # Keep only necessary columns
    df = df[['features', 'umi', 'cb']]

    # Drop rows where 'features', 'umi', or 'cb' are null or empty
    df = df.dropna(subset=['features', 'umi', 'cb'])
    df = df[(df['features'] != '') & (df['umi'] != '') & (df['cb'] != '')]

    # Split the feature strings into lists
    df['features'] = df['features'].str.split(',')

    # Group by cell barcode (CB) and UMI, aggregate features into a list
    df_grouped = df.groupby(['cb', 'umi'])['features'].apply(list)

    # Calculate the intersection of features within each UMI
    df_grouped = df_grouped.apply(intersect_lists)

    # Convert back to a DataFrame
    df_grouped = df_grouped.reset_index()

    # Identify rows where the intersection resulted in an empty list
    empty_intersection_rows = df_grouped['features'].apply(lambda x: len(x) == 0)

    # Count these rows and print the number
    empty_intersection_count = empty_intersection_rows.sum()
    print(f"Dropped {empty_intersection_count} counts due to empty intersections")

    # Drop these rows from the DataFrame
    df_grouped = df_grouped[~empty_intersection_rows]

    # Join the intersected features back into a string
    df_grouped['features'] = df_grouped['features'].apply(lambda x: ','.join(x))

    # Rename columns
    df_grouped.columns = ['cell_barcode', 'umi', 'feature']

    # Count unique UMIs per cell_barcode-feature pair
    df_counts = df_grouped.groupby(['cell_barcode', 'feature']).size().reset_index()

    # Rename count column
    df_counts.columns = ['cell_barcode', 'feature', 'count']

    # Reorder the columns
    df_counts = df_counts.reindex(['feature', 'count', 'cell_barcode'], axis=1)

    # Write to output file
    df_counts.to_csv(output, sep='\t', index=False, header=False)

def sort_input_bam(file_tuple, cores):
    print("Sorting input .bam")
    tmp_dir = os.environ.get("TMPDIR")
    create_tmp_dir = False

    bam = ""
    sorted_bam = ""

    if len(file_tuple) > 1:
        bam = file_tuple[0] + "/" + file_tuple[1]
        sorted_bam = file_tuple[0] + "/sorted-" + file_tuple[1]
    else:
        bam = file_tuple[0]
        sorted_bam = "./sorted-" + file_tuple[0]

    print("Sorting " + bam + " Outputting to " + sorted_bam)
    sys.stdout.flush()

    if os.path.isfile(sorted_bam):
        print("Sorted bam file already exists, skipping the sorting step.")
        sys.stdout.flush()
        return

    if tmp_dir and not os.path.exists(tmp_dir):
        try:
            os.makedirs(tmp_dir)
            create_tmp_dir = True
            print(f"Created temporary directory {tmp_dir}")
        except OSError as e:
            print(f"Could not create temporary directory {tmp_dir}: {e}")
            sys.stdout.flush()

    if tmp_dir:
        pysam.sort('-t', 'UR', '-n', '-o', sorted_bam, '-@', str(cores), '-T', tmp_dir, bam)
    else:
        pysam.sort('-t', 'UR', '-n', '-o', sorted_bam, '-@', str(cores), bam)

    sort_log = pysam.sort.get_messages()

    if (sort_log):
        print("samtools messages: " + sort_log)

    if create_tmp_dir:
        try:
            shutil.rmtree(tmp_dir)
            print(f"Deleted temporary directory {tmp_dir}")
        except Exception as e:
            print(f"Could not delete temporary directory {tmp_dir}: {e}")
        sys.stdout.flush()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='nimble align')
    subparsers = parser.add_subparsers(title='subcommands', dest='subcommand')

    download_parser = subparsers.add_parser('download')
    download_parser.add_argument('--release', help='The release to download.', type=str, default=[])

    generate_parser = subparsers.add_parser('generate')
    generate_parser.add_argument('--file', help='The file to process.', type=str, required=True)
    generate_parser.add_argument('--opt-file', help='The optional file to process.', type=str, default=None)
    generate_parser.add_argument('--output_path', help='The path to the output file.', type=str, required=True)

    align_parser = subparsers.add_parser('align')
    align_parser.add_argument('--reference', help='The reference genome to align to.', type=str, required=True)
    align_parser.add_argument('--output', help='The path to the output file.', type=str, required=True)
    align_parser.add_argument('--input', help='The input reads.', type=str, required=True, nargs='+')
    align_parser.add_argument('--alignment_path', help='The path to the alignment file.', type=str, default=None)
    align_parser.add_argument('--log_path', help='The path to the log file.', type=str, default=None)
    align_parser.add_argument('-c', '--num_cores', help='The number of cores to use for alignment.', type=int, default=1)
    align_parser.add_argument('--strand_filter', help='Filter reads based on strand information.', type=str, default="unstranded")

    report_parser = subparsers.add_parser('report')
    report_parser.add_argument('-i', '--input', help='The input file.', type=str, required=True)
    report_parser.add_argument('-o', '--output', help='The path to the output file.', type=str, required=True)

    args = parser.parse_args()

    if args.subcommand == 'download':
        download(args.release)
    elif args.subcommand == 'generate':
        generate(args.file, args.opt_file, args.output_path)
    elif args.subcommand == 'align':
        sys.exit(align(args.reference, args.output, args.input, args.alignment_path, args.log_path, args.num_cores, args.strand_filter))
    elif args.subcommand == 'report':
        report(args.input, args.output)


import logging

from ..models.formularios.form_accion_sort import Form_Accion_Sorteable
from ..models.formularios.form_campo import Form_Campo
from ..models.formularios.form_campo_sort import Form_Campo_Sorteable
from ..models.formularios.form_elemento import Form_Elemento
from ..models.formularios.form_filtro_param_sort import Form_Filtro_Params_Sorteable
from ..models.formularios.form_filtro_param import Form_Filtro_Param
from ..models.formularios.form_filtro import Form_Filtro
from ..models.formularios.form_formularios import Form_Formulario
from ..models.formularios.form_imagen import Form_Imagen
from ..models.formularios.form_multi import Form_Multi_Option
from ..models.formularios.form_multi_sort import Form_Multi_Options_Sorteable
from ..models.formularios.form_texto import Form_Texto

logger = logging.getLogger(name=__name__)

##################################################################################
# OBTENER FORM
##################################################################################
def obtener_form(solicitud, post=None, files=None, initial=False):

    try:
        instancia = Form_Formulario.objects.get(solicitud = solicitud)

        form = form_to_dict(instancia)

        return form

    except Exception as e:
        logger.error(e)
        return "No se encontro el formulario especificado"

##################################################################################
# GENERA EL FORMATO DE LOS DATOS EN LOS CAMPOS QUE SE PUEDEN SORTEAR
##################################################################################
def get_datos_campos_sorteable(modelo, valor):
    options_sort = modelo.objects.filter(form_id = valor.id)

    dict_base= {}

    if not options_sort:
        return []
    else:
        for multi in options_sort:
            dict_base.update({ multi.clave : multi.valor })
        return dict_base

##################################################################################
# OBTIENE LA CLAVE SEGUN EL TIPO
##################################################################################
def get_clave_tipo(tipo, text, image):
    if tipo == 'Texto':
        return f'texto_{text}'
    elif tipo == 'Imagen':
        return f'imagen_{image}'
    else:
        return False

##################################################################################
# FORMATEA EL FORMULARIO PARA QUE SEA UN DICT PASABLE A JSON
##################################################################################
def form_to_dict(instancia : Form_Formulario):

    form = {}

    campos_sort = Form_Campo_Sorteable.objects.filter(form_id = instancia.id)

    count_text = 0

    count_image = 0

    for campo_sort in campos_sort:

        tipo = campo_sort.get_type()

        #####################################################
        # CHEQUEAMOS SI ALGUN CAMPO TIENE MAS DE UNA OPCION
        #####################################################
        if tipo == 'Error':
            return "Uno de los campos sorteables tiene seleccionado mas de un tipo, solo aceptan una opcion"

        #####################################################
        # CHEQUEAMOS SI ALGUN CAMPO ESTA TODO VACIO
        #####################################################
        elif tipo == 'None':
            return "Uno de los campos sorteables esta completamente vacio, eligale una opcion"

        #####################################################
        # CARGAMOS EL TEXTO
        #####################################################
        elif tipo == 'Texto':

            count_text += 1

            obj_texto = Form_Texto.objects.get(id = campo_sort.text_id)

            dict_text = {}

            for elemento in obj_texto._meta.get_fields():
                if elemento.name != 'form_campo_sorteable' and elemento.name != 'id':
                    valor = getattr(obj_texto, elemento.name)
                    dict_text.update( { elemento.name : valor } )

            dict_text["elemento_html"] = {"tipo": "text"}
            dict_text["filtro"] = 8

            form.update( { f"texto_{count_text}" : dict_text } )

        #####################################################
        # CARGAMOS EL IMAGEN
        #####################################################
        elif tipo == 'Imagen':

            count_image += 1

            obj_imagen = Form_Imagen.objects.get(id = campo_sort.image_id)

            dict_imagen = {}

            for elemento in obj_imagen._meta.get_fields():
                if elemento.name != 'form_campo_sorteable' and elemento.name != 'id':
                    valor = getattr(obj_imagen, elemento.name)
                    dict_imagen.update( { elemento.name : valor } )

            dict_imagen["elemento_html"] = {"tipo": "text"}
            dict_imagen["filtro"] = 8

            form.update( { f"imagen_{count_image}" : dict_imagen } )

        #####################################################
        # CARGAMOS EL ELEMENTO
        #####################################################
        elif tipo == 'Campo':

            obj_campo = Form_Campo.objects.get(id = campo_sort.field_id)

            dict_campo = {}

            for elemento in obj_campo._meta.get_fields():
                if elemento.name != 'form_campo_sorteable' and elemento.name != 'id' and elemento.name != 'key':

                    valor = getattr(obj_campo, elemento.name)

                    if isinstance(valor, Form_Elemento):
                        dict_campo.update( { elemento.name : valor.get_dict()})

                    elif isinstance(valor, Form_Filtro):
                        dict_campo.update( { elemento.name : valor.get_valor()})

                    elif isinstance(valor, Form_Filtro_Param):

                        res = get_datos_campos_sorteable(Form_Filtro_Params_Sorteable, valor)

                        dict_campo.update( { elemento.name : res } )

                    elif isinstance(valor, Form_Multi_Option):

                        res = get_datos_campos_sorteable(Form_Multi_Options_Sorteable, valor)

                        dict_campo.update( { elemento.name : res } )

                    elif valor is not None:
                        dict_campo.update( { elemento.name : valor } )

            form.update( { obj_campo.key : dict_campo } )

        #####################################################
        # CARGAMOS ARCHIVO
        #####################################################
        if campo_sort.file == True:
            clave = get_clave_tipo(tipo, count_text, count_image)
            if clave :
                form[clave].update( { "archivo" : True } )
            else:
                form[obj_campo.key].update( {"archivo" : True } )

        #####################################################
        # CARGAMOS FORM ERROR
        #####################################################
        if campo_sort.form_error == True:
            clave = get_clave_tipo(tipo, count_text, count_image)
            if clave :
                form[clave].update( { "form_error" : True } )
            else:
                form[obj_campo.key].update( {"form_error" : True } )

        #####################################################
        # CARGAMOS FORM MENSAJE
        #####################################################
        if campo_sort.form_message == True:
            clave = get_clave_tipo(tipo, count_text, count_image)
            if clave :
                form[clave].update( { "form_message" : True } )
            else:
                form[obj_campo.key].update( {"form_message" : True } )

        #####################################################
        # CARGAMOS MISMA LINEA QUE ELEMENTO ANTERIOR
        #####################################################
        if campo_sort.form_same_line:
            clave = get_clave_tipo(tipo, count_text, count_image)
            if clave :
                form[clave].update( { "form_same_line" : campo_sort.form_same_line } )
            else:
                form[obj_campo.key].update( {"form_same_line" : campo_sort.form_same_line } )

        #####################################################
        # CARGAMOS ACCIONES
        #####################################################
        if campo_sort.get_action():
            acciones_sort = Form_Accion_Sorteable.objects.filter(form_id = campo_sort.action_id)

            dict_act = {}

            for act in acciones_sort:
                dict_act.update( { act.elemento : act.accion } )

            clave = get_clave_tipo(tipo, count_text, count_image)
            if clave :
                form[clave].update( { "acciones" : dict_act } )
            else:
                form[obj_campo.key].update( {"acciones" : dict_act } )

    return form


from djoser.views import UserViewSet
from rest_framework.pagination import LimitOffsetPagination
from rest_framework.permissions import IsAuthenticated

class DjoserUserViewSet(UserViewSet):
    pagination_class = LimitOffsetPagination
    permission_classes = [IsAuthenticated]

    def me(self, request, *args, **kwargs):
        response = super().me(request, *args, **kwargs)
        response.data['is_subscribed'] = False
        return response


from collections import Counter

class Solution:
    def permuteUnique(self, nums: List[int]) -> List[List[int]]:
        if len(nums) == 1:
            return [nums]

        counter = Counter(nums)
        result = []
        def backtracking(perm):
            if len(perm) == len(nums):
                result.append(perm[:])
                return
            for num in counter:
                if counter[num] <= 0:
                    continue
                counter[num] -= 1
                perm.append(num)
                backtracking(perm)
                counter[num] += 1
                perm.pop()

        backtracking([])
        return result