"""
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 plot_funcs_der, plot_funcs
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."
            )
