import streamlit as st

from PIL import Image

import numpy as np

from zipfile import ZipFile

import os

import cv2

from src.production import read_files, get_setup, make_masks, create_folder, make_legend

@st.cache

def cached_get_setup():

    return get_setup()

def main():

    models, transforms = cached_get_setup()

    st.markdown(

        f"""

    <style>

        .sidebar .sidebar-content {{

            background: url("https://i.ibb.co/BL3qFQW/background.png");

            background-repeat: repeat;

            background-size: 100% auto;

    }}

        .reportview-container {{

            background: url("https://i.ibb.co/BL3qFQW/background.png");

            background-repeat: repeat;

            background-size: 100% auto;

        }}

        .reportview-container .main .block-container{{

            max-width: 850px;

            padding-top: 0rem;

            padding-right: 0rem;

            padding-left: 0rem;

            padding-bottom: 0rem;

        }}

    </style>

    """,

        unsafe_allow_html=True,

    )

    for folder in ['segmentations/', 'images/']:

        create_folder(folder)

    st.title('Сегментация поражения легких коронавирусной пневмонией')

    st.subheader("Загрузка файлов")

    filenames = st.file_uploader('Выберите или ператащите сюда снимки', type=['png', 'jpeg', 'jpg', '.nii', '.nii.gz'],

                                 accept_multiple_files=True)

    multi_class = st.checkbox(label='Мульти-классовая сегментация', value=False)

    show_legend = st.checkbox(label='Легенда на картинке', value=False)

    if st.button('Загрузить') and filenames:

        paths, folder_name = read_files(filenames)

        if not paths:

            st.error('Неправильный формат или название файла')

        else:

            user_dir = "segmentations/" + folder_name

            # creating folders

            create_folder(user_dir)

            create_folder(os.path.join(user_dir, 'segmentations'))

            create_folder(os.path.join(user_dir, 'annotations'))

            zip_obj = ZipFile(user_dir + 'segmentations.zip', 'w')

            with st.expander("Информация о каждом фото"):

                info = st.info('Делаем предсказания, пожалуйста, подождите')

                for _paths in paths:

                    for img, annotation, original_path in make_masks(_paths, models, transforms, multi_class):

                        name = original_path.split('/')[-1].split('.')[0]

                        name = name.replace('\\', '/')

                        # saving annotation

                        annotation_path = os.path.join(user_dir, 'annotations', name + '_annotation.txt')

                        with open(annotation_path, mode='w') as f:

                            f.write(annotation)

                        info.empty()

                        # name and annotation

                        st.markdown(f'<h3>{name}</h3>', unsafe_allow_html=True)

                        if not show_legend:

                            if len(annotation.split('\n')) == 3:

                                st.markdown('[red] - consolidation')

                                st.markdown('[green] - ground-glass')

                            else:

                                st.markdown('[yellow] - disease')

                            for line in annotation.split('\n'):

                                st.markdown(line)

                        col1, col2 = st.columns(2)

                        # original image

                        original = np.array(Image.open(original_path))

                        col1.header("Оригинал")

                        col1.image(original, width=350)

                        # refactoring image

                        if show_legend:

                            img = make_legend(img, annotation)

                        # saving image

                        path = os.path.join(user_dir, 'segmentations', name + '_mask.png')

                        cv2.imwrite(path, img)

                        # adding in zip

                        zip_obj.write(path)

                        zip_obj.write(annotation_path)

                        # show segmentation

                        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

                        img = img / 255  # to [0;1] range

                        # print(img.shape, img.dtype, img)

                        col2.header("Сегментация")

                        col2.image(img, width=350)

                        st.markdown('<br />', unsafe_allow_html=True)

                zip_obj.close()

            # download segmentation zip

            with st.expander("Скачать сегментации"):

                with open(os.path.join(user_dir, 'segmentations.zip'), 'rb') as file:

                    st.download_button(

                        label="Архив сегментаций",

                        data=file,

                        file_name="segmentations.zip")

if __name__ == '__main__':

    main()