text: import numpy as np
  import numpy.random as random
  from utils import Struct, ToTensor
  from rewardmodel import BaseRewardModel,DynamicRewardModel
  import torch
  import torch.optim as optim
  from torch.distributions.categorical import Categorical
  
  class MCEM(object):
  
      def __init__(
      self, 
      game,
      true_rewards,
      model, 
      linmodel_solutions,
      example_samples,
      n_samples, 
      clusters,
      alpha,
      lr=0.001
      ):
  
          self.game = game
          self.true_rewards = true_rewards
          self.model = model
          self.mdpmodel_solutions = linmodel_solutions
          self.example_samples = example_samples
          self.n_samples = n_samples
          self.clusters = clusters
          self.alpha = alpha
          self.lr = lr
          self.mu_s, self.mu_sa, self.initD, self.F, self.P_k_0 = self.initial_SVF(clusters)
          
      def momaxentrun(self, maxIter = 500):
  
          K = self.clusters
          M, T = np.shape(self.example_samples)[0:2]
          actions = self.game.actions
          states = self.game.states
          features = np.shape(self.F)[1]
          baserewardmodel = BaseRewardModel(features)
          baserewardmodel.train()
          # print('reward model:')
          # print(baserewardmodel)
          baseoptimizer = optim.Adam(baserewardmodel.parameters(), lr=0.001)
          dynets = []
          dyoptims = []
          for k in range(K):
              dynet = DynamicRewardModel()
              dynet.train()
              dynets.append(dynet)
              dyoptims.append(optim.Adam(dynet.parameters(), lr=self.lr)) 
  
          
  
  
          P_k_true = []
          for i in range(len(self.true_rewards)):
              P = np.zeros((self.n_samples[0],len(self.true_rewards)))
              P[:,i] = 1
              P_k_true.append(P)
          P_k_true_reward = np.concatenate(P_k_true, axis=0)
          P_k_true_policy = np.copy(P_k_true_reward)
          EV_n_true,_ = self.Espected_value_per_n(self.mdpmodel_solutions, P_k_true_reward, P_k_true_policy)
  
          min_EVD = np.inf
          min_r = 0
          min_r_seq = np.arange(K)
          
  
          for i in range(maxIter):
  
  
              
              reward_f = baserewardmodel(ToTensor(self.F))
              policies = []
              logpolicies = []
              rewards = []
              D_k = []
              solutions = []
              
  
              for k in range(self.clusters):
                  dynet = dynets[k]
                  r_k = dynet(reward_f)
                  _, _, policy, logpolicy = self.model.linearvalueiteration(r_k.detach().numpy())
                  policies.append(policy)
                  logpolicies.append(logpolicy)
                  rewards.append(r_k)
  
              
              M_k, mu_s_k, mu_sa_k, initD_k,P_k, rewards, policies, logpolicies, dynets, dyoptims = self.posterior_cluster(
                  rewards, policies, logpolicies, self.P_k_0, dynets, dyoptims, reward_f
              )
              self.clusters = len(policies)
              self.P_k_0 = P_k
              st_logpolicies = np.stack(logpolicies,axis=2)
              val = np.sum(np.expand_dims(np.sum(st_logpolicies*mu_sa_k, axis=(0,1)), axis=0)/(M*T))
  
              baseoptimizer.zero_grad()
              for k in range(self.clusters):
                  curr_optim = dyoptims[k]
                  curr_optim.zero_grad()
                  solution = Struct()
                  solution.p = policies[k]
  
                  curr_initD = np.expand_dims(initD_k[:,k], axis=1)
                  curr_D = self.model.linearmdpfrequency(solution, curr_initD)
                  D_k.append(curr_D)
                  solutions.append(solution)
  
              D_k = np.concatenate(D_k, axis=1)
              rewards = torch.cat(rewards, axis=1)
              flat_rewards = rewards.view(-1,1)
              grad = -ToTensor(mu_s_k - D_k).view(-1,1)
              flat_rewards.backward(grad)
              for k in range(self.clusters):
                  curr_optim = dyoptims[k]
                  curr_optim.step()
  
              baseoptimizer.step()
  
              EV_n_pred, r_seq = self.Espected_value_per_n(solutions, P_k_true_reward, P_k)
              EVD = np.mean(abs(EV_n_true - EV_n_pred))
              Loss = -val
              Est_phi = M_k/M
                  
              if EVD < min_EVD:
                  min_EVD = EVD
                  min_r = rewards.detach().numpy()
                  min_r_seq = r_seq
                  min_baserewardmodel = baserewardmodel
                  min_dynets = list.copy(dynets)
                  min_clusters = self.clusters
                  min_Est_phi = Est_phi
              if (i+1) % 50 == 0:
                  print('{:03d}/{:03d}    |   loss = {:.5f}    |   EVD = {:.5f}    |   min_EVD = {:.5f}'.format(
                      i+1,maxIter,Loss.item(),EVD.item(),min_EVD
                  ))
                  print(Est_phi)
  
  
          irl_solutions = []
          R = min_r
          for k in range(min_clusters):
              r_k = np.expand_dims(R[:,k], axis=1)
              r = np.tile(r_k, (1, actions))
              soln = self.model.linearmdpsolve(r)
              v = soln.v
              q = soln.q
              p = soln.p
  
              irl_solution = Struct()
              irl_solution.r = r
              irl_solution.v = v
              irl_solution.q = q
              irl_solution.p = p
              irl_solution.baserewardmodel = min_baserewardmodel
              irl_solution.dynet = min_dynets[k]
              irl_solution.min_Est_phi = min_Est_phi
              irl_solutions.append(irl_solution)
  
          return irl_solutions, min_EVD, min_r_seq
  
  
      def initial_SVF(self,clusters):
  
          sa_s = self.game.sa_s
          sa_p = self.game.sa_p
          discount = self.game.discount
          states, actions, transitions = np.shape(sa_p)
          F = self.game.feature_data.splittable
          F = np.concatenate((F, np.ones((states,1))), axis=1)
          N = len(self.example_samples)
          T = len(self.example_samples[0])
          mu_s = np.zeros((states,N))
          mu_sa = np.zeros((states, actions, N))
          for n in range(N):
              for t in range(T):
                  s = self.example_samples[n][t][0]
                  a = self.example_samples[n][t][1]
                  mu_s[s,n] += 1
                  mu_sa[s,a,n] += 1
          initD = np.copy(mu_s)
          for n in range(N):
              for t in range(T):
                  s = self.example_samples[n][t][0]
                  a = self.example_samples[n][t][1]
                  for k in range(transitions):
                      sp = sa_s[s,a,k]
                      ap = sa_p[s,a,k]
                      initD[sp, n] += -discount*ap
  
          K = self.clusters       
          P_k_0 = np.zeros((N,K))
          for n in range (N):
              r = random.randint(K)
              P_k_0[n,r] = 1
          return mu_s, mu_sa, initD, F, P_k_0
  
  
      def posterior_cluster(
          self, rewards, policies, logpolicies, P_k_0, dynets, dyoptims, reward_f
      ):
          
          N = len(self.example_samples)
          T = len(self.example_samples[0])
          states = self.game.states
          actions = self.game.actions
          
          for n in range(N):
              prior_n, rewards, policies, logpolicies, dynets, dyoptims = self.prior_cluster(
                  P_k_0, rewards, policies, logpolicies, dynets,dyoptims,n, reward_f
              )
              K = len(policies)
              P_tau_n = np.ones((1,K))
              st_policies = np.stack(policies,axis=2)
              for k in range(K):
                  # if prior_n[k] == 0:
                  #     P_tau_n[0,k] = 0
                  #     continue 
                  for t in range(T):
                      s = self.example_samples[n][t][0]
                      a = self.example_samples[n][t][1]
                      P_tau_n[0,k] *= st_policies[s, a, k]
  
              
              
              
              prior_catdis = Categorical(torch.from_numpy(np.transpose(prior_n)))
              new_cluster_n = prior_catdis.sample().numpy()
              old_cluster_n = np.where(P_k_0[n,:] == 1)[0]
              old_P_tau_n = P_tau_n[0,old_cluster_n]
              new_P_tau_n = P_tau_n[0,new_cluster_n]
              replacement_prob = np.min((1, new_P_tau_n/old_P_tau_n))
              if random.rand() <= replacement_prob:
                  eta_n = new_cluster_n
              else:
                  eta_n = old_cluster_n
              
              # P_k_n = np.transpose(prior_n)*P_tau_n
              # sum_k = np.expand_dims(np.sum(P_k_n, axis=1), axis=1)
              # sum_k = np.tile(sum_k, (1,K))
              # P_k_n = P_k_n/sum_k
              # catdis = Categorical(torch.from_numpy(P_k_n))
              # eta_n = catdis.sample()
              
              P_k_0, rewards, policies, logpolicies, dynets, dyoptims = self.update_P_k(
                  P_k_0, eta_n, n, rewards, policies, logpolicies, dynets, dyoptims
              )
          K = len(policies)
          P_k = P_k_0
          M_k = np.expand_dims(np.sum(P_k, axis=0), axis=1)
  
          temp1 = np.tile(np.expand_dims(self.mu_s, axis=2), (1,1,K))
          temp2 = np.tile(np.expand_dims(P_k, axis=0), (states,1,1))
          mu_s_k = temp1*temp2
          mu_s_k = np.sum(mu_s_k, axis=1)
  
          temp1 = np.tile(np.expand_dims(self.initD, axis=2), (1,1,K))
          temp2 = np.tile(np.expand_dims(P_k, axis=0), (states,1,1))
          initD_k = temp1*temp2
          initD_k = np.sum(initD_k, axis=1)
  
  
          temp1 = np.tile(np.expand_dims(self.mu_sa, axis=3), (1,1,1,K))
          temp2 = np.tile(P_k[np.newaxis, np.newaxis, :, :], (states,actions,1,1))
          mu_sa_k = temp1*temp2
          mu_sa_k = np.sum(mu_sa_k, axis=2)
  
  
          return M_k, mu_s_k, mu_sa_k, initD_k, P_k, rewards, policies, logpolicies, dynets, dyoptims
  
  
      def prior_cluster(
          self, P_k, rewards, policies, logpolicies, dynets,dyoptims,n, reward_f
      ):
          curr_P_k = np.copy(P_k)
          P_k_minus_n = np.delete(curr_P_k,n,0)
          M_k_minus_n = np.expand_dims(np.sum(P_k_minus_n, axis=0), axis=1)
  
          M = len(self.example_samples)
          alpha = self.alpha
          if len(dynets) >= 3*len(self.true_rewards):
              alpha = 0
          curr_priors = M_k_minus_n/(M - 1 + alpha)
          new_prior = np.expand_dims(np.array([alpha/(M - 1 + alpha)]),axis=1)
          prior_n = np.concatenate( (curr_priors, new_prior), axis=0 )
  
          dynet = DynamicRewardModel()
          dynet.train()
          dyoptim = optim.Adam(dynet.parameters(),  lr=self.lr)
          r_new_k = dynet(reward_f)
          _, _, policy, logpolicy = self.model.linearvalueiteration(r_new_k.detach().numpy())
          rewards.append(r_new_k)
          policies.append(policy)
          logpolicies.append(logpolicy)
          dynets.append(dynet)
          dyoptims.append(dyoptim) 
          return prior_n, rewards, policies, logpolicies, dynets, dyoptims
  
      def update_P_k(self, P_k, eta_n, n, rewards, policies,logpolicies,dynets,dyoptims):
  
          M,curr_k = np.shape(P_k)
          if eta_n > curr_k - 1:
              extra_k = np.zeros((M,1))
              P_k = np.concatenate((P_k,extra_k),axis=1)
              P_k_n = np.zeros((1,curr_k+1))
              P_k_n[0,eta_n] = 1
              P_k[n,:] = P_k_n
          else:
              P_k_n = np.zeros((1,curr_k))
              P_k_n[0,eta_n] = 1
              P_k[n,:] = P_k_n
              del rewards[-1]
              del policies[-1]
              del logpolicies[-1]
              del dynets[-1]
              del dyoptims[-1]
  
          M_k = np.expand_dims(np.sum(P_k, axis=0), axis=1)
          for i in range(len(M_k)):
              if M_k[i] == 0:
                  P_k = np.delete(P_k,i,1)
                  del rewards[i]
                  del policies[i]
                  del logpolicies[i]
                  del dynets[i]
                  del dyoptims[i]
                  break
  
          return P_k, rewards, policies, logpolicies, dynets, dyoptims
  
  
  
      def Espected_value_per_n(self, solutions, P_k_reward, P_k_policy):
          M, T = np.shape(self.example_samples)[0:2]
          R = len(self.true_rewards)
          S = len(solutions)
          policy_reward_value = np.zeros((S,R)) 
          for s in range(S):
              for r in range(R):
                  policy_reward_value[s,r] = self.model.stdvalue_from_policy(self.true_rewards[r],solutions[s].p)
  
          EV_n = np.zeros((M))
          for n in range(M):
              reward_idx = np.where(P_k_reward[n,:] == 1)[0][0]
              policy_idx = np.where(P_k_policy[n,:] == 1)[0][0]
              EV_n[n] = policy_reward_value[policy_idx,reward_idx]
          
          EV_argtemp = np.argsort(policy_reward_value, axis=0)
          EV_arg = EV_argtemp[-1,:]
  
          return EV_n, EV_arg
          
  
  
  
  
  
  
  
  
  
  
  
  
  