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 info(species_name):
  status = check_species_status(species_name)
  if status == '':
    info = ''
  return info

def get_name(en_name, lan):
  """
  Return name in Vietnamese 
  """
  return data[data[f'name_en'] == en_name][f'name_{lan}'].to_list()[0]

def get_fun_fact(en_name, lan):
  """
  Return fun fact of the species
  """
  return data[data[f'name_en'] == en_name][f'fun_fact_{lan}'].to_list()[0]

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

def get_habitat(en_name, lan):
  return data[data['name_en'] == en_name][f'habitat_{lan}'].to_list()[0]

def get_conservation_status(en_name, lan):
  status = data[data['name_en'] == en_name]['conservation_status'].to_list()[0]
  status_dict = {
      'NE': {
          'vi': "Không được đánh giá",
          'en': "Not Evaluated"
      },
      'DD': {
          'vi': "Thiếu dữ liệu",
          'en': "Data Deficient"
      },
      'LC': {
          'vi': "Ít quan tâm",
          'en': "Least Concern"
      },
      'NT': {
          'vi': "Sắp bị đe dọa",
          'en': "Near Threatened"
      },
      'VU': {
          'vi': "Sắp nguy cấp",
          'en': "Vulnerable"
      },
      'EN': {
          'vi': "Nguy cấp",
          'en': "Endangered"
      },
      'CR': {
          'vi': "Cực kỳ nguy cấp",
          'en': "Critically Endangered"
      },
      'EW': {
          'vi': "Tuyệt chủng trong tự nhiên",
          'en': "Extinct in the Wild"
      },
      'EX': {
          'vi': "Tuyệt chủng",
          'en': "Extinct"
      },
  }
  return status_dict[status][lan]
  
def predict(image, language="Tiếng Việt"):

  # 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()

  sorted_index = np.argsort(pred).tolist()[0]

  print(sorted_index)

  display_index = []
  for i in sorted_index[::-1]:
    if pred[0][i] > 0.01:
      display_index.append(i)

  with open('labels.txt','r') as f:
    labels = f.readlines()
  
  en_name = format_label(labels[sorted_index[-1]]).strip()

  lan = ''
  if language == "Tiếng Việt":
    lan = 'vi'
  else:
    lan = 'en'

  result = {get_name(format_label(labels[i]), lan): round(pred[0][i],2) for i in display_index}
  cites, nd06 = get_law(en_name)
  info = ""

  fun_fact = ""
  if lan == 'vi':
    fun_fact += "💡 **Thông tin lý thú** \n\n"
  if lan == 'en':
    fun_fact += "💡 **Fun fact** \n\n"

  fun_fact += f"{get_fun_fact(en_name, lan)}."

  status = ""
  if lan == 'vi':
    status += "**🐢 Tình trạng bảo tồn** \n\n"
  if lan == 'en':
    status += "**🐢 Conservation status** \n\n"

  status += f"{get_conservation_status(en_name, lan)}"
  
  law = ""
  if lan == 'vi':
    law += "**👮‍♀️ Luật** \n\n"
  if lan == 'en':
    law += "**👮‍♀️ Law** \n\n"
  law += f'CITES: {cites}, NĐ06: {nd06}'

  if lan == 'vi' and str(nd06) != "":
    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, lan)}**. \n\n"
    info += f"Liên hệ Chương trình Bảo tồn Rùa châu Á ([ATP](https://asianturtleprogram.org/vi/)) qua email: info@asianturtleprogram.org hoặc [Facebook](https://www.facebook.com/search/top?q=asian%20turtle%20program) để được hướng dẫn trong các trường hợp cụ thể"

  if lan == 'en' and str(nd06) != "":
    info += "This species is protected by law. All acts of trafficking and captive breeding without \
    [authority permisison](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) are illegal \n"
    info += "- If you happen to come across this species being traded without a license, \
    do not buy for any purpose (e.g., to release). \
    Instead, you should report the suspicious act to ENV via their hotline **1800-1522**. \n"
    info += "- If you are raising them, you should hand them over to the authorities to return to the wild. Find your local wild-animal receiver \
    [here](https://drive.google.com/file/d/1K2ZWcHKGEsNudh_LtHgHJOXlVw-GQ6AZ/view). \n"
    info += f"- If you find them in your garden, consider releasing them back into their habitat. Please make sure the place you release is \
     **{get_habitat(en_name, lan)}**. \n\n"
    info += f"Contact Asian Turtle Program ([ATP](https://asianturtleprogram.org/)) via email (info@asianturtleprogram.org) or [Facebook](https://www.facebook.com/search/top?q=asian%20turtle%20program) for instructions in specific cases."
  
  return result, fun_fact, status, law, 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 \n\n
**(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 Trung Bộ** \n\n
và **5** loài ngoại lai \n\n
**(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"), gr.Radio(value="Tiếng Việt", choices=["Tiếng Việt", "English"], label="Language")],
             outputs=[gr.Label(label="Predictions"), gr.Markdown(), gr.Markdown(), gr.Markdown(), gr.Markdown()],
             live=True,
             title='VNTurtle - Toy Model',
             description=description,
             examples=examples,
             article=article).launch()