import gradio as gr
import matplotlib.pyplot as plt
from io import BytesIO
from PIL import Image
import random

# ---------------- Function to add random noise (simulate eavesdropper) ----------------
def add_noise_to_key(key_str, flip_percent=0.1):
    """
    Flips 'flip_percent' of bits in the key to simulate eavesdropper noise.
    """
    bits = list(key_str.strip())
    total_bits = len(bits)
    num_flips = int(total_bits * flip_percent)
    
    # Randomly select positions to flip
    flip_indices = random.sample(range(total_bits), num_flips)

    for i in flip_indices:
        bits[i] = '1' if bits[i] == '0' else '0'  # Flip the bit

    return ''.join(bits), flip_indices  # Also return which bits were flipped


# ---------------- Function to compare and plot ----------------
def compare_original_vs_noisy(key_str):
    """
    Compares original and noisy QKD keys with a visual plot.
    Highlights flipped bits for clarity.
    """
    key_str = ''.join([b for b in key_str.strip() if b in '01'])  # Clean input
    if not key_str:
        return None, "⚠️ Invalid input: Please enter a binary key."

    noisy_key, flipped_indices = add_noise_to_key(key_str)

    total_bits = len(key_str)
    heights = [1] * total_bits  # Fixed height bars
    x = range(total_bits)

    fig, axs = plt.subplots(2, 1, figsize=(12, 3.8), sharex=True)

    # ---------------- First plot: Original Key (all green) ----------------
    axs[0].bar(x, heights, color='green', edgecolor='black', linewidth=0.2)
    axs[0].set_title("Before Noise: Original QKD Key", fontsize=11)
    axs[0].set_yticks([])
    axs[0].set_ylabel("Bit")

    # ---------------- Second plot: Noisy Key (red where flipped, green where same) ----------------
    colors = ['red' if i in flipped_indices else 'green' for i in range(total_bits)]
    axs[1].bar(x, heights, color=colors, edgecolor='black', linewidth=0.2)
    axs[1].set_title("After Noise: Key with Eavesdropper Flips", fontsize=11)
    axs[1].set_xlabel("Bit Index")
    axs[1].set_yticks([])
    axs[1].set_ylabel("Bit")

    plt.tight_layout()

    # Save to buffer and return
    buf = BytesIO()
    plt.savefig(buf, format='png')
    plt.close()
    buf.seek(0)

    # Generate human-readable summary
    corruption_rate = (len(flipped_indices) / total_bits) * 100
    summary = f"⚠️ {len(flipped_indices)} bits flipped out of {total_bits} — {corruption_rate:.2f}% corruption detected."

    return Image.open(buf), summary


def get_tab6_noise_simulation():
    with gr.Tab("🔍 Eavesdropper Noise Simulation"):
        original_input = gr.Textbox(label="Enter QKD Key (binary)", lines=3)
        simulate_btn = gr.Button("Simulate Eavesdropping")
        output_graph = gr.Image(label="Original vs Noisy Key")
        summary_output = gr.Textbox(label="Noise Impact Summary", lines=2)

        simulate_btn.click(
            fn=compare_original_vs_noisy,
            inputs=[original_input],
            outputs=[output_graph, summary_output]
        )