# The MIT License (MIT)
#
# Copyright (c) 2015 University of East Anglia, Norwich, UK
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#
# Developed by Geoffrey French in collaboration with Dr. M. Fisher and
# Dr. M. Mackiewicz.
import click


def flask_labeller(label_classes, labelled_images, tasks=None, colour_schemes=None, anno_controls=None,
                   config=None, dextr_fn=None, use_reloader=True, debug=True, port=None):
    import json
    import uuid
    import numpy as np

    from flask import Flask, render_template, request, make_response, send_file
    try:
        from flask_socketio import SocketIO, emit as socketio_emit
    except ImportError:
        SocketIO = None
        socketio_emit = None

    from image_labelling_tool import labelling_tool

    # Generate image IDs list
    image_ids = [str(i)   for i in range(len(labelled_images))]
    # Generate images table mapping image ID to image so we can get an image by ID
    images_table = {image_id: img   for image_id, img in zip(image_ids, labelled_images)}
    # Generate image descriptors list to hand over to the labelling tool
    # Each descriptor provides the image ID, the URL and the size
    image_descriptors = []
    for image_id, img in zip(image_ids, labelled_images):
        height, width = img.image_source.image_size
        image_descriptors.append(labelling_tool.image_descriptor(
            image_id=image_id, url='/image/{}'.format(image_id),
            width=width, height=height
        ))


    app = Flask(__name__, static_folder='static')
    if SocketIO is not None:
        print('Using web sockets')
        socketio = SocketIO(app)
    else:
        socketio = None


    def apply_dextr_js(image, dextr_points_js):
        image_for_dextr = image.image_source.image_as_array_or_pil()
        dextr_points = np.array([[p['y'], p['x']] for p in dextr_points_js])
        if dextr_fn is not None:
            mask = dextr_fn(image_for_dextr, dextr_points)
            regions = labelling_tool.PolygonLabel.mask_image_to_regions_cv(mask, sort_decreasing_area=True)
            regions_js = labelling_tool.PolygonLabel.regions_to_json(regions)
            return regions_js
        else:
            return []




    if config is None:
        config = labelling_tool.DEFAULT_CONFIG


    @app.route('/')
    def index():
        label_classes_json = [(cls.to_json() if isinstance(cls, labelling_tool.LabelClassGroup) else cls)
                               for cls in label_classes]
        if anno_controls is not None:
            anno_controls_json = [c.to_json() for c in anno_controls]
        else:
            anno_controls_json = []
        return render_template('labeller_page.jinja2',
                               tasks=tasks,
                               colour_schemes=colour_schemes,
                               label_class_groups=label_classes_json,
                               image_descriptors=image_descriptors,
                               initial_image_index=0,
                               anno_controls=anno_controls_json,
                               labelling_tool_config=config,
                               dextr_available=dextr_fn is not None,
                               use_websockets=socketio is not None)


    if socketio is not None:
        @socketio.on('get_labels')
        def handle_get_labels(arg_js):
            image_id = arg_js['image_id']

            image = images_table[image_id]

            wrapped_labels = image.labels_store.get_wrapped_labels()

            label_header = dict(image_id=image_id,
                                labels=wrapped_labels.labels_json,
                                completed_tasks=wrapped_labels.completed_tasks,
                                timeElapsed=0.0,
                                state='editable',
                                session_id=str(uuid.uuid4()),
            )

            socketio_emit('get_labels_reply', label_header)


        @socketio.on('set_labels')
        def handle_set_labels(arg_js):
            label_header = arg_js['label_header']

            image_id = label_header['image_id']

            image = images_table[image_id]

            wrapped_labels = image.labels_store.get_wrapped_labels()
            wrapped_labels.labels_json = label_header['labels']
            wrapped_labels.completed_tasks = label_header['completed_tasks']
            image.labels_store.update_wrapped_labels(wrapped_labels)

            socketio_emit('set_labels_reply', '')


        @socketio.on('dextr')
        def handle_dextr(dextr_js):
            if 'request' in dextr_js:
                dextr_request_js = dextr_js['request']
                image_id = dextr_request_js['image_id']
                dextr_id = dextr_request_js['dextr_id']
                dextr_points = dextr_request_js['dextr_points']

                image = images_table[image_id]

                regions_js = apply_dextr_js(image, dextr_points)

                dextr_labels = dict(image_id=image_id, dextr_id=dextr_id, regions=regions_js)
                dextr_reply = dict(labels=[dextr_labels])

                socketio_emit('dextr_reply', dextr_reply)
            elif 'poll' in dextr_js:
                dextr_reply = dict(labels=[])
                socketio_emit('dextr_reply', dextr_reply)
            else:
                dextr_reply = {'error': 'unknown_command'}
                socketio_emit('dextr_reply', dextr_reply)


    else:
        @app.route('/labelling/get_labels/<image_id>')
        def get_labels(image_id):
            image = images_table[image_id]
            wrapped_labels = image.labels_store.get_wrapped_labels()

            label_header = {
                'image_id': image_id,
                'labels': wrapped_labels.labels_json,
                'completed_tasks': wrapped_labels.completed_tasks,
                'timeElapsed': 0.0,
                'state': 'editable',
                'session_id': str(uuid.uuid4()),
            }

            r = make_response(json.dumps(label_header))
            r.mimetype = 'application/json'
            return r


        @app.route('/labelling/set_labels', methods=['POST'])
        def set_labels():
            label_header = json.loads(request.form['labels'])
            image_id = label_header['image_id']

            image = images_table[image_id]

            wrapped_labels = image.labels_store.get_wrapped_labels()
            wrapped_labels.labels_json = label_header['labels']
            wrapped_labels.completed_tasks = label_header['completed_tasks']
            image.labels_store.update_wrapped_labels(wrapped_labels)

            return make_response('')


        @app.route('/labelling/dextr', methods=['POST'])
        def dextr():
            dextr_js = json.loads(request.form['dextr'])
            if 'request' in dextr_js:
                dextr_request_js = dextr_js['request']
                image_id = dextr_request_js['image_id']
                dextr_id = dextr_request_js['dextr_id']
                dextr_points = dextr_request_js['dextr_points']

                image = images_table[image_id]
                regions_js = apply_dextr_js(image, dextr_points)

                dextr_labels = dict(image_id=image_id, dextr_id=dextr_id, regions=regions_js)
                dextr_reply = dict(labels=[dextr_labels])

                return make_response(json.dumps(dextr_reply))
            elif 'poll' in dextr_js:
                dextr_reply = dict(labels=[])
                return make_response(json.dumps(dextr_reply))
            else:
                return make_response(json.dumps({'error': 'unknown_command'}))


    @app.route('/image/<image_id>')
    def get_image(image_id):
        image_source = images_table[image_id].image_source
        local_path = image_source.local_path
        if local_path is not None:
            return send_file(str(local_path))
        else:
            bin_image, mimetype = image_source.image_binary_and_mime_type()
            r = make_response(bin_image)
            r.mimetype = mimetype
            return r



    if socketio is not None:
        socketio.run(app, debug=debug, port=port, use_reloader=use_reloader)
    else:
        app.run(debug=debug, port=port, use_reloader=use_reloader)



@click.command()
@click.option('--images_dir', type=click.Path(dir_okay=True, file_okay=False, exists=True), default='./images')
@click.option('--images_pat', type=str, default='*.png|*.jpg')
@click.option('--labels_dir', type=click.Path(dir_okay=True, file_okay=False, writable=True))
@click.option('--readonly', is_flag=True, default=False, help='Don\'t persist changes to disk')
@click.option('--update_label_object_ids', is_flag=True, default=False, help='Update object IDs in label JSON files')
@click.option('--enable_dextr', is_flag=True, default=False)
@click.option('--dextr_weights', type=click.Path())
def run_app(images_dir, images_pat, labels_dir, readonly, update_label_object_ids,
            enable_dextr, dextr_weights):
    import pathlib
    import json
    import uuid
    from image_labelling_tool import labelled_image, labelling_tool

    if enable_dextr or dextr_weights is not None:
        from dextr.model import DextrModel
        import torch

        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

        if dextr_weights is not None:
            dextr_weights = pathlib.Path(dextr_weights).expanduser()
            dextr_model = torch.load(dextr_weights, map_location=device)
        else:
            dextr_model = DextrModel.pascalvoc_resunet101().to(device)

        dextr_model.eval()

        dextr_fn = lambda image, points: dextr_model.predict([image], points[None, :, :])[0] >= 0.5
    else:
        dextr_fn = None


    # Colour schemes
    # The user may select different colour schemes for different tasks.
    # If you have a lot of classes, it will be difficult to select colours that are easily distinguished
    # from one another. For one task e.g. segmentation, design a colour scheme that highlights the different
    # classes for that task, while another task e.g. fine-grained classification would use another scheme.
    # Each colour scheme is a dictionary containing the following:
    #   name: symbolic name (Python identifier)
    #   human_name: human readable name for UI
    # These colour schemes are going to split the classes by 'default' (all), natural, and artificial.
    # Not really useful, but demonstrates the feature.
    colour_schemes = [
        dict(name='default', human_name='All'),
        dict(name='natural', human_name='Natural'),
        dict(name='artificial', human_name='Artifical')
    ]

    # Specify our label classes, organised in groups.
    # `LabelClass` parameters are:
    #   symbolic name (Python identifier)
    #   human readable name for UI
    #   and colours by colour scheme, as a dict mapping colour scheme name to RGB value as a list
    # The label classes are arranged in groups and will be displayed as such in the UI.
    # `LabelClassGroup` parameters are:
    #   human readable name for UI
    #   label class (`LabelClass` instance) list
    label_classes = [
        labelling_tool.LabelClassGroup('Natural', [
            labelling_tool.LabelClass('tree', 'Trees', dict(default=[0, 255, 192], natural=[0, 255, 192],
                                                            artificial=[128, 128, 128])),
            labelling_tool.LabelClass('lake', 'Lake', dict(default=[0, 128, 255], natural=[0, 128, 255],
                                                           artificial=[128, 128, 128])),
            labelling_tool.LabelClass('flower', 'Flower', dict(default=[255, 96, 192], natural=[255, 192, 96],
                                                               artificial=[128, 128, 128])),
            labelling_tool.LabelClass('leaf', 'Leaf', dict(default=[65, 255, 0], natural=[65, 255, 0],
                                                           artificial=[128, 128, 128])),
            labelling_tool.LabelClass('stem', 'Stem', dict(default=[128, 64, 0], natural=[128, 64, 0],
                                                           artificial=[128, 128, 128])),
        ]),
        labelling_tool.LabelClassGroup('Artificial', [
            labelling_tool.LabelClass('building', 'Buildings', dict(default=[255, 128, 0], natural=[128, 128, 128],
                                                                   artificial=[255, 128, 0])),
            labelling_tool.LabelClass('wall', 'Wall', dict(default=[0, 128, 255], natural=[128, 128, 128],
                                                           artificial=[0, 128, 255])),
        ])]

    # Annotation controls
    # Labels may also have optional meta-data associated with them
    # You could use this for e.g. indicating if an object is fully visible, mostly visible or significantly obscured.
    # You could also indicate quality (e.g. blurriness, etc)
    # There are four types of annotation. They have some common properties:
    #   - name: symbolic name (Python identifier)
    #   - label_text: label text in UI
    #   Check boxes, radio buttons and popup menus also have:
    #     - visibility_label_text: [optional] if provided, label visibility can be filtered by this annotation value,
    #       in which case a drop down will appear in the UI allowing the user to select a filter value
    #       that will hide/show labels accordingly
    # Control types:
    # Check box (boolean value):
    #   `labelling_tool.AnnoControlCheckbox`; only the 3 common parameters listed above
    # Radio button (choice from a list):
    #   `labelling_tool.AnnoControlRadioButtons`; the 3 common parameters listed above and:
    #       - choices: list of `labelling_tool.AnnoControlRadioButtons.choice` that provide:
    #           - value: symbolic value name for choice
    #           - tooltip: extra information for user
    #       - label_on_own_line [optional]: if True, place the label and the buttons on a separate line in the UI
    # Popup menu (choice from a grouped list):
    #   `labelling_tool.AnnoControlPopupMenu`; the 3 common parameters listed above and:
    #       - groups: list of groups `labelling_tool.AnnoControlPopupMenu.group`:
    #           - label_text: label text in UI
    #           - choices: list of `labelling_tool.AnnoControlPopupMenu.choice` that provide:
    #               - value: symbolic value name for choice
    #               - label_text: choice label text in UI
    #               - tooltip: extra information for user
    # Text (free form plain text):
    #   `labelling_tool.AnnoControlText`; only the 2 common parameters listed above and:
    #       - multiline: boolean; if True a text area will be used, if False a single line text entry
    anno_controls = [
        labelling_tool.AnnoControlCheckbox('good_quality', 'Good quality',
                                           visibility_label_text='Filter by good quality'),
        labelling_tool.AnnoControlRadioButtons('visibility', 'Visible', choices=[
            labelling_tool.AnnoControlRadioButtons.choice(value='full', label_text='Fully',
                                                          tooltip='Object is fully visible'),
            labelling_tool.AnnoControlRadioButtons.choice(value='mostly', label_text='Mostly',
                                                          tooltip='Object is mostly visible'),
            labelling_tool.AnnoControlRadioButtons.choice(value='obscured', label_text='Obscured',
                                                          tooltip='Object is significantly obscured'),
        ], label_on_own_line=False, visibility_label_text='Filter by visibility'),
        labelling_tool.AnnoControlPopupMenu('material', 'Material', groups=[
            labelling_tool.AnnoControlPopupMenu.group(label_text='Artifical/buildings', choices=[
                labelling_tool.AnnoControlPopupMenu.choice(value='concrete', label_text='Concrete',
                                                           tooltip='Concrete objects'),
                labelling_tool.AnnoControlPopupMenu.choice(value='plastic', label_text='Plastic',
                                                           tooltip='Plastic objects'),
                labelling_tool.AnnoControlPopupMenu.choice(value='asphalt', label_text='Asphalt',
                                                           tooltip='Road, pavement, etc.'),
            ]),
            labelling_tool.AnnoControlPopupMenu.group(label_text='Flat natural', choices=[
                labelling_tool.AnnoControlPopupMenu.choice(value='grass', label_text='Grass',
                                                           tooltip='Grass covered ground'),
                labelling_tool.AnnoControlPopupMenu.choice(value='water', label_text='Water', tooltip='Water/lake')]),
            labelling_tool.AnnoControlPopupMenu.group(label_text='Vegetation', choices=[
                labelling_tool.AnnoControlPopupMenu.choice(value='trees', label_text='Trees', tooltip='Trees'),
                labelling_tool.AnnoControlPopupMenu.choice(value='shrubbery', label_text='Shrubs',
                                                           tooltip='Shrubs/bushes'),
                labelling_tool.AnnoControlPopupMenu.choice(value='flowers', label_text='Flowers',
                                                           tooltip='Flowers'),
                labelling_tool.AnnoControlPopupMenu.choice(value='ivy', label_text='Ivy', tooltip='Ivy')]),
        ], visibility_label_text='Filter by material'),
        # labelling_tool.AnnoControlText('comment', 'Comment', multiline=False),
    ]

    image_pats = images_pat.split('|')

    # Load in .JPG images from the 'images' directory.
    labelled_images = labelled_image.LabelledImage.for_directory(
        images_dir, image_filename_patterns=image_pats, readonly=readonly)
    print('Loaded {0} images'.format(len(labelled_images)))

    if update_label_object_ids:
        n_updated = 0
        for limg in labelled_images:
            if limg.labels_store.labels_path.exists():
                label_js = json.load(limg.labels_store.labels_path.open('r'))
                prefix = str(uuid.uuid4())
                modified = labelling_tool.ensure_json_object_ids_have_prefix(
                    label_js, id_prefix=prefix)
                if modified:
                    with open(limg.labels_store.labels_path, 'w') as f_out:
                        json.dump(label_js, f_out, indent=3)
                    n_updated += 1
        print('Updated object IDs in {} files'.format(n_updated))



    # For documentation of the configuration, please see the comment above `labelling_tool.DEFAULT_CONFIG`
    config = labelling_tool.DEFAULT_CONFIG

    tasks = [
        dict(name='finished', human_name='[old] finished'),
        dict(name='segmentation', human_name='Outlines'),
        dict(name='classification', human_name='Classification'),
    ]

    flask_labeller(label_classes, labelled_images, tasks=tasks, colour_schemes=colour_schemes,
                   anno_controls=anno_controls, config=config, dextr_fn=dextr_fn)


if __name__ == '__main__':
    run_app()