# ---------------------------------------------------------------------------------------------------------------------
"""
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 to a large extend (except EOL treatment) the Reuse model of Thiebaud et al. (2017).
   It consists of two function, one that extends cohort matrices based and one that does the rest.
   A flexible number of use cycles can be considered, given that the required data is passed to the model as well."""

import pandas as pd
import scipy.stats


# supporting functions
def extend_cohort_matrix(cohort_matrix, inflow, distribution, location, scale, shape, iteration, n_years,
                         defined_distributions_pd):
    """
    This function extends cohort matrices based on the given inputs. Type of extetntion is necessary due to the iterative
    nature of the MaTrace model
    :param cohort_matrix: n_years x n_years numpy matrix
    :param inflow: float of material entering the stock in the given year
    :param weibull_scale: float Weibull scale
    :param weibull_shape: float Weibull shape
    :param iteration: integer indicating the current iteration or the current year (starting year = 0)
    :param n_years: number of considered years
    :return: cohort matrix (numpy), current stock and current outflow
    """
    # Fill cohort matrix
    if distribution == 'weibull':
        cohort_matrix[iteration:, iteration] = inflow * scipy.stats.weibull_min.sf(x=range(n_years - iteration),
                                                                                   c=shape, loc=location,
                                                                                   scale=scale)
    if distribution == 'normal':
        cohort_matrix[iteration:, iteration] = inflow * scipy.stats.norm.sf(x=range(n_years - iteration),
                                                                            loc=location, scale=scale)
    if distribution == 'gamma':
        cohort_matrix[iteration:, iteration] = inflow * scipy.stats.gamma.sf(x=range(n_years - iteration),
                                                                             a=shape, loc=location, scale=scale)
    if distribution == 'gompertz':
        cohort_matrix[iteration:, iteration] = inflow * scipy.stats.gompertz.sf(x=range(n_years - iteration),
                                                                                c=shape, loc=location,
                                                                                scale=scale)
    if distribution == 'lognormal':
        cohort_matrix[iteration:, iteration] = 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]

        cohort_matrix[iteration:, iteration] = inflow * defined_distributions_pd[defined_distribution_name].to_numpy()[
                                                        :(n_years - iteration)]

        # Calculate stock of current year
    stock = cohort_matrix[iteration, :].sum()

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

    ## calculate netflow
    netflow = stock - previous_stock
    ## calculate outflow
    outflow = inflow - netflow

    return cohort_matrix, stock, outflow


# Function returning output dataframe,

def reuse_model_single_cohort(inflow_pd, data_use_dic, data_storage_dic, cohort_matrices_dic, iteration, n_years,
                              considered_use_cycles, defined_distributions_pd):
    """
    This function calculates stocks and flows and extends all cohort matrices of the reuse model for one year.
    :param inflow_pd: pandas entailing the inflows over products (as index)
    :param data_use_dic: dictionary entailing pandas with all parameters as transfer coefficients and lifetimes over products
    :param data_storage_dic: dictionary entailing pandas with hoarding parameters over products
    :param cohort_matrices_dic: dictionary entailing all use and storage cohort matrics
    :param iteration: integer indicating the current iteration or the current year (starting year = 0)
    :param n_years: integer indicating the current iteration or the current year (starting year = 0)
    :param considered_use_cycles: integer number of considered use cycles
    :return: output_pd pandas with all stocks and flows, cohort_matrices_dic dictionary with all cohort matrices
    """

    # Create empty output data frame
    products = inflow_pd.index
    output_pd = pd.DataFrame({'Products': products})
    output_pd = output_pd.set_index('Products')

    # Populating dictionaries, later the dictionaries will have lists. Each entry refers to one product
    # flow dictionaries
    flow_use_to_next_use_dic = {}
    flow_use_to_storage_dic = {}
    flow_use_to_disposal_dic = {}
    flow_storage_to_next_use_dic = {}
    flow_storage_to_disposal_dic = {}

    # stock dictionaries
    stock_use_dic = {}
    stock_storage_dic = {}

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

        flow_use_to_storage_dic[str(use_cycle)] = []
        flow_use_to_disposal_dic[str(use_cycle)] = []
        flow_storage_to_disposal_dic[str(use_cycle)] = []

        stock_use_dic[str(use_cycle)] = []
        stock_storage_dic[str(use_cycle)] = []

        if use_cycle != considered_use_cycles:
            # last stage does not have flow to next use phase
            flow_use_to_next_use_dic[str(use_cycle)] = []
            flow_storage_to_next_use_dic[str(use_cycle)] = []

    # Iterate over products
    for product in products:

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

            if use_cycle == 1:
                inflow = inflow_pd.loc[product, 'Inflow']

            # Use pahse ############################
            cohort_matrices_dic['service_time_' + str(use_cycle)][product], use_stock, \
            use_outflow = extend_cohort_matrix(cohort_matrices_dic['service_time_' + str(use_cycle)][product],
                                               inflow, data_use_dic['Reuse_service_time_{}'.format(use_cycle)].loc[
                                                   product, 'distribution'],
                                               data_use_dic['Reuse_service_time_{}'.format(use_cycle)].loc[
                                                   product, 'location'],
                                               data_use_dic['Reuse_service_time_{}'.format(use_cycle)].loc[
                                                   product, 'scale'],
                                               data_use_dic['Reuse_service_time_{}'.format(use_cycle)].loc[
                                                   product, 'shape'],
                                               iteration, n_years,
                                               defined_distributions_pd)

            # append the use stock

            stock_use_dic[str(use_cycle)].append(use_stock)

            # Create and append the flows
            # Outflow use to storage
            to_storage = use_outflow * data_use_dic['Reuse_service_time_{}'.format(use_cycle)].loc[
                product, 'to_storage']
            flow_use_to_storage_dic[str(use_cycle)].append(to_storage)

            if use_cycle != considered_use_cycles:
                # Outflow use to next use
                use_to_next_use = use_outflow * data_use_dic['Reuse_service_time_{}'.format(use_cycle)].loc[
                    product, 'to_use']
                flow_use_to_next_use_dic[str(use_cycle)].append(use_to_next_use)

                # Outflow to disposal
                flow_use_to_disposal_dic[str(use_cycle)].append(use_outflow * (
                        1 - data_use_dic['Reuse_service_time_{}'.format(use_cycle)].loc[product, 'to_storage'] -
                        data_use_dic['Reuse_service_time_{}'.format(use_cycle)].loc[product, 'to_use']))

            else:
                flow_use_to_disposal_dic[str(use_cycle)].append(use_outflow * (
                        1 - data_use_dic['Reuse_service_time_{}'.format(use_cycle)].loc[product, 'to_storage']))

            # Storage pahse ############################

            cohort_matrices_dic['storage_time_' + str(use_cycle)][
                product], storage_stock, storage_outflow = extend_cohort_matrix(
                    cohort_matrices_dic['storage_time_' + str(use_cycle)][product],
                    to_storage,
                    data_storage_dic['Reuse_storage_time_{}'.format(use_cycle)].loc[product, 'distribution'],
                    data_storage_dic['Reuse_storage_time_{}'.format(use_cycle)].loc[product, 'location'],
                    data_storage_dic['Reuse_storage_time_{}'.format(use_cycle)].loc[product, 'scale'],
                    data_storage_dic['Reuse_storage_time_{}'.format(use_cycle)].loc[product, 'shape'],
                    iteration, n_years, defined_distributions_pd)

            # append storage stock
            stock_storage_dic[str(use_cycle)].append(storage_stock)

            # create and append hibernating stock outflows
            if use_cycle != considered_use_cycles:

                storage_to_next_use = storage_outflow * data_storage_dic['Reuse_storage_time_{}'.format(use_cycle)].loc[
                    product, 'to_use']
                flow_storage_to_next_use_dic[str(use_cycle)].append(storage_to_next_use)

                flow_storage_to_disposal_dic[str(use_cycle)].append(storage_outflow - storage_to_next_use)

                # overwrite inflow, inflow is what goes into the next use stage

                inflow = storage_to_next_use + use_to_next_use


            else:
                flow_storage_to_disposal_dic[str(use_cycle)].append(storage_outflow)

    # Combine the output to totals
    ## Total use stock
    total_use_stock = []
    total_hoarding_stock = []
    to_disposal_flow = []

    for product_index in range(len(products)):
        use_stock_collector = []
        hoarding_stock_collector = []
        disposal_flow_collector = []

        for i in range(considered_use_cycles):
            use_cycle = i + 1
            # putting each stage in list
            use_stock_collector.append(stock_use_dic[str(use_cycle)][product_index])

            hoarding_stock_collector.append(stock_storage_dic[str(use_cycle)][product_index])

            disposal_flow_collector.append(flow_use_to_disposal_dic[str(use_cycle)][product_index] + \
                                           flow_storage_to_disposal_dic[str(use_cycle)][product_index])

        total_use_stock.append(sum(use_stock_collector))
        total_hoarding_stock.append(sum(hoarding_stock_collector))
        to_disposal_flow.append(sum(disposal_flow_collector))

    output_pd['total_use_stock'] = total_use_stock
    output_pd['total_hoarding_stock'] = total_hoarding_stock
    output_pd['to_disposal_flow'] = to_disposal_flow

    # Adding stocks and flows
    for i in range(considered_use_cycles):
        use_cycle = i + 1

        # Stocks
        output_pd['use_stock_{}'.format(use_cycle)] = stock_use_dic[str(use_cycle)]
        output_pd['storage_stock_{}'.format(use_cycle)] = stock_storage_dic[str(use_cycle)]

        # Flows
        output_pd['use_{}_to_storage_{}_flow'.format(use_cycle, use_cycle)] = flow_use_to_storage_dic[str(use_cycle)]
        output_pd['use_{}_to_disposal_flow'.format(use_cycle)] = flow_use_to_disposal_dic[str(use_cycle)]
        output_pd['storage_{}_to_disposal_flow'.format(use_cycle)] = flow_storage_to_disposal_dic[str(use_cycle)]

        if use_cycle != considered_use_cycles:
            output_pd['use_{}_to_use_{}_flow'.format(use_cycle, use_cycle + 1)] = flow_use_to_next_use_dic[
                str(use_cycle)]
            output_pd['storage_{}_to_use_{}_flow'.format(use_cycle, use_cycle + 1)] = flow_storage_to_next_use_dic[
                str(use_cycle)]

    return output_pd, cohort_matrices_dic
