from pathlib import Path

import cv2
import numpy as np
from bm3d import BM3DProfileLC, bm3d_rgb
from PIL import Image


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


source = np.asarray(Image.open(SOURCE).convert("RGB"), dtype=np.float32) / 255.0
restored = bm3d_rgb(source, sigma_psd=0.055, profile=BM3DProfileLC())
restored = np.clip(restored, 0.0, 1.0)

# Strong restoration over the background. Keep a small amount of the source so
# the broad photographic tonal structure survives without restoring the fine
# repeating texture.
result = restored * 0.96 + source * 0.04

# Preserve the ring and only its immediate contact shadow from the source.
h, w = source.shape[:2]
yy, xx = np.mgrid[0:h, 0:w]
cx, cy = 840.0, 486.0
rx, ry = 54.0, 40.0
distance = np.sqrt(((xx - cx) / rx) ** 2 + ((yy - cy) / ry) ** 2)
ring_mask = np.clip((1.22 - distance) / 0.38, 0.0, 1.0)[..., None]
result = result * (1.0 - ring_mask) + source * ring_mask

result_u8 = np.clip(result * 255.0 + 0.5, 0, 255).astype(np.uint8)
Image.fromarray(result_u8).save(OUTPUT_1X)

upscaled = cv2.resize(result_u8, (3840, 2160), interpolation=cv2.INTER_LANCZOS4)
Image.fromarray(upscaled).save(OUTPUT_4K)