"""
evo_plot_mp4.py  –  Visualize VIGS-SLAM intermediate trajectories as an MP4/GIF.

Usage
-----
python evo_plot_mp4.py \
    --output_folder  outputs/my_sequence \
    --gt_file        data/my_sequence/groundtruth.txt \
    [--align_scale] \
    [--plot_est_only] \
    [--stride 20] \
    [--fps 30] \
    [--gif_scale 800]

The script expects the following layout under --output_folder:
    <output_folder>/traj/          # TUM-format .txt files (one per keyframe snapshot)
    <output_folder>/traj_kf_beforeBA.txt  # final full trajectory (for summary plot)
"""

import argparse
import os
import subprocess


# ---------------------------------------------------------------------------
# helpers
# ---------------------------------------------------------------------------

def _run(cmd):
    subprocess.run(cmd, check=False)


def generate_plots(traj_folder, gt_file, output_dir,
                   align_scale=False, plot_est_only=False, stride=20):
    os.makedirs(output_dir, exist_ok=True)

    # --- per-snapshot plots ---
    traj_files = sorted(f for f in os.listdir(traj_folder) if f.endswith(".txt"))
    for traj_file in traj_files[::stride]:
        print(traj_file)
        traj_path = os.path.join(traj_folder, traj_file)
        plot_path = os.path.join(output_dir, traj_file.replace('.txt', '.png'))

        if plot_est_only:
            _run(["evo_traj", "tum", traj_path,
                  "--plot", "--plot_mode=xy",
                  "--save_plot", plot_path])
        elif align_scale:
            _run(["evo_ape", "tum", gt_file, traj_path,
                  "-vas", "--plot_mode=xyz", "--pose_relation", "trans_part",
                  "--save_plot", plot_path])
        else:
            _run(["evo_ape", "tum", gt_file, traj_path,
                  "-a", "--plot_mode=xy", "--pose_relation", "trans_part",
                  "--save_plot", plot_path])


def create_video_ffmpeg(image_dir, output_video, fps=30):
    _run(["ffmpeg", "-y", "-framerate", str(fps),
          "-pattern_type", "glob", "-i", os.path.join(image_dir, "*.png"),
          "-vf", "pad=ceil(iw/2)*2:ceil(ih/2)*2",
          "-c:v", "libx264", "-pix_fmt", "yuv420p", output_video])
    print(f"Video saved as {output_video}")


def convert_mp4_to_gif(input_video, output_gif, fps=15, scale=800):
    _run(["ffmpeg", "-y", "-i", input_video,
          "-vf", f"fps={fps},scale={scale}:-1:flags=lanczos,"
                 "split[s0][s1];[s0]palettegen[p];[s1][p]paletteuse",
          output_gif])
    print(f"Converted {input_video} to {output_gif}")


# ---------------------------------------------------------------------------
# main
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description="Generate trajectory animation MP4/GIF from VIGS-SLAM outputs.")
    parser.add_argument("--output_folder", required=True,
                        help="Path to the sequence output directory.")
    parser.add_argument("--gt_file", required=True,
                        help="Ground-truth trajectory in TUM format.")
    parser.add_argument("--align_scale", action="store_true",
                        help="Align with scale (monocular).")
    parser.add_argument("--plot_est_only", action="store_true",
                        help="Plot estimated trajectory only (no APE comparison).")
    parser.add_argument("--stride", type=int, default=20,
                        help="Use every Nth snapshot file (default: 20).")
    parser.add_argument("--fps", type=int, default=30,
                        help="Frame rate for the output video (default: 30).")
    parser.add_argument("--gif_scale", type=int, default=800,
                        help="Width in pixels for the output GIF (default: 800).")
    args = parser.parse_args()

    traj_folder = os.path.join(args.output_folder, "traj")
    suffix = "align_scale" if args.align_scale else "xy"
    if args.plot_est_only:
        output_dir = os.path.join(args.output_folder, "animation_est_only")
        suffix = "est_only"
    else:
        output_dir = os.path.join(args.output_folder, f"animation_{suffix}")

    output_video = os.path.join(args.output_folder, f"traj_animation_{suffix}.mp4")
    output_gif   = os.path.join(args.output_folder, f"traj_animation_{suffix}.gif")

    # clean previous frames
    if os.path.isdir(output_dir):
        import shutil
        shutil.rmtree(output_dir)

    generate_plots(traj_folder, args.gt_file, output_dir,
                   align_scale=args.align_scale,
                   plot_est_only=args.plot_est_only,
                   stride=args.stride)
    create_video_ffmpeg(output_dir, output_video, fps=args.fps)
    convert_mp4_to_gif(output_video, output_gif, fps=args.fps, scale=args.gif_scale)


if __name__ == "__main__":
    main()
