text: import numpy as np
  import numpy.linalg as la
  import numpy.random as npr
  import matplotlib.pyplot as plt
  import copy
  import os
  
  from problem_data_gen import gen_double_spring_mass, example_system_erdos_renyi
  from policy_iteration import policy_iteration, value_iteration, get_initial_gains, verify_gare
  from ltimult import gdlyap
  from matrixmath import mdot, specrad, dlyap, dare_gain, is_pos_def
  
  
  def set_numpy_decimal_places(places, width=0):
      set_np = '{0:' + str(width) + '.' + str(places) + 'f}'
      np.set_printoptions(formatter={'float': lambda x: set_np.format(x)})
  
  
  def get_problem_data(problem_type, problem_data_id=None, seed=None):
      """Get problem data_files"""
      from time import time
      from code_files.data_io import save_problem_data, load_problem_data
      from code_files.problem_data_gen import gen_rand_problem_data
  
      if problem_type == 'load':
          problem_data = load_problem_data(problem_data_id)
      elif problem_type == 'gen':
          problem_data = gen_rand_problem_data(n=2, m=1, p=1, rho=0.9, seed=seed)
          problem_data_id = int(time())
          save_problem_data(problem_data_id, problem_data)
  
      return problem_data
  
  
  def model_check(A, B, Q, R, K, op=True, model_type=None):
      """Check stability of system (A,B) with costs Q, R under the feedback gain K """
      # A, B, C = sample_ABCrand(problem_data_true)
      Qbar = Q + mdot(K.T, R, K)
      Abar = A + mdot(B,K)
      # print("Feedback gains test ")
      if model_type != None:
          print(model_type)
      if is_pos_def(dlyap(Abar,Qbar))==1: #if pos_def => gain stabilizes
          if op == True:
              print("Policy Iteration on model stabilizes the system - D-Lyap is Positive Definite.\n")
          ret = True
      else:
          if op == True:
              print("Policy Iteration on model does not stabilize the system - D-Lyap is NOT Positive Definite.\n")
          ret = False
      return ret
  
  
  def ctrb(A, B):
      """Controllabilty matrix
      Parameters
      ----------
      A, B: array_like
          Dynamics and input matrix of the system
      Returns
      -------
      C: matrix
          Controllability matrix
      Examples
      --------
      >>> C = ctrb(A, B)
      """
  
      n = np.shape(A)[0]
      C = np.hstack([B] + [np.dot(la.matrix_power(A, i), B) for i in range(1, n)])
      return C
  
  
  def mtpl(A, B, K_pi, Kn_pi, Km_pi, Q, R, init=1, multiplier_accuracy=10**(-8), multiplier_bound=10**(5), iteration_steps=10**10):
      """ Calculate a multiplying/scaling factor on the dynamics such that the system is not stabilized by gains from nominal game or LQR with multiplicative noise
       and only by gains from game with multiplicative noise """
      At_base = copy.deepcopy(A)
      Bt = copy.deepcopy(B)
      mpl = copy.deepcopy(init)
  
      n = A.shape[1]
      m = B.shape[1]
  
      #standard controllability check
      if la.matrix_rank(ctrb(At_base, Bt)) < len(At_base):
          print("Given A,B matrix pair for the true system is not controllable.")
          return 0
  
      iter_check = 0
      check_gains = 0
      mpl_his = 0
  
      while check_gains == 0:
          iter_check += 1
          At = copy.deepcopy(At_base)
          At *= mpl
          c_gm = model_check(At, Bt, K_pi, np.eye(n,n), np.eye(m,m), False)
          c_n = model_check(At, Bt, Kn_pi, np.eye(n,n), np.eye(m,m), False)
          c_m = model_check(At, Bt, Km_pi, np.eye(n,n), np.eye(m,m), False)
          # print(c_gm,":",c_n,":",c_m,":",mpl,"\n")
          if c_gm == 0 and c_n + c_m >0:
              check_gains = 2
              break
          elif c_gm ==1 and c_n + c_m ==0:
              check_gains = 1
              break
          elif abs(mpl-mpl_his)<multiplier_accuracy:
              check_gains = 3
              break
          elif mpl>multiplier_bound:
              check_gains = 4
              break
          elif iter_check>iteration_steps:
              check_gains = 5
              break
          else:
              mpl_his = mpl
              if c_gm == 1 and c_n + c_m > 0:
                  mpl = mpl*3/2
              elif c_gm + c_n + c_m == 0:
                  mpl = mpl/2
  
      print("\tMultiplier:",mpl)
      # print("check_gains:",check_gains)
      if check_gains == 1:
          print("\tGain of game w/ noise stabilizes while other two gains fail")
      elif check_gains == 2:
          print("\tGains of nominal game w/o m-noise or lqr w/ m-noise stabilize the true system, game w/ m-noise fails")
      elif check_gains ==3:
          print("\tFailed check. Multiplier out of threshold or too many iterations.")
      return mpl
  
  
  def model_based_robust_stabilization_experiment():
      seed = 1
      npr.seed(seed)
  
      problem_data_true, problem_data = gen_double_spring_mass()
  
      problem_data_keys = ['A', 'B', 'C', 'Ai', 'Bj', 'Ck', 'varAi', 'varBj', 'varCk', 'Q', 'R', 'S']
      A, B, C, Ai, Bj, Ck, varAi, varBj, varCk, Q, R, S = [problem_data[key] for key in problem_data_keys]
  
      n, m, p = [M.shape[1] for M in [A, B, C]]
      q, r, s = [M.shape[0] for M in [Ai, Bj, Ck]]
  
      # Synthesize controllers using various uncertainty modeling terms
  
      # Modify problem data_files
      # LQR w/ game adversary
      # Setting varAi, varBj, varCk = 0 => no multiplicative noise on the game
      problem_data_model_n = copy.deepcopy(problem_data)
      problem_data_model_n['varAi'] *= 0
      problem_data_model_n['varBj'] *= 0
      problem_data_model_n['varCk'] *= 0
  
      # LQR w/ multiplicative noise
      # Setting C = 0 and varCk = 0 => no game adversary
      problem_data_model_m = copy.deepcopy(problem_data)
      problem_data_model_m['C'] *= 0
      problem_data_model_m['varCk'] *= 0
  
      # Simulation options
      sim_options = None
      num_iterations = 50
      problem_data_known = True
  
      # Policy iteration on LQR w/ game adversary and multiplicative noise
      K0, L0 = get_initial_gains(problem_data, initial_gain_method='dare')
      print("LQR w/ game adversary and multiplicative noise")
      P_pi, K_pi, L_pi, H_pi, P_history_pi, K_history_pi, L_history_pi, c_history_pi, H_history_pi = policy_iteration(problem_data, problem_data_known, K0, L0, sim_options, num_iterations)
      verify_gare(problem_data, P_pi, algo_str='Policy iteration - Game w/ Multiplicative noise')
  
      # Check concavity condition
      Qvv_pi = -S + mdot(C.T, P_pi, C) + np.sum([varCk[k]*mdot(Ck[k].T, P_pi, Ck[k]) for k in range(s)], axis=0)
      if not is_pos_def(-Qvv_pi):
          raise Exception('Problem fails the concavity condition, adjust adversary strength')
  
      # Check positive definiteness condition
      QKL_pi = Q + mdot(K_pi.T, R, K_pi) - mdot(L_pi.T, S, L_pi)
      if not is_pos_def(QKL_pi):
          raise Exception('Problem fails the positive definiteness condition, adjust adversary strength')
      print(QKL_pi)
  
      # Policy Iteration on LQR w/ game adversary
      K0n, L0n = get_initial_gains(problem_data_model_n, initial_gain_method='dare')
      print("LQR w/ game adversary")
      Pn_pi, Kn_pi, Ln_pi, Hn_pi, Pn_history_pi, Kn_history_pi, Ln_history_pi, cn_history_pi, Hn_history_pi = policy_iteration(problem_data_model_n, problem_data_known, K0n, L0n, sim_options, num_iterations)
      verify_gare(problem_data_model_n, Pn_pi, algo_str='Policy iteration - Game w/o Multiplicative noise')
  
      # Policy Iteration on LQR w/ multiplicative noise
      K0m, L0m = get_initial_gains(problem_data_model_m, initial_gain_method='dare')
      print("LQR w/ multiplicative noise")
      Pm_pi, Km_pi, Lm_pi, Hm_pi, Pm_history_pi, Km_history_pi, Lm_history_pi, cm_history_pi, Hm_history_pi = policy_iteration(problem_data_model_m, problem_data_known, K0m, L0m, sim_options, num_iterations)
      verify_gare(problem_data_model_m, Pm_pi, algo_str='Policy iteration - LQR w/ Multiplicative noise')
  
      # LQR on true system
      A_true, B_true, Q_true, R_true = [problem_data_true[key] for key in ['A', 'B', 'Q', 'R']]
      n_true, m_true = [M.shape[1] for M in [A_true, B_true]]
      Pare_true, Kare_true = dare_gain(A_true, B_true, Q_true, R_true)
  
      # LQR on nominal system, no explicit robust control design
      Pce, Kce = dare_gain(A, B, Q, R)
  
      # Check if synthesized controllers stabilize the true system
      K_pi_true = np.hstack([K_pi, np.zeros([m, n])])
      Kn_pi_true = np.hstack([Kn_pi, np.zeros([m, n])])
      Km_pi_true = np.hstack([Km_pi, np.zeros([m, n])])
      Kce_true = np.hstack([Kce, np.zeros([m, n])])
      Kol_true = np.zeros_like(Kce_true)
  
      control_method_strings = ['open-loop     ', 'cert equiv    ', 'noise         ', 'game          ', 'noise + game  ', 'optimal       ']
      K_list = [Kol_true, Kce_true, Km_pi_true, Kn_pi_true, K_pi_true, Kare_true]
      AK_list = [A_true + np.dot(B_true, K) for K in K_list]
      QK_list = [Q_true + mdot(K.T, R_true, K) for K in K_list]
      specrad_list = [specrad(AK) for AK in AK_list]
      cost_list = [np.trace(dlyap(AK.T, QK)) if sr < 1 else np.inf for AK, QK, sr in zip(AK_list, QK_list, specrad_list)]
  
      set_numpy_decimal_places(1)
  
      output_text_filename = 'results_model_based_robust_stabilization_experiment.txt'
      output_text_path = os.path.join('..', 'results', output_text_filename)
      with open(output_text_path, 'w') as f:
          header_str = 'method       |  specrad  |   cost   |  gains'
          print(header_str, file=f)
          for control_method_string, sr, cost, K in zip(control_method_strings, specrad_list, cost_list, K_list):
              line_str = '%s  %.3f %8s      %s' % (control_method_string, sr, '%10.0f' % cost, K)
              print(line_str, file=f)
  
  
      # # Plot closed-loop response of true and nominal models
      # nt = 100
      # x_hist = np.zeros([nt, n])
      # x_hist_true = np.zeros([nt, n_true])
      # u0 = np.ones(m)
      #
      # for i in range(nt-1):
      #     if i == 0:
      #         u = u0
      #         u_true = u0
      #     else:
      #         u = np.dot(Kare_true[:,0:n], x_hist[i])
      #         u_true =  np.dot(Kare_true, x_hist_true[i])
      #
      #     x_hist[i+1] = np.dot(A, x_hist[i]) + np.dot(B, u)
      #     x_hist_true[i+1] = np.dot(A_true, x_hist_true[i]) + np.dot(B_true, u_true)
      #
      #
      #
      # fig, ax = plt.subplots(nrows=2)
      #
      # ax[0].plot(x_hist)
      # ax[1].plot(x_hist_true)
      #
      #
      # t_history = np.arange(num_iterations)+1
      #
      #
      # # Cost over iterations
      # fig, ax = plt.subplots()
      # ax.plot(t_history, c_history_pi)
      # ax.plot(t_history, cn_history_pi)
      # ax.plot(t_history, cm_history_pi)
      # plt.legend(['Combined', 'Adversary only', 'Mult. noise only'])
      # plt.xlabel('Iteration')
      # plt.ylabel('Cost')
      # if num_iterations <= 20:
      #     plt.xticks(np.arange(num_iterations)+1)
      # plt.show()
  
  
  def model_free_network_slq_game_experiment(seed=2):
      npr.seed(seed)
  
      problem_data = example_system_erdos_renyi(n=3, m=2, p=2, seed=seed)
      # from data_io import load_problem_data
      # data_files = load_problem_data(4)
  
      problem_data_keys = ['A', 'B', 'C', 'Ai', 'Bj', 'Ck', 'varAi', 'varBj', 'varCk', 'Q', 'R', 'S']
      A, B, C, Ai, Bj, Ck, varAi, varBj, varCk, Q, R, S = [problem_data[key] for key in problem_data_keys]
      n, m, p = [M.shape[1] for M in [A, B, C]]
      q, r, s = [M.shape[0] for M in [Ai, Bj, Ck]]
  
      # Initial gains
      K0, L0 = get_initial_gains(problem_data, initial_gain_method='zero')
  
      # Simulation options
      # Std deviation for initial state, defender inputs, attacker inputs, and additive noise
      xstd, ustd, vstd, wstd = 1.0, 1.0, 1.0, 0.0
  
      # Rollout length
      nt = 5000
  
      # Number of rollouts
      nr = 1
  
      # Rollout computation type
      group_option = 'group'
  
      # Q-function estimation scheme
      # qfun_estimator = 'direct'
      # qfun_estimator = 'lsadp'
      qfun_estimator = 'lstdq'
  
      sim_options_keys = ['xstd', 'ustd', 'vstd', 'wstd', 'nt', 'nr', 'group_option', 'qfun_estimator']
      sim_options_values = [xstd, ustd, vstd, wstd, nt, nr, group_option, qfun_estimator]
      sim_options = dict(zip(sim_options_keys, sim_options_values))
  
      num_iterations = 10
  
      # Policy iteration
      problem_data_known = True
      all_data = policy_iteration(problem_data, problem_data_known, K0, L0, sim_options, num_iterations)
      P, K, L, H, P_history, K_history, L_history, c_history, H_history = all_data
      verify_gare(problem_data, P, algo_str='Policy iteration')
  
      all_data_list = []
      problem_data_known = False
      num_trials = 10
      for i in range(num_trials):
          all_data_list.append(policy_iteration(problem_data, problem_data_known, K0, L0, sim_options, num_iterations))
  
      def norm_history_plotter(ax, M_history, M_history_list, ylabel_str):
          # Plot the history of the error norm from the single trial in the model-based case
          ax.plot(la.norm(M_history-M_history[-1], ord=2, axis=(1, 2)))
  
          # Plot the history of the error norm from each of the individual trials in the model-free case
          for i in range(num_trials):
              ax.plot(la.norm(M_history_list[i]-M_history[-1], ord=2, axis=(1, 2)), color='k', alpha=0.5)
          ax.set_ylabel(ylabel_str, rotation=0, labelpad=10)
  
      fig, ax = plt.subplots(nrows=4, figsize=(6, 8))
      data_idx_list = [4, 5, 6, 8]
      ylabel_list = ['P', 'K', 'L', 'H']
      for i in range(4):
          j = data_idx_list[i]
          norm_history_plotter(ax[i], all_data[j], [all_data_list[i][j] for i in range(num_trials)], ylabel_list[i])
          ax[i].set_yscale('log')
          ax[i].legend(['Model-based', 'Model-free'])
      # plt.tight_layout()
      ax[0].set_title('Relative norm of error vs iteration')
      plt.xlabel('Iteration')
      plt.show()
      figure_out_path = os.path.join('..', 'results', 'results_model_free_network_slq_game_experiment.png')
      plt.savefig(figure_out_path, dpi=300)
  
  
  def generic_experiment(seed=1):
      # Model-based or model-free solution to the linear-quadratic game with stochastic parameters for generic systems
      npr.seed(seed)
  
      # problem_data_id = 1581199445 # 5-state random system
      problem_data_id = 1581378883 # 2-state example system
      # problem_data_id = 3 # 3-state example system
      # problem_data_id = 2 # 2-state example system
      problem_data = get_problem_data(problem_type='load', problem_data_id=problem_data_id)
      problem_data_keys = ['A', 'B', 'C', 'Ai', 'Bj', 'Ck', 'varAi', 'varBj', 'varCk', 'Q', 'R', 'S']
      A, B, C, Ai, Bj, Ck, varAi, varBj, varCk, Q, R, S = [problem_data[key] for key in problem_data_keys]
      n, m, p = [M.shape[1] for M in [A, B, C]]
      q, r, s = [M.shape[0] for M in [Ai, Bj, Ck]]
  
      # Initial gains
      K0, L0 = get_initial_gains(problem_data, initial_gain_method='zero')
  
      # Simulation options
      # Std deviation for initial state, defender inputs, attacker inputs, and additive noise
      xstd, ustd, vstd, wstd = 1.0, 1.0, 1.0, 0.0
  
      # Rollout length
      nt = 4000
  
      # Number of rollouts
      nr = 1
  
      # Rollout computation type
      group_option = 'group'
  
      # Q-function estimation scheme
      # qfun_estimator = 'direct'
      # qfun_estimator = 'lsadp'
      qfun_estimator = 'lstdq'
  
      sim_options_keys = ['xstd', 'ustd', 'vstd', 'wstd', 'nt', 'nr', 'group_option', 'qfun_estimator']
      sim_options_values = [xstd, ustd, vstd, wstd, nt, nr, group_option, qfun_estimator]
      sim_options = dict(zip(sim_options_keys, sim_options_values))
  
      num_iterations = 20
  
      # Policy iteration
      problem_data_known = False
      P_pi, K_pi, L_pi, H_pi, P_history_pi, K_history_pi, L_history_pi, c_history_pi, H_history_pi = policy_iteration(problem_data, problem_data_known, K0, L0, sim_options, num_iterations)
      verify_gare(problem_data, P_pi, algo_str='Policy iteration')
  
      # Value iteration
      # Start value iteration at the same initial P as from policy iteration
      P0 = gdlyap(problem_data, K0, L0)
      P_vi, K_vi, L_vi, P_history_vi, c_history_vi = value_iteration(problem_data, P0, num_iterations)
      verify_gare(problem_data, P_vi, algo_str='Value iteration')
  
      # Plotting
      plt.close('all')
      t_history = np.arange(num_iterations)+1
  
      # Cost-to-go matrix
      fig, ax = plt.subplots(ncols=2)
      plt.suptitle('Value matrix (P)')
      ax[0].imshow(P_pi)
      ax[1].imshow(P_vi)
      ax[0].set_title('Policy iteration')
      ax[1].set_title('Value iteration')
  
      # Cost over iterations
      fig, ax = plt.subplots()
      ax.plot(t_history, c_history_pi)
      ax.plot(t_history, c_history_vi)
      plt.legend(['Policy iteration', 'Value iteration'])
      plt.xlabel('Iteration')
      plt.ylabel('Cost')
      if num_iterations <= 20:
          plt.xticks(np.arange(num_iterations)+1)
  
      plt.show()
  
  
  if __name__ == "__main__":
      # Model-based robust stabilization of the double-mass-spring system described in the paper
      # "Policy iteration for linear quadratic games with stochastic parameters"
      model_based_robust_stabilization_experiment()
  
      # Model-free solution to the linear-quadratic game with stochastic parameters of a diffusion network
      # (not reported explicitly in the paper)
      model_free_network_slq_game_experiment()
  