# SPDX-License-Identifier: Apache-2.0
"""Synthetic predictions with known shift, scale and blur, made from the full-size
published ink preview, plus a pure-noise control. Used by tests/test_bench.py and the
browser end-to-end test."""
import json, os, sys, math
import numpy as np
from PIL import Image
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'tools'))
import inkbench as ib
ROOT = os.path.join(os.path.dirname(__file__), '..')

def make(seg, tx, ty, sx, sy, sigma, c=1.5, seed=0, noise=0.0):
    idx = json.load(open(os.path.join(ROOT, 'data/bench/index.json')))['segments'][seg]
    W, H = idx['width'], idx['height']
    hi = np.asarray(Image.open(os.path.join(ROOT, 'data/hi/%s.jpg' % seg)).convert('L'), float) / 255
    Hh, Wh = hi.shape
    if sigma > 0:
        hi = ib.blur(hi, sigma * Wh / W)
    Wp, Hp = round(c * W), round(c * H)
    u = np.arange(Wp) + 0.0; v = np.arange(Hp) + 0.0
    x = ((u + 0.5) * W / Wp - W / 2) * sx + W / 2 - 0.5 + tx
    y = ((v + 0.5) * H / Hp - H / 2) * sy + H / 2 - 0.5 + ty
    X = (x + 0.5) * Wh / W - 0.5; Y = (y + 0.5) * Hh / H - 0.5
    ok = (X >= 0) & (X <= Wh - 1); okY = (Y >= 0) & (Y <= Hh - 1)
    X = np.clip(X, 0, Wh - 1.001); Y = np.clip(Y, 0, Hh - 1.001)
    x0 = np.floor(X).astype(int); y0 = np.floor(Y).astype(int); fx = X - x0; fy = Y - y0
    a = hi[y0[:, None], x0[None]] * (1 - fx) + hi[y0[:, None], x0[None] + 1] * fx
    b = hi[y0[:, None] + 1, x0[None]] * (1 - fx) + hi[y0[:, None] + 1, x0[None] + 1] * fx
    out = a * (1 - fy[:, None]) + b * fy[:, None]
    out = np.where(okY[:, None] & ok[None], out, 0)
    if noise:
        out = out + np.random.default_rng(seed).normal(0, noise, out.shape)
    return np.clip(out, 0, 1)

def noise(seg, seed=1, c=1.5):
    idx = json.load(open(os.path.join(ROOT, 'data/bench/index.json')))['segments'][seg]
    return np.random.default_rng(seed).random((round(c * idx['height']), round(c * idx['width'])))

if __name__ == '__main__':
    out = sys.argv[1]; os.makedirs(out, exist_ok=True)
    cases = {'shift_blur': ('w046', 23.0, -17.0, 1.0, 1.0, 2.0), 'scale_shift': ('w048', -12.5, 9.25, 1.12, 1.12, 1.0)}
    for k, (seg, *p) in cases.items():
        ib.save_png16(make(seg, *p), os.path.join(out, k + '.png'))
    Image.fromarray((noise('w046') * 255).astype(np.uint8)).save(os.path.join(out, 'noise.png'))
    json.dump(cases, open(os.path.join(out, 'cases.json'), 'w'))