# vs_psnr_variants.py
import os
import subprocess
import sys
import traceback

import vapoursynth as vs
core = vs.core

# ---------- CONFIG ----------
SRC = "/run/media/e/Almacenamiento/Media/Anime/Raws/[BDMV] Ojamajo Doremi (1999)/BIXA_9550_2/BDMV/STREAM/00001.m2ts"   # <- pon aquí tu archivo
FRAME = 267            # <- pon aquí el frame entero
W = 1920               # opcional: fuerza ancho si lo conoces (ej. 1920)
H = 1080               # opcional: fuerza alto si lo conoces (ej. 1080)
# ----------------------------

# Intentar abrir con los plugins más comunes
clip = None
for opener in ("ffms2.Source", "lsmas.LWLibavSource", "bs.VideoSource"):
    try:
        if opener == "ffms2.Source":
            clip = core.ffms2.Source(SRC)
        elif opener == "lsmas.LWLibavSource":
            clip = core.lsmas.LWLibavSource(SRC)
        else:
            clip = core.bs.VideoSource(SRC)
        print("Opened with", opener)
        break
    except Exception:
        clip = None

if clip is None:
    raise RuntimeError("No se pudo abrir el archivo con ffms2/lsmas/bs.VideoSource")

# Forzar metadata base (opcional pero recomendable)
clip = core.std.SetFrameProps(clip, _Matrix=1, _Transfer=1, _Primaries=1, _ColorRange=1)
single = core.std.Trim(clip, FRAME, FRAME)

# Listas de valores a probar (puedes ajustar)
matrices = ["709", "601", "smpte170m", "bt2020"]
transfers = ["709", "srgb", "smpte170m", "bt2020"]
ranges = ["limited", "full"]

# Resamplers a probar (nombre, función)
resamplers = []
if hasattr(core, "resize") and hasattr(core.resize, "Bilinear"):
    resamplers.append(("Bilinear", core.resize.Bilinear))
if hasattr(core, "resize") and hasattr(core.resize, "Bicubic"):
    resamplers.append(("Bicubic", core.resize.Bicubic))
if hasattr(core, "resize") and hasattr(core.resize, "Point"):
    resamplers.append(("Point", core.resize.Point))
if not resamplers:
    raise RuntimeError("No se detectaron resamplers en core.resize")

out_dir = "vs_param_bruteforce"
os.makedirs(out_dir, exist_ok=True)
log_path = os.path.join(out_dir, "bruteforce.log")

successes = []
failures = []

def safe_write(rgb_clip, out_template):
    """
    Escribe usando imwri (requiere %d). imwri escribe en el directorio actual.
    Luego movemos el archivo escrito al out_dir con nombre final.
    """
    try:
        core.imwri.Write(rgb_clip, imgformat="PNG", filename=out_template).get_frame(0)
    except Exception as e:
        return False, f"imwri write failed: {e}"
    written = out_template % 0  # nombre que imwri crea (ej: X_0.png)
    if not os.path.exists(written):
        return False, f"expected file not found: {written}"
    final_name = out_template.replace("%d", str(FRAME))
    dst = os.path.join(out_dir, final_name)
    try:
        os.replace(written, dst)
    except Exception as e:
        return False, f"move failed: {e}"
    return True, dst

# Probar combinaciones
for res_name, res_func in resamplers:
    for m in matrices:
        for t in transfers:
            for r in ranges:
                combo = f"{res_name}_m{m}_t{t}_r{r}"
                out_template = f"{combo}_%d.png"
                try:
                    rgb = res_func(
                        single,
                        format=vs.RGB24,
                        matrix_in_s=m,
                        transfer_in_s=t,
                        primaries_in_s="709",
                        range_in_s=r,
                        range_s="full",
                        dither_type="none"
                    )
                    ok, info = safe_write(rgb, out_template)
                    if ok:
                        msg = f"OK: {combo} -> {info}"
                        print(msg)
                        successes.append(msg)
                    else:
                        msg = f"WRITE_FAIL: {combo} -> {info}"
                        print(msg)
                        failures.append(msg)
                except Exception:
                    tb = traceback.format_exc().splitlines()[-1]
                    msg = f"ERR: {combo} -> {tb}"
                    print(msg)
                    failures.append(msg)

# Guardar log
with open(log_path, "w") as fh:
    fh.write("=== SUCCESSES ===\n")
    for s in successes: fh.write(s + "\n")
    fh.write("\n=== FAILURES ===\n")
    for f in failures: fh.write(f + "\n")

print("\nDone. Results in folder:", out_dir)
print("Log:", log_path)