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

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

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))
ccx, ccy = cx - RX, cy - RY                      # centro marca en region
ox, oy = RX - X0, RY - Y0                         # offset region dentro del crop ProPainter

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

J = np.empty((N, RH, RW, 3), np.float32)
B = np.empty((N, RH, RW, 3), np.float32)
for i in range(N):
    J[i] = cv2.imread(Jf[i])[RY:RY+RH, RX:RX+RW].astype(np.float32)
    B[i] = cv2.imread(Bf[i])[oy:oy+RH, ox:ox+RW].astype(np.float32)   # ProPainter (sin marca)

def estimate(B, J, sm=0.4, lp=2.5):
    if lp>0:
        B=np.stack([cv2.GaussianBlur(b,(0,0),lp) for b in B]); J=np.stack([cv2.GaussianBlur(j,(0,0),lp) for j in J])
    D = J - B
    mB = B.mean(0); mD = D.mean(0)
    varB = B.var(0).sum(2)
    covBD = ((B - mB) * (D - mD)).mean(0).sum(2)
    alpha = np.clip(-covBD / (varB + 1e-3), 0.0, 0.95)
    if sm > 0:
        alpha = cv2.GaussianBlur(alpha.astype(np.float32), (0, 0), sigmaX=sm)
    return alpha.astype(np.float32), (mD + alpha[..., None] * mB).astype(np.float32)

Bcur = B.copy()
for k in range(K_ITERS):
    alpha, aW = estimate(Bcur, J, sm=0.4)
    oneA = (1.0 - alpha)[..., None]
    for t in range(N):
        Bcur[t] = np.clip((J[t] - aW) / oneA, 0, 255)
    print(f"  {v} iter {k}: alpha max={alpha.max():.3f}")

yy, xx = np.mgrid[0:RH, 0:RW]
ell = (((xx - ccx) / 52.0) ** 2 + ((yy - ccy) / 62.0) ** 2) <= 1.0
alpha = np.where(ell, alpha, 0.0).astype(np.float32)
aW = np.where(ell[..., None], aW, 0.0).astype(np.float32)
alpha = cv2.GaussianBlur(alpha, (0, 0), sigmaX=0.4)
aW = cv2.GaussianBlur(aW, (0, 0), sigmaX=0.4)
oneA = (1.0 - alpha)[..., 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), sigmaX=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) / 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_orig_eq.png", eq(J[:120].mean(0)))
cv2.imwrite(f"{W}/{v}/res_rec_eq.png", eq(accr/nr))
cv2.imwrite(f"{W}/{v}/alpha.png", (np.clip(alpha,0,0.85)/0.85*255).astype(np.uint8))
print(f"{v}: hecho -> {out_dir}")