"""Executable Python script for running Python actions.

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
"""

'''
Some is based on Ildoo Kim's code (https://github.com/ildoonet/tf-openpose) and https://gist.github.com/alesolano/b073d8ec9603246f766f9f15d002f4f4
and derived from the OpenPose Library (https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/master/LICENSE)
'''

import tensorflow as tf
import numpy as np
from PIL import Image
from tensorflow.core.framework import graph_pb2
import urllib3
import certifi
import os
import shutil

from common import estimate_pose, crop_image, draw_humans

import time


def print_time(message, start):
    print(message, "{:10.4f}".format(time.time() - start))
    return time.time()


class SmartBodyCrop:
    initialized = False
    tmp_path = '/tmp/'
    tmpfs_path = '/mnt/action/'

    def __init__(self, model_url):
        self.model_url = model_url

    def read_img(self, imgpath, width, height):
        img = Image.open(imgpath)
        orig_width, orig_height = img.size
        # resize the image to match openpose's training data
        # https://github.com/ildoonet/tf-pose-estimation#inference-time
        img.thumbnail((width, height))
        thumbnail_w, thumbnail_h = img.size
        #val_img = val_img.resize((width, height))
        val_img = np.asarray(img, dtype=np.float32)
        val_img = val_img.reshape([1, thumbnail_h, thumbnail_w, 3])
        # val_img = val_img.astype(float)
        val_img = val_img * (2.0 / 255.0) - 1.0  # normalization

        return val_img, img, orig_width, orig_height

    def _download_model(self):
        # check if the model is a ref to local file path
        if type(self.model_url) is str:
            if not self.model_url.startswith('http'):
                return self.model_url

        start = time.time()
        local_model_path = SmartBodyCrop.tmp_path + 'optimized_openpose.pb'
        tmpfs_model_path = SmartBodyCrop.tmpfs_path + 'optimized_openpose.pb'

        if (os.path.isfile(local_model_path)):
            print_time("model was found in the local storage: " +
                       local_model_path, start)
            return local_model_path

        # check if this model was downloaded by another invocation in the tmpfs path
        if (os.path.isfile(tmpfs_model_path)):
            print_time("model was found in the tmpfs storage: " +
                       tmpfs_model_path, start)
            shutil.copy(tmpfs_model_path, local_model_path)
            print_time("model copied FROM tmpfs:" + tmpfs_model_path, start)
            return local_model_path

        http = urllib3.PoolManager(
            cert_reqs='CERT_REQUIRED',
            ca_certs=certifi.where(),
            headers={
                'Accept': 'application/octet-stream',
                'Content-Type': 'application/octet-stream'
            })
        urllib3.disable_warnings()

        r = http.request('GET', self.model_url,
                         preload_content=False,
                         retries=urllib3.Retry(5, redirect=5))

        with open(local_model_path, 'wb') as out:
            while True:
                data = r.read(8192)  # 64 # 8192
                if not data:
                    break
                out.write(data)

        r.release_conn()
        print_time("model downloaded in :", start)

        # copy the file to the tmpfs_model_path to be reused by other actions
        # this seems to work concurrently as per: https://stackoverflow.com/questions/35605463/why-is-concurrent-copy-of-a-file-not-failing
        if (os.path.isdir(SmartBodyCrop.tmpfs_path)):
            shutil.copy(local_model_path, tmpfs_model_path)
            print_time("model copied to tmpfs:" + tmpfs_model_path, start)

        return local_model_path

    def _download_image(self, image):
        start = time.time()
        headers = {}
        image_url = image
        local_image_path = SmartBodyCrop.tmp_path + 'image'
        if type(image) is dict:
            headers = image.get('headers')
            image_url = image.get('uri')
        # check if the image is a local file path
        if type(image) is str:
            if not image.startswith('http'):
                return image

        http = urllib3.PoolManager(
            cert_reqs='CERT_REQUIRED',
            ca_certs=certifi.where(),
            headers=headers)
        urllib3.disable_warnings()

        r = http.request('GET', image_url,
                         preload_content=False,
                         retries=urllib3.Retry(5, redirect=5))

        with open(local_image_path, 'wb') as out:
            while True:
                data = r.read(1024)  # 8192
                if not data:
                    break
                out.write(data)

        r.release_conn()
        print_time("image downloaded in :", start)
        return local_image_path

    def load_graph_def(self):
        start = time.time()

        local_model_path = self._download_model()

        tf.reset_default_graph()
        graph_def = graph_pb2.GraphDef()
        with open(local_model_path, 'rb') as f:
            graph_def.ParseFromString(f.read())
        tf.import_graph_def(graph_def, name='')

        start = print_time("model imported in :", start)
        start = time.time()

        # SmartBodyCrop.initialized = True

    def infer(self, image, upper_body, lower_body):
        start = time.time()

        imgpath = self._download_image(image)
        image, thumbnail, input_width, input_height = self.read_img(
            imgpath, 368, 368)
        start = print_time("image (" + str(input_width) +
                           "x" + str(input_height) + ") loaded in: ", start)

        if not SmartBodyCrop.initialized:
            print("Loading the model...")
            self.load_graph_def()

        with tf.Session() as sess:
            inputs = tf.get_default_graph().get_tensor_by_name('inputs:0')
            heatmaps_tensor = tf.get_default_graph().get_tensor_by_name(
                'Mconv7_stage6_L2/BiasAdd:0')
            pafs_tensor = tf.get_default_graph().get_tensor_by_name(
                'Mconv7_stage6_L1/BiasAdd:0')

            heatMat, pafMat = sess.run(
                [heatmaps_tensor, pafs_tensor], feed_dict={inputs: image})

            start = print_time("tf session executed in: ", start)

            humans = estimate_pose(heatMat[0], pafMat[0])
            start = print_time("pose estimated in: ", start)
            # send the thumbnail to render an initial crop
            img, crop_position, crop_size = crop_image(
                thumbnail, humans, upper_body, lower_body)
            # scale back the crop_coordinates to match the original picture size
            scale_factor_w = input_width / thumbnail.size[0]
            scale_factor_h = input_height / thumbnail.size[1]
            crop_coordinates = {
                'x':      crop_position[0] * scale_factor_w,
                'y':      crop_position[1] * scale_factor_h,
                'width':  crop_size[0] * scale_factor_w,
                'height': crop_size[1] * scale_factor_h
            }

            start = print_time("image cropped in: ", start)

            sess.close()
            return img, crop_coordinates, imgpath

    def detect_parts(self, image):
        start = time.time()

        imgpath = self._download_image(image)
        image, thumbnail, input_width, input_height = self.read_img(
            imgpath, 368, 368)
        start = print_time("image loaded in: ", start)

        if not SmartBodyCrop.initialized:
            print("Loading the model...")
            self.load_graph_def()

        with tf.Session() as sess:
            inputs = tf.get_default_graph().get_tensor_by_name('inputs:0')
            heatmaps_tensor = tf.get_default_graph().get_tensor_by_name(
                'Mconv7_stage6_L2/BiasAdd:0')
            pafs_tensor = tf.get_default_graph().get_tensor_by_name(
                'Mconv7_stage6_L1/BiasAdd:0')

            heatMat, pafMat = sess.run(
                [heatmaps_tensor, pafs_tensor], feed_dict={inputs: image})

            start = print_time("tf session executed in: ", start)

            humans = estimate_pose(heatMat[0], pafMat[0])
            start = print_time("pose estimated in: ", start)

            # display
            img1 = draw_humans(thumbnail, humans)
            return img1