# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: hark-env
#     language: python
#     name: hark-env
# ---

# +
import HARK.ConsumptionSaving.ConsPortfolioFrameModel as cpfm
import HARK.ConsumptionSaving.ConsPortfolioModel as cpm
import numpy as np

from HARK.utilities import (
    CRRAutility,
)
# -

# The `FrameAgentType` is an alternative way to specify a model.
#
# The library contains a demonstration of this form of model, `ConsPortfolioFrameModel`, which is a replica of the `ConsPortfolioModel`.
#
# This notebook compares the results of simulations of the two models.

# +
pct = cpm.PortfolioConsumerType(T_sim=5000, AgentCount=200)
pct.cycles = 0

# Solve the model under the given parameters

pct.solve()
pct.track_vars += [
    "mNrm",
    "cNrm",
    "Share",
    "aNrm",
    "Risky",
    "Adjust",
    "PermShk",
    "TranShk",
    "bNrm",
    "who_dies"
]

pct.initialize_sim()

pct.simulate()
# -



shock_history = {shock : pct.history[shock] for shock in pct.shocks}

shock_history['who_dies'] = pct.history['who_dies']

# +
pcft = cpfm.PortfolioConsumerFrameType(
    T_sim=5000,
    AgentCount=200,
    read_shocks = True
)

pcft.shock_history = shock_history

pcft.cycles = 0

# Solve the model under the given parameters

pcft.solve()
pcft.track_vars += [
    "mNrm",
    "cNrm",
    "Share",
    "aNrm",
    "Adjust",
    "PermShk",
    "TranShk",
    "bNrm",
    'U'
]

pcft.initialize_sim()

pcft.simulate()
# -

pcft.history['TranShk'].shape

# +
import matplotlib.pyplot as plt

plt.plot(range(5000), pct.history['PermShk'].mean(axis=1), label = 'original')
plt.plot(range(5000), pcft.history['PermShk'].mean(axis=1), label = 'frames', alpha = 0.5)
plt.legend()
# -

plt.plot(range(5000), pct.history['TranShk'].mean(axis=1), label = 'original')
plt.plot(range(5000), pcft.history['TranShk'].mean(axis=1), label = 'frames', alpha = 0.5)
plt.legend()

plt.plot(range(5000), pct.history['bNrm'].mean(axis=1), label = 'original')
plt.plot(range(5000), pcft.history['bNrm'].mean(axis=1), label = 'frames', alpha = 0.5)
plt.legend()

# +
#plt.plot(range(5000), pct.history['Risky'].mean(axis=1), label = 'original')
#plt.plot(range(5000), pcft.history['Risky'].mean(axis=1), label = 'frames', alpha = 0.5)
#plt.legend()
# -

plt.plot(range(5000), pct.history['aNrm'].mean(axis=1), label = 'original')
plt.plot(range(5000), pcft.history['aNrm'].mean(axis=1), label = 'frames', alpha = 0.5)
plt.legend()

plt.plot(range(5000), pct.history['mNrm'].mean(axis=1), label = 'original')
plt.plot(range(5000), pcft.history['mNrm'].mean(axis=1), label = 'frames', alpha = 0.5)
plt.legend()

plt.plot(range(5000), pct.history['cNrm'].mean(axis=1), label = 'original')
plt.plot(range(5000), pcft.history['cNrm'].mean(axis=1), label = 'frames', alpha = 0.5)
plt.legend()

# **TODO**: Handly Risky as an aggregate value.

# +
#pct.history['Risky'][:3, :3]

# +
#pcft.history['Risky'][:3, :3]
# -

plt.plot(range(5000), pct.history['Share'].mean(axis=1), label = 'original')
plt.plot(range(5000), pcft.history['Share'].mean(axis=1), label = 'frames', alpha = 0.5)
plt.legend()

plt.plot(range(5000), pcft.history['cNrm'].mean(axis=1), label = 'frames - cNrm', alpha = 0.5)
plt.plot(range(5000), pcft.history['U'].mean(axis=1), label = 'frames - U', alpha = 0.5)
plt.legend()

pcft.history['U']

pcft.history['U'].mean(axis=1)

pcft.history['U'][0,:]

pcft.history['cNrm'][0,:]

pcft.parameters['CRRA']

CRRAutility(pcft.history['cNrm'][0,:], 5)

pcft.frames[4].parents

# # Visualizing the Transition Equations

import networkx as nx

# +
g = nx.DiGraph()

g.add_nodes_from([
    (frame.name(), 
     {'control' : frame.control, 'reward' : frame.reward, 'aggregate' : frame.aggregate})
    for frame in pcft.frames
])

for frame in pcft.frames:
    for child in frame.children:
        g.add_nodes_from([
            (child.name(), 
             {
                 'control' : child.control,
                 'reward' : child.reward,
                 'aggregate' : child.aggregate
             })])
        g.add_edge(frame.name(), child.name())

# +
pos = nx.drawing.layout.kamada_kawai_layout(g)

node_options = {
    "node_size": 2500,
    "node_color": "white",
    "edgecolors": "black",
    "linewidths": 1,
    "pos" : pos
}

edge_options = {
    "node_size": 2500,
    "width": 2,
    "pos" : pos
}

label_options = {
    "font_size": 12,
     #"labels" : {node : str(node[0]) if len(node) == 1 else str(node) for node in g.nodes},
    "pos" : pos
}


reward_nodes = [k for k,v in g.nodes(data = True) if v['reward']]
control_nodes = [k for k,v in g.nodes(data = True) if v['control']]
aggregate_nodes = [k for k,v in g.nodes(data = True) if v['aggregate']]

chance_nodes = [node for node in g.nodes() 
                if node not in reward_nodes 
                and node not in control_nodes
                and node not in aggregate_nodes
               ]

plt.figure(figsize=(16,12))

nx.draw_networkx_nodes(g, nodelist = chance_nodes, node_shape = 'o', **node_options)
nx.draw_networkx_nodes(g, nodelist = reward_nodes, node_shape = 'd', **node_options)
nx.draw_networkx_nodes(g, nodelist = control_nodes, node_shape = 's', **node_options)
nx.draw_networkx_nodes(g, nodelist = aggregate_nodes, node_shape = 'h', **node_options)
nx.draw_networkx_edges(g, **edge_options)

nx.draw_networkx_labels(g, **label_options)


# -

# Note that in the HARK `ConsIndShockModel`, from which the `ConsPortfolio` model inherits, the aggregate permanent shocks are considered to be portions of the permanent shocks experienced by the agents, not additions to those idiosyncratic shocks. Hence, they do not show up directly in the problem solved by the agent. This explains why the aggregate income levels are in a separarte component of the graph.

# # Building the Solver [INCOMPLETE]

# Preliminery work towards a generic solver for FramedAgentTypes.

controls = [frame for frame in pcft.frames if frame.control]


def get_expected_return_function(control):
    # Input: a control frame
    # Returns: function of the control variable (control frame target)
    #      that returns the expected return, which is
    #          the sum of:
    #              - direct rewards
    #              - expected value of next-frame states (not yet implemented)
    #
    
    rewards = [child for child in control.children if child.reward]
    expected_values = [] # TODO
    
    ## note: function signature is what's needed for scipy.optimize
    def expected_return_function(x, *args):
        ##   returns the sum of 
        ##     the reward functions evaluated in context of
        ##       - parameters
        ##       - the control variable input
        
        # x - array of inputs, here the control frame target
        # args - a tuple of other parameters needed to complete the function
        
        expected_return = 0
        
        for reward in rewards:
            ## TODO: figuring out the ordering of `x` and `args` needed for multiple downstream scopes
            
            
            local_context = {}
            
            # indexing through the x and args values
            i = 0
            num_control_vars = None
            
            # assumes that all frame scopes list model variables first, parameters later
            # should enforce or clarify at the frame level.
            for var in reward.scope:
                if var in control.target:
                    local_context[var] = x[i]
                    i = i + 1
                elif var in pcft.parameters:
                    if num_control_vars is None:
                        num_control_vars = i
                    
                    local_context[var] = args[i - num_control_vars]
                    i = i + 1
                    
            # can `self` be implicit here?
            expected_return += reward.transition(reward, **local_context)
        
        return expected_return
    
    return expected_return_function            


def optimal_policy_function(control):
    
    erf = get_expected_return_function(control)
    constraints = control.constraints ## these will reference the context of the control transition, including scope
    
    ## Returns function:
    ##   input: control frame scope
    ##   output: result of scipy.optimize of the erf with respect to constraints
    ##           getting the optimal input (control variable) value
    return func
    


def approximate_optimal_policy_function(control, grid):
    ## returns a new function:
    ##   that is an interpolation over optimal_policy_function
    ##   over the grid
    
    return func

# # Solving methods on the Frame, FrameAgentType



# # Using Frame Solution in simulation and comparing result.



StarcoderdataJupyterScriptsDedupFiltered