from statistics import mode

import streamlit as st
from fastai.vision.all import *
from PIL import Image

from Processor import Processor


@st.experimental_singleton
def initialize_app():
    return Processor(load_learner('model.pkl'))


def process_images(images, processor: Processor):
    filtered_images = []
    result = []
    class_names = list(
        map(lambda name: {name: 0}, processor.inference.dls.vocab))

    for image in images:
        image = Image.open(image)
        if processor.filter_image(image):
            filtered_images.append(np.asarray(image))

    for img in filtered_images:
        result.append(processor.classify_image(img)[0])

    if len(result) == 0:
        return None

    for res_name in result:
        for idx, class_name in enumerate(class_names):
            for key, value in class_name.items():
                if res_name == key:
                    class_names[idx][key] = value + 1

    outfit = mode(result)

    with open(f'./texts/{outfit}.txt') as text:
        personality = text.read()

    return {'outfit': outfit.title(), 'personality': personality,
            'chart': class_names}


# Streamlit UI

processor = initialize_app()

st.title('Instagram Clothes Psychology (Photos)')
uploaded_photos = st.file_uploader(label="Upload photos", type=[
    'jpg', 'jpeg'], accept_multiple_files=True)

photos_empty = True if len(uploaded_photos) == 0 else False

is_clicked = st.button(label='Predict Personality',
                       disabled=photos_empty)

if is_clicked:
    with st.spinner('Please wait...'):
        result = process_images(uploaded_photos, processor)
        if result is None:
            st.write('Tidak ditemukan gambar yang valid')
        else:
            st.header('Your personality is..')
            st.subheader(result['outfit'])
            st.markdown(result['personality'])
            st.bar_chart(result['chart'])