# ---------------------------------------------------------------------------------------------------------------------
"""
Author: Raphael Andreas Elbing
Last Modified: 25/08/2022
License: This work is licensed under a Creative Commons Attribution-ShareAlike 4.0 International (CC BY-SA 4.0) License.
"""
# --------------------------------------------------------------------------------------------------------------------
"""This file contains the used MaTrace model by Godoy León et al. (2020). There are four functions. Three of them
evaluate the single clusters of the model use (calculate_use_stocks_flows), end-of-life (calculate_recycling_stocks_flows,
and production (calculate production stocks and flows). The last function evaluate cohort calls the functions before
and sets the model together. In this context cohort refers to the fact, that only one inflow is considered."""

import pandas as pd
import numpy as np
import scipy.stats


def calculate_use_stocks_flows(inflow, use_lifetime_pd, hoarding_pd, use_cohort_matrices, hoarding_cohort_matrices,
                               iteration, n_years, defined_distributions_pd):
    """
    This function calculates the stocks and flows of the use phase of the MaTrace model for one year. Cohort
    matrices are passed to this function, extended and returned

    :param inflow: pandas of inflow over product categories. The index are the product categories.
    :param use_lifetime_pd: pandas of survival curve data per product (normally Weibull shape and scale
    :param hoarding_pd: pandas indicating the hoarding time and the hoarding rate
    :param use_cohort_matrices: dictionary containing one cohort matrix per product category.
    :param hoarding_cohort_matrices: dictionary containing one cohort matrix per product category.
    :param iteration: number of the current iteration (starting with 0)
    :param n_years: number of the considered years
    :return: output_pd (contains magnitude of stocks and flows), use_cohort_matrices (see above),
             hoarding_cohort_matrices (see above)
    """
    # Create empty output data frame
    output_pd = pd.DataFrame({'Products': use_lifetime_pd.index})
    output_pd = output_pd.set_index('Products')

    ############################### U.A Use ########################################################################

    # Expend cohort matrix of use stock, calculate stock and outflow
    use_stock = []
    use_outflow = []

    # To test
    total_privious_use_stock = 0
    total_privious_hoarding_stock = 0

    for product in inflow.index:

        # get inflow
        use_inflow = inflow.loc[product]

        # get matrix for product
        use_cohort_matrix = use_cohort_matrices[product]

        # get scale and shape
        distribution = use_lifetime_pd.loc[product, 'distribution']
        location = use_lifetime_pd.loc[product, 'location']
        shape = use_lifetime_pd.loc[product, 'shape']
        scale = use_lifetime_pd.loc[product, 'scale']

        # extend cohort matrix
        if distribution == 'weibull':
            use_cohort_matrix[iteration:, iteration] = use_inflow * scipy.stats.weibull_min.sf(
                x=range(n_years - iteration),
                c=shape, loc=location, scale=scale)
        if distribution == 'normal':
            use_cohort_matrix[iteration:, iteration] = use_inflow * scipy.stats.norm.sf(x=range(n_years - iteration),
                                                                                        loc=location, scale=scale)
        if distribution == 'gamma':
            use_cohort_matrix[iteration:, iteration] = use_inflow * scipy.stats.gamma.sf(x=range(n_years - iteration),
                                                                                         a=shape, loc=location,
                                                                                         scale=scale)
        if distribution == 'gompertz':
            use_cohort_matrix[iteration:, iteration] = use_inflow * scipy.stats.gompertz.sf(
                x=range(n_years - iteration),
                c=shape, loc=location,
                scale=scale)
        if distribution == 'lognormal':
            use_cohort_matrix[iteration:, iteration] = use_inflow * scipy.stats.lognorm.sf(x=range(n_years - iteration),
                                                                                           s=shape, loc=location,
                                                                                           scale=scale)

        if distribution.split('_', 2)[0] == 'defined':
            defined_distribution_name = distribution.split('_', 2)[2]

            use_cohort_matrix[iteration:, iteration] = use_inflow * defined_distributions_pd[
                                                                        defined_distribution_name].to_numpy()[
                                                                    :(n_years - iteration)]

        # write it back to dictionary
        use_cohort_matrices[product] = use_cohort_matrix

        # Calculate stock of current year
        current_use_stock = use_cohort_matrix[iteration, :].sum()
        use_stock.append(current_use_stock)

        # Calculate outflow
        ## calculate previous stock
        if iteration == 0:
            previous_use_stock = 0
        else:
            previous_use_stock = use_cohort_matrix[iteration - 1, :].sum()

        # to test
        total_privious_use_stock += previous_use_stock

        ## calculate netflow
        use_netflow = current_use_stock - previous_use_stock

        use_outflow.append(use_inflow - use_netflow)
        ### End for loop

    # Add things to table
    output_pd['U.A use stock'] = use_stock
    output_pd['U.2 use outflow'] = use_outflow

    #################### U.B Hoarding and EoL products #############################################################

    # Calculate U.3 hoarding inflow and U4 no hoarding flow
    hoarding_stock = []
    hoarding_outflow = []

    output_pd['U.3 hoarding inflow'] = output_pd['U.2 use outflow'] * hoarding_pd['hoarding rate']
    output_pd['U.4 no hoarding flow'] = output_pd['U.2 use outflow'] * (1 - hoarding_pd['hoarding rate'])

    # Extend cohort matrix of hoarding stock, calculate stock and outflow

    for product in inflow.index:
        # get inflow
        hoarding_inflow = output_pd.loc[product, 'U.3 hoarding inflow']
        hoarding_time = hoarding_pd.loc[product, 'hoarding time']

        # Not all products have a share going to hoarding, but they might have an outflow in special cases
        # Therefore, the code needs to be executed for all products.

        # get matrix for product
        hoarding_cohort_matrix = hoarding_cohort_matrices[product]

        if hoarding_inflow != 0:
            length_entry = n_years - iteration

            if hoarding_time <= length_entry:
                hoarding_vec = [hoarding_inflow for i in range(hoarding_time)]

                if hoarding_time < length_entry:
                    hoarding_vec = hoarding_vec + [0 for i in range(length_entry - hoarding_time)]

            # In case that the hoarding length exceeds the time model covers
            else:
                hoarding_vec = [hoarding_inflow for i in range(length_entry)]

            # add to cohort matrix
            hoarding_cohort_matrix[iteration:, iteration] = hoarding_vec
            # Write the matrix back
            hoarding_cohort_matrices[product] = hoarding_cohort_matrix

        # Calculate stock of current year
        current_hoarding_stock = hoarding_cohort_matrix[iteration, :].sum()
        hoarding_stock.append(current_hoarding_stock)

        # Calculate outflow
        ## calculate previous stock
        if iteration == 0:
            previous_hoarding_stock = 0
        else:
            previous_hoarding_stock = hoarding_cohort_matrix[iteration - 1, :].sum()

        # to test
        total_privious_hoarding_stock += previous_hoarding_stock

        ## calculate netflow
        hoarding_netflow = current_hoarding_stock - previous_hoarding_stock
        hoarding_outflow.append(hoarding_inflow - hoarding_netflow)

    output_pd['U.B hoarding stock'] = hoarding_stock
    output_pd['U.5 hoarding outflow'] = hoarding_outflow

    output_pd['U.6 eol products'] = output_pd['U.5 hoarding outflow'] + output_pd['U.4 no hoarding flow']

    return output_pd, use_cohort_matrices, hoarding_cohort_matrices


# %%
########################################## Recycling ################################################################

def calculate_recycling_stocks_flows(output_pd, pretreatment_pd, recycling_B_pd):
    """
    This function calculates the stocks and flows of the end-of-life/recycling phase of the MaTrace model for one year.

    :param output_pd: pandas entailing the stocks and flows over product categories of the use phase
    :param pretreatment_pd: pandas entailing transfer coefficients of pretreatment
    :param recycling_B_pd: pandas entailing allocation of product categories to recycling processes and efficiency of the
           processes
    :return: output_pd: pandas, same as input but extended by stocks and flows of end-of-life phase.
    """

    # Getting the inflow from the previous cluster
    eol_inflow = output_pd['U.6 eol products']

    # Collection node
    output_pd['E.2 exported eol products'] = eol_inflow * pretreatment_pd['fraction export eol products']
    output_pd['E.1 to waste treatment'] = eol_inflow * (1 - pretreatment_pd['fraction export eol products'])

    # Wast to pretreatment
    output_pd['E.3 to pretreatment'] = output_pd['E.1 to waste treatment'] * pretreatment_pd[
        'collection to recycling rate']
    output_pd['E.4 E.5 non-selective collection'] = output_pd['E.1 to waste treatment'] * (
                1 - pretreatment_pd['collection to recycling rate'])

    # Pretreatment node
    output_pd['E.6 to recycling'] = output_pd['E.3 to pretreatment'] * pretreatment_pd['pre-treatment efficiency']
    output_pd['E.7 pretreatment waste'] = output_pd['E.3 to pretreatment'] * (
                1 - pretreatment_pd['pre-treatment efficiency'])

    # Recycling
    recycling_split_pd = recycling_B_pd.loc[pretreatment_pd.index, :]
    recycling_efficiency_pd = recycling_B_pd.loc[
        'Efficiency', ['Chemical', 'Zn']]  # One efficiency per process, not per product class

    ## Downcycling
    output_pd['E.12 downcycling'] = output_pd['E.6 to recycling'] * recycling_split_pd['Downcycling']

    ## Zn
    output_pd['recycled w-co powder'] = output_pd['E.6 to recycling'] * recycling_split_pd['Zn'] * \
                                        recycling_efficiency_pd['Zn']

    ## Chemical
    output_pd['co metal compound'] = output_pd['E.6 to recycling'] * recycling_split_pd['Chemical'] * \
                                     recycling_efficiency_pd['Chemical']

    output_pd['E.8 recycling waste'] = output_pd['E.6 to recycling'] - output_pd['co metal compound'] - output_pd[
        'recycled w-co powder'] - output_pd['E.12 downcycling']

    return output_pd


# %%
################################################# Production ########################################################

def calculate_production_stocks_flows(output_pd, allocation_D, production_pd):
    """
    This function calculates the stocks and flows of the production phase of the MaTrace model for one year.
    :param output_pd: pandas entailing the stocks and flows of the previous clusters
    :param allocation_D: pandas with allocation of material to product categories and export and production rate.
    :param production_pd: pandas entailing the production parameters and transfer coefficients.
    :return: output_pd, pandas including all stocks and flows for the currently treated year.
    """

    inflow_pd = output_pd[['co metal compound', 'recycled w-co powder']]

    # Creation of efficiency matrix N
    efficiency_N = production_pd['processing scrap recovery'] * (1 - production_pd['processing yield']) + \
                   production_pd['manufacturing scrap recovery'] * (1 - production_pd['manufacturing yield']) * \
                   production_pd['processing yield']

    # Calculate total inflow of co metal compund and w-co powder
    inflow_co_metal_compound = inflow_pd['co metal compound'].sum()
    inflow_w_co_powder = inflow_pd['recycled w-co powder'].sum()

    # Calculate export flow per product, since output table needs this format
    inflow_co_metal_compound_export = inflow_pd['co metal compound'] * allocation_D.loc[
        'export rate', 'Co metal or compound']
    inflow_w_co_powder_export = inflow_pd['recycled w-co powder'] * allocation_D.loc['export rate', 'W-Co powder']

    # Drop export row since the rest are the products.
    allocation_D = allocation_D.drop(['export rate', 'to production rate'])

    inflow_co_metal_compound_production = inflow_co_metal_compound - inflow_co_metal_compound_export.sum()
    inflow_w_co_powder_production = inflow_w_co_powder - inflow_w_co_powder_export.sum()

    # print(inflow_co_metal_compound + inflow_w_co_powder == inflow_w_co_powder_production +inflow_co_metal_compound_production +inflow_w_co_powder_export.sum() +inflow_co_metal_compound_export.sum())

    # Calculate export
    output_pd['E.11 exported recycled materials'] = inflow_co_metal_compound_export + inflow_w_co_powder_export

    # Calculation of material for recycling vector m
    recycled_material_m = inflow_co_metal_compound_production * allocation_D[
        'Co metal or compound'] + inflow_w_co_powder_production * allocation_D['W-Co powder']

    # Total product outflow

    output_pd['P.1 total recycled products'] = production_pd['processing yield'] * production_pd[
        'manufacturing yield'] * (1 / (1 - efficiency_N)) * recycled_material_m

    # Export and inflow
    output_pd['P.8 export recycled products'] = production_pd['export of products'] * output_pd[
        'P.1 total recycled products']
    output_pd['U.1 product inflow'] = output_pd['P.1 total recycled products'] - output_pd[
        'P.8 export recycled products']

    # processing
    ## Calculation of processing waste
    output_pd['P.7 processing waste'] = (1 - production_pd['processing scrap recovery']) * (
                1 - production_pd['processing yield']) * (1 / (1 - efficiency_N)) * recycled_material_m

    ## Calcolation of downcycled scrap of production
    output_pd['P.5p downcycled scrap'] = output_pd['P.7 processing waste'] * production_pd[
        'processing downcycled scrap']

    ## Calculation disposed material of production
    output_pd['P.4p disposed scrap'] = output_pd['P.7 processing waste'] - output_pd['P.5p downcycled scrap']

    # manufacturing
    ## calculation of manufacturing waste
    output_pd['P.2 manufacturing waste'] = (1 - production_pd['manufacturing scrap recovery']) * \
                                           (1 - production_pd['manufacturing yield']) * production_pd[
                                               'processing yield'] * (1 / (1 - efficiency_N)) * recycled_material_m

    ## Calculation of downcycled scrap of manufacturing
    output_pd['P.5m downcycled scrap'] = output_pd['P.2 manufacturing waste'] * production_pd[
        'manufacturing downcycled scrap']

    ## Calculation of disposed material of manufacturing
    output_pd['P.4m disposed scrap'] = output_pd['P.2 manufacturing waste'] - output_pd['P.5m downcycled scrap']

    # Sum of P.4 and P.5
    output_pd['P.4 disposed scrap'] = output_pd['P.4m disposed scrap'] + output_pd['P.4p disposed scrap']
    output_pd['P.5 downcycled scrap'] = output_pd['P.5m downcycled scrap'] + output_pd['P.5p downcycled scrap']

    total_inflow = inflow_pd['co metal compound'].sum() + inflow_pd['recycled w-co powder'].sum()

    return output_pd


# %%
def evaluate_cohort(initial_inflow, use_lifetime_pd, hoarding_pd, pretreatment_pd, recycling_B_pd, allocation_D,
                    production_pd, n_years, start_year, defined_distributions_pd, print_state=True):
    """
    This function executes the whole MaTrace model by calling the function calculate_use_stocks_flows,
    calculate_recycling_stocks_flows, and calculate_production_stocks_flows.
    :param initial_inflow: pandas of inflow over product categories. The index are the product categories.
    :param use_lifetime_pd: pandas of survival curve data per product (normally Weibull shape and scale
    :param hoarding_pd: pandas indicating the hoarding time and the hoarding rate
    :param pretreatment_pd: pandas entailing transfer coefficients of pretreatment
    :param recycling_B_pd: pandas entailing allocation of product categories to recycling processes and efficiency of the
           processes
    :param allocation_D: pandas with allocation of material to product categories and export and production rate.
    :param production_pd: pandas entailing the production parameters and transfer coefficients.
    :param n_years: integer number of considered years/iterations
    :param start_year: integer start year
    :param print_state: boolean, if True the number of the current iteration will be printed
    :return: data_dic, ditctionary with the years (as string) as key. Each entry holds the respective output_pd dataframe
             graph_data_pd, pandas holding all stocks and accumulated outflows in order to easily create a stacked
             area chart.
    """

    # Data collectors
    data_dic = {}

    # defining the collumns of graph_data_pd
    categories_graph = ['Recycling losses', 'Pre-treatment losses', 'Production losses',
                        'Non-selective collection',
                        'Downcycled', 'Exported', 'Hoarded', 'Superalloys', 'Other metallic uses', 'Magnets',
                        'Hard metals',
                        'Dissibative uses', 'Catalysts', 'Mobility batteries', 'Portable batteries']

    categories_graph.reverse()

    graph_data_dict = dict.fromkeys(categories_graph)
    for category in categories_graph:
        graph_data_dict[category] = [0]

    # Write initial_inflow in inflow
    inflow = initial_inflow['share']

    # create empty cohort matrices
    products_list = initial_inflow.index
    use_cohort_matrices = {}
    hibernating_cohort_matrices = {}

    for product in products_list:
        use_cohort_matrices[product] = np.zeros((n_years, n_years))
        hibernating_cohort_matrices[product] = np.zeros((n_years, n_years))

    # Loop over considered years
    for year in range(n_years):

        if print_state:
            print('Year {} of {}'.format(year + 1, n_years))
        # evaluate use stock
        output_pd, use_cohort_matrices, hibernating_cohort_matrices = calculate_use_stocks_flows(inflow,
                                                                                                 use_lifetime_pd,
                                                                                                 hoarding_pd,
                                                                                                 use_cohort_matrices,
                                                                                                 hibernating_cohort_matrices,
                                                                                                 year, n_years,
                                                                                                 defined_distributions_pd)

        # evaluate recycling
        output_pd = calculate_recycling_stocks_flows(output_pd, pretreatment_pd, recycling_B_pd)

        # evaluate production
        output_pd = calculate_production_stocks_flows(output_pd, allocation_D, production_pd)

        # change inflow to new inflow
        inflow = output_pd['U.1 product inflow']

        # update variables
        data_dic[str(year)] = output_pd

        # Collect data for graph
        graph_data_dict['Recycling losses'].append(
            graph_data_dict['Recycling losses'][-1] + output_pd['E.8 recycling waste'].sum())
        graph_data_dict['Pre-treatment losses'].append(
            graph_data_dict['Pre-treatment losses'][-1] + output_pd['E.7 pretreatment waste'].sum())

        graph_data_dict['Production losses'].append(
            graph_data_dict['Production losses'][-1] + output_pd['P.4 disposed scrap'].sum())
        graph_data_dict['Non-selective collection'].append(
            graph_data_dict['Non-selective collection'][-1] + output_pd['E.4 E.5 non-selective collection'].sum())
        graph_data_dict['Downcycled'].append(
            graph_data_dict['Downcycled'][-1] + output_pd['E.12 downcycling'].sum() + output_pd[
                'P.5 downcycled scrap'].sum())
        graph_data_dict['Exported'].append(
            graph_data_dict['Exported'][-1] + output_pd['E.11 exported recycled materials'].sum() + output_pd[
                'E.2 exported eol products'].sum() + output_pd['P.8 export recycled products'].sum())
        graph_data_dict['Hoarded'].append(output_pd['U.B hoarding stock'].sum())

        graph_data_dict['Superalloys'].append(output_pd.loc['superalloys', 'U.A use stock'] + inflow.loc['superalloys'])
        graph_data_dict['Other metallic uses'].append(
            output_pd.loc['other metallic uses', 'U.A use stock'] + inflow.loc['other metallic uses'])
        graph_data_dict['Magnets'].append(output_pd.loc['magnets', 'U.A use stock'] + inflow.loc['magnets'])
        graph_data_dict['Hard metals'].append(output_pd.loc['hard metals', 'U.A use stock'] + inflow.loc['hard metals'])
        graph_data_dict['Dissibative uses'].append(
            output_pd.loc['dissipative uses', 'U.A use stock'] + inflow.loc['dissipative uses'])
        graph_data_dict['Catalysts'].append(
            output_pd.loc['hydroprocessing catalysts coke', 'U.A use stock'] + output_pd.loc[
                'hydroprocessing catalysts poisoning', 'U.A use stock'] +
            output_pd.loc['hydroformylation catalysts', 'U.A use stock'] + output_pd.loc[
                'pet precursors catalysts', 'U.A use stock'] +
            inflow.loc['hydroprocessing catalysts coke'] + inflow.loc['hydroprocessing catalysts poisoning'] +
            inflow.loc['hydroformylation catalysts'] + inflow.loc['pet precursors catalysts'])
        graph_data_dict['Mobility batteries'].append(
            output_pd.loc['mobility batteries', 'U.A use stock'] + inflow.loc['mobility batteries'])
        graph_data_dict['Portable batteries'].append(
            output_pd.loc['portable batteries', 'U.A use stock'] + inflow.loc['portable batteries'])

    ############### end loop over years ##################################

    # Treat table of dictionary. Remove zeros from first entry.
    for category in categories_graph:
        graph_data_dict[category] = graph_data_dict[category][1:]

    # add year to dataframe
    graph_data_dict['Year'] = [i + start_year for i in range(n_years)]

    # transform dictionary to table
    graph_data_pd = pd.DataFrame(graph_data_dict)
    graph_data_pd = graph_data_pd.set_index('Year')

    # Return tables

    return data_dic, graph_data_pd
