#!/usr/bin/env python3

import base64
import json
import os
import io
import datetime
import time
import cv2
import boto3
import botocore.response

threshold = 0.10

object_categories = [
    'tug',
    'bridge',
    'aircraft',
    'beltloader',
    'fueltruck',
    'towbar',
    'towtug',
    'truck',
    'liftloader',
    'gpu'
]


def video_to_frames(video, path_output_dir):
    capture = cv2.VideoCapture(video)
    frames = 0
    while True:
        success, image = capture.read()
        if success:
            cv2.imwrite(os.path.join(path_output_dir,
                                     'frame-%04d.jpg') % frames, image)
            frames += 1
        else:
            break
    capture.release()


def read_file(file):
    with open(file, "rb") as input:
        return input.read()


def save(filepath, fig=None):
    '''Save the current image with no whitespace
    Example filepath: "myfig.png" or r"C:\myfig.pdf"
    '''
    import matplotlib.pyplot as plt
    if not fig:
        fig = plt.gcf()

    plt.subplots_adjust(0,0,1,1,0,0)
    for ax in fig.axes:
        ax.axis('off')
        ax.margins(0,0)
        ax.xaxis.set_major_locator(plt.NullLocator())
        ax.yaxis.set_major_locator(plt.NullLocator())

    basedir  = os.path.dirname(filepath)
    basefile = os.path.basename(filepath)
    basename = os.path.splitext(basefile)[0]
    baseext = os.path.splitext(basefile)[1]

    fig.savefig(os.path.join(basedir, f"{basename}-detections{baseext}"), pad_inches = 0, bbox_inches='tight')

def visualize_detection(img_file, dets, classes=[], thresh=0.6):
    import random
    import matplotlib.pyplot as plt
    import matplotlib.image as mpimg
    img = mpimg.imread(img_file)
    IMAGE_SIZE = (24, 16)
    plt.figure(figsize=IMAGE_SIZE, frameon=False)
    plt.imshow(img)
    height = img.shape[0]
    width = img.shape[1]
    colors = dict()
    for det in dets:
        (klass, score, x0, y0, x1, y1) = det
        if score < thresh:
            continue
        cls_id = int(klass)
        if cls_id not in colors:
            colors[cls_id] = (0, 1, 0)
        xmin = int(x0 * width)
        ymin = int(y0 * height)
        xmax = int(x1 * width)
        ymax = int(y1 * height)
        rect = plt.Rectangle((xmin, ymin), xmax - xmin,
                             ymax - ymin, fill=False,
                             edgecolor=colors[cls_id],
                             linewidth=1.5)
        plt.gca().add_patch(rect)
        class_name = str(cls_id)
        if classes and len(classes) > cls_id:
            class_name = classes[cls_id]
        plt.gca().text(xmin, ymin - 2,
                       '{:s} {:.3f}'.format(class_name, score),
                       bbox=dict(facecolor=colors[cls_id], alpha=0.5),
                       fontsize=12, color='white')
    save(img_file, plt.gcf())
    plt.close('all')

def main():
    # Video to frames
    # source = 'resources/target.mp4'
    # destination = 'out'
    # os.makedirs(destination, exist_ok=True)
    # video_to_frames(source, destination)

    # Frames to predictions
    sagemaker = boto3.client('sagemaker-runtime')

    for subdir, dirs, files in os.walk('out'):
        for file in sorted(files):
            if file.startswith(("frame")) and file.endswith((".jpg")):
                frame = os.path.join(subdir, file)

                print(frame)

                response = sagemaker.invoke_endpoint(
                    EndpointName='f360-cavu-prod',
                    Body=read_file(frame),
                    ContentType='application/x-image',
                    Accept='application/json',
                )
                body = response["Body"]
                body = body.read().decode('utf-8')
                body = json.loads(body)
                visualize_detection(frame, body['prediction'], object_categories, 0.10)


if __name__ == '__main__':
    main()
    exit(0)