from keras.models import load_model
from PIL import Image, ImageOps
import numpy as np
import gradio as gr
import pandas as pd

data = pd.read_csv('species_info.csv')

def format_label(label):
  """
  From '0 rùa khác\n' to 'rùa khác'
  """
  return label[label.find(" ")+1:-1]

def check_species_status(species_name):
  status = ''

  return status

def info(species_name):
  status = check_species_status(species_name)
  if status == '':
    info = ''
  return info

def get_vi_name(en_name):
  """
  Return name in Vietnamese
  """
  print(en_name)
  return data[data['en_name'] == en_name]['vi_name'].to_list()[0]

def get_law(en_name):
  cites = data[data['en_name'] == en_name]['CITES'].to_list()[0]
  nd06 = data[data['en_name'] == en_name]['ND06'].to_list()[0]
  return cites, nd06

def get_habitat(en_name):
  return data[data['en_name'] == en_name]['habitat'].to_list()[0]
  
def predict(image):

  # Load the model
  model = load_model('keras_model.h5')

  # Create the array of the right shape to feed into the keras model
  # The 'length' or number of images you can put into the array is
  # determined by the first position in the shape tuple, in this case 1.
  data = np.ndarray(shape=(1, 224, 224, 3), dtype=np.float32)
  
  #resize the image to a 224x224 with the same strategy as in TM2:
  #resizing the image to be at least 224x224 and then cropping from the center
  size = (224, 224)
  image = ImageOps.fit(image, size, Image.ANTIALIAS)

  #turn the image into a numpy array
  image_array = np.asarray(image)
  # Normalize the image
  normalized_image_array = (image_array.astype(np.float32) / 127.0) - 1
  # Load the image into the array
  data[0] = normalized_image_array

  # run the inference
  pred = model.predict(data)
  pred = pred.tolist()

  with open('labels.txt','r') as f:
    labels = f.readlines()
  
  en_name = format_label(labels[pred.index(max(pred))]).strip()

  result = {get_vi_name(format_label(labels[i])): round(pred[0][i],2) for i in range(len(pred[0]))}
  cites, nd06 = get_law(en_name)
  info = ""
  if str(nd06) != "":
    info += f'CITES: {cites}, NĐ06: {nd06} \n \n'
    info += "Đây là loài được pháp luật bảo vệ. Mọi hành vi buôn bán, nuôi nhốt không có \
    [giấy phép](https://thuvienphapluat.vn/van-ban/Tai-nguyen-Moi-truong/Nghi-dinh-06-2019-ND-CP-quan-ly-thuc-vat-rung-dong-vat-rung-nguy-cap-quy-hiem-405883.aspx) đều vi phạm pháp luật \n"
    info += "- Nếu bạn vô tình bắt gặp loài này bị buôn bán mà không có giấy phép, \
     tuyệt đối không mua nhằm bất kỳ mục đích gì (ví dụ để phóng sinh) \
     mà nên báo cáo vi phạm tại đường dây nóng bảo vệ DVHD của ENV **1800-1522**. \n"
    info += "- Nếu bạn đang nuôi thì nên giao nộp cho cơ quan chức năng để trả về tự nhiên. Tham khảo đơn vị tiếp nhận DVHD ở địa phương \
    bạn tại [đây](https://drive.google.com/file/d/1K2ZWcHKGEsNudh_LtHgHJOXlVw-GQ6AZ/view). \n"
    info += f"- Nếu bạn bắt gặp trong vườn nhà thì có thể xem xét thả chúng về môi trường sống. Hãy đảm bảo nơi bạn thả là\
     **{get_habitat(en_name)}**."
  
  return result, info
  
  
description="""
VNTurtle nhận diện các loài rùa Việt Nam. Mô hình mẫu này có thể nhận diện 10 loại rùa thường xuất hiện ở VN gồm **5** loài bản địa
**(1) Rùa núi viền**, **(2) Rùa núi vàng**, **(3) Rùa ba gờ**, **(4) Rùa răng**, và **(5) Rùa sa nhân**, 
và **5** loài ngoại lai **(1) Rùa Sulcata**, **(2) Rùa bản đồ**, **(3) Rùa cá sấu**, **(4) Rùa tai đỏ**, và **(5) Rùa ninja**
"""
article = "© Hình ảnh minh hoạ được cung cấp bởi ATP"
examples = [ ['test1.jpg'], ['test2.jpg'], ['test3.jpg']  ]
gr.Interface(fn=predict, 
             inputs=gr.Image(type="pil", label="Input Image"),
             outputs=[gr.Label(label="Predictions"), gr.Markdown()],
             live=True,
             title='VNTurtle - Toy Model',
             description=description,
             examples=examples,
             article=article).launch()