import cvxpy as cp
import pandas as pd


# Load data
data = pd.read_csv('new_file_extended.csv')
demand = data['Q_tot_mpc'].to_list()
demand_watts = [x * 1000 for x in demand]
t_out = data['T_out'].to_list()
p_electricity = data['Electricity_Price'].to_list()
data['Datetime'] = pd.to_datetime(data['Datetime'])

cp_water = 4182  # J/kgK

# Heat Pump Sizing
hp_heating_cap = max(demand_watts) * 0.7
hp_nominal_cop = 2.5
hp_delta_t = 5
hp_cop_curve_coefficients = [1.039924, 0.0146, 6e-06, -0.05026, 0.000635, -0.000154]
# TES Sizing
tank_volume = round(max(demand_watts) * 3.6e3 / (1000 * cp_water * 15))
ua = 0.28
time_step_size = 900  # Time step size in seconds (15 minutes)
control_horizon = 24  # Control horizon (96 time steps = 24 hours)
initial_temperature = 40
lower_limit = 40
upper_limit = 55

variable_names = ["t_sup_hp", "t_tank", "t_ret", "m_ch", "m_dis", "q_hp", "hp_cop", "hp_electricity",
                  "electricity_cost"]
num_hours = len(demand)
variables = {name: [0] * num_hours for name in variable_names}
t_sup_hp, t_tank, t_ret, m_ch, m_dis, q_hp, hp_cop, hp_electricity, electricity_cost = [variables[name] for name in variable_names]
t_tank[0] = 40


# Define the optimization function
def optimize_heating_system(heating_demand, initial_tank_temp, electricity_price, time_step_size=900, cp_water=4182,
                            volume=tank_volume, hp_nominal_efficiency=hp_nominal_cop, hp_cap=hp_heating_cap,
                            lower_limit_tes=lower_limit, upper_limit_tes=upper_limit):
  time_horizon = len(heating_demand)
  # Define problem variables
  storage_charge = cp.Variable(time_horizon, nonneg=True)  # Energy from heat pump to tank
  storage_discharge = cp.Variable(time_horizon, nonneg=True)  # Energy from tank to house
  heat_pump_energy = cp.Variable(time_horizon, nonneg=True)  # Heat pump energy
  tank_temperature = cp.Variable(time_horizon)  # Tank temperature (°C)

  # Calculate the change in tank temperature
  temperature_change = (time_step_size / (1000 * volume * cp_water)) * (storage_charge[:-1] - storage_discharge[:-1])

  # Calculate the electricity cost
  electricity_cost = ((heat_pump_energy[:-1] * time_step_size) / (hp_nominal_efficiency * 3600)) @ electricity_price[
                                                                                                   :-1]

  # Define the objective function
  objective = cp.Minimize(cp.sum_squares(heating_demand - storage_discharge)) + cp.Minimize(
    cp.sum(electricity_cost)) + cp.Minimize(cp.sum(heat_pump_energy))

  # Define the constraints
  constraints = [heat_pump_energy <= hp_cap, tank_temperature >= lower_limit_tes,
                 tank_temperature <= upper_limit_tes, tank_temperature[0] == initial_tank_temp,
                 storage_charge <= heat_pump_energy]

  # Specify the tank temperature in the next time step
  constraints.extend(
    [tank_temperature[i + 1] == tank_temperature[i] + temperature_change[i] for i in range(time_horizon - 1)])

  # Create the optimization problem
  problem = cp.Problem(objective, constraints)

  # Solve the problem
  problem.solve(solver=cp.GUROBI)

  # Get the optimized values
  optimized_storage_charge = storage_charge.value
  optimized_storage_discharge = storage_discharge.value
  optimized_heat_pump_energy = heat_pump_energy.value
  optimized_tank_temperature = tank_temperature.value

  return optimized_heat_pump_energy, optimized_storage_discharge, optimized_tank_temperature

# Rolling optimization loop
for j in range(len(demand_watts) - control_horizon):  # Adjust loop to avoid exceeding data length

  # Ensure that we don't exceed the data length
  end_time_step = min(j + control_horizon, len(demand_watts))

  # Adjust the slicing window to the available data
  demand_window = demand_watts[j:end_time_step]
  p_electricity_window = p_electricity[j:end_time_step]

  initial_tank_temperature = t_tank[j]

  # Call the optimization function with adjusted window size
  hp_output, storage_output, storage_temperature = optimize_heating_system(demand_window, initial_tank_temperature,
                                                                           p_electricity_window)

  # Perform necessary calculations using the first values from the optimization output
  q_hp[j] = hp_output[0]

  if q_hp[j] > 0:
    m_ch[j] = q_hp[j] / (cp_water * hp_delta_t)
    t_sup_hp[j] = (q_hp[j] / (m_ch[j] * cp_water)) + t_tank[j]
    t_out_fahrenheit = 1.8 * t_out[j] + 32
    t_tank_fahrenheit = 1.8 * t_tank[j] + 32
    hp_cop[j] = (1 / (hp_cop_curve_coefficients[0] +
                      hp_cop_curve_coefficients[1] * t_tank_fahrenheit +
                      hp_cop_curve_coefficients[2] * t_tank_fahrenheit ** 2 +
                      hp_cop_curve_coefficients[3] * t_out_fahrenheit +
                      hp_cop_curve_coefficients[4] * t_out_fahrenheit ** 2 +
                      hp_cop_curve_coefficients[5] * t_tank_fahrenheit * t_out_fahrenheit)) * hp_nominal_cop
    hp_electricity[j] = q_hp[j] / hp_cop[j]
    electricity_cost[j] = hp_electricity[j] * p_electricity[j] / 4000

  # Update storage discharge and tank temperature for next step
  if storage_output[0] > 0.5 * max(demand_window):
    factor = 6
  else:
    factor = 4
  m_dis[j] = (max(demand_window)) / (cp_water * factor)
  t_ret[j] = t_tank[j] - (demand_watts[j]) / (m_dis[j] * cp_water)
  t_tank[j + 1] = t_tank[j] + ((m_ch[j] * (t_sup_hp[j] - t_tank[j])) + (ua * (t_out[j] - t_tank[j])) / cp_water -
                               m_dis[j] * (t_tank[j] - t_ret[j])) * (time_step_size / (1000 * tank_volume))

# Output results to a CSV file
output = pd.DataFrame(index=data['Datetime'])
output["demand"] = demand_watts
output["q_hp"] = q_hp
output["hp_cop"] = hp_cop
output["hp_electricity_consumption"] = hp_electricity
output["electricity_cost"] = electricity_cost
output["m_ch"] = m_ch
output["m_dis"] = m_dis
output["t_sup_hp"] = t_sup_hp
output["t_tank"] = t_tank
output["t_return"] = t_ret
output.to_csv("results_rolling_window.csv")