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

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

a0 = np.load(f"{W}/marca/alpha.npy").astype(np.float32)     # matte canonico plato_1
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)

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); Pmean = (accB / N).astype(np.float32)

# huella medida de la marca en este video = media_original - media_ProPainter (sin marca)
F = cv2.cvtColor(np.clip(Jmean - Pmean + 128, 0, 255).astype(np.uint8), cv2.COLOR_BGR2GRAY).astype(np.float32) - 128.0
T = cv2.cvtColor(np.clip(aW0 + 128, 0, 255).astype(np.uint8), cv2.COLOR_BGR2GRAY).astype(np.float32) - 128.0
# correlacion de fase (sub-pixel) para alinear el template T a la huella F
win = cv2.createHanningWindow((RW, RH), cv2.CV_32F)
(sx, sy), resp = cv2.phaseCorrelate(T * win, F * win)
print(f"{v}: shift fase = ({sx:.2f},{sy:.2f}) resp={resp:.3f}")

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)

# ajuste de escala s por minimos cuadrados sobre la huella (constrenido cerca de 1)
def err_for(dx, dy, s):
    a = np.clip(s * shift(a0, dx, dy), 0, 0.95); aw = s * shift(aW0, dx, dy)
    recm = (Jmean - aw) / (1 - a)[..., None]
    hp = recm - cv2.GaussianBlur(recm, (0, 0), 6.0)
    g = (cv2.GaussianBlur(shift(a0, dx, dy), (0,0), 1.0) > 0.05).astype(np.float32)[..., None]
    return float(((hp ** 2) * g).sum())

# refinar localmente alrededor del shift de fase
best = (1e18, sx, sy, 1.0)
for ddy in np.arange(-1.5, 1.51, 0.5):
    for ddx in np.arange(-1.5, 1.51, 0.5):
        for s in np.arange(0.85, 1.21, 0.05):
            e = err_for(sx + ddx, sy + ddy, s)
            if e < best[0]:
                best = (e, sx + ddx, sy + ddy, float(s))
_, dx, dy, s = best
print(f"{v}: final dx={dx:.2f} dy={dy:.2f} s={s:.2f}")

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)
ccx, ccy = 90 + dx, 100 + dy
yy, xx = np.mgrid[0:RH, 0:RW]
ell = (((xx - ccx) / 56.0) ** 2 + ((yy - ccy) / 66.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_clean"; 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]
accF = np.zeros((RH, RW, 3), np.float64)
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)
    out = np.clip(reg * (1 - fe) + rec * fe, 0, 255)
    full[RY:RY+RH, RX:RX+RW] = out.astype(np.uint8)
    cv2.imwrite(f"{out_dir}/{i+1:05d}.png", full)
    accF += out.astype(np.float64)

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_align_eq.png", eq(accF/N))
print(f"{v}: align done -> {out_dir}")