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

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

fs = sorted(glob.glob(f"{W}/{v}/frames_full/*.png"))
N = len(fs); assert N == 240, N

# --- 1) media temporal (incremental) y deteccion del centro de la marca ---
acc = np.zeros((1280, 720, 3), np.float64)
for f in fs:
    acc += cv2.imread(f).astype(np.float64)
M = (acc / N)
gM = cv2.cvtColor(np.clip(M, 0, 255).astype(np.uint8), cv2.COLOR_BGR2GRAY).astype(np.float32)
# buscar el pico brillante (sparkle) en ventana alrededor de la posicion conocida
ys, ye, xs, xe = 1120, 1220, 550, 660
win = gM[ys:ye, xs:xe].copy()
win_b = cv2.GaussianBlur(win, (0, 0), sigmaX=2.0)
_, _, _, maxloc = cv2.minMaxLoc(win_b)
cx, cy = xs + maxloc[0], ys + maxloc[1]
print(f"{v}: centro marca detectado = ({cx},{cy})")

RX, RY = cx - 90, cy - 100            # esquina sup-izq de la region
RX = max(0, min(RX, 720 - RW)); RY = max(0, min(RY, 1280 - RH))
ccx, ccy = cx - RX, cy - RY           # centro en coords de region

# --- 2) regresor inicial B_t por inpaint Telea de la region ---
mask = np.zeros((RH, RW), np.uint8)
cv2.ellipse(mask, (ccx, ccy), (40, 50), 0, 0, 360, 255, -1)
mask = cv2.dilate(mask, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7)))

J = np.empty((N, RH, RW, 3), np.float32)
B = np.empty((N, RH, RW, 3), np.float32)
for i, f in enumerate(fs):
    reg = cv2.imread(f)[RY:RY+RH, RX:RX+RW]
    J[i] = reg.astype(np.float32)
    B[i] = cv2.inpaint(reg, mask, 4, cv2.INPAINT_TELEA).astype(np.float32)

# --- 3) descomposicion EM: estimar alpha y aW ---
def estimate(B, J, sm=0.4):
    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"  iter {k}: alpha max={alpha.max():.3f}")

# gate eliptico (limpia ruido lejos de la estrella)
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]

# --- 4) aplicar recuperacion a todos los frames y componer ---
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, f in enumerate(fs):
    full = cv2.imread(f)
    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

# --- 5) test de residuo: media temporal recuperado vs original (eq) ---
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}")