import pandas as pd
import calendar


def enrich(city):
  """
  Enrich the city by using the heating and cooling demand data from Excel files.
  """


  # Read the heating and cooling Excel files
  heating_df = pd.read_excel('./input_files/heating.xlsx')
  cooling_df = pd.read_excel('./input_files/cooling.xlsx')

  # Ensure that the data frames have the same number of buildings as city.buildings
  assert len(heating_df) == len(city.buildings), "Mismatch in number of buildings in heating data."
  assert len(cooling_df) == len(city.buildings), "Mismatch in number of buildings in cooling data."

  # Helper function to compute monthly totals
  def get_monthly_totals(hourly_values):
    monthly_totals = []
    hour_index = 0
    for month in range(1, 13):
      days_in_month = calendar.monthrange(2023, month)[1]
      hours_in_month = days_in_month * 24
      monthly_total = sum(hourly_values[hour_index:hour_index + hours_in_month])
      monthly_totals.append(monthly_total)
      hour_index += hours_in_month
    return monthly_totals

  # Iterate over buildings and data frames
  for idx, (building, heating_row, cooling_row) in enumerate(
      zip(city.buildings, heating_df.itertuples(index=False), cooling_df.itertuples(index=False))):
    # Skip the 'building' column (assuming it's the first column)
    heating_demand_hourly = list(heating_row)[1:]  # Exclude the 'building' column
    cooling_demand_hourly = list(cooling_row)[1:]  # Exclude the 'building' column

    # Convert demands to Joules (1 kWh = 3.6e6 J)
    heating_demand_hourly_joules = [value * 3.6e6 for value in heating_demand_hourly]
    cooling_demand_hourly_joules = [value * 3.6e6 for value in cooling_demand_hourly]

    # Store the demands in the building object
    building.heating_demand = {}
    building.cooling_demand = {}

    building.heating_demand['hour'] = heating_demand_hourly_joules
    building.cooling_demand['hour'] = cooling_demand_hourly_joules

    # Compute monthly demands
    building.heating_demand['month'] = get_monthly_totals(heating_demand_hourly_joules)
    building.cooling_demand['month'] = get_monthly_totals(cooling_demand_hourly_joules)

    # Compute yearly demands
    building.heating_demand['year'] = [sum(building.heating_demand['month'])]
    building.cooling_demand['year'] = [sum(building.cooling_demand['month'])]