# 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 = 'magi-0.stylelens.io'
GPU_PORT = '50051'

TMP_FILE = 'img.jpg'

class ExtractFeature(object):
  def __init__(self, use_gpu=False):
    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()

    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

    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
