!python3 -m pip install -q qiskit
!python3 -m pip install -q qiskit_ibm_runtime
!python3 -m pip install qiskit-aer

import numpy as np
import pandas as pd
from qiskit import IBMQ, BasicAer, Aer
from qiskit_ibm_runtime import QiskitRuntimeService, Session, Options, Sampler, Estimator
from qiskit.circuit.quantumcircuit import QuantumCircuit
from qiskit.quantum_info.operators import Operator
from qiskit.extensions.unitary import UnitaryGate
from qiskit.execute_function import execute
from typing import Dict, Optional
from sympy import Matrix
from sympy.physics.quantum import TensorProduct

num_players = 2

# define the payoff matrices. Assume, the 1st matrix is Alice's payoff, the secomd matrix is Bob's payoff
payoff_matrix =[[[3, 0],
                [5, 1]],
                [[3, 5],
                [0, 1]]]

# calculating payoffs

def get_payoff(counts):
  """
    Calculate the reward for the players after the game ends.
  """
  payoff_bob = []
  payoff_alice = []

  for strategy, prob in counts.items():
    strategy_bob = int(strategy[1])
    strategy_alice = int(strategy[0])
    payoff_bob.append(prob * payoff_matrix[0][strategy_alice][strategy_bob])
    payoff_alice.append(prob * payoff_matrix[0][strategy_bob][strategy_alice])

  return sum(payoff_alice), sum(payoff_bob)

# define the simulator
simulator = Aer.get_backend('statevector_simulator')

def build_qcir_two_by_two_parameters(theta1, phi1, theta2, phi2):

  alice = np.array([
      [np.exp(1.0j*phi1)*np.cos(theta1), np.sin(theta1)],
      [-np.sin(theta1), np.exp(-1.0j*phi1)*np.cos(theta1)]])

  bob = np.array([
      [np.exp(1.0j*phi2)*np.cos(theta2), np.sin(theta2)],
      [-np.sin(theta2), np.exp(-1.0j*phi2)*np.cos(theta2)]])

  qc = QuantumCircuit(num_players)

  J = 1/np.sqrt(2)*np.array([[1.0, 0.0, 0.0, 1.0j],
                             [0.0, 1.0, 1.0j, 0.0],
                             [0.0, 1.0j, 1.0, 0.0],
                             [1.0j, 0.0, 0.0, 1.0]]).astype(complex)

  J_unitary = UnitaryGate(Operator(J))
  qc.append(J_unitary, [0,1])
  qc.barrier()

  unitary_alice = UnitaryGate(Operator(alice))
  unitary_bob = UnitaryGate(Operator(bob))

  qc.append(unitary_alice, [0])
  qc.append(unitary_bob, [1])
  qc.barrier()

  Jdagger_unitary = UnitaryGate(Operator(J.conj().T))
  qc.append(Jdagger_unitary, [0,1])

  return qc



space_size = 4

actions_list = []

"quantum games theta and phi ranges"
for theta1 in np.linspace(0, np.pi/2, space_size):
  for phi1 in np.linspace(0, np.pi/2, space_size):
      actions_list.append((theta1,phi1))



def epsilon_Greedy_Policy_2x2(eps, playerName):

  p = np.random.random()

  if p < eps:
    index = np.random.choice(len(actions_list))
    action_chosen = actions_list[index]

  else:
    if(playerName == "Alice"):
      index = np.argmax(Q_Alice)
      action_chosen = actions_list[index]
    elif(playerName == "Bob"):
      index = np.argmax(Q_Bob)
      action_chosen = actions_list[index]

  return (action_chosen[0],action_chosen[1], index)

epsilonArray = []
payoffPercentage3_3 = []
eps = np.arange(0.0, 0.1, 0.01)

for e in eps:
  count = 0
  for k in range(100):

    Q_Alice = [0]*len(actions_list)
    Q_Bob = [0]*len(actions_list)

    T = 20 #assuming t = 100

    alice_reward = [0]*len(actions_list)
    bob_reward = [0]*len(actions_list)

    # Initialize parameters
    #gamma = 0.5 # Discount factor
    alpha = 0.9 # Learning rate
    #epsilon = 0.04

    for t in range(T):

        theta1, phi1, action_index_Alice = epsilon_Greedy_Policy_2x2(e,"Alice")
        theta2, phi2, action_index_Bob = epsilon_Greedy_Policy_2x2(e, "Bob")

        qc = build_qcir_two_by_two_parameters(theta1, phi1, theta2, phi2)
        results = execute(qc, simulator, shots=1024).result().get_counts()
        payoff_alice, payoff_bob = get_payoff(results)

        alice_reward[action_index_Alice] = payoff_alice
        bob_reward[action_index_Bob] = payoff_bob

        Q_Alice[action_index_Alice] = Q_Alice[action_index_Alice] + alpha * (alice_reward[action_index_Alice]-(Q_Alice[action_index_Alice]))
        Q_Bob[action_index_Bob] = Q_Bob[action_index_Bob] + alpha * (bob_reward[action_index_Bob]-(Q_Bob[action_index_Bob]))

        #Q_Alice[action_index_Alice] = Q_Alice[action_index_Alice] + alpha * (alice_reward[action_index_Alice] + (gamma*max(Q_Alice)) -(Q_Alice[action_index_Alice]))
        #Q_Bob[action_index_Bob] = Q_Bob[action_index_Bob] + alpha * (bob_reward[action_index_Bob] + (gamma*max(Q_Bob)) - (Q_Bob[action_index_Bob]))

    for i in range(len(Q_Alice)):
      if(Q_Alice[i] == 3 and Q_Bob[i] ==3):
            count+=1

  epsilonArray.append(e)
  payoffPercentage3_3.append(count)


print(epsilonArray)
print(payoffPercentage3_3)

import matplotlib.pyplot as plt

x = epsilonArray
y = payoffPercentage3_3

plt.scatter(x, y)
plt.show()