from pathlib import Path

import cv2
import numpy as np
from PIL import Image


SOURCE = Path(".runtime/imagegen/shore-a-restored-bm3d-strong-v1.png")
ORIGINAL = Path(
    "/home/vela/.codex-account-2/generated_images/"
    "01a018e9-399a-7293-8554-1323cc25a974/"
    "exec-6fb9c831-d2ea-4598-938d-99a02185da0b.png"
)


base = np.asarray(Image.open(SOURCE).convert("RGB"), dtype=np.uint8)
original = np.asarray(Image.open(ORIGINAL).convert("RGB"), dtype=np.uint8)
h, w = base.shape[:2]
yy, xx = np.mgrid[0:h, 0:w]

# Preserve the ring and immediate contact shadow exactly.
distance = np.sqrt(((xx - 840.0) / 54.0) ** 2 + ((yy - 486.0) / 40.0) ** 2)
ring_mask = np.clip((1.22 - distance) / 0.38, 0.0, 1.0)[..., None].astype(np.float32)

# Keep the main surf band somewhat more photographic than the wet foreground.
surf_keep = np.clip((220.0 - yy) / 90.0, 0.0, 1.0)[..., None].astype(np.float32)

settings = {
    "balanced": (24, 10, 0.18),
    "clean": (36, 14, 0.10),
    "cleanest": (48, 18, 0.05),
}

for name, (sigma_color, sigma_space, source_mix) in settings.items():
    filtered = cv2.bilateralFilter(
        base,
        d=0,
        sigmaColor=sigma_color,
        sigmaSpace=sigma_space,
        borderType=cv2.BORDER_REFLECT,
    ).astype(np.float32)
    base_f = base.astype(np.float32)
    result = filtered * (1.0 - source_mix) + base_f * source_mix
    result = result * (1.0 - 0.12 * surf_keep) + base_f * (0.12 * surf_keep)
    result = result * (1.0 - ring_mask) + original.astype(np.float32) * ring_mask
    result_u8 = np.clip(result + 0.5, 0, 255).astype(np.uint8)
    Image.fromarray(result_u8).save(
        f".runtime/imagegen/shore-a-restored-{name}-v1.png"
    )
    upscaled = cv2.resize(result_u8, (3840, 2160), interpolation=cv2.INTER_LANCZOS4)
    Image.fromarray(upscaled).save(
        f".runtime/imagegen/shore-a-restored-{name}-4k-v1.png"
    )