"""
Search over cube states guided by a learned distance estimate.

The policy beam search in serve.py explores sequences the *policy* ranks highly.
This explores states the *value function* thinks are closest to solved, which is
a different and stronger thing: it can prefer a move the policy never considered,
because the ranking comes from an independent estimate rather than from the
policy's own confidence.

This is DeepCubeA's shape in miniature -- batch-weighted search with the engine
as the transition function and a network supplying the heuristic. The weighting
`h + lambda * g` trades solution length against search effort: lambda=0 is pure
greedy on the heuristic (fast, longer solutions), higher lambda behaves more like
uniform-cost search.

Two kinds of value model can drive this, and `load_value_model` tells them apart
by what the checkpoint carries:

- **supervised** (`train_value.py`): a distribution over distances, trained on
  Kociemba lengths. Those are upper bounds rather than true optimal distances, so
  the heuristic is inadmissible and no shortest-solution guarantee holds.
- **value iteration** (`train_value_iteration.py`): a scalar, bootstrapped from
  the solved state with no solver involved.

Either way the engine verifies that whatever comes back actually solves, and
finding *a* solution at depth 20 is the open problem, not finding the shortest.
"""
import sys
from pathlib import Path

import torch

sys.path.insert(0, str(Path(__file__).parent))
import cube_tokenizer as T
from gen_data import SOLVED, MOVES, apply_move


def load_value_model(path, device):
    """Loads either value-model flavour, detected from the checkpoint's own keys.

    The two trainers save under different keys ("state_dict" vs "model"), which is
    the only reliable discriminator: both carry identical hidden/layers/heads, so
    shape alone cannot distinguish a 27-way distribution head from a scalar one
    without guessing. The flavour is recorded on the model so `heuristic` reads
    each correctly -- feeding a scalar head through softmax would silently return
    a constant and turn the search into an untargeted walk.
    """
    ckpt = torch.load(path, map_location=device)
    if "state_dict" in ckpt:
        from train_value import ValueNet
        model = ValueNet(ckpt["hidden"], ckpt["layers"], ckpt["heads"]).to(device)
        model.load_state_dict(ckpt["state_dict"])
        kind = "distribution"
    elif "model" in ckpt:
        from train_value_iteration import ValueNet as ValueNetScalar
        model = ValueNetScalar(ckpt["hidden"], ckpt["layers"], ckpt["heads"]).to(device)
        model.load_state_dict(ckpt["model"])
        kind = "scalar"
    else:
        raise ValueError(f"{path}: no 'state_dict' or 'model' key; not a value "
                         f"checkpoint (keys: {sorted(ckpt)[:8]})")
    model.eval()
    model.value_kind = kind
    return model


@torch.no_grad()
def heuristic(model, states, device, batch_size=1024):
    """Estimated distance-to-solved for each state.

    For a distribution head, the expectation over the predicted distribution
    rather than the argmax: the extra resolution matters when ranking siblings
    that all round to the same integer, which is most of the frontier. A scalar
    head already has that resolution and is read directly.
    """
    kind = getattr(model, "value_kind", "distribution")
    out = []
    for i in range(0, len(states), batch_size):
        chunk = states[i:i + batch_size]
        ids = torch.tensor([T.encode_state(s) for s in chunk], device=device)
        raw = model(ids).float()
        if kind == "scalar":
            out.extend(raw.clamp(min=0.0).tolist())
        else:
            probs = torch.softmax(raw, dim=-1)
            values = torch.arange(probs.shape[-1], device=device, dtype=torch.float32)
            out.extend((probs * values).sum(-1).tolist())
    return out


def solve_value_beam(model, state, device, beam=64, max_depth=26, lam=0.0,
                     batch_size=1024):
    """Beam search over states, ranked by `h + lam * g`.

    Returns the move list that solves, or [] if the beam is exhausted. Visited
    states are tracked globally: revisiting one cannot help, and on a group this
    symmetric the frontier collapses onto duplicates quickly without it.
    """
    if state == SOLVED:
        return []
    frontier = [(state, [])]
    seen = {state}

    for _ in range(max_depth):
        children, paths = [], []
        for s, path in frontier:
            for mv in MOVES:
                nxt = apply_move(s, mv)
                if nxt in seen:
                    continue
                if nxt == SOLVED:
                    return path + [mv]
                seen.add(nxt)
                children.append(nxt)
                paths.append(path + [mv])
        if not children:
            return []
        h = heuristic(model, children, device, batch_size)
        scored = sorted(zip(h, children, paths), key=lambda x: x[0] + lam * len(x[2]))
        frontier = [(c, p) for _, c, p in scored[:beam]]
    return []