#!/usr/bin/env python3
"""Measure what the planner's field actually encodes.

Fits a ridge probe from the planner's per-cell features to the true BFS
distance and reports R^2. This is the diagnostic that caught the central
training bug: with only the sparse decision questions supervising it, the
recurrence scored R^2 = 0.21 -- it was running, it was differentiable, its
accuracy numbers looked reasonable, and it was not computing distances.

Also reports the field head's own predictions, which is what the decision
heads actually read.
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

import torch

import envs.maze as mz
from jevon.grid import bfs_field, encode_maze, pad_boards
from jevon.inference import resolve_weights
from jevon.modeling import JevonConfig, JevonModel
from jevon.planner import default_iterations


def ridge_r2(features: torch.Tensor, target: torch.Tensor) -> tuple[float, float]:
    design = torch.cat([features, torch.ones(len(features), 1)], 1)
    gram = design.T @ design + 1e-3 * torch.eye(design.shape[1])
    weights = torch.linalg.lstsq(gram, design.T @ target.unsqueeze(1)).solution
    pred = (design @ weights).squeeze(1)
    ss_res = ((pred - target) ** 2).sum()
    ss_tot = ((target - target.mean()) ** 2).sum()
    return float(1 - ss_res / ss_tot), float((pred - target).abs().mean())


def descent_accuracy(head: torch.Tensor, truth: torch.Tensor) -> float:
    """Fraction of cells whose lowest-predicted neighbour is a correct step.

    R^2 scores the field's absolute values, but greedy descent only ever
    compares a cell's four neighbours, so a field can be numerically mediocre
    and still route perfectly -- or fit well on average and still invert the
    one comparison that matters.  This is the quantity that predicts solve
    rate, so it is measured rather than inferred from R^2.
    """
    head = head.detach()
    hits = total = 0
    h, w = truth.shape
    for r in range(h):
        for c in range(w):
            here = float(truth[r, c])
            if here <= 0:                      # blocked, unreachable, or the goal
                continue
            options = [(rr, cc) for rr, cc in
                       ((r - 1, c), (r + 1, c), (r, c - 1), (r, c + 1))
                       if 0 <= rr < h and 0 <= cc < w and truth[rr, cc] >= 0]
            if not options:
                continue
            pick = min(options, key=lambda x: float(head[x[0], x[1]]))
            # Correct iff that neighbour lies on *a* shortest path.
            hits += float(truth[pick]) == here - 1
            total += 1
    return hits / max(total, 1)


def descent_outcomes(head: torch.Tensor, truth: torch.Tensor,
                     cap_mult: int = 8) -> float:
    """Fraction of cells from which greedy descent actually reaches the goal.

    Per-step accuracy is not the quantity a controller lives or dies by.  A
    field that is right 99% of the time but contains one spurious basin traps
    forever; a field that is right 70% of the time with no basin still arrives,
    just by a longer route.  Solve rate equals *this* number, so report it next
    to the per-step figure and let the gap between them speak.
    """
    head = head.detach()
    h, w = truth.shape
    goal = min(((r, c) for r in range(h) for c in range(w)
                if float(truth[r, c]) == 0.0), default=None)
    if goal is None:
        return 0.0
    reached = total = 0
    cap = cap_mult * h * w
    for r in range(h):
        for c in range(w):
            if float(truth[r, c]) < 0:
                continue
            total += 1
            cur, seen = (r, c), set()
            for _ in range(cap):
                if cur == goal:
                    reached += 1
                    break
                if cur in seen:
                    break                       # a cycle: descent never escapes
                seen.add(cur)
                options = [(rr, cc) for rr, cc in
                           ((cur[0] - 1, cur[1]), (cur[0] + 1, cur[1]),
                            (cur[0], cur[1] - 1), (cur[0], cur[1] + 1))
                           if 0 <= rr < h and 0 <= cc < w and truth[rr, cc] >= 0]
                if not options:
                    break
                cur = min(options, key=lambda x: float(head[x[0], x[1]]))
    return reached / max(total, 1)


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--checkpoint", default="runs/jevon-final")
    # None means "prefer balanced.pt", the checkpoint every shipped number
    # comes from. See jevon.inference.resolve_weights.
    p.add_argument("--weights", default=None)
    p.add_argument("--sizes", default="11,21,31,51")
    p.add_argument("--episodes", type=int, default=4)
    p.add_argument("--device", default="cpu")
    p.add_argument("--out", default=None)
    args = p.parse_args()

    root = Path(args.checkpoint)
    cfg = JevonConfig(**json.loads((root / "config.json").read_text()))
    model = JevonModel(cfg)
    model.load_state_dict(torch.load(resolve_weights(root, args.weights),
                                     map_location="cpu"))
    model.eval().to(args.device)

    report = {}
    print(f"{'size':>5s} {'cells':>7s} {'probe R2':>9s} {'probe MAE':>10s} "
          f"{'head R2':>8s} {'descent':>8s} {'reach':>7s} "
          f"{'diameter':>9s} {'T':>6s}")
    for size in [int(s) for s in args.sizes.split(",")]:
        feats, heads, dists, descents, reaches = [], [], [], [], []
        diameter = 0
        steps = default_iterations(size)
        with torch.no_grad():
            for i in range(args.episodes):
                topology = mz.TOPOLOGIES[i % len(mz.TOPOLOGIES)]
                # Seeds ending in 9 are the held-out test split.
                state = mz.generate_maze(size, topology, 70_009 + i * 10)
                board, _, target = encode_maze(state)
                boards, passable, _ = pad_boards([board])
                cells = model.encode_boards(boards.to(args.device),
                                            passable.to(args.device), steps, 0)
                head = model.predict_field(
                    cells, boards.to(args.device),
                    passable.to(args.device), iterations=steps)[0, 0]
                truth = bfs_field(board, target)
                mask = truth >= 0
                diameter = max(diameter, int(truth.max()))
                descents.append(descent_accuracy(head.cpu(), truth))
                reaches.append(descent_outcomes(head.cpu(), truth))
                feats.append(cells[0, :, mask].T.cpu())
                heads.append(head[mask].cpu())
                dists.append(truth[mask])
        features = torch.cat(feats)
        truth = torch.cat(dists)
        r2, mae = ridge_r2(features, truth)
        head_r2, _ = ridge_r2(torch.cat(heads).unsqueeze(1), truth)
        descent = sum(descents) / max(len(descents), 1)
        reach = sum(reaches) / max(len(reaches), 1)
        report[size] = {"cells": len(truth), "probe_r2": r2, "probe_mae": mae,
                        "head_r2": head_r2, "descent_acc": descent,
                        "reach_rate": reach,
                        "diameter": diameter, "iterations": steps}
        print(f"{size:5d} {len(truth):7d} {r2:9.4f} {mae:10.2f} {head_r2:8.4f} "
              f"{descent:8.4f} {reach:7.3f} {diameter:9d} {steps:6d}")

    if args.out:
        Path(args.out).parent.mkdir(parents=True, exist_ok=True)
        Path(args.out).write_text(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()