from __future__ import division

import sys
import warnings;

warnings.filterwarnings('ignore')
from tensorflow.keras.preprocessing.image import save_img, img_to_array, array_to_img
from os.path import join, exists, dirname, realpath
from tensorflow.keras.applications.inception_v3 import preprocess_input
from tensorflow.keras.applications import InceptionV3
from sklearn.metrics import pairwise_distances_argmin_min
import subprocess


def install(package):
    subprocess.check_call([sys.executable, "-m", "pip", "install", package])

try:

    from keras.backend.tensorflow_backend import set_session
except BaseException as e:
    print("Keras Tensorflor Backend Failled With Error: ", e)
    print("Using most recent tensorflow keras backend...")
    from tensorflow.compat.v1.keras.backend import set_session

    print("Success! :)")


try:
    from pointgrid import align_points_to_grid

except BaseException as e:

    install("pointgrid")
    from pointgrid import align_points_to_grid
try:
    from iiif_downloader import Manifest
except BaseException as e:

    install("iiif_downloader")
    from iiif_downloader import Manifest

try:
    from rasterfairy import coonswarp
except BaseException as e:

    install("rasterfairy")
    from rasterfairy import coonswarp

try:
    from hdbscan import HDBSCAN
except BaseException as e:

    install("hdbscan")
    from hdbscan import HDBSCAN
try:
    import lap
except BaseException as e:

    install("lap")
    import lap

from dateutil.parser import parse as parse_date
from sklearn.preprocessing import minmax_scale
from keras_preprocessing.image import load_img

from scipy.spatial.distance import cdist
from distutils.dir_util import copy_tree
from sklearn.decomposition import PCA
from scipy.spatial import ConvexHull
# from iiif_downloader import Manifest
from collections import defaultdict

import matplotlib.pyplot as plt
from tensorflow.keras.models import Model
from scipy.stats import kde
# from hdbscan import HDBSCAN
import tensorflow as tf

# import tensorflow as tf
tf.python.control_flow_ops = tf
from umap import UMAP
import multiprocessing
import pkg_resources
import rasterfairy
import numpy as np
import datetime
import operator
import argparse
import random
import shutil
import glob2
import uuid
import math
import gzip
import json
# import lap

import csv
import os

from PIL import ImageFile

ImageFile.LOAD_TRUNCATED_IMAGES = True

try:
    from MulticoreTSNE import MulticoreTSNE as TSNE
except:
    from sklearn.manifold import TSNE

try:
    from urllib.parse import unquote  # python 3
except:
    from urllib import unquote  # python 2

# handle dynamic GPU memory allocation
tf_config = tf.compat.v1.ConfigProto()
tf_config.gpu_options.allow_growth = True
tf_config.log_device_placement = True
sess = tf.compat.v1.Session(config=tf_config)

'''
NB: Keras Image class objects return image.size as w,h
    Numpy array representations of images return image.shape as h,w,c
'''

config = {
    'images': None,
    'metadata': None,
    'out_dir': 'output',
    'max_images': None,
    'use_cache': True,
    'encoding': 'utf8',
    'min_cluster_size': 20,
    'atlas_size': 2048,
    'cell_size': 32,
    'lod_cell_height': 128,
    'n_neighbors': 6,
    'min_distance': 0.001,
    'metric': 'correlation',
    'pointgrid_fill': 0.05,
    'square_cells': False,
    'gzip': False,
    'plot_id': str(uuid.uuid1()),
    'seed': 24,
}


##
# Entry
##


def process_images(**kwargs):
    '''Main method for processing user images and metadata'''
    np.random.seed(kwargs['seed'])
    tf.compat.v1.set_random_seed(kwargs['seed'])
    copy_web_assets(**kwargs)
    kwargs['out_dir'] = join(kwargs['out_dir'], 'data')
    kwargs['image_paths'], kwargs['metadata'] = filter_images(**kwargs)
    kwargs['atlas_dir'] = get_atlas_data(**kwargs)
    get_manifest(**kwargs)
    write_images(**kwargs)
    print(' * done!')


def copy_web_assets(**kwargs):
    '''Copy the /web directory from the pixplot_notebook source to the users cwd'''
    src = join(dirname(realpath(__file__)), 'web')
    dest = join(os.getcwd(), kwargs['out_dir'])
    copy_tree(src, dest)
    # write version numbers into output
    for i in ['index.html', os.path.join('assets', 'js', 'tsne.js')]:
        path = os.path.join(dest, i)
        with open(path, 'r') as f:
            f = f.read().replace('VERSION_NUMBER', get_version())
            with open(path, 'w') as out:
                out.write(f)
    if kwargs['copy_web_only']: sys.exit()


##
# Images
##


def filter_images(**kwargs):
    '''Main method for filtering images given user metadata (if provided)'''
    image_paths = []
    for i in stream_images(image_paths=get_image_paths(**kwargs)):
        # get image height and width
        w, h = i.original.size
        # remove images with 0 height or width when resized to lod height
        if (h == 0) or (w == 0):
            print(' * skipping {} because it contains 0 height or width'.format(i.path))
            continue
        # remove images that have 0 height or width when resized
        try:
            resized = i.resize_to_max(kwargs['lod_cell_height'])
        except ValueError:
            print(' * skipping {} because it contains 0 height or width when resized'.format(i.path))
            continue
        # remove images that are too wide for the atlas
        if (w / h) > (kwargs['atlas_size'] / kwargs['cell_size']):
            print(' * skipping {} because its dimensions are oblong'.format(i.path))
            continue
        image_paths.append(i.path)
    # handle the case user provided no metadata
    if not kwargs.get('metadata', False):
        return [
            limit_image_count(image_paths, **kwargs),
            [],
        ]
    # handle user metadata: retain only records with image and metadata
    l = get_metadata_list(**kwargs)
    meta_bn = set([clean_filename(i.get('filename', '')) for i in l])
    img_bn = set([clean_filename(i) for i in image_paths])
    # identify images with metadata and those without metadata
    meta_present = img_bn.intersection(meta_bn)
    meta_missing = list(img_bn - meta_bn)
    # notify the user of images that are missing metadata
    if meta_missing:
        print(' ! Some images are missing metadata:\n  -', '\n  - '.join(meta_missing[:10]))
        if len(meta_missing) > 10: print(' ...', len(meta_missing) - 10, 'more')
        with open('missing-metadata.txt', 'w') as out:
            out.write('\n'.join(meta_missing))
    # get the sorted lists of images and metadata
    d = {clean_filename(i['filename']): i for i in l}
    images = []
    metadata = []
    for i in image_paths:
        if clean_filename(i) in meta_present:
            images.append(i)
            metadata.append(d[clean_filename(i)])
    kwargs['metadata'] = metadata
    write_metadata(**kwargs)
    return [
        limit_image_count(images, **kwargs),
        limit_image_count(metadata, **kwargs),
    ]


def limit_image_count(arr, **kwargs):
    '''If the user passed a max_images value, return [:max_images] from arr'''
    if kwargs.get('max_images', False):
        return arr[:kwargs['max_images']]
    return arr


def get_image_paths(**kwargs):
    '''Called once to provide a list of image paths--handles IIIF manifest input'''
    # handle case where --images points to iiif manifest
    image_paths = None
    if not kwargs['images']:
        print('\nError: please provide an images argument, e.g.:')
        print('pixplot --images "cat_pictures/*.jpg"\n')
        sys.exit()
    # handle list of IIIF image inputs
    if os.path.exists(kwargs['images']):
        with open(kwargs['images']) as f:
            f = [i.strip() for i in f.read().split('\n') if i.strip()]
            if [i.startswith('http') for i in f]:
                for i in f: Manifest(url=i).save_images(limit=1)
                image_paths = sorted(glob2.glob(os.path.join('iiif-downloads', 'images', '*')))
    # handle case where images flag points to a glob of images
    if not image_paths:
        image_paths = sorted(glob2.glob(kwargs['images']))
    # handle case user provided no images
    if not image_paths:
        print('\nError: No input images were found. Please check your --images glob\n')
        sys.exit()
    # optional shuffle that mutates image_paths
    if kwargs['shuffle']:
        print(' * shuffling input images')
        random.shuffle(image_paths)
    return image_paths


def stream_images(**kwargs):
    '''Read in all images from args[0], a list of image paths'''
    for idx, i in enumerate(kwargs['image_paths']):
        try:
            metadata = None
            if kwargs.get('metadata', False) and kwargs['metadata'][idx]:
                metadata = kwargs['metadata'][idx]
            yield Image(i, metadata=metadata)
        except Exception as exc:
            print(' * image', i, 'could not be processed --', exc)


def clean_filename(s):
    '''Given a string that points to a filename, return a clean filename'''
    return unquote(os.path.basename(s))


##
# Metadata
##


def get_metadata_list(**kwargs):
    '''Return a list of objects with image metadata'''
    if not kwargs['metadata']: return []
    # handle csv metadata
    l = []
    if kwargs['metadata'].endswith('.csv'):
        with open(kwargs['metadata']) as f:
            reader = csv.reader(f)
            headers = [i.lower() for i in next(reader)]
            for i in reader:
                l.append({headers[j]: i[j] if i[j] else '' for j, _ in enumerate(headers)})
    # handle json metadata
    else:
        for i in glob2.glob(kwargs['metadata']):
            with open(i) as f:
                l.append(json.load(f))
    return l


def write_metadata(metadata, **kwargs):
    '''Write list `metadata` of objects to disk'''
    if not metadata: return
    out_dir = join(kwargs['out_dir'], 'metadata')
    for i in ['filters', 'options', 'file']:
        out_path = join(out_dir, i)
        if not exists(out_path): os.makedirs(out_path)
    # create the lists of images with each tag
    d = defaultdict(list)
    for i in metadata:
        filename = clean_filename(i['filename'])
        i['tags'] = [j.strip() for j in i.get('tags', '').split('|')]
        for j in i['tags']: d['__'.join(j.split())].append(filename)
        write_json(os.path.join(out_dir, 'file', filename + '.json'), i, **kwargs)
    write_json(os.path.join(out_dir, 'filters', 'filters.json'), [{
        'filter_name': 'select',
        'filter_values': list(d.keys()),
    }], **kwargs)
    # create the options for the category dropdown
    for i in d:
        write_json(os.path.join(out_dir, 'options', i + '.json'), d[i], **kwargs)
    # create the map from date to images with that date (if dates present)
    date_d = defaultdict(list)
    for i in metadata:
        date = i.get('year', '')
        if date:
            date_d[date].append(clean_filename(i['filename']))
    # find the min and max dates to show on the date slider
    dates = np.array([int(i.strip()) for i in date_d if is_number(i)])
    domain = {'min': float('inf'), 'max': -float('inf')}
    mean = np.mean(dates)
    std = np.std(dates)
    for i in dates:
        # update the date domain with all non-outlier dates
        if abs(mean - i) < (std * 4):
            domain['min'] = int(min(i, domain['min']))
            domain['max'] = int(max(i, domain['max']))
    # write the dates json
    if len(date_d) > 1:
        write_json(os.path.join(out_dir, 'dates.json'), {
            'domain': domain,
            'dates': date_d,
        }, **kwargs)


def is_number(s):
    '''Return a boolean indicating if a string is a number'''
    try:
        int(s)
        return True
    except:
        return False


##
# Main
##


def get_manifest(**kwargs):
    '''Create and return the base object for the manifest output file'''
    # load the atlas data
    atlas_data = json.load(open(join(kwargs['atlas_dir'], 'atlas_positions.json')))
    # store each cell's size and atlas position
    atlas_ids = set([i['idx'] for i in atlas_data])
    sizes = [[] for _ in atlas_ids]
    pos = [[] for _ in atlas_ids]
    for idx, i in enumerate(atlas_data):
        sizes[i['idx']].append([i['w'], i['h']])
        pos[i['idx']].append([i['x'], i['y']])
    # obtain the paths to each layout's JSON positions
    layouts = get_layouts(**kwargs)
    # create a heightmap for the umap layout
    if 'umap' in layouts and layouts['umap']:
        get_heightmap(layouts['umap']['layout'], 'umap', **kwargs)
    # specify point size scalars
    point_sizes = {}
    point_sizes['min'] = 0
    point_sizes['grid'] = 1 / math.ceil(len(kwargs['image_paths']) ** (1 / 2))
    point_sizes['max'] = point_sizes['grid'] * 1.2
    point_sizes['scatter'] = point_sizes['grid'] * .2
    point_sizes['initial'] = point_sizes['scatter']
    # fetch the date distribution data for point sizing
    if 'date' in layouts and layouts['date']:
        date_layout = read_json(layouts['date']['labels'], **kwargs)
        point_sizes['date'] = 1 / ((date_layout['cols'] + 1) * len(date_layout['labels']))
    # create manifest json
    manifest = {
        'version': get_version(),
        'plot_id': kwargs['plot_id'],
        'layouts': layouts,
        'initial_layout': 'umap',
        'point_sizes': point_sizes,
        'imagelist': get_path('imagelists', 'imagelist', **kwargs),
        'atlas_dir': kwargs['atlas_dir'],
        'metadata': True if kwargs['metadata'] else False,
        'default_hotspots': get_hotspots(vecs=read_json(layouts['umap']['layout'], **kwargs), **kwargs),
        'custom_hotspots': get_path('hotspots', 'user_hotspots', add_hash=False, **kwargs),
        'config': {
            'sizes': {
                'atlas': kwargs['atlas_size'],
                'cell': kwargs['cell_size'],
                'lod': kwargs['lod_cell_height'],
            },
        },
        'creation_date': datetime.datetime.today().strftime('%d-%B-%Y-%H:%M:%S'),
    }
    path = get_path('manifests', 'manifest', **kwargs)
    write_json(path, manifest, **kwargs)
    path = get_path(None, 'manifest', add_hash=False, **kwargs)
    write_json(path, manifest, **kwargs)
    # create images json
    imagelist = {
        'cell_sizes': sizes,
        'images': [clean_filename(i) for i in kwargs['image_paths']],
        'atlas': {
            'count': len(atlas_ids),
            'positions': pos,
        },
    }
    write_json(manifest['imagelist'], imagelist)


##
# Atlases
##


def get_atlas_data(**kwargs):
    '''
    Generate and save to disk all atlases to be used for this visualization
    If square, center each cell in an nxn square, else use uniform height
    '''
    # if the atlas files already exist, load from cache
    out_dir = os.path.join(kwargs['out_dir'], 'atlases', kwargs['plot_id'])
    if os.path.exists(out_dir) and kwargs['use_cache'] and not kwargs.get('shuffle', False):
        print(' * loading saved atlas data')
        return out_dir
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)
    # else create the atlas images and store the positions of cells in atlases
    print(' * creating atlas files')
    n = 0  # number of atlases
    x = 0  # x pos in atlas
    y = 0  # y pos in atlas
    positions = []  # l[cell_idx] = atlas data
    atlas = np.zeros((kwargs['atlas_size'], kwargs['atlas_size'], 3))
    for idx, i in enumerate(stream_images(**kwargs)):
        if kwargs['square_cells']:
            cell_data = i.resize_to_square(kwargs['cell_size'])
        else:
            cell_data = i.resize_to_height(kwargs['cell_size'])
        _, v, _ = cell_data.shape
        appendable = False
        if (x + v) <= kwargs['atlas_size']:
            appendable = True
        elif (y + (2 * kwargs['cell_size'])) <= kwargs['atlas_size']:
            y += kwargs['cell_size']
            x = 0
            appendable = True
        if not appendable:
            save_atlas(atlas, out_dir, n)
            n += 1
            atlas = np.zeros((kwargs['atlas_size'], kwargs['atlas_size'], 3))
            x = 0
            y = 0
        atlas[y:y + kwargs['cell_size'], x:x + v] = cell_data
        # find the size of the cell in the lod canvas
        lod_data = i.resize_to_max(kwargs['lod_cell_height'])
        h, w, _ = lod_data.shape  # h,w,colors in lod-cell sized image `i`
        positions.append({
            'idx': n,  # atlas idx
            'x': x,  # x offset of cell in atlas
            'y': y,  # y offset of cell in atlas
            'w': w,  # w of cell at lod size
            'h': h,  # h of cell at lod size
        })
        x += v
    save_atlas(atlas, out_dir, n)
    out_path = os.path.join(out_dir, 'atlas_positions.json')
    with open(out_path, 'w') as out:
        json.dump(positions, out)
    return out_dir


def save_atlas(atlas, out_dir, n):
    '''Save an atlas to disk'''
    out_path = join(out_dir, 'atlas-{}.jpg'.format(n))
    save_img(out_path, atlas)


##
# Layouts
##


def get_layouts(**kwargs):
    '''Get the image positions in each projection'''
    vecs = vectorize_images(**kwargs)
    umap = get_umap_layout(vecs=vecs, **kwargs)
    linear_assignment = get_lap_layout(umap=umap, **kwargs)
    grid = get_grid_layout(**kwargs)
    umap_jittered = get_pointgrid_layout(umap, 'umap', **kwargs)
    categorical = get_categorical_layout(**kwargs)
    date = get_date_layout(**kwargs)
    layouts = {
        'umap': {
            'layout': umap,
            'jittered': umap_jittered,
        },
        'alphabetic': {
            'layout': grid,
        },
        'grid': {
            'layout': linear_assignment,
        },
        'categorical': categorical,
        'date': date,
    }
    return layouts


def vectorize_images(**kwargs):
    '''Create and return vector representation of Image() instances'''
    print(' * preparing to vectorize {} images'.format(len(kwargs['image_paths'])))
    vector_dir = os.path.join(kwargs['out_dir'], 'image-vectors')
    if not os.path.exists(vector_dir): os.makedirs(vector_dir)
    base = InceptionV3(include_top=True, weights='imagenet', )
    model = Model(inputs=base.input, outputs=base.get_layer('avg_pool').output)
    print(' * creating image array')
    vecs = []
    for idx, i in enumerate(stream_images(**kwargs)):
        vector_path = os.path.join(vector_dir, os.path.basename(i.path) + '.npy')
        if os.path.exists(vector_path) and kwargs['use_cache']:
            vec = np.load(vector_path)
        else:
            im = preprocess_input(img_to_array(i.original.resize((299, 299))))
            vec = model.predict(np.expand_dims(im, 0)).squeeze()
            np.save(vector_path, vec)
        vecs.append(vec)
        print(' * vectorized {}/{} images'.format(idx + 1, len(kwargs['image_paths'])))
    return np.array(vecs)


def get_umap_layout(**kwargs):
    '''Get the x,y positions of images passed through a umap projection'''
    print(' * creating UMAP layout')
    out_path = get_path('layouts', 'umap', **kwargs)
    if os.path.exists(out_path) and kwargs['use_cache']: return out_path
    model = UMAP(n_neighbors=kwargs['n_neighbors'],
                 min_dist=kwargs['min_distance'],
                 metric=kwargs['metric'])
    # run PCA to reduce dimensionality of image vectors
    w = PCA(n_components=min(100, len(kwargs['vecs']))).fit_transform(kwargs['vecs'])
    # fetch categorical labels for images (if provided)
    y = []
    if kwargs.get('metadata', False):
        labels = [i.get('label', None) for i in kwargs['metadata']]
        # if the user provided labels, integerize them
        if any([i for i in labels]):
            d = defaultdict(lambda: len(d))
            for i in labels:
                if i == None:
                    y.append(-1)
                else:
                    y.append(d[i])
            y = np.array(y)
    # project the PCA space down to 2d for visualization
    z = model.fit(w, y=y if np.any(y) else None).embedding_
    return write_layout(out_path, z, **kwargs)


def get_tsne_layout(**kwargs):
    '''Get the x,y positions of images passed through a TSNE projection'''
    print(' * creating TSNE layout with ' + str(multiprocessing.cpu_count()) + ' cores...')
    out_path = get_path('layouts', 'tsne', **kwargs)
    if os.path.exists(out_path) and kwargs['use_cache']: return out_path
    model = TSNE(perplexity=kwargs.get('perplexity', 2), n_jobs=multiprocessing.cpu_count())
    z = model.fit_transform(kwargs['vecs'])
    return write_layout(out_path, z, **kwargs)


def get_rasterfairy_layout(**kwargs):
    '''Get the x, y position of images passed through a rasterfairy projection'''
    print(' * creating rasterfairy layout')
    out_path = get_path('layouts', 'rasterfairy', **kwargs)
    if os.path.exists(out_path) and kwargs['use_cache']: return out_path
    umap = np.array(read_json(kwargs['umap'], **kwargs))
    umap = (umap + 1) / 2  # scale 0:1
    try:
        umap = coonswarp.rectifyCloud(umap,  # stretch the distribution
                                      perimeterSubdivisionSteps=4,
                                      autoPerimeterOffset=False,
                                      paddingScale=1.05)
    except:
        print(' * coonswarp rectification could not be performed')
    pos = rasterfairy.transformPointCloud2D(umap)[0]
    return write_layout(out_path, pos, **kwargs)


def get_lap_layout(**kwargs):
    print(' * creating linear assignment layout')
    out_path = get_path('layouts', 'assignment', **kwargs)
    if os.path.exists(out_path) and kwargs['use_cache']: return out_path
    # load the umap layout
    umap = np.array(read_json(kwargs['umap'], **kwargs))
    umap = (umap + 1) / 2  # scale 0:1
    # determine length of each side in square grid
    side = math.ceil(umap.shape[0] ** (1 / 2))
    # create square grid 0:1 in each dimension
    grid_x, grid_y = np.meshgrid(np.linspace(0, 1, side), np.linspace(0, 1, side))
    grid = np.dstack((grid_x, grid_y)).reshape(-1, 2)
    # compute pairwise distance costs
    cost = cdist(grid, umap, 'sqeuclidean')
    # increase cost
    cost = cost * (10000000. / cost.max())
    # run the linear assignment
    min_cost, row_assignments, col_assignments = lap.lapjv(np.copy(cost), extend_cost=True)
    # use the assignment vals to determine gridified positions of `arr`
    pos = grid[col_assignments]
    return write_layout(out_path, pos, **kwargs)


def get_grid_layout(**kwargs):
    '''Get the x,y positions of images in a grid projection'''
    print(' * creating grid layout')
    out_path = get_path('layouts', 'grid', **kwargs)
    if os.path.exists(out_path) and kwargs['use_cache']: return out_path
    paths = kwargs['image_paths']
    n = math.ceil(len(paths) ** (1 / 2))
    l = []  # positions
    for i, _ in enumerate(paths):
        x = i % n
        y = math.floor(i / n)
        l.append([x, y])
    z = np.array(l)
    return write_layout(out_path, z, **kwargs)


def get_pointgrid_layout(path, label, **kwargs):
    '''Gridify the positions in `path` and return the path to this new layout'''
    print(' * creating {} pointgrid'.format(label))
    out_path = get_path('layouts', label + '-jittered', **kwargs)
    if os.path.exists(out_path) and kwargs['use_cache']: return out_path
    arr = np.array(read_json(path, **kwargs))
    z = align_points_to_grid(arr, fill=0.025)
    return write_layout(out_path, z, **kwargs)


##
# Date layout
##


def get_date_layout(cols=3, bin_units='years', **kwargs):
    '''
    Get the x,y positions of input images based on their dates
    @param int cols: the number of columns to plot for each bar
    @param str bin_units: the temporal units to use when creating bins
    '''
    date_vals = [kwargs['metadata'][i].get('year', False) for i in range(len(kwargs['metadata']))]
    if not kwargs['metadata'] or not any(date_vals): return False
    # if the data layouts have been cached, return them
    positions_out_path = get_path('layouts', 'timeline', **kwargs)
    labels_out_path = get_path('layouts', 'timeline-labels', **kwargs)
    if os.path.exists(positions_out_path) and \
            os.path.exists(labels_out_path) and \
            kwargs['use_cache']:
        return {
            'layout': positions_out_path,
            'labels': labels_out_path,
        }
    # date layout is not cached, so fetch dates and process
    print(' * creating date layout with {} columns'.format(cols))
    datestrings = [i.metadata.get('year', 'no_date') for i in stream_images(**kwargs)]
    dates = [datestring_to_date(i) for i in datestrings]
    rounded_dates = [round_date(i, bin_units) for i in dates]
    # create d[formatted_date] = [indices into datestrings of dates that round to formatted_date]
    d = defaultdict(list)
    for idx, i in enumerate(rounded_dates):
        d[i].append(idx)
    # determine the number of distinct grid positions in the x and y axes
    n_coords_x = (cols + 1) * len(d)
    n_coords_y = 1 + max([len(d[i]) for i in d]) // cols
    if n_coords_y > n_coords_x: return get_date_layout(cols=int(cols * 2), **kwargs)
    # create a mesh of grid positions in clip space -1:1 given the time distribution
    grid_x = (np.arange(0, n_coords_x) / (n_coords_x - 1)) * 2
    grid_y = (np.arange(0, n_coords_y) / (n_coords_x - 1)) * 2
    # divide each grid axis by half its max length to center at the origin 0,0
    grid_x = grid_x - np.max(grid_x) / 2.0
    grid_y = grid_y - np.max(grid_y) / 2.0
    # make dates increase from left to right by sorting keys of d
    d_keys = np.array(list(d.keys()))
    seconds = np.array([date_to_seconds(dates[d[i][0]]) for i in d_keys])
    d_keys = d_keys[np.argsort(seconds)]
    # determine which images will fill which units of the grid established above
    coords = np.zeros((len(datestrings), 2))  # 2D array with x, y clip-space coords of each date
    for jdx, j in enumerate(d_keys):
        for kdx, k in enumerate(d[j]):
            x = jdx * (cols + 1) + (kdx % cols)
            y = kdx // cols
            coords[k] = [grid_x[x], grid_y[y]]
    # find the positions of labels
    label_positions = np.array([[grid_x[i * (cols + 1)], grid_y[0]] for i in range(len(d))])
    # move the labels down in the y dimension by a grid unit
    dx = (grid_x[1] - grid_x[0])  # size of a single cell
    label_positions[:, 1] = label_positions[:, 1] - dx
    # quantize the label positions and label positions
    image_positions = round_floats(coords)
    label_positions = round_floats(label_positions.tolist())
    # write and return the paths to the date based layout
    return {
        'layout': write_json(positions_out_path, image_positions, **kwargs),
        'labels': write_json(labels_out_path, {
            'positions': label_positions,
            'labels': d_keys.tolist(),
            'cols': cols,
        }, **kwargs),
    }


def datestring_to_date(datestring):
    '''
    Given a string representing a date return a datetime object
    '''
    try:
        return parse_date(str(datestring), fuzzy=True, default=datetime.datetime(9999, 1, 1))
    except Exception as exc:
        print(' * could not parse datestring {}'.format(datestring))
        return datestring


def date_to_seconds(date):
    '''
    Given a datetime object return an integer representation for that datetime
    '''
    if isinstance(date, datetime.datetime):
        return (date - datetime.datetime.today()).total_seconds()
    else:
        return - float('inf')


def round_date(date, unit):
    '''
    Return `date` truncated to the temporal unit specified in `units`
    '''
    if not isinstance(date, datetime.datetime): return 'no_date'
    formatted = date.strftime('%d %B %Y -- %X')
    if unit in set(['seconds', 'minutes', 'hours']):
        date = formatted.split('--')[1].strip()
        if unit == 'seconds':
            date = date
        elif unit == 'minutes':
            date = ':'.join(d.split(':')[:-1]) + ':00'
        elif unit == 'hours':
            date = date.split(':')[0] + ':00:00'
    elif unit in set(['days', 'months', 'years', 'decades', 'centuries']):
        date = formatted.split('--')[0].strip()
        if unit == 'days':
            date = date
        elif unit == 'months':
            date = ' '.join(date.split()[1:])
        elif unit == 'years':
            date = date.split()[-1]
        elif unit == 'decades':
            date = str(int(date.split()[-1]) // 10) + '0'
        elif unit == 'centuries':
            date = str(int(date.split()[-1]) // 100) + '00'
    return date


##
# Metadata layout
##


def get_categorical_layout(null_category='Other', margin=2, **kwargs):
    '''
    Return a numpy array with shape (n_points, 2) with the point
    positions of observations in box regions determined by
    each point's category metadata attribute (if applicable)
    '''
    if not kwargs.get('metadata', False): return False
    # determine the out path and return from cache if possible
    out_path = get_path('layouts', 'categorical', **kwargs)
    labels_out_path = get_path('layouts', 'categorical-labels', **kwargs)
    if os.path.exists(out_path): return out_path
    # accumulate d[category] = [indices of points with category]
    categories = [i.get('category', None) for i in kwargs['metadata']]
    if not any(categories): return False
    d = defaultdict(list)
    for idx, i in enumerate(categories): d[i].append(idx)
    # store the number of observations in each group
    keys_and_counts = [{'key': i, 'count': len(d[i])} for i in d]
    keys_and_counts.sort(key=operator.itemgetter('count'), reverse=True)
    # get the box layout then subdivide into discrete points
    boxes = get_categorical_boxes([i['count'] for i in keys_and_counts], margin=margin)
    points = get_categorical_points(boxes)
    # sort the points into the order of the observations in the metadata
    counts = {i['key']: 0 for i in keys_and_counts}
    offsets = {i['key']: 0 for i in keys_and_counts}
    for idx, i in enumerate(keys_and_counts):
        offsets[i['key']] += sum([j['count'] for j in keys_and_counts[:idx]])
    sorted_points = []
    for idx, i in enumerate(stream_images(**kwargs)):
        category = i.metadata.get('category', null_category)
        sorted_points.append(points[offsets[category] + counts[category]])
        counts[category] += 1
    sorted_points = np.array(sorted_points)
    # add to the sorted points the anchors for the text labels for each group
    text_anchors = np.array([[i.x, i.y - margin / 2] for i in boxes])
    # add the anchors to the points - these will be removed after the points are projected
    sorted_points = np.vstack([sorted_points, text_anchors])
    # scale -1:1 using the largest axis as the scaling metric
    _max = np.max(sorted_points)
    for i in range(2):
        _min = np.min(sorted_points[:, i])
        sorted_points[:, i] -= _min
        sorted_points[:, i] /= (_max - _min)
        sorted_points[:, i] -= np.max(sorted_points[:, i]) / 2
        sorted_points[:, i] *= 2
    # separate out the sorted points and text positions
    text_anchors = sorted_points[-len(text_anchors):]
    sorted_points = sorted_points[:-len(text_anchors)]
    z = round_floats(sorted_points.tolist())
    return {
        'layout': write_json(out_path, z, **kwargs),
        'labels': write_json(labels_out_path, {
            'positions': round_floats(text_anchors.tolist()),
            'labels': [i['key'] for i in keys_and_counts],
        }, **kwargs)
    }


def get_categorical_boxes(group_counts, margin=2):
    '''
    @arg [int] group_counts: counts of the number of images in each
      distinct level within the metadata's caetgories
    @kwarg int margin: space between boxes in the 2D layout
    @returns [Box] an array of Box() objects; one per level in `group_counts`
    '''
    boxes = []
    for i in group_counts:
        x = y = math.ceil(i ** (1 / 2))
        boxes.append(Box(i, x, y, None, None))
    # find the position along x axis where we want to create a break
    wrap = sum([i.cells for i in boxes]) ** (1 / 2)
    # find the valid positions on the y axis
    y = margin
    y_spots = []
    for idx, i in enumerate(boxes):
        if (y + i.h) < wrap:
            y_spots.append(y)
            y += i.h + margin
    y_spot_index = 0
    for idx, i in enumerate(boxes):
        # find the y position
        y = y_spots[y_spot_index]
        # find members with this y position
        row_members = [j.x + j.w for j in boxes if j.y == y]
        # assign the y position
        i.y = y
        y_spot_index = (y_spot_index + 1) % len(y_spots)
        # assign the x position
        i.x = max(row_members) + margin if row_members else margin
    return boxes


def get_categorical_points(arr, unit_size=None):
    '''Given an array of Box() objects, return a 2D distribution with shape (n_cells, 2)'''
    points_arr = []
    for i in arr:
        area = i.w * i.h
        per_unit = (area / i.cells) ** (1 / 2)
        x_units = math.ceil(i.w / per_unit)
        y_units = math.ceil(i.h / per_unit)
        if not unit_size: unit_size = min(i.w / x_units, i.h / y_units)
        for j in range(i.cells):
            x = j % x_units
            y = j // x_units
            points_arr.append([
                i.x + x * unit_size,
                i.y + y * unit_size,
            ])
    return np.array(points_arr)


class Box:
    '''Store the width, height, and x, y coords of a box'''

    def __init__(self, *args):
        self.cells = args[0]
        self.w = args[1]
        self.h = args[2]
        self.x = None if len(args) < 4 else args[3]
        self.y = None if len(args) < 5 else args[4]


##
# Helpers
##


def get_path(*args, **kwargs):
    '''Return the path to a JSON file with conditional gz extension'''
    sub_dir, filename = args
    out_dir = join(kwargs['out_dir'], sub_dir) if sub_dir else kwargs['out_dir']
    if kwargs.get('add_hash', True):
        filename += '-' + kwargs['plot_id']
    path = join(out_dir, filename + '.json')
    return path + '.gz' if kwargs.get('gzip', False) else path


def write_layout(path, obj, **kwargs):
    '''Write layout json `obj` to disk and return the path to the saved file'''
    obj = (minmax_scale(obj) - 0.5) * 2  # scale -1:1
    obj = round_floats(obj)
    return write_json(path, obj, **kwargs)


def round_floats(obj, digits=5):
    '''Return 2D array obj with rounded float precision'''
    return [[round(float(j), digits) for j in i] for i in obj]


def write_json(path, obj, **kwargs):
    '''Write json object `obj` to disk and return the path to that file'''
    out_dir, filename = os.path.split(path)
    if not os.path.exists(out_dir): os.makedirs(out_dir)
    if kwargs.get('gzip', False):
        with gzip.GzipFile(path, 'w') as out:
            out.write(json.dumps(obj, indent=4).encode(kwargs['encoding']))
        return path
    else:
        with open(path, 'w') as out:
            json.dump(obj, out, indent=4)
        return path


def read_json(path, **kwargs):
    '''Read and return the json object written by the current process at `path`'''
    if kwargs.get('gzip', False):
        with gzip.GzipFile(path, 'r') as f:
            return json.loads(f.read().decode(kwargs['encoding']))
    with open(path) as f:
        return json.load(f)


def get_hotspots(**kwargs):
    '''Return the stable clusters from the condensed tree of connected components from the density graph'''
    print(' * HDBSCAN clustering data with ' + str(multiprocessing.cpu_count()) + ' cores...')
    config = {
        'core_dist_n_jobs': multiprocessing.cpu_count(),
        'min_cluster_size': kwargs['min_cluster_size'],
        'cluster_selection_epsilon': 0.01,
        'min_samples': 1,
        'approx_min_span_tree': False,
    }
    v = kwargs['vecs']
    z = HDBSCAN(**config).fit(v)
    # find the points in each cluster
    d = defaultdict(list)
    for idx, i in enumerate(z.labels_):
        d[i].append(v[idx])
    # find the convex hull for each cluster's points
    convex_hulls = []
    for i in d:
        hull = ConvexHull(d[i])
        points = [hull.points[j] for j in hull.vertices]
        # the last convex hull simplex needs to connect back to the first point
        convex_hulls.append(np.vstack([points, points[0]]))
    # find the centroids for each cluster
    centroids = []
    for i in d:
        x, y = np.array(d[i]).T
        centroids.append(np.array([np.mean(x), np.mean(y)]))
    # identify the number of points in each cluster
    lens = [len(d[i]) for i in d]
    # combine data into cluster objects
    closest, _ = pairwise_distances_argmin_min(centroids, v)
    paths = [kwargs['image_paths'][i] for i in closest]
    clusters = [{
        'img': clean_filename(paths[idx]),
        'convex_hull': convex_hulls[idx].tolist(),
        'n_images': lens[idx],
    } for idx, i in enumerate(closest)]
    # remove massive clusters
    retained = []
    for idx, i in enumerate(clusters):
        x, y = np.array(i['convex_hull']).T
        area = 0.5 * np.abs(np.dot(x, np.roll(y, 1)) - np.dot(y, np.roll(x, 1)))
        if area < 0.2:
            retained.append(i)
    # sort the clusers by size
    clusters = sorted(retained, key=lambda i: i['n_images'], reverse=True)
    for idx, i in enumerate(clusters):
        i['label'] = 'Cluster {}'.format(idx + 1)
    # save the hotspots to disk and return the path to the saved json
    print(' * found', len(clusters), 'hotspots')
    return write_json(get_path('hotspots', 'hotspot', **kwargs), clusters, **kwargs)


def get_heightmap(path, label, **kwargs):
    '''Create a heightmap using the distribution of points stored at `path`'''
    X = read_json(path, **kwargs)
    if 'positions' in X: X = X['positions']
    X = np.array(X)
    # create kernel density estimate of distribution X
    nbins = 200
    x, y = X.T
    xi, yi = np.mgrid[x.min():x.max():nbins * 1j, y.min():y.max():nbins * 1j]
    zi = kde.gaussian_kde(X.T)(np.vstack([xi.flatten(), yi.flatten()]))
    # create the plot
    fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(5, 5))
    fig.subplots_adjust(0, 0, 1, 1)
    plt.pcolormesh(xi, yi, zi.reshape(xi.shape), shading='gouraud', cmap=plt.cm.gray)
    plt.axis('off')
    # save the plot
    out_dir = os.path.join(kwargs['out_dir'], 'heightmaps')
    if not os.path.exists(out_dir): os.makedirs(out_dir)
    out_path = os.path.join(out_dir, label + '-heightmap.png')
    plt.savefig(out_path, pad_inches=0)


def write_images(**kwargs):
    '''Write all originals and thumbs to the output dir'''
    for i in stream_images(**kwargs):
        filename = clean_filename(i.path)
        # copy original for lightbox
        out_dir = join(kwargs['out_dir'], 'originals')
        if not exists(out_dir): os.makedirs(out_dir)
        out_path = join(out_dir, filename)
        shutil.copy(i.path, out_path)
        # copy thumb for lod texture
        out_dir = join(kwargs['out_dir'], 'thumbs')
        if not exists(out_dir): os.makedirs(out_dir)
        out_path = join(out_dir, filename)
        img = array_to_img(i.resize_to_max(kwargs['lod_cell_height']))
        save_img(out_path, img)


def get_version():
    '''Return the version of pixplot installed'''
    return pkg_resources.get_distribution('pixplot').version


class Image:
    def __init__(self, *args, **kwargs):
        self.path = args[0]
        self.original = load_img(self.path)
        self.metadata = kwargs['metadata']

    def resize_to_max(self, n):
        '''
        Resize self.original so its longest side has n pixels (maintain proportion)
        '''
        w, h = self.original.size
        size = (n, int(n * h / w)) if w > h else (int(n * w / h), n)
        return img_to_array(self.original.resize(size))

    def resize_to_height(self, height):
        '''
        Resize self.original into an image with height h and proportional width
        '''
        w, h = self.original.size
        if (w / h * height) < 1:
            resizedwidth = 1
        else:
            resizedwidth = int(w / h * height)
        size = (resizedwidth, height)
        return img_to_array(self.original.resize(size))

    def resize_to_square(self, n, center=False):
        '''
        Resize self.original to an image with nxn pixels (maintain proportion)
        if center, center the colored pixels in the square, else left align
        '''
        a = self.resize_to_max(n)
        h, w, c = a.shape
        pad_lr = int((n - w) / 2)  # left right pad
        pad_tb = int((n - h) / 2)  # top bottom pad
        b = np.zeros((n, n, 3))
        if center:
            b[pad_tb:pad_tb + h, pad_lr:pad_lr + w, :] = a
        else:
            b[:h, :w, :] = a
        return b


##
# Entry Point
##


def parse():
    '''Read command line args and begin data processing'''
    description = 'Generate the data required to create a PixPlot viewer'
    parser = argparse.ArgumentParser(description=description, formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--images', type=str, default=config['images'], help='path to a glob of images to process',
                        required=False)
    parser.add_argument('--metadata', type=str, default=config['metadata'],
                        help='path to a csv or glob of JSON files with image metadata (see readme for format)',
                        required=False)
    parser.add_argument('--max_images', type=int, default=config['max_images'],
                        help='maximum number of images to process from the input glob', required=False)
    parser.add_argument('--use_cache', type=bool, default=config['use_cache'],
                        help='given inputs identical to prior inputs, load outputs from cache', required=False)
    parser.add_argument('--encoding', type=str, default=config['encoding'], help='the encoding of input metadata',
                        required=False)
    parser.add_argument('--min_cluster_size', type=int, default=config['min_cluster_size'],
                        help='the minimum number of images in a cluster', required=False)
    parser.add_argument('--out_dir', type=str, default=config['out_dir'],
                        help='the directory to which outputs will be saved', required=False)
    parser.add_argument('--cell_size', type=int, default=config['cell_size'], help='the size of atlas cells in px',
                        required=False)
    parser.add_argument('--n_neighbors', type=int, default=config['n_neighbors'],
                        help='the n_neighbors argument for UMAP')
    parser.add_argument('--min_distance', type=float, default=config['min_distance'],
                        help='the min_distance argument for umap')
    parser.add_argument('--metric', type=str, default=config['metric'], help='the metric argument for umap')
    parser.add_argument('--pointgrid_fill', type=float, default=config['pointgrid_fill'],
                        help='float 0:1 that determines sparsity of jittered distributions (lower means more sparse)')
    parser.add_argument('--copy_web_only', action='store_true', help='update ./output/web without reprocessing data')
    parser.add_argument('--gzip', action='store_true', help='save outputs with gzip compression')
    parser.add_argument('--shuffle', action='store_true', help='shuffle the input images before data processing begins')
    parser.add_argument('--plot_id', type=str, default=config['plot_id'],
                        help='unique id for a plot; useful for resuming processing on a started plot')
    parser.add_argument('--seed', type=int, default=config['seed'], help='seed for random processes')
    config.update(vars(parser.parse_args()))
    process_images(**config)


if __name__ == '__main__':
    parse()
# python pixplot_notebook.py --images "images/*.jpg"
