"""
HeatPumpExport exports heatpump coefficient into several formats
SPDX - License - Identifier: LGPL - 3.0 - or -later
Copyright © 2021 Project Author Peter Yefi peteryefi@gmail.com
"""
import numpy as np
from scipy.optimize import curve_fit
from typing import Dict, Tuple
import pandas as pd


class HeatPumpExport:
  """
  Exports heat pump values as coefficients
  of some defined function
  """

  def __init__(self, base_path, city):
    self._base_path = (base_path / 'heat_pumps/coefficients.xlsx')
    self._city = city

  def export_xlsx(self):
    """
    Writes the coefficients computed from heat performance
    and cooling performance data to excel sheet
    :return: None
    """
    writer = pd.ExcelWriter(self._base_path)
    heat_column_names = ["a1", "a2", "a3", "a4", "a5", "a6"]
    cool_column_names = ["b1", "b2", "b3", "b4", "b5", "b6"]
    heat_coff, cool_coff = self._compute_coefficients()
    for (k_cool, v_cool), (k_heat, v_heat) in \
        zip(heat_coff.items(), cool_coff.items()):
      heat_df = pd.DataFrame([v_heat["heat_cap"], v_heat["comp_power"]], columns=heat_column_names,
                             index=["Heat Capacity", "Compressor Power"])
      heat_df.to_excel(writer, sheet_name=k_heat)
      cool_df = pd.DataFrame([v_heat["cool_cap"], v_heat["comp_power"]], columns=cool_column_names,
                             index=["Cooling Capacity", "Compressor Power"])
      cool_df.to_excel(writer, sheet_name=k_cool, startrow=10)

      writer.save()

  def _compute_coefficients(self) -> Tuple[Dict, Dict]:
    """
    Compute heat output and electrical demand coefficients
    from heating and cooling performance data
    :return: Tuple[Dict, Dict]
    """
    out_temp = [25, 30, 32, 35, 40, 45] * 6
    heat_x_values = np.repeat([-5, 0, 7, 10, 15], 6)
    cool_x_values = np.repeat([6, 7, 8, 9, 10, 11], 6)
    cooling_coff = {}
    heating_coff = {}
    for energy_system in self._city.energy_systems:
      # Compute heat output coefficients
      heating_cap_popt, _ = curve_fit(self._objective_function, [heat_x_values, out_temp],
                                      energy_system.heat_pump.heating_capacity)
      heating_comp_power_popt, _ = curve_fit(self._objective_function, [heat_x_values, out_temp],
                                             energy_system.heat_pump.heating_comp_power)
      # Compute electricity demand coefficients
      cooling_cap_popt, _ = curve_fit(self._objective_function, [cool_x_values, out_temp],
                                      energy_system.heat_pump.cooling_capacity)
      cooling_comp_power_popt, _ = curve_fit(self._objective_function, [cool_x_values, out_temp],
                                             energy_system.heat_pump.cooling_comp_power)

      heating_coff[energy_system.heat_pump.model] = {"heat_cap": heating_cap_popt.tolist(),
                                                     "comp_power": heating_comp_power_popt.tolist()}
      cooling_coff[energy_system.heat_pump.model] = {"cool_cap": cooling_cap_popt.tolist(),
                                                     "comp_power": cooling_comp_power_popt.tolist()}
    return heating_coff, cooling_coff

  def _objective_function(self, xdata, a1, a2, a3, a4, a5, a6):
    """
    Objective function for computing coefficients
    :param xdata:
    :param a1: float
    :param a2: float
    :param a3: float
    :param a4: float
    :param a5: float
    :param a6: float
    :return:
    """
    x, y = xdata
    return (a1 * x ** 2) + (a2 * x) + (a3 * x * y) + (a4 * y) + (a5 * y ** 2) + a6