"""Training-set augmentation for mnist_learn_example datasets.

Input file layout (as in ~/run/exp): <train images><eval images>, 784 bytes per image.
Output: each of the first `--base` training images is kept and followed by `--copies` random
variants (shift up to ±2 px, rotation up to ±12°, scale 0.9–1.1, optional 1 px thicker/thinner
stroke); the training part is shuffled; the eval part is copied unchanged.

usage: python augment_mnist.py <src.bin> <src.target> <out_prefix> --train 30000 --base 30000 --copies 2
Writes <out_prefix><N>.bin/.target, where N = base * (copies + 1) is the new training size,
so that run.sh can use DATA=<out_prefix> and n=N.
"""
import argparse

import numpy as np
from PIL import Image

SIDE = 28
CROSS = [(0, 1), (0, -1), (1, 0), (-1, 0)]


def shift(a, dy, dx):
    return np.roll(np.roll(a, dy, 0), dx, 1)


def augment(img, rng):
    angle = rng.uniform(-12, 12)
    scale = rng.uniform(0.9, 1.1)
    dy, dx = rng.integers(-2, 3, size=2)
    out = Image.fromarray(img).rotate(angle, resample=Image.BILINEAR, center=(13.5, 13.5))
    if scale != 1:
        size = max(1, round(SIDE * scale))
        scaled = out.resize((size, size), Image.BILINEAR)
        canvas = Image.new("L", (SIDE, SIDE))
        off = (SIDE - size) // 2
        canvas.paste(scaled, (off, off)) if size <= SIDE else canvas.paste(scaled.crop((-off, -off, -off + SIDE, -off + SIDE)))
        out = canvas
    a = np.asarray(out, dtype=np.uint8)
    stroke = rng.choice(["same", "thicker", "thinner"], p=[0.5, 0.25, 0.25])
    if stroke == "thicker":
        a = np.maximum.reduce([a] + [shift(a, y, x) for y, x in CROSS])
    elif stroke == "thinner":
        a = np.minimum.reduce([a] + [shift(a, y, x) for y, x in CROSS])
    # Shift without wrap-around.
    res = np.zeros_like(a)
    ys, yd = (slice(dy, None), slice(None, SIDE - dy)) if dy >= 0 else (slice(None, dy), slice(-dy, None))
    xs, xd = (slice(dx, None), slice(None, SIDE - dx)) if dx >= 0 else (slice(None, dx), slice(-dx, None))
    res[ys, xs] = a[yd, xd]
    return res


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("images")
    ap.add_argument("labels")
    ap.add_argument("out")
    ap.add_argument("--train", type=int, required=True, help="training images in the source file")
    ap.add_argument("--base", type=int, required=True, help="how many source training images to use")
    ap.add_argument("--copies", type=int, default=2)
    ap.add_argument("--seed", type=int, default=20261014)
    a = ap.parse_args()
    rng = np.random.default_rng(a.seed)
    imgs = np.fromfile(a.images, dtype=np.uint8).reshape(-1, SIDE, SIDE)
    labels = open(a.labels).read().split()
    train_imgs, train_labels = [], []
    for i in range(a.base):
        train_imgs.append(imgs[i])
        train_labels.append(labels[i])
        for _ in range(a.copies):
            train_imgs.append(augment(imgs[i], rng))
            train_labels.append(labels[i])
    order = rng.permutation(len(train_imgs))
    out_imgs = np.concatenate([np.stack(train_imgs)[order], imgs[a.train:]])
    out_labels = [train_labels[j] for j in order] + labels[a.train:]
    n = len(train_imgs)
    out_imgs.tofile(f"{a.out}{n}.bin")
    with open(f"{a.out}{n}.target", "w") as f:
        f.write("\n".join(out_labels) + "\n")
    print(f"{a.out}{n}: {n} training + {len(imgs) - a.train} eval images")


if __name__ == "__main__":
    main()