"""Render Sudoku grids as images of MNIST digits and build training/eval batches.

Image: 1 x 144 x 144 in [-1, 1] (9 x 9 cells of 16 px, the paper's training resolution).
Conditioning: the solution image with non-clue cells blacked out, plus a pixel mask of clue cells.
"""
import numpy as np
import torch
import torch.nn.functional as F

from . import sudoku

CELL = 16
IMG = 9 * CELL


class DigitBank:
    """MNIST digits resized to CELL x CELL, grouped by class: images[d] is (N_d, CELL, CELL) in [-1, 1]."""

    def __init__(self, train=True, data_dir="data", images=None, labels=None):
        if images is None:
            from torchvision.datasets import MNIST
            ds = MNIST(data_dir, train=train, download=True)
            images, labels = ds.data, ds.targets
        x = F.interpolate(images[:, None].float() / 255.0, size=(CELL, CELL), mode="area")[:, 0]
        x = x * 2 - 1
        n = min(int((labels == d).sum()) for d in range(10))
        self.images = torch.stack([x[labels == d][:n] for d in range(10)])  # (10, n, CELL, CELL)

    def to(self, device):
        self.images = self.images.to(device)
        return self


def render(grids, bank, generator=None):
    """grids: (B, 9, 9) long tensor of digits 1-9 (0 renders as black). Returns (B, 1, IMG, IMG)."""
    B = grids.shape[0]
    flat = grids.reshape(B, 81).to(bank.images.device)
    idx = torch.randint(bank.images.shape[1], (B, 81), generator=generator, device="cpu").to(flat.device)
    cells = bank.images[flat, idx]  # (B, 81, CELL, CELL)
    cells = torch.where((flat == 0)[..., None, None], torch.full_like(cells, -1.0), cells)
    return cells.reshape(B, 9, 9, CELL, CELL).permute(0, 1, 3, 2, 4).reshape(B, 1, IMG, IMG)


def cells(images):
    """Inverse of the layout in render: (B, 1, IMG, IMG) -> (B, 81, CELL, CELL)."""
    B = images.shape[0]
    return images.reshape(B, 9, CELL, 9, CELL).permute(0, 1, 3, 2, 4).reshape(B, 81, CELL, CELL)


def pixel_mask(puzzles):
    """(B, 9, 9) puzzles -> (B, 1, IMG, IMG) float mask, 1 on clue cells."""
    m = (puzzles > 0).float()
    return m.repeat_interleave(CELL, 1).repeat_interleave(CELL, 2)[:, None]


def make_batch(puzzles, solutions, bank, generator=None):
    """Render a batch. Clue pixels in cond are the same digit images as in x0."""
    puzzles, solutions = torch.as_tensor(puzzles).long(), torch.as_tensor(solutions).long()
    x0 = render(solutions, bank, generator)
    mask = pixel_mask(puzzles).to(x0.device)
    cond = torch.where(mask > 0, x0, torch.full_like(x0, -1.0))
    return dict(x0=x0, cond=cond, mask=mask, puzzles=puzzles, solutions=solutions)


class PuzzleSource:
    """Samples (puzzles, solutions) for 'extreme' (TRM csv), 'hard' (0-26 random clues of a real solution)
    or 'solutions' (no clues; Painter data). Solutions always come from the sudoku-extreme split, so train and
    test grids are disjoint; training samples get a random validity-preserving augmentation."""

    def __init__(self, mode, split="train", data_dir="data", seed=0):
        if mode not in ("extreme", "hard", "solutions"):
            raise ValueError(mode)
        self.mode, self.split, self.rng = mode, split, np.random.default_rng(seed)
        self.puzzles, self.solutions = sudoku.load_extreme(split, data_dir)

    def sample(self, n):
        i = self.rng.integers(len(self.solutions), size=n)
        p, s = sudoku.augment(self.rng, self.puzzles[i], self.solutions[i])
        if self.mode == "hard":
            p = np.stack([sudoku.hard_puzzle(self.rng, g)[0] for g in s])
        elif self.mode == "solutions":
            p = np.zeros_like(p)
        return p, s

    def fixed(self, n, seed=1234):
        """A deterministic evaluation set (no augmentation)."""
        rng = np.random.default_rng(seed)
        i = rng.choice(len(self.solutions), size=n, replace=False)
        p, s = self.puzzles[i].astype(np.int64), self.solutions[i].astype(np.int64)
        if self.mode == "hard":
            p = np.stack([sudoku.hard_puzzle(rng, g)[0] for g in s])
        return p, s