"""Randomized degradation that made the LR images of the PSRGAN Sentinel-2 dataset.

Each HR image gets one random blur (Gaussian, horizontal or vertical motion, defocus)
and is then downsampled 4x with a random interpolation (bicubic, bilinear, area).

    python degradation.py hr_dir lr_dir --seed 0

Writes one LR PNG per HR image and degradations.csv with the drawn parameters.
"""
import argparse
import csv
import random
from pathlib import Path

import cv2
import numpy as np

INTERP_NAMES = {cv2.INTER_CUBIC: "bicubic", cv2.INTER_LINEAR: "bilinear", cv2.INTER_AREA: "area"}


def degrade(hr, scale=4, rng=random):
    """hr: (H, W, 3) float32 in [0, 1]. Returns the LR image and the drawn parameters."""
    h, w = hr.shape[:2]
    filter_type = rng.choice(["gaussian", "motion", "defocus"])
    sigma = 0.0

    if filter_type == "gaussian":
        # Isotropic Gaussian blur
        k = rng.choice([3, 5, 7])
        sigma = rng.uniform(0.5, 1.5)
        blurred = cv2.GaussianBlur(hr, (k, k), sigmaX=sigma)

    elif filter_type == "motion":
        # Line kernel, horizontal or vertical
        k = rng.choice([5, 7, 9])
        kernel = np.zeros((k, k), np.float32)
        if rng.random() < 0.5:
            kernel[k // 2, :] = 1.0
            filter_type = "motion_h"
        else:
            kernel[:, k // 2] = 1.0
            filter_type = "motion_v"
        blurred = cv2.filter2D(hr, -1, kernel / k)

    else:
        # Defocus as a disk kernel
        k = rng.choice([5, 7, 9])
        c = k // 2
        yy, xx = np.mgrid[:k, :k]
        kernel = ((yy - c) ** 2 + (xx - c) ** 2 <= c ** 2).astype(np.float32)
        blurred = cv2.filter2D(hr, -1, kernel / kernel.sum())

    interpolation = rng.choice([cv2.INTER_CUBIC, cv2.INTER_LINEAR, cv2.INTER_AREA])
    lr = cv2.resize(blurred, (w // scale, h // scale), interpolation=interpolation)
    return lr, {"blur": filter_type, "blur_kernel": k, "blur_sigma": sigma,
                "interpolation": INTERP_NAMES[interpolation]}


def main():
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument("hr_dir")
    p.add_argument("lr_dir")
    p.add_argument("--scale", type=int, default=4)
    p.add_argument("--seed", type=int, default=None)
    args = p.parse_args()

    rng = random.Random(args.seed)
    out = Path(args.lr_dir)
    out.mkdir(parents=True, exist_ok=True)
    rows = []
    for path in sorted(Path(args.hr_dir).iterdir()):
        img = cv2.imread(str(path), cv2.IMREAD_COLOR)
        if img is None:
            continue
        # BGR stays BGR, the filters act per channel
        hr = img.astype(np.float32) / 255
        lr, params = degrade(hr, args.scale, rng)
        lr = np.round(np.clip(lr, 0, 1) * 255).astype(np.uint8)
        cv2.imwrite(str(out / f"{path.stem}.png"), lr)
        rows.append({"file": path.name, **params})

    with open(out / "degradations.csv", "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=["file", "blur", "blur_kernel", "blur_sigma", "interpolation"])
        w.writeheader()
        w.writerows(rows)
    print(f"{len(rows)} LR images in {out}")


if __name__ == "__main__":
    main()