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

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
width = 0.25  # the width of the bars

# Prettier colors for each scenario
colors = ['#66B2FF', '#e74c3c']  # Blue, Red, Green


# Plotting heating data for all buildings in a 2x5 grid
fig, axes = plt.subplots(2, 5, figsize=(20, 10), dpi=96)
fig.suptitle('Monthly DHW Consumption Comparison Across Buildings', fontsize=16, weight='bold', alpha=0.8)
axes = axes.flatten()

for idx, building_name in enumerate(base_case.keys()):
    heating_data = [list(data["monthly_dhw_consumption_kWh"].values()) for data in
                    [base_case[building_name], water[building_name]]]

    ax = axes[idx]
    for i, data in enumerate(heating_data):
        ax.bar(x + (i - 1) * width, data, width, label=f'Scenario {i+1}', color=colors[i], zorder=2)

    # Grid settings
    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 labels and title
    ax.set_title(building_name, fontsize=14, weight='bold', alpha=0.8, pad=10)
    ax.set_xticks(x)
    ax.set_xticklabels(month_names, rotation=45, ha='right')
    ax.xaxis.set_major_locator(MaxNLocator(integer=True))
    ax.yaxis.set_major_locator(MaxNLocator(integer=True))
    if idx % 5 == 0:
        ax.set_ylabel('DHW Consumption (kWh)', fontsize=12, labelpad=10)

fig.legend(['Base Case', 'Scenario 1&2'], loc='upper right', ncol=3)
plt.tight_layout(rect=[0, 0.03, 1, 0.95])
plt.savefig(output_path / 'monthly_dhw.png')