#!/usr/bin/env python3

#
# slides-alpha.py
#
# Generates an alpha layer for slides (in PNG format), trying to replace the
# white background while leaving a small border around any element so they are
# still visible / readable over anything dark
#
# Copyright (c) 2023 Sylvain Munaut <tnt@246tNt.com>
# SPDX-License-Identifier: MIT
#

import argparse
import os
import pathlib
import re

from collections import namedtuple

from PIL import Image

import numpy as np

import scipy.ndimage
import scipy.signal


CropSpec = namedtuple("CropSpec", "w h x y")

def croptype(s):
	m = re.match("^([0-9]+)x([0-9]+)(\+([0-9]+)\+([0-9]+))?$", s)
	if not m:
		raise argparse.ArgumentTypeError(f"{s} is not a valid crop specification")

	return CropSpec(
		int(m.group(1)),
		int(m.group(2)),
		int(m.group(4) or 0),
		int(m.group(5) or 0),
	)


def slide_process_single(in_file, out_file, opts):
	# Debug
	print(f"Processing {str(in_file):s} -> {str(out_file):s}")

	# Open image
	img = Image.open(in_file)

	# Convert to NP array
	ia = np.array(img.convert('RGBA'))

	# Crop
	if opts.crop:
		c = opts.crop
		ia = ia[c.y:c.y+c.h,c.x:c.x+c.w]

	# Create alpha channel
	iaf = ia.astype(np.float32) / 255.0
	x = 1.0 - (iaf[:,:,0] + iaf[:,:,1] + iaf[:,:,2]) / 3.0
	y = np.where(x > 0.1, 1.0, 0.0)

	# Borders
	if opts.top_border:
		y[0:opts.top_border,    :] = np.zeros([opts.top_border,    y.shape[1]], dtype=np.float32)

	if opts.bottom_border:
		y[-opts.bottom_border:, :] = np.zeros([opts.bottom_border, y.shape[1]], dtype=np.float32)

	# Expand outward
	r = np.max(y.shape) // 384
	ye = scipy.signal.convolve2d(y, np.ones([r,r]), mode='same', boundary='fill', fillvalue=0)
	ye = ye.clip(0.0, 1.0)

	# Blur it
	r = np.max(y.shape) / 1280
	ye = scipy.ndimage.gaussian_filter(ye, r, mode='nearest')

	# Recombine
	y = (y + ye).clip(0.0, 1.0)
	#y = np.maximum(y, x)

	# Borders
	if opts.top_border:
		y[0:opts.top_border,    :] = np.ones([opts.top_border,    y.shape[1]], dtype=np.float32)

	if opts.bottom_border:
		y[-opts.bottom_border:, :] = np.ones([opts.bottom_border, y.shape[1]], dtype=np.float32)

	# Inject in original image
	ia[:,:,3] = (y * 255).astype(np.uint8)

	# Save
	img = Image.fromarray(ia)
	img.save(out_file)


def slide_process_directory(in_path, out_path, opts):
	for spe in os.listdir(in_path):
		# Create in/out path
		spi = in_path  / spe
		spo = out_path / spe

		# Skip hidden files
		if spe.startswith('.'):
			continue

		# Is this a PNG file ?
		if spe.endswith('.png') and spi.is_file():
			slide_process_single(spi, spo, opts)

		# Or a sub directory ?
		if spi.is_dir():
			spo.mkdir(exist_ok=opts.force)
			slide_process_directory(spi, spo, opts)


def parse_opt():
	parser = argparse.ArgumentParser(
		prog = 'slides-alpha.py',
		description = 'Alpha channel generator for slides',
	)

	parser.add_argument('-i', '--input',  required=True, type=pathlib.Path,
		help="Path to either input file or input directory to scan")
	parser.add_argument('-o', '--output', required=True, type=pathlib.Path,
		help="Path to either output file or output dirctory to write results to")
	parser.add_argument('-c', '--crop', type=croptype,
		help="Crop (WxH[+X+Y])")
	parser.add_argument('-t', '--top-border', type=int, default=0,
		help="Size in pixel of the solid top border")
	parser.add_argument('-b', '--bottom-border', type=int, default=0,
		help="Size in pixel of the solid bottom border")
	parser.add_argument('-f', '--force', action='store_true',
		help="Overwrite destination if it exists")

	return parser.parse_args()


def main():
	# Parse option
	opts = parse_opt()

	# Check single or recurse mode
	if opts.input.is_dir():
		# Directory
		if not opts.output.is_dir():
			raise RuntimeError("Output directory is not valid")

		slide_process_directory(opts.input, opts.output, opts)

	elif opts.input.is_file():
		# Single file
		if opts.output.exists():
			if not opts.output.is_file():
				raise RuntimeError("Output path already exists and is not a file")
			elif not opts.force:
				raise RuntimeError("Output path already exists")

		slide_process_single(opts.input, opts.output, opts)

	else:
		# Huh ???
		raise RuntimeError("No input file ?")


if __name__ == '__main__':
	main()