# ---------------------------------------------------------------------------------------------------------------------
"""
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 Reuse-MaTrace model which is combined out of the Reuse model of Thiebaud et al. (2017) and
   the MaTrace model of Godoy León et al. (2020). The code consists of one large function which uses functions of
   the files base_matrace_model.py and base_reuse_model.py."""

from base_matrace_model import calculate_use_stocks_flows, calculate_recycling_stocks_flows, \
    calculate_production_stocks_flows
from base_reuse_model import reuse_model_single_cohort

import numpy as np
import pandas as pd


#
def ordinal(n):
    """
    Function for converting numbers (1 to '1st', 2 to '2nd') source:
    https://stackoverflow.com/questions/9647202/ordinal-numbers-replacement
    :param n: integer
    :return: string converting number as described.
    """

    return "%d%s" % (
        n, "tsnrhtdd"[(n // 10 % 10 != 1) * (n % 10 < 4) * n % 10::4])


def evaluate_cohort_combined_model(data_dic, n_years, start_year, defined_distributions_pd,
                                   print_state=True, separate_reuse_graph=True, considered_use_cycles=3):
    """Executes and evaluates the extended MaTrace model. The created data is returned.

    :param data_dic: Dictionary containing pandas Dataframes with needed data (see next two code chunks)
    :param n_years: Integer of years the model will forcast the wherabouts of the material
    :param start_year: Integer start year for data frame indexing
    :param print_state: Boolean if true calculation progress is printed in console
    :param separate_reuse_graph: Boolean if true are splitted up in 1st, 2nd, and 3rd in return graph_data_pd
    :param considered_use_cycles: integer indicating the number of considered use cycles
    :return: data_collector_dic; entails all stocks and flows of MaTrace model,
             data_collector_reuse_dic; entails all stocks and flows of reuse model,
             graph_data_pd; entails all stocks and cummulated flows leaving the system
    """

    ################# Initialization of input data, data structures to run the model and to be returned ##############

    # Input data
    # Variables to be fed into original matrace model
    initial_inflow = data_dic['MaTrace_initial_inflow']
    use_lifetime_pd = data_dic['MaTrace_in_use_stock']
    hoarding_pd = data_dic['MaTrace_hibernating_stock']
    pretreatment_pd = data_dic['MaTrace_end_of_life']
    recycling_B_pd = data_dic['MaTrace_B_recycling']
    allocation_D = data_dic['MaTrace_D_secondary_material']
    production_pd = data_dic['MaTrace_production']

    product_for_reuse = initial_inflow.index[0]

    product_categories = initial_inflow.index.to_list()
    product_categories = list(reversed(product_categories))[:-1]

    # Calculate weibull parameter for use lifetime
    # The following code is old. the input data was changed. it should still work.

    # use_lifetime_pd['Weibull scale'] = [weibull_scale_median(mean, scale) for mean, scale in
    #                                     zip(use_lifetime_pd['Lifetime (years)'], use_lifetime_pd['Weibull shape'])]

    # Variables to be fed into reuse model
    data_split_reuse = data_dic['Reuse_inflow_split']['split']

    # Population of reuse dictionaries
    reuse_use_dic = {}
    reuse_storage_dic = {}

    for i in range(considered_use_cycles):
        use_cycle = i + 1

        # use phase data
        reuse_use_dic['Reuse_service_time_{}'.format(
                use_cycle)] = data_dic['Reuse_service_time_{}'.format(use_cycle)]
        # storage phase
        reuse_storage_dic['Reuse_storage_time_{}'.format(
                use_cycle)] = data_dic['Reuse_storage_time_{}'.format(use_cycle)]

    # Ouptut data
    # Data collectors
    data_collector_dic = {}
    data_collector_reuse_dic = {}

    # data graph data frame
    categories_graph = ['Recycling losses', 'Pre-treatment losses', 'Production losses',
                        'Non-selective collection', 'Downcycled', 'Exported', 'Hoarded'] + [i.capitalize() for i in
                                                                                            product_categories]

    if separate_reuse_graph:
        categories_graph = categories_graph + ['{} {} use'.format(product_for_reuse.capitalize(), ordinal(
                considered_use_cycles - i)) for i in range(considered_use_cycles)]

    else:
        categories_graph = categories_graph + [product_for_reuse.capitalize()]

    categories_graph.reverse()

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

    # Required data structure to run the model
    # create empty cohort matrices MaTrace model
    products_list = initial_inflow.index
    use_cohort_matrices = {}
    hibernating_cohort_matrices = {}

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

    # create empty cohort matrices reuse model
    reuse_cohort_matrices = {}

    for i in range(considered_use_cycles):
        use_cycle = i + 1

        for storage in ['service_time_', 'storage_time_']:

            actual_matrices = {}
            for product in data_dic['Reuse_inflow_split'].index:
                actual_matrices[product] = np.zeros((n_years, n_years))

            # putting dictionary into a dictionary
            reuse_cohort_matrices[storage + str(use_cycle)] = actual_matrices

    # Write initial_inflow in inflow - inflow will be overwritten each year in the following year.
    inflow = initial_inflow['share']

    ###################################### Execution of model #########################################################
    # Loop over considered years

    for year in range(n_years):

        if print_state:
            print('Year {} of {}'.format(year + 1, n_years))

        ### Use Phase ###
        # split portable batteries to into categories considered in reuse model
        reuse_inflow = pd.DataFrame({})
        reuse_inflow['Inflow'] = inflow.loc[product_for_reuse] * \
                                 data_split_reuse

        # Rest to original matrace model
        normal_inflow = inflow.iloc[1:]

        # evaluate use stock original MaTrace model
        output_pd, use_cohort_matrices, hibernating_cohort_matrices = calculate_use_stocks_flows(
                normal_inflow, use_lifetime_pd.iloc[1:], hoarding_pd, use_cohort_matrices, hibernating_cohort_matrices,
                year,
                n_years, defined_distributions_pd)

        # evaluate reuse model
        output_reuse_pd, reuse_cohort_matrices = reuse_model_single_cohort(reuse_inflow, reuse_use_dic,
                                                                           reuse_storage_dic,
                                                                           reuse_cohort_matrices, year, n_years,
                                                                           considered_use_cycles,
                                                                           defined_distributions_pd)

        # Combining the two different use phases to one dataframe
        # total stock
        total_use_stock = output_reuse_pd['total_use_stock'].sum()
        total_hoarding_stock = output_reuse_pd['total_hoarding_stock'].sum()
        to_disposal_flow = output_reuse_pd['to_disposal_flow'].sum()

        reuse_summary_pd = pd.DataFrame(
                np.array([[total_use_stock, None, None, None,
                           total_hoarding_stock, None, to_disposal_flow]]),
                columns=output_pd.columns)

        # setting index
        reuse_summary_pd['Products'] = product_for_reuse
        reuse_summary_pd = reuse_summary_pd.set_index('Products')

        output_pd = pd.concat([reuse_summary_pd, output_pd])

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

        # Production Phase
        # 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']

        ############################### Updating data to be collected ##############################################

        # update variables
        data_collector_dic[str(year)] = output_pd
        data_collector_reuse_dic[str(year)] = output_reuse_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())

        for product_category in product_categories:
            graph_data_dict[product_category.capitalize()].append(
                output_pd.loc[product_category, 'U.A use stock'] + inflow.loc[product_category])

        if separate_reuse_graph:

            for i in range(considered_use_cycles):
                lifecycle = i + 1

                graph_data_dict['{} {} use'.format(product_for_reuse.capitalize(), ordinal(lifecycle))].append(
                        output_reuse_pd['use_stock_{}'.format(lifecycle)].sum())

        else:
            graph_data_dict[product_for_reuse.capitalize()].append(
                    output_pd.loc[product_for_reuse, 'U.A use stock'] + inflow.loc[product_for_reuse])

    ############### 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_collector_dic, data_collector_reuse_dic, graph_data_pd
