"""Verify the frozen LeWM loads, encodes, and gives gradients w.r.t. actions."""

import sys
from pathlib import Path

import torch

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

from lejepa_control.world_model import load_lewm  # noqa: E402


def main():
    import stable_worldmodel as swm

    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    model = load_lewm(device=device)
    print('loaded LeWM ok')
    print('  predictor.num_frames =', model.predictor.num_frames)
    print('  action_encoder.input_dim =', model.action_encoder.input_dim)

    ds = swm.data.load_dataset('pusht_smoke')
    print('  dataset columns =', ds.column_names)
    print('  episodes =', len(ds.lengths))

    ep = ds.load_episode(0)
    for k, v in ep.items():
        if hasattr(v, 'shape'):
            print(f'  ep[{k}] {tuple(v.shape)} {v.dtype}')

    pixels_key = 'pixels' if 'pixels' in ep else 'obs.pixels'
    frames = ep[pixels_key]
    if not torch.is_tensor(frames):
        frames = torch.as_tensor(frames)
    if frames.shape[-1] in (1, 3):  # NHWC -> NCHW
        frames = frames.permute(0, 3, 1, 2)
    frames = frames.float() / 255.0 if frames.dtype == torch.uint8 else frames

    T = model.predictor.num_frames
    pixels = frames[:T].unsqueeze(0).to(device)  # (1, T, C, H, W)
    with torch.no_grad():
        info = model.encode({'pixels': pixels})
    emb = info['emb']
    print('  emb', tuple(emb.shape))

    # gradient check: does d(pred)/d(action) flow?
    action_dim = model.action_encoder.input_dim
    action = torch.zeros(1, T, action_dim, device=device, requires_grad=True)
    act_emb = model.action_encoder(action)
    pred = model.predict(emb, act_emb)
    print('  pred', tuple(pred.shape))
    pred.sum().backward()
    g = action.grad
    print('  action.grad norm =', float(g.norm()))
    assert g.norm() > 0, 'no gradient reached the action'
    print('OK: gradients flow through the frozen predictor to actions')


if __name__ == '__main__':
    main()