from pathlib import Path

import numpy as np
import torch
from PIL import Image, ImageOps
from torch import nn


class DigitCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(64 * 7 * 7, 128),
            nn.ReLU(),
            nn.Dropout(0.25),
            nn.Linear(128, 10),
        )

    def forward(self, x):
        return self.net(x)


ARTIFACT = torch.load(
    Path(__file__).with_name("mnist_cnn.pt"),
    map_location="cpu",
    weights_only=False,
)
MODEL = DigitCNN()
MODEL.load_state_dict(ARTIFACT["model_state_dict"])
MODEL.eval()
ACCURACY = ARTIFACT["accuracy"]


def _to_image(image):
    if image is None:
        raise ValueError("Please draw or upload a digit image.")

    if isinstance(image, dict):
        image = image.get("composite") or image.get("image") or image.get("background")

    if isinstance(image, (str, Path)):
        return Image.open(image)

    return Image.fromarray(np.asarray(image))


def preprocess_image(image):
    img = _to_image(image).convert("L")
    arr = np.asarray(img).astype("float32")

    # Convert common black-on-white handwriting into MNIST-style white-on-black.
    if arr.mean() > 127:
        arr = 255.0 - arr

    arr[arr < 25] = 0.0

    rows, cols = np.where(arr > 0)
    if len(rows) == 0 or len(cols) == 0:
        raise ValueError("No digit-like strokes were detected.")

    cropped = arr[rows.min() : rows.max() + 1, cols.min() : cols.max() + 1]
    cropped_img = Image.fromarray(cropped.astype("uint8"), mode="L")
    cropped_img.thumbnail((20, 20), Image.Resampling.LANCZOS)

    canvas = Image.new("L", (28, 28), 0)
    left = (28 - cropped_img.width) // 2
    top = (28 - cropped_img.height) // 2
    canvas.paste(cropped_img, (left, top))

    preview = np.asarray(canvas).astype("float32") / 255.0
    features = torch.from_numpy(preview.reshape(1, 1, 28, 28)).float()
    return features, preview


def predict_digit(image):
    features, preview = preprocess_image(image)
    with torch.no_grad():
        probabilities = torch.softmax(MODEL(features), dim=1).numpy()[0]

    digit = int(np.argmax(probabilities))
    confidence = float(probabilities[digit])
    ranking = {
        str(index): float(score)
        for index, score in sorted(
            enumerate(probabilities),
            key=lambda item: item[1],
            reverse=True,
        )
    }
    return str(digit), confidence, ranking, preview