#!/usr/bin/env python3

import os

import sys

import json

import numpy as np

from PIL import Image

from pycocotools.mask import encode, decode

def get_outfile(infile, outdir):

    outfile = os.path.basename(infile)

    outsubdir = os.path.basename(os.path.dirname(infile))

    outfile = os.path.join(outdir, outsubdir, outfile)

    return outfile

def filegen(indir, ext='png'):

    for dd in os.scandir(indir):

        if not os.path.isdir(dd.path):

            continue

        for ff in os.scandir(dd.path):

            if not ff.name.endswith(ext):

                continue

            yield ff.path

def subsample_coco_mask(mask, factor):

    mask = Image.fromarray(mask)

    sz = np.asarray(mask.size)

    out_sz = (sz / factor).astype(int)

    mask.thumbnail(size, Image.ANTIALIAS)

    mask = np.asarray(mask, order='F')

    return mask

def subsample_verts(verts, factor):

    verts = (np.asarray(verts)//factor).tolist()

    vvprev = [None, None]

    newverts = []

    for vv in verts:

        if np.all(vv == vvprev):

           continue

        else:

            newverts.append(vv)

        vvprev = vv.copy()

    return newverts

def subsample_roi(fn, fnout, factor):

    with open(fn) as fh:

        rois = json.load(fh)

    if len(rois) == 0:

        print("NO ROIS IN\t%s" % fn)

    for roi in rois:

        mask = decode(roi)

        mask = subsample_coco_mask(mask, factor)

        cocomask = encode(mask)

        cocomask['counts'] = cocomask['counts'].decode()

        roi.update(cocomask)

        verts = roi["vertices"]

        roi.update({"vertices": subsample_verts(verts, factor)})

        if 'zoom' in roi:

            roi.pop('zoom')

    with open(fnout, 'w+') as fh:

        json.dump(rois, fh)

####################################################################

if __name__ == '__main__':

    import sys

    import argparse

    parser = argparse.ArgumentParser()

    parser.add_argument(

        'indir', type=str, help = 'input directory')

    parser.add_argument(

        '--outdir',

        default='',

        type=str,

        help = 'output directory')

    parser.add_argument(

      '--original-side',

      type=int,

      default=1024,

      help='Original side (in pixels) of the image patches')

    parser.add_argument('--img', dest='img', action='store_true')

    parser.add_argument('--no-img', dest='img', action='store_false')

    parser.set_defaults(img=True)

    group = parser.add_mutually_exclusive_group(required=True)

    group.add_argument(

      '--target-side',

      type=int,

      default=512,

      help='Requested side (in pixels) of the image patches')

    group.add_argument(

      '--factor',

      type=int,

      default=None,

      help='Requested subsampling rate')

    prms = parser.parse_args()

    if prms.factor is None:

        factor = prms.original_side // prms.target_side 

        side = prms.target_side

    else:

        factor = prms.factor

        side = prms.original_side // prms.factor

    size = side, side

    if len(prms.outdir)==0:

        basedir, set_ = os.path.split(prms.indir.rstrip('/'))

        basedir = os.path.dirname(basedir)

        #basedir = os.path.dirname(basedir)

        outdir = "{}/data_{}_subsample_{:d}x/fullsplit/all".format(basedir, side, factor)

    else:

        outdir = prms.outdir

    print("SAVING TO:", outdir, sep='\t')

    os.makedirs(outdir, exist_ok = True)

    if prms.img:

        print("SUBSAMPLING IMAGES")

        for infile in filegen(prms.indir):

            outfile = get_outfile(infile, outdir=outdir)

            os.makedirs(os.path.dirname(outfile), exist_ok=True)

            if infile != outfile:

                try:

                    im = Image.open(infile)

                    im.thumbnail(size, Image.ANTIALIAS)

                    im.save(outfile, "png")

                except IOError as ee:

                    print( "cannot create thumbnail for '%s'" % infile)

                    print(ee)

    print("SUBSAMPLING MASKS")

    for infile in filegen(prms.indir, ext = '.json'):

        outfile = get_outfile(infile, outdir=outdir)

        if infile != outfile:

            try:

                subsample_roi(infile, outfile, factor)

            except IOError as ee:

                print( "cannot create thumbnail for '%s'" % infile)

                print(ee)

            except Exception as ee:

                print( "cannot create thumbnail for '%s'" % infile)

                raise ee