import gradio as gr

from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.messages import HumanMessage, SystemMessage


import tensorflow as tf
from tensorflow import keras

import json
import numpy as np
from PIL import Image


bird_cnn = keras.models.load_model('bird_model.h5')

with open('class_names.json', 'r') as class_name:
    class_names = json.load(class_name)

llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash-latest", convert_system_message_to_human=True)


def get_info(language, img):
    img = img.resize((224, 224), resample=Image.NEAREST)
    img_array = keras.utils.img_to_array(img)
    img_array = tf.expand_dims(img_array, 0)

    predictions = bird_cnn.predict(img_array)
    score = tf.nn.softmax(predictions[0])

    species = class_names[np.argmax(score)]

    result = llm(
      [
          SystemMessage(content=f'Anda adalah sistem penulis ensiklopedia yang memberikan informasi teraktual dan '
                                f'bersifat fakta mengenai dunia aves.'
                                f'Setiap kali memberikan informasi, berikan dengan struktur yang terformat layaknya '
                                f'buku ensiklopedia.'
                                f'Informasi yang diberikan antara lain adalah taksonomi, habitat, populasi, tingkah '
                                f'laku, serta sebarannya diseluruh dunia.'
                                f'Berikan Output yang detail mengenai Species Burungnya tambahkan Infomasi yang '
                                f'berkaitan dengan Informasi yang diinginkan sebelumnya'
                                f'Jadikan Struktur output rapih dan bersih agar lebih mudah dibaca dan dipahami'
                                f'Tampilkan Domain, Kingdom, Phylum, Class, Order, Family, Genus, dan Species dari '
                                f'Species Burung sebagai infomasi Taksonomi'
                                f'Berikan output tanpa format teks dan hanya berupa raw string dengan garis baru dan '
                                f'pemaragrafan seperti docs'
                                f'Output tidak perlu menggunakan markdown'
                                f'Nama Species Burung (Nama Latin)'
                                f''
                                f'Taksonomi'
                                f'- Domain :'
                                f'- Kingdom :'
                                f'- Phylum :'
                                f'- Class :'
                                f'- Order :'
                                f'- Family :'
                                f'- Genus :'
                                f'- Species :'
                                f''
                                f'Habitat'
                                f''
                                f'Populasi'
                                f''
                                f'Tingkah Laku'
                                f''
                                f'Penyebaran'
                                f'Jadikan struktur diatas menjadi default untuk output dari species burung dan selalu '
                                f'gunakan struktur ini tanpa perlu markdown'
                                f'Jika jenis burung bukan merupakan spesies spesifik, '
                                f'berikan informasi umum mengenai burung dan '
                                f'pada kolom species berikan setiap species yang termasuk dalam jenis burung tersebut'
                                f'lalu pada nama spesies burung, '
                                f'berikan nama latin untuk genus saja. Kasuari (Casuarius). '
                                f'Nama burung disesuaikan dengan bahasa yang dipilih oleh pengguna, jika pengguna '
                                f'memilih bahasa indonesia, maka nama burung harus disesuaikan dengan bahasa indonesia'
                                f'Jangan tambahkan simbol sebagain pembeda judul dan sub-judul dengan teks biasa'
                                f'Tambahkan Referensi darimana informasi didapatkan dalam bentuk journal, website '
                                f'atau referensi manapun, gunakan formatting seperti mendeley untuk referensi'),
          HumanMessage(content=f'Berikan informasi taksonomi, habitat, Populasi, Tingkah Laku, dan Penyebaran '
                               f'mengenai burung {species} menggunakan bahasa {language}'),
      ]
    )
    return result.content


title = "Aviary Vision"
description = """
<div style="display: flex; flex-direction: column; align-items: center;">
    <h1 style="margin: 0; font-size: 60px;">Aviary Vision</h1>
    <p style="margin: 0; font-size: 20px;">By: A2 Si Paling Data</p>
</div>
<p>Aviary Vision is a system that can identify bird species from an image and provide 
information about the species in the selected language.</p>
<i>Note: The information provided is AI generated and may not be accurate. Use at your own risk.</i>
"""
examples = [
    ['English', 'bali_starling.png'],
    ['Bahasa Indonesia', 'cassowary.png']
]

demo = gr.Interface(
    fn=get_info,
    inputs=[
        gr.Dropdown(choices=[
                        'English',
                        'Bahasa Indonesia',
                        '日本語',
                        '中文',
                        '한국어',
                        'Español',
                        'Français',
                        'Deutsch',
                        'Русский'
                    ],
                    label="Select Language",
                    type="value",
                    info="Select the language you want to use"),
        gr.Image(type="pil",
                 label="Upload Image")
    ],
    outputs=['text'],
    allow_flagging="never",
    title=title,
    description=description,
    examples=examples,
    cache_examples=False
)

demo.launch()