# Suppress warnings for a clean output
import warnings; warnings.filterwarnings("ignore")
import numpy as np
np.random.seed(42)

# Data setup
from tensorflow.keras.datasets import mnist
(X_train, y_train), (X_test, y_test) = mnist.load_data()
print(X_train.shape, y_train.shape)
print(X_test.shape, y_test.shape)

# Let us preview an image and its label
import matplotlib.pyplot as plt
plt.imshow(X_train[0], cmap="gray")
plt.title(f"Label: {y_train[0]}")
plt.axis("off")
plt.show()

# Data normalization for better learning
X_train = X_train / 255.0
X_test = X_test / 255.0

# Flatten the 28x28 images into 1D vectors
X_train_flat = X_train.reshape(X_train.shape[0], -1)
X_test_flat = X_test.reshape(X_test.shape[0], -1)

# Let us confirm the new shape
print(X_train_flat.shape)
print(X_test_flat.shape)

# Build a simple neural network model
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
model = Sequential([
    Dense(128, activation="relu", input_shape=(784,)),
    Dense(64, activation="relu"),
    Dense(10, activation="softmax")
])
model.summary()

# Compile the model with loss, optimizer, and metrics
model.compile(
    loss="sparse_categorical_crossentropy",
    optimizer="adam",
    metrics=["accuracy"]
)

# Train the neural network
history = model.fit(X_train_flat, y_train, epochs=5, batch_size=32, validation_split=0.1)

# Plot training and validation accuracy
import matplotlib.pyplot as plt
plt.plot(history.history["accuracy"], label="Training Accuracy")
plt.plot(history.history["val_accuracy"], label="Validation Accuracy")
plt.xlabel("Epoch")
plt.ylabel("Accuracy")
plt.title("Model Training Progress")
plt.legend()
plt.show()

# Evaluate on test data to check real-world performance
test_loss, test_acc = model.evaluate(X_test_flat, y_test)
print(f"Test accuracy: {test_acc:.2f}")

# Predict a digit and show the result
import numpy as np
idx = np.random.randint(0, X_test_flat.shape[0])
image = X_test_flat[idx]
label = y_test[idx]
pred = model.predict(image.reshape(1, -1))
predicted_digit = np.argmax(pred)
plt.imshow(X_test[idx], cmap="gray")
plt.title(f"Predicted: {predicted_digit}, True: {label}")
plt.axis("off")
plt.show()

# Make multiple predictions and check overall performance
preds = model.predict(X_test_flat)
predicted_labels = np.argmax(preds, axis=1)
from sklearn.metrics import classification_report, confusion_matrix
print(classification_report(y_test, predicted_labels))

# Show a confusion matrix of real vs predicted digits
import seaborn as sns
cm = confusion_matrix(y_test, predicted_labels)
plt.figure(figsize=(8,6))
sns.heatmap(cm, annot=True, fmt="d", cmap="Blues")
plt.xlabel("Predicted Digit")
plt.ylabel("True Digit")
plt.title("Confusion Matrix")
plt.show()

# Mini project challenge: Try with fewer data or more epochs
choice = input("Type '1' to use half the data, or '2' to train for 10 epochs: ")
if choice == '1':
    print("Training with half the data...")
    subset_X = X_train_flat[:30000]
    subset_y = y_train[:30000]
    model.fit(subset_X, subset_y, epochs=5, batch_size=32, validation_split=0.1)
elif choice == '2':
    print("Training for 10 epochs...")
    model.fit(X_train_flat, y_train, epochs=10, batch_size=32, validation_split=0.1)
else:
    print("Try rerunning and enter 1 or 2.")