from HARK.ConsumptionSaving.ConsIndShockModel import PerfForesightConsumerType
import numpy as np
import unittest


class testPerfForesightConsumerType(unittest.TestCase):
    def setUp(self):
        self.agent = PerfForesightConsumerType()
        self.agent_infinite = PerfForesightConsumerType(cycles=0)
        self.agent_constrained = PerfForesightConsumerType(cycles = 0, BoroCnstArt = 0.0)
        
        PF_dictionary = {
            "CRRA": 2.5,
            "DiscFac": 0.96,
            "Rfree": 1.03,
            "LivPrb": [0.98],
            "PermGroFac": [1.01],
            "T_cycle": 1,
            "cycles": 0,
            "AgentCount": 10000,
        }
        self.agent_alt = PerfForesightConsumerType(**PF_dictionary)

    def test_default_solution(self):
        self.agent.solve()
        c = self.agent.solution[0].cFunc

        self.assertEqual(c.x_list[0], -0.9805825242718447)
        self.assertEqual(c.x_list[1], 0.01941747572815533)
        self.assertEqual(c.y_list[0], 0)
        self.assertEqual(c.y_list[1], 0.511321002804608)
        self.assertEqual(c.decay_extrap, False)

    def test_another_solution(self):
        self.agent_alt.DiscFac = 0.90
        self.agent_alt.solve()
        self.assertAlmostEqual(
            self.agent_alt.solution[0].cFunc(10).tolist(), 3.9750093524820787
        )

    def test_checkConditions(self):
        self.agent_infinite.checkConditions()
        self.assertTrue(self.agent_infinite.conditions["AIC"])
        self.assertTrue(self.agent_infinite.conditions["GICPF"])
        self.assertTrue(self.agent_infinite.conditions["RIC"])
        self.assertTrue(self.agent_infinite.conditions["FHWC"])

    def test_simulation(self):

        self.agent_infinite.solve()

        # Create parameter values necessary for simulation
        SimulationParams = {
            "AgentCount": 10000,  # Number of agents of this type
            "T_sim": 120,  # Number of periods to simulate
            "aNrmInitMean": -6.0,  # Mean of log initial assets
            "aNrmInitStd": 1.0,  # Standard deviation of log initial assets
            "pLvlInitMean": 0.0,  # Mean of log initial permanent income
            "pLvlInitStd": 0.0,  # Standard deviation of log initial permanent income
            "PermGroFacAgg": 1.0,  # Aggregate permanent income growth factor
            "T_age": None,  # Age after which simulated agents are automatically killed
        }

        self.agent_infinite(
            **SimulationParams
        )  # This implicitly uses the assignParameters method of AgentType

        # Create PFexample object
        self.agent_infinite.track_vars = ["mNrmNow"]
        self.agent_infinite.initializeSim()
        self.agent_infinite.simulate()

        self.assertAlmostEqual(
            np.mean(self.agent_infinite.history["mNrmNow"], axis=1)[40],
            -23.008063500363942,
        )

        self.assertAlmostEqual(
            np.mean(self.agent_infinite.history["mNrmNow"], axis=1)[100],
            -27.164608851546927,
        )

        ## Try now with the manipulation at time step 80

        self.agent_infinite.initializeSim()
        self.agent_infinite.simulate(80)

        # This actually does nothing because aNrmNow is
        # epiphenomenal. Probably should change mNrmNow instead
        self.agent_infinite.state_now['aNrmNow'] += (
            -5.0
        )
        self.agent_infinite.simulate(40)

        self.assertAlmostEqual(
            np.mean(self.agent_infinite.history["mNrmNow"], axis=1)[40],
            -23.008063500363942,
        )

        self.assertAlmostEqual(
            np.mean(self.agent_infinite.history["mNrmNow"], axis=1)[100],
            -29.140261331951606,
        )
        
    def test_stable_points(self):
        
        # Solve the constrained agent. Stable points exists only with a
        # borrowing constraint.
        self.agent_constrained.solve()
        
        # Check against pre-computed values.
        self.assertEqual(self.agent_constrained.solution[0].mNrmSS , 1.0)
        # Check that they are both the same, since the problem is deterministic
        self.assertEqual(self.agent_constrained.solution[0].mNrmSS,
                         self.agent_constrained.solution[0].mNrmTrg)

@@ -305,98 +305,8 @@
 =0)%0A
-        self.agent_constrained = PerfForesightConsumerType(cycles = 0, BoroCnstArt = 0.0)%0A
     
@@ -3789,35 +3789,118 @@
 nt.%0A        
-self.agent_
+constrained_agent = PerfForesightConsumerType(cycles = 0, BoroCnstArt = 0.0)%0A        %0A        
 constrained.
@@ -3894,24 +3894,30 @@
  constrained
+_agent
 .solve()%0A   
@@ -3988,35 +3988,24 @@
 assertEqual(
-self.agent_
 constrained.
@@ -3995,32 +3995,38 @@
 qual(constrained
+_agent
 .solution%5B0%5D.mNr
@@ -4141,27 +4141,16 @@
 rtEqual(
-self.agent_
 constrai
@@ -4144,32 +4144,38 @@
 qual(constrained
+_agent
 .solution%5B0%5D.mNr
@@ -4204,27 +4204,16 @@
         
-self.agent_
 constrai
@@ -4215,16 +4215,22 @@
 strained
+_agent
 .solutio

Fix PerForesight