import cv2, numpy as np, glob, sys, os

W = "/Users/danicosta/Desktop/AI-Instagram/_inpaint"
v = sys.argv[1]
RW, RH = 180, 200

# matte canonico limpio de plato_1 (region 200x180, centro 90,100)
a0 = np.load(f"{W}/marca/alpha.npy").astype(np.float32)
aW0 = np.load(f"{W}/marca/aW.npy").astype(np.float32)

cx, cy, X0, Y0, CW, CH = map(int, open(f"{W}/{v}/geom.txt").read().split())
RX = max(0, min(cx - 90, 720 - RW)); RY = max(0, min(cy - 100, 1280 - RH))
ox, oy = RX - X0, RY - Y0

Jf = sorted(glob.glob(f"{W}/{v}/frames_full/*.png"))
Bf = sorted(glob.glob(f"{W}/{v}/ppresults/crop/frames/*.png"))
N = len(Jf)

# medias (la recuperada media = (Jmean - s*aW)/(1-s*a) porque a,aW no dependen de t)
accJ = np.zeros((RH, RW, 3), np.float64)
accB = np.zeros((RH, RW, 3), np.float64)
for i in range(N):
    accJ += cv2.imread(Jf[i])[RY:RY+RH, RX:RX+RW].astype(np.float64)
    accB += cv2.imread(Bf[i])[oy:oy+RH, ox:ox+RW].astype(np.float64)
Jmean = (accJ / N).astype(np.float32)
Bmean = (accB / N).astype(np.float32)          # referencia sin marca (ProPainter)

# mascara de la estrella (dilatada para incluir el anillo de borde)
starbin = (cv2.GaussianBlur(a0, (0,0), 1.0) > 0.05).astype(np.uint8)
starbin = cv2.dilate(starbin, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(9,9)))
starm = starbin[..., None].astype(np.float32)

def shift(img, dx, dy):
    M = np.float32([[1, 0, dx], [0, 1, dy]])
    return cv2.warpAffine(img, M, (RW, RH), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT)

# Metrica: energia de ALTA FRECUENCIA de la media recuperada en la zona de la estrella.
# El fondo (movido) queda borroso en la media; solo un residuo ESTATICO de marca es nitido.
best = (1e18, 0.0, 0.0, 1.0)
for dy in np.arange(-8, 8.01, 0.5):
    for dx in np.arange(-8, 8.01, 0.5):
        a_sh = shift(a0, dx, dy); aW_sh = shift(aW0, dx, dy)
        for s in np.arange(0.55, 1.45, 0.05):
            a = np.clip(s * a_sh, 0, 0.95); aw = s * aW_sh
            recm = (Jmean - aw) / (1 - a)[..., None]
            hp = recm - cv2.GaussianBlur(recm, (0, 0), 6.0)
            err = float(((hp ** 2) * starm).sum())
            if err < best[0]:
                best = (err, float(dx), float(dy), float(s))
_, dx, dy, s = best
print(f"{v}: mejor dx={dx} dy={dy} s={s:.2f} err={best[0]:.0f}")

# matte final alineado/escalado
a_sh = np.clip(s * shift(a0, dx, dy), 0, 0.95).astype(np.float32)
aW_sh = (s * shift(aW0, dx, dy)).astype(np.float32)
# gate eliptico generoso + pluma
yy, xx = np.mgrid[0:RH, 0:RW]
ell = (((xx - 90 - dx) / 54.0) ** 2 + ((yy - 100 - dy) / 64.0) ** 2) <= 1.0
a_sh = np.where(ell, a_sh, 0.0).astype(np.float32)
aW_sh = np.where(ell[..., None], aW_sh, 0.0).astype(np.float32)
a_sh = cv2.GaussianBlur(a_sh, (0,0), 0.4); aW_sh = cv2.GaussianBlur(aW_sh, (0,0), 0.4)
oneA = (1 - a_sh)[..., None]

out_dir = f"{W}/{v}/frames_final"; os.makedirs(out_dir, exist_ok=True)
fe = np.zeros((RH, RW), np.float32); cv2.rectangle(fe, (6,6), (RW-6,RH-6), 1.0, -1)
fe = cv2.GaussianBlur(fe, (0,0), 4.0)[..., None]
accr = np.zeros((RH, RW, 3), np.float64); nr = 0
for i in range(N):
    full = cv2.imread(Jf[i])
    reg = full[RY:RY+RH, RX:RX+RW].astype(np.float32)
    rec = np.clip((reg - aW_sh) / oneA, 0, 255)
    full[RY:RY+RH, RX:RX+RW] = np.clip(reg*(1-fe) + rec*fe, 0, 255).astype(np.uint8)
    cv2.imwrite(f"{out_dir}/{i+1:05d}.png", full)
    if i < 120: accr += rec.astype(np.float64); nr += 1

def eq(img):
    g = cv2.cvtColor(np.clip(img,0,255).astype(np.uint8), cv2.COLOR_BGR2GRAY).astype(np.float32)
    return ((g-g.min())/(g.max()-g.min()+1e-6)*255).astype(np.uint8)
cv2.imwrite(f"{W}/{v}/res_rec_eq.png", eq(accr/nr))
cv2.imwrite(f"{W}/{v}/res_orig_eq.png", eq(Jmean))
print(f"{v}: canonical done -> {out_dir}")