import pandas as pd 
import numpy as np

from matplotlib import pyplot as plt 
import matplotlib.dates as mdates

import argparse
import yaml

from mix_simul.scenarios import Scenario 

parser = argparse.ArgumentParser()
parser.add_argument("--begin", help="begin date (YYYY-MM-DD), between 1985-01-01 and 2015-01-01", required=True)
parser.add_argument("--end", help="end date (YYYY-MM-DD), between 1985-01-01 and 2015-01-01", required=True)
parser.add_argument("--flexibility", help="enable load flexibility modeling", action="store_true")
parser.add_argument("--scenarios", help="path to scenarios parameters yml file", default="scenarios/rte.yml")
args = parser.parse_args()

with open(args.scenarios, 'r') as f:
    scenarios = yaml.load(f, Loader=yaml.FullLoader)

potential = pd.read_parquet("data/potential.parquet")
potential = potential.loc['1985-01-01 00:00:00':'2015-01-01 00:00:00', :]
potential.fillna(0, inplace=True)

begin = args.begin
end = args.end
flexibility = args.flexibility

potential = potential.loc[(
    slice(f'{begin} 00:00:00', f'{end} 00:00:00'), 'FR'), :]

# intermittent sources potential
p = potential[["onshore", "offshore", "solar"]].to_xarray().to_array()
p = np.insert(p, 3, 0.7, axis=0)  # nuclear power-like

times = potential.index.get_level_values(0)

n_scenarios = len(scenarios)

fig, axes = plt.subplots(nrows=n_scenarios, ncols=2, sharex="col", sharey=True)
w, h = fig.get_size_inches()
fig.set_figwidth(w*1.5)
fig.set_figheight(h*1.5)

fig_storage, axes_storage = plt.subplots(nrows=n_scenarios, ncols=2, sharex="col", sharey=True)
fig_storage.set_figwidth(w*1.5)
fig_storage.set_figheight(h*1.5)

fig_dispatch, axes_dispatch = plt.subplots(nrows=n_scenarios, ncols=2, sharex="col", sharey=True)
fig_dispatch.set_figwidth(w*1.5)
fig_dispatch.set_figheight(h*1.5)

fig_gap_distribution, axes_gap_distribution = plt.subplots(nrows=int(np.ceil(n_scenarios/2)), ncols=2, sharex="col", sharey=True)
fig_gap_distribution.set_figwidth(w*1.5)
fig_gap_distribution.set_figheight(h*1.5)

months = [
    "Février",
    "Juin"
]

labels = [
    "load (GW)",
    "production (GW)",
    "available power (production-storage) (GW)",
    "power deficit"
]

labels_storage = [
    "Batteries (TWh)",
    "STEP (TWh)",
    "P2G (TWh)"
]

labels_dispatch = [
    "Hydro (GW)",
    "Biomass (GW)",
    "Thermal (GW)"
]

date_fmt = mdates.DateFormatter('%d/%m')

# for step in np.linspace(start, stop, 2050-2022, True)[::-1]:
row = 0
for scenario in scenarios:
    if not flexibility:
        scenarios[scenario]["flexibility_power"] = 0
    scenario_model = Scenario(**scenarios[scenario])
    S, load, production, gap, storage, dp = scenario_model.run(times, p)

    print(f"{scenario}:", S, gap.max(), np.quantile(gap, 0.95))
    print(f"exports: {np.minimum(np.maximum(-gap, 0), 39).sum()/1000} TWh; imports: {np.minimum(np.maximum(gap, 0), 39).sum()/1000} TWh")
    print(f"dispatchable: " + ", ".join([f"{dp[i].sum()/1000:.2f} TWh" for i in range(dp.shape[0])]))

    potential['load'] = load
    potential['production'] = production
    potential['available'] = production-np.diff(storage.sum(axis=0), append=0)

    for i in range(3):
        potential[f"storage_{i}"] = np.diff(storage[i,:], append=0)#storage[i,:]/1000
        potential[f"storage_{i}"] = storage[i,:]/1000

    for i in range(dp.shape[0]):
        potential[f"dispatch_{i}"] = dp[i,:]

    potential["dispatch"] = dp.sum(axis=0)

    data = [
        potential.loc[(slice('2013-02-01 00:00:00',
                       '2013-03-01 00:00:00'), 'FR'), :],
        potential.loc[(slice('2013-06-01 00:00:00',
                       '2013-07-01 00:00:00'), 'FR'), :]
    ]

    for col in range(2):
        ax = axes[row, col] if axes.ndim > 1 else axes[col]
        ax.plot(data[col].index.get_level_values(0), data[col]
                ["load"], label="adjusted load (GW)", lw=1)
        ax.plot(data[col].index.get_level_values(0), data[col]
                ["production"], label="production (GW)", ls="dotted", lw=1)
        ax.plot(data[col].index.get_level_values(0), data[col]["available"],
                label="available power (production-d(storage)/dt) (GW)", lw=1)

        ax.fill_between(
            data[col].index.get_level_values(0),
            data[col]["available"],
            data[col]["load"],
            where=data[col]["load"] > data[col]["available"],
            color='red',
            alpha=0.15
        )

        ax.xaxis.set_major_formatter(date_fmt)
        ax.text(
            0.5, 0.87, f"Scénario {scenario} ({months[col]})", ha='center', transform=ax.transAxes)

        ax.set_ylim(10, 210)

        ax = axes_storage[row, col] if axes.ndim > 1 else axes_storage[col]
        for i in np.arange(3):
            if i == 2:
                base = 0
            else:
                base = np.sum([data[col][f"storage_{j}"] for j in np.arange(i+1,3)], axis=0)
            
            ax.fill_between(data[col].index.get_level_values(0), base, base+data[col]
                [f"storage_{i}"], label=f"storage {i}", alpha=0.5)

            ax.plot(data[col].index.get_level_values(0), base+data[col]
                [f"storage_{i}"], label=f"storage {i}", lw=0.25)

        ax.xaxis.set_major_formatter(date_fmt)
        ax.text(
            0.5, 0.87, f"Scénario {scenario} ({months[col]})", ha='center', transform=ax.transAxes)

        ax = axes_dispatch[row, col] if axes.ndim > 1 else axes_dispatch[col]
        for i in range(dp.shape[0]):
            if i == 0:
                base = 0
            else:
                base = np.sum([data[col][f"dispatch_{j}"] for j in np.arange(i)], axis=0)
            
            ax.fill_between(data[col].index.get_level_values(0), base, base+data[col]
                [f"dispatch_{i}"], label=f"dispatch {i}", alpha=0.5)

            ax.plot(data[col].index.get_level_values(0), base+data[col]
                [f"dispatch_{i}"], label=f"dispatch {i}", lw=0.25)

        ax.xaxis.set_major_formatter(date_fmt)
        ax.text(
            0.5, 0.87, f"Scénario {scenario} ({months[col]})", ha='center', transform=ax.transAxes)

        div = int(np.ceil(n_scenarios/2))
        ax = axes_gap_distribution[row%div, row//div] if axes_gap_distribution.ndim > 1 else axes_gap_distribution[row%div]
        hist, bin_edges = np.histogram(-gap, bins=1000)
        hist = np.cumsum(hist)
        hist = 100*(hist - hist.min()) / hist.ptp()

        keep = np.abs(bin_edges[:-1]) < 50
        ax.plot(bin_edges[:-1][keep], hist[keep], label="power gap")
        ax.text(
            0.5, 0.87, f"Scénario {scenario}", ha='center', transform=ax.transAxes)
        

    row += 1

for axs in [axes, axes_dispatch, axes_storage]:
    if axes.ndim > 1:
        for label in axs[-1, 0].get_xmajorticklabels() + axs[-1, 1].get_xmajorticklabels():
            label.set_rotation(30)
            label.set_horizontalalignment("right")
    else:
        for label in axs[0].get_xmajorticklabels() + axs[-1].get_xmajorticklabels():
            label.set_rotation(30)
            label.set_horizontalalignment("right")

flex = "With" if flexibility else "Without"

plt.subplots_adjust(wspace=0, hspace=0)
fig.suptitle(f"Simulations based on {begin}--{end} weather data.\n{flex} consumption flexibility; no nuclear seasonality (unrealistic)")
fig.text(1, 0, 'Lucas Gautheron', ha="right")
fig.legend(labels, loc='lower right', bbox_to_anchor=(1, -0.1),
           ncol=len(labels), bbox_transform=fig.transFigure)
fig.savefig("output/load_supply.png", bbox_inches="tight", dpi=200)

fig_storage.suptitle(f"Simulations based on {begin}--{end} weather data.\n{flex} consumption flexibility; no nuclear seasonality (unrealistic)")
fig_storage.text(1, 0, 'Lucas Gautheron', ha="right")
fig_storage.legend(labels_storage, loc='lower right', bbox_to_anchor=(1, -0.1),
           ncol=len(labels_storage), bbox_transform=fig_storage.transFigure)
fig_storage.savefig("output/storage.png", bbox_inches="tight", dpi=200)

fig_dispatch.suptitle(f"Simulations based on {begin}--{end} weather data.\n{flex} consumption flexibility; no nuclear seasonality (unrealistic)")
fig_dispatch.text(1, 0, 'Lucas Gautheron', ha="right")
fig_dispatch.legend(labels_dispatch, loc='lower right', bbox_to_anchor=(1, -0.1),
           ncol=len(labels_dispatch), bbox_transform=fig_dispatch.transFigure)
fig_dispatch.savefig("output/dispatch.png", bbox_inches="tight", dpi=200)

fig_gap_distribution.suptitle(f"Power gap cumulative distribution (%)\nSimulations based on {begin}--{end} weather data.\n{flex} consumption flexibility; no nuclear seasonality (unrealistic)")
fig_gap_distribution.legend(["Power gap (available-load) (GW)"], loc='lower right', bbox_to_anchor=(1, -0.1),
           ncol=1, bbox_transform=fig_dispatch.transFigure)
fig_gap_distribution.text(1, 0, 'Lucas Gautheron', ha="right")
fig_gap_distribution.savefig("output/gap_distribution.png", bbox_inches="tight", dpi=200)

plt.show()
