#!/usr/bin/env python3
"""SSIM + PSNR: A/B vs C (sage-off reference) on matched frames."""
import subprocess, numpy as np

def frames(path):
    raw = subprocess.run(["ffmpeg","-v","error","-i",path,"-vf","format=gray,scale=608:336",
                          "-f","rawvideo","-"],capture_output=True).stdout
    n=len(raw)//(336*608)
    return np.frombuffer(raw[:n*336*608],np.uint8).reshape(n,336,608).astype(np.float32)

def ssim(a,b):
    C1,C2=(0.01*255)**2,(0.03*255)**2
    a=a-a.mean(); b=b-b.mean()
    sxy=(a*b).mean(); sx=a.std()**2; sy=b.std()**2
    return ((2*sxy+C1)*(sx+sy+C2))/((sx+sy)*(sx+sy)+C2) if False else (((2*sxy)* (C1)) ), 0

def psnr(a,b):
    mse=((a-b)**2).mean()
    return 10*np.log10((255**2)/mse) if mse>0 else 99

def ssim2(a,b):
    from skimage.metrics import structural_similarity as sk_ssim
    return sk_ssim(a,b,data_range=255)
try:
    import skimage  # noqa
    HAVE_SK=True
except ImportError:
    HAVE_SK=False

A=frames("bcHD36_A_00001_.mp4"); B=frames("bcHD36_B_00001_.mp4"); C=frames("bcHD36_C_00001_.mp4")
n=min(len(A),len(B),len(C))
# skip first seconds (prompt divergence); compare f30..end mean per-frame metrics
sl=slice(30,n)
for name,X in [("B_vs_C(BC vs sageoff)",B[sl]),("A_vs_C(sage vs sageoff)",A[sl])]:
    p=np.mean([psnr(X[i],C[sl][i]) for i in range(len(range(*sl.indices(n))))])
    print(f"{name}: PSNR={p:.2f} dB")
if HAVE_SK:
    for name,X in [("B_vs_C",B[sl]),("A_vs_C",A[sl])]:
        vals=[ssim2(X[i].astype(np.uint8),C[sl][i].astype(np.uint8)) for i in range(len(range(*sl.indices(n))))]
        print(f"{name}: SSIM={np.mean(vals):.4f}")
else:
    print("skimage not available, SSIM skipped")