"""Label-free raw prediction for one or more Visual Sudoku boards."""

from pathlib import Path

import numpy as np
import torch
from omegaconf import DictConfig
from softasp_solver.sudoku.softasp_sudoku import SoftASPSudoku

from softasp.instantiations.checkpoint import read_checkpoint, validate_checkpoint
from softasp.sudoku_reporting import file_sha256


def load_model(
    checkpoint: Path, device: str = "auto", *, cfg: DictConfig | None = None
) -> tuple[SoftASPSudoku, torch.device]:
    """Load public inference weights; optimizer state and labels are unnecessary."""
    target = torch.device(
        ("cuda" if torch.cuda.is_available() else "cpu") if device == "auto" else device
    )
    saved = read_checkpoint(checkpoint)
    if cfg is not None:
        validate_checkpoint(
            cfg,
            saved,
            legacy_model="softasp_sudoku",
            legacy_constraints={"_target_": "softasp_solver.sudoku.constraints.TPOperatorLoss"},
        )
    model_config = saved["config"]["model"]
    if model_config["name"] != "softasp_sudoku":
        raise ValueError("Expected a Sudoku AS2 checkpoint")
    params = dict(model_config["params"])
    if params.get("oracle_perception", False):
        raise ValueError("Oracle checkpoints are not supported")
    model = SoftASPSudoku(**params)
    if cfg is not None:
        from hydra.utils import instantiate

        model.constraint_loss = instantiate(cfg.instantiation.constraints)
    model.load_state_dict(saved["model_state_dict"], strict=True)
    return model.to(target).eval(), target


@torch.no_grad()
def predict_file(
    checkpoint: Path, input_path: Path, device: str = "auto", *, cfg: DictConfig | None = None
) -> dict:
    """Read images and clue positions only; output raw digits 1 through 9."""
    with np.load(input_path, allow_pickle=False) as data:
        images = data["images"].copy()
        mask = data["puzzle_mask"].copy()
    if images.ndim != 4 or images.shape[1:] != (252, 252, 1) or len(images) == 0:
        raise ValueError("images must have shape (N, 252, 252, 1), N > 0")
    if images.dtype != np.uint8 and not (
        images.dtype == np.float32
        and np.isfinite(images).all()
        and np.all((images >= 0) & (images <= 1))
    ):
        raise ValueError("images must be uint8 [0,255] or float32 [0,1]")
    if mask.shape != (len(images), 81) or not np.isin(mask, [0, 1]).all():
        raise ValueError("puzzle_mask must have shape (N,81) and contain only 0/1")
    model, target = load_model(checkpoint, device, cfg=cfg)
    predictions = []
    # Bound inference batches without requiring all images to fit on the device.
    for start in range(0, len(images), 32):
        out = model.predict(
            torch.from_numpy(images[start : start + 32]).to(target),
            torch.from_numpy(mask[start : start + 32].astype(np.int64)).to(target),
        )
        logits = out["logits"]
        if logits.shape[1:] != (81, 9) or not torch.isfinite(logits).all():
            raise ValueError("Expected finite raw Sudoku logits")
        predictions.append(logits.argmax(-1).cpu())
    digits = torch.cat(predictions).numpy().reshape(-1, 9, 9) + 1
    expected = np.arange(1, 10)
    valid = []
    for board in digits:
        boxes = board.reshape(3, 3, 3, 3).transpose(0, 2, 1, 3).reshape(9, 9)
        valid.append(
            bool(
                np.all(np.sort(board, axis=1) == expected)
                and np.all(np.sort(board.T, axis=1) == expected)
                and np.all(np.sort(boxes, axis=1) == expected)
            )
        )
    return {
        "decoder": "raw_argmax",
        "digits": digits.tolist(),
        "structurally_valid": valid,
        "checkpoint_sha256": file_sha256(checkpoint),
        "input_sha256": file_sha256(input_path),
    }