"""Recolour dark lines red (--redshift) for photos with black or grey grids.

Darkness d = 1 - brightness / local paper brightness; blue and green darken by d, red by 0.25 d.
Only the networks see the recoloured photo.
"""

import cv2
import numpy as np

PAPER_KERNEL = 25
PAPER_BLUR = 15
RED_KEEP = 0.25


def redshift(photo):
    image = photo.astype(np.float32)
    gray = image.mean(2)
    paper = cv2.dilate(gray, np.ones((PAPER_KERNEL, PAPER_KERNEL), np.uint8))
    paper = cv2.GaussianBlur(paper, (0, 0), PAPER_BLUR)
    darkness = np.clip(1.0 - gray / np.maximum(paper, 1.0), 0.0, 1.0)[..., None]
    base = image / np.maximum(gray, 1.0)[..., None] * paper[..., None]
    out = base * (1.0 - darkness * np.array([1.0, 1.0, RED_KEEP], np.float32))
    return np.clip(out, 0, 255).astype(np.uint8)