"""
This file tests whether HARK's models are initialized correctly.
"""


# Bring in modules we need
import unittest
import numpy as np
from HARK.ConsumptionSaving.ConsIndShockModel import PerfForesightConsumerType
from HARK.ConsumptionSaving.ConsIndShockModel import KinkedRconsumerType, init_kinked_R
from HARK.ConsumptionSaving.ConsIndShockModel import IndShockConsumerType, init_idiosyncratic_shocks, init_lifecycle
from HARK.ConsumptionSaving.ConsMarkovModel import MarkovConsumerType
from HARK.utilities import plotFuncsDer, plotFuncs
from copy import copy

class testInitialization(unittest.TestCase):
    # We don't need a setUp method for the tests to run, but it's convenient
    # if we want to test various things on the same model in different test_*
    # methods.
    def test_PerfForesightConsumerType(self):
        try:
            model = PerfForesightConsumerType()
        except:
            self.fail("PerfForesightConsumerType failed to initialize with Params.init_perfect_foresight.")

    def test_IndShockConsumerType(self):
        try:
            model = IndShockConsumerType(**init_lifecycle)
        except:
            self.fail("IndShockConsumerType failed to initialize with Params.init_lifecycle.")

    def test_KinkedRconsumerType(self):
        try:
            model = KinkedRconsumerType(**init_kinked_R)
        except:
            self.fail("KinkedRconsumerType failed to initialize with Params.init_kinked_R.")

    def test_MarkovConsumerType(self):
        try:
            unemp_length = 5 # Averange length of unemployment spell
            urate_good = 0.05        # Unemployment rate when economy is in good state
            urate_bad = 0.12         # Unemployment rate when economy is in bad state
            bust_prob = 0.01         # Probability of economy switching from good to bad
            recession_length = 20    # Averange length of bad state
            p_reemploy =1.0/unemp_length
            p_unemploy_good = p_reemploy*urate_good/(1-urate_good)
            p_unemploy_bad = p_reemploy*urate_bad/(1-urate_bad)
            boom_prob = 1.0/recession_length
            MrkvArray = np.array([[(1-p_unemploy_good)*(1-bust_prob),p_unemploy_good*(1-bust_prob),
                                   (1-p_unemploy_good)*bust_prob,p_unemploy_good*bust_prob],
                                  [p_reemploy*(1-bust_prob),(1-p_reemploy)*(1-bust_prob),
                                   p_reemploy*bust_prob,(1-p_reemploy)*bust_prob],
                                  [(1-p_unemploy_bad)*boom_prob,p_unemploy_bad*boom_prob,
                                   (1-p_unemploy_bad)*(1-boom_prob),p_unemploy_bad*(1-boom_prob)],
                                  [p_reemploy*boom_prob,(1-p_reemploy)*boom_prob,
                                   p_reemploy*(1-boom_prob),(1-p_reemploy)*(1-boom_prob)]])

            # Make a consumer with serially correlated unemployment, subject to boom and bust cycles
            init_serial_unemployment = copy(init_idiosyncratic_shocks)
            init_serial_unemployment['MrkvArray'] = [MrkvArray]
            init_serial_unemployment['UnempPrb'] = 0 # to make income distribution when employed
            init_serial_unemployment['global_markov'] = False
            SerialUnemploymentExample = MarkovConsumerType(**init_serial_unemployment)
        except:
            self.fail("MarkovConsumerType failed to initialize with boom/bust unemployment.")
