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

# ---------------- Function to plot bit distribution (Graph 1) ----------------
def plot_key_bits(key_str):
    # Extract only 0s and 1s from input
    bits = [int(b) for b in key_str.strip() if b in '01']
    total_bits = len(bits)

    if total_bits == 0:
        return None, "⚠️ Invalid input: Please enter a binary key."

    # Count zeros and ones
    zero_count = bits.count(0)
    one_count = bits.count(1)
    diff = abs(zero_count - one_count)

    # Fixed height bars for visibility
    heights = [1] * total_bits
    # Blue for 0, Orange for 1
    colors = ['#1f77b4' if b == 0 else '#ff7f0e' for b in bits]

    # Create figure
    fig, ax = plt.subplots(figsize=(12, 2.5))
    ax.bar(range(total_bits), heights, color=colors, edgecolor='black', linewidth=0.2)

    # Add title, labels, legend
    ax.set_title("Distribution of 0s and 1s in Your Key", fontsize=12)
    ax.set_xlabel("Bit Position in Key", fontsize=10)
    ax.set_ylabel("Bit (0 = Blue, 1 = Orange)", fontsize=10)
    ax.set_yticks([])

    legend_handles = [
        Patch(color='#1f77b4', label='0 (Blue)'),
        Patch(color='#ff7f0e', label='1 (Orange)')
    ]
    ax.legend(handles=legend_handles, loc='upper right')
    plt.tight_layout()

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

    # Verdict
    if diff <= total_bits * 0.1:
        verdict = f"✅ Balanced Key: {zero_count} zeros & {one_count} ones — good randomness!"
    else:
        verdict = f"⚠️ Skewed Key: {zero_count} zeros & {one_count} ones — consider regenerating."

    return Image.open(buf), verdict

# ---------------- Add Noise to Key (simulate eavesdropping) ----------------
def add_noise_to_key(key_str, flip_percent=0.1):
    bits = list(key_str.strip())
    total_bits = len(bits)
    num_flips = int(total_bits * flip_percent)

    flip_indices = random.sample(range(total_bits), num_flips)
    for i in flip_indices:
        bits[i] = '1' if bits[i] == '0' else '0'

    return ''.join(bits), flip_indices

# ---------------- Compare Original vs Noisy Key (Graph 2) ----------------
def compare_original_vs_noisy(key_str):
    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
    x = range(total_bits)

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

    # First plot: Original Key
    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
    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()
    buf = BytesIO()
    plt.savefig(buf, format='png')
    plt.close()
    buf.seek(0)

    # 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

# ---------------- Combined Tab ----------------
def get_tab5_randomness():
    with gr.Tab("📊 QKD Key Analyzer"):
        gr.Markdown(
            "<h2 style='text-align: center;'>🔑 Analyze Your QKD Key for Randomness & Noise Simulation</h2>"
        )

        # Input & Button
        binary_input = gr.Textbox(label="🔑 Enter Your QKD Key (binary)", placeholder="Example: 0101011001...", lines=3)
        analyze_btn = gr.Button("Analyze Key")

        # Graph 1 Heading
        gr.Markdown("<h3 style='text-align: center;'>📈 Graph 1: Bit Randomness Distribution</h3>")
        randomness_graph = gr.Image(label="🧪 Distribution of 0s and 1s in your key")
        randomness_text = gr.Textbox(label="Randomness Insight", lines=2)

        # Graph 2 Heading
        gr.Markdown("<h3 style='text-align: center;'>🧨 Graph 2: Noise Simulation – Flipped Bits due to Eavesdropping</h3>")
        noise_graph = gr.Image(label="🔍 Original vs Flipped Bits")
        noise_summary = gr.Textbox(label="Noise Impact Summary", lines=2)

        # Combined callback
        def analyze_key_combined(key_str):
            graph1, verdict = plot_key_bits(key_str)
            graph2, summary = compare_original_vs_noisy(key_str)
            return graph1, verdict, graph2, summary

        analyze_btn.click(
            fn=analyze_key_combined,
            inputs=[binary_input],
            outputs=[randomness_graph, randomness_text, noise_graph, noise_summary]
        )