"""Energy carbon and cost model for natural gas."""

from typing import Sequence

from absl import logging
import gin
import numpy as np
import pandas as pd

from smart_control.models.base_energy_cost import BaseEnergyCost
from smart_control.utils import constants

# Source: https://www.eia.gov/dnav/ng/hist/n3035ca3m.htm
# For 2020, units: Dollars per Thousand Cubic Feet
GAS_PRICE_BY_MONTH_SOURCE = (
    9.02,
    8.35,
    7.77,
    7.26,
    6.69,
    6.86,
    6.77,
    6.76,
    6.99,
    7.19,
    7.96,
    8.98,
)


@gin.configurable()
class NaturalGasEnergyCost(BaseEnergyCost):
  """Energy cost and carbon emission model for reward function.

  Attributes:
    gas_price_per_month: Cost/energy consumed [$/1000 cubic feet] by month.
  """

  def __init__(
      self, gas_price_by_month: Sequence[float] = GAS_PRICE_BY_MONTH_SOURCE
  ):
    assert (
        len(gas_price_by_month) == 12
    ), 'Gas price per month must have exactly 12 values.'
    # Convert the month-by-month gas price from $/1000 cubic feet to $/Joule.
    self._month_gas_price = (
        np.array(gas_price_by_month)
        / constants.KWH_PER_KFT3_GAS
        / constants.JOULES_PER_KWH
    )

    # Convert the carbon from kg Carbon / 1000 cubic feet to kgC / J.
    self._carbon_rate = (
        constants.GAS_CO2
        / constants.KWH_PER_KFT3_GAS
        / constants.JOULES_PER_KWH
    )

  def cost(
      self, start_time: pd.Timestamp, end_time: pd.Timestamp, energy_rate: float
  ) -> float:
    """Returns the cost of energy from this time step.

    Args:
      start_time: start of window
      end_time: end of window
      energy_rate: thermal power in W applied during the window.

    Returns:
      cost of energy consumed in window in USD.

    Raises ValueError if the energy is negative since natural gas can only be
    applied for heating.
    """

    if energy_rate < 0.0:
      logging.warning(
          'Negative natural gas energy %3.2f. Setting to 0.', energy_rate
      )
      energy_rate = 0.0

    # Get the total time in seconds.
    dt = (end_time - start_time).total_seconds()

    # Get the energy price in $ / J.
    energy_price = self._month_gas_price[start_time.month - 1]

    # Get the energy in J.
    energy_consumed = energy_rate * dt

    # Compute the cost of the energy:
    # $ = energy_price [$ / J] * energy_consumed [J]
    return energy_price * energy_consumed

  def carbon(
      self, start_time: pd.Timestamp, end_time: pd.Timestamp, energy_rate: float
  ) -> float:
    """Returns the carbon produced in this time step.

    Args:
      start_time: start of window
      end_time: end of window
      energy_rate: thermal power in W applied during the window.

    Returns:
      carbon mass consumed in kg.

    Raises ValueError if the energy is negative since natural gas can only be
    applied for heating.
    """

    if energy_rate < 0.0:
      logging.warning(
          'Negative natural gas energy %3.2f. Setting to 0.', energy_rate
      )
      energy_rate = 0.0

    dt = (end_time - start_time).total_seconds()
    # Get the energy in J.
    energy_consumed = energy_rate * dt
    # Return carbon mass [kg] generated by the energy consumed [J].
    return self._carbon_rate * energy_consumed
