!pip install qiskit
!pip install matplotlib
!pip install sympy
from tqdm import tqdm

import numpy as np
import math
import qiskit
from qiskit import circuit
from qiskit.circuit.random import random_circuit
import copy
import matplotlib as mpl
import matplotlib.pyplot as plt
from qiskit.quantum_info import PTM, Chi, Statevector, DensityMatrix, partial_trace
from qiskit import transpile, QuantumCircuit, QuantumRegister
import qiskit.quantum_info as qi
from qiskit.providers.aer import AerSimulator
from qiskit.providers.aer.noise import NoiseModel, amplitude_damping_error
from qiskit.tools.visualization import plot_histogram
import sympy as sp
from sympy import linsolve, sympify, var, Eq, solve, solve_linear_system, Matrix, symbols
import random

def twirl_qubit(circ, dist=None, qubit=0, twirling_gate=None, r=None):
  if not twirling_gate:
    twirling_gate = (r.choice([i for i in range(4)], size=1, p=dist) if dist else r.choice([i for i in range(4)], size=1))[0]
  
  # twirling_gate = random.randint(0, 3)
  if twirling_gate == 1:
    circ.x(qubit)
  elif twirling_gate == 2:
    circ.y(qubit)
  elif twirling_gate == 3:
    circ.z(qubit)

  return twirling_gate


def sim(num_gates=100, noise=False, twirl=False, twirl_dist=[1, 0, 0, 0], circ=None, noise_AD=math.pi/9, noise_dephasing=math.pi/9, circ_seed=1, twirl_seed=1):
  # twirl_set = {0: 'I', 1: 'X', 2: 'Y', 3: 'Z'}
  np.random.seed(0)
  circ_seed = 7654
  rand_circ = np.random.RandomState()
  rand_twirl = np.random.RandomState()

  rand_circ.seed(circ_seed)
  rand_twirl.seed(twirl_seed)

  if not circ:
    circ = QuantumCircuit(2)
    circ.initialize([1, 0], 0)
    circ.initialize([1, 0], 1)

  random_gate_set = [i for i in range(3)] # and whatever other gates you want for the actual circuit 
  twirl_set = [i for i in range(3)] # and whatever other twirling gates you want
  special_reg_AD = 1
  special_reg_dephasing = 2
  for i in range(num_gates):
    random_gate = rand_circ.choice(random_gate_set)
    random_theta = rand_circ.uniform(low=0, high=np.pi)

    if random_gate == 0:
      circ.rx(random_theta, 0)
    elif random_gate == 1:
      circ.ry(random_theta, 0)
    elif random_gate == 2:
      circ.rz(random_theta, 0)

    #print(random_theta, end=' ')

    # random_gate = np.random.choice([i for i in range(4)])

    # if random_gate == 0:
    #   # circ.h(0)
    #   circ.z(0)
    # elif random_gate == 1:
    #   circ.x(0)
    # elif random_gate == 2:
    #   circ.y(0)
    # elif random_gate == 3:
    #   circ.z(0)
    # else:
    #   circ.s(0)

    if noise:
      if twirl:
          twirling_gate = twirl_qubit(circ, dist=twirl_dist, r=rand_twirl)
          # print(twirling_gate, end=' ')

      # simulate noise on circuit
      circ.cry(noise_AD, 0, special_reg_AD)
      circ.cnot(special_reg_AD, 0)
      special_reg_AD += 2
      if twirl:
          twirl_qubit(circ, dist=twirl_dist, twirling_gate=twirling_gate, r=rand_twirl)

          twirling_gate = twirl_qubit(circ, dist=twirl_dist, r=rand_twirl)

      circ.ry(noise_dephasing, special_reg_dephasing)
      circ.cz(special_reg_dephasing, 0)
      special_reg_dephasing += 2

      if twirl:
          twirl_qubit(circ, dist=twirl_dist, twirling_gate=twirling_gate, r=rand_twirl)

  # boom
  # circ.draw()
  return circuit

def state_tomography(p):
    I = np.trace(np.matmul(np.array([[1, 0],[0, 1]]), np.array(p)))
    X = np.trace(np.matmul(np.array([[0, 1],[1, 0]]), np.array(p)))
    Y = np.trace(np.matmul(np.array([[0, 0- 1j],[1j, 0]]), np.array(p)))
    Z = np.trace(np.matmul(np.array([[1, 0],[0, -1]]), np.array(p)))
    return np.array([I, X, Y, Z])

def equation_PTM(input_list, output_list):
    I, X, Y, Z = var('I X Y Z')
    I_o, X_o, Y_o, Z_o = var('I_o X_o Y_o Z_o')
    eq = sp.Function('eq')
    first = float(np.real(input_list[0])).__round__(4)*I
    second = float(np.real(input_list[1])).__round__(4)*X
    third = float(np.real(input_list[2])).__round__(4)*Y
    fourth = float(np.real(input_list[3])).__round__(4)*Z
    fifth = float(np.real(output_list[0])).__round__(4)*I_o
    sixth = float(np.real(output_list[1])).__round__(4)*X_o
    seventh = float(np.real(output_list[2])).__round__(4)*Y_o
    eighth = float(np.real(output_list[3])).__round__(4)*Z_o
    eq = Eq(first + second + third + fourth, fifth + sixth + seventh + eighth)
    return eq

import random

# Function returns a list of 4 random distributions as a list - [i_dist, x_dist, y_dist, z_dist]
def gen_dist():
  r = [random.random() for i in range(4)]
  r_sum = sum(r)
  r = [(i/r_sum) for i in r]
  return r

def conjug_dists():
  dists = []
  for i in range(4):
    dist = [0 for _ in range(4)]
    dist[i] = 1
    dists.append(dist)
  return dists
conjug_dists()

runs = 1
noise_AD = math.pi/100
noise_dephasing = math.pi/100

def error():
  t = np.array(PTM(True))
  f = np.array(PTM(False))
  ret = t - f
  return ret

def fid(twirl_dist):
  twirled, ideal = (PTM(noisy=True, twirl=True, twirl_dist=twirl_dist, noise_AD=noise_AD, noise_dephasing=noise_dephasing, ideal=False))
  # print(t[0]) # t is a list of R matrices, of length [runs]

  fid_sum = 0
  for R_twirled, R_ideal in zip(twirled, ideal):
    fid_sum += (np.trace(np.matmul(np.transpose(R_ideal), R_twirled)) + 2) / 6

  fid_sum /= runs
  # print(fid_sum)

  return fid_sum

def ket(k):
  k = str(k)

  ket_dict = {'0': [1, 0], '1': [0, 1], '+': [1/math.sqrt(2), 1/math.sqrt(2)], '-': [1/math.sqrt(2), -1/math.sqrt(2)], 'i': [1/math.sqrt(2), 1j/math.sqrt(2)], '-i': [1/math.sqrt(2), -1j/math.sqrt(2)]}
  return ket_dict[k]

def init_circ(num_qubits=13, qubit_0_init='0'):
  circ0 = QuantumCircuit(num_qubits)
  circ0.initialize(ket(qubit_0_init), 0)
  for i in range(1, num_qubits):
    circ0.initialize(ket('0'), i)
  
  return circ0

def run(noisy, state='0', num_qubits=13, twirl=False, twirl_dist=None, noise_AD=0, noise_dephasing=0, circ_seed=1):
  # print('state:', state)
  partials = np.zeros(shape=(2, 2), dtype='complex128')
  n = 0

  partials = []

  while n < runs:
      circ0 = init_circ(num_qubits, qubit_0_init=state)
      sim(num_gates=6, noise=noisy, circ=circ0, noise_AD=noise_AD, twirl=twirl, noise_dephasing=noise_dephasing, twirl_dist=twirl_dist, circ_seed=1, twirl_seed=10000+n)
      state_in0 = Statevector.from_int(0, 2 ** 13) 
      state_out0 = state_in0.evolve(circ0)
      rho_out0 = DensityMatrix(state_out0)
      # print(np.array(partial_trace(rho_out0, [i for i in range(1, 13)])))
      partials.append(np.array(partial_trace(rho_out0, [i for i in range(1, 13)])))
      n += 1

  return partials

def PTM(noisy=False, twirl=False, twirl_dist=None, noise_AD=noise_AD, noise_dephasing=noise_dephasing, ideal=False, circ_seed=1):
    #print(locals())
    # ket 0
    partial_rho_out0s = run(noisy=noisy, state='0', num_qubits=13, twirl=twirl, twirl_dist=twirl_dist, noise_AD=noise_AD, noise_dephasing=noise_dephasing, circ_seed=circ_seed)
    partial_rho_in0 = np.array([[1, 0], [0, 0]])

    # ket 1
    partial_rho_out1s = run(noisy=noisy, state='1', num_qubits=13, twirl=twirl, twirl_dist=twirl_dist, noise_AD=noise_AD, noise_dephasing=noise_dephasing, circ_seed=circ_seed)
    partial_rho_in1 = np.array([[0, 0], [0, 1]])

    # ket +
    partial_rho_outXs = run(noisy=noisy, state='+', num_qubits=13, twirl=twirl, twirl_dist=twirl_dist, noise_AD=noise_AD, noise_dephasing=noise_dephasing, circ_seed=circ_seed)
    partial_rho_inX = np.array([[0.5, 0.5], [0.5, 0.5]])

    # ket i
    partial_rho_outYs = run(noisy=noisy, state='i', num_qubits=13, twirl=twirl, twirl_dist=twirl_dist, noise_AD=noise_AD, noise_dephasing=noise_dephasing, circ_seed=circ_seed)
    partial_rho_inY = np.array([[0.5, -0.5j], [0.5j, 0.5]])

    eqs = []

    for partial_rho_out0, partial_rho_out1, partial_rho_outX, partial_rho_outY in zip(partial_rho_out0s, partial_rho_out1s, partial_rho_outXs, partial_rho_outYs):
      Eq1 = equation_PTM(state_tomography(partial_rho_in0), state_tomography(partial_rho_out0))
      # print("1", Eq1)
      Eq2 = equation_PTM(state_tomography(partial_rho_in1), state_tomography(partial_rho_out1))
      # print("2", Eq2)
      Eq3 = equation_PTM(state_tomography(partial_rho_inX), state_tomography(partial_rho_outX))
      # print("3", Eq3)
      Eq4 = equation_PTM(state_tomography(partial_rho_inY), state_tomography(partial_rho_outY))
      # print("4", Eq4)
      eqs.append((Eq1, Eq2, Eq3, Eq4))

    I, X, Y, Z = var('I X Y Z')
    I_o, X_o, Y_o, Z_o = var('I_o X_o Y_o Z_o')

    solution_set = [linsolve([Eq1, Eq2, Eq3, Eq4], (I, X, Y, Z)) for Eq1, Eq2, Eq3, Eq4 in eqs]

    I_o = np.array([[1, 0],[0, 1]])
    X_o = np.array([[0, 1],[1, 0]])
    Y_o = np.array([[0, 0- 1j],[1j, 0]])
    Z_o = np.array([[1, 0],[0, -1]])

    R_list = []

    for solutions in solution_set:
      solutions_list = list(solutions.args[0])

      #convert sympy to numpy
      solutions_list = eval(str(solutions_list))
      # print(solutions_list)
      reference_list = [I_o, X_o, Y_o, Z_o]

      R = [[0 for i in range(4)] for j in range(4)]
      for i in range(4):
          for j in range(4):
              sol_j = solutions_list[j]
              ref_i = reference_list[i]
              R[i][j] = 0.5*np.real(np.matmul(solutions_list[j], reference_list[i]).trace())
      # print(R)
      R_list.append(R)
    # print('\nR:')
    # print(R)
    # print()

    if not ideal:
      return (R_list, PTM(noisy=True, twirl=True, twirl_dist=twirl_dist, noise_AD=0, noise_dephasing=0, ideal=True))
    else:
      return R_list


noise_AD *= 5
noise_dephasing *= 5
runs *= 2
data = {}

# runs = 1
# ideal_matrix = PTM(noisy=False, twirl=True, noise_AD=0, noise_dephasing=0, ideal=True)
# ideal_matrix

# dists = [gen_dist() for _ in range(5)]
# dists = conjug_dists() + [[0.25 for _ in range(4)]]
# dists = [[0.25 for i in range(4)]] + conjug_dists()

# for dist in (dists):
#   data[tuple(dist)] = fid(dist)
#   print(dist, data[tuple(dist)])

# sorted_data = list(map(lambda x: str(x) + ': ' + str(data[x]), reversed(sorted(data.keys(), key=lambda x: data[x]))))
# print(data)



def print_mat(mat):
  for row in mat:
    print(row)

runs = 1
#seed = 1

matrices = []
storage = []
for circ_seed in range(234, 235):
  runs=1
  ideal_matrix = PTM(noisy=False, twirl=True, noise_AD=0, noise_dephasing=0, ideal=True)[0]
  print('ideal matrix for circ_seed', circ_seed, ':')
  print_mat(ideal_matrix)
  print()

  runs=1
  for noise_mult in range(5, 26, 5):
    noise_AD = noise_mult * math.pi/100
    #print("AD error angle =", noise_mult/100)
    for noise_dephase_mult in range(5, 26, 5):
      #print("Dephasing error angle=", noise_dephase_mult/100)
      noise_dephasing = noise_dephase_mult * math.pi/100

      noisy_mat = PTM(noisy=True, twirl=True, noise_AD=noise_AD, noise_dephasing=noise_dephasing, ideal=True)[0]
      # just_AD_mat = PTM(noisy=True, twirl=True, noise_AD=noise_AD, noise_dephasing=0, ideal=True)[0]
      
      matrices.append(noisy_mat)
      storage.append((matrices[-1]))
      #print(matrices[-1])
print(storage)
#returns a data set of matrices and takes in a numpy array a
 

# x = [1, 2, 3, 4, 5, 6, 7, 8]
# y = [0.6282, 0.5894, 0.485, 0.3457, 0.2083, 0.1015, 0.0366, 0.0079]
# plt.plot(x, y)
# plt.show()


dists = [gen_dist() for _ in range(5)]

for dist in (dists):
  data[tuple(dist)] = fid(dist)
  # print(dist, data[tuple(dist)])

sorted_data = list(map(lambda x: str(x) + ': ' + str(data[x]), reversed(sorted(data.keys(), key=lambda x: data[x]))))
#print(data)


# training = []
# dists = conjug_dists()
# for seed in range(10):
#   for dist in dists:
#     training.append(fid(twirl_dist=dist))

# training

