#!/usr/bin/env python3
"""Record landing videos of the 14-parameter LunarLander policy (headless).

Writes `lunarlander-14params.mp4` (a few episodes) and `lunarlander-14params.gif`
(one episode, downscaled) next to this file.
"""
import os

os.environ.setdefault("SDL_VIDEODRIVER", "dummy")
os.environ.setdefault("SDL_AUDIODRIVER", "dummy")

import numpy as np
import gymnasium as gym
import imageio.v2 as imageio

from model import TinyLunarPolicy

FPS = 50
N_MP4_EPISODES = 3
HERE = os.path.dirname(os.path.abspath(__file__))
SEEDS = [40122, 42940, 51305, 55316, 60944, 61197, 93978, 96327, 103680, 157789]


def record(policy, n_episodes, min_return=250.0):
    env = gym.make("LunarLander-v3", render_mode="rgb_array")
    clips = []
    for seed in SEEDS:
        for ep in range(100):
            o, _ = env.reset(seed=seed * 1000 + ep)
            frames, tot, last, term, trunc = [], 0.0, 0.0, False, False
            frames.append(env.render())
            while not (term or trunc):
                o, r, term, trunc, _ = env.step(policy.act(o))
                tot += r
                last = r
                frames.append(env.render())
            if term and not trunc and last == 100.0 and tot >= min_return:
                clips.append((seed, ep, tot, frames))
                print("  клип: seed %d ep %d  возврат %.1f  кадров %d" % (seed, ep, tot, len(frames)),
                      flush=True)
                break                      # один клип на сид -> разные старты в видео
            if len(clips) >= n_episodes:
                env.close()
                return clips
    env.close()
    return clips


def main():
    policy = TinyLunarPolicy.from_safetensors(os.path.join(HERE, "model.safetensors"))
    print("запись эпизодов...", flush=True)
    clips = record(policy, N_MP4_EPISODES)
    if not clips:
        raise SystemExit("не нашёл подходящих эпизодов")

    # mp4: склеиваем клипы с короткой белой паузой между ними
    white = np.full_like(clips[0][3][0], 255)
    frames = []
    for i, (s, e, r, fr) in enumerate(clips):
        frames += fr
        frames += [white] * 15
    mp4 = os.path.join(HERE, "lunarlander-14params.mp4")
    with imageio.get_writer(mp4, fps=FPS, codec="libx264", quality=8,
                            pixelformat="yuv420p", macro_block_size=8) as w:
        for f in frames:
            w.append_data(np.asarray(f, np.uint8))
    print("saved:", mp4, "| кадров", len(frames), "| ~%.1f c" % (len(frames) / FPS))

    # gif: один клип, уменьшенный вдвое и 25 fps
    fr = clips[0][3]
    small = [np.asarray(f, np.uint8)[::2, ::2] for f in fr[::2]]
    gif = os.path.join(HERE, "lunarlander-14params.gif")
    imageio.mimsave(gif, small, fps=25)
    print("saved:", gif, "| кадров", len(small), "| размер", small[0].shape)


if __name__ == "__main__":
    main()