"""Make a page image look like a photocopy of a photocopy.

`intensity` scales the whole effect: 1.0 is a badly worn copy of a copy, 0.5 a
tired one, 0 leaves the image alone.

Each "generation" re-scans the previous copy: a little skew and shift, blur from
the optics, a steeper toner curve, uneven toner, speckle, dust streaks down the
page and dark shadows where the lid did not sit flat. Colour survives only as
faint warmth, so coffee stains stay recognisable but lose their punch.
"""
import numpy as np
from PIL import Image, ImageDraw, ImageFilter


def _lowfreq(rng, w, h, cells_x, cells_y, amp):
    """Smooth random field in [-amp, amp], roughly cells_x by cells_y blobs."""
    small = rng.normal(0, 1, (cells_y + 2, cells_x + 2)).astype(np.float32)
    small = (small - small.min()) / (np.ptp(small) + 1e-6) * 2 - 1
    img = Image.fromarray(small, mode="F").resize((w, h), Image.BICUBIC)
    return np.asarray(img, dtype=np.float32) * amp


def _speckles(rng, w, h, count, rmin, rmax):
    mask = Image.new("L", (w, h), 0)
    d = ImageDraw.Draw(mask)
    count = max(0, count)
    for _ in range(count):
        x, y = rng.uniform(0, w), rng.uniform(0, h)
        r = rng.uniform(rmin, rmax)
        d.ellipse((x - r, y - r, x + r, y + r), fill=int(rng.uniform(120, 255)))
    return np.asarray(mask.filter(ImageFilter.GaussianBlur(0.5)), dtype=np.float32) / 255.0


def _streaks(rng, arr, gen, t):
    h, w = arr.shape
    rows = np.arange(h)
    for _ in range(int(rng.integers(2, 5 + gen))):
        x0 = rng.uniform(0, w)
        xs = x0 + np.cumsum(rng.normal(0, 0.12, h))          # the streak wanders a little
        width = int(rng.integers(1, 4))
        strength = rng.uniform(0.03, 0.13) * t
        # dust comes and goes along the page
        fade = 0.55 + 0.45 * np.sin(rows / rng.uniform(60, 400) + rng.uniform(0, 6.28))
        fade *= np.clip(_1d(rng, h) + 0.6, 0, 1)
        for dx in range(width):
            cols = np.clip(xs.astype(int) + dx, 0, w - 1)
            arr[rows, cols] -= strength * fade * (1.0 if dx == 0 else 0.55)
    # one broad, faint band where the roller was tired
    x0 = int(rng.uniform(0.1, 0.9) * w)
    arr[:, x0:x0 + int(rng.uniform(20, 60))] *= 1 - (1 - rng.uniform(0.985, 0.997)) * t


def _1d(rng, n):
    pts = rng.normal(0, 0.5, n // 40 + 3).astype(np.float32)
    return np.interp(np.arange(n), np.linspace(0, n - 1, len(pts)), pts)


def _edge_shadows(rng, arr, t):
    h, w = arr.shape
    for edge in rng.permutation(4)[: int(rng.integers(1, 3))]:
        depth = int(rng.uniform(8, 26))
        strength = rng.uniform(0.18, 0.42) * t
        ramp = np.linspace(strength, 0, depth, dtype=np.float32) ** 1.6
        jag = 1 + 0.35 * _1d(rng, max(h, w))
        if edge == 0:
            arr[:, :depth] -= ramp[None, :] * jag[:h, None]
        elif edge == 1:
            arr[:, w - depth:] -= ramp[::-1][None, :] * jag[:h, None]
        elif edge == 2:
            arr[:depth, :] -= ramp[:, None] * jag[None, :w]
        else:
            arr[h - depth:, :] -= ramp[::-1][:, None] * jag[None, :w]


def _generation(rng, luma, gen, t):
    """One trip through the copier. luma is a float array in [0, 1]."""
    h, w = luma.shape
    im = Image.fromarray((np.clip(luma, 0, 1) * 255).astype(np.uint8), "L")

    # placed slightly crooked on the glass; the lid backing shows in the wedges
    im = im.rotate(rng.uniform(-0.28, 0.28) * t, resample=Image.BICUBIC, fillcolor=70,
                   translate=(rng.uniform(-4, 4) * t, rng.uniform(-4, 4) * t))
    # toner spreads: thin strokes fatten a little on every pass
    fat = np.asarray(im.filter(ImageFilter.MinFilter(3)), dtype=np.float32) / 255.0
    thin = np.asarray(im, dtype=np.float32) / 255.0
    im = Image.fromarray((np.clip((1 - 0.14 * t) * thin + 0.14 * t * fat, 0, 1) * 255).astype(np.uint8), "L")
    if t > 0:
        im = im.filter(ImageFilter.GaussianBlur(t * (0.55 + 0.05 * gen)))    # optics
    a = np.asarray(im, dtype=np.float32) / 255.0

    # a faint ghost from the previous pass being a hair out of register
    dy, dx = int(rng.integers(-2, 3)), int(rng.integers(-2, 3))
    if dx or dy:
        a = (1 - 0.12 * t) * a + 0.12 * t * np.roll(a, (dy, dx), axis=(0, 1))

    # toner curve: blacker blacks, paler paper, but keep mid-tones for stains
    k, c = 4.6 + 0.4 * gen, 0.50
    s = lambda x: 1 / (1 + np.exp(-k * (x - c)))
    a = (1 - t) * a + t * (s(a) - s(0.0)) / (s(1.0) - s(0.0))

    # uneven toner and grain
    a *= 1 + _lowfreq(rng, w, h, 5, 7, 0.022 * t) + _lowfreq(rng, w, h, w // 45, h // 45, 0.014 * t)
    a += rng.normal(0, 0.018 * t, a.shape).astype(np.float32)

    # dust and toner specks; a few white dropouts inside the ink
    a = np.minimum(a, 1 - 0.65 * _speckles(rng, w, h, int((150 + 90 * gen) * t), 0.5, 1.5))
    a = np.maximum(a, 0.85 * _speckles(rng, w, h, int(110 * t), 0.5, 1.3) * (a < 0.35))

    _streaks(rng, a, gen, t)
    _edge_shadows(rng, a, t)
    return np.clip(a, 0, 1)


def photocopy(img, seed, copies=3, intensity=1.0):
    """Return a degraded RGB copy of `img` (an RGB PIL image)."""
    rng = np.random.default_rng(seed)
    ycc = img.convert("YCbCr")
    luma = np.asarray(ycc.split()[0], dtype=np.float32) / 255.0
    for g in range(copies):
        luma = _generation(rng, luma, g, intensity)

    # keep only a whisper of the original colour, smeared
    _, cb, cr = ycc.split()
    soften = lambda ch: ch.filter(ImageFilter.GaussianBlur(5)).point(lambda v: 128 + (v - 128) * (1 - 0.15 * intensity))
    y = Image.fromarray((luma * 255).astype(np.uint8), "L")
    return Image.merge("YCbCr", (y, soften(cb), soften(cr))).convert("RGB")