"""Image degradations that mimic scans, phone photos, faxes and faded thermal receipts.

Geometric transforms return the homography so boxes can be moved with the pixels.
"""

import cv2
import numpy as np

from ..geometry import transform_quads


def _smooth_field(rng, h, w, cells=4):
    f = rng.random((cells, cells)).astype(np.float32)
    return cv2.resize(f, (w, h), interpolation=cv2.INTER_CUBIC)


def photometric(img, rng, strength=1.0, thermal=False):
    """img: uint8 HxW or HxWx3. Returns a degraded copy of the same shape."""
    out = img.astype(np.float32)
    h, w = out.shape[:2]
    if thermal or rng.random() < 0.15 * strength:          # faded print: lift ink toward paper, unevenly
        fade = 0.25 + 0.6 * _smooth_field(rng, h, w, int(rng.integers(2, 6)))
        if out.ndim == 3:
            fade = fade[..., None]
        paper = np.percentile(out, 90)
        out = paper - (paper - out) * (1 - fade * rng.uniform(0.3, 0.9))
    if rng.random() < 0.5 * strength:                      # contrast / brightness
        out = (out - 128) * rng.uniform(0.6, 1.3) + 128 + rng.uniform(-40, 30)
    if rng.random() < 0.3 * strength:                      # lighting gradient / shadow
        g = 0.55 + 0.45 * _smooth_field(rng, h, w, int(rng.integers(2, 4)))
        out *= g[..., None] if out.ndim == 3 else g
    if out.ndim == 3 and rng.random() < 0.3 * strength:   # colour cast
        out *= rng.uniform(0.85, 1.1, 3)
    out = np.clip(out, 0, 255).astype(np.uint8)
    if rng.random() < 0.3 * strength:                      # resolution loss
        f = rng.uniform(0.35, 0.8)
        small = cv2.resize(out, (max(1, int(w * f)), max(1, int(h * f))), interpolation=cv2.INTER_AREA)
        out = cv2.resize(small, (w, h), interpolation=rng.choice([cv2.INTER_LINEAR, cv2.INTER_NEAREST, cv2.INTER_CUBIC]))
    if rng.random() < 0.4 * strength:                      # blur
        if rng.random() < 0.8:
            out = cv2.GaussianBlur(out, (0, 0), rng.uniform(0.3, 1.3))
        else:
            k = int(rng.integers(3, 7))
            kern = np.zeros((k, k), np.float32)
            kern[k // 2, :] = 1.0 / k
            m = cv2.getRotationMatrix2D((k / 2 - 0.5, k / 2 - 0.5), rng.uniform(0, 180), 1)
            out = cv2.filter2D(out, -1, cv2.warpAffine(kern, m, (k, k)))
    if rng.random() < 0.5 * strength:                      # sensor / scan noise
        noise = rng.normal(0, rng.uniform(2, 14), out.shape)
        out = np.clip(out.astype(np.float32) + noise, 0, 255).astype(np.uint8)
    if rng.random() < 0.08 * strength:                     # salt & pepper speckle
        mask = rng.random(out.shape[:2])
        out[mask < 0.002] = 0
        out[mask > 0.998] = 255
    if rng.random() < 0.05 * strength:                     # fax-like binarisation
        g = out if out.ndim == 2 else cv2.cvtColor(out, cv2.COLOR_RGB2GRAY)
        _, b = cv2.threshold(g, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        out = b if out.ndim == 2 else cv2.cvtColor(b, cv2.COLOR_GRAY2RGB)
    if rng.random() < 0.6 * strength:                      # JPEG
        ok, enc = cv2.imencode(".jpg", out, [cv2.IMWRITE_JPEG_QUALITY, int(rng.integers(25, 95))])
        out = cv2.imdecode(enc, cv2.IMREAD_UNCHANGED)
    return out


def _background(rng, h, w):
    r = rng.random()
    if r < 0.4:
        base = rng.integers(20, 200, 3)
        bg = np.ones((h, w, 3), np.float32) * base
        bg += (rng.random((h // 8 + 1, w // 8 + 1, 3)) * 40 - 20).repeat(8, 0).repeat(8, 1)[:h, :w]
    elif r < 0.7:  # wood-ish stripes
        base = rng.integers(60, 180, 3).astype(np.float32)
        stripes = np.sin(np.linspace(0, rng.uniform(20, 80), w))[None, :, None] * rng.uniform(5, 25)
        bg = np.ones((h, w, 3), np.float32) * base + stripes
    else:
        bg = np.ones((h, w, 3), np.float32) * rng.integers(0, 255, 3)
        bg *= (0.6 + 0.4 * _smooth_field(rng, h, w, 3))[..., None]
    return np.clip(bg, 0, 255).astype(np.uint8)


def geometric(img, rng, mode=None):
    """Returns (image, 3x3 homography from input to output coordinates).

    modes: none | rotate (scanner skew) | photo (page on a background with perspective)
    """
    h, w = img.shape[:2]
    mode = mode or rng.choice(["none", "rotate", "photo"], p=[0.35, 0.4, 0.25])
    if mode == "none":
        return img, np.eye(3)
    if mode == "rotate":
        ang = rng.normal(0, 2.0) if rng.random() < 0.8 else rng.uniform(-8, 8)
        m = cv2.getRotationMatrix2D((w / 2, h / 2), ang, 1.0)
        cos, sin = abs(m[0, 0]), abs(m[0, 1])
        nw, nh = int(h * sin + w * cos), int(h * cos + w * sin)
        m[0, 2] += nw / 2 - w / 2
        m[1, 2] += nh / 2 - h / 2
        border = tuple(int(v) for v in np.median(img.reshape(-1, img.shape[-1]) if img.ndim == 3 else img.reshape(-1, 1), 0))
        out = cv2.warpAffine(img, m, (nw, nh), flags=cv2.INTER_LINEAR, borderValue=border)
        return out, np.vstack([m, [0, 0, 1]])
    # photo: page occupies 55-95% of a larger frame, corners jittered
    s = rng.uniform(1.05, 1.6)
    W, H = int(w * s), int(h * s)
    ox, oy = rng.uniform(0, W - w), rng.uniform(0, H - h)
    j = 0.04 * min(w, h)
    src = np.float32([[0, 0], [w, 0], [w, h], [0, h]])
    dst = np.float32([[ox, oy], [ox + w, oy], [ox + w, oy + h], [ox, oy + h]]) + rng.uniform(-j, j, (4, 2)).astype(np.float32)
    m = cv2.getPerspectiveTransform(src, dst)
    bg = _background(rng, H, W) if img.ndim == 3 else cv2.cvtColor(_background(rng, H, W), cv2.COLOR_RGB2GRAY)
    warped = cv2.warpPerspective(img, m, (W, H), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=0)
    mask = cv2.warpPerspective(np.full((h, w), 255, np.uint8), m, (W, H), flags=cv2.INTER_NEAREST)
    out = np.where((mask > 0)[..., None] if img.ndim == 3 else mask > 0, warped, bg)
    return out, m


def augment_page(img, quads, ignore_quads, rng, strength=1.0, thermal=False, geo_mode=None):
    img, m = geometric(img, rng, geo_mode)
    quads = transform_quads(quads, m)
    ignore_quads = transform_quads(ignore_quads, m)
    img = photometric(img, rng, strength, thermal)
    return img, quads, ignore_quads