"""Small MNIST CNN on 16x16 cells, used to read generated Sudoku images back into grids."""
import torch
import torch.nn as nn
import torch.nn.functional as F

from .data import cells


def conv_bn(cin, cout):
    return [nn.Conv2d(cin, cout, 3, padding=1), nn.BatchNorm2d(cout), nn.ReLU()]


class DigitCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.f = nn.Sequential(
            *conv_bn(1, 64), *conv_bn(64, 64), nn.MaxPool2d(2),
            *conv_bn(64, 128), *conv_bn(128, 128), nn.MaxPool2d(2),
            *conv_bn(128, 256), nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Dropout(0.3), nn.Linear(256, 10))

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


@torch.no_grad()
def read_grids(clf, images):
    """(B, 1, 144, 144) images -> (B, 9, 9) predicted digits in 1..9."""
    c = cells(images)
    B = c.shape[0]
    logits = clf(c.reshape(B * 81, 1, *c.shape[-2:]))
    return (logits[:, 1:].argmax(-1) + 1).reshape(B, 9, 9)


def augment(x):
    """Random small rotation, scale and shift (cell digits are not perfectly centred)."""
    B = len(x)
    a = (torch.rand(B, device=x.device) - 0.5) * 0.4
    s = 1 + (torch.rand(B, device=x.device) - 0.5) * 0.2
    theta = torch.zeros(B, 2, 3, device=x.device)
    theta[:, 0, 0], theta[:, 0, 1] = s * a.cos(), -s * a.sin()
    theta[:, 1, 0], theta[:, 1, 1] = s * a.sin(), s * a.cos()
    theta[:, :, 2] = (torch.rand(B, 2, device=x.device) - 0.5) * 0.25
    grid = F.affine_grid(theta, x.shape, align_corners=False)
    return F.grid_sample(x + 1, grid, align_corners=False) - 1  # pad with background (-1)


def train(train_bank, test_bank, device, epochs=10, batch=128, lr=2e-3, seed=0):
    """Train on DigitBank tensors (all 10 classes). Returns (model, test accuracy)."""
    torch.manual_seed(seed)
    def flat(bank):
        x = bank.images
        y = torch.arange(10, device=x.device)[:, None].expand(x.shape[:2])
        return x.reshape(-1, 1, *x.shape[-2:]).to(device), y.reshape(-1).to(device)
    x, y = flat(train_bank)
    xt, yt = flat(test_bank)
    model = DigitCNN().to(device)
    opt = torch.optim.AdamW(model.parameters(), lr, weight_decay=5e-4)
    sched = torch.optim.lr_scheduler.OneCycleLR(opt, lr, total_steps=epochs * ((len(x) + batch - 1) // batch))
    for ep in range(epochs):
        model.train()
        perm = torch.randperm(len(x), device=device)
        for i in range(0, len(x), batch):
            idx = perm[i:i + batch]
            loss = F.cross_entropy(model(augment(x[idx])), y[idx], label_smoothing=0.05)
            opt.zero_grad()
            loss.backward()
            opt.step()
            sched.step()
    model.eval()
    with torch.no_grad():
        acc = sum((model(xt[i:i + 2048]).argmax(-1) == yt[i:i + 2048]).sum().item() for i in range(0, len(xt), 2048)) / len(xt)
    return model, acc