# Do the necessary import for our program
%matplotlib inline
from qiskit import *
from qiskit.aqua.algorithms import Grover
from qiskit.aqua.components.oracles import LogicalExpressionOracle
from qiskit.tools.visualization import plot_histogram
import ipywidgets as widgets
import matplotlib.pyplot as plt
from IPython.display import clear_output

SHOW_CIRCUIT = True

def create_expanded_button(description, button_style):
    return widgets.Button(description=description, button_style=button_style, 
                  layout=widgets.Layout(height='auto', width='5%'))

header_button = create_expanded_button('Header', 'success')
left_button = create_expanded_button('Left', 'info')
center_button = create_expanded_button('Center', 'warning')
right_button = create_expanded_button('Right', 'info')
footer_button = create_expanded_button('Footer', 'success')

app_ui = widgets.AppLayout(header=None,
          left_sidebar=None,
          center=center_button,
          right_sidebar=None,
          footer=None)

display(app_ui)

log_expr_widget = widgets.Text(
    value='',
    description='Logical expression:',
    disabled=True
)

out1 = widgets.Output()
out2 = widgets.Output()

rel_matrix = None

names = ["A", "B", "C", "D"]
num_names = len(names)


REL_NONE = ""
REL_GOOD = "good"
REL_BAD = "bad"

# Compute logical expression from dropdowns
def compute_log_expr():
    log_expr_str = ""
    good_rel_exist = False
    bad_rel_exist = False
    for col_num in range(num_names + 1):
        for row_num in range(num_names + 1):
            if col_num > row_num and row_num > 0:
                if rel_matrix[col_num][row_num].value == 1: #1 is good
                    good_rel_exist = True
                    break
                    
    if good_rel_exist:                
        log_expr_str = "("
        for col_idx in range(num_names + 1):
            for row_idx in range(num_names + 1):
                if col_idx > row_idx and row_idx > 0:
                    if rel_matrix[col_idx][row_idx].value == 1: #1 is good
                        log_expr_str += "("
                        log_expr_str += rel_matrix[0][row_idx].value
                        log_expr_str += " & "
                        log_expr_str += rel_matrix[col_idx][0].value
                        log_expr_str += ")|"
        if log_expr_str[-1] == '|':
            log_expr_str = log_expr_str[:-1]
        log_expr_str += ")"
        
    for col_num in range(num_names + 1):
        for row_num in range(num_names + 1):
            if col_num > row_num and row_num > 0:
                if rel_matrix[col_num][row_num].value == 2: #2 is bad
                    bad_rel_exist = True
                    break

    if bad_rel_exist:
        if good_rel_exist:
            log_expr_str += "&"
        bad_log_expr_str = "("
        for col_idx in range(num_names + 1):
            for row_idx in range(num_names + 1):
                if col_idx > row_idx and row_idx > 0:
                    if rel_matrix[col_idx][row_idx].value == 2: #2 is bad
                        bad_log_expr_str += "~("
                        bad_log_expr_str += rel_matrix[0][row_idx].value
                        bad_log_expr_str += " & "
                        bad_log_expr_str += rel_matrix[col_idx][0].value
                        bad_log_expr_str += ")&"
        if bad_log_expr_str[-1] == '&':
            bad_log_expr_str = bad_log_expr_str[:-1]
        bad_log_expr_str += ")"
        log_expr_str += bad_log_expr_str
        
    log_expr_widget.value = log_expr_str   

#def handle_log_expr_text_entered(text_widget):
def handle_log_expr_text_entered(some_widget):
    with out1:
        clear_output()
        
    compute_log_expr()
    algorithm = Grover(LogicalExpressionOracle(log_expr_widget.value))

    # Run the algorithm on a QASM simulator, printing the most frequently occurring result
    qasm_simulator_backend = BasicAer.get_backend('qasm_simulator')
    qasm_simulator_result = algorithm.run(qasm_simulator_backend)
    
    if SHOW_CIRCUIT:
        circuit = qasm_simulator_result['circuit']
    
    with out1:
        display(plot_histogram(qasm_simulator_result['measurement']))
    
    if SHOW_CIRCUIT:
        with out2:
            clear_output()
            display(circuit.draw(output='mpl'))


log_expr_button = widgets.Button(
    description = "Run",
)


# Relationship model for the grid
def compute_rel_matrix_entry(col_num, row_num):
    entry = None
    if col_num == 0:
        if row_num > 0:
            entry = widgets.Label(value=names[row_num - 1])
        else:
            entry = widgets.Label(value='')
    elif row_num == 0: 
        entry = widgets.Label(value=names[col_num - 1])
    elif col_num < row_num: 
        entry = widgets.Dropdown(
            options=[('', 0), (REL_GOOD, 1), (REL_BAD, 2)],
            value=0,
            layout=widgets.Layout(width='auto', height='auto')
        )
    else:
        entry = widgets.Label(value='')
    return entry

rel_matrix = [[compute_rel_matrix_entry(i, j) for i in range(num_names + 1)] for j in range(num_names + 1)]
for col_num in range(num_names + 1):
    for row_num in range(num_names + 1):
        if col_num > row_num and row_num > 0:
            rel_matrix[col_num][row_num].observe(handle_log_expr_text_entered, names='value')            

# Make Grid with name headings and dropdown relationships
rel_gridbox = widgets.GridBox(children=[rel_matrix[col_idx][row_idx] for row_idx in [0, 1, 2, 3] for col_idx in [0, 2, 3, 4]],
        layout=widgets.Layout(
            width='100%',
            grid_template_columns='70px ' * len(names),
            grid_template_rows='30px ' * len(names), 
            grid_gap='10px 10px')
       )

tab = widgets.Tab(children = [out1, out2], layout=widgets.Layout(height='auto', width='auto'))
tab.set_title(0, 'Probabilities')

if SHOW_CIRCUIT:
    tab.set_title(1, 'Circuit')
    
app_header = widgets.Label(value = "Grover Search Party")

center_pane = widgets.VBox([app_header, rel_gridbox, widgets.VBox([log_expr_widget, tab])])

app_ui.center = center_pane






