from copy import copy
from HARK.ConsumptionSaving.ConsGenIncProcessModel import \
    IndShockExplicitPermIncConsumerType, IndShockConsumerType, PersistentShockConsumerType
import numpy as np
import unittest

GenIncDictionary = { 
    "CRRA": 2.0,                           # Coefficient of relative risk aversion
    "Rfree": 1.03,                         # Interest factor on assets
    "DiscFac": 0.96,                       # Intertemporal discount factor
    "LivPrb" : [0.98],                     # Survival probability
    "AgentCount" : 10000,                  # Number of agents of this type (only matters for simulation)
    "aNrmInitMean" : 0.0,                  # Mean of log initial assets (only matters for simulation)
    "aNrmInitStd"  : 1.0,                  # Standard deviation of log initial assets (only for simulation)
    "pLvlInitMean" : 0.0,                  # Mean of log initial permanent income (only matters for simulation)
    "pLvlInitStd"  : 0.0,                  # Standard deviation of log initial permanent income (only matters for simulation)
    "PermGroFacAgg" : 1.0,                 # Aggregate permanent income growth factor (only matters for simulation)
    "T_age" : None,                        # Age after which simulated agents are automatically killed
    "T_cycle" : 1,                         # Number of periods in the cycle for this agent type
# Parameters for constructing the "assets above minimum" grid
    "aXtraMin" : 0.001,                    # Minimum end-of-period "assets above minimum" value
    "aXtraMax" : 30,                       # Maximum end-of-period "assets above minimum" value               
    "aXtraExtra" : [0.005,0.01],           # Some other value of "assets above minimum" to add to the grid
    "aXtraNestFac" : 3,                    # Exponential nesting factor when constructing "assets above minimum" grid
    "aXtraCount" : 48,                     # Number of points in the grid of "assets above minimum"
# Parameters describing the income process
    "PermShkCount" : 7,                    # Number of points in discrete approximation to permanent income shocks
    "TranShkCount" : 7,                    # Number of points in discrete approximation to transitory income shocks
    "PermShkStd" : [0.1],                  # Standard deviation of log permanent income shocks
    "TranShkStd" : [0.1],                  # Standard deviation of log transitory income shocks
    "UnempPrb" : 0.05,                     # Probability of unemployment while working
    "UnempPrbRet" : 0.005,                 # Probability of "unemployment" while retired
    "IncUnemp" : 0.3,                      # Unemployment benefits replacement rate
    "IncUnempRet" : 0.0,                   # "Unemployment" benefits when retired
    "tax_rate" : 0.0,                      # Flat income tax rate
    "T_retire" : 0,                        # Period of retirement (0 --> no retirement)
    "BoroCnstArt" : 0.0,                  # Artificial borrowing constraint; imposed minimum level of end-of period assets
    "CubicBool" : False,                  # Use cubic spline interpolation when True, linear interpolation when False
    "vFuncBool" : True,                   # Whether to calculate the value function during solution    
# More parameters specific to "Explicit Permanent income" shock model
    "pLvlPctiles" : np.concatenate(([0.001, 0.005, 0.01, 0.03], np.linspace(0.05, 0.95, num=19),[0.97, 0.99, 0.995, 0.999])),
    "PermGroFac" : [1.0],                  # Permanent income growth factor - long run permanent income growth doesn't work yet    
}

class testIndShockExplicitPermIncConsumerType(unittest.TestCase):

    def setUp(self):
        self.agent = IndShockExplicitPermIncConsumerType(
            cycles=1,
            **GenIncDictionary)
        self.agent.solve()

    def test_solution(self):
        pLvlGrid = self.agent.pLvlGrid[0]
        self.assertAlmostEqual(self.agent.pLvlGrid[0][0],
                               1.0)

        self.assertAlmostEqual(
            self.agent.solution[0].mLvlMin(pLvlGrid[0]),
            0.0)

        self.assertAlmostEqual(
            self.agent.solution[0].cFunc(10,pLvlGrid[5]).tolist(),
            5.6030075768585075)

class testPersistentShockConsumerType(unittest.TestCase):

    def setUp(self):
        # "persistent idiosyncratic shocks" model
        PrstIncCorr = 0.98       # Serial correlation coefficient for persistent income
        persistent_shocks = copy(GenIncDictionary)
        persistent_shocks['PrstIncCorr'] = PrstIncCorr

        # "persistent idisyncratic shocks" consumer
        self.agent = PersistentShockConsumerType(cycles=1,
            **persistent_shocks)
        self.agent.solve()

    def test_solution(self):
        pLvlGrid = self.agent.pLvlGrid[0]

        self.assertAlmostEqual(
            self.agent.solution[0].cFunc(10,pLvlGrid[1]).tolist(),
            5.6030075768585075)

    def test_simulation(self):
        self.agent.T_sim = 25
        self.agent.track_vars = ['mLvlNow',
                                 'cLvlNow',
                                 'pLvlNow']
        self.agent.initializeSim()
        self.agent.simulate()

        self.assertAlmostEqual(
            np.mean(self.agent.history['mLvlNow']),
            1.2043946738813716)
