import numpy as np
import gradio as gr
from PIL import Image

def ReLU(Z):
    return np.maximum(Z, 0)

def softmax(Z):
    A = np.exp(Z) / np.sum(np.exp(Z), axis=0)
    return A

def init_params():
    W1 = np.random.rand(10, 784) - 0.5
    b1 = np.random.rand(10, 1) - 0.5
    W2 = np.random.rand(10, 10) - 0.5
    b2 = np.random.rand(10, 1) - 0.5
    return W1, b1, W2, b2

def forward_prop(W1, b1, W2, b2, X):
    Z1 = W1.dot(X) + b1
    A1 = ReLU(Z1)
    Z2 = W2.dot(A1) + b2
    A2 = softmax(Z2)
    return Z1, A1, Z2, A2

def get_predictions(A2):
    return np.argmax(A2, axis=0)

def make_predictions(X, W1, b1, W2, b2):
    _, _, _, A2 = forward_prop(W1, b1, W2, b2, X)
    predictions = get_predictions(A2)
    return predictions

def predict_digit(img):
    # Load the trained parameters
    params = np.load("trained_params.npz", allow_pickle=True)
    W1, b1, W2, b2 = params["W1"], params["b1"], params["W2"], params["b2"]

    # Convert the sketchpad drawing to grayscale and resize it to (28, 28)
    img_pil = Image.fromarray(np.uint8(img * 255)).convert("L")

    res=img_pil.resize((28, 28))
    
    # Convert the image to a NumPy array and normalize it
    X = np.array(res).reshape((784, 1)) / 255.

    # Get the prediction
    prediction = make_predictions(X, W1, b1, W2, b2)

    return int(prediction)

iface = gr.Interface(
    fn=predict_digit,
    inputs="sketchpad",
    outputs=gr.outputs.Label(num_top_classes=3),
    live=True,
    capture_session=True,
    title="Handwritten Digit Recognizer",
    description="Draw a digit using your mouse, and the model will try to recognize it.",
)

if __name__ == "__main__":
    iface.launch()