from flask import Flask, request, jsonify, render_template
import io, base64
from PIL import Image, ImageOps
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
import numpy as np


# --------- Your Model definition ---------
class MNIST(nn.Module):
    def __init__(self, num_classes: int=10):
        super().__init__()
        #1st conv layer
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.bn1 = nn.BatchNorm2d(num_features=32)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)


        #2nd conv layer
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        self.bn2 = nn.BatchNorm2d(num_features=64)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)


        #dynamic feature computation
        with torch.no_grad():
            dummy = torch.zeros(1,1,28,28)
            dummy = self.pool1(F.relu(self.bn1(self.conv1(dummy))))
            dummy = self.pool2(F.relu(self.bn2(self.conv2(dummy))))
            flattened = int(torch.prod(torch.tensor(dummy.shape[1:])))


        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(in_features=flattened, out_features=64)
        self.bn_l = nn.BatchNorm1d(64)
        self.dropout = nn.Dropout(0.4)
        self.out = nn.Linear(in_features=64, out_features=num_classes)


    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = F.relu(x)
        x = self.pool1(x)


        x = self.conv2(x)
        x = self.bn2(x)
        x = F.relu(x)
        x = self.pool2(x)


        x = self.flatten(x)


        x = self.fc1(x)
        x = self.bn_l(x)
        x = self.dropout(x)
        x = self.out(x)
        return x


# -------------------- Load model -------------------------------------
MODEL_PATH = "mnist_model.pth"
device = torch.device('cpu')
model = MNIST()
state = torch.load(MODEL_PATH, map_location=device)
model.load_state_dict(state)
model.to(device)
model.eval()
app = Flask(__name__)

# -------------------- Compatibility for Pillow resampling --------------
# Image.ANTIALIAS was removed in newer Pillow; use Resampling.LANCZOS with fallback.
RESAMPLE_FILTER = getattr(Image, 'Resampling', Image).LANCZOS

# -------------------- Preprocessing transform -------------------------
# Ensure the same normalization you used during training. Adjust if needed.
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])


@app.route('/')
def index():
    return render_template("index.html")


@app.route('/predict', methods=['POST'])
def predict():
    try:
        data = request.get_json()
        img_b64 = data.get('image', '')
        header, encoded = img_b64.split(',', 1)
        img_bytes = base64.b64decode(encoded)
        img = Image.open(io.BytesIO(img_bytes)).convert('L')

        # invert so drawn black strokes become bright as in MNIST preprocessing
        img = ImageOps.invert(img)

        # center-crop/fit to 28x28 using the modern resample filter
        img = ImageOps.fit(img, (28,28), method=RESAMPLE_FILTER)

        # apply torchvision transforms
        tensor = transform(img).unsqueeze(0)  # shape: 1x1x28x28

        with torch.no_grad():
            outputs = model(tensor.to(device))
            probs = torch.softmax(outputs, dim=1).cpu().numpy().flatten().tolist()
            pred = int(np.argmax(probs))
        return jsonify({'pred': pred, 'probs': probs})
    except Exception as e:
        return jsonify({'error': str(e)})


if __name__ == '__main__':
    app.run(debug=True)