from sympy import sqrt, I, pi, exp, N, simplify
from math import log2

class Matrix:
  def __init__(self, m=None):
    if m is None:
      self.state = [{}]
    else:
      #print(matrix)
      self.shape = (len(m), len(m[0]))
      #self.state = [{}] * self.shape[1]
      self.state = []
      for i in range(self.shape[1]):
        self.state.append({})
      #print(self.state[0])
      for i in range(self.shape[0]):
        for j in range(self.shape[1]):
          #print(self.state, m[i][j])
          if m[i][j] != 0:
            self.state[j][i] = m[i][j]
  def __repr__(self):
    str = '['
    for i in range(self.shape[1]):
      str += '{\n'
      for j in sorted(self.state[i]):
        str += f'  |{j:0{int(log2(self.shape[0]))}b}>: {N(abs(self.state[i][j]) ** 2) * 100:8.4f}% [{self.state[i][j]}],\n'
      str = str.strip(',\n')
      str += '\n}, \n'
    str = str.strip(', \n')
    str += ']'
    return str
  def get_matrix(self):
    m = [[0 for j in range(self.shape[1])] for i in range(self.shape[0])]
    for i in range(self.shape[0]):
      for j in range(self.shape[1]):
        if i in self.state[j]:
          m[i][j] = self.state[j][i]
    return m
  def kron(self, m):
    #print(m.state)
    #print(n.state)
    #print(m.shape[0], m.shape[1], n.shape[0], n.shape[1])
    C = zeros(self.shape[0] * m.shape[0], self.shape[1] * m.shape[1])
    for i in range(self.shape[0]):
      for j in self.state[i]: # range(m.shape[1]):
        for k in range(m.shape[0]):
          for l in m.state[k]: # range(n.shape[1]):
            #print(i,j,k,l)
            #if i in m.state[j] and k in n.state[l]:
            C.state[i * m.shape[0] + k][j * m.shape[1] + l] = self.state[i][j] * m.state[k][l]
            #print(C.state)
    return C
  def dot(self, m):
    #print(self.shape[0], self.shape[1], m.shape[0], m.shape[1])
    C = zeros(self.shape[0], m.shape[1])
    for j in range(m.shape[1]):
      for k in m.state[j]: #range(m.shape[1]):
        for i in self.state[k]: # range(m.shape[0]):
          #print(i,j,k)
          #if i in m.state[k] and k in n.state[j]:
          if not i in C.state[j]:
            C.state[j][i] = 0
          C.state[j][i] += self.state[k][i] * m.state[j][k]
          C.state[j][i] = simplify(C.state[j][i])
          if C.state[j][i] == 0:
            del C.state[j][i]
    #print(C.state)
    return C

class GateMatrix:
  def __init__(self, m=None):
    if m is None:
      self.state = [eye(2)]
    else:
      #print(matrix)
      self.shape = (len(m), 2, 2)
      #self.state = [{}] * self.shape[1]
      self.state = m
      #for i in range(self.shape[1]):
      #  self.state.append({})
      #print(self.state[0])
      #for i in range(self.shape[0]):
      #  for j in range(self.shape[1]):
          #print(self.state, m[i][j])
      #    if m[i][j] != 0:
      #      self.state[j][i] = m[i][j]
  def __repr__(self):
    s = str(self.state)
  #  str = '['
  #  for i in range(self.shape[1]):
  #    str += '{\n'
  #    for j in self.state[i]:
  #      str += f'  |{j:0{int(log2(self.shape[0]))}b}>: {self.state[i][j]},\n'
  #    str = str.strip(',\n')
  #    str += '\n}, \n'
  #  str = str.strip(', \n')
  #  str += ']'
    return s
  def get_matrix(self):
    m = [[0 for j in range(self.shape[1])] for i in range(self.shape[0])]
    for i in range(self.shape[0]):
      for j in range(self.shape[1]):
        if i in self.state[j]:
          m[i][j] = self.state[j][i]
    return m
  def kron(self, m):
    #print(m.state)
    #print(n.state)
    #print(m.shape[0], m.shape[1], n.shape[0], n.shape[1])
    C = zeros(self.shape[0] * m.shape[0], self.shape[1] * m.shape[1])
    for i in range(self.shape[0]):
      for j in self.state[i]: # range(m.shape[1]):
        for k in range(m.shape[0]):
          for l in m.state[k]: # range(n.shape[1]):
            #print(i,j,k,l)
            #if i in m.state[j] and k in n.state[l]:
            C.state[i * m.shape[0] + k][j * m.shape[1] + l] = self.state[i][j] * m.state[k][l]
            #print(C.state)
    return C
  def dot(self, m):
    #print(self.shape[0], self.shape[1], m.shape[0], m.shape[1])
    C = zeros(self.shape[0], m.shape[1])
    for j in range(m.shape[1]):
      for k in m.state[j]: #range(m.shape[1]):
        for i in self.state[k]: # range(m.shape[0]):
          #print(i,j,k)
          #if i in m.state[k] and k in n.state[j]:
          if not i in C.state[j]:
            C.state[j][i] = 0
          C.state[j][i] += self.state[k][i] * m.state[j][k]
          if C.state[j][i] == 0:
            del C.state[j][i]
    #print(C.state)
    return C

def zeros(x, y):
  m = Matrix()
  m.shape = (x, y)
  m.state = []
  for i in range(y):
    m.state.append({})
  return m

def eye(x):
  m = Matrix()
  m.shape = (x, x)
  m.state = []
  for i in range(x):
    m.state.append({})
    m.state[i][i] = 1
  return m

#def kron(m, n):
#  #print(m.state)
#  #print(n.state)
#  #print(m.shape[0], m.shape[1], n.shape[0], n.shape[1])
#  C = zeros(m.shape[0] * n.shape[0], m.shape[1] * n.shape[1])
#  for i in range(m.shape[0]):
#    for j in m.state[i]: # range(m.shape[1]):
#      for k in range(n.shape[0]):
#        for l in n.state[k]: # range(n.shape[1]):
#          #print(i,j,k,l)
#          #if i in m.state[j] and k in n.state[l]:
#          C.state[i * n.shape[0] + k][j * n.shape[1] + l] = m.state[i][j] * n.state[k][l]
#          #print(C.state)
#  return C

#def dot(m, n):
#  #print(m.shape[0], m.shape[1], n.shape[0], n.shape[1])
#  C = zeros(m.shape[0], n.shape[1])
#  for j in range(n.shape[1]):
#    for k in n.state[j]: #range(m.shape[1]):
#      for i in m.state[k]: # range(m.shape[0]):
#        #print(i,j,k)
#        #if i in m.state[k] and k in n.state[j]:
#        if not i in C.state[j]:
#          C.state[j][i] = 0
#        C.state[j][i] += m.state[k][i] * n.state[j][k]
#        if C.state[j][i] == 0:
#          del C.state[j][i]
#  #print(C.state)
#  return C

#I = eye(2)
#X = Matrix([[0,1],[1,0]])

GATES = {
  'I':   Matrix([[1, 0], [0, 1]]),
  'X':   Matrix([[0, 1], [1, 0]]),
  'Y':   Matrix([[0, -I], [I, 0]]), 
  'Z':   Matrix([[1, 0], [0, -1]]),
  'S':   Matrix([[1, 0], [0, I]]),
  'SDG': Matrix([[1, 0], [0, -I]]),
  'T':   Matrix([[1, 0], [0, exp(I * pi / 4)]]),
  'TDG': Matrix([[1, 0], [0, exp(-I * pi / 4)]]),
  'H':   Matrix([[sqrt(2) / 2, sqrt(2) / 2], [sqrt(2) / 2, -sqrt(2) / 2]]),
}