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

# Lint as: python3
"""Mortality task implementation to be used in experiments."""

from typing import Dict, List, Mapping, Optional, Tuple, Union

from absl import logging
from ehr_prediction_modeling import types
from ehr_prediction_modeling.tasks import base_task
from ehr_prediction_modeling.tasks import task_data
from ehr_prediction_modeling.tasks import task_masks
from ehr_prediction_modeling.utils import batches
from ehr_prediction_modeling.utils import label_utils
from ehr_prediction_modeling.utils import loss_utils
from ehr_prediction_modeling.utils import mask_utils
import tensorflow.compat.v1 as tf

from ehr_prediction_modeling import configdict


class MortalityRisk(base_task.Task):
  """Task implementation for mortality risk."""

  task_type = types.TaskTypes.MORTALITY

  @property
  def default_masks(self) -> List[str]:
    return [mask_utils.IGNORE_MASK, mask_utils.PADDED_EVENT_MASK]

  @property
  def _supported_train_masks(self) -> Dict[str, List[str]]:
    return {
        **super()._supported_train_masks,
        **{
            task_masks.Train.SINCE_EVENT:
                self.default_masks + [mask_utils.SINCE_EVENT_TRAIN_MASK],
        }
    }

  @property
  def _unformatted_supported_eval_masks(self) -> Dict[str, List[str]]:
    """See base class."""
    return {
        task_masks.Eval.BASE:
            self.default_masks,
        task_masks.Eval.SINCE_EVENT:
            self.default_masks + [mask_utils.SINCE_EVENT_EVAL_MASK],
    }

  def __init__(self, config: configdict.ConfigDict):
    label_keys = self._target_names = [
        label_utils.MORTALITY_LOOKAHEAD_LABEL_BASE.format(
            hours=str(hours)) for hours in self.window_hours
    ]
    if config.mortality_during_admission:
      label_keys.append(label_utils.MORTALITY_IN_ADMISSION_LABEL)
      self._target_names.append(label_utils.MORTALITY_IN_ADMISSION_LABEL)
    logging.info("Targets are: %s", str(self._target_names))
    super().__init__(config, label_keys=label_keys)

  @property
  def num_targets(self) -> int:
    return len(self._label_keys)

  @property
  def target_names(self) -> List[str]:
    return self._target_names

  @property
  def window_hours(self) -> List[int]:
    return [d * 24 for d in self._config.window_days]

  @property
  def prediction_task_type(self) -> str:
    return types.TaskType.BINARY_CLASSIFICATION


  def get_label_dicts(
      self
  ) -> Tuple[Dict[Optional[str], Union[
      tf.FixedLenSequenceFeature, tf.FixedLenFeature]], Dict[Union[
          str, None], Union[tf.FixedLenSequenceFeature, tf.FixedLenFeature]]]:
    """Gets the feature dictionaries to parse a tf.SequenceExample.

    These dicts are used in the parse op (either batched or non-batched):
    https://www.tensorflow.org/api_docs/python/tf/io/parse_single_sequence_example#args
    They should include any feature that is used by this task.

    Returns:
      context_d: Dictionaries of TF features to read in context.
      sequence_d: Dictionaries of TF features to read in sequence.
    """
    context_d, sequence_d = mask_utils.get_labels_for_masks(
        self._config.train_mask, self._config.eval_masks,
        self._all_supported_masks, self._config.time_since_event_label_key)

    sequence_d[label_utils.IGNORE_LABEL] = tf.FixedLenSequenceFeature([1],
                                                                      tf.int64)
    for label_key in self._label_keys:
      sequence_d[label_key] = tf.FixedLenSequenceFeature([1], tf.int64)

    return context_d, sequence_d

  def _get_all_task_variables(
      self, batch: batches.TFBatch,
      model_output: tf.Tensor) -> task_data.TaskVariables:
    """Computes variables for MortalityRisk task.

    Args:
      batch: tf.NextQueuedSequenceBatch, containing a batch of data.
      model_output: Tensor, the output from the model, shape wnt [num_unroll,
        batch_size, dim_model_output].

    Returns:
      task_data.TaskVariables with all the variables from this task.
    """
    logits = self.layer.get_logits(model_output)

    # This requires the targets to be sorted by window.
    if self._config.accumulate_logits:
      logits = tf.cumsum(logits, axis=-1)

    targets = self.get_targets(batch)
    train_loss_mask = self.get_train_mask(batch)

    loss = loss_utils.loss_fn(
        logits,
        targets,
        train_loss_mask,
        scale_pos_weight=self._config.get("scale_pos_weight", 1.0),
        loss_type=types.TaskLossType.CE)
    eval_mask_dict = self.get_eval_mask_dict(batch)
    predictions = tf.sigmoid(logits)

    return task_data.TaskVariables(
        loss=loss,
        targets=targets,