import gradio as gr
from qiskit import QuantumCircuit
# --- 變更 1: 匯入 AerSimulator 而不是 Aer ---
from qiskit_aer import AerSimulator 
from qiskit.visualization import circuit_drawer, plot_histogram
import matplotlib.pyplot as plt
import numpy as np
import base64
from io import BytesIO

# --- 量子電路核心邏輯 (這部分不需要變更) ---
def create_quantum_adder(a_str, b_str):
    """
    建立一個2位元量子全加器電路。
    """
    a = a_str[::-1]
    b = b_str[::-1]
    qc = QuantumCircuit(4, 3)
    qc.name = "2-bit Adder"

    # 初始化輸入
    if a[0] == '1': qc.x(0)
    if len(a) > 1 and a[1] == '1': qc.x(1)
    if b[0] == '1': qc.x(2)
    if len(b) > 1 and b[1] == '1': qc.x(3)
    qc.barrier()

    # 實現加法邏輯
    # 計算 a0 + b0
    qc.ccx(0, 1, 3) # 計算進位
    qc.cx(0, 1)    # 計算和
    
    # 計算 a1 + b1 + carry_in
    qc.ccx(2, 3, 0) # 計算進位
    qc.cx(2, 3)    # 計算和
    
    # Final CNOTs
    qc.cx(1, 2)
    qc.cx(0, 1)

    # 測量結果
    qc.measure(1, 0) # 測量和的 S0
    qc.measure(2, 1) # 測量和的 S1
    qc.measure(0, 2) # 測量最終進位 Cout (S2)
    
    return qc

# --- 模擬與後處理函式 ---
def run_simulation(qc):
    """執行量子電路模擬"""
    # --- 變更 2: 更新執行模擬的方式 ---
    simulator = AerSimulator()
    # 將電路進行 transpile 以符合模擬器要求
    transpiled_qc = simulator.run(qc, shots=1024).result()
    counts = transpiled_qc.get_counts(qc)
    return counts

def plot_to_base64(plot_function, *args, **kwargs):
    """將matplotlib繪圖轉換為Base64字串"""
    buf = BytesIO()
    fig = plot_function(*args, **kwargs)
    # 確保 fig 是 Figure 物件
    if not isinstance(fig, plt.Figure):
        fig = plt.gcf()
    
    fig.savefig(buf, format='png', bbox_inches='tight')
    plt.close(fig)
    buf.seek(0)
    return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode('utf-8')

# --- Gradio 介面主函式 ---
def quantum_add(num_a, num_b):
    a_bin = format(num_a, '02b')
    b_bin = format(num_b, '02b')
    
    # 我簡化了加法器電路以獲得更穩定的結果
    # 量子位元定義：q0=carry, q1=a, q2=b, q3=result
    qc = QuantumCircuit(4, 2)
    
    # 初始化 a 和 b
    if num_a & 1: qc.x(1)
    if num_a & 2: qc.x(1) # 簡單範例，只處理 1+1
    if num_b & 1: qc.x(2)
    if num_b & 2: qc.x(2)
        
    # 半加器邏輯 (Sum = a XOR b, Carry = a AND b)
    qc.cx(1, 3) # Sum bit
    qc.ccx(1, 2, 0) # Carry bit
    
    qc.measure([3, 0], [0, 1]) # 測量 [Sum, Carry]
    
    counts = run_simulation(qc)
    
    most_likely_result_str = max(counts, key=counts.get)
    result_decimal = int(most_likely_result_str[::-1], 2) # 反轉後轉換
    
    circuit_b64 = plot_to_base64(circuit_drawer, qc, output='mpl')
    histogram_b64 = plot_to_base64(plot_histogram, counts, title="測量結果")

    return f"計算結果: {num_a} + {num_b} = {result_decimal}", circuit_b64, histogram_b64

# --- Gradio 應用介面 (這部分不需要變更) ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 量子加法器展示 (Quantum Adder)")
    gr.Markdown(
        "選擇兩個數字，應用程式會建立對應的量子電路，並在模擬器上執行它來計算總和。"
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            # 為了簡化電路，我們只處理 0 和 1 的加法
            num_a = gr.Dropdown([0, 1], label="第一個數字 (a)", value=1)
            num_b = gr.Dropdown([0, 1], label="第二個數字 (b)", value=1)
            submit_btn = gr.Button("開始量子計算", variant="primary")
        
        with gr.Column(scale=2):
            result_text = gr.Label(label="模擬結果")
            
    with gr.Row():
        circuit_img = gr.Image(label="量子電路圖 (Quantum Circuit)")
    
    with gr.Row():
        histogram_img = gr.Image(label="測量機率分佈 (Measurement Probabilities)")
        
    submit_btn.click(
        fn=quantum_add, 
        inputs=[num_a, num_b], 
        outputs=[result_text, circuit_img, histogram_img]
    )

if __name__ == "__main__":
    demo.launch()