# coding: utf-8
"""
    stylelens-feature

    This is a API document for Object Detection on fashion items\"

    OpenAPI spec version: 0.0.1
    Contact: devops@bluehack.net
    Generated by: https://github.com/swagger-api/swagger-codegen.git
"""

from __future__ import absolute_import

import os
import time
import grpc
from PIL import Image
import numpy as np
# import tensorflow as tf

from .grpc import feature_extract_pb2_grpc
from .grpc import feature_extract_pb2

CHUNK_SIZE = 4 * 1024

GPU_HOST = os.environ['FEATURE_GRPC_HOST']
GPU_PORT = os.environ['FEATURE_GRPC_PORT']

TMP_FILE = 'img.jpg'

class ExtractFeature(object):
  def __init__(self, use_gpu=True):
    self.use_gpu = use_gpu

    # if use_gpu == False:
    #   try:
    #     MODEL = os.environ['CLASSIFY_GRAPH']
    #   except:
    #     print("!!! Need to define environment variable: CLASSIFY_GRAPH=/path/to/model.pb")
    #
    #   print(MODEL)
    #   with tf.gfile.FastGFile(MODEL, 'rb') as f:
    #     graph_def = tf.GraphDef()
    #     graph_def.ParseFromString(f.read())
    #     _ = tf.import_graph_def(graph_def, name='')
    #
    #   with tf.Session() as self.sess:
    #     self.pool3 = self.sess.graph.get_tensor_by_name('pool_3:0')

  def extract_feature(self, file):

    im = Image.open(file)
    size = 300, 300
    im.thumbnail(size, Image.ANTIALIAS)
    im.save(TMP_FILE)
    # with tf.gfile.GFile(TMP_FILE, 'rb') as fid:
    #   image_data = fid.read()
    with open(TMP_FILE, 'rb') as fid:
      image_data = fid.read()

    if self.use_gpu:
      print('use GPU')
      start_time = time.time()
      channel = grpc.insecure_channel(GPU_HOST + ':' + GPU_PORT)
      stub = feature_extract_pb2_grpc.ExtractStub(channel)
      response = stub.GetFeature(feature_extract_pb2.FeatureRequest(file_data=image_data))
      # arr = np.fromstring(response.vector, dtype=np.float32)
      elapsed_time = time.time() - start_time
      print(elapsed_time)
      print('done')
      # return arr
      return response.vector

    # else:
    #   print('use CPU')
    #   start_time = time.time()
    #   pool3_features = self.sess.run(self.pool3,{'DecodeJpeg/contents:0': image_data})
    #   elapsed_time = time.time() - start_time
    #   print(elapsed_time)
    #   feature = np.squeeze(pool3_features)
    #   return feature
