# 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
                                         split_name):
  """Adds regression tasks predictions to RegressionMetricsData objects.

  Args:
    metrics_coordinator: MetricsCoordinator, used to accumulate metrics for each
      target.
    target_names: list of str, the names of the targets in the task.
    predictions: array of predictions in time-major shape wnct [num_unroll,
      batch_size, channels, num_targets].
    targets: array of float 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)
  predictions_per_target = np.split(
      predictions, indices_or_sections=num_targets, axis=3)
  target_list = np.split(
      targets, indices_or_sections=num_targets, axis=3)
  for mask_name, mask in eval_mask_dict.items():
    masks_per_target = np.split(
        mask, indices_or_sections=num_targets, axis=3)
    for idx, target_name in enumerate(target_names):
      predictions = predictions_per_target[idx][masks_per_target[idx].astype(
          bool)]
      targets = target_list[idx][masks_per_target[idx].astype(bool)]
      weights = np.ones_like(predictions)
      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

                                         split_name):
  """Adds regression tasks predictions to RegressionMetricsData objects.

  Args:
    metrics_coordinator: MetricsCoordinator, used to accumulate metrics for each
      target.
    target_names: list of str, the names of the targets in the task.
    predictions: array of predictions in time-major shape wnct [num_unroll,
      batch_size, channels, num_targets].
    targets: array of float 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)
  predictions_per_target = np.split(
      predictions, indices_or_sections=num_targets, axis=3)
  target_list = np.split(
      targets, indices_or_sections=num_targets, axis=3)
  for mask_name, mask in eval_mask_dict.items():
    masks_per_target = np.split(
        mask, indices_or_sections=num_targets, axis=3)
    for idx, target_name in enumerate(target_names):
      predictions = predictions_per_target[idx][masks_per_target[idx].astype(
          bool)]
      targets = target_list[idx][masks_per_target[idx].astype(bool)]
      weights = np.ones_like(predictions)
      metrics_coordinator.add_data(

          metrics_data.RegressionMetricsData, (predictions, targets, weights),