from HARK.ConsumptionSaving.ConsRepAgentModel import (
    RepAgentConsumerType,
    RepAgentMarkovConsumerType,
)
import numpy as np
import unittest

class testRepAgentConsumerType(unittest.TestCase):

    def setUp(self):
        self.agent = RepAgentConsumerType()
        self.agent.solve()

    def test_solution(self):
        self.assertAlmostEqual(
            self.agent.solution[0].cFunc(10).tolist(),
            1.7130553407923501
        )

    def test_simulation(self):
        # Simulate the representative agent model
        self.agent.T_sim = 100
        self.agent.track_vars = ["cNrmNow", "mNrmNow", "Rfree", "wRte"]
        self.agent.initializeSim()
        self.agent.simulate()

class testRepAgentMarkovConsumerType(unittest.TestCase):

    def setUp(self):
        self.agent = RepAgentMarkovConsumerType()
        self.agent.IncomeDstn[0] = 2 * [self.agent.IncomeDstn[0]]
        self.agent.solve()

    def test_solution(self):
         self.assertAlmostEqual(
             self.agent.solution[0].cFunc[0](10).tolist(),
             1.3829466326248048)

    def test_simulation(self):
        # Simulate the representative agent model
        self.agent.T_sim = 100
        self.agent.track_vars = ["cNrmNow", "mNrmNow", "Rfree", "wRte", "MrkvNow"]
        self.agent.initializeSim()
        self.agent.simulate()
