import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

# EMNIST Letters: labels 1-26 (A-Z), but torchvision returns 1-26
# We remap to 0-25 for a clean 26-class output

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

print("Downloading EMNIST Letters dataset...")
train_set = torchvision.datasets.EMNIST(
    root='./emnist_data', split='letters', train=True,
    download=True, transform=transform
)
test_set = torchvision.datasets.EMNIST(
    root='./emnist_data', split='letters', train=False,
    download=True, transform=transform
)

train_loader = DataLoader(train_set, batch_size=128, shuffle=True, num_workers=0)
test_loader = DataLoader(test_set, batch_size=128, shuffle=False, num_workers=0)

class LetterCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 32, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(32, 32, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Dropout2d(0.25),

            nn.Conv2d(32, 64, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 64, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Dropout2d(0.25),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(64 * 7 * 7, 128),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(128, 26),
        )

    def forward(self, x):
        x = self.features(x)
        x = self.classifier(x)
        return x

model = LetterCNN()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Training on {device}")
model.to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)

# EMNIST letters labels are 1-26, remap to 0-25
for epoch in range(10):
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0
    for images, labels in train_loader:
        labels = labels - 1  # remap 1-26 to 0-25
        images, labels = images.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        _, predicted = outputs.max(1)
        total += labels.size(0)
        correct += predicted.eq(labels).sum().item()

    scheduler.step()
    acc = 100.0 * correct / total
    print(f"Epoch {epoch+1}/10 - loss: {running_loss/len(train_loader):.4f}, train acc: {acc:.2f}%")

# Test accuracy
model.eval()
correct = 0
total = 0
with torch.no_grad():
    for images, labels in test_loader:
        labels = labels - 1
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        _, predicted = outputs.max(1)
        total += labels.size(0)
        correct += predicted.eq(labels).sum().item()

print(f"Test accuracy: {100.0 * correct / total:.2f}%")

# Export to ONNX
model.eval()
model.to('cpu')
dummy = torch.randn(1, 1, 28, 28)
onnx_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "public", "models", "emnist-letters.onnx")

torch.onnx.export(
    model, dummy, onnx_path,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input': {0: 'batch'}, 'output': {0: 'batch'}},
    opset_version=13
)

import os
size_kb = os.path.getsize(onnx_path) / 1024
print(f"Saved ONNX model to {onnx_path} ({size_kb:.1f} KB)")

# Verify with onnx
import onnx
m = onnx.load(onnx_path)
onnx.checker.check_model(m)
print("ONNX model verified OK")
print(f"Input: {m.graph.input[0].type.tensor_type.shape}")
print(f"Output: {m.graph.output[0].type.tensor_type.shape}")