%matplotlib notebook

# import pygame to run the PewPew emulator from a notebook
import pygame

# add the src directory to the search path to load modules
import sys
sys.path.insert(0, "../src/")

qiskit_images = (
    (
        (0, 3, 3, 3, 3, 3, 3, 0),
        (3, 2, 0, 1, 1, 0, 0, 3),
        (3, 0, 2, 0, 0, 1, 0, 3),
        (3, 1, 0, 2, 0, 0, 1, 3),
        (3, 3, 0, 0, 2, 0, 3, 3),
        (3, 0, 3, 3, 3, 3, 0, 3),
        (3, 0, 0, 0, 0, 0, 2, 3),
        (0, 3, 3, 3, 3, 3, 3, 0)
    ),
    (
        (0, 3, 3, 3, 3, 3, 3, 0),
        (3, 0, 0, 2, 1, 0, 0, 3),
        (3, 0, 1, 2, 0, 1, 0, 3),
        (3, 1, 0, 0, 2, 0, 1, 3),
        (3, 3, 0, 0, 2, 0, 3, 3),
        (3, 0, 3, 3, 3, 3, 0, 3),
        (3, 0, 0, 0, 0, 2, 0, 3),
        (0, 3, 3, 3, 3, 3, 3, 0)
    ),
    (
        (0, 3, 3, 3, 3, 3, 3, 0),
        (3, 0, 0, 1, 2, 0, 0, 3),
        (3, 0, 1, 0, 2, 1, 0, 3),
        (3, 1, 0, 2, 0, 0, 1, 3),
        (3, 3, 0, 2, 0, 0, 3, 3),
        (3, 0, 3, 3, 3, 3, 0, 3),
        (3, 0, 2, 0, 0, 0, 0, 3),
        (0, 3, 3, 3, 3, 3, 3, 0)
    ),
    (
        (0, 3, 3, 3, 3, 3, 3, 0),
        (3, 0, 0, 1, 1, 0, 2, 3),
        (3, 0, 1, 0, 0, 2, 0, 3),
        (3, 1, 0, 0, 2, 0, 1, 3),
        (3, 3, 0, 2, 0, 0, 3, 3),
        (3, 0, 3, 3, 3, 3, 0, 3),
        (3, 2, 0, 0, 0, 0, 0, 3),
        (0, 3, 3, 3, 3, 3, 3, 0)
    )
)

import pew

i = 0      # frame counter
value = 0  # selected level

pew.init() # initialize the PewPew console

# main loop
while not value:
    # check for pressed keys
    keys = pew.keys()
    if keys & pew.K_UP:
        value = 1
    elif keys & pew.K_RIGHT:
        value = 2
    elif keys & pew.K_DOWN:
        value = 3
    elif keys & pew.K_LEFT:
        value = 4
    
    # display the next frame
    animation = (0,0,1,1,2,2,3,3,2,2,1,1)
    i = (i + 1) % len(animation)
    screen = qiskit_images[animation[i]]
    pew.show(pew.Pix.from_iter(screen))
    
    # wait 0.1 seconds
    pew.tick(0.1)
    
# freeze the screen while a button is pressed
while pew.keys():
    pew.tick(0.1)

# the following is just necessary in a jupyter notebook:
pygame.display.quit()
pygame.quit()

import random
import time

# initialize the random-number generator
random.seed(int(time.monotonic()*1000))

# unload unnecessary objects to save RAM
del qiskit_images

# load the new main loop 
from loop import main_loop

# execute the new main loop passing to it the seleced level
if value == 1:
    from rotations import instruction_set_XYZ
    main_loop(instruction_set_XYZ())
else:
    from instruction_sets import InstructionSet
    from displays import IBMQ
    main_loop(InstructionSet(level=value-2, goal=IBMQ))
    
# the following is just necessary in a jupyter notebook:
pygame.display.quit()
pygame.quit()

def main_loop(ins):
    # initialize PewPew console
    pew.init()

    # Load start screens
    for start_screen in start_screens:
        pew.show(pew.Pix.from_iter(start_screen))
        pew.tick(0.2)
    pew.show(pew.Pix.from_iter(blank_screen))
    pew.tick(0.5)
    
    # display the first frame of the level
    pew.show(ins.get_current_screen())

    # flags used throughout the loop
    bool_loop = True
    old_keys = 0
    
    # new main loop
    while bool_loop:
        keys = pew.keys()
    
        if keys != 0 and keys != old_keys:
            # old_keys is necessary to debounce the buttons
            old_keys = keys
    
            # dispatch the pushed buttons
            if keys & pew.K_X:
                value = pew.K_X
            elif keys & pew.K_DOWN:
                value = pew.K_DOWN
            elif keys & pew.K_LEFT:
                value = pew.K_LEFT
            elif keys & pew.K_RIGHT:
                value = pew.K_RIGHT
            elif keys & pew.K_UP:
                value = pew.K_UP
            elif keys & pew.K_O:
                # the key "O" ("Z" in the emulator) will terminate the game
                value = pew.K_O
                bool_loop = False
            else:
                value = 0

            # send the pressed key to the instruction set
            ins.key_pressed(value)
    
        elif keys == 0:
            # this is necessary to be able to push 
            # a button twice in a row
            old_keys = keys
    
        # update the screen and wait for 20ms
        pew.show(ins.get_current_screen())
        pew.tick(0.02)
        
    # the program has been terminated. 
    # display the final sequence
    for final_screen in final_screens:
        pew.show(pew.Pix.from_iter(final_screen))
        pew.tick(0.2)
    pew.show(pew.Pix.from_iter(blank_screen))
    pew.tick(0.2)

class instruction_set_XYZ:
    
    def __init__(self):
        # history of pushed keys
        self.key_hist = []
        # current state vector
        self.state = random_state()

    def key_pressed(self, key):
        if key == pew.K_UP:
            # forget all pushed buttons
            self.key_hist = []
        elif key == pew.K_LEFT or key == pew.K_DOWN or key == pew.K_RIGHT:
            # append button to history
            self.key_hist.append(key)
            
            # if two buttons have been pressed, determine the corresponding transformation
            if len(self.key_hist) == 2:
                if self.key_hist[0] == pew.K_LEFT:
                    gate = 'x'
                elif self.key_hist[0] == pew.K_DOWN:
                    gate = 'y'
                else:
                    gate = 'z'
                
                if self.key_hist[1] == pew.K_LEFT:
                    gate = gate + 'x'
                elif self.key_hist[1] == pew.K_DOWN:
                    gate = gate + 'y'
                else:
                    gate = gate + 'z'
                
                # update the state vector
                self.state = propagate_statevector(self.state, make_circuit(gate))
                
                # clear the history of pushed buttons
                self.key_hist = []
        elif key == pew.K_X:
            # restart the level
            self.__init__()

    def get_current_screen(self):
        return make_image(self.state)

from aether import QuantumCircuit

pi4 = 0.785398
pi2 = 1.570796

def make_circuit(gate):
    qc = QuantumCircuit(2)
    
    if gate[0] == 'x':
        qc.h(0)
    elif gate[0] == 'y':
        qc.rx(pi2,0)
    
    if gate[1] == 'x':
        qc.h(1)
    elif gate[1] == 'y':
        qc.rx(pi2,1)
    
    qc.cx(0,1)
    qc.h(1)
    qc.rx(pi2,1)
    qc.h(1)
    qc.cx(0,1)
    
    if gate[0] == 'x':
        qc.h(0)
    elif gate[0] == 'y':
        qc.rx(-pi2,0)
    
    if gate[1] == 'x':
        qc.h(1)
    elif gate[1] == 'y':
        qc.rx(-pi2,1)
    
    return qc

from aether import QuantumCircuit, simulate

def propagate_statevector(vec,qc):
    qc_i = QuantumCircuit(2,0)
    qc_i.initialize(vec)

    return simulate(qc_i + qc, get='statevector')

def random_state():
    state = [[1.0,0.0],[0.0,0.0],[0.0,0.0],[0.0,0.0]]
    
    for i in range(5):
        gate = ['xx','xy','xz','yx','yz','yy','zx','zy','zz'][randint(0,8)]
        qc = make_circuit(gate)
        state = propagate_statevector(state, qc)
    return state

def make_image(state):
    blocks = []
    for num in state:
        blocks.append(make_block(num))
    
    image = pew.Pix()
    
    for i in range(2):
        for j in range(4):
            tmp = blocks[2*i][j] + blocks[2*i+1][j]
            for x in range(8):
                image.pixel(x,4*i+j,tmp[x])
    
    return image 

def make_block(c_num):
    amp = sqrt(c_num[0]*c_num[0] + c_num[1]*c_num[1])
    phi = atan2(c_num[1], c_num[0])
    
    if amp < 0.01:
        phi = 0
    
    scenario = 0
    phases = [0.0, pi4, -pi4, 2.0*pi4, -2.0*pi4, 3.0*pi4, -3.0*pi4, 4.0*pi4, -4.0*pi4]
    scenarios = [1,  -1,   -2,   4,        2,      -4,       -3,       3,        3      ]
    for i in range(9):
        if (phi - phases[i])*(phi - phases[i]) < 0.001:
            scenario = scenarios[i]
            continue
    
    if amp < 0.25:
        block = [[0,0,0,0],[0,2,2,0],[0,2,2,0],[0,0,0,0]]
    elif amp < 0.6:
        if scenario > 0:
            block = [[0,0,2,0],[0,0,0,2],[0,0,0,2],[0,0,2,0]]
        else:
            block = [[0,0,0,0],[0,2,2,0],[0,0,2,0],[0,0,0,0]]
    elif amp < 0.9:
        if scenario > 0:
            block = [[0,0,0,2],[0,0,2,0],[0,0,2,0],[0,0,0,2]]
        else:
            block = [[0,0,2,0],[0,0,2,2],[0,0,0,0],[0,0,0,0]]
    else:
        if scenario > 0:
            block = [[0,0,0,2],[0,0,0,2],[0,0,0,2],[0,0,0,2]]
        else:
            block = [[0,0,2,2],[0,0,0,2],[0,0,0,0],[0,0,0,0]]
    
    if scenario != 1 and scenario != -1:
        for r in range(abs(scenario) - 1):
            block = rot90(block)
    
    return block

def rot90(block):
    res = []
    for i in range(len(block)):
        transposed = []
        for col in block:
            transposed.append(col[i])
        transposed.reverse()
        res.append(transposed)
    return res