"""
Trains a value function by approximate value iteration -- no solver in the loop.

## Why this replaces the supervised version

`train_value.py` learns distance-to-solved from Kociemba's solution lengths.
That reached MAE 1.333 and, with beam search, solved canonical depth 10 at 100%
-- but depth 15/20 stayed at 28%/16%, and the reason turned out to be the
labels, not the model:

    Kociemba length on uniform random states: 47% are exactly 21, 79% in 20-22.

They are two-phase *upper bounds*, and they cluster. Above distance ~15 there is
almost no label variance to learn from, so nothing can learn to tell a
distance-19 state from a distance-21 one. Measured directly: in the 15+ band
(85.5% of that holdout) the supervised model scored MAE 1.44 against 0.57 for a
constant that always answers 21 -- 2.5x worse than ignoring the cube. Search
worked where the heuristic was informative and wandered where it was not.

More capacity, steps or data cannot fix a label that carries no information.

## The method (DeepCubeA's DAVI, Agostinelli et al., Nature MI 2019)

Bootstrap the targets instead of labelling them. For a state s, its cost-to-go
is one move plus the best of its successors:

    J(s) = 0                                  if s is solved
    J(s) = min over a of (1 + J_target(s'))    otherwise

Start from a randomly initialised net and iterate. Accuracy radiates outward
from the solved state: once distance-1 states are right, distance-2 states
become learnable, and so on. The targets are real distances the net discovers,
not a solver's upper bounds, so they stay informative at every distance -- which
is precisely what the supervised labels stopped being past 15.

Two details that make it converge rather than chase itself:

- **A target network.** Targets come from a frozen copy, refreshed only when the
  training loss falls below `--update-threshold`. Bootstrapping against the
  live net is the standard way this diverges.
- **Regression, not classification.** `train_value.py` argues for
  classification because a squared-error model collapses toward the 18-21 mode
  of the *Kociemba label* distribution. That argument does not transfer: these
  targets are bootstrapped and spread across the whole range, and DAVI's `min`
  needs a scalar it can compare and add 1 to.

## What to watch

Not the loss -- it is measured against a moving target, so it says little about
quality. **Watch the greedy solve rate**, reported per canonical depth bucket.
That is the heuristic doing the actual job: descend to the lowest-valued child
each step and see whether it reaches solved. It needs no labels, which is the
point, since out here there is no ground-truth distance to compare against.
"""
import argparse, json, sys, time
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F

sys.path.insert(0, str(Path(__file__).parent))
import cube_ops as C
import cube_tokenizer as T


class ValueNet(torch.nn.Module):
    """Llama encoder over the 54-facelet state, mean-pooled, to one scalar."""

    def __init__(self, hidden, layers, heads):
        super().__init__()
        from transformers import LlamaConfig, LlamaModel
        cfg = LlamaConfig(
            vocab_size=T.VOCAB_SIZE, hidden_size=hidden, intermediate_size=hidden * 4,
            num_hidden_layers=layers, num_attention_heads=heads,
            num_key_value_heads=heads, max_position_embeddings=T.MAX_SEQ,
            pad_token_id=T.PAD,
        )
        self.encoder = LlamaModel(cfg)
        self.head = torch.nn.Linear(hidden, 1)

    def forward(self, ids):
        h = self.encoder(input_ids=ids).last_hidden_state.mean(dim=1)
        return self.head(h).squeeze(-1)


# Colour-index -> token-id lookup, taken from the tokenizer so the array path and
# T.encode_state produce byte-identical sequences (verified in tests).
_COLOR_IDS = np.array(T.COLOR_IDS, dtype=np.int64)


# The tokenizer's encode_state wraps a state in BOS/SEP for the seq2seq policy.
# Reused unchanged so a state has one encoding across every model here.
def encode_batch(states: np.ndarray, device):
    """(N, 54) uint8 facelet indices -> (N, MAX) token ids on `device`."""
    toks = _COLOR_IDS[states.astype(np.int64)]
    bos = np.full((len(states), 1), T.BOS, dtype=np.int64)
    sep = np.full((len(states), 1), T.SEP, dtype=np.int64)
    return torch.from_numpy(np.concatenate([bos, toks, sep], axis=1)).to(device)


@torch.no_grad()
def value_of(net, states: np.ndarray, device, batch_size=4096, amp=True):
    """Scalar value per state, in eval mode, batched."""
    net.eval()
    out = np.empty(len(states), dtype=np.float32)
    for i in range(0, len(states), batch_size):
        chunk = states[i:i + batch_size]
        ids = encode_batch(chunk, device)
        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=amp and device == "cuda"):
            v = net(ids)
        out[i:i + len(chunk)] = v.float().cpu().numpy()
    net.train()
    # A solved state is worth 0 by definition, never by prediction: the search
    # and the targets both lean on that being exact.
    out[C.is_solved(states)] = 0.0
    return np.maximum(out, 0.0)


def davi_targets(target_net, states: np.ndarray, device, amp=True):
    """J(s) = min over the 18 moves of (1 + J_target(s')), and 0 where s is solved."""
    kids = C.all_children(states)                      # (N, 18, 54)
    flat = kids.reshape(-1, 54)
    v = value_of(target_net, flat, device, amp=amp).reshape(len(states), C.NUM_MOVES)
    # A move onto the solved state costs exactly 1 and nothing more, regardless
    # of what the net thinks of the solved state.
    solved_kid = C.is_solved(flat).reshape(len(states), C.NUM_MOVES)
    v = np.where(solved_kid, 0.0, v)
    tgt = (1.0 + v).min(axis=1)
    tgt[C.is_solved(states)] = 0.0
    return tgt.astype(np.float32)


@torch.no_grad()
def greedy_solve_rate(net, device, depths, n_per_depth, rng, max_steps=40, amp=True):
    """Fraction of scrambles solved by descending to the lowest-valued child.

    The honest metric for this trainer: no labels involved, and it is the
    heuristic doing the job search will ask of it. Every scramble at every depth
    is stepped in one batch, so this costs a handful of forward passes rather
    than one per puzzle.
    """
    out = {}
    for d in depths:
        states = _walk_exact(n_per_depth, d, rng)
        done = C.is_solved(states)
        for _ in range(max_steps):
            live = np.where(~done)[0]
            if not len(live):
                break
            kids = C.all_children(states[live])
            flat = kids.reshape(-1, 54)
            v = value_of(net, flat, device, amp=amp).reshape(len(live), C.NUM_MOVES)
            best = v.argmin(axis=1)
            states[live] = kids[np.arange(len(live)), best]
            done[live] = C.is_solved(states[live])
        out[d] = float(done.mean())
    return out


def _walk_exact(n, depth, rng):
    """n states each from a walk of exactly `depth` moves."""
    states = np.repeat(C.SOLVED_ARR[None, :], n, axis=0)
    prev = np.full(n, -1)
    for _ in range(depth):
        mv = rng.integers(0, C.NUM_MOVES, size=n)
        clash = (mv // 3) == prev
        while clash.any():
            mv[clash] = rng.integers(0, C.NUM_MOVES, size=int(clash.sum()))
            clash = (mv // 3) == prev
        states = states[np.arange(n)[:, None], C.PERMS[mv]]
        prev = mv // 3
    return states


def main():
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--hidden", type=int, default=512)
    p.add_argument("--layers", type=int, default=8)
    p.add_argument("--heads", type=int, default=8)
    p.add_argument("--batch-size", type=int, default=1024)
    p.add_argument("--lr", type=float, default=3e-4)
    p.add_argument("--max-steps", type=int, default=100000)
    p.add_argument("--scramble-depth", type=int, default=30,
                   help="States are sampled from walks of 1..this many moves. Past "
                        "~20 every state is fully mixed (God's Number), so a larger "
                        "value buys nothing.")
    p.add_argument("--update-threshold", type=float, default=0.15, metavar="MOVES",
                   help="Refresh the target network once the net fits its current "
                        "target to within this mean absolute error, IN MOVES. "
                        "Deliberately not an MSE threshold: target values grow from "
                        "~1 early to ~20 at full distance and MSE scales with their "
                        "square, so the achievable loss floor rises through the run "
                        "and no fixed MSE value works at both ends (measured: 0.0006 "
                        "at step 400, 0.02-0.03 by step 1500). MAE in moves is "
                        "interpretable and roughly scale-stable.")
    p.add_argument("--min-update-steps", type=int, default=100,
                   help="Floor on the refresh interval, so one lucky batch cannot "
                        "ratchet the target forward.")
    p.add_argument("--max-update-steps", type=int, default=400,
                   help="Backstop: refresh after this many steps even if the "
                        "threshold is unmet. Each refresh propagates accuracy about "
                        "one move further from solved, so a threshold that stops "
                        "being reachable stalls the frontier outright while training "
                        "still looks healthy -- observed at 1300 steps with no "
                        "refresh. This bounds the run at ~max*30 steps.")
    p.add_argument("--eval-every", type=int, default=2000)
    p.add_argument("--eval-n", type=int, default=200)
    p.add_argument("--eval-depths", default="3,6,10,15,20")
    p.add_argument("--out", default="checkpoints/value_iter")
    p.add_argument("--hub-repo", default="")
    p.add_argument("--upload-every", type=int, default=1,
                   help="Push to the Hub every N evals that improve on the best "
                        "mean greedy solve rate. The point is that a run can be "
                        "stopped the moment its frontier plateaus without losing "
                        "the model -- an end-of-run-only upload makes stopping "
                        "early cost the whole run.")
    p.add_argument("--seed", type=int, default=0)
    args = p.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    amp = device == "cuda"
    torch.manual_seed(args.seed)
    rng = np.random.default_rng(args.seed)

    net = ValueNet(args.hidden, args.layers, args.heads).to(device)
    target_net = ValueNet(args.hidden, args.layers, args.heads).to(device)
    target_net.load_state_dict(net.state_dict())
    for q in target_net.parameters():
        q.requires_grad_(False)

    params = sum(q.numel() for q in net.parameters())
    print(f"value-iteration net: {params/1e6:.1f}M params | {device} | "
          f"batch {args.batch_size} | walk depth 1..{args.scramble_depth}", flush=True)

    opt = torch.optim.AdamW(net.parameters(), lr=args.lr)
    eval_depths = [int(d) for d in args.eval_depths.split(",")]
    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)

    api = None
    if args.hub_repo:
        import os
        from huggingface_hub import HfApi
        api = HfApi(token=os.environ.get("HF_TOKEN"))

    def publish(tag):
        if api is None:
            return
        try:
            api.upload_folder(repo_id=args.hub_repo, folder_path=str(out_dir),
                              path_in_repo="checkpoints/value_iter")
            print(f"  uploaded ({tag})", flush=True)
        except Exception as e:
            # A transient Hub failure must not kill a run that is otherwise fine;
            # the next improvement retries, and the local copy is still on disk.
            print(f"  upload failed ({tag}): {type(e).__name__}: {e}", flush=True)

    best = -1.0
    improved = 0
    updates = 0
    last_update = 0
    run = []
    t0 = time.time()
    for step in range(1, args.max_steps + 1):
        states, _ = C.random_walk(args.batch_size, args.scramble_depth, rng)
        tgt = davi_targets(target_net, states, device, amp=amp)

        ids = encode_batch(states, device)
        y = torch.from_numpy(tgt).to(device)
        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=amp):
            pred = net(ids)
            loss = F.mse_loss(pred.float(), y)
        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(net.parameters(), 1.0)
        opt.step()

        l = loss.item()
        with torch.no_grad():
            mae = (pred.float() - y).abs().mean().item()
        since = step - last_update
        ready = mae < args.update_threshold and since >= args.min_update_steps
        forced = since >= args.max_update_steps
        if ready or forced:
            target_net.load_state_dict(net.state_dict())
            updates += 1
            last_update = step
            print(f"  step {step}: target net updated (#{updates}, "
                  f"mae {mae:.3f} moves{', forced' if forced and not ready else ''})",
                  flush=True)

        if step % 100 == 0:
            print(f"step {step}/{args.max_steps} loss {l:.4f} mae {mae:.3f} "
                  f"tgt_updates {updates} {step/(time.time()-t0):.1f} it/s", flush=True)

        if step % args.eval_every == 0 or step == args.max_steps:
            rates = greedy_solve_rate(net, device, eval_depths, args.eval_n,
                                      np.random.default_rng(12345), amp=amp)
            parts = "  ".join(f"d{d}:{r:.0%}" for d, r in rates.items())
            mean = float(np.mean(list(rates.values())))
            print(f"  step {step}  greedy solve  {parts}  | mean {mean:.1%}", flush=True)
            run.append({"step": step, "loss": l, "target_updates": updates,
                        "greedy": rates, "mean": mean})
            torch.save({"model": net.state_dict(), "hidden": args.hidden,
                        "layers": args.layers, "heads": args.heads, "step": step},
                       out_dir / "value_iter.pt")
            (out_dir / "history.json").write_text(json.dumps(run, indent=2))
            if mean > best:
                best = mean
                improved += 1
                torch.save({"model": net.state_dict(), "hidden": args.hidden,
                            "layers": args.layers, "heads": args.heads, "step": step},
                           out_dir / "value_iter_best.pt")
                if improved % args.upload_every == 0:
                    publish(f"step {step}, mean {mean:.1%}")

    print(f"done. best mean greedy solve rate {best:.1%}. weights in {out_dir}",
          flush=True)

    publish("final")


if __name__ == "__main__":
    main()