'''
------------------------------------------------------------------------
Functions for generating demographic objects necessary for the OG-ITA
model

This module defines the following function(s):
    get_fert()
    get_imm()
    get_mort()
    pop_rebin()
    immsolve()
    get_pop_objs()
------------------------------------------------------------------------
'''
# Import libraries, packages, and modules
import os
import numpy as np
import scipy.interpolate as si
import scipy.optimize as opt
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.ticker import MultipleLocator


'''
------------------------------------------------------------------------
Define functions
------------------------------------------------------------------------
'''


def get_fert(tot_model_pers, graph=False):
    '''
    --------------------------------------------------------------------
    This function generates a vector of fertility rates for Italy by
    model period age that corresponds to the fertility rate data by age
    in years (Source: PopulationDataAll_ITA_2.xlsx)
    --------------------------------------------------------------------
    INPUTS:
    tot_model_pers = integer >= 3, total number of agent life periods
                     (E+S) from time period of birth (s=1) to latest
                     possible period of death (s=E+S)
    graph   = boolean, =True if want graphical output

    OTHER FUNCTIONS AND FILES CALLED BY THIS FUNCTION:
        data/demographics/PopulationDataAll_ITA_2.xlsx

    OBJECTS CREATED WITHIN FUNCTION:
    min_yr           = integer >= 0, age in years at which agents are
                       born, minimum age. This is hard coded at 1
                       because source data begins at age 0
    max_yr           = integer >= 4, age in years at which agents die
                       with certainty, maximum age. This is hard coded
                       at 100 because source data ends at age 99
    cur_path         = string, path in which calling file resides
    pop_file         = string, path of population data source .xlsx file
    popworksheet     = string, name of Excel worksheet within target
                       spreadsheet
    pop_data         = 100 x 5 DataFrame, Age, Population, Fertility
                       Rate (f), Net Immigration Rate (i), Surv Rates
                       (1-rho)
    cur_pop_2015     = (100,) Series, population level by age
    curr_pop_2015pct = (100,) Series, population percentage by age
    fert_rates_yrly  = (100,) Series, fertility rates by age
    ages_yrly        = (100,) Series, ages from 0 to 99
    age_midp         = (100,) Series, midpoint age of age bins
                       corresponding to ages_yrly
    fert_func        = function, generated by interp1d function, takes
                       ages and returns the interpolated fertility rates
    binsize          = scalar > 0, size of each model period bin in data
                       years
    num_sub_bins     = scalar, an arbitrarily and deliberately large
                       number of sub-bins that each population bin will
                       be broken up into
    sub_binsize      = scalar, number of years represented by each sub-
                       bin (should be a small fraction of a year)
    len_sub_bins     = scalar, length of a model period in data sub-bins
    age_sub          = (num_sub_bins*100,) vector, midpoint ages of each
                       data sub-bin
    curr_pop_sub     = (num_sub_bins*100,) vector, population linearly
                       interpolated from data in each sub-bin
    fert_rates_sub   = (num_sub_bins*100,) vector, fertility rates by
                       sub-bin interpolated from fert_func()
    pred_ind         = (num_sub_bins*100,) boolean vector, =True if
                       period is one that must be interpolated
    age_pred         = (num_sub_bins*100-some,) vector, midpoint age in
                       years corresponding to each period to be
                       interpolated
    fert_rates       = (tot_model_pers,) vector, fertility rates for
                       each model period of life
    i                = integer >= 0, index of model period being
                       computed
    beg_sub_bin      = integer >= 0, index of beginning sub-bin for
                       calculation of average fertility rate of given
                       model period
    end_sub_bin      = integer >= 0, index of ending sub-bin + 1 for
                       calculation of average fertility rate of given
                       model period

    FILES CREATED BY THIS FUNCTION:
        fert_rates.png

    RETURNS: fert_rates
    --------------------------------------------------------------------
    '''
    min_yr = 1
    max_yr = 100

    # Get current Italy population data (2015)
    cur_path = os.path.split(os.path.abspath(__file__))[0]
    pop_file = os.path.join(cur_path,
                            'data/demographics/' +
                            'PopulationDataAll_ITA_2015.xlsx')
    popworksheet = 'popdata'
    pop_data = pd.ExcelFile(pop_file).parse(popworksheet)
    curr_pop_2015 = pop_data['Population (ω)']
    curr_pop_2015pct = curr_pop_2015 / curr_pop_2015.sum()
    fert_rates_yrly = pop_data['Fertility Rate (f)']
    ages_yrly = pop_data['Age']
    age_midp = ages_yrly + 0.5

    # Generate interpolation functions for fertility rates
    fert_func = si.interp1d(age_midp, fert_rates_yrly, kind='cubic')
    # Calculate average fertility rate in each age bin using trapezoid
    # method with a large number of points in each bin.
    binsize = (max_yr - min_yr + 1) / tot_model_pers
    num_sub_bins = int(1000)
    sub_binsize = 1 / num_sub_bins
    len_sub_bins = ((np.float64((max_yr - min_yr + 1) * num_sub_bins)) /
                    tot_model_pers)
    age_sub = (np.linspace(min_yr - 1, max_yr - sub_binsize,
                           int((max_yr - min_yr + 1) * num_sub_bins)) +
               0.5 * sub_binsize)
    curr_pop_sub = np.repeat(np.float64(curr_pop_2015pct) /
                             num_sub_bins, num_sub_bins)
    fert_rates_sub = np.zeros_like(curr_pop_sub)
    pred_ind = np.logical_and(age_sub > age_midp[min_yr - 1],
                              age_sub < age_midp[max_yr - 1])
    age_pred = age_sub[pred_ind]
    fert_rates_sub[pred_ind] = np.float64(fert_func(age_pred))
    fert_rates = np.zeros(tot_model_pers)
    end_sub_bin = 0
    for i in range(tot_model_pers):
        beg_sub_bin = int(end_sub_bin)
        end_sub_bin = int(np.rint((i + 1) * len_sub_bins))
        fert_rates[i] = \
            ((curr_pop_sub[beg_sub_bin:end_sub_bin] *
              fert_rates_sub[beg_sub_bin:end_sub_bin]).sum() /
             curr_pop_sub[beg_sub_bin:end_sub_bin].sum())

    if graph:
        '''
        ----------------------------------------------------------------
        cur_path       = string, path name of current directory
        image_fldr     = string, folder in current path to save files
        image_dir      = string, total path of images folder
        image_path     = string, path of file name of figure to be saved
        age_fine_pred  = (300,) vector, equally spaced support of ages
                         between the minimum and maximum interpolating
                         ages
        fert_fine_pred = (300,) vector, interpolated fertility rates
                         based on age_fine_pred
        age_fine       = (302,) vector of ages including leading and
                         trailing zeros
        fert_fine      = (302,) vector of fertility rates including
                         leading and trailing zeros
        age_mid_new    = (tot_model_pers,) vector, midpoint age of each
                         model period age bin
        ----------------------------------------------------------------
        '''
        # Create directory if images directory does not already exist
        cur_path = os.path.split(os.path.abspath(__file__))[0]
        image_fldr = 'images/demographics'
        image_dir = os.path.join(cur_path, image_fldr)
        if not os.access(image_dir, os.F_OK):
            os.makedirs(image_dir)

        # Generate finer age vector and fertility rate vector for
        # graphing cubic spline interpolating function
        age_fine_pred = np.linspace(age_midp[min_yr - 1],
                                    age_midp[max_yr - 1], 300)
        fert_fine_pred = fert_func(age_fine_pred)
        age_fine = np.hstack((min_yr - 1, age_fine_pred, max_yr))
        fert_fine = np.hstack((0, fert_fine_pred, 0))
        age_mid_new = (np.linspace(min_yr - 1, max_yr - binsize,
                                   tot_model_pers) + 0.5 * binsize)

        fig, ax = plt.subplots()
        plt.scatter(age_midp, fert_rates_yrly, s=70, c='blue',
                    marker='o', label='Data')
        plt.scatter(age_mid_new, fert_rates, s=40, c='red', marker='d',
                    label='Model period (integrated)')
        plt.plot(age_fine, fert_fine, label='Cubic spline')
        # for the minor ticks, use no labels; default NullFormatter
        minorLocator = MultipleLocator(1)
        ax.xaxis.set_minor_locator(minorLocator)
        plt.grid(b=True, which='major', color='0.65', linestyle='-')
        # plt.title('Fitted fertility rate function by age ($f_{s}$)',
        #     fontsize=20)
        plt.xlabel(r'Age $s$')
        plt.ylabel(r'Fertility rate $f_{s}$')
        plt.xlim((min_yr - 1, max_yr + 1))
        plt.ylim((-0.15 * fert_fine_pred.max(),
                  1.15 * fert_fine_pred.max()))
        plt.legend(loc='upper right')
        plt.text(-5, -0.018, 'Source: PopulationDataAll_ITA_2015.xlsx.',
                 fontsize=9)
        plt.tight_layout(rect=(0, 0.03, 1, 1))
        image_path = os.path.join(image_dir, 'fert_rates')
        plt.savefig(image_path)
        # plt.show()
        plt.close()

    return fert_rates


def get_imm(tot_model_pers, graph=False):
    '''
    --------------------------------------------------------------------
    This function generates a vector of immigration rates for Italy by
    model period age that corresponds to the immigration rate data by
    age in years (Source: PopulationDataAll_ITA_2.xlsx)
    --------------------------------------------------------------------
    INPUTS:
    tot_model_pers = integer >= 3, total number of agent life periods
                     (E+S) from time period of birth (s=1) to latest
                     possible period of death (s=E+S)
    graph   = boolean, =True if want graphical output

    OTHER FUNCTIONS AND FILES CALLED BY THIS FUNCTION:
        data/demographics/PopulationDataAll_ITA_2.xlsx

    OBJECTS CREATED WITHIN FUNCTION:
    min_yr           = integer >= 0, age in years at which agents are
                       born, minimum age. This is hard coded at 1
                       because source data begins at age 0
    max_yr           = integer >= 4, age in years at which agents die
                       with certainty, maximum age. This is hard coded
                       at 100 because source data ends at age 99
    cur_path         = string, path in which calling file resides
    pop_file         = string, path of population data source .xlsx file
    popworksheet     = string, name of Excel worksheet within target
                       spreadsheet
    pop_data         = 100 x 5 DataFrame, Age, Population, Fertility
                       Rate (f), Net Immigration Rate (i), Surv Rates
                       (1-rho)
    cur_pop_2015     = (100,) Series, population level by age
    curr_pop_2015pct = (100,) Series, population percentage by age
    imm_rates_yrly   = (100,) Series, immigration rates by age
    ages_yrly        = (100,) Series, ages from 0 to 99
    age_midp         = (100,) Series, midpoint age of age bins
                       corresponding to ages_yrly
    imm_func         = function, generated by interp1d function, takes
                       ages and returns the interpolated immigration
                       rates
    binsize          = scalar > 0, size of each model period bin in data
                       years
    num_sub_bins     = scalar, an arbitrarily and deliberately large
                       number of sub-bins that each population bin will
                       be broken up into
    sub_binsize      = scalar, number of years represented by each sub-
                       bin (should be a small fraction of a year)
    len_sub_bins     = scalar, length of a model period in data sub-bins
    age_sub          = (num_sub_bins*100,) vector, midpoint ages of each
                       data sub-bin
    curr_pop_sub     = (num_sub_bins*100,) vector, population linearly
                       interpolated from data in each sub-bin
    imm_rates_sub    = (num_sub_bins*100,) vector, immigration rates by
                       sub-bin interpolated from imm_func()
    pred_ind         = (num_sub_bins*100,) boolean vector, =True if
                       period is one that must be interpolated
    age_pred         = (num_sub_bins*100-some,) vector, midpoint age in
                       years corresponding to each period to be
                       interpolated
    imm_rates        = (tot_model_pers,) vector, immigration rates for
                       each model period of life
    i                = integer >= 0, index of model period being
                       computed
    beg_sub_bin      = integer >= 0, index of beginning sub-bin for
                       calculation of average immigration rate of given
                       model period
    end_sub_bin      = integer >= 0, index of ending sub-bin + 1 for
                       calculation of average immigration rate of given
                       model period

    FILES CREATED BY THIS FUNCTION:
        imm_rates.png

    RETURNS: imm_rates
    --------------------------------------------------------------------
    '''
    min_yr = 1
    max_yr = 100

    # Get current Italy population data (2015)
    cur_path = os.path.split(os.path.abspath(__file__))[0]
    pop_file = os.path.join(cur_path,
                            'data/demographics/' +
                            'PopulationDataAll_ITA_2015.xlsx')
    popworksheet = 'popdata'
    pop_data = pd.ExcelFile(pop_file).parse(popworksheet)
    curr_pop_2015 = pop_data['Population (ω)']
    curr_pop_2015pct = curr_pop_2015 / curr_pop_2015.sum()
    imm_rates_yrly = pop_data['Net Immigration Rate (i)']
    ages_yrly = pop_data['Age']
    age_midp = ages_yrly + 0.5

    # Generate interpolation functions for immigration rates
    imm_func = si.interp1d(age_midp, imm_rates_yrly, kind='cubic')
    # Calculate average immigration rate in each age bin using trapezoid
    # method with a large number of points in each bin.
    binsize = (max_yr - min_yr + 1) / tot_model_pers
    num_sub_bins = int(1000)
    sub_binsize = 1 / num_sub_bins
    len_sub_bins = ((np.float64((max_yr - min_yr + 1) * num_sub_bins)) /
                    tot_model_pers)
    age_sub = (np.linspace(min_yr - 1, max_yr - sub_binsize,
                           int((max_yr - min_yr + 1) * num_sub_bins)) +
               0.5 * sub_binsize)
    curr_pop_sub = np.repeat(np.float64(curr_pop_2015pct) /
                             num_sub_bins, num_sub_bins)
    imm_rates_sub = np.zeros_like(curr_pop_sub)
    pred_ind = np.logical_and(age_sub > age_midp[min_yr - 1],
                              age_sub < age_midp[max_yr - 1])
    age_pred = age_sub[pred_ind]
    imm_rates_sub[pred_ind] = np.float64(imm_func(age_pred))
    imm_rates = np.zeros(tot_model_pers)
    end_sub_bin = 0
    for i in range(tot_model_pers):
        beg_sub_bin = int(end_sub_bin)
        end_sub_bin = int(np.rint((i + 1) * len_sub_bins))
        imm_rates[i] = \
            ((curr_pop_sub[beg_sub_bin:end_sub_bin] *
              imm_rates_sub[beg_sub_bin:end_sub_bin]).sum() /
             curr_pop_sub[beg_sub_bin:end_sub_bin].sum())

    if graph:
        '''
        ----------------------------------------------------------------
        cur_path      = string, path name of current directory
        image_fldr    = string, folder in current path to save files
        image_dir     = string, total path of images folder
        image_path    = string, path of file name of figure to be saved
        age_fine_pred = (300,) vector, equally spaced support of ages
                        between the minimum and maximum interpolating
                        ages
        imm_fine_pred = (300,) vector, interpolated immigration rates
                        based on age_fine_pred
        age_fine      = (302,) vector of ages including leading and
                        trailing zeros
        imm_fine      = (302,) vector of immigration rates including
                        leading and trailing zeros
        age_mid_new   = (tot_model_pers,) vector, midpoint age of each
                         model period age bin
        ----------------------------------------------------------------
        '''
        # Create directory if images directory does not already exist
        cur_path = os.path.split(os.path.abspath(__file__))[0]
        image_fldr = 'images/demographics'
        image_dir = os.path.join(cur_path, image_fldr)
        if not os.access(image_dir, os.F_OK):
            os.makedirs(image_dir)

        # Generate finer age vector and fertility rate vector for
        # graphing cubic spline interpolating function
        age_fine_pred = np.linspace(age_midp[min_yr - 1],
                                    age_midp[max_yr - 1], 300)
        imm_fine_pred = imm_func(age_fine_pred)
        age_fine = np.hstack((min_yr - 1, age_fine_pred, max_yr))
        imm_fine = np.hstack((0, imm_fine_pred, 0))
        age_mid_new = (np.linspace(min_yr - 1, max_yr - binsize,
                                   tot_model_pers) + 0.5 * binsize)

        fig, ax = plt.subplots()
        plt.scatter(age_midp, imm_rates_yrly, s=70, c='blue',
                    marker='o', label='Data')
        plt.scatter(age_mid_new, imm_rates, s=40, c='red', marker='d',
                    label='Model period (integrated)')
        plt.plot(age_fine, imm_fine, label='Cubic spline')
        # for the minor ticks, use no labels; default NullFormatter
        minorLocator = MultipleLocator(1)
        ax.xaxis.set_minor_locator(minorLocator)
        plt.grid(b=True, which='major', color='0.65', linestyle='-')
        # plt.title('Fitted fertility rate function by age ($f_{s}$)',
        #     fontsize=20)
        plt.xlabel(r'Age $s$')
        plt.ylabel(r'Immigration rate $i_{s}$')
        plt.xlim((min_yr - 1, max_yr + 1))
        plt.ylim((-0.15 * imm_fine_pred.max(),
                  1.15 * imm_fine_pred.max()))
        plt.legend(loc='upper right')
        plt.text(-5, -0.018, 'Source: PopulationDataAll_ITA_2015.xlsx.',
                 fontsize=9)
        plt.tight_layout(rect=(0, 0.03, 1, 1))
        image_path = os.path.join(image_dir, 'imm_rates')
        plt.savefig(image_path)
        # plt.show()
        plt.close()

    return imm_rates


def get_mort(tot_model_pers, graph=False):
    '''
    --------------------------------------------------------------------
    This function generates a vector of mortality rates for Italy by
    model period age that corresponds to the mortality rate data by age
    age in years (Source: PopulationDataAll_ITA_2.xlsx)
    --------------------------------------------------------------------
    INPUTS:
    tot_model_pers = integer >= 3, total number of agent life periods
                     (E+S) from time period of birth (s=1) to latest
                     possible period of death (s=E+S)
    graph          = boolean, =True if want graphical output

    OTHER FUNCTIONS AND FILES CALLED BY THIS FUNCTION:
        data/demographics/PopulationDataAll_ITA_2.xlsx

    OBJECTS CREATED WITHIN FUNCTION:
    min_yr           = integer >= 0, age in years at which agents are
                       born, minimum age. This is hard coded at 1
                       because source data begins at age 0
    max_yr           = integer >= 4, age in years at which agents die
                       with certainty, maximum age. This is hard coded
                       at 100 because source data ends at age 99
    cur_path         = string, path in which calling file resides
    pop_file         = string, path of population data source .xlsx file
    popworksheet     = string, name of Excel worksheet within target
                       spreadsheet
    pop_data         = 100 x 5 DataFrame, Age, Population, Fertility
                       Rate (f), Net Immigration Rate (i), Surv Rates
                       (1-rho)
    surv_rates_yrly  = (100,) Series, survival rates by age 0 to 99
    mort_rates_yrly  = (100,) Series, mortality rates by age 0 to 99
    mort_rates_intrp = (100,) Series, mortality rates by age 0 to 99,
                       with the years 94 to 98 linearly interpolated
                       because of a weird blip in the data
    ages_yrly        = (100,) Series, age in mortality rate data
    infmort_rate     = scalar > 0, infant mortality rate from 2016 est.:
                       3.3 deaths per 1000.
                       'https://www.indexmundi.com/italy/' +
                       'infant_mortality_rate.html'
    binsize          = scalar > 0, size of each model period bin in data
                       years
    num_sub_bins     = scalar, an arbitrarily and deliberately large
                       number of sub-bins that each population bin will
                       be broken up into
    len_subbins      = scalar, length of a model period in data sub-bins
    mort_rates_sub   = (num_sub_bins*100,) vector, mortality rates by
                       sub-bin implied by mort_rates_intrp
    mort_rates       = (tot_model_pers,) vector, mortality rates that
                       correspond to each period of life
    i               = integer >= 0, index of model period being computed
    beg_sub_bin     = integer >= 0, index of beginning sub-bin for
                      calculation of cumulative mortality rate of given
                      model period
    end_sub_bin     = integer >= 0, index of ending sub-bin + 1 for
                      calculation of cumulative mortality rate of given
                      model period

    FILES CREATED BY THIS FUNCTION:
        mort_rates.png

    RETURNS: mort_rates, infmort_rate
    --------------------------------------------------------------------
    '''
    min_yr = 1
    max_yr = 100

    # Get current Italy population data (2015)
    cur_path = os.path.split(os.path.abspath(__file__))[0]
    pop_file = os.path.join(cur_path,
                            'data/demographics/' +
                            'PopulationDataAll_ITA_2015.xlsx')
    popworksheet = 'popdata'
    pop_data = pd.ExcelFile(pop_file).parse(popworksheet)
    surv_rates_yrly = pop_data['Surv Rates (1-ρ)']
    mort_rates_yrly = 1 - surv_rates_yrly

    # Linearly interpolate the spike in mortality rates between age 94
    # and 98
    mort_rates_intrp = mort_rates_yrly.copy()
    mort_rates_intrp[94:99] = \
        np.linspace(mort_rates_yrly[94], mort_rates_yrly[98],
                    98 - 94 + 1)

    ages_yrly = pop_data['Age']
    infmort_rate = 3.3 / 1000  # taken from 2016 est. (see docstring)

    # Calculate implied mortality rates in sub-bins of mort_rates_intrp
    binsize = (max_yr - min_yr + 1) / tot_model_pers
    num_sub_bins = int(100)
    len_subbins = ((np.float64((max_yr - min_yr + 1) * num_sub_bins)) /
                   tot_model_pers)
    mort_rates_sub = np.zeros(num_sub_bins * (max_yr - min_yr + 1),
                              dtype=float)
    for i in range(max_yr - min_yr + 1):
        mort_rates_sub[i * num_sub_bins:(i + 1) * num_sub_bins] = \
            (1 - ((1 - mort_rates_intrp[i]) ** (1.0 / num_sub_bins)))
    mort_rates = np.zeros(tot_model_pers)
    end_sub_bin = 0
    for i in range(tot_model_pers):
        beg_sub_bin = int(end_sub_bin)
        end_sub_bin = int(np.rint((i + 1) * len_subbins))
        mort_rates[i] = \
            (1 - (1 - (mort_rates_sub[beg_sub_bin:end_sub_bin])).prod())
    mort_rates[-1] = 1.0  # Mortality rate in last period is set to 1

    if graph:
        '''
        ----------------------------------------------------------------
        cur_path    = string, path name of current directory
        image_fldr  = string, folder in current path to save files
        image_dir   = string, total path of images folder
        image_path  = string, path of file name of figure to be saved
        age_mid_new = (tot_model_pers,) vector, midpoint age of each
                      model period age bin
        ----------------------------------------------------------------
        '''
        # Create directory if images directory does not already exist
        cur_path = os.path.split(os.path.abspath(__file__))[0]
        image_fldr = 'images/demographics'
        image_dir = os.path.join(cur_path, image_fldr)
        if not os.access(image_dir, os.F_OK):
            os.makedirs(image_dir)

        age_mid_new = (np.linspace(min_yr - 1, max_yr - binsize,
                                   tot_model_pers) + 0.5 * binsize)

        fig, ax = plt.subplots()
        plt.scatter(np.hstack([0, ages_yrly + 1]),
                    np.hstack([infmort_rate, mort_rates_intrp]), s=20,
                    c='blue', marker='o', label='Data')
        plt.scatter(np.hstack([0, age_mid_new]),
                    np.hstack([infmort_rate, mort_rates]), s=40,
                    c='red', marker='d',
                    label='Model period (cumulative)')
        plt.plot(np.hstack([0, ages_yrly + 1]),
                 np.hstack([infmort_rate, mort_rates_intrp]))
        # for the minor ticks, use no labels; default NullFormatter
        minorLocator = MultipleLocator(1)
        ax.xaxis.set_minor_locator(minorLocator)
        plt.grid(b=True, which='major', color='0.65', linestyle='-')
        # plt.title('Fitted mortality rate function by age ($rho_{s}$)',
        #     fontsize=20)
        plt.xlabel(r'Age $s$')
        plt.ylabel(r'Mortality rate $\rho_{s}$')
        plt.xlim((min_yr - 2, max_yr + 1))
        plt.ylim((-0.05, 1.05))
        plt.legend(loc='upper left')
        plt.text(-5, -0.2, 'Source: PopulationDataAll_ITA_2.xlsx.',
                 fontsize=9)
        plt.tight_layout(rect=(0, 0.03, 1, 1))
        image_path = os.path.join(image_dir, 'mort_rates')
        plt.savefig(image_path)
        # plt.show()
        plt.close()

    return mort_rates, infmort_rate


def pop_rebin(curr_pop_dist, totpers_new):
    '''
    --------------------------------------------------------------------
    For cases in which tot_model_pers (E+S) is less than the number of
    periods in the population distribution data, this function
    calculates a new population distribution vector with tot_model_pers
    (E+S) elements.
    --------------------------------------------------------------------
    INPUTS:
    curr_pop_dist = (N,) vector, population distribution over N periods
    totpers_new   = integer >= 4, number of periods to which we are
                    transforming the population distribution

    OTHER FUNCTIONS AND FILES CALLED BY THIS FUNCTION: None

    OBJECTS CREATED WITHIN FUNCTION:
    totpers_orig = integer >= 4, number of periods in curr_pop_dist
                   original data
    curr_pop_new = (totpers,) vector, new population distribution over
                   totpers (E+S) periods that approximates curr_pop_dist
    num_sub_bins = integer > 1, an arbitrarily and deliberately large
                   number of sub-bins that each population bin will be
                   broken up into
    len_subbins  = scalar > 1, number of bins (in terms of sub-bins)
                   that equal the length of one of the new periods
    i            = integer >= 0, index of new bin number
    beg_sub_bin  = integer >= 0, index of beginning bin in sum
    end_sub_bin  = integer >= 0, index of ending bin plus 1 in sum

    FILES CREATED BY THIS FUNCTION: None

    RETURNS: curr_pop_new
    --------------------------------------------------------------------
    '''
    totpers_orig = len(curr_pop_dist)
    if int(totpers_new) == totpers_orig:
        curr_pop_new = curr_pop_dist
    elif int(totpers_new) < totpers_orig:
        num_sub_bins = float(10000)
        curr_pop_sub = np.repeat(np.float64(curr_pop_dist) /
                                 num_sub_bins, num_sub_bins)
        len_subbins = ((np.float64(totpers_orig * num_sub_bins)) /
                       totpers_new)
        curr_pop_new = np.zeros(totpers_new, dtype=np.float64)
        end_sub_bin = 0
        for i in range(totpers_new):
            beg_sub_bin = int(end_sub_bin)
            end_sub_bin = int(np.rint((i + 1) * len_subbins))
            curr_pop_new[i] = \
                curr_pop_sub[beg_sub_bin:end_sub_bin].sum()
        # Return curr_pop_new to single precision float (float32)
        # datatype
        curr_pop_new = np.float32(curr_pop_new)

    return curr_pop_new


def immsolve(imm_rates, *args):
    '''
    --------------------------------------------------------------------
    This function generates a vector of errors representing the
    difference in two consecutive periods stationary population
    distributions. This vector of differences is the zero-function
    objective used to solve for the immigration rates vector, similar to
    the original immigration rates vector from get_imm(), that sets the
    steady-state population distribution by age equal to the population
    distribution in period T1
    --------------------------------------------------------------------
    INPUTS:
    imm_rates = (totpers,) vector, adjusted immigration rates by age s
    args      = length 5 tuple, (fert_rates, mort_rates, infmort_rate,
                omega_cur, g_n_SS)

    OTHER FUNCTIONS AND FILES CALLED BY THIS FUNCTION: None

    OBJECTS CREATED WITHIN FUNCTION:
    fert_rates    = (totpers,) vector, fertility rates by age s
    mort_rates    = (totpers,) vector, mortality rates by age s
    infmort_rate  = scalar > 0, infant mortality rate from 2016 est.
    omega_cur_lev = (totpers,) vector, population distribution (levels)
                    by age in current period
    g_n_SS        = scalar, original steady-state population growth rate
    omega_cur_pct = (totpers,) vector, population distribution (pct) by
                    age in current period
    totpers       = integer >= 3, number of agent life periods (E+S)
    OMEGA         = (totpers, totpers) matrix, transition matrix for
                    population distribution law of motion
    omega_new_pct = (totpers,) vector, population distribution (pct) by
                    age in next period
    omega_errs    = (totpers,) vector, difference between omega_new_pct
                    and omega_cur_pct

    FILES CREATED BY THIS FUNCTION: None

    RETURNS: omega_errs
    --------------------------------------------------------------------
    '''
    fert_rates, mort_rates, infmort_rate, omega_cur_lev, g_n_SS = args
    omega_cur_pct = omega_cur_lev / omega_cur_lev.sum()
    totpers = len(fert_rates)
    OMEGA = np.zeros((totpers, totpers))
    OMEGA[0, :] = ((1 - infmort_rate) * fert_rates +
                   np.hstack((imm_rates[0], np.zeros(totpers - 1))))
    OMEGA[1:, :-1] += np.diag(1 - mort_rates[:-1])
    OMEGA[1:, 1:] += np.diag(imm_rates[1:])
    omega_new_pct = np.dot(OMEGA, omega_cur_pct) / (1 + g_n_SS)
    omega_errs = omega_new_pct - omega_cur_pct

    return omega_errs


def get_pop_objs(E, S, T1, T2, curr_year, GraphDiag=False):
    '''
    --------------------------------------------------------------------
    This function produces the demographics objects to be used in the
    OG-USA model package.
    --------------------------------------------------------------------
    INPUTS:
    E         = integer >= 1, number of model periods in which agent is
                not economically active
    S         = integer >= 3, number of model periods in which agent is
                economically active
    T1        = integer < T, number of periods after which the
                population dynamics are forced to be at the steady-state
                by adjusting immigration rates. This value should be far
                enough in the future that the demographics are close to
                the steady state
    T2        = integer > 2*S, number of periods after which the
                endogenous model variables should have reached the
                steady state
    curr_year = integer >= 2016, current year for which analysis will
                begin
    GraphDiag = boolean, =True if want graphical output and printed
                diagnostics

    OTHER FUNCTIONS AND FILES CALLED BY THIS FUNCTION:
        get_fert()
        get_mort()
        get_imm()
        pop_rebin()
        immsolve()
        data/demographics/PopulationDataAll_ITA_2.xlsx

    OBJECTS CREATED WITHIN FUNCTION:
    data_year       = int >= 2015, year represented by most recent data
    fert_rates      = (E+S,) vector, fertility rates that correspond to
                      each model period of life
    mort_rates      = (E+S,) vector, mortality rates that correspond to
                      each model period of life
    infmort_rate    = scalar > 0, infant mortality rate from 2016 est.
    mort_rates_S    = (S,) vector, mortality rates that correspond to
                      each economically active model period of life
    imm_rates_orig  = (E+S,) vector, immigration rates that correspond
                      to each model period of life
    OMEGA_orig      = (E+S, E+S) matrix, transition matrix for
                      population distribution law of motion
    eigvalues       = (E+S,) vector, eigenvalues of OMEGA matrix
    eigvectors      = (E+S, E+S) matrix, matrix of eigenvectors of OMEGA
                      where each column is the eigenvector that goes
                      with the corresponding eigenvalue in eigvalues
    g_n_SS          = scalar, steady-state population growth rate, which
                      is the largest real part of the eigenvalues
                      (normalized)
    eigvec_raw      = (E+S,) vector, nonnormalized eigenvector
                      corresponding to the largest real-part eigenvalue
    omega_SS_orig   = (E+S,) vector, steady-state population
                      distribution which is normalized by eigvec_raw
    omega_path_lev  = (E+S, T2+S-1) matrix, time path of population
                      distribution by age in levels (non-stationary)
    cur_path        = string, path in which calling file resides
    pop_file        = string, path of population data source .xlsx file
    popworksheet    = string, name of Excel worksheet within target
                      spreadsheet
    pop_data        = 100 x 5 DataFrame, Age, Population, Fertility Rate
                      (f), Net Immigration Rate (i), Surv Rates (1-rho)
    pop_2015        = (100,) vector, population for ages 0 to 99 in 2015
    pop_2015pct     = (100,) vector, population distribution (pct) for
                      ages 0 to 99 in 2015
    age_per_EpS     = (E+S,) vector, period numbers 1 through E+S
    pop_2015_EpS    = (E+S,) vector, population distribution by model
                      periods E + S in levels
    per             = int >= 0, index of period
    pop_curr        = (E+S,) vector, current-period population
                      distribution in levels
    pop_next        = (E+S,) vector, next-period population distribution
                      in levels
    g_n_curr        = scalar, population growth rate of economically
                      relevant population between pop_curr and pop_next
    pop_past        = (E+S,) vector, previous-period population
                      distribution in levels. We use this to generate
                      period-before-start-year population distribution
    fert_rates_100  = (100,) vector, fertility rates for each age from
                      data
    mort_rates_100  = (100,) vector, mortality rates for each age from
                      data
    imm_rates_100   = (100,) vector, immigration rates for each age from
                      data
    OMEGA_100       = (100, 100) matrix, OMEGA matrix with same
                      dimensions as data
    pop_curr_100    = (100,) vector, current population distribution
                      (levels) by age from data
    pop_next_100    = (100,) vector, next period population distribution
                      (levels) by age from data
    omega_cur_80    = (80,) vector, current population distribution for
                      economically relevant ages from data
    imm_tol         = scalar > 0, tolerance for fsolve in immsolve()
    omega_SSfx      = (E+S,) vector, adjusted steady-state distribution
                      (pct) at which the economy remains after period T1
    imm_args        = length 5 tuple, arguments to pass in to immsolve()
    imm_fulloutput  = length 2 dict, fsolve output
    imm_rates_adj   = (E+S,) vector, adjusted immigration rates that
                      induce the adjusted steady-state after period T1
    imm_diagdict    = length 6 dict, diagnostic output from fsolve
    omega_path_S    = (E+S, T2+S-1) matrix, time path of population
                      distribution by age in percents
    g_n_path        = (T2+S-1,) vector, time path of population growth
                      rate of economically relevant population
    omega_S_preTP   = (E+S,) vector, population distribution in percents
                      in period before initial period

    FILES CREATED BY THIS FUNCTION:
        images/demographics/OrigVsFixSSpop.png
        images/demographics/g_n_path.png
        images/demographics/OrigVsAdjImm.png
        images/demographics/PopDistPath.png

    RETURNS: omega_path_S.T, g_n_SS,
             omega_SSfx[-S:] / omega_SSfx[-S:].sum(), mort_rates_S,
             g_n_path, imm_rates_orig[-S:], imm_rates_adj[-S:],
             omega_S_preTP, omega_cur_80
    --------------------------------------------------------------------
    '''
    data_year = int(2015)

    # Get fertility, mortality, and immigration rates
    fert_rates = get_fert(E + S, graph=False)
    mort_rates, infmort_rate = get_mort(E + S, graph=False)
    mort_rates_S = mort_rates[-S:]
    imm_rates_orig = get_imm(E + S, graph=False)
    OMEGA_orig = np.zeros((E + S, E + S))
    OMEGA_orig[0, :] = \
        ((1 - infmort_rate) * fert_rates +
         np.hstack((imm_rates_orig[0], np.zeros(E + S - 1))))
    OMEGA_orig[1:, :-1] += np.diag(1 - mort_rates[:-1])
    OMEGA_orig[1:, 1:] += np.diag(imm_rates_orig[1:])

    # Solve for steady-state population growth rate and steady-state
    # population distribution by age using eigenvalue and eigenvector
    # decomposition
    eigvalues, eigvectors = np.linalg.eig(OMEGA_orig)
    g_n_SS = (eigvalues[np.isreal(eigvalues)].real).max() - 1
    eigvec_raw = \
        eigvectors[:,
                   (eigvalues[np.isreal(eigvalues)].real).argmax()].real
    omega_SS_orig = eigvec_raw / eigvec_raw.sum()

    # Generate time path of the nonstationary population distribution
    omega_path_lev = np.zeros((E + S, T2 + S - 1))
    cur_path = os.path.split(os.path.abspath(__file__))[0]
    pop_file = os.path.join(cur_path,
                            'data/demographics/' +
                            'PopulationDataAll_ITA_2015.xlsx')
    popworksheet = 'popdata'
    pop_data = pd.ExcelFile(pop_file).parse(popworksheet)
    pop_2015 = pop_data['Population (ω)']
    pop_2015pct = pop_2015 / pop_2015.sum()

    # Generate the current population distribution given that E+S might
    # be less than max_yr-min_yr+1
    age_per_EpS = np.arange(1, E + S + 1)
    pop_2015_EpS = pop_rebin(pop_2015, E + S)

    # Age most recent population level data to the current year of
    # analysis
    pop_curr = pop_2015_EpS.copy()
    for per in range(0, curr_year - data_year):
        pop_next = np.dot(OMEGA_orig, pop_curr)
        g_n_curr = ((pop_next[-S:].sum() - pop_curr[-S:].sum()) /
                    pop_curr[-S:].sum())
        pop_past = pop_curr.copy()
        pop_curr = pop_next.copy()

    # Age most recent population level data to the current year of
    # analysis fixed at 80 economically relevant periods
    fert_rates_100 = get_fert(100, graph=False)
    mort_rates_100 = get_mort(100, graph=False)[0]
    imm_rates_100 = get_imm(100, graph=False)
    OMEGA_100 = np.zeros((100, 100))
    OMEGA_100[0, :] = \
        ((1 - infmort_rate) * fert_rates_100 +
         np.hstack((imm_rates_100[0], np.zeros(100 - 1))))
    OMEGA_100[1:, :-1] += np.diag(1 - mort_rates_100[:-1])
    OMEGA_100[1:, 1:] += np.diag(imm_rates_100[1:])
    pop_curr_100 = pop_2015.copy()
    for per in range(0, curr_year - data_year):
        pop_next_100 = np.dot(OMEGA_100, pop_curr_100)
        pop_curr_100 = pop_next_100.copy()
    omega_cur_80 = pop_curr_100[-80:] / pop_curr_100[-80:].sum()

    # Generate time path of the population distribution in levels
    omega_path_lev[:, 0] = pop_curr.copy()
    for per in range(1, T2 + S - 1):
        pop_next = np.dot(OMEGA_orig, pop_curr)
        omega_path_lev[:, per] = pop_next.copy()
        pop_curr = pop_next.copy()

    # Force the population distribution after T1 periods to be the
    # steady-state distribution by adjusting immigration rates, holding
    # constant mortality, fertility, and SS growth rates
    imm_tol = 1e-14
    omega_SSfx = (omega_path_lev[:, T1] / omega_path_lev[:, T1].sum())
    imm_args = (fert_rates, mort_rates, infmort_rate,
                omega_path_lev[:, T1], g_n_SS)
    imm_fulloutput = opt.fsolve(immsolve, imm_rates_orig,
                                args=(imm_args), full_output=True,
                                xtol=imm_tol)
    imm_rates_adj = imm_fulloutput[0]
    imm_diagdict = imm_fulloutput[1]

    omega_path_S = (omega_path_lev[-S:, :] /
                    np.tile(omega_path_lev[-S:, :].sum(axis=0), (S, 1)))
    omega_path_S[:, T1:] = \
        np.tile(omega_path_S[:, T1].reshape((S, 1)),
                (1, T2 + S - 1 - T1))
    g_n_path = np.zeros(T2 + S - 1)
    g_n_path[0] = g_n_curr.copy()
    g_n_path[1:T1] = ((omega_path_lev[-S:, 1:T1].sum(axis=0) /
                       omega_path_lev[-S:, :T1 - 1].sum(axis=0)) - 1)
    omega_S_preTP = ((pop_past.copy()[-S:]) /
                     (pop_past.copy()[-S:].sum()))
    g_n_path[T1:] = g_n_SS

    if GraphDiag:
        '''
        ----------------------------------------------------------------
        image_fldr       = string, folder in current path to save files
        image_dir        = string, total path of images folder
        image_path       = string, path of file name of figure to be
                           saved
        omegaSSmaxdif    = scalar > 0, maximum absolute distance between
                           original steady-state population distribution
                           and population distribution in period T2
        omegaSSvTmaxdiff = scalar > 0, maximum absolute distance between
                           original steady-state population distribution
                           and adjusted steady-state distribution (T1)
        immtol_solved    = scalar, maximum absolute error between period
                           T1 distribution and distribution implied by
                           adjusted immigration rates
        OMEGA_adj        = (E+S, E+S) matrix, adjusted OMEGA matrix with
                           adjusted immigration rates
        eigvalues_adj    = (E+S,) vector, eigenvalues of OMEGA_adj
                           matrix
        eigvectors_adj   = (E+S, E+S) matrix, matrix of eigenvectors of
                           OMEGA_adj where each column is the
                           eigenvector that corresponds to eigenvalue in
                           eigvalues_adj
        g_n_SS_adj       = scalar, steady-state population growth rate,
                           largest real part of the eigenvalues
                           (normalized)
        omega_new        = (E+S,) vector, next-period population
                           distribution (pcts) when curr distribution is
                           steady-state
        abs_omega_errs   = (E+S,) vector, absolute value of difference
                           between steady-stated population distribution
                           and omega_new
        immratesmaxdiff  = scalar > 0, maximum absolute difference
                           between imm_rates and imm_rates_adj
        ----------------------------------------------------------------
        '''
        # Create directory if images directory does not already exist
        image_fldr = 'images/demographics'
        image_dir = os.path.join(cur_path, image_fldr)
        if not os.access(image_dir, os.F_OK):
            os.makedirs(image_dir)

        # Check whether original SS population distribution is close to
        # the period-T population distribution
        omegaSSmaxdif = \
            np.absolute(omega_SS_orig -
                        (omega_path_lev[:, T2] /
                         omega_path_lev[:, T2].sum())).max()
        if omegaSSmaxdif > 0.0003:
            print('POP. WARNING: Max. abs. dist. between original SS ' +
                  'pop. distn. and period-T2 pop. distn. is greater ' +
                  'than 0.0003. It is ' + str(omegaSSmaxdif) + '.')
        else:
            print('POP. SUCCESS: orig. SS pop. dist is very close to ' +
                  'period-T2 pop. distn. The maximum absolute ' +
                  'difference is ' + str(omegaSSmaxdif) + '.')
        print()

        # Plot the adjusted steady-state population distribution versus
        # the original steady-state population distribution. The
        # difference should besmall
        omegaSSvTmaxdiff = np.absolute(omega_SS_orig - omega_SSfx).max()
        if omegaSSvTmaxdiff > 0.0003:
            print('POP. WARNING: The maximum absolute difference ' +
                  'between any two corresponding points in the ' +
                  'original and adjusted steady-state population ' +
                  'distributions is ' + str(omegaSSvTmaxdiff) + ', ' +
                  'which is greater than 0.0003.')
        else:
            print('POP. SUCCESS: The maximum absolute difference ' +
                  'between any two corresponding points in the ' +
                  'original and adjusted steady-state population ' +
                  'distributions is ' + str(omegaSSvTmaxdiff) + '.')
        print()
        fig, ax = plt.subplots()
        plt.plot(age_per_EpS, omega_SS_orig,
                 label='Original SS Dist\'n')
        plt.plot(age_per_EpS, omega_SSfx, label='Fixed SS Dist\'n')
        # for the minor ticks, use no labels; default NullFormatter
        minorLocator = MultipleLocator(1)
        ax.xaxis.set_minor_locator(minorLocator)
        plt.grid(b=True, which='major', color='0.65', linestyle='-')
        # plt.title(
        #     'Original steady-state population distribution vs. fixed',
        #     fontsize=15)
        plt.xlabel(r'Age $s$')
        plt.ylabel(r'Pop. dist\'n $\omega_{s}$')
        plt.xlim((0, E + S + 1))
        plt.legend(loc='lower left')
        image_path = os.path.join(image_dir, 'OrigVsFixSSpop')
        plt.savefig(image_path)
        # plt.show()
        plt.close()

        # Print whether or not the adjusted immigration rates solved the
        # zero condition
        immtol_solved = \
            np.absolute(imm_diagdict['fvec'].max()) < imm_tol
        if immtol_solved:
            print('POP. SUCCESS: Adjusted immigration rates solved ' +
                  'with maximum absolute error of ' +
                  str(np.absolute(imm_diagdict['fvec'].max())) +
                  ', which is less than the tolerance of ' +
                  str(imm_tol) + '.')
        else:
            print('POP. WARNING: Adjusted immigration rates did not ' +
                  'solve. Maximum absolute error of ' +
                  str(np.absolute(imm_diagdict['fvec'].max())) +
                  ' is greater than the tolerance of ' + str(imm_tol) +
                  '.')
        print()
        # Test whether the steady-state growth rates implied by the
        # adjusted OMEGA matrix equals the steady-state growth rate of
        # the original OMEGA matrix
        OMEGA_adj = np.zeros((E + S, E + S))
        OMEGA_adj[0, :] = \
            ((1 - infmort_rate) * fert_rates +
             np.hstack((imm_rates_adj[0], np.zeros(E + S - 1))))
        OMEGA_adj[1:, :-1] += np.diag(1 - mort_rates[:-1])
        OMEGA_adj[1:, 1:] += np.diag(imm_rates_adj[1:])
        eigvalues_adj, eigvectors_adj = np.linalg.eig(OMEGA_adj)
        g_n_SS_adj = \
            (eigvalues_adj[np.isreal(eigvalues_adj)].real).max() - 1
        if np.max(np.absolute(g_n_SS_adj - g_n_SS)) > 10 ** (-8):
            print('FAILURE: The steady-state population growth rate ' +
                  'from adjusted OMEGA is different (diff is ' +
                  str(g_n_SS_adj - g_n_SS) + ') than the steady-' +
                  'state population growth rate from the original ' +
                  'OMEGA.')
        elif np.max(np.absolute(g_n_SS_adj - g_n_SS)) <= 10 ** (-8):
            print('SUCCESS: The steady-state population growth rate ' +
                  'from adjusted OMEGA is close to (diff is ' +
                  str(g_n_SS_adj - g_n_SS) + ') the steady-' +
                  'state population growth rate from the original ' +
                  'OMEGA.')
        print()

        # Plot the time path of the population growth rate
        fig, ax = plt.subplots()
        plt.plot(np.arange(1, T2 + S), g_n_path, label='original')
        plt.axvline(x=T1, color='red')
        plt.axvline(x=T2, color='red')
        plt.axhline(y=g_n_SS, xmin=0, xmax=75, color='red',
                    linestyle='dashed', linewidth=1)
        minorLocator = MultipleLocator(10)
        ax.xaxis.set_minor_locator(minorLocator)
        plt.grid(b=True, which='major', color='0.65', linestyle='-')
        # plt.title('Time path of population growth', fontsize=15)
        plt.xlabel(r'Period $t$')
        plt.ylabel(r'Pop. growth rate $\tilde{g}_{n,t}$')
        plt.xlim((0, T2 + S + 1))
        plt.tight_layout()
        image_path = os.path.join(image_dir, 'g_n_path')
        plt.savefig(image_path)
        # plt.show()
        plt.close()

        # Do another test of the adjusted immigration rates. Create the
        # new OMEGA matrix implied by the new immigration rates. Plug in
        # the adjusted steady-state population distribution. Hit it with
        # the new OMEGA transition matrix and it should return the new
        # steady-state population distribution
        omega_new = np.dot(OMEGA_adj, omega_SSfx) / (1 + g_n_SS)
        abs_omega_errs = np.absolute(omega_new - omega_SSfx)
        if abs_omega_errs.max() > 1e-6:
            print('POPULATION FAILURE: The maximum absolute diff' +
                  'erence between the original steady-state pop. ' +
                  'dist\'n and a next period dist\'n starting from ' +
                  'the steady state is ' + str(abs_omega_errs.max()) +
                  '>1e-6.')
        elif abs_omega_errs.max() <= 1e-6:
            print('POPULATION SUCCESS: The maximum absolute diff' +
                  'erence between the original steady-state pop. ' +
                  'dist\'n and a next period dist\'n starting from ' +
                  'the steady state is ' + str(abs_omega_errs.max()) +
                  '<=1e-6.')
        print()

        # Plot the original immigration rates versus the adjusted
        # immigration rates
        immratesmaxdiff = \
            np.absolute(imm_rates_orig - imm_rates_adj).max()
        print('The maximum absolute distance between any two points ' +
              'of the original immigration rates and adjusted ' +
              'immigration rates is ' + str(immratesmaxdiff) + '.')
        print()
        fig, ax = plt.subplots()
        plt.plot(age_per_EpS, imm_rates_orig,
                 label='Original Imm. Rates')
        plt.plot(age_per_EpS, imm_rates_adj, label='Adj. Imm. Rates')
        # for the minor ticks, use no labels; default NullFormatter
        minorLocator = MultipleLocator(1)
        ax.xaxis.set_minor_locator(minorLocator)
        plt.grid(b=True, which='major', color='0.65', linestyle='-')
        # plt.title('Original immigration rates vs. adjusted',
        #           fontsize=15)
        plt.xlabel(r'Age $s$')
        plt.ylabel(r'Imm. rates $i_{s}$')
        plt.xlim((0, E + S + 1))
        plt.legend(loc='upper center')
        image_path = os.path.join(image_dir, 'OrigVsAdjImm')
        plt.savefig(image_path)
        # plt.show()
        plt.close()

        # Plot population distributions for data_year, curr_year,
        # curr_year+0.5*S, curr_year + S, and omega_SSfx
        fig, ax = plt.subplots()
        plt.plot(age_per_EpS, pop_2015pct, label='2015 pop.')
        plt.plot(age_per_EpS,
                 (omega_path_lev[:, 0] / omega_path_lev[:, 0].sum()),
                 label=str(curr_year) + ' pop.')
        plt.plot(age_per_EpS, (omega_path_lev[:, int(0.5 * S)] /
                               omega_path_lev[:, int(0.5 * S)].sum()),
                 label=str(curr_year + int(0.5 * S)) + ' pop.')
        plt.plot(age_per_EpS, (omega_path_lev[:, int(S)] /
                               omega_path_lev[:, int(S)].sum()),
                 label=str(curr_year + int(S)) + ' pop.')
        plt.plot(age_per_EpS, omega_SSfx, label='Adj. SS pop.')
        # for the minor ticks, use no labels; default NullFormatter
        minorLocator = MultipleLocator(1)
        ax.xaxis.set_minor_locator(minorLocator)
        plt.grid(b=True, which='major', color='0.65', linestyle='-')
        # plt.title('Population distribution at points in time path',
        #           fontsize=15)
        plt.xlabel(r'Age $s$')
        plt.ylabel(r'Pop. dist\'n $\omega_{s}$')
        plt.xlim((0, E + S + 1))
        plt.legend(loc='lower left')
        image_path = os.path.join(image_dir, 'PopDistPath')
        plt.savefig(image_path)
        # plt.show()
        plt.close()

    return (omega_path_S.T, g_n_SS,
            omega_SSfx[-S:] / omega_SSfx[-S:].sum(), mort_rates_S,
            g_n_path, imm_rates_orig[-S:], imm_rates_adj[-S:],
            omega_S_preTP, omega_cur_80)
