import json
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.ticker import MaxNLocator
from matplotlib.patches import Patch

output_path = (Path(__file__).parent / 'out_files').resolve()

# File paths for the three JSON files
file1 = output_path / 'base_case_buildings_data.json'
file2 = output_path / 'air_to_air_hp_buildings_data.json'
file3 = output_path / 'air_to_water_hp_buildings_data.json'

# Opening and reading all three JSON files at the same time
with open(file1) as f1, open(file2) as f2, open(file3) as f3:
    base_case = json.load(f1)
    air = json.load(f2)
    water = json.load(f3)

month_names = ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec']
x = np.arange(len(month_names))  # the label locations


# Scenario labels and color palette
scenarios = ['Scenario 1', 'Scenario 2']
colors = ['#66B2FF', '#e74c3c']  # Blue for Scenario 1, Red for Scenario 2
width = 0.25  # Width for each bar

# Creating the grid for peak load comparisons across buildings
fig, axes = plt.subplots(2, 5, figsize=(20, 10), dpi=96)
fig.suptitle('Yearly Heating and Cooling Peak Load Comparison Across Buildings', fontsize=16, weight='bold', alpha=0.8)
axes = axes.flatten()

for idx, building_name in enumerate(base_case.keys()):
    # Extracting heating and cooling peak loads for each scenario
    heating_peak_load = [
        air[building_name]["heating_peak_load_kW"],
        water[building_name]["heating_peak_load_kW"]
    ]
    cooling_peak_load = [
        air[building_name]["cooling_peak_load_kW"],
        water[building_name]["cooling_peak_load_kW"]
    ]

    ax = axes[idx]
    x = np.arange(2)  # X locations for the "Heating" and "Cooling" groups

    # Plotting each scenario for heating and cooling
    for i in range(len(scenarios)):
        ax.bar(x[0] - width + i * width, heating_peak_load[i], width, color=colors[i], zorder=2)
        ax.bar(x[1] - width + i * width, cooling_peak_load[i], width, color=colors[i], zorder=2)

    # Grid and styling
    ax.grid(which="major", axis='x', color='#DAD8D7', alpha=0.5, zorder=1)
    ax.grid(which="major", axis='y', color='#DAD8D7', alpha=0.5, zorder=1)

    # Axis and title settings
    ax.set_title(building_name, fontsize=14, weight='bold', alpha=0.8, pad=10)
    ax.set_xticks(x)
    ax.set_xticklabels(['Heating Peak Load', 'Cooling Peak Load'])
    ax.xaxis.set_major_locator(MaxNLocator(integer=True))
    ax.yaxis.set_major_locator(MaxNLocator(integer=True))
    if idx % 5 == 0:
        ax.set_ylabel('Peak Load (kW)', fontsize=12, labelpad=10)

# Custom legend handles to ensure color match with scenarios
legend_handles = [Patch(color=colors[i], label=scenarios[i]) for i in range(len(scenarios))]

# Global legend and layout adjustments
fig.legend(handles=legend_handles, loc='upper right', ncol=1)
plt.tight_layout(rect=[0, 0.03, 1, 0.95])
plt.savefig(output_path / 'peak_loads.png')