import numpy as np
import pandas as pd
import scipy.interpolate as si
from scipy.integrate import quad
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.ticker import MultipleLocator
import os

def get_consump(periods, min_yr, max_yr, graph = False):
    '''
    --------------------------------------------------------------------
    This function generates a vector of consumption by model period
    age that corresponds to the expenditures data by age in years
    (Source: Consumer Expenditure Survey (CEX), 2016
    https://www.bls.gov/cex/2016/combined/age.pdf)
    --------------------------------------------------------------------
    INPUTS:
    periods = integer >= 3, total number of periods agents live
    min_yr  = integer >= 0, age at which agents are born
    max_yr  = integer >= 4, age at which agents die
    graph   = boolean, =True if want graphical output
    OTHER FUNCTIONS AND FILES CALLED BY THIS FUNCTION:

    OBJECTS CREATED WITHIN FUNCTION:
    cur_path       = string, path in which calling file resides
    consump_data   = (9,) vector, average consumption for given age bins;
                     The first two and the last one are not from data
    age_midp       = (9,) vector, midpoint age of age bins for consumption data
    consump_tilde  = function, generated by interp1d function, takes ages and
                     returns the interpolated consumption
    population     = (101, ) vector, population data in December 2016 for age 0
                     to 100 individuals
    population_dist= (101, ) vector, population distribution in percentile for
                     age 0 to 100 individuals
    data_mean      = scalar, average consumption in the population from data
    consump_curve  = function, consump_tilde - consump_data[0]
    consump_wt     = (101, ) vector, weighted consumption for each age
    consump_mean   = scalar, weighted mean consumption indicated by the curve
                     (-consump_data[0])
    factor_c       = scalar, scaling variable to match average consumption
                     indicated by curve with data_mean
    consump_func   = function, generated by scaling consump_tilde, to match
                     data_mean
    age_bin        = (periods, ) vector, age for each model period of life
    consump        = (periods, ) vector, consumption for each model period of
                     life

    FILES CREATED BY THIS FUNCTION:
        cons_by_age.png
    RETURNS: consump
    --------------------------------------------------------------------
    '''
    cur_path = os.path.split(os.path.abspath(__file__))[0]
    # Obtain consumption data by age bins
    consump_data = np.array([29000, 29200, 34438, 52838, 66444, 71166, 61346, 50873, 38691, 29000])
    age_midp = np.array([-0.5, 2.5, 13, 30, 40, 50, 60, 70, 87, 105.5])
    # Generate interpolation functions for consumption
    consump_tilde = si.interp1d(age_midp, consump_data, kind='cubic')
    # Obtain population data by each age
    population = pd.read_table(os.path.join(cur_path, 'data', 'population.csv'), sep=',')
    population = np.array(population['Population'], dtype = float)
    population_dist = population / population.sum()
    # Adjust the curve to match mean consumption in the data
    data_mean = 57311
    consump_curve = lambda x: consump_tilde(x) - 29000
    consump_wt = np.zeros(101)
    for i in range(100):
        consump_wt[i] = quad(consump_curve, i, i + 1)[0] * population_dist[i]
    curve_mean = consump_wt.sum()
    factor_c = (data_mean - 29000) / curve_mean
    consump_func = lambda x: consump_curve(x) * factor_c + 29000
    # Find consumption by age
    consump = np.zeros(periods)
    age = np.zeros(periods)
    age_bin = np.linspace(min_yr, max_yr, periods + 1)
    for i in range(periods):
        consump[i] = quad(consump_func, age_bin[i], age_bin[i + 1])[0]

    if graph:
        # Create directory if images directory does not already exist
        cur_path = os.path.split(os.path.abspath(__file__))[0]
        output_fldr = "images"
        output_dir = os.path.join(cur_path, output_fldr)
        if not os.access(output_dir, os.F_OK):
            os.makedirs(output_dir)
        fig , ax = plt .subplots ()
        plt.scatter(age_midp[2: -1], consump_data[2: -1], s=70, c='blue', marker ='o', label ='Data')
        plt.scatter([age_midp[0], age_midp[1], age_midp[ -1]], [consump_data[0], consump_data[1], consump_data[ -1]], s=70,
                      c='red', marker ='o', label ='Added data ')
        plt.plot(np.linspace(-0.05, 105.5, 1000), consump_tilde(np.linspace(-0.05, 105.5, 1000)), c='green', linestyle ='--',
                 label ='Cubic spline ' + r'$\tilde{c}(s)$')
        plt.plot(np.linspace(-0.05, 105.5, 1000), consump_func(np.linspace(-0.05, 105.5, 1000)), c='black', linestyle ='-',
                   label ='$c(s)$')
        minorLocator = MultipleLocator(1)
        ax.xaxis.set_minor_locator(minorLocator)
        plt.grid(b=True , which ='major ', color ='0.65 ', linestyle ='-')
        plt.title('Average consumption expenditure by age $c(s^i)$', fontsize =15)
        plt.xlabel(r'Age $s$ ')
        plt.ylabel(r' Consumption $c(s^i)$')
        plt.xlim((-3 , 113))
        plt.ylim((25000 , 1.05 * consump_func(np.linspace(-0.05, 105.5, 1000)).max ()))
        plt.legend(loc='upper right ')
        plt.annotate('Source : Consumer Expenditure ' + 'Survey (CEX ), 2016 Summary Data .', (0,0), (0, -32), xycoords='axes fraction', textcoords='offset points', va='top')
        plt.tight_layout( rect =(0 , 0.03 , 1, 1))
        output_path = os.path.join(output_dir, "cons_by_age")
        plt.savefig(output_path)
        plt.close()

    return consump
