from sklearn.datasets import fetch_openml
from sklearn.model_selection import train_test_split
import numpy as np
import streamlit as st
import tensorflow as tf
import numpy as np
import keras
import matplotlib.pyplot as plt
import pandas as pd
from PIL import Image
import streamlit as st
from streamlit_drawable_canvas import st_canvas
import cv2
import torch
from skorch import NeuralNetClassifier
from torch import nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
mnist = fetch_openml('mnist_784', as_frame=False, cache=False)
X = mnist.data.astype('float32')
y = mnist.target.astype('int64')
X /= 255.0


#device = 'cuda' if torch.cuda.is_available() else 'cpu'
XCnn = X.reshape(-1, 1, 28, 28)
XCnn_train, XCnn_test, y_train, y_test = train_test_split(XCnn, y, test_size=0.25, random_state=42)

from PIL import Image
import torchvision.transforms as transforms
class Cnn(nn.Module):
    def __init__(self, dropout=0.5):
        super(Cnn, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3)
        self.conv2_drop = nn.Dropout2d(p=dropout)
        self.fc1 = nn.Linear(1600, 100) # 1600 = number channels * width * height
        self.fc2 = nn.Linear(100, 10)
        self.fc1_drop = nn.Dropout(p=dropout)

    def forward(self, x):
        x = torch.relu(F.max_pool2d(self.conv1(x), 2))
        x = torch.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))

        # flatten over channel, height and width = 1600
        x = x.view(-1, x.size(1) * x.size(2) * x.size(3))

        x = torch.relu(self.fc1_drop(self.fc1(x)))
        x = torch.softmax(self.fc2(x), dim=-1)
        return x
torch.manual_seed(0)


# # Create an instance of your model
# model = NeuralNetClassifier(
#     Cnn,
#     max_epochs=10,
#     lr=0.002,
#     optimizer=torch.optim.Adam,
#     device=device,
import streamlit as st
from st_canvas import st_canvas
import torch
from PIL import Image
import cv2
import numpy as np
from your_model_module import Cnn  # Import your model architecture

# Create an instance of your model (Cnn model)
model = Cnn()

# Specify the path to the saved model weights
model_weights_path = 'model_weights.pth'

# Load the model weights onto a CPU device (if you want to use the CPU)
model.load_state_dict(torch.load(model_weights_path, map_location=torch.device('cpu')))

# Set the model to evaluation mode for inference
model.eval()

stroke_width = st.sidebar.slider("Stroke width: ", 1, 35, 32)
stroke_color = st.sidebar.color_picker("Stroke color hex: ")
bg_color = st.sidebar.color_picker("Background color hex: ", "#eee")
bg_image = st.sidebar.file_uploader("Background image:", type=["png", "jpg"])
drawing_mode = st.sidebar.selectbox(
    "Drawing tool:", ("freedraw", "line", "rect", "circle", "transform", "polygon")
)
realtime_update = st.sidebar.checkbox("Update in realtime", True)

# Create a canvas component
canvas_result = st_canvas(
    fill_color="rgba(255, 165, 0, 0.3)",  # Fixed fill color with some opacity
    stroke_width=stroke_width,
    stroke_color=stroke_color,
    background_color=bg_color,
    background_image=Image.open(bg_image) if bg_image else None,
    update_streamlit=realtime_update,
    height=300,
    drawing_mode=drawing_mode,
    display_toolbar=st.sidebar.checkbox("Display toolbar", True),
    key="full_app",
)

# Do something interesting with the image data and paths
if canvas_result.image_data is not None:
    image = canvas_result.image_data
    image1 = image.copy()
    image1 = image1.astype('uint8')
    image1 = cv2.cvtColor(image1, cv2.COLOR_BGR2GRAY)
    image1 = cv2.resize(image1, (28, 28))
    st.image(image1)

    # Convert the image for prediction (assuming image1 is in the right format)
    image1 = image1[np.newaxis, np.newaxis, ...]  # Add batch and channel dimensions

    # Perform prediction using the pre-trained model
    with torch.no_grad():
        tensor_image = torch.tensor(image1, dtype=torch.float32)
        prediction = model(tensor_image)

    # Display the predicted class
    predicted_class = prediction.argmax().item()
    st.title(f"Predicted Class: {predicted_class}")