import numpy as np
import streamlit as st
from PIL import Image
import urllib.request
import io
import tensorflow as tf
from utils import preprocess_image

# Initialize labels and model
labels = ['cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash']
model = tf.keras.models.load_model('classify_model.h5')

# Streamlit UI
st.markdown('''
 <div style="padding-bottom: 20px; padding-top: 20px; padding-left: 5px; padding-right: 5px">
 <center><h1>EcoIdentify (Test)</h1></center>
 </div>
''', unsafe_allow_html=True)

st.markdown('''
 <div>
 <center><h3>Please upload Waste Image to find its Category</h3></center>
 </div>
''', unsafe_allow_html=True)

opt = st.selectbox(
    "How do you want to upload the image for classification?",
    ("Please Select", "Upload image via link", "Upload image from device"),
)

image = None

if opt == 'Upload image from device':
  file = st.file_uploader('Select', type=['jpg', 'png', 'jpeg'])
  if file:
      image = preprocess_image(file)
    
elif opt == 'Upload image via link':
  img_url = st.text_input('Enter the Image Address')
  if st.button('Submit'):
    try:
      response = urllib.request.urlopen(img_url)
      image = preprocess_image(response)
    except ValueError:
      st.error("Please Enter a valid Image Address!")


if image is not None:
    st.image(image, width=256, caption='Uploaded Image')
if st.button('Predict'):
  prediction = model.predict(image[np.newaxis, ...])

  print("---------------img-array---------------------")
  print(image[np.newaxis, ...])
  print("------------summary------------------------")
  print(model.summary())
  print("------------------------------------")
  print(prediction)

  st.info('Hey! The uploaded image has been classified as " {} waste " '.format(labels[np.argmax(prediction[0], axis=-1)]))

  def message(img):
    if img == 'paper' or 'cardboard' or 'metal' or 'glass':
      return (" therefore your item is recyclable. Please refer to https://www.wm.com/us/en/drop-off-locations to find a drop-off location near you.")
    elif img == 'plastic':
      return ("therefore your item may have a chance of being recyclable.")