"""Inverse-brighten a PARA eval frame to remove the in-render heatmap overlay.

`libero/eval.py:render_eval_frame` composites every PARA visualization frame as
    vis = clip(frame * 0.55 + heat_red * 0.45) * 255
which uniformly dims the underlying LIBERO render to 55%. ACT frames don't get
this treatment (eval.py:588 saves ACT as raw RGB), so PARA frames look noticeably
darker than ACT frames in the qualitative panels.

This script takes a baked-in-overlay PNG and divides RGB by 0.55 to recover the
original render. Where the heatmap was concentrated (red channel boost), there
will be faint red residue — but for diffuse heatmap distributions (the typical
case in our eval frames) that residue is invisible.

Usage:
    python3 clean_para_heatmap_overlay.py <input.png> <output.png>
"""
import argparse
from pathlib import Path
import numpy as np
from PIL import Image


def clean(in_path: Path, out_path: Path, dim_factor: float = 0.55):
    src = np.array(Image.open(in_path).convert("RGB")).astype(np.float32)
    out = np.clip(src / dim_factor, 0, 255).astype(np.uint8)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    Image.fromarray(out).save(out_path)
    print(f"wrote {out_path}")


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("input")
    ap.add_argument("output")
    ap.add_argument("--dim", type=float, default=0.55,
                    help="Dim factor used by render_eval_frame (default 0.55)")
    args = ap.parse_args()
    clean(Path(args.input), Path(args.output), args.dim)