text: """
  Filename: oligopoly.py
  Authors: Chase Coleman
  
  This is an example for the lecture dyn_stack.rst from the QuantEcon
  series of lectures by Tom Sargent and John Stachurski.
  
  We deal with a large monopolistic firm who faces costs:
  
  C_t = e Q_t + .5 g Q_t^2 + .5 c (Q_{t+1} - Q_t)^2
  
  where the fringe firms face:
  
  sigma_t = d q_t + .5 h q_t^2 + .5 c (q_{t+1} - q_t)^2
  
  Additionally, there is a linear inverse demand curve of the form:
  
  p_t = A_0 - A_1 (Q_t + \bar{q_t}) + \eta_t,
  
  where:
  
  .. math
      \eta_{t+1} = ρ \eta_t + C_{\varepsilon} \varepsilon_{t+1};
      \varepsilon_{t+1} ∼ N(0, 1)
  
  For more details, see the lecture.
  """
  import numpy as np
  import scipy.linalg as la
  from quantecon import LQ
  from quantecon.matrix_eqn import solve_discrete_lyapunov
  from scipy.optimize import root
  
  
  def setup_matrices(params):
      """
      This function sets up the A, B, R, Q for the oligopoly problem
      described in the lecture.
  
      Parameters
      ----------
      params : Array(Float, ndim=1)
          Contains the parameters that describe the problem in the order
          [a0, a1, rho, c_eps, c, d, e, g, h, beta]
  
      Returns
      -------
      (A, B, Q, R) : Array(Float, ndim=2)
          These matrices describe the oligopoly problem.
      """
  
      # Left hand side of (37)
      Alhs = np.eye(5)
      Alhs[4, :] = np.array([a0-d, 1., -a1, -a1-h, c])
      Alhsinv = la.inv(Alhs)
  
      # Right hand side of (37)
      Brhs = np.array([[0., 0., 1., 0., 0.]]).T
      Arhs = np.eye(5)
      Arhs[1, 1] = rho
      Arhs[3, 4] = 1.
      Arhs[4, 4] = c / beta
  
      # R from equation (40)
      R = np.array([[0., 0., (a0-e)/2., 0., 0.],
                    [0., 0., 1./2., 0., 0.],
                    [(a0-e)/2., 1./2, -a1 - .5*g, -a1/2, 0.],
                    [0., 0., -a1/2, 0., 0.],
                    [0., 0., 0., 0., 0.]])
      Q = np.array([[c/2]])
  
      A = Alhsinv.dot(Arhs)
      B = Alhsinv.dot(Brhs)
  
      return A, B, Q, R
  
  
  def find_PFd(A, B, Q, R, beta=.95):
      """
      Taking the parameters A, B, Q, R as found in the `setup_matrices`,
      we find the value function of the optimal linear regulator problem.
      This is steps 2 and 3 in the lecture notes.
  
      Parameters
      ----------
      (A, B, Q, R) : Array(Float, ndim=2)
          The matrices that describe the oligopoly problem
  
      Returns
      -------
      (P, F, d) : Array(Float, ndim=2)
          The matrix that describes the value function of the optimal
          linear regulator problem.
  
      """
  
      lq = LQ(Q, -R, A, B, beta=beta)
      P, F, d = lq.stationary_values()
  
      return P, F, d
  
  
  def solve_for_opt_policy(params, eta0=0., Q0=0., q0=0.):
      """
      Taking the parameters as given, solve for the optimal decision rules
      for the firm.
  
      Parameters
      ----------
      params : Array(Float, ndim=1)
          This holds all of the model parameters in an array
  
      Returns
      -------
      out :
  
      """
      # Step 1/2: Formulate/Solve the optimal linear regulator
      (A, B, Q, R) = setup_matrices(params)
      (P, F, d) = find_PFd(A, B, Q, R, beta=beta)
  
      # Step 3: Convert implementation into state variables (Find coeffs)
      P22 = P[-1, -1]
      P21 = P[-1, :-1]
      P22inv = P22**(-1)
      dotmat = np.empty((5, 5))
      upper = np.eye(4, 5)  # Gives me 4x4 identity with a column of 0s
      lower = np.hstack([-P22inv*P21, P22inv])
      dotmat[:-1, :] = upper
      dotmat[-1, :] = lower
  
      coeffs = np.dot(-F, dotmat)
  
      # Step 4: Find optimal x_0 and \mu_{x, 0}
      z0 = np.array([1., eta0, Q0, q0])
      x0 = -P22inv*np.dot(P21, z0)
  
      # Do some rearranging for convenient representation of policy
      # TODO: Finish getting the equations into the from
      # u_t = rho u_{t-1} + gamma_1 z_t + gamma_2 z_{t-1}
  
      part1 = np.vstack([np.eye(4, 5), P[-1, :]])
      part2 = A - np.dot(B, F)
      part3 = dotmat
      m = np.dot(part1, part2).dot(part3)
      m12 = m[-1, :-1]
      m22 = m[-1, -1]
  
      f = np.dot(-F, dotmat)
      f11 = f[-1, :-1]
      f12 = f[-1, -1]
  
      coeff_utm1 = f12*m22*f12**(-1)
      coeff_zt = coeffs[0, :-1]
      coeff_ztm1 = f12*(m12 - f12**(-1)*m22*f11)
  
      return coeffs, x0, (coeff_utm1, coeff_zt, coeff_ztm1)
  
  
  # Parameter values
  a0 = 100.
  a1 = 1.
  rho = .8
  c_eps = .2
  c = 1.
  d = 20.
  e = 20.
  g = .2
  h = .2
  beta = .95
  params = np.array([a0, a1, rho, c_eps, c, d, e, g, h, beta])
  
  
  coefficients, x0, alt_coeffs = solve_for_opt_policy(params)
  
  print("The original coefficients are")
  print("u_t = {} [z_t mu_[x, t]]'".format(coefficients))
  print("or in other terms")
  print("u_t = {} u_[t-1] \n + {} z_t \n + {} z_[t-1]".format(alt_coeffs[0],
                                                              alt_coeffs[1],
                                                              alt_coeffs[2]))