"""Release 1 label-free, single-process Sudoku evaluation with raw argmax only."""

from __future__ import annotations

import inspect
import json
import os
from pathlib import Path

import torch
from omegaconf import DictConfig, OmegaConf
from softasp_domain.registry import get_domain
from softasp_solver.registry import get_model
from torch.utils.data import DataLoader

from softasp.instantiations.checkpoint import read_checkpoint, validate_checkpoint
from softasp.sudoku_data import select_hdf5
from softasp.sudoku_reporting import file_sha256, write_raw_report


@torch.no_grad()
def evaluate_raw(model, loader, destination: Path, provenance: dict, device: torch.device) -> dict:
    """Labels go to scoring only; all batches use the model's label-free interface."""
    if model.oracle_perception:
        raise ValueError("Release 1 prohibits oracle perception")
    model.to(device).eval()
    predictions, labels, masks = [], [], []
    for batch in loader:
        output = model.predict(batch["images"].to(device), batch["puzzle_mask"].to(device))
        logits = output["logits"]
        if logits.shape != (len(batch["images"]), 81, 9) or not torch.isfinite(logits).all():
            raise ValueError("Expected finite raw Sudoku logits with shape (B, 81, 9)")
        predictions.append(logits.argmax(dim=-1).cpu())
        labels.append(batch["labels"].cpu())
        masks.append(batch["puzzle_mask"].cpu())
    if not predictions:
        raise ValueError("Cannot evaluate an empty Sudoku split")
    return write_raw_report(
        destination, torch.cat(predictions), torch.cat(labels), torch.cat(masks), provenance
    )


def evaluate_release(cfg: DictConfig, device: str = "auto") -> dict:
    """Use the pinned dataset and strict checkpoint loading; never invoke a decoder."""
    if cfg.dataset.name != "sudoku" or cfg.model.name != "softasp_sudoku":
        raise ValueError("Release 1 requires the Sudoku AS2 model")
    if cfg.model.params.get("oracle_perception", False):
        raise ValueError("Release 1 prohibits oracle perception")
    if int(os.environ.get("WORLD_SIZE", "1")) != 1:
        raise ValueError("Release 1 evaluation requires one process")
    checkpoint_value = cfg.evaluator.get("checkpoint")
    if not checkpoint_value:
        raise ValueError("evaluator.checkpoint must be set for evaluation")
    checkpoint_path = Path(checkpoint_value).resolve()
    hdf5_path = select_hdf5(cfg)
    splits = list(cfg.dataset.splits)
    if (
        not splits
        or len(set(splits)) != len(splits)
        or any(split not in ("train", "val", "test") for split in splits)
    ):
        raise ValueError("Expected unique nonempty Sudoku splits: train, val, test")
    destination = Path(cfg.paths.output_dir) / "eval"
    if destination.exists():
        raise FileExistsError(f"Use a fresh evaluation output directory: {destination}")
    batch_size = cfg.evaluator.get("batch_size_per_gpu", 64)
    workers = cfg.evaluator.get("num_workers", 0)
    if batch_size < 1 or workers < 0:
        raise ValueError("Expected positive batch size and nonnegative worker count")
    resolved_device = torch.device(
        ("cuda" if torch.cuda.is_available() else "cpu") if device == "auto" else device
    )
    checkpoint = read_checkpoint(checkpoint_path)
    if cfg.get("instantiation") is not None:
        validate_checkpoint(
            cfg,
            checkpoint,
            legacy_model="softasp_sudoku",
            legacy_constraints={"_target_": "softasp_solver.sudoku.constraints.TPOperatorLoss"},
        )
    checkpoint_cfg = checkpoint.get("config", {})
    if checkpoint_cfg.get("model", {}).get("params", {}).get("oracle_perception", False):
        raise ValueError("Release 1 prohibits oracle checkpoint configurations")
    model = get_model(cfg.model.name, **dict(cfg.model.params))
    if cfg.get("instantiation") is not None:
        from hydra.utils import instantiate

        model.constraint_loss = instantiate(cfg.instantiation.constraints)
    model.load_state_dict(checkpoint["model_state_dict"], strict=True)
    identity = OmegaConf.to_container(cfg.resolved_dataset, resolve=True)
    provenance = {
        "checkpoint": str(checkpoint_path),
        "checkpoint_sha256": file_sha256(checkpoint_path),
        "checkpoint_config": checkpoint_cfg,
        "evaluation_config": OmegaConf.to_container(cfg, resolve=True),
        "dataset_identity": identity,
        "device": str(resolved_device),
        "precision": "float32; autocast disabled",
        "torch_version": str(torch.__version__),
        "postprocessing": "none",
        "oracle_perception": False,
        "source_sha256": {
            "sudoku_evaluation.py": file_sha256(Path(__file__)),
            "sudoku_reporting.py": file_sha256(Path(__file__).with_name("sudoku_reporting.py")),
            "softasp_sudoku.py": file_sha256(Path(inspect.getfile(type(model)))),
        },
    }
    dataset_cls = get_domain("sudoku").dataset_cls
    results = {}
    for split in splits:
        dataset = dataset_cls(hdf5_path, split=split)
        loader = DataLoader(dataset, batch_size=batch_size, num_workers=workers, shuffle=False)
        results[split] = evaluate_raw(
            model, loader, destination / split, {**provenance, "split": split}, resolved_device
        )
    (destination / "results.json").write_text(json.dumps(results, indent=2, allow_nan=False) + "\n")
    (destination / "dataset-identity.json").write_text(json.dumps(identity, indent=2) + "\n")
    return results