#!/usr/bin/env python

import argparse
from time import time

import numpy as np
from PIL import Image

from internetarchivepdf.mrc import partial_blur, partial_boxblur
from optimiser import optimise_gray, optimise_rgb, optimise_gray2, optimise_rgb2

if __name__ == '__main__':
    parser = argparse.ArgumentParser('PDF recoder partial blurrer')
    parser.add_argument('--use-boxblur', default=False, action='store_true')
    parser.add_argument('--use-optimise', default=False, action='store_true')
    parser.add_argument('--sigma', default=1, type=int)
    parser.add_argument('--size', default=1, type=int)
    parser.add_argument('--optimise-size', default=3, type=int)
    parser.add_argument('infile', nargs='?', default=None)
    parser.add_argument('maskfile', nargs='?', default=None)
    parser.add_argument('outfile', nargs='?', default=None)

    args = parser.parse_args()

    img = Image.open(args.infile)
    mask = Image.open(args.maskfile)

    img_arr = np.array(img)
    mask_arr = np.array(mask)

    t = time()
    if args.use_optimise:
        height, width = img_arr.shape[0:2]
        if img.mode == 'L':
            arr = optimise_gray2(mask_arr, img_arr, width, height, args.optimise_size)
        else:
            arr = optimise_rgb2(mask_arr, img_arr, width, height, args.optimise_size)
    elif args.use_boxblur:
        arr = partial_boxblur(mask_arr, img_arr, size=args.size, mode=img.mode)
    else:
        arr = partial_blur(mask_arr, img_arr, sigma=args.sigma, mode=img.mode)
    print('Blur took:', time()-t)

    outimg = Image.fromarray(arr)
    outimg.save(args.outfile)
