from datetime import datetime

import matplotlib.dates as mdates
import pandas
from data import Country, DataCO2, DataConsumption, DataPower, Index, DataPower_ENTSOE
from dateutil.relativedelta import relativedelta
from figure.Figure import Figure, Subplot
from figure.genericfigures import LinearPlot
from lib.pandas import PandasOperation, PandasSelection

# Dates
start = pandas.Timestamp(datetime(2023, 1, 1), tz="UTC")
end = pandas.Timestamp(datetime(2023, 1, 2), tz="UTC")
step = relativedelta(days=1)
# Countries
listCountries = Country.getAllCountries()
listCountries = [Country.PT, Country.ES, Country.FR]
# Categories
listCategory = None


if False:
    tmp = start
    while tmp < end:
        # ENTSO-E data
        generation_ENTSOE = DataPower.Generation(
            dataframe=pandas.DataFrame(), start=tmp, end=tmp + step, listCountry=listCountries
        )
        generation_ENTSOE.downloader = DataPower_ENTSOE.Generation_ENTSOE()

        load_ENTSOE = DataPower.Load(
            dataframe=pandas.DataFrame(), start=tmp, end=tmp + step, listCountry=listCountries
        )
        load_ENTSOE.downloader = DataPower_ENTSOE.Load_ENTSOE()

        export_ENTSOE = DataPower.Export(
            dataframe=pandas.DataFrame(), start=tmp, end=tmp + step, listCountry=Country.getAllCountries()
        )
        export_ENTSOE.downloader = DataPower_ENTSOE.Export_ENTSOE()

        import_ENTSOE = DataPower.Import(
            dataframe=pandas.DataFrame(), start=tmp, end=tmp + step, listCountry=Country.getAllCountries()
        )
        import_ENTSOE.downloader = DataPower_ENTSOE.Import_ENTSOE()

        installedCapacity_ENTSOE = DataPower.InstalledCapacity(
            dataframe=pandas.DataFrame(), start=tmp, end=tmp + step, listCountry=listCountries
        )
        installedCapacity_ENTSOE.downloader = DataPower_ENTSOE.InstalledCapacity_ENTSOE()

        price_ENTSOE = DataPower.DayAheadPrices(
            dataframe=pandas.DataFrame(), start=tmp, end=tmp + step, listCountry=listCountries
        )
        price_ENTSOE.downloader = DataPower_ENTSOE.DayAheadPrices_ENTSOE()

        # # Download data
        # generation_ENTSOE.download()
        # load_ENTSOE.download()
        # export_ENTSOE.download()
        # import_ENTSOE.download()
        # installedCapacity_ENTSOE.download()
        # price_ENTSOE.download()

        # Consumption
        consumption = DataConsumption.DataConsumptionMarkov(
            dataframe=pandas.DataFrame(), start=tmp, end=tmp + step
        )
        consumption.downloader = DataConsumption.DataConsumptionMarkov_Downloader()
        consumption.download()

        CO2Emission = DataCO2.ConsumptionCO2Emission(
            consumption=consumption,
            start=start,
            end=end,
        )

        consumption = DataConsumption.DataConsumptionLeontief(
            dataframe=pandas.DataFrame(), start=tmp, end=tmp + step
        )
        consumption.downloader = DataConsumption.DataConsumptionLeontief_Downloader()
        consumption.download()

        tmp += step


def formatCountry(df):
    df = df.reset_index(level=Index.country)
    df[Index.country] = df[Index.country].str.replace('_load', '')
    df = df.set_index(Index.country, append=True)
    
    return df

generation = DataPower.GenerationBalance(start=start, end=end, listCountry=listCountries)
balance = DataPower.Balance(start=start, end=end, listCountry=listCountries)
load = DataPower.Load(start=start, end=end, listCountry=listCountries)

consumption_Leontief = DataConsumption.DataConsumptionLeontief(
    start=start, end=end, listCountry=listCountries
)
consumption_Leontief.dataframe = formatCountry(consumption_Leontief.dataframe)
consumption1_Leontief = consumption_Leontief.copy()
consumption1_Leontief.dataframe = PandasOperation.sumOver(data=consumption1_Leontief.dataframe, levels=[Index.country, Index.category, Index.datetime])
consumption2_Leontief = consumption_Leontief.copy()
consumption2_Leontief.dataframe = PandasOperation.sumOver(data=consumption2_Leontief.dataframe, levels=[Index.country, Index.source, Index.datetime])

consumption_Markov = DataConsumption.DataConsumptionMarkov(
    start=start, end=end, listCountry=listCountries
)
consumption_Markov.dataframe = formatCountry(consumption_Markov.dataframe)
consumption1_Markov = consumption_Markov.copy()
consumption1_Markov.dataframe = PandasOperation.sumOver(data=consumption1_Markov.dataframe, levels=[Index.country, Index.category, Index.datetime])
# consumption1_Markov = consumption1_Markov.grouping(category="category")
consumption2_Markov = consumption_Markov.copy()
consumption2_Markov.dataframe = PandasOperation.sumOver(data=consumption2_Markov.dataframe, levels=[Index.country, Index.source, Index.datetime])
# consumption2_Markov = consumption2_Markov.grouping(category="source")

summedGeneration = generation.copy()
summedGeneration.dataframe = PandasOperation.sumOver(
    data=summedGeneration.dataframe, levels=[Index.datetime, Index.country]
)
summedGeneration.dataframe[Index.category] = "Generation"
summedGeneration.dataframe = summedGeneration.dataframe.set_index([Index.category], append=True)


summedLoad = load.copy()
summedLoad.dataframe[summedLoad.dataframe < 0] = None
summedLoad.dataframe = PandasOperation.sumOver(data=load.dataframe, levels=[Index.datetime, Index.country])
summedLoad.dataframe[Index.category] = "Load"
summedLoad.dataframe = summedLoad.dataframe.set_index([Index.category], append=True)

yLimTab = {"ES": [-10000, 35000],
"PT": [-500, 8000],
"FR": [-5000, 80000],
"DE": [-10000, 80000],}


def getLeontiefSubplot(figure, nrows, ncols, index_1, index_2, yLim=[-3000, 12000]):
    s1 = Subplot(
        functionPlot=LinearPlot.stackplot,
        dataIndex=[0],
        title="Leontief",
        abscisseIndex=Index.datetime,
        categoryIndex=Index.category,
        majorLocatorInterval=mdates.HourLocator(interval=1),
        majorFormatter=mdates.DateFormatter("%H:%M"),
        nrows=nrows,
        ncols=ncols,
        index=index_1,
        yLim=yLim,
    )
    s2 = Subplot(
        functionPlot=LinearPlot.linearplot,
        dataIndex=[6],
        title="Leontief",
        abscisseIndex=Index.datetime,
        categoryIndex=Index.category,
        majorLocatorInterval=mdates.HourLocator(interval=1),
        majorFormatter=mdates.DateFormatter("%H:%M"),
        nrows=nrows,
        ncols=ncols,
        index=index_1,
        yLim=yLim,
    )
    s3 = Subplot(
        functionPlot=LinearPlot.linearplot,
        dataIndex=[7],
        title="Leontief",
        abscisseIndex=Index.datetime,
        categoryIndex=Index.category,
        majorLocatorInterval=mdates.HourLocator(interval=1),
        majorFormatter=mdates.DateFormatter("%H:%M"),
        nrows=nrows,
        ncols=ncols,
        index=index_1,
        yLim=yLim,
    )
    s4 = Subplot(
        functionPlot=LinearPlot.stackplot,
        dataIndex=[1],
        title="Leontief",
        abscisseIndex=Index.datetime,
        categoryIndex=Index.source,
        majorLocatorInterval=mdates.HourLocator(interval=1),
        majorFormatter=mdates.DateFormatter("%H:%M"),
        nrows=nrows,
        ncols=ncols,
        index=index_2,
        yLim=yLim,
    )
    s5 = Subplot(
        functionPlot=LinearPlot.linearplot,
        dataIndex=[6],
        title="Leontief",
        abscisseIndex=Index.datetime,
        categoryIndex=Index.category,
        majorLocatorInterval=mdates.HourLocator(interval=1),
        majorFormatter=mdates.DateFormatter("%H:%M"),
        nrows=nrows,
        ncols=ncols,
        index=index_2,
        yLim=yLim,
    )
    s6 = Subplot(
        functionPlot=LinearPlot.linearplot,
        dataIndex=[7],
        title="Leontief",
        abscisseIndex=Index.datetime,
        categoryIndex=Index.category,
        majorLocatorInterval=mdates.HourLocator(interval=1),
        majorFormatter=mdates.DateFormatter("%H:%M"),
        nrows=nrows,
        ncols=ncols,
        index=index_2,
        yLim=yLim,
    )
    figure.addSubplots(subplots=[s1, s2, s3, s4, s5, s6])


def getMarkovSubplot(figure, nrows, ncols, index_1, index_2, yLim=[-3000, 12000]):
    s1 = Subplot(
        functionPlot=LinearPlot.stackplot,
        dataIndex=[2],
        title="Markov",
        abscisseIndex=Index.datetime,
        categoryIndex=Index.category,
        majorLocatorInterval=mdates.HourLocator(interval=1),
        majorFormatter=mdates.DateFormatter("%H:%M"),
        nrows=nrows,
        ncols=ncols,
        index=index_1,
        yLim=yLim,
    )
    s2 = Subplot(
        functionPlot=LinearPlot.linearplot,
        dataIndex=[6],
        title="Markov",
        abscisseIndex=Index.datetime,
        categoryIndex=Index.category,
        majorLocatorInterval=mdates.HourLocator(interval=1),
        majorFormatter=mdates.DateFormatter("%H:%M"),
        nrows=nrows,
        ncols=ncols,
        index=index_1,
        yLim=yLim,
    )
    s3 = Subplot(
        functionPlot=LinearPlot.linearplot,
        dataIndex=[7],
        title="Markov",
        abscisseIndex=Index.datetime,
        categoryIndex=Index.category,
        majorLocatorInterval=mdates.HourLocator(interval=1),
        majorFormatter=mdates.DateFormatter("%H:%M"),
        nrows=nrows,
        ncols=ncols,
        index=index_1,
        yLim=yLim,
    )
    s4 = Subplot(
        functionPlot=LinearPlot.stackplot,
        dataIndex=[3],
        title="Markov",
        abscisseIndex=Index.datetime,
        categoryIndex=Index.source,
        majorLocatorInterval=mdates.HourLocator(interval=1),
        majorFormatter=mdates.DateFormatter("%H:%M"),
        nrows=nrows,
        ncols=ncols,
        index=index_2,
        yLim=yLim,
    )
    s5 = Subplot(
        functionPlot=LinearPlot.linearplot,
        dataIndex=[6],
        title="Markov",
        abscisseIndex=Index.datetime,
        categoryIndex=Index.category,
        majorLocatorInterval=mdates.HourLocator(interval=1),
        majorFormatter=mdates.DateFormatter("%H:%M"),
        nrows=nrows,
        ncols=ncols,
        index=index_2,
        yLim=yLim,
    )
    s6 = Subplot(
        functionPlot=LinearPlot.linearplot,
        dataIndex=[7],
        title="Markov",
        abscisseIndex=Index.datetime,
        categoryIndex=Index.category,
        majorLocatorInterval=mdates.HourLocator(interval=1),
        majorFormatter=mdates.DateFormatter("%H:%M"),
        nrows=nrows,
        ncols=ncols,
        index=index_2,
        yLim=yLim,
    )
    figure.addSubplots(subplots=[s1, s2, s3, s4, s5, s6])


def getGenerataionSubplots(figure, nrows, ncols, index_1, index_2, yLim=[-3000, 12000]):
    s1 = Subplot(
        functionPlot=LinearPlot.stackplot,
        dataIndex=[4],
        title="Generation",
        abscisseIndex=Index.datetime,
        categoryIndex=Index.category,
        majorLocatorInterval=mdates.HourLocator(interval=1),
        majorFormatter=mdates.DateFormatter("%H:%M"),
        nrows=nrows,
        ncols=ncols,
        index=index_1,
        yLim=yLim,
    )
    s2 = Subplot(
        functionPlot=LinearPlot.linearplot,
        dataIndex=[6],
        title="Generation",
        abscisseIndex=Index.datetime,
        categoryIndex=Index.category,
        majorLocatorInterval=mdates.HourLocator(interval=1),
        majorFormatter=mdates.DateFormatter("%H:%M"),
        nrows=nrows,
        ncols=ncols,
        index=index_1,
        yLim=yLim,
    )
    s3 = Subplot(
        functionPlot=LinearPlot.linearplot,
        dataIndex=[7],
        title="Generation",
        abscisseIndex=Index.datetime,
        categoryIndex=Index.category,
        majorLocatorInterval=mdates.HourLocator(interval=1),
        majorFormatter=mdates.DateFormatter("%H:%M"),
        nrows=nrows,
        ncols=ncols,
        index=index_1,
        yLim=yLim,
    )
    s4 = Subplot(
        functionPlot=LinearPlot.stackplot,
        dataIndex=[5],
        title="Load",
        abscisseIndex=Index.datetime,
        categoryIndex=Index.country,
        majorLocatorInterval=mdates.HourLocator(interval=1),
        majorFormatter=mdates.DateFormatter("%H:%M"),
        nrows=nrows,
        ncols=ncols,
        index=index_2,
        yLim=yLim,
    )
    s5 = Subplot(
        functionPlot=LinearPlot.linearplot,
        dataIndex=[6],
        title="Load",
        abscisseIndex=Index.datetime,
        categoryIndex=Index.category,
        majorLocatorInterval=mdates.HourLocator(interval=1),
        majorFormatter=mdates.DateFormatter("%H:%M"),
        nrows=nrows,
        ncols=ncols,
        index=index_2,
        yLim=yLim,
    )
    s6 = Subplot(
        functionPlot=LinearPlot.linearplot,
        dataIndex=[7],
        title="Load",
        abscisseIndex=Index.datetime,
        categoryIndex=Index.category,
        majorLocatorInterval=mdates.HourLocator(interval=1),
        majorFormatter=mdates.DateFormatter("%H:%M"),
        nrows=nrows,
        ncols=ncols,
        index=index_2,
        yLim=yLim,
    )
    figure.addSubplots(subplots=[s1, s2, s3, s4, s5, s6])


# Balance
if False:
    # Linear figures
    figure = Figure(
        filename="./output/test/consumption_balance.png",
        data=[balance],
        title="Balance",
    )
    s1 = Subplot(
        functionPlot=LinearPlot.linearplot,
        abscisseIndex=Index.datetime,
        categoryIndex=Index.country,
        majorLocatorInterval=mdates.MonthLocator(interval=1),
        nrows=2,
        ncols=2,
        index=1,
    )
    s2 = Subplot(
        functionPlot=LinearPlot.stackplot,
        abscisseIndex=Index.datetime,
        categoryIndex=Index.country,
        majorLocatorInterval=mdates.MonthLocator(interval=1),
        nrows=2,
        ncols=2,
        index=2,
    )
    s3 = Subplot(
        functionPlot=LinearPlot.linearplot,
        abscisseIndex=Index.datetime,
        categoryIndex=Index.category,
        majorLocatorInterval=mdates.MonthLocator(interval=1),
        nrows=2,
        ncols=2,
        index=3,
    )
    s4 = Subplot(
        functionPlot=LinearPlot.stackplot,
        abscisseIndex=Index.datetime,
        categoryIndex=Index.category,
        majorLocatorInterval=mdates.MonthLocator(interval=1),
        nrows=2,
        ncols=2,
        index=4,
    )
    figure.addSubplots(subplots=[s1, s2, s3, s4])
    figure.plot(plotIfExist=True)


for c in listCountries:

    data = [
                PandasSelection.selectCountry(consumption1_Leontief.copy(), country=[c.iso]),
                PandasSelection.selectCountry(consumption2_Leontief.copy(), country=[c.iso]),
                PandasSelection.selectCountry(consumption1_Markov.copy(), country=[c.iso]),
                PandasSelection.selectCountry(consumption2_Markov.copy(), country=[c.iso]),
                PandasSelection.selectCountry(generation.copy(), country=[c.iso]),
                PandasSelection.selectCountry(load.copy(), country=[c.iso]),
                PandasSelection.selectCountry(summedGeneration.copy(), country=[c.iso]),
                PandasSelection.selectCountry(summedLoad.copy(), country=[c.iso]),
            ]

    if True:
        figure = Figure(
            filename=f"./output/test/consumption_{c.iso}.png",
            data=data,
            title=f"Consumption of {c.name} [{c.iso}]",
            size=(72, 32),
        )
        getLeontiefSubplot    (figure, nrows=2, ncols=3, index_1=1, index_2=4, yLim=yLimTab[c.iso])
        getMarkovSubplot      (figure, nrows=2, ncols=3, index_1=2, index_2=5, yLim=yLimTab[c.iso])
        getGenerataionSubplots(figure, nrows=2, ncols=3, index_1=3, index_2=6, yLim=yLimTab[c.iso])
        figure.plot(plotIfExist=True)

    # Leontief
    if True:
        figure = Figure(
            filename=f"./output/test/consumption_{c.iso}_1.png",
            data=data,
            title=f"Consumption of {c.name} [{c.iso}]",
            size=(24, 32),
        )
        getLeontiefSubplot(figure, nrows=2, ncols=1, index_1=1, index_2=2, yLim=yLimTab[c.iso])
        figure.plot(plotIfExist=True)

    # Markov
    if True:
        figure = Figure(
            filename=f"./output/test/consumption_{c.iso}_2.png",
            data=data,
            title=f"Consumption of {c.name} [{c.iso}]",
            size=(24, 32),
        )
        getMarkovSubplot(figure, nrows=2, ncols=1, index_1=1, index_2=2, yLim=yLimTab[c.iso])
        figure.plot(plotIfExist=True)

    # Generation
    if True:
        figure = Figure(
            filename=f"./output/test/consumption_{c.iso}_3.png",
            data=data,
            title=f"Consumption of {c.name} [{c.iso}]",
            size=(24, 32),
        )
        getGenerataionSubplots(figure, nrows=2, ncols=1, index_1=1, index_2=2, yLim=yLimTab[c.iso])
        figure.plot(plotIfExist=True)

print("end")