"""
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.

The heuristic is Kociemba length, an upper bound rather than true optimal
distance, so it is inadmissible and the search is not guaranteed to return a
shortest solution. That is fine: 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):
    from train_value import ValueNet
    ckpt = torch.load(path, map_location=device)
    model = ValueNet(ckpt["hidden"], ckpt["layers"], ckpt["heads"]).to(device)
    model.load_state_dict(ckpt["state_dict"])
    model.eval()
    return model


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

    Uses 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.
    """
    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)
        probs = torch.softmax(model(ids).float(), 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 []