import gradio as gr
import cv2
import numpy as np
import matplotlib.pyplot as plt
from tensorflow.keras.datasets import mnist

# Functions for MNIST processing steps
def load_mnist():
    (x_train, y_train), (x_test, y_test) = mnist.load_data()
    return x_test, y_test

def get_grayscale(image):
    return cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

def thresholding(src):
    return cv2.threshold(src, 127, 255, cv2.THRESH_BINARY)[1]

def gaussian_blur(image):
    return cv2.GaussianBlur(image, (5, 5), 0)

def edge_detection(image):
    return cv2.Canny(image, 100, 200)

def process_mnist_image(img, steps):
    original_img = img.copy()
    step_images = {'Original': original_img}

    for step in steps:
        if step == "Grayscale Conversion":
            img = get_grayscale(img)
        elif step == "Thresholding":
            img = thresholding(img)
        elif step == "Gaussian Blur":
            img = gaussian_blur(img)
        elif step == "Edge Detection":
            img = edge_detection(img)
        
        step_images[step] = img

    return step_images

def visualize_steps(img, steps):
    step_images = process_mnist_image(img, steps)
    
    fig, axes = plt.subplots(1, len(step_images), figsize=(15, 5))
    for ax, (step, img) in zip(axes, step_images.items()):
        ax.imshow(img, cmap='gray')
        ax.set_title(step)
        ax.axis('off')
    
    plt.tight_layout()
    plt.savefig('mnist_processing_steps.png')
    return 'mnist_processing_steps.png'

# Interactive tutorial steps
tutorial_steps = [
    "Grayscale Conversion",
    "Thresholding",
    "Gaussian Blur",
    "Edge Detection"
]

# Interactive questions
questions = [
    {
        "question": "What is the first step in preprocessing MNIST images?",
        "options": ["Gaussian Blur", "Grayscale Conversion", "Edge Detection"],
        "answer": "Grayscale Conversion"
    },
    {
        "question": "What does thresholding do in image preprocessing?",
        "options": ["Detects edges", "Blurs the image", "Binarizes the image"],
        "answer": "Binarizes the image"
    },
    {
        "question": "What library is used to load the MNIST dataset?",
        "options": ["OpenCV", "TensorFlow", "Pytorch"],
        "answer": "TensorFlow"
    },
    {
        "question": "What color space conversion is used in Grayscale Conversion?",
        "options": ["BGR to RGB", "RGB to GRAY", "BGR to GRAY"],
        "answer": "BGR to GRAY"
    },
    {
        "question": "What is the purpose of Gaussian Blur?",
        "options": ["Detect edges", "Reduce noise", "Convert to grayscale"],
        "answer": "Reduce noise"
    }
]

def quiz_interface():
    def check_answer(question_idx, selected):
        if questions[question_idx]["answer"] == selected:
            return "Correct!"
        else:
            return "Incorrect. The correct answer is: " + questions[question_idx]["answer"]

    interfaces = []
    for idx, question in enumerate(questions):
        radio = gr.Radio(choices=question["options"], label=question["question"])
        button = gr.Button("Submit")
        output = gr.Textbox(label="Result")

        def create_submit_fn(idx):
            def submit(selected):
                return check_answer(idx, selected)
            return submit

        interfaces.append(gr.Interface(
            create_submit_fn(idx),
            radio,
            output,
            live=True
        ))

    return interfaces

# Explanation text
explanation_text = """
**Welcome to the MNIST Processing Tutorial!**
This tutorial will guide you through the basic steps of preprocessing images from the MNIST dataset.
**Steps in the MNIST Image Processing:**
1. **Grayscale Conversion:** Converts the image to grayscale to simplify the data.
2. **Thresholding:** Converts the grayscale image into a binary image to distinguish the digits more clearly.
3. **Gaussian Blur:** Applies a blur to reduce noise and detail in the image.
4. **Edge Detection:** Detects the edges of the digits to enhance the features for further processing or recognition tasks.
**Interactive Tutorial:**
Please upload an MNIST image and select the preprocessing steps to visualize their effects.
"""

(x_test, y_test) = load_mnist()
image = gr.Image(shape=(28, 28), image_mode='L')
steps = gr.CheckboxGroup(choices=tutorial_steps, label="Select and order the steps for MNIST processing")
output = gr.Image(type='file', label="Processing Steps Visualization")
explanation = gr.Markdown(explanation_text)

mnist_app = gr.Interface(
    fn=visualize_steps,
    inputs=[image, steps],
    outputs=output,
    title="MNIST Image Processing",
    description=explanation_text,
    css=".gradio-container {background-color: lightgray} #radio_div {background-color: #FFD8B4; font-size: 40px;}"
)

quiz_app = gr.TabbedInterface(
    [mnist_app] + quiz_interface(),
    ["MNIST Tool"] + [f"Question {i+1}" for i in range(len(questions))],
    title="MNIST Tutorial and Quiz"
)

quiz_app.launch()