app.py

import numpy as np
import random
import io
import duckdb
import gradio as gr
import math
from datetime import datetime
import PIL
import matplotlib.pyplot as plt
from PIL import Image
import pennylane as qml
import base64
from qutip import *
from qutip.qip.operations import *
from qutip.qip.circuit import QubitCircuit, Gate
import pennylane as qml
from math import pi

# Define a Pennylane device
dev = qml.device('default.qubit', wires=10)

# Hugging Face and DuckDB function placeholders
def store_in_hf_dataset(data):
    pass

def load_from_hf_dataset():
    return []

# Function to buffer the plot and return as PIL image
def buffer_plot_and_get(fig):
    buf = io.BytesIO()
    fig.savefig(buf, format='png')
    buf.seek(0)
    return PIL.Image.open(buf)

# Store image in bytes for DuckDB
def pil_image_to_bytes(image):
    img_byte_arr = io.BytesIO()
    image.save(img_byte_arr, format='PNG')
    return img_byte_arr.getvalue()

# Encode the image in base64 to display in HTML
def encode_image_from_blob(blob):
    img_buffer = io.BytesIO(blob)
    image = Image.open(img_buffer)
    img_str = base64.b64encode(img_buffer.getvalue()).decode("utf-8")
    return f'<img src="data:image/png;base64,{img_str}" style="max-width:500px;"/>'

# Function to generate a random Hamiltonian
def generate_random_hamiltonian(num_qubits):
    terms = []
    for _ in range(random.randint(1, 5)):
        coeff = round(random.uniform(-1, 1), 2)
        pauli_ops = [random.choice(['I', 'X', 'Y', 'Z']) for _ in range(num_qubits)]
        term = f"{coeff} * {' '.join(pauli_ops)}"
        terms.append(term)
    return " + ".join(terms)

# Function to convert Hamiltonian to QASM code
def hamiltonian_to_qasm(hamiltonian, num_qubits):
    qasm_code = f"OPENQASM 2.0;\ninclude \"qelib1.inc\";\nqreg q[{num_qubits}];\n"
    rotations = {i: 0.0 for i in range(num_qubits)}
    terms = hamiltonian.split(" + ")
    
    for term in terms:
        coeff, paulis = term.split(" * ")
        paulis = paulis.split()
        coeff = float(coeff)

        for i, pauli in enumerate(paulis):
            if pauli == "X":
                qasm_code += f"x q[{i}];\n"
            elif pauli == "Y":
                qasm_code += f"ry(pi/2) q[{i}];\n"
            elif pauli == "Z":
                rotations[i] += coeff

    for i, angle in rotations.items():
        if angle != 0:
            angle_degrees = round(angle * 180 / math.pi, 2)
            qasm_code += f"rz({angle_degrees}) q[{i}];\n"
    
    return qasm_code

# Function to parse QASM code and create Pennylane circuit
def qasm_to_pennylane(qasm_code):
    qasm_lines = qasm_code.split("\n")
    num_qubits = int(qasm_lines[2].split('[')[1].split(']')[0])  # Extract number of qubits from QASM
    
    @qml.qnode(dev)
    def circuit():
        for line in qasm_lines:
            if "x" in line:
                qml.PauliX(int(line.split('q[')[1].split(']')[0]))
            elif "rz" in line:
                angle = float(line.split('(')[1].split(')')[0])
                qml.RZ(angle, int(line.split('q[')[1].split(']')[0]))
            elif "ry" in line:
                qml.RY(pi / 2, int(line.split('q[')[1].split(']')[0]))
        return qml.state()
    
    return circuit

# Store data in DuckDB
def store_in_duckdb(data, db_file='quantum_hamiltonians.duckdb'):
    conn = duckdb.connect(database=db_file)
    conn.execute("""CREATE TABLE IF NOT EXISTS hamiltonians (
        id INTEGER,
        plot BLOB,
        hamiltonian VARCHAR,
        qasm_code VARCHAR,
        trotter_code VARCHAR,
        num_qubits INTEGER,
        trotter_order INTEGER,
        timestamp TIMESTAMP
    )""")
    conn.executemany("""INSERT INTO hamiltonians (id, plot, hamiltonian, qasm_code, trotter_code, num_qubits, trotter_order, timestamp)
                        VALUES (?, ?, ?, ?, ?, ?, ?, ?)""", data)
    conn.close()

# Function to load results from DuckDB
def load_from_duckdb(db_file='quantum_hamiltonians.duckdb'):
    conn = duckdb.connect(database=db_file)
    df = conn.execute("SELECT * FROM hamiltonians").df()
    conn.close()

    # Convert results to HTML with images
    html_content = []
    for index, row in df.iterrows():
        plot_blob = row['plot']
        encoded_img = encode_image_from_blob(plot_blob)

        html_content.append(f"""
            <table style='width: 100%; border-collapse: collapse; margin: 10px;'>
                <tr>
                    <td style='width: 30%; text-align: center;'>
                        <h3>Circuit {index + 1}</h3>
                        {encoded_img}  <!-- Display the image -->
                    </td>        
                    <td style='padding: 10px;'>
                        <table style='width: 100%; border-collapse: collapse;'>
                            <tr>
                                <td><strong>Hamiltonian:</strong></td><td>{row['hamiltonian']}</td>
                            </tr>
                            <tr>
                                <td><strong>QASM Representation:</strong></td><td>{row['qasm_code']}</td>
                            </tr>
                            <tr>
                                <td><strong>Trotter Decomposition:</strong></td><td>{row['trotter_code']}</td>
                            </tr>
                            <tr>
                                <td><strong>Number of Qubits:</strong></td><td>{row['num_qubits']}</td>
                            </tr>
                            <tr>
                                <td><strong>Trotter Order:</strong></td><td>{row['trotter_order']}</td>
                            </tr>
                            <tr>
                                <td><strong>Timestamp:</strong></td><td>{row['timestamp']}</td>
                            </tr>
                        </table>
                    </td>
                </tr>
            </table>
        """)

    return "".join(html_content)

# Function to generate Hamiltonians
def generate_hamiltonians(num_hamiltonians, selected_qubits, selected_order, write_to_hf, write_to_duckdb):
    results_table = []
    timestamp = datetime.now()

    for i in range(num_hamiltonians):
        num_qubits = random.choice(selected_qubits)
        order = selected_order
        hamiltonian = generate_random_hamiltonian(num_qubits)
        qasm_code = hamiltonian_to_qasm(hamiltonian, num_qubits)
        trotter_code = trotter_decomposition(hamiltonian, order)

        # Generate Pennylane circuit from QASM code
        circuit = qasm_to_pennylane(qasm_code)

        # Draw the Pennylane circuit and save as an image
        fig, ax = qml.draw_mpl(circuit)()
        circuit_plot_image = buffer_plot_and_get(fig)
        circuit_plot_bytes = pil_image_to_bytes(circuit_plot_image)

        # Append data to results table
        results_table.append((i + 1, circuit_plot_bytes, hamiltonian, qasm_code, trotter_code, num_qubits, order, timestamp))

    # Write data to Hugging Face dataset if selected
    if write_to_hf:
        store_in_hf_dataset(results_table)
    
    # Write data to DuckDB if selected
    if write_to_duckdb:
        store_in_duckdb(results_table)

# Function to load results from either DuckDB or Hugging Face dataset
def load_results(load_from_hf, load_from_duckdb1):
    if load_from_hf:
        return load_from_hf_dataset()
    if load_from_duckdb1:
        return load_from_duckdb()

# Function for Trotter decomposition
def trotter_decomposition(hamiltonian, order):
    terms = hamiltonian.split(" + ")
    trotter_steps = []
    
    for term in terms:
        coeff, *pauli_ops = term.split(" * ")
        coeff = float(coeff)
        for _ in range(order):
            trotter_steps.append(f"exp({coeff / order}) * ({' * '.join(pauli_ops)})")
        for _ in range(order):
            trotter_steps.append(f"exp({-coeff / order}) * ({' * '.join(pauli_ops)})")
    
    return " + ".join(trotter_steps)


# Gradio app
with gr.Blocks() as app:
    gr.Markdown("# Quantum Hamiltonian Generator")
    
    with gr.Tab("Generate Hamiltonians"):
        num_hamiltonians = gr.Dropdown(label="Select number of Hamiltonians to generate", choices=[1, 10, 20, 100], value=20)
        qubit_choices = [1, 2, 3, 4, 5, 6]
        selected_qubits = gr.CheckboxGroup(label="Select number of qubits", choices=qubit_choices, value=[1])
        order_choices = [1, 2, 3, 4, 5]
        selected_order = gr.Dropdown(label="Select Trotter order", choices=order_choices, value=1)

        # Checkboxes for writing to HF dataset and DuckDB
        write_to_hf = gr.Checkbox(label="Write to Hugging Face dataset", value=False)
        write_to_duckdb = gr.Checkbox(label="Write to DuckDB", value=True)

        generate_button = gr.Button("Generate Hamiltonians")
        status = gr.Markdown("Click 'Generate Hamiltonians' to start the process.")

        def update_status(num, qubits, order, write_hf, write_duckdb):
            generate_hamiltonians(num, qubits, order, write_hf, write_duckdb)
            return "Data stored as per selection."

        generate_button.click(update_status, inputs=[num_hamiltonians, selected_qubits, selected_order, write_to_hf, write_to_duckdb], outputs=status)

    with gr.Tab("View Results"):
        load_from_hf = gr.Checkbox(label="Load from Hugging Face dataset", value=False)
        load_from_duckdb1 = gr.Checkbox(label="Load from DuckDB", value=True)

        load_button = gr.Button("Load Results")
        output_display = gr.HTML()

        load_button.click(load_results, inputs=[load_from_hf, load_from_duckdb1], outputs=output_display)

app.launch()