import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import classification_report
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.datasets import mnist
import streamlit as st

# Load the MNIST dataset
(train_images, train_labels), (test_images, test_labels) = mnist.load_data()

# Preprocess the data
train_images = train_images.reshape((60000, 28, 28, 1)).astype("float32") / 255
test_images = test_images.reshape((10000, 28, 28, 1)).astype("float32") / 255

# Convert labels to categorical format
train_labels = keras.utils.to_categorical(train_labels, 10)
test_labels = keras.utils.to_categorical(test_labels, 10)

# Define the CNN model
def create_model():
    model = keras.Sequential([
        layers.Conv2D(32, (3, 3), activation="relu", input_shape=(28, 28, 1)),
        layers.MaxPooling2D((2, 2)),
        layers.Conv2D(64, (3, 3), activation="relu"),
        layers.MaxPooling2D((2, 2)),
        layers.Conv2D(64, (3, 3), activation="relu"),
        layers.Flatten(),
        layers.Dense(64, activation="relu"),
        layers.Dense(10, activation="softmax")
    ])
    
    model.compile(optimizer="adam", loss="categorical_crossentropy", metrics=["accuracy"])
    return model

# Streamlit UI
st.title("CNN for MNIST Classification")

if st.button("Train Model"):
    model = create_model()
    with st.spinner("Training..."):
        history = model.fit(train_images, train_labels, validation_data=(test_images, test_labels), epochs=10, batch_size=64)
        
    # Plot training loss and accuracy
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))

    ax1.plot(history.history["loss"], label="Train Loss")
    ax1.plot(history.history["val_loss"], label="Val Loss")
    ax1.set_title("Training and Validation Loss")
    ax1.set_xlabel("Epoch")
    ax1.set_ylabel("Loss")
    ax1.legend()

    ax2.plot(history.history["accuracy"], label="Train Accuracy")
    ax2.plot(history.history["val_accuracy"], label="Val Accuracy")
    ax2.set_title("Training and Validation Accuracy")
    ax2.set_xlabel("Epoch")
    ax2.set_ylabel("Accuracy")
    ax2.legend()

    st.pyplot(fig)

    # Evaluate the model on test data
    test_preds = np.argmax(model.predict(test_images), axis=1)
    true_labels = np.argmax(test_labels, axis=1)

    # Classification report
    report = classification_report(true_labels, test_preds, digits=4)
    st.text("Classification Report:")
    st.text(report)

# Testing with a specific index
index = st.number_input("Enter an index (0-9999) to test:", min_value=0, max_value=9999, step=1)

def test_index_prediction(index):
    image = test_images[index].reshape(28, 28)
    st.image(image, caption=f"True Label: {true_labels[index]}", use_column_width=True)

    prediction = model.predict(test_images[index].reshape(1, 28, 28, 1))
    predicted_class = np.argmax(prediction)
    st.write(f"Predicted Class: {predicted_class}")

if st.button("Test Index"):
    if 'model' in locals() and model:
        test_index_prediction(index)
    else:
        st.error("Train the model first.")