# Path: ehr_prediction_modeling/eval/metrics_data.py
# MAX_METRICS_KEY_LENGTH = 64
# class MetricsData():
# class BinaryMetricsData(MetricsData):
# class RegressionMetricsData(MetricsData):
#   def get_metrics(self) -> Mapping[str, Any]:
#   def add_data(self):
#   def log_metrics(self, target_name: str, current_step: int) -> None:
#   def __init__(
#       self,
#       positive_predictions: Optional[List[float]] = None,
#       negative_predictions: Optional[List[float]] = None,
#       positive_weights: Optional[List[float]] = None,
#       negative_weights: Optional[List[float]] = None):
#   def __eq__(self, other):
#   def __str__(self):
#   def add_data(self, positive_predictions, negative_predictions,
#                positive_weights, negative_weights):
#   def get_metrics(self) -> Mapping[str, Any]:
#   def __init__(
#       self,
#       predictions: Optional[List[float]] = None,
#       targets: Optional[List[float]] = None,
#       weights: Optional[List[float]] = None) -> None:
#   def __eq__(self, other):
#   def __str__(self):
#   def add_data(self, predictions, targets, weights):
#   def get_metrics(self) -> Mapping[str, Any]:


import collections
import os
import dataclasses
import numpy as np
import tensorflow.compat.v1 as tf
from typing import List, Mapping, Optional, Tuple, Type
from absl import logging
from ehr_prediction_modeling.eval import metrics_data
    predictions: array of predictions in time-major shape wnct [num_unroll,
      batch_size, channels, num_targets].
    binary_targets: array of binary targets in time-major shape wnct [
      num_unroll, batch_size, channels, num_targets].
    eval_mask_dict: dict of string mask name to array, the loss masks to be used
      in evaluation, in time-major shape wnct [num_unroll, batch_size, channels,
      num_targets].
    split_name: str, name of the data split this batch is from.
  """
  num_targets = len(target_names)
  # Split predictions by target into a list of length num_targets.
  predictions_per_target = np.split(
      predictions, indices_or_sections=num_targets, axis=3)

  for mask_name, mask in eval_mask_dict.items():
    positive_filter_and_mask = (mask * binary_targets).astype(bool)
    negative_filter_and_mask = (
        mask * (np.ones_like(binary_targets) - binary_targets)).astype(bool)

    positive_masks_per_target = np.split(
        positive_filter_and_mask, indices_or_sections=num_targets, axis=3)
    negative_masks_per_target = np.split(
        negative_filter_and_mask, indices_or_sections=num_targets, axis=3)

    for idx, target_name in enumerate(target_names):
      positives = predictions_per_target[idx][positive_masks_per_target[idx]]
      positive_weights = np.ones_like(positives)
      negatives = predictions_per_target[idx][negative_masks_per_target[idx]]
      negative_weights = np.ones_like(negatives)
      metrics_coordinator.add_data(

# Path: ehr_prediction_modeling/eval/metrics_data.py
# MAX_METRICS_KEY_LENGTH = 64
# class MetricsData():
# class BinaryMetricsData(MetricsData):
# class RegressionMetricsData(MetricsData):
#   def get_metrics(self) -> Mapping[str, Any]:
#   def add_data(self):
#   def log_metrics(self, target_name: str, current_step: int) -> None:
#   def __init__(
#       self,
#       positive_predictions: Optional[List[float]] = None,
#       negative_predictions: Optional[List[float]] = None,
#       positive_weights: Optional[List[float]] = None,
#       negative_weights: Optional[List[float]] = None):
#   def __eq__(self, other):
#   def __str__(self):
#   def add_data(self, positive_predictions, negative_predictions,
#                positive_weights, negative_weights):
#   def get_metrics(self) -> Mapping[str, Any]:
#   def __init__(
#       self,
#       predictions: Optional[List[float]] = None,
#       targets: Optional[List[float]] = None,
#       weights: Optional[List[float]] = None) -> None:
#   def __eq__(self, other):
#   def __str__(self):
#   def add_data(self, predictions, targets, weights):
#   def get_metrics(self) -> Mapping[str, Any]:


# Path: ehr_prediction_modeling/eval/metrics_coordinator.py
import collections
import os
import dataclasses
import numpy as np
import tensorflow.compat.v1 as tf
from typing import List, Mapping, Optional, Tuple, Type
from absl import logging
from ehr_prediction_modeling.eval import metrics_data

    predictions: array of predictions in time-major shape wnct [num_unroll,
      batch_size, channels, num_targets].
    binary_targets: array of binary targets in time-major shape wnct [
      num_unroll, batch_size, channels, num_targets].
    eval_mask_dict: dict of string mask name to array, the loss masks to be used
      in evaluation, in time-major shape wnct [num_unroll, batch_size, channels,
      num_targets].
    split_name: str, name of the data split this batch is from.
  """
  num_targets = len(target_names)
  # Split predictions by target into a list of length num_targets.
  predictions_per_target = np.split(
      predictions, indices_or_sections=num_targets, axis=3)

  for mask_name, mask in eval_mask_dict.items():
    positive_filter_and_mask = (mask * binary_targets).astype(bool)
    negative_filter_and_mask = (
        mask * (np.ones_like(binary_targets) - binary_targets)).astype(bool)

    positive_masks_per_target = np.split(
        positive_filter_and_mask, indices_or_sections=num_targets, axis=3)
    negative_masks_per_target = np.split(
        negative_filter_and_mask, indices_or_sections=num_targets, axis=3)

    for idx, target_name in enumerate(target_names):
      positives = predictions_per_target[idx][positive_masks_per_target[idx]]
      positive_weights = np.ones_like(positives)
      negatives = predictions_per_target[idx][negative_masks_per_target[idx]]
      negative_weights = np.ones_like(negatives)
      metrics_coordinator.add_data(

          metrics_data.BinaryMetricsData,