# coding=utf-8
# Copyright 2022 The Google Research Authors.
#
# 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.

"""Model constructor for Tensorflow state-level models."""
from typing import Dict, List
import numpy as np
import tensorflow as tf
from covid_epidemiology.src import constants
from covid_epidemiology.src.models import generic_seir_model_constructor
from covid_epidemiology.src.models import losses
from covid_epidemiology.src.models.shared import model_utils


class StateModelConstructor(generic_seir_model_constructor.ModelConstructor):
  """Constructs a state Tensorflow model, to be used in tf_seir."""

  def __init__(self, model_spec, random_seed=0):
    super(StateModelConstructor, self).__init__(model_spec, random_seed)
    self.num_states = 17

  def extract_prediction(self, all_states):
    """Extract the death and confirmed predictions."""

    confirmed_all = list()
    death_all = list()

    for curr_state in all_states:
      # pylint: disable=unused-variable
      (exposed_t, infected_d_t, infected_ud_t, recovered_d_t, recovered_ud_t,
       hospitalized_t, hospitalized_cumulative_t, hospitalized_increase_t,
       icu_t, ventilator_t, death_t, population_t, reinfectable_d_t,
       reinfectable_ud_t, reinfectable_vaccine_t, vaccine_immuned_t,
       infected_ud_increase_t) = tf.unstack(curr_state)

      # Include ICU and Ventilator since they are separate compartments.
      confirmed_t = (
          infected_d_t + recovered_d_t + hospitalized_t + icu_t + ventilator_t +
          death_t + reinfectable_d_t)

      confirmed_all.append(confirmed_t)
      death_all.append(death_t)

    return {"confirmed": confirmed_all, "death": death_all}

  def compute_coef(self,
                   ground_truth_timeseries,
                   ground_truth_state,
                   num_train_steps,
                   num_known_steps,
                   power=2.0):
    """Compute train/valid coefficients for loss computation.

    Args:
      ground_truth_timeseries: ground truth compartments
      ground_truth_state: ground truth state level compartments
      num_train_steps: number of timesteps for training
      num_known_steps: number of known timesteps
      power: 2 for MSE and 1 for MAE

    Returns:
      train_coefs: training coeffcients for each compartment
      valid_coefs: valid coeffcients for each compartment
    """
    (_, gt_list, gt_indicator, _, _) = ground_truth_timeseries

    # Recovered
    recovered_train, recovered_valid = model_utils.compartment_base(
        gt_list["recovered"], gt_indicator["recovered"], num_train_steps,
        num_known_steps)
    # Death
    death_train, death_valid = model_utils.compartment_base(
        gt_list["death"], gt_indicator["death"], num_train_steps,
        num_known_steps)
    # Confirmed
    confirmed_train, confirmed_valid = model_utils.compartment_base(
        gt_list["confirmed"], gt_indicator["confirmed"], num_train_steps,
        num_known_steps)
    # Hospitalized
    hospitalized_train, hospitalized_valid = model_utils.compartment_base(
        gt_list["hospitalized"], gt_indicator["hospitalized"], num_train_steps,
        num_known_steps)
    # Hospitalized cumulative
    hospitalized_cumulative_train, hospitalized_cumulative_valid = model_utils.compartment_base(
        gt_list["hospitalized_cumulative"],
        gt_indicator["hospitalized_cumulative"], num_train_steps,
        num_known_steps)
    # ICU
    icu_train, icu_valid = model_utils.compartment_base(gt_list["icu"],
                                                        gt_indicator["icu"],
                                                        num_train_steps,
                                                        num_known_steps)
    # Ventilator
    ventilator_train, ventilator_valid = model_utils.compartment_base(
        gt_list["ventilator"], gt_indicator["ventilator"], num_train_steps,
        num_known_steps)

    train_coefs = [
        0, (death_train / recovered_train)**power, 1,
        (death_train / confirmed_train)**power,
        (death_train / hospitalized_train)**power,
        (death_train / hospitalized_cumulative_train)**power,
        (death_train / icu_train)**power,
        (death_train / ventilator_train)**power
    ]

    valid_coefs = [
        0, (death_valid / recovered_valid)**power, 1,
        (death_valid / confirmed_valid)**power,
        (death_valid / hospitalized_valid)**power,
        (death_valid / hospitalized_cumulative_valid)**power,
        (death_valid / icu_valid)**power,
        (death_valid / ventilator_valid)**power
    ]

    train_coefs = np.nan_to_num(train_coefs).tolist()
    valid_coefs = np.nan_to_num(valid_coefs).tolist()

    return train_coefs, valid_coefs

  def seir_dynamics(self, current_state, seir_variables):
    """Model dynamics."""

    (first_dose_vaccine_ratio_per_day, second_dose_vaccine_ratio_per_day,
     average_contact_id, average_contact_iud, reinfectable_rate, alpha,
     diagnosis_rate, recovery_rate_id, recovery_rate_iud, recovery_rate_h,
     recovery_rate_i, recovery_rate_v, hospitalization_rate, icu_rate,
     ventilator_rate, death_rate_id, death_rate_h, death_rate_i,
     death_rate_v) = seir_variables

    # pylint: disable=unused-variable
    (exposed_t, infected_d_t, infected_ud_t, recovered_d_t, recovered_ud_t,
     hospitalized_t, hospitalized_cumulative_t, hospitalized_increase_t, icu_t,
     ventilator_t, death_t, population_t, reinfectable_d_t, reinfectable_ud_t,
     reinfectable_vaccine_t, vaccine_immuned_t,
     infected_ud_increase_t) = tf.unstack(current_state)

    # Setting the susceptible so that the population adds up to a constant.
    normalized_susceptible_t = 1.0 - (
        exposed_t + infected_d_t + infected_ud_t + recovered_d_t +
        recovered_ud_t + hospitalized_t + icu_t + ventilator_t + death_t +
        vaccine_immuned_t) / population_t
    normalized_susceptible_t = tf.nn.relu(normalized_susceptible_t)

    # Differential change on vaccine immuned.
    d_vaccine_immuned_dt = (
        first_dose_vaccine_ratio_per_day * population_t +
        second_dose_vaccine_ratio_per_day * population_t -
        reinfectable_vaccine_t - vaccine_immuned_t)

    # Differential change on reinfectable after vaccination.
    d_reinfectable_vaccine_dt = vaccine_immuned_t * 1.0 / constants.VACCINE_IMMUNITY_DURATION

    # Differential change on exposed
    d_exposed_dt = (average_contact_id * infected_d_t +
                    average_contact_iud * infected_ud_t
                   ) * normalized_susceptible_t - alpha * exposed_t

    # Differential change on infected, documented and undocumented
    d_infected_d_dt = (
        diagnosis_rate * infected_ud_t - recovery_rate_id * infected_d_t -
        death_rate_id * infected_d_t - hospitalization_rate * infected_d_t)

    d_infected_ud_dt = (
        alpha * exposed_t - diagnosis_rate * infected_ud_t -
        recovery_rate_iud * infected_ud_t)

    d_infected_ud_increase_dt = alpha * exposed_t - infected_ud_increase_t

    # Differential change on recovered, documented and undocumented
    d_recovered_d_dt = (
        recovery_rate_id * infected_d_t + recovery_rate_h * hospitalized_t -
        reinfectable_rate * recovered_d_t)

    d_recovered_ud_dt = (
        recovery_rate_iud * infected_ud_t - reinfectable_rate * recovered_ud_t)

    # Differential change on hospitalized
    d_hospitalized_d_dt = (
        hospitalization_rate * infected_d_t -
        (death_rate_h + recovery_rate_h + icu_rate) * hospitalized_t +
        recovery_rate_i * icu_t)

    d_hospitalized_cumulative_d_dt = (hospitalization_rate * infected_d_t)

    d_hospitalized_increase_d_dt = (
        hospitalization_rate * infected_d_t - hospitalized_increase_t)

    # Differential change on icu
    d_icu_d_dt = (
        icu_rate * hospitalized_t -
        (death_rate_i + recovery_rate_i + ventilator_rate) * icu_t +
        recovery_rate_v * ventilator_t)

    # Differential change on ventilator
    d_ventilator_d_dt = (
        ventilator_rate * icu_t -
        (death_rate_v + recovery_rate_v) * ventilator_t)

    # Differential change on death, documented
    d_death_d_dt = (
        death_rate_id * infected_d_t + death_rate_h * hospitalized_t +
        death_rate_i * icu_t + death_rate_v * ventilator_t)

    # Differential change on recovered, who may get the disease again.
    d_reinfectable_d_dt = reinfectable_rate * recovered_d_t

    d_reinfectable_ud_dt = reinfectable_rate * recovered_ud_t

    all_state_derivatives = [
        d_exposed_dt, d_infected_d_dt, d_infected_ud_dt, d_recovered_d_dt,
        d_recovered_ud_dt, d_hospitalized_d_dt, d_hospitalized_cumulative_d_dt,
        d_hospitalized_increase_d_dt, d_icu_d_dt, d_ventilator_d_dt,
        d_death_d_dt, -d_death_d_dt, d_reinfectable_d_dt, d_reinfectable_ud_dt,
        d_reinfectable_vaccine_dt, d_vaccine_immuned_dt,
        d_infected_ud_increase_dt
    ]

    return tf.stack(all_state_derivatives)

  def compute_losses(self,
                     hparams,
                     train_coefs,
                     valid_coefs,
                     propagated_states,
                     ground_truth_timeseries,
                     r_eff,
                     train_start_index,
                     train_end_index,
                     valid_start_index,
                     valid_end_index,
                     num_forecast_steps,
                     quantiles=None):

    train_loss_coefs = hparams["train_loss_coefs"]
    valid_loss_coefs = hparams["valid_loss_coefs"]
    time_scale_weight = hparams["time_scale_weight"]
    width_coef_train = hparams["width_coef_train"]
    width_coef_valid = hparams["width_coef_valid"]
    quantile_cum_viol_coef = hparams["quantile_cum_viol_coef"]
    increment_loss_weight = hparams["increment_loss_weight"]
    train_crps_weight = hparams["train_crps_weight"]
    valid_crps_weight = hparams["valid_crps_weight"]

    (_, gt_list, gt_indicator, _, _) = ground_truth_timeseries

    unstacked_propagated_states = tf.unstack(propagated_states, axis=1)
    pred_infected = unstacked_propagated_states[1]
    pred_recovered = unstacked_propagated_states[3]
    pred_hospitalized = unstacked_propagated_states[5]
    pred_hospitalized_cumulative = unstacked_propagated_states[6]
    pred_icu = unstacked_propagated_states[8]
    pred_ventilator = unstacked_propagated_states[9]
    pred_death = unstacked_propagated_states[10]
    pred_reinfected = unstacked_propagated_states[12]

    pred_confirmed = (
        pred_infected + pred_recovered + pred_death + pred_hospitalized +
        pred_icu + pred_ventilator + pred_reinfected)

    train_start_index = tf.identity(train_start_index)
    train_end_index = tf.identity(train_end_index)
    valid_start_index = tf.identity(valid_start_index)
    valid_end_index = tf.identity(valid_end_index)
    if quantiles is not None:
      quantiles = tf.constant(quantiles, dtype=tf.float32)

    # Use quantile loss if the value of quantiles are given
    def loss(pred_states,
             gt_list,
             gt_indicator,
             train_start_index,
             train_end_index,
             valid_start_index,
             valid_end_index,
             time_scale_weight=0,
             is_training=True):
      if quantiles is not None:

        if is_training:
          train_loss = losses.weighted_interval_loss(
              quantile_pred_states=pred_states,
              tau_list=quantiles,
              gt_list=gt_list,
              gt_indicator=gt_indicator,
              begin_timestep=train_start_index,
              end_timestep=train_end_index,
              time_scale_weight=time_scale_weight,
              width_coef=width_coef_train)
          valid_loss = losses.weighted_interval_loss(
              quantile_pred_states=pred_states,
              tau_list=quantiles,
              gt_list=gt_list,
              gt_indicator=gt_indicator,
              begin_timestep=valid_start_index,
              end_timestep=valid_end_index,
              time_scale_weight=time_scale_weight,
              width_coef=width_coef_train)
        else:
          train_loss = losses.weighted_interval_loss(
              quantile_pred_states=pred_states,
              tau_list=quantiles,
              gt_list=gt_list,
              gt_indicator=gt_indicator,
              begin_timestep=train_start_index,
              end_timestep=train_end_index,
              time_scale_weight=time_scale_weight,
              width_coef=width_coef_valid)
          valid_loss = losses.weighted_interval_loss(
              quantile_pred_states=pred_states,
              tau_list=quantiles,
              gt_list=gt_list,
              gt_indicator=gt_indicator,
              begin_timestep=valid_start_index,
              end_timestep=valid_end_index,
              time_scale_weight=time_scale_weight,
              width_coef=width_coef_valid)
        train_loss += train_crps_weight * losses.crps_loss(
            quantile_pred_states=pred_states,
            tau_list=quantiles,
            gt_list=gt_list,
            gt_indicator=gt_indicator,
            begin_timestep=train_start_index,
            end_timestep=train_end_index,
            time_scale_weight=time_scale_weight)
        valid_loss += valid_crps_weight * losses.crps_loss(
            quantile_pred_states=pred_states,
            tau_list=quantiles,
            gt_list=gt_list,
            gt_indicator=gt_indicator,
            begin_timestep=valid_start_index,
            end_timestep=valid_end_index,
            time_scale_weight=time_scale_weight)

      else:
        train_loss = losses.state_estimation_loss(
            pred_states=pred_states,
            gt_list=gt_list,
            gt_indicator=gt_indicator,
            begin_timestep=train_start_index,
            end_timestep=train_end_index,
            time_scale_weight=time_scale_weight,
            increment_loss_weight=increment_loss_weight,
            num_forecast_steps=num_forecast_steps)
        valid_loss = losses.state_estimation_loss(
            pred_states=pred_states,
            gt_list=gt_list,
            gt_indicator=gt_indicator,
            begin_timestep=valid_start_index,
            end_timestep=valid_end_index,
            time_scale_weight=time_scale_weight,
            increment_loss_weight=increment_loss_weight,
            num_forecast_steps=num_forecast_steps)
      return train_loss, valid_loss

    infected_doc_train_loss, infected_doc_valid_loss = loss(
        pred_infected,
        gt_list["infected"],
        gt_indicator["infected"],
        train_start_index,
        train_end_index,
        valid_start_index,
        valid_end_index,
        time_scale_weight=time_scale_weight)
    recovered_doc_train_loss, recovered_doc_valid_loss = loss(
        pred_recovered + pred_reinfected,
        gt_list["recovered"],
        gt_indicator["recovered"],
        train_start_index,
        train_end_index,
        valid_start_index,
        valid_end_index,
        time_scale_weight=time_scale_weight)
    death_train_loss, death_valid_loss = loss(
        pred_death,
        gt_list["death"],
        gt_indicator["death"],
        train_start_index,
        train_end_index,
        valid_start_index,
        valid_end_index,
        time_scale_weight=time_scale_weight)
    hospitalized_train_loss, hospitalized_valid_loss = loss(
        pred_hospitalized + pred_icu + pred_ventilator,
        gt_list["hospitalized"],
        gt_indicator["hospitalized"],
        train_start_index,
        train_end_index,
        valid_start_index,
        valid_end_index,
        time_scale_weight=time_scale_weight)
    hospitalized_cumulative_train_loss, hospitalized_cumulative_valid_loss = loss(
        pred_hospitalized_cumulative,
        gt_list["hospitalized_cumulative"],
        gt_indicator["hospitalized_cumulative"],
        train_start_index,
        train_end_index,
        valid_start_index,
        valid_end_index,
        time_scale_weight=time_scale_weight)
    icu_train_loss, icu_valid_loss = loss(
        pred_icu + pred_ventilator,
        gt_list["icu"],
        gt_indicator["icu"],
        train_start_index,
        train_end_index,
        valid_start_index,
        valid_end_index,
        time_scale_weight=time_scale_weight)
    ventilator_train_loss, ventilator_valid_loss = loss(
        pred_ventilator,
        gt_list["ventilator"],
        gt_indicator["ventilator"],
        train_start_index,
        train_end_index,
        valid_start_index,
        valid_end_index,
        time_scale_weight=time_scale_weight)
    confirmed_train_loss, confirmed_valid_loss = loss(
        pred_confirmed,
        gt_list["confirmed"],
        gt_indicator["confirmed"],
        train_start_index,
        train_end_index,
        valid_start_index,
        valid_end_index,
        time_scale_weight=time_scale_weight)

    train_loss_overall = (
        train_coefs[0] * train_loss_coefs[0] * infected_doc_train_loss +
        train_coefs[1] * train_loss_coefs[1] * recovered_doc_train_loss +
        train_coefs[2] * train_loss_coefs[2] * death_train_loss +
        train_coefs[3] * train_loss_coefs[3] * confirmed_train_loss +
        train_coefs[4] * train_loss_coefs[4] * hospitalized_train_loss +
        train_coefs[5] *
        (train_loss_coefs[5] * hospitalized_cumulative_train_loss) +
        train_coefs[6] * train_loss_coefs[6] * icu_train_loss +
        train_coefs[7] * train_loss_coefs[7] * ventilator_train_loss)

    valid_loss_overall = (
        valid_coefs[0] * valid_loss_coefs[0] * infected_doc_valid_loss +
        valid_coefs[1] * valid_loss_coefs[1] * recovered_doc_valid_loss +
        valid_coefs[2] * valid_loss_coefs[2] * death_valid_loss +
        valid_coefs[3] * valid_loss_coefs[3] * confirmed_valid_loss +
        valid_coefs[4] * valid_loss_coefs[4] * hospitalized_valid_loss +
        valid_coefs[5] *
        (valid_loss_coefs[5] * hospitalized_cumulative_valid_loss) +
        valid_coefs[6] * valid_loss_coefs[6] * icu_valid_loss +