### ------------------------------- ###
###            libraries            ###
### ------------------------------- ###

from tensorflow.keras.models import load_model
import gradio as gr # remove later
import numpy as np
import os
from yattag import Doc
# import h5py # remove later

### -------------------------------- ###
###           model loading          ###
### -------------------------------- ###

model = load_model('model.h5') # single file model from colab
labels = ['please upload categories.txt' for i in range(10)] # placeholder


## --------------------------------- ###
###     reading: categories.txt      ###
### -------------------------------- ###
if os.path.isfile("categories.txt"):
    # open info.txt in read mode
    categories = open("categories.txt", "r")
    labels = categories.readline().split()
    print(labels)


## --------------------------------- ###
###        reading: info.txt         ###
### -------------------------------- ###
# placeholders in case info.txt does not exist
placeholder =  "please create an info.txt to customize this text"
title = bkgd = data_collection = priv_cons = bias_cons = ident_cons = img_src = membs = placeholder
description = "An AI project created by [name], [name], and [name]"
# check if info.txt is present
if os.path.isfile("info.txt"):
    # open info.txt in read mode
    info = open("info.txt", "r")
 
    # each line to a string
    title = info.readline()
    bkgd = info.readline()
    data_collection = info.readline()
    priv_cons = info.readline()
    bias_cons = info.readline()
    ident_cons = info.readline()
    img_src = info.readline()
    membs = info.readline()
    
    # close file
    info.close()
    
# use yattag library to generate html
doc, tag, text, line = Doc().ttl()
# create html based on info.txt
with tag('div'):
  with tag('div', klass='my-div'):
    line('h2', 'Project Background')
    line('p', bkgd)
  with tag('div', klass='my-div'):
    line('h2', 'Data Collection')
    line('p', data_collection)
  with tag('div', klass='my-div'):
    line('h2', 'Ethical Considerations')
    with tag('ul'):
      line('li', priv_cons)
      line('li', bias_cons)
      line('li', ident_cons)
  with tag('div', klass='my-div'):
    line('h2', 'Our Team')
    line('p', membs)
    doc.stag('img', src=img_src)
    
my_css = '''
.my-div {
  border: 2px solid black;
  text-align: center;
  margin: 10px;
  padding: 5%;
}
ul {
  display: inline-block;
  text-align: left;
}
img {
  display: block;
  margin: auto;
}
.description {
  text-align: center;
}
'''
### ------------------------------- ###
###        interface creation       ###
### ------------------------------- ###

def preprocess(image):
  image = np.array(image) / 255
  image = np.expand_dims(image, axis=0)
  return image
  
def predict_image(image):
  pred = model.predict(preprocess(image))
  results = {}
  for row in pred:
    for idx, item in enumerate(row):
      results[labels[idx]] = float(item)
  return results

image = gr.inputs.Image(shape=(300, 300), label="Upload Your Image Here")
label = gr.outputs.Label(num_top_classes=len(labels))

gr.Interface(fn=predict_image, inputs=image, outputs=label, capture_session=True, article=doc.getvalue(), css=my_css, theme='huggingface', title=title, allow_flagging=False, description=description).launch(debug=True)