text: # -*- coding: utf-8 -*-
  '''
  Classes to solve the steady state of liquid and illiquid assets model
  '''
  from __future__ import print_function
  
  
  import sys 
  sys.path.insert(0,'../')
  
  import numpy as np
  import scipy as sc
  from scipy.stats import norm 
  from scipy.interpolate import interp1d, interp2d, griddata, RegularGridInterpolator
  from scipy import sparse as sp
  import time
  from .SharedFunc3 import Transition, ExTransitions, GenWeight, MakeGridkm, Tauchen, Fastroot
  
  
  class SteadyStateTwoAsset:
  
      '''
      Classes to solve the steady state of liquid and illiquid assets model
      '''
  
      def __init__(self, par, mpar, grid):
           
          self.par = par
          self.mpar = mpar
          self.grid = grid
         
      def SolveSteadyState(self):
          '''
          solve for steady state
          
          returns
          ----------
           par : dict
               parametres
           mpar : dict
               parametres
           grid: dict
               grid for solution
           Output : float
               steady state output  
           targets : dict
              steady state stats
           Vm : np.array
              marginal value of assets m
           Vk : np.array
              marginal value of assets m
           joint_distr : np.array
              joint distribution of m and h
           Copula : dict
              points for interpolation of joint distribution
           c_a_star : np.array
              policy function for consumption w/ adjustment
           c_n_star : np.array
              policy function for consumption w/o adjustment
           psi_star : np.array
              continuation value of holding capital
           m_a_star : np.array   
              policy function for asset m w/ adjustment
           m_n_star : np.array   
              policy function for asset m w/o adjustment
           mutil_c_a : np.array
              marginal utility of c w/ adjustment
           mutil_c_n : np.array
              marginal utility of c w/o adjustment   
           mutil_c : np.array
              marginal utility of c w/ & w/o adjustment   
           P_H : np.array
              transition probability    
              
          
          '''
         
          ## Set grid h
          grid = self.grid
          resultStVar=self.StochasticsVariance(self.par, self.mpar, grid)
         
          P_H = resultStVar['P_H'].copy()
          grid = resultStVar['grid'].copy()
          par = resultStVar['par'].copy()
         
          grid = MakeGridkm(self.mpar, grid, grid['k_min'], grid['k_max'], grid['m_min'], grid['m_max'])
          meshes = {}
          meshes['m'], meshes['k'], meshes['h'] =  np.meshgrid(grid['m'],grid['k'],grid['h'],indexing='ij')
          
          ## Solve for steady state capital by bi-section
          result_SS = self.SteadyState(P_H, grid, meshes, self.mpar, par)
          
          c_n_guess = result_SS['c_n_guess'].copy()
          m_n_star = result_SS['m_n_star'].copy()
          c_a_guess = result_SS['c_a_guess'].copy()
          m_a_star = result_SS['m_a_star'].copy()
          cap_a_star = result_SS['cap_a_star'].copy()
          psi_guess = result_SS['psi_guess'].copy()
          joint_distr = result_SS['joint_distr'].copy()
          
          R_fc = result_SS['R_fc']
          W_fc = result_SS['W_fc']
          Profits_fc = result_SS['Profits_fc']
          Output = result_SS['Output']
          grid = result_SS['grid'].copy()
          
          ## SS stats       
          mesh ={}
          mesh['m'],mesh['k'] =np.meshgrid(grid['m'].copy(),grid['k'].copy(), indexing = 'ij')
          
          targets = {}
          targets['ShareBorrower'] = np.sum((grid['m']<0)*np.transpose(np.sum(np.sum(joint_distr.copy(),axis = 1), axis = 1)))
          targets['K'] = np.sum(grid['k'].copy()*np.sum(np.sum(joint_distr.copy(),axis =0),axis=1))
          targets['B'] = np.dot(grid['m'].copy(),np.sum(np.sum(joint_distr.copy(),axis = 1),axis = 1))
          grid['K'] = targets['K']
          grid['B'] = targets['B']
  
          JDredux = np.sum(joint_distr.copy(),axis =2)
          targets['BoverK']     = targets['B']/targets['K']
  
          targets['L'] = grid['N']*np.sum(np.dot(grid['h'].copy(),np.sum(np.sum(joint_distr.copy(),axis=0),axis=0)))
          targets['KY'] = targets['K']/Output
          targets['BY'] = targets['B']/Output
          targets['Y'] = Output
          BCaux_M = np.sum(np.sum(joint_distr.copy(),axis =1), axis=1)
          targets['m_bc'] = BCaux_M[0].copy()
          targets['m_0'] = float(BCaux_M[grid['m']==0].copy())
          BCaux_K = np.sum(np.sum(joint_distr.copy(),axis=0),axis=1)
          targets['k_bc'] = BCaux_K[0].copy()
          aux_MK = np.sum(joint_distr.copy(),axis=2)
          
          targets['WtH_b0']=np.sum(aux_MK[(mesh['m']==0)*(mesh['k']>0)].copy())
          targets['WtH_bnonpos']=np.sum(aux_MK[(mesh['m']<=0)*(mesh['k']>0)].copy())
  
          targets['T'] =(1.0-par['tau'])*W_fc*grid['N'] +(1.0-par['tau'])*Profits_fc
          par['G']=targets['B']*(1.0-par['RB']/par['PI'])+targets['T']
          par['R']=R_fc
          par['W']=W_fc
          par['PROFITS']=Profits_fc
          par['N']=grid['N']
          targets['GtoY']=par['G']/Output
  
          ## Ginis
          # Net worth Gini
          mplusk=mesh['k'].copy().flatten('F')*par['Q']+mesh['m'].copy().flatten('F')
          
          IX = np.argsort(mplusk.copy())
          mplusk = mplusk[IX.copy()].copy()
          
          moneycapital_pdf   = JDredux.flatten(order='F')[IX].copy()
          moneycapital_cdf   = np.cumsum(moneycapital_pdf.copy())
          targets['NegNetWorth']= np.sum((mplusk.copy()<0)*moneycapital_pdf.copy())
  
          S                  = np.cumsum(moneycapital_pdf.copy()*mplusk.copy())
          
          S                  = np.concatenate(([0.], S.copy()))
          targets['GiniW']      = 1.0-(np.sum(moneycapital_pdf.copy()*(S[:-1].copy()+S[1:].copy()).transpose())/S[-1])
  
          # Liquid Gini
          IX = np.argsort(mesh['m'].copy().flatten('F'))
          liquid_sort = mesh['m'].copy().flatten('F')[IX.copy()].copy()
          liquid_pdf         = JDredux.flatten(order='F')[IX.copy()].copy()
          liquid_cdf         = np.cumsum(liquid_pdf.copy())
          targets['Negliquid']  = np.sum((liquid_sort.copy()<0)*liquid_pdf.copy())
  
          S                  = np.cumsum(liquid_pdf.copy()*liquid_sort.copy())
          S                  = np.concatenate(([0.], S.copy()))
          targets['GiniLI']      = 1.0-(np.sum(liquid_pdf.copy()*(S[:-1].copy()+S[1:].copy()))/S[-1].copy())
  
          # Illiquid Gini
          IX = np.argsort(mesh['k'].copy().flatten('F'))
          illiquid_sort = mesh['k'].copy().flatten('F')[IX.copy()].copy()
          illiquid_pdf        = JDredux.flatten(order='F')[IX.copy()].copy()
          illiquid_cdf        = np.cumsum(illiquid_pdf.copy());
          targets['Negliquid']   = np.sum((illiquid_sort.copy()<0)*illiquid_pdf.copy())
  
          S                   = np.cumsum(illiquid_pdf.copy()*illiquid_sort.copy())
          S                   = np.concatenate(([0.], S.copy()))
          targets['GiniIL']      = 1.-(np.sum(illiquid_pdf.copy()*(S[:-1].copy()+S[1:].copy()))/S[-1].copy())
  
          ##   MPCs
          meshesm, meshesk, meshesh =  np.meshgrid(grid['m'],grid['k'],grid['h'],indexing='ij')
  
          NW = par['gamma']/(1.+par['gamma'])*(par['N']/par['H'])*par['W']
          WW = NW*np.ones((self.mpar['nm'],self.mpar['nk'],self.mpar['nh']))  # Wages
          WW[:,:,-1]=par['PROFITS']*par['profitshare']
          # MPC
          WW_h=np.squeeze(WW[0,0,:].copy().flatten('F'))
          WW_h_mesh=np.squeeze(WW.copy()*meshes['h'].copy())
  
          grid_h_aux=grid['h']
  
          MPC_a_m = np.zeros((self.mpar['nm'],self.mpar['nk'],self.mpar['nh']))
          MPC_n_m = np.zeros((self.mpar['nm'],self.mpar['nk'],self.mpar['nh']))
  
          for kk in range(0 ,self.mpar['nk']) :
             for hh in range(0, self.mpar['nh']) :
              MPC_a_m[:,kk,hh]=np.gradient(np.squeeze(c_a_guess[:,kk,hh].copy()))/np.gradient(grid['m'].copy()).transpose()
              MPC_n_m[:,kk,hh]=np.gradient(np.squeeze(c_n_guess[:,kk,hh].copy()))/np.gradient(grid['m'].copy()).transpose()
      
  
          MPC_a_m = MPC_a_m.copy()*(WW_h_mesh.copy()/c_a_guess.copy())
          MPC_n_m = MPC_n_m.copy()*(WW_h_mesh.copy()/c_n_guess.copy())
  
          MPC_a_h = np.zeros((self.mpar['nm'],self.mpar['nk'],self.mpar['nh']))
          MPC_n_h = np.zeros((self.mpar['nm'],self.mpar['nk'],self.mpar['nh']))
  
          for mm in range(0, self.mpar['nm']) :
             for kk in range(0, self.mpar['nk']) :
               MPC_a_h[mm,kk,:] = np.gradient(np.squeeze(np.log(c_a_guess[mm,kk,:].copy())))/np.gradient(np.log(WW_h.copy().transpose()*grid_h_aux.copy())).transpose()
               MPC_n_h[mm,kk,:] = np.gradient(np.squeeze(np.log(c_n_guess[mm,kk,:].copy())))/np.gradient(np.log(WW_h.copy().transpose()*grid_h_aux.copy())).transpose()
      
  
          EMPC_h = np.dot(joint_distr.copy().flatten('F'),(par['nu']*MPC_a_h.copy().flatten('F')+(1.-par['nu'])*MPC_n_h.copy().flatten('F')))
          EMPC_m = np.dot(joint_distr.copy().flatten('F'),(par['nu']*MPC_a_m.copy().flatten('F')+(1.-par['nu'])*MPC_n_m.copy().flatten('F')))
  
          EMPC_a_h = np.dot(joint_distr.copy().flatten('F'), MPC_a_h.copy().flatten('F'))
          EMPC_a_m = np.dot(joint_distr.copy().flatten('F'), MPC_a_m.copy().flatten('F'))
  
          EMPC_n_h = np.dot(joint_distr.copy().flatten('F'), MPC_n_h.copy().flatten('F'))
          EMPC_n_m = np.dot(joint_distr.copy().flatten('F'), MPC_n_m.copy().flatten('F'))
  
          targets['Insurance_coeff']=np.concatenate((np.concatenate(([[1.-EMPC_h]], [[1.-EMPC_m]]), axis =1),
                                                  np.concatenate(([[1.-EMPC_a_h]],[[ 1.-EMPC_a_m]]), axis =1),
                                                  np.concatenate(([[1.-EMPC_n_h]], [[1.-EMPC_n_m]]), axis =1)) , axis =0)
                           
                           
                           
          
          ## Calculate Value Functions
          # Calculate Marginal Values of Capital (k) and Liquid Assets(m)
          RBRB = par['RB']/par['PI'] + (meshes['m']<0)*(par['borrwedge']/par['PI'])
  
          # Liquid Asset
          mutil_c_n = 1./(c_n_guess.copy()**par['xi']) # marginal utility at consumption policy no adjustment
          mutil_c_a = 1./(c_a_guess.copy()**par['xi']) # marginal utility at consumption policy adjustment
          mutil_c = par['nu']*mutil_c_a.copy() + (1-par['nu'])*mutil_c_n.copy() # Expected marginal utility at consumption policy (w &w/o adjustment)
          Vm = RBRB.copy()*mutil_c.copy()  # take return on money into account
          Vm = np.reshape(np.reshape(Vm.copy(),(self.mpar['nm']*self.mpar['nk'], self.mpar['nh']),order ='F'),(self.mpar['nm'],self.mpar['nk'], self.mpar['nh']),order ='F')
  
          # Capital
          Vk = par['nu']*(par['R']+par['Q'])*mutil_c_a.copy() + (1-par['nu'])*par['R']*mutil_c_n.copy() + (1-par['nu'])*psi_guess.copy() # Expected marginal utility at consumption policy (w &w/o adjustment)
          Vk = np.reshape(np.reshape(Vk.copy(),(self.mpar['nm']*self.mpar['nk'], self.mpar['nh']),order = 'F'),(self.mpar['nm'],self.mpar['nk'], self.mpar['nh']), order='F') 
  
          ## Produce non-parametric Copula
          cum_dist = np.cumsum(np.cumsum(np.cumsum(joint_distr.copy(), axis=0),axis=1),axis=2)
          marginal_m = np.cumsum(np.squeeze(np.sum(np.sum(joint_distr.copy(),axis=1),axis=1)))
          marginal_k = np.cumsum(np.squeeze(np.sum(np.sum(joint_distr.copy(),axis=0),axis=1)))
          marginal_h = np.cumsum(np.squeeze(np.sum(np.sum(joint_distr.copy(),axis=1),axis=0)))
   
              
     
          Cgridm, Cgridk, Cgridh = np.meshgrid(marginal_m.copy(),marginal_k.copy(),marginal_h.copy(),indexing='ij')
          Cpoints = np.concatenate(( [Cgridm.flatten(order='F')],[Cgridk.flatten(order='F')],[Cgridh.flatten(order='F')]),axis=0).T
          
          
          #Copula_aux = griddata((marginal_m.copy(),marginal_k.copy(),marginal_h.copy()),cum_dist.copy().transpose(),(points,points,points),method ='cubic')
          #Copula_aux = griddata(Cpoints,cum_dist.copy().flatten(order='F'),(points0,points1,points2))
          #Copula = RegularGridInterpolator((spm,spk,sph),np.reshape(Copula_aux,(200,200,20),order='F'),bounds_error = False, fill_value = None)
          Copula ={}
          Copula['grid'] = Cpoints.copy()
          Copula['value'] = cum_dist.flatten(order = 'F').copy()
  
          
         
          return {'par':par,
                  'mpar':self.mpar,
                  'grid':grid,
                  'Output':Output,
                  'targets':targets,
                  'Vm': Vm,
                  'Vk': Vk,
                  'joint_distr': joint_distr,
                  'Copula': Copula,
                  'c_n_guess': c_n_guess,
                  'c_a_guess': c_a_guess,
                  'psi_guess': psi_guess,
                  'm_n_star': m_n_star,
                  'm_a_star': m_a_star,
                  'cap_a_star':cap_a_star,
                  'mutil_c_n': mutil_c_n,
                  'mutil_c_a': mutil_c_a,
                  'mutil_c': mutil_c,
                  'P_H' : P_H
                  }
         
      def JDiteration(self, joint_distr, m_n_star, m_a_star, cap_a_star,P_H, par, mpar, grid):
          '''
          Iterates the joint distribution over m,k,h using a transition matrix
          obtained from the house distributing the households optimal choices. 
          It distributes off-grid policies to the nearest on grid values.
          
          parameters
          ------------
          m_a_star :np.array
              optimal m func
          m_n_star :np.array
              optimal m func
          cap_a_star :np.array
              optimal a func    
          P_H : np.array
              transition probability    
          par : dict
               parameters    
          mpar : dict
               parameters    
          grid : dict
               grids
               
          returns
          ------------
          joint_distr : np.array
              joint distribution of m and h
          
          '''
          ## Initialize matirces
                  
          weight11  = np.empty((mpar['nm']*mpar['nk'], mpar['nh'],mpar['nh']))
          weight12  = np.empty((mpar['nm']*mpar['nk'], mpar['nh'],mpar['nh']))
          weight21  = np.empty((mpar['nm']*mpar['nk'], mpar['nh'],mpar['nh']))
          weight22  = np.empty((mpar['nm']*mpar['nk'], mpar['nh'],mpar['nh']))
          
      
          # Find next smallest on-grid value for money and capital choices
          resultGWa = GenWeight(m_a_star, grid['m'])
          resultGWn = GenWeight(m_n_star, grid['m'])
          resultGWk = GenWeight(cap_a_star, grid['k'])
          Dist_m_a = resultGWa['weight'].copy()
          idm_a = resultGWa['index'].copy()
          Dist_m_n = resultGWn['weight'].copy()
          idm_n = resultGWn['index'].copy()
          Dist_k = resultGWk['weight'].copy()
          idk_a = resultGWk['index'].copy()
          idk_n = np.tile(np.ones((mpar['nm'],1))*np.arange(mpar['nk']),(1,1,mpar['nh']))
          
          ## Transition matrix adjustment case
          idm_a = np.tile(idm_a.copy().flatten(order='F'),(1, mpar['nh']))
          idk_a = np.tile(idk_a.copy().flatten(order='F'),(1, mpar['nh']))
          idh = np.kron(np.arange(mpar['nh']),np.ones((1,mpar['nm']*mpar['nk']*mpar['nh'])))
          
          idm_a = idm_a.copy().astype(int)
          idk_a = idk_a.copy().astype(int)
          idh = idh.copy().astype(int)
                  
          index11 = np.ravel_multi_index([idm_a.flatten(order='F'), idk_a.flatten(order='F'), idh.flatten(order='F')],(mpar['nm'],mpar['nk'],mpar['nh']),order='F')
          index12 = np.ravel_multi_index([idm_a.flatten(order='F'), idk_a.flatten(order='F')+1, idh.flatten(order='F')],(mpar['nm'],mpar['nk'],mpar['nh']),order='F')
          index21 = np.ravel_multi_index([idm_a.flatten(order='F')+1, idk_a.flatten(order='F'), idh.flatten(order='F')],(mpar['nm'],mpar['nk'],mpar['nh']),order='F')
          index22 = np.ravel_multi_index([idm_a.flatten(order='F')+1, idk_a.flatten(order='F')+1, idh.flatten(order='F')],(mpar['nm'],mpar['nk'],mpar['nh']),order='F')
          
          ## Policy Transition Matrix for no-adjustment case
                  
          weight13 = np.empty((mpar['nm']*mpar['nk'],mpar['nh'],mpar['nh']))
          weight23 = np.empty((mpar['nm']*mpar['nk'],mpar['nh'],mpar['nh']))
          
          idm_n = np.tile(idm_n.copy().flatten(order='F'),(1,mpar['nh']))
          idk_n = np.tile(idk_n.copy().flatten(order='F'),(1,mpar['nh']))
          
          idm_n = idm_n.copy().astype(int)
          idk_n = idk_n.copy().astype(int)
          
          index13 = np.ravel_multi_index([idm_n.flatten(order='F'), idk_n.flatten(order='F'), idh.flatten(order='F')],(mpar['nm'],mpar['nk'],mpar['nh']),order='F')
          index23 = np.ravel_multi_index([idm_n.flatten(order='F')+1, idk_n.flatten(order='F'), idh.flatten(order='F')],(mpar['nm'],mpar['nk'],mpar['nh']),order='F')
          
          
          for hh in range(mpar['nh']):
          
              # Corresponding weights
              weight21_aux = Dist_m_a[:,:,hh].copy()*(1.-Dist_k[:,:,hh].copy())
              weight11_aux = (1.-Dist_m_a[:,:,hh].copy())*(1-Dist_k[:,:,hh].copy())
              weight22_aux = Dist_m_a[:,:,hh].copy()*Dist_k[:,:,hh].copy()
              weight12_aux = (1.-Dist_m_a[:,:,hh].copy())*Dist_k[:,:,hh].copy()
              
              weight23_aux = Dist_m_n[:,:,hh].copy()
              weight13_aux = (1.-Dist_m_n[:,:,hh].copy())
      
              # Dimensions (mxk,h',h)   
              weight11[:,:,hh]=np.outer(weight11_aux.flatten(order='F'),P_H[hh,:].copy())
              weight12[:,:,hh]=np.outer(weight12_aux.flatten(order='F'),P_H[hh,:].copy())
              weight21[:,:,hh]=np.outer(weight21_aux.flatten(order='F'),P_H[hh,:].copy())
              weight22[:,:,hh]=np.outer(weight22_aux.flatten(order='F'),P_H[hh,:].copy())
              
              weight13[:,:,hh]=np.outer(weight13_aux.flatten(order='F'),P_H[hh,:].copy())
              weight23[:,:,hh]=np.outer(weight23_aux.flatten(order='F'),P_H[hh,:].copy())
              
          
          # Dimensions (m*k,h,h')
          weight11 = np.ndarray.transpose(weight11.copy(),(0,2,1))
          weight12 = np.ndarray.transpose(weight12.copy(),(0,2,1))
          weight21 = np.ndarray.transpose(weight21.copy(),(0,2,1))
          weight22 = np.ndarray.transpose(weight22.copy(),(0,2,1))
                  
          rowindex = np.tile(range(mpar['nm']*mpar['nk']*mpar['nh']),(1,4*mpar['nh']))
          
          
          H_a = sp.coo_matrix((np.concatenate((weight11.flatten(order='F'),weight21.flatten(order='F'),weight12.flatten(order='F'),weight22.flatten(order='F'))), 
                         (rowindex.flatten(order='F'), np.concatenate((index11.flatten(order='F'),index21.flatten(order='F'),index12.flatten(order='F'),index22.flatten(order='F'))))),
                         shape=(mpar['nm']*mpar['nk']*mpar['nh'], mpar['nm']*mpar['nk']*mpar['nh'])) # mu'(h',k'), a without interest
         
          
          weight13 = np.ndarray.transpose(weight13.copy(),(0,2,1))
          weight23 = np.ndarray.transpose(weight23.copy(),(0,2,1))
          
          rowindex = np.tile(range(mpar['nm']*mpar['nk']*mpar['nh']),(1,2*mpar['nh']))
          
          H_n = sp.coo_matrix((np.concatenate((weight13.flatten(order='F'),weight23.flatten(order='F'))), 
                         (rowindex.flatten(order='F'), np.concatenate((index13.flatten(order='F'),index23.flatten(order='F'))))),
                         shape=(mpar['nm']*mpar['nk']*mpar['nh'], mpar['nm']*mpar['nk']*mpar['nh'])) # mu'(h',k'), a without interest
          
          ## Joint transition matrix and transitions
          H = par['nu']*H_a.copy() +(1.-par['nu'])*H_n.copy()
          
          ## Joint transition matrix and transitions
          
          distJD = 9999.
          countJD = 1
          joint_distr = (joint_distr.copy().flatten(order='F')).T
          joint_distr_next = joint_distr.copy().dot(H.copy().todense())
          joint_distr_next = joint_distr_next.copy()/joint_distr_next.copy().sum(axis=1)
  
          distJD = np.max((np.abs(joint_distr_next.copy().flatten(order='F')-joint_distr.copy().flatten(order='F'))))
                            
          if distJD > 10**(-9):
              eigen, joint_distr = sp.linalg.eigs(H.transpose(), k=1, which='LM')
              joint_distr = joint_distr.copy().real  
              joint_distr = joint_distr.copy().transpose()/joint_distr.copy().sum()
  
          distJD = 9999.
          
          while (distJD > 10**(-14) or countJD<50) and countJD<10000:
          
              joint_distr_next = joint_distr.copy().dot(H.copy().todense())
              joint_distr_next = joint_distr_next.copy()/joint_distr_next.copy().sum(axis=1)
              
              distJD = np.max((np.abs(joint_distr_next.copy().flatten(order='F')-joint_distr.copy().flatten(order='F'))))
              
              countJD += 1
              joint_distr = joint_distr_next.copy()
                   
          joint_distr = np.array(joint_distr.copy())
          
          return {'joint_distr': joint_distr, 'distJD': distJD}
             
  
      def PoliciesSS(self, c_a_guess, c_n_guess, psi_guess, grid, inc, RR, RBRB, P, mpar, par, meshes):
          
          distC_n = 99999
          distPSI = distC_n
          distC_a = distPSI
          mutil_c_n = 1./(c_n_guess.copy()**par['xi']) # marginal utility at consumption policy no adjustment
          mutil_c_a = 1./(c_a_guess.copy()**par['xi']) # marginal utility at consumption policy adjustment
          mutil_c = par['nu']*mutil_c_a.copy() + (1-par['nu'])*mutil_c_n.copy() # Expected marginal utility at consumption policy (w &w/o adjustment)
  
          count=0
          while max(distC_n, distC_a, distPSI)>mpar['crit'] and count<100000:
          
          
              
              
              count=count+1
      
              # Step 1: Update policies for only money adjustment
              mutil_c=RBRB*mutil_c.copy() # take return on money into account
              aux=np.reshape(np.ndarray.transpose(mutil_c.copy(),(2, 0, 1)), (mpar['nh'], mpar['nm']*mpar['nk']),order='F')
              # form expectations
              EMU_aux = par['beta']*np.ndarray.transpose(np.reshape(P.copy().dot(aux.copy()),(mpar['nh'], mpar['nm'], mpar['nk']),order='F'),(1, 2, 0))
              
              c_n_aux = 1./(EMU_aux.copy()**(1./par['xi']))
              
              # Take borrowing constraint into account
              results_EGM_Step1_b=self.EGM_Step1_b(grid,inc,c_n_aux,mpar,par,meshes)
              c_n_new=results_EGM_Step1_b['c_update'].copy()
              m_n_star=results_EGM_Step1_b['m_update'].copy()
              
              m_n_star[m_n_star.copy()>grid['m'][-1]] = grid['m'][-1] # not extrapolation
                                      
              # Step 2: Find for every k on grid some off-grid m*(k')
              m_a_star_aux = self.EGM_Step2_SS(mutil_c_n,mutil_c_a, psi_guess, grid,P,RBRB,RR,par,mpar)
              m_a_star_aux = m_a_star_aux['mstar'].copy()
              
              # Step 3: Solve for initial resources / consumption in  adjustment case
              results_EGM_Step3 = self.EGM_Step3(EMU_aux,grid,inc,m_a_star_aux,c_n_aux,mpar,par)
              cons_list = results_EGM_Step3['cons_list']
              res_list  = results_EGM_Step3['res_list']
              mon_list = results_EGM_Step3['mon_list']
              cap_list = results_EGM_Step3['cap_list']
                                      
              
              # Step 4: Interpolate Consumption Policy
              results_EGM_Step4 = self.EGM_Step4( cons_list,res_list, mon_list,cap_list,inc,mpar,grid ) 
              c_a_new = results_EGM_Step4['c_a_new'].copy()
              m_a_star = results_EGM_Step4['m_a_star'].copy()
              cap_a_star = results_EGM_Step4['cap_a_star'].copy()
              
              # a = cap_a_star>grid['k'].T[-1]
              # log_index = indices(a, lambda x: x==1)
              
              cap_a_star[cap_a_star.copy()>grid['k'][-1]] = grid['k'][-1] # not extrapolation
              m_a_star[m_a_star.copy()>grid['m'][-1]] = grid['m'][-1] # not extrapolation
              
                          
              # Step 5: Update ~psi
              mutil_c_n = 1./(c_n_new.copy()**par['xi']) # marginal utility at consumption policy no adjustment
              mutil_c_a = 1./(c_a_new.copy()**par['xi']) # marginal utility at consumption policy adjustment
              mutil_c = par['nu']*mutil_c_a.copy() + (1-par['nu'])*mutil_c_n.copy() # Expected marginal utility at consumption policy (w &w/o adjustment)
              
              
              # VFI analogue in updating psi
              term1=((par['nu']* mutil_c_a.copy() *(par['Q'] + RR)) + ((1.-par['nu'])* mutil_c_n.copy()* RR) + (1.-par['nu'])* psi_guess.copy())
              
              aux = np.reshape(np.ndarray.transpose(term1.copy(),(2, 0, 1)),(mpar['nh'], mpar['nm']*mpar['nk']),order='F')
              E_rhs_psi = par['beta']*np.ndarray.transpose(np.reshape(P.copy().dot(aux.copy()),(mpar['nh'], mpar['nm'], mpar['nk']),order='F'),(1, 2, 0))
              
              E_rhs_psi=np.reshape( E_rhs_psi.copy(), (mpar['nm'], mpar['nk']*mpar['nh']), order='F' )
              m_n_star=np.reshape( m_n_star.copy(), (mpar['nm'], mpar['nk']*mpar['nh']), order='F' )
              
                 
              
              # Interpolation of psi-function at m*_n(m,k)
              index = np.digitize(m_n_star.copy(),grid['m'])-1 # find indexes on grid next smallest to optimal policy
              index[m_n_star <= grid['m'][0]] = 0 # if below minimum
              index[m_n_star >= grid['m'][-1]] = len(grid['m'])-2 # if above maximum
              
              step = np.squeeze(np.diff(grid['m'])) # Stepsize on grid
              s = (np.asmatrix(m_n_star.copy()) - np.squeeze(grid['m'].T[index]))/step[index] # Distance of optimal policy to next grid point
                                      
              aux_index = np.ones((mpar['nm'],1))*np.arange(0, mpar['nk']*mpar['nh'])*mpar['nm'] # aux for linear indexes
              E_rhs_psi = E_rhs_psi.flatten(order='F').copy()
              aux3 = E_rhs_psi[(index.flatten(order='F').copy()+aux_index.flatten(order='F').copy()).astype(int)] # calculate linear indexes
              
              psi_new = aux3.copy() + np.squeeze(np.asarray(s.flatten(order='F').copy()))*(E_rhs_psi[(index.flatten(order='F')+aux_index.flatten(order='F')).astype(int)+1].copy()-aux3.copy()) # linear interpolation
              psi_new = np.reshape( psi_new.copy(), (mpar['nm'], mpar['nk'], mpar['nh']), order='F' )
              m_n_star = np.reshape( m_n_star.copy(), (mpar['nm'], mpar['nk'], mpar['nh']), order='F' )
              distPSI = max( (abs(psi_guess.flatten(order='F').copy()-psi_new.flatten(order='F').copy())) )
              
                          
              # Step 6: Check convergence of policies
              distC_n = max( (abs(c_n_guess.flatten(order='F').copy()-c_n_new.flatten(order='F').copy())) )
              distC_a = max( (abs(c_a_guess.flatten(order='F').copy()-c_a_new.flatten(order='F').copy())) )
              
              # Update c policy guesses
              c_n_guess = c_n_new.copy()
              c_a_guess = c_a_new.copy()
              psi_guess = psi_new.copy()
              
             
          #distPOL=(distC_n, distC_a, distPSI)
          distPOL=np.array((distC_n.copy(), distC_a.copy(), distPSI.copy()))
          print(max(distC_n, distC_a, distPSI))
          print(count)
          return {'c_n_guess':c_n_new, 
                  'm_n_star':m_n_star,
                  'c_a_guess':c_a_new,
                  'm_a_star':m_a_star,
                  'cap_a_star':cap_a_star,
                  'psi_guess':psi_new,
                  'distPOL':distPOL}
      
      def EGM_Step1_b(self, grid,inc,c_n_aux,mpar,par,meshes):
          ## EGM_Step1_b computes the optimal consumption and corresponding optimal money
           # holdings in case the capital stock cannot be adjusted by taking the budget constraint into account.
           # c_update(m,k,h):    Update for consumption policy under no-adj.
           # m_update(m,k,h):    Update for money policy under no-adj.
          m_star_n = (c_n_aux.copy() + meshes['m'] - inc['labor'] - inc['rent'] - inc['profits'])
          m_star_n = (m_star_n.copy() < 0) * m_star_n.copy() / ((par['RB']+par['borrwedge'])/par['PI']) + (m_star_n.copy() >= 0) * m_star_n.copy()/(par['RB']/par['PI'])
  
          # Identify binding constraints
          binding_constraints = meshes['m'] < np.tile(m_star_n[0,:,:].copy(),(mpar['nm'], 1, 1))
  
          # Consumption when drawing assets m' to zero: Eat all Resources
          Resource = inc['labor'] + inc['rent'] + inc['money'] + inc['profits']
  
          ## Next step: Interpolate w_guess and c_guess from new k-grids
           # using c(s,h,k',K), k(s,h,k',K)
  
          m_star_n = np.reshape(m_star_n.copy(),(mpar['nm'], mpar['nk']*mpar['nh']), order='F')
          c_n_aux= np.reshape(c_n_aux.copy(),(mpar['nm'], mpar['nk']*mpar['nh']), order='F')
  
          # Interpolate grid.m and c_n_aux defined on m_star_n over grid.m
          # [c_update, m_update]=egm1b_aux_mex(grid.m,m_star_n,c_n_aux);
          c_update=np.zeros((mpar['nm'], mpar['nk']*mpar['nh']))
          m_update=np.zeros((mpar['nm'], mpar['nk']*mpar['nh']))
          
          for hh in range(mpar['nk']*mpar['nh']):
              
              Savings=interp1d(m_star_n[:,hh].copy(),grid['m'], fill_value='extrapolate') # generate savings function a(s,a*)=a'
              m_update[:,hh] = Savings(grid['m']) # Obtain m'(m,h) by Interpolation
              Consumption = interp1d(m_star_n[:,hh].copy(),c_n_aux[:,hh],fill_value='extrapolate') # generate consumption function c(s,a*(s,a'))
              c_update[:,hh] = Consumption(grid['m'])  # Obtain c(m,h) by interpolation (notice this is out of grid, used linear interpolation)
  
          c_update = np.reshape(c_update,(mpar['nm'], mpar['nk'], mpar['nh']), order='F')
          m_update = np.reshape(m_update,(mpar['nm'], mpar['nk'], mpar['nh']), order='F')
  
          c_update[binding_constraints] = Resource[binding_constraints].copy()-grid['m'].T[0]
          m_update[binding_constraints] = min(grid['m'].T)
          
          return {'c_update':c_update, 'm_update':m_update}
      
      
      def EGM_Step2_SS(self, mutil_c_n,mutil_c_a, psi_guess, grid,P,RBRB,RR,par,mpar):
          
          term1 = ((par['nu'] * mutil_c_a.copy() * (par['Q'] + RR))+((1-par['nu']) * mutil_c_n.copy() * RR)+(1-par['nu']) * psi_guess.copy())
          aux = np.reshape( np.ndarray.transpose(term1.copy(),(2, 0, 1)),(mpar['nh'], mpar['nm']*mpar['nk']), order='F' )
          #term1 = par['beta']*np.ndarray.transpose( np.reshape(np.array(np.matrix(P.copy())*np.matrix(aux.copy())),(mpar['nh'], mpar['nm'], mpar['nk']), order='F'), (1, 2, 0) )
          term1 = par['beta']*np.ndarray.transpose( np.reshape(P.copy().dot(aux.copy()),(mpar['nh'], mpar['nm'], mpar['nk']), order='F'), (1, 2, 0) )
          
          term2 = RBRB*( par['nu'] * mutil_c_a.copy() +(1-par['nu']) * mutil_c_n.copy() )
          aux = np.reshape( np.ndarray.transpose(term2.copy(), (2, 0, 1)), (mpar['nh'], mpar['nm']*mpar['nk']), order='F' )
          #term2 = par['beta']*np.ndarray.transpose( np.reshape(np.array(np.matrix(P.copy())*np.matrix(aux.copy())),(mpar['nh'], mpar['nm'], mpar['nk']), order='F'), (1, 2, 0) )
          term2 = par['beta']*np.ndarray.transpose( np.reshape(P.copy().dot(aux.copy()),(mpar['nh'], mpar['nm'], mpar['nk']), order='F'), (1, 2, 0) )
  
  
          # Equation (59) in Appedix B.4.
          E_return_diff=term1.copy()/par['Q']-term2.copy()
  
          # Find an m*_n for given k' that solves the difference equation (59)
          mstar = Fastroot(grid['m'], E_return_diff)
          mstar = np.maximum(mstar.copy(),grid['m'].T[0]) # Use non-negativity constraint and monotonicity
          mstar = np.minimum(mstar.copy(),grid['m'].T[-1]) # Do not allow for extrapolation
          mstar = np.reshape(mstar.copy(), (mpar['nk'], mpar['nh']), order='F')
          
          return {'mstar':mstar}
  
  #    xgrid=grid['m']
  #    fx = E_return_diff 
  # np.savetxt('fx.csv', fx, delimiter=',')
          
  
      
      def EGM_Step3(self, EMU,grid,inc,m_a_star,c_n_aux,mpar,par):
          
          # EGM_Step3 returns the resources (res_list), consumption (cons_list)
          # and money policy (mon_list) for given capital choice (cap_list).
          # For k'=0, there doesn't need to be a unique corresponding m*. We get a
          # list of consumption choices and resources for money choices m'<m* (mon_list) and cap
          # choices k'=0 (cap_list) and merge them with consumption choices and
          # resources that we obtain if capital constraint doesn't bind next period.
  
          # c_star: optimal consumption policy as function of k',h (both
          # constraints do not bind)
          # Resource: required resource for c_star
  
          # cons_list: optimal consumption policy if a) only k>=0 binds and b) both
          # constraints do not bind
  
          # res_list: Required resorces for cons_list
          # c_n_aux: consumption in t as function of t+1 grid (constrained version)
  
          # Constraints for money and capital are not binding
          EMU=np.reshape(EMU.copy(), (mpar['nm'], mpar['nk']*mpar['nh']),order='F')
          m_a_star=m_a_star.flatten(order='F').copy()
          # Interpolation of psi-function at m*_n(m,k)
          index = np.digitize(np.asarray(m_a_star.copy()),np.squeeze(grid['m']))-1 # find indexes on grid next smallest to optimal policy
          index[m_a_star<=grid['m'].T[0]] = 0 # if below minimum
          index[m_a_star>=grid['m'].T[-1]] = mpar['nm']-2 # if above maximum
          step = np.squeeze(np.diff(grid['m'])) # Stepsize on grid
          s = (np.asmatrix(m_a_star.T) - grid['m'].T[index].T)/step.T[index].T # Distance of optimal policy to next grid point
  
          aux_index=np.arange(0,(mpar['nk']*mpar['nh']),1)*mpar['nm'] # aux for linear indexes
          EMU=EMU.flatten(order='F').copy()
          aux3=EMU[index.flatten(order='F')+aux_index.flatten(order='F')].copy() # calculate linear indexes
  
          # Interpolate EMU(m',k',h') over m*_n(k'), m-dim is dropped
          EMU_star = ( aux3.copy() + np.asarray(s.copy())*np.asarray( EMU[index.flatten(order='F').copy() + aux_index.flatten(order='F').copy()+1].copy() - aux3.copy() ) ).T # linear interpolation
          
          c_star = 1./(EMU_star.copy()**(1/par['xi']))
          cap_expenditure = np.squeeze(inc['capital'][0,:,:])
          auxL = np.squeeze(inc['labor'][0,:,:])
          auxP = inc['profits']
          # Resources that lead to capital choice k'
          # = c + m*(k') + k' - w*h*N = value of todays cap and money holdings
          Resource = c_star.flatten(order='F').copy() + m_a_star.flatten(order='F').copy() + cap_expenditure.flatten(order='F').copy() - auxL.flatten(order='F').copy() - auxP
  
          c_star = np.reshape( c_star.copy(), (mpar['nk'], mpar['nh']), order='F' )
          Resource = np.reshape( Resource.copy(), (mpar['nk'], mpar['nh']), order='F' )
  
          # Money constraint is not binding, but capital constraint is binding
          m_a_star = np.reshape(m_a_star.copy(), (mpar['nk'], mpar['nh']), order='F')
          m_star_zero = np.squeeze(m_a_star[0,:].copy()) # Money holdings that correspond to k'=0:  m*(k=0)
  
          # Use consumption at k'=0 from constrained problem, when m' is on grid
          aux_c = np.reshape(c_n_aux[:,0,:].copy(), (mpar['nm'], mpar['nh']), order='F')
          aux_inc = np.reshape( inc['labor'][0,0,:].copy() + inc['profits'], (1, mpar['nh']), order='F' )
          cons_list = []
          res_list = []
          mon_list = []
          cap_list = []
  
          # j=0
          for j in range(mpar['nh']):
              # When choosing zero capital holdings, HHs might still want to choose money
              # holdings smaller than m*(k'=0)
              if m_star_zero[j] > grid['m'].T[0]:
                  # Calculate consumption policies, when HHs chooses money holdings
                  # lower than m*(k'=0) and capital holdings k'=0 and save them in cons_list
                  a = (grid['m']< m_star_zero[j]).astype(int).T
                  log_index = self.indices(a, lambda x: x==1)
                  # log_index = np.squeeze((np.cumsum(log_index.copy())-1)[:-1])
                  # aux_c is the consumption policy under no cap. adj. (fix kï¿½=0), for mï¿½<m_a*(k'=0)
                  c_k_cons=aux_c[log_index, j].T.copy()
                  cons_list.append(c_k_cons.copy()) # Consumption at k'=0, m'<m_a*(0)
                  # Required Resources: Money choice + Consumption - labor income
                  # Resources that lead to k'=0 and m'<m*(k'=0)
                 
                  res_list.append(np.squeeze(grid['m'].T[log_index]) + c_k_cons.copy() - aux_inc.T[j])
                  mon_list.append(np.squeeze(grid['m'].T[log_index]))
                  
                  #log_index = (grid['m']< m_star_zero[j]).astype(int)
                  cap_list.append(np.zeros((np.sum(a.copy()),1)))
  
  
          # Merge lists
          c_star = np.reshape(c_star.copy(),(mpar['nk'], mpar['nh']), order='F')
          m_a_star = np.reshape(m_a_star.copy(),(mpar['nk'], mpar['nh']), order='F')
          Resource = np.reshape(Resource.copy(),(mpar['nk'], mpar['nh']), order='F')
          
          cons_list_1=[]
          res_list_1=[]
          mon_list_1=[]
          cap_list_1=[]
          
          for j in range(mpar['nh']):
              if m_star_zero[j] > grid['m'].T[0]:
                 cons_list_1.append( np.vstack((np.asmatrix(cons_list[j]).T, np.asmatrix(c_star[:,j]).T)) )
                 res_list_1.append( np.vstack((np.asmatrix(res_list[j]).T, np.asmatrix(Resource[:,j]).T)) )
                 mon_list_1.append( np.vstack((np.asmatrix(mon_list[j]).T, np.asmatrix(m_a_star[:,j]).T)) )
                 cap_list_1.append( np.vstack((np.asmatrix(cap_list[j]), np.asmatrix(grid['k']).T)) )
              else:
                  cons_list_1.append(np.asmatrix(c_star[:,j]).T)
                  res_list_1.append(np.asmatrix(Resource[:,j]).T)
                  mon_list_1.append(np.asmatrix(m_a_star[:,j]).T)
                  cap_list_1.append( np.asmatrix(grid['k'].T))
                  
          return {'c_star': c_star, 'Resource': Resource, 'cons_list':cons_list_1, 'res_list':res_list_1, 'mon_list':mon_list_1, 'cap_list':cap_list_1}
  
  
      def indices(self, a, func):
          return [i for (i, val) in enumerate(a) if func(val)] 
  
              
      def EGM_Step4(self, cons_list, res_list, mon_list, cap_list, inc, mpar, grid ):
          # EGM_Step4 obtains consumption, money, and capital policy under adjustment.
          # The function uses the {(cons_list{j},res_list{j})} as measurement
          # points. The consumption function in (m,k) can be obtained from
          # interpolation by using the total resources available at (m,k): R(m,k)=qk+m/pi.
          # c_a_new(m,k,h): Update for consumption policy under adjustment
          # m_a_new(m,k,h): Update for money policy under adjustment
          # k_a_new(m,k,h): Update for capital policy under adjustment
          
          c_a_new=np.empty((mpar['nm']*mpar['nk'], mpar['nh']))
          m_a_new=np.empty((mpar['nm']*mpar['nk'], mpar['nh']))
          k_a_new=np.empty((mpar['nm']*mpar['nk'], mpar['nh']))
          Resource_grid=np.reshape(inc['capital']+inc['money']+inc['rent'], (mpar['nm']*mpar['nk'], mpar['nh']),order='F')
          labor_inc_grid=np.reshape(inc['labor'] + inc['profits'], (mpar['nm']*mpar['nk'], mpar['nh']), order='F')
  
          for j in range(mpar['nh']):
              a = (Resource_grid[:,j].copy() < res_list[j][0]).astype(int).T
              log_index = self.indices(a, lambda x: x==1)
              # when at most one constraint binds:
              #     [c_a_new(:,j), m_a_new(:,j),k_a_new(:,j)] = ...
              #         myinter1m_mex(res_list{j},Resource_grid(:,j),cons_list{j},mon_list{j},cap_list{j});
              
              cons = interp1d(np.squeeze(np.asarray(res_list[j].copy())), np.squeeze(np.asarray(cons_list[j].copy())), fill_value='extrapolate')
              c_a_new[:,j] = cons(Resource_grid[:,j])
              mon = interp1d(np.squeeze(np.asarray(res_list[j].copy())), np.squeeze(np.asarray(mon_list[j].copy())), fill_value='extrapolate')
              m_a_new[:,j] = mon(Resource_grid[:,j])
              cap = interp1d(np.squeeze(np.asarray(res_list[j].copy())), np.squeeze(np.asarray(cap_list[j].copy())), fill_value='extrapolate')
              k_a_new[:,j] = cap(Resource_grid[:,j])
      
              # Lowest value of res_list corresponds to m_a'=0 and k_a'=0.
              # Any resources on grid smaller then res_list imply that HHs consume all
              # resources plus income.
              # When both constraints are binding:
              c_a_new[log_index,j] = Resource_grid[log_index,j].copy() + labor_inc_grid[log_index,j].copy() - grid['m'].T[0]
              m_a_new[log_index,j] = grid['m'].T[0]
              k_a_new[log_index,j] = 0
  
          c_a_new = np.reshape(c_a_new.copy(),(mpar['nm'], mpar['nk'], mpar['nh']), order='F')
          k_a_new = np.reshape(k_a_new.copy(),(mpar['nm'], mpar['nk'], mpar['nh']), order='F')
          m_a_new = np.reshape(m_a_new.copy(),(mpar['nm'], mpar['nk'], mpar['nh']), order='F')
          
          return {'c_a_new':c_a_new, 'm_a_star':m_a_new, 'cap_a_star':k_a_new}
    
   
  
         
      def PolicyGuess(self, meshes, WW, RR, RBRB, par, mpar):
          '''
          policyguess returns autarky policy guesses (in the first period only):
          c_a_guess, c_n_guess, psi_guess
          as well as income matrices later on used in EGM: inc
    
          Consumption is compositite leisure and physical consumption (x_it) in the
          paper, therefore labor income is reduced by the fraction of leisure consumed
         
          parameters
          -----------
          meshes : dict
               meshes for m and h
          par : dict
               parameters
          mpar : dict
               parameters
          WW : np.array
              wage for each m and h
          RR : float    
              rental rate    
          RBRB : float    
              interest rate
              
          returns
          -----------
          c_a_guess : np.array
              guess for c func
          c_n_guess : np.array
          psi_guess : np.array
          inc : dict
              guess for incomes
          '''
          inc = { }
          
          inc['labor'] = par['tau']*WW.copy()*meshes['h'].copy()
          inc['rent']  = RR*meshes['k'].copy()
          inc['money'] = RBRB.copy()*meshes['m'].copy()
          inc['capital'] = par['Q']*meshes['k'].copy()
          inc['profits'] = 0. # lumpsum profits
          
          ## Initial policy guesses: Autarky policies as guess
          # consumption guess
          c_a_guess = inc['labor'].copy() + inc['rent'].copy() + inc['capital'].copy() + np.maximum(inc['money'].copy(),0.) + inc['profits']
          c_n_guess = inc['labor'].copy() + inc['rent'].copy() +                         np.maximum(inc['money'].copy(),0.) + inc['profits']
      
          # initially guessed marginal continuation value of holding capital
          psi_guess = np.zeros((mpar['nm'],mpar['nk'],mpar['nh']))            
              
          return {'c_a_guess':c_a_guess, 'c_n_guess':c_n_guess, 'psi_guess':psi_guess, 'inc': inc}       
           
      def FactorReturns(self, meshes, grid, par, mpar):
          '''
          return factors for steady state
          
          parameters
          -----------
          meshes : dict
               meshes for m and h
          par : dict
               parameters
          mpar : dict
               parameters
          grid : dict
               grids
               
          returns
          ----------
          N : float
             aggregate labor
          w : float
              wage   
          Profits_fc : float
              profit of firm
          WW : np.array
              wage for each m and h
          RBRB : float    
              interest rate
          '''
          ## GHH preferences
          mc = par['mu'] - ((par['beta'] - 1.)*np.log(par['PI']))/par['kappa']        
          N = (par['tau']*par['alpha']*grid['K']**(1.-par['alpha'])*mc)**(1./(1.-par['alpha']+par['gamma']))
          W_fc = par['alpha'] *mc *(grid['K']/N)**(1.-par['alpha'])
          
          # Before tax return on capital
          R_fc = (1.-par['alpha'])*mc *(N/grid['K'])**par['alpha'] - par['delta'] 
              
          Y = N**par['alpha']*grid['K']**(1.-par['alpha'])
          Profits_fc = (1.-mc)*Y - Y*(1./(1.-par['mu'])) /par['kappa'] /2. *np.log(par['PI'])**2.
      
          NW = par['gamma']/(1.+par['gamma'])*N/par['H'] *W_fc
      
          WW = NW*np.ones((mpar['nm'], mpar['nk'] ,mpar['nh'])) # Wages
          WW[:,:,-1] = Profits_fc * par['profitshare']
          RR = R_fc # Rental rates
          RBRB = (par['RB']+(meshes['m']<0)*par['borrwedge'])/par['PI']
      
          return {'N':N, 'R_fc':R_fc, 'W_fc':W_fc, 'Profits_fc':Profits_fc, 'WW':WW, 'RR':RR, 'RBRB':RBRB,'Y':Y}
  
  
      def StochasticsVariance(self, par, mpar, grid):
          '''
          generates transition probabilities for h: P_H
          
          parameters
          -------------
          par : dict
               parameters
          mpar : dict
               parameters
          grid : dict
               grids
               
          return
          -----------
          P_H : np.array
              transition probability
          grid : dict
              grid
          par : dict
              parameters
          '''
      
          # First for human capital
          TauchenResult = Tauchen(par['rhoH'], mpar['nh']-1, 1., 0., mpar['tauchen'])
          hgrid = TauchenResult['grid'].copy()
          P_H = TauchenResult['P'].copy()
          boundsH = TauchenResult['bounds'].copy()
      
          # correct long run variance for human capital
          hgrid = hgrid.copy()*par['sigmaH']/np.sqrt(1-par['rhoH']**2)
          hgrid = np.exp(hgrid.copy()) # levels instead of logs
      
          grid['h'] = np.concatenate((hgrid,[1]), axis=0)
          
          P_H = Transition(mpar['nh']-1, par['rhoH'], np.sqrt(1-par['rhoH']**2), boundsH)
      
          # Transitions to enterpreneur state
          P_H = np.concatenate((P_H.copy(),np.tile(mpar['in'],(mpar['nh']-1,1))), axis=1)
          lastrow = np.concatenate((np.tile(0.,(1,mpar['nh']-1)),[[1-mpar['out']]]), axis=1)
          lastrow[0,int(np.ceil(mpar['nh']/2))-1] = mpar['out'] 
          P_H = np.concatenate((P_H.copy(),lastrow), axis=0)
          P_H = P_H.copy()/np.transpose(np.tile(np.sum(P_H.copy(),1),(mpar['nh'],1)))
      
          Paux = np.linalg.matrix_power(P_H.copy(),1000)
          hh = Paux[0,:mpar['nh']-1].copy().dot(grid['h'][:mpar['nh']-1].copy())
      
          par['H'] = hh # Total employment
          par['profitshare'] = Paux[-1,-1]**(-1) # Profit per household
          grid['boundsH'] = boundsH.copy()
       
      
          return {'P_H': P_H, 'grid':grid, 'par':par}
      
      
      def SteadyState(self, P_H, grid, meshes, mpar, par):
          '''
          Prepare items for EGM
          Layout of matrices:
           Dimension 1: money m
           Dimension 2: capital k
           Dimension 3: stochastic sigma s x stochastic human capital h
           Dimension 4: aggregate money M
           Dimension 5: aggregate capital K
          '''
          
          # 1)  Construct relevant return matrices
          joint_distr = np.ones((mpar['nm'],mpar['nk'],mpar['nh']))/(mpar['nh']*mpar['nk']*mpar['nm'])
          fr_result = self.FactorReturns(meshes, grid, par, mpar)
          WW = fr_result['WW'].copy()
          RR = fr_result['RR']
          RBRB = fr_result['RBRB'].copy()
          
          # 2) Guess initial policies
          pg_result = self.PolicyGuess(meshes, WW, RR, RBRB, par, mpar)
          c_a_guess = pg_result['c_a_guess'].copy()
          c_n_guess = pg_result['c_n_guess'].copy()
          psi_guess = pg_result['psi_guess'].copy()
          
          KL = 0.5*grid['K']
          KH = 1.5*grid['K']
          
          r0_excessL = self.ExcessK(KL, c_a_guess, c_n_guess, psi_guess, joint_distr, grid, P_H, mpar, par, meshes)
          r0_excessH = self.ExcessK(KH, c_a_guess, c_n_guess, psi_guess, joint_distr, grid, P_H, mpar, par, meshes)
          excessL = r0_excessL['excess']
          excessH = r0_excessH['excess']
          
          if np.sign(excessL) == np.sign(excessH):
              print('ERROR! Sign not diff')
          
          ## Brent  
          fa = excessL
          fb = excessH
          a = KL
          b = KH
          
          if fa*fb > 0. :
              print('Error! f(a) and f(b) should have different signs')
  
          c = a
          fc = fa            
          d = b-1.0
          e = d
          
          
          iter = 0
          maxiter = 1000
          
          
          while iter<maxiter :
              iter += 1
              print(iter)
              
              if fb*fc>0:
                c = a
                fc = fa
                d = b-a
                e = d
      
      
              if np.abs(fc) < np.abs(fb) :
                a = b
                b = c
                c = a
                
                fa = fb
                fb = fc
                fc = fa
  
              eps = np.spacing(1)     
              tol = 2*eps*np.abs(b) + mpar['crit']
              m = (c-b)/2. # tolerance
      
              if (np.abs(m)>tol) and (np.abs(fb)>0) :
          
                if (np.abs(e)<tol) or (np.abs(fa)<=np.abs(fb)) :
                  d = m
                  e = m
                
                else:
                  s = fb/fa
                  
                  if a==c:
                     p=2*m*s
                     q=1-s
                  else :
                     q=fa/fc
                     r=fb/fc
                     p=s*(2*m*q*(q-r)-(b-a)*(r-1))
                     q=(q-1)*(r-1)*(s-1)
              
              
                  if p>0 :
                     q=-q
                  else :
                     p=-p
              
                  s=e
                  e=d
                   
                  if ( 2*p<3*m*q-np.abs(tol*q) ) and (p<np.abs(s*q/2)) :
                     d=p/q
                  else:
                     d=m
                     e=m
              
                a=b
                fa=fb
          
                if np.abs(d)>tol :
                  b=b+d
                else :
                  if m>0 :
                    b=b+tol
                  else:
                    b=b-tol
              
              else:
                break
              
              
              r_excessK = self.ExcessK(b,c_a_guess,c_n_guess,psi_guess,joint_distr, grid,P_H,mpar,par,meshes)
              fb = r_excessK['excess']
              c_n_guess = r_excessK['c_n_guess'].copy()
              c_a_guess = r_excessK['c_a_guess'].copy()
              psi_guess = r_excessK['psi_guess'].copy()
              joint_distr = r_excessK['joint_distr'].copy()
              
              
          Kcand = b
          grid['K'] = b
         
          ## Update
          r1_excessK=self.ExcessK(Kcand,c_a_guess,c_n_guess,psi_guess,joint_distr, grid,P_H,mpar,par,meshes)
          excess = r1_excessK['excess']
          N = r1_excessK['N']
         
          grid['N'] = N
          
          return {'c_n_guess':r1_excessK['c_n_guess'], 'm_n_star':r1_excessK['m_n_star'], 'c_a_guess':r1_excessK['c_a_guess'],
                  'm_a_star':r1_excessK['m_a_star'], 'cap_a_star':r1_excessK['cap_a_star'], 'psi_guess':r1_excessK['psi_guess'], 
                  'joint_distr':r1_excessK['joint_distr'],
                  'R_fc':r1_excessK['R_fc'], 'W_fc':r1_excessK['W_fc'], 'Profits_fc':r1_excessK['Profits_fc'], 
                  'Output':r1_excessK['Output'], 'grid':grid }
  
      def ExcessK(self,K,c_a_guess,c_n_guess,psi_guess,joint_distr, grid,P_H,mpar,par,meshes):
          
          grid['K'] = K
          fr_ek_result = self.FactorReturns(meshes, grid, par, mpar)
          N    = fr_ek_result['N']
          R_fc = fr_ek_result['R_fc']
          W_fc = fr_ek_result['W_fc']
          Profits_fc = fr_ek_result['Profits_fc']
          WW = fr_ek_result['WW'].copy()
          RR = fr_ek_result['RR']
          RBRB = fr_ek_result['RBRB'].copy()
          Output = fr_ek_result['Y']
                  
          pg_ek_result = self.PolicyGuess(meshes, WW, RR, RBRB, par, mpar)
          c_a_guess = pg_ek_result['c_a_guess'].copy()
          c_n_guess = pg_ek_result['c_n_guess'].copy()
          psi_guess = pg_ek_result['psi_guess'].copy()
          inc = pg_ek_result['inc'].copy()
          
                  
          # solve policies and joint distr
          print('Solving Household problem by EGM')
          start_time = time.clock()
                              
          pSS_ek_result = self.PoliciesSS(c_a_guess,c_n_guess,psi_guess, grid, inc, RR,RBRB,P_H,mpar,par,meshes)
          c_n_guess = pSS_ek_result['c_n_guess'].copy()
          m_n_star = pSS_ek_result['m_n_star'].copy()
          c_a_guess = pSS_ek_result['c_a_guess'].copy()
          m_a_star = pSS_ek_result['m_a_star'].copy()
          cap_a_star = pSS_ek_result['cap_a_star'].copy()
          psi_guess = pSS_ek_result['psi_guess'].copy()
          distPOL = pSS_ek_result['distPOL'].copy()
          
          
          print(distPOL)
          
          end_time = time.clock()
          print('Elapsed time is ',  (end_time-start_time), ' seconds.')
          
          print('Calc Joint Distr')
          start_time = time.clock()
          jd_ek_result = self.JDiteration(joint_distr, m_n_star, m_a_star, cap_a_star, P_H, par, mpar, grid)
          joint_distr = np.reshape(jd_ek_result['joint_distr'].copy(),(mpar['nm'],mpar['nk'],mpar['nh']),order='F')
          print(jd_ek_result['distJD'])
          
          end_time = time.clock()
          print('Elapsed time is ',  (end_time-start_time), ' seconds.')
          
          AggregateCapitalDemand = np.sum(grid['k']*np.sum(np.sum(joint_distr.copy(),axis = 0),axis=1))
          excess = grid['K']-AggregateCapitalDemand
          
          return{'excess':excess,'c_n_guess':c_n_guess,'m_n_star':m_n_star,'c_a_guess':c_a_guess,
                 'm_a_star':m_a_star,'cap_a_star':cap_a_star,'psi_guess':psi_guess,
                 'joint_distr':joint_distr, 'R_fc':R_fc,'W_fc':W_fc,'Profits_fc':Profits_fc,'Output':Output,'N':N}
  
  ###############################################################################
  
  if __name__ == '__main__':
      
      import defineSSParametersTwoAsset as Params
      from copy import copy
      import time
      import pickle
      
      EX3param = copy(Params.parm_TwoAsset)
      
      start_time0 = time.clock()        
      EX3 = SteadyStateTwoAsset(**EX3param)
      
      EX3SS = EX3.SolveSteadyState()
      
      end_time0 = time.clock()
      print('Elapsed time is ',  (end_time0-start_time0), ' seconds.')
      
      pickle.dump(EX3SS, open("EX3SS_20.p", "wb"))
      
      
      #EX3SS=pickle.load(open("EX3SS_ag.p", "rb"))