# Sources:
# https://pymrio.readthedocs.io/en/latest/
# https://pymrio.readthedocs.io/en/latest/notebooks/working_with_exiobase.html
# https://pymrio.readthedocs.io/en/latest/notebooks/buildflowmatrix.html

import pymrio
import pandas as pd

# %% Import Exiobase

exio3_path = '/IOT_2021_pxp'

exio3 = pymrio.parse_exiobase3(path=exio3_path)

# Production matrix F (see: https://pymrio.readthedocs.io/en/latest/math.html)
production_matrix_F = exio3.satellite.F 

# %% Determine satellites

# Get blue water satellites
water_satellites_blue = production_matrix_F.index[
    production_matrix_F.index.str.contains('Water Consumption Blue')
    ]

# Get green water satellites
water_satellites_green = production_matrix_F.index[
    production_matrix_F.index.str.contains('Water Consumption Green')
    ]

# %% Calculate the green and blue water footprint of Germany

wf_results_consumption_based = {}

# For all green and blue water satellites

for i, satellite in enumerate(
        list(water_satellites_green)
        + list(water_satellites_blue)
        ):
    
    print(
        '\nProcessing satellite {} of {}:\n{}'.format(
            i+1,
            len(
                list(water_satellites_green)
                + list(water_satellites_blue)
                ),
            satellite
            )
        )

    et1_diag = exio3.satellite.diag_stressor(
        (satellite),
        name = satellite
        )

    # Connect back to the system 
    exio3.et1_diag = et1_diag

    # Calulate all the stressor accounts 
    exio3.calc_all()

    # The consumption-based grouped footprint of a region
    wf_results_consumption_based[satellite] = exio3.et1_diag.D_cba.sum(
        axis=0
        ).DE

# %% Turn into DF and save

consumption_based_footprint_DE = pd.DataFrame.from_dict(
    wf_results_consumption_based
    )

consumption_based_footprint_DE.to_excel(
    'consumption_based_footprint_DE.xlsx'
    )