from pathlib import Path
import fire
import cv2
import numpy as np
from tqdm import tqdm


def xDoG(img, sigma, k_sigma, p, epsilon, phi):
    sigma_large = sigma * k_sigma
    sharped_img = (1 + p) * cv2.GaussianBlur(img, (0, 0), sigma) - p * cv2.GaussianBlur(img, (0, 0), sigma_large)
    img = np.multiply(img, sharped_img)

    t = np.zeros(img.shape)
    t[img >= epsilon] = 1.0
    img_dark_indices = img < epsilon
    t[img_dark_indices] = 1.0 + np.tanh(phi * (img[img_dark_indices] - epsilon))

    return t * 256


def transform_single(origin, to, anime: bool = True):
    img = cv2.imread(str(origin), cv2.IMREAD_GRAYSCALE)
    if anime:
        r = xDoG(img / 256, 0.3, 4.5, 19, 0.01, 10 ^ 8 * 2)
    else:
        r = xDoG(img / 256, 0.7, 5, 20, 0.02, 10 ^ 8 * 2)
    cv2.imwrite(str(to), r)


def transform(origin, to, anime=True):
    origin = Path(origin)
    to = Path(to)
    if origin.is_dir() and to.is_dir():
        for f in tqdm(origin.glob("*")):
            transform_single(f, to / f.name, anime)
    elif origin.is_file():
        transform_single(origin, to, anime)


if __name__ == '__main__':
    fire.Fire(transform)