"""Synthetic clean microscopy-style image + a noisy version (shot + read noise).

Bright blobs (cells) and thin filaments on a smooth background. The clean image is
the ground truth, so denoisers can be scored with PSNR / SSIM.
"""
from __future__ import annotations

import numpy as np


def clean_image(H: int = 256, W: int = 256, seed: int = 0) -> np.ndarray:
    rng = np.random.default_rng(seed)
    yy, xx = np.mgrid[0:H, 0:W]
    img = 0.08 + 0.05 * np.sin(xx / W * np.pi) * np.cos(yy / H * np.pi)

    for _ in range(rng.integers(14, 22)):                 # cells / blobs
        cy, cx = rng.uniform(12, H - 12), rng.uniform(12, W - 12)
        a, b = rng.uniform(5, 12), rng.uniform(5, 12)
        img += rng.uniform(0.4, 0.9) * np.exp(-(((xx - cx) ** 2) / (2 * a ** 2) +
                                                ((yy - cy) ** 2) / (2 * b ** 2)))

    for _ in range(rng.integers(4, 8)):                   # filaments (thin lines)
        x0, y0 = rng.uniform(0, W), rng.uniform(0, H)
        ang = rng.uniform(0, np.pi)
        t = np.linspace(0, rng.uniform(40, 120), 200)
        for s in t:
            px, py = x0 + s * np.cos(ang), y0 + s * np.sin(ang)
            if 1 <= px < W - 1 and 1 <= py < H - 1:
                img[int(py), int(px)] += 0.5
    return np.clip(img, 0, 1).astype(np.float32)


def add_noise(clean: np.ndarray, sigma: float = 0.12, peak: float = 40.0,
              seed: int = 0) -> np.ndarray:
    rng = np.random.default_rng(seed)
    shot = rng.poisson(np.clip(clean, 0, None) * peak) / peak     # Poisson (shot)
    noisy = shot + rng.normal(0, sigma, clean.shape)              # Gaussian (read)
    return np.clip(noisy, 0, 1).astype(np.float32)