"""E3 -- token-latent geometry probe (exp7 explainability battery, sec 5).

Only meaningful in the 192-identity cells (A, C): with ``width ==
latent_dim``, plan tokens ``y_j`` live in literally the same coordinate
system as world-model latents, so they can be compared to ``x_hat_j`` (the
state block ``j`` actually causes) with a plain distance/cosine, and decoded
*directly* through the pixel decoder -- no analogous move exists in any
256-d variant, which is exactly what "identity embedding buys" (sec 1).
"""

import argparse
import json
import sys
from pathlib import Path

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

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

from lejepa_control.data import LatentGoalDataset, split_episodes  # noqa: E402
from lejepa_control.decoder import load_decoder  # noqa: E402
from lejepa_control.rollout import rollout_plan  # noqa: E402
from lejepa_control.solver import load_controller  # noqa: E402
from lejepa_control.world_model import load_lewm  # noqa: E402

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / 'tools'))
from decode_rollout import panel, to_uint8  # noqa: E402


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument('--controller', default='data/runs/exp7/fused192/controller.pt')
    p.add_argument('--decoder', default='data/runs/decoder/decoder.pt')
    p.add_argument('--latents', default='data/latents')
    p.add_argument('--samples', type=int, default=256)
    p.add_argument('--horizon', type=int, default=5)
    p.add_argument('--decode-episodes', type=int, default=3)
    p.add_argument('--tag', default='main')
    p.add_argument('--out', default='data/runs/probe_geometry')
    return p.parse_args()


@torch.no_grad()
def geometry_trace(controller, model, ctx, past, goal):
    """Per-refinement ``y_k`` (raw plan tokens), ``x_hat_k`` (the latent each
    block causes), ``||y - x_hat||`` and cosine. ``K+1`` entries, same
    recursion as ``IterativeController.forward``.
    """
    cond = controller.condition(ctx, goal)
    tokens = controller.initial_plan(cond)

    ys, xhats, dists, coss = [], [], [], []
    for k in range(controller.refinements + 1):
        actions = controller.to_actions(tokens)
        pred = rollout_plan(model, ctx, past, actions)  # x_hat_j, (B, H, D)

        ys.append(tokens.clone())
        xhats.append(pred.clone())
        dists.append((tokens - pred).norm(dim=-1))            # (B, H)
        coss.append(F.cosine_similarity(tokens, pred, dim=-1))  # (B, H)

        if k == controller.refinements:
            break
        delta = controller.refine(tokens, cond, pred, goal)
        idx = min(k, controller.step_logit.numel() - 1)
        tokens = tokens + torch.sigmoid(controller.step_logit[idx]) * delta

    return torch.stack(ys), torch.stack(xhats), torch.stack(dists), torch.stack(coss)


@torch.no_grad()
def decode_panel(decoder, ys, sample_idx, out_path):
    """One panel per sample: rows = refinement k, cols = horizon slot j,
    each cell = ``decoder(y_k[sample, j])`` -- what that plan slot "intends".
    """
    rows = [to_uint8(decoder(ys[k, sample_idx])) for k in range(ys.size(0))]
    labels = [f'k={k}' for k in range(ys.size(0))]
    titles = [f'slot {j}' for j in range(ys.size(2))]
    return panel(rows, labels, titles, out_path)


def main():
    args = parse_args()
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    torch.manual_seed(0)

    stats = json.loads((Path(args.latents) / 'stats.json').read_text())
    model = load_lewm(device=device)
    controller, ckpt = load_controller(
        args.controller, latent_dim=stats['latent_dim'], device=device
    )
    controller.eval()
    if not controller.no_latent_proj:
        print(f'{args.controller}: no_latent_proj=False (width != latent_dim) '
              '-- E3 is only defined for the 192-identity cells (A, C). Skipping.')
        return
    print(f'controller step {ckpt["step"]}, K={controller.refinements}, '
          f'fused={controller.fused}')

    _, val_eps = split_episodes(stats['n_episodes'])
    val = LatentGoalDataset(
        args.latents, max_offset=5, episodes=val_eps, horizon=args.horizon
    )
    idx = np.random.default_rng(0).choice(len(val), args.samples, replace=False)
    batch = {
        k: torch.stack([val[int(i)][k] for i in idx]).to(device)
        for k in ('context', 'past_actions', 'goal')
    }

    ys, xhats, dists, coss = geometry_trace(
        controller, model, batch['context'], batch['past_actions'], batch['goal']
    )

    per_k_dist = dists.mean(dim=(1, 2))
    per_k_cos = coss.mean(dim=(1, 2))
    print('\n=== E3: plan token vs. its own consequence, ||y_j - x_hat_j|| ===')
    print(f'{"k":>3}{"mean ||.||":>12}{"mean cos":>11}')
    for k in range(dists.size(0)):
        print(f'{k:>3}{per_k_dist[k].item():>12.4f}{per_k_cos[k].item():>11.4f}')

    report = {
        'tag': args.tag,
        'checkpoint': args.controller,
        'step': int(ckpt['step']),
        'fused': bool(controller.fused),
        'mean_dist_per_k': [round(float(v), 5) for v in per_k_dist],
        'mean_cos_per_k': [round(float(v), 5) for v in per_k_cos],
    }
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    with (out / 'probe_geometry.jsonl').open('a') as f:
        f.write(json.dumps(report) + '\n')
    print(f'wrote {out / "probe_geometry.jsonl"}')

    if args.decode_episodes > 0 and Path(args.decoder).exists():
        decoder, dec_ckpt = load_decoder(args.decoder, device=device)
        latent_kind = dec_ckpt['meta']['latent_source']
        if latent_kind != 'emb':
            print(f"decoder trained on '{latent_kind}' latents, not 'emb' "
                  '-- skipping the direct plan-token decode.')
        else:
            for i in range(min(args.decode_episodes, ys.size(1))):
                path = decode_panel(
                    decoder, ys, i, out / f'{args.tag}_sample{i}.png'
                )
                print(f'  -> {path}')
    elif args.decode_episodes > 0:
        print(f'no decoder at {args.decoder} -- skipping the direct plan-token decode.')


if __name__ == '__main__':
    main()