"""Recover the per-q validation distance profile for checkpoints that predate
the profile logging in ``train_controller.py``.

The original controller was trained before ``evaluate()`` recorded
``val_profile``, so its arrival-timing signature — the thing Figure 2 shows —
is missing from its checkpoint while every later run has it. This recomputes
it with the same function, the same held-out split, and the same batch order,
and writes it next to the diagnostics so the figure script can read it.

    python report/recover_profiles.py --controllers data/runs/controller/controller.pt
"""

import argparse
import json
import sys
from pathlib import Path

import torch
from torch.utils.data import DataLoader

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

from lejepa_control.data import LatentGoalDataset, split_episodes  # 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] / 'scripts'))
from train_controller import evaluate  # noqa: E402


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--controllers', nargs='+',
                   default=['data/runs/controller/controller.pt'])
    p.add_argument('--latents', default='data/latents')
    p.add_argument('--batch-size', type=int, default=128)
    p.add_argument('--out', default='data/runs/diagnostics/profiles_posthoc.json')
    args = p.parse_args()

    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    stats = json.loads((Path(args.latents) / 'stats.json').read_text())
    model = load_lewm(device=device)

    _, val_eps = split_episodes(stats['n_episodes'])
    val_set = LatentGoalDataset(
        args.latents, max_offset=5, episodes=val_eps, horizon=5
    )

    out = {}
    for path in args.controllers:
        name = Path(path).parent.name
        controller, ckpt = load_controller(
            path, latent_dim=stats['latent_dim'], device=device
        )
        # same batches for every checkpoint, so profiles are comparable
        loader = DataLoader(
            val_set, batch_size=args.batch_size, shuffle=True,
            generator=torch.Generator().manual_seed(0)
        )
        val = evaluate(controller, model, loader, device)
        out[name] = {
            'step': int(ckpt['step']),
            'terminal': val['terminal'],
            'first': val['first'],
            'arrival': val['arrival'],
            'profile': val['profile'],
        }
        print(f'{name}: step {ckpt["step"]}  terminal {val["terminal"]:.4f}  '
              f'arrival {val["arrival"]:.4f}')
        for q, prof in val['profile'].items():
            print(f'  q={q}: {prof}')

    Path(args.out).parent.mkdir(parents=True, exist_ok=True)
    Path(args.out).write_text(json.dumps(out, indent=2))
    print(f'\nwrote {args.out}')


if __name__ == '__main__':
    main()