import streamlit as st
import moduls.Texture as texture
from PIL import Image
import mmcv
import moduls.model
from moduls.texture_for_point import texture_transformation_1
from main import model
import numpy as np 
import cv2
import os
from streamlit_image_select import image_select
from streamlit_image_annotation import pointdet
from moduls.conturs import conturs_on_image
from streamlit_image_comparison import image_comparison
from moduls.settings import path_for_streamlit
st.set_page_config(layout="wide")


@st.cache_data
def load_image(image_file):
    img = Image.open(image_file)
    img = mmcv.imread(np.array(img))
    return img 
def save_file_st(path_file):
    file_name = path_file.name
    img = load_image(path_file)
    im = Image.fromarray(img)
    output_filename = 'processing.jpg' 
    im.save(output_filename)
    return output_filename
@st.dialog("Выбор текстуры пола", width =  'large')
def select_floor(img_list_1, image_1, img_list):
    tex_floor = image_select("Выбор текстуры пола", img_list_1, return_value= 'index')
    if st.button("Выбрать"):
        st.session_state.json_list[0] = img_list[tex_floor]
        modificate_image('floor')
        st.session_state.select_floor = [tex_floor]
        st.rerun()
@st.dialog("Выбор текстуры стен", width =  'large')
def select_wall(img_list_1, image_1, img_list):
    tex_wall = image_select("Выбор текстуры стен", img_list_1, return_value= 'index')
    if st.button("Выбрать"):
        st.session_state.json_list[1] = img_list[tex_wall]
        modificate_image('wall')
        st.session_state.select_wall = [tex_wall]
        st.rerun()
@st.dialog("Выбор текстуры потолка", width =  'large')
def select_ceiling(img_list_1, image_1, img_list):
    tex_ceiling = image_select("Выбор текстуры потолка", img_list_1, return_value= 'index')
    if st.button("Выбрать"):
        st.session_state.json_list[2] = img_list[tex_ceiling]
        modificate_image('ceiling')
        st.session_state.select_ceiling = [tex_ceiling]
        st.rerun()
def points_sorter(points):
    messages_list = ['Недостаточно точек. Добавьте точек!',
                    'Точек не четное количество. Добавьте или уберите точку!',
                    'Точек слишком много. Уберите лишние точки!',
                    'Точки успешно отсортированы!',
                    'Пол или потолок имеют обратную перспективу. Скорректируйте точки!']
    print('len(points)', len(points))
    print(points)
    if (num_points := len(points)) < 6:
        return [], []
    elif num_points%2 != 0:
        return [], []
    elif num_points > 8:
        return [], []
    else:
        # Сортируем список по y, а при равных y по х
        sorted_list = sorted(points, key=lambda k: [k[1], k[0]])
        print('sorted_list', sorted_list)
        # Делим список пополам и сортируем по x
        ceiling_points = sorted(sorted_list[:num_points//2], key=lambda k: [k[0], k[1]])
        floor_points = sorted(sorted_list[num_points//2:], key=lambda k: [k[0], k[1]])
        bad_points = []
        print('ceiling_points, floor_points', ceiling_points, floor_points)
        # Проверяем направление перспективы потолка
        if ceiling_points[0][1] > ceiling_points[1][1]:
            bad_points.append(ceiling_points[0])
            bad_points.append(ceiling_points[1])
        if ceiling_points[-1][1] > ceiling_points[-2][1]:
            bad_points.append(ceiling_points[-1])
            bad_points.append(ceiling_points[-2])
        # Проверем направление перспективы пола
        if floor_points[0][1] < floor_points[1][1]:
            bad_points.append(floor_points[0])
            bad_points.append(floor_points[1])
        if floor_points[-1][1] < floor_points[-2][1]:
            bad_points.append(floor_points[-1])
            bad_points.append(floor_points[-2])
        c_f_point = [ceiling_points, floor_points]
    return c_f_point
    #if len(bad_points) != 0:
    #  return ceiling_points, floor_points
    #else:
    #  return ceiling_points, floor_points
#@st.dialog('Определение углов комнаты', width='large')
def corner_change_2(image_path):
    list_1 = ['Угол']
    list_label_1 = []
    list_label_2 = []
    list_point_1 = []
    list_point_2 = []
    if 'ceiling_point' in st.session_state and 'floor_point' in st.session_state and st.session_state.ceiling_point is not None and st.session_state.floor_point is not None:
        list_point_1 = st.session_state.ceiling_point
        print('list_point_1', list_point_1)
        for i in range(len(st.session_state.ceiling_point)):
            list_label_1.append(0)
        list_point_2 = st.session_state.floor_point
        for i in range(len(st.session_state.floor_point)):
            list_label_2.append(0)
    else:
            list_point_1 = []
            list_point_2 = []
    list_point = list_point_1 + list_point_2
    list_label = list_label_1 + list_label_2
    print('list_point', list_point)
    print('list_label', list_label)
    point = pointdet(image_path , list_1, 
                                    points=list_point,
                                    #points = [],
                                    labels=list_label,
                                    #labels = [],
                                    use_space=True)
    if point is not None:
        st.session_state.point = point
        st.session_state.floor_point = []
        st.session_state.ceiling_point = []
    #if 'points' not in st.session_state:
        st.session_state.points = []
    if point is not None:
        for i in st.session_state.point:
            if i['label_id'] == 0:
                #st.session_state.ceiling_point.append([int(i['point'][0]), int(i['point'][1])])
                st.session_state.points.append([int(i['point'][0]), int(i['point'][1])])
            else:
                #st.session_state.floor_point.append([int(i['point'][0]), int(i['point'][1])])
                st.session_state.points.append([int(i['point'][0]), int(i['point'][1])])
        print('st.session_state.points', st.session_state.points)
        c_f_point = points_sorter(st.session_state.points)
        st.session_state.ceiling_point = c_f_point[0]
        st.session_state.floor_point = c_f_point[1]
        #print('a', a)
        #st.session_state.ceiling_point = sort_list_ceiling(st.session_state.ceiling_point)
        print(st.session_state.ceiling_point)
        #st.session_state.floor_point = sort_list_floor(st.session_state.floor_point)
        #st.session_state.ceiling_point = cv2.convexHull(np.array(st.session_state.ceiling_point))
        #st.session_state.floor_point = cv2.convexHull(np.array(st.session_state.floor_point))
        #print('st.session_state.ceiling_point', st.session_state.ceiling_point)
        #print('st.session_state.floor_point', st.session_state.floor_point)
        #st.rerun()
def sort_list_ceiling(no_change_list):
    if len(no_change_list) >= 4:
        no_change_list.sort()
        if no_change_list[0][1] > no_change_list[1][1]:
            no_change_list[0], no_change_list[1] = no_change_list[1], no_change_list[0]
        if no_change_list[-1][1] > no_change_list[-2][1]:
            no_change_list[-1], no_change_list[-2] = no_change_list[-2], no_change_list[-1]
    if len(no_change_list) == 3:
        no_change_list.sort()
        if no_change_list[0][1] > no_change_list[1][1]:
            no_change_list[0], no_change_list[1] = no_change_list[1], no_change_list[0]
    return no_change_list
def sort_list_floor(no_change_list):
    if len(no_change_list) >= 4:
        no_change_list.sort()
        if no_change_list[0][1] < no_change_list[1][1]:
            no_change_list[0], no_change_list[1] = no_change_list[1], no_change_list[0]
        if no_change_list[-1][1] < no_change_list[-2][1]:
            no_change_list[-1], no_change_list[-2] = no_change_list[-2], no_change_list[-1]
    if len(no_change_list) == 3:
        no_change_list.sort()
        if no_change_list[0][1] < no_change_list[1][1]:
            no_change_list[0], no_change_list[1] = no_change_list[1], no_change_list[0]
    return no_change_list
#@st.dialog('Определение углов комнаты', width='large')
def corner_change(image_path):
    list_1 = ['Угол потолка', 'Угол пола']
    list_label_1 = []
    list_label_2 = []
    list_point_1 = []
    list_point_2 = []
    if 'ceiling_point' in st.session_state and 'floor_point' in st.session_state and st.session_state.ceiling_point is not None and st.session_state.floor_point is not None:
        list_point_1 = st.session_state.ceiling_point
        for i in range(len(st.session_state.ceiling_point)):
            list_label_1.append(0)
        list_point_2 = st.session_state.floor_point
        for i in range(len(st.session_state.floor_point)):
            list_label_2.append(1)
    else:
            list_point_1 = []
            list_point_2 = []
    list_point = list_point_1 + list_point_2
    list_label = list_label_1 + list_label_2
    print('list_point', list_point)
    print('list_label', list_label)
    point = pointdet(image_path , list_1, 
                                    points=list_point,
                                    #points = [],
                                    labels=list_label,
                                    #labels = [],
                                    use_space=True)
    if point is not None:
        st.session_state.point = point
        st.session_state.floor_point = []
        st.session_state.ceiling_point = []
        
    if point is not None:
        for i in st.session_state.point:
            if i['label_id'] == 0:

                st.session_state.ceiling_point.append([int(i['point'][0]), int(i['point'][1])])
            else:
                st.session_state.floor_point.append([int(i['point'][0]), int(i['point'][1])])
        st.session_state.ceiling_point = sort_list_ceiling(st.session_state.ceiling_point)
        print(st.session_state.ceiling_point)
        st.session_state.floor_point = sort_list_floor(st.session_state.floor_point)
        #st.session_state.ceiling_point = cv2.convexHull(np.array(st.session_state.ceiling_point))
        #st.session_state.floor_point = cv2.convexHull(np.array(st.session_state.floor_point))
        print('st.session_state.ceiling_point', st.session_state.ceiling_point)
        print('st.session_state.floor_point', st.session_state.floor_point)
        #st.rerun()
    return point
@st.dialog("modification_photo")
def modification_photo(image_1):
    st.session_state.modification_photo = [image_1]
    st.rerun()
def a(path_photo):
    st.session_state.a = path_photo
def predict(path_photo, model):
    path_file =save_file_st(path_photo)
    _, seg_map, seg_all, image = moduls.model.predict_and_visualize(model, path_file)
    print(seg_map.shape)
    st.session_state.predict =[seg_all, image]
    contur_image, simplified_contour = conturs_on_image(image, seg_map)
    st.session_state.contur_image =contur_image
    st.session_state.simplified_contour = simplified_contour
    return image
def modificate_image(surface):
    original_img = cv2.cvtColor(st.session_state.predict[1], cv2.COLOR_BGR2RGB)
    if st.session_state.f_h !=0:
        f_h = st.session_state.f_h
    else:
        f_h = None
    if st.session_state.w_h !=0:
        w_h = st.session_state.w_h
    else:
        w_h = None
    if st.session_state.c_h !=0:
        c_h = st.session_state.c_h
    else:
        c_h = None
    image_1 = texture_transformation_1(original_img, st.session_state.predict[0],
                                    st.session_state.room_l*100, st.session_state.room_w*100, st.session_state.room_h*100,
                                    st.session_state.ceiling_point, st.session_state.floor_point,
                                    st.session_state.floor_angel, st.session_state.wall_angel, st.session_state.ceiling_angel,
                                    st.session_state.json_list, surface = st.session_state.surface,
                                    h_f= f_h, h_w = w_h, h_c = c_h,
                                    rescale = 1.)
    st.session_state.image_1 = image_1
    st.session_state.modification_photo = [image_1]
def save_angel():
    st.session_state.floor_angel = floor_angel
    st.session_state.ceiling_angel = ceiling_angel
    st.session_state.wall_angel = wall_angel
path_icon, path_json, path_logo = path_for_streamlit()
if 'json_list' not in st.session_state:
    st.session_state.json_list = [None, None, None]
img_list_1 = []
img_list = []
for path_file in os.listdir(path_icon):
    img_list_1.append(path_icon + f'\{path_file}')
for path_json_file in os.listdir(path_json):
    if path_json_file.endswith(('.json')):
        img_list.append(path_json + f'\{path_json_file}')
img_list_1.sort()
img_list.sort()
title_col1, title_col2 = st.columns([1,3])
with title_col1:
    st.image(path_logo)
with title_col2:
    st.title(":orange[Замена текстур на фотографии вашей комнаты]",)

change_col_1, change_col_2, change_col_3 = st.columns([1, 4, 1])
# Загрузка фотографиии
with change_col_1:
    path_photo = st.file_uploader("Загрузка фото", type=['png','jpeg','jpg'], accept_multiple_files=False, key=None,
                  help=None, on_change=None, args=None, kwargs=None,
                  disabled=False, label_visibility = 'collapsed')
    if path_photo == None:
        st.session_state.a = None
        st.session_state.pp = None
    if path_photo is not None:
        if st.session_state.a != path_photo:
            st.session_state.ceiling_point = None
            st.session_state.floor_point = None
            st.session_state.pp = None
            st.session_state.a = path_photo
            predict(st.session_state.a, model)
            st.session_state.modification_photo = [cv2.cvtColor(st.session_state.predict[1], cv2.COLOR_BGR2RGB)]
# Выбор угловых точек на фотографии
with change_col_2:
    if path_photo is not None:
            cv2.imwrite('corners.png',st.session_state.contur_image)
            #st.session_state.corners = corner_change('corners.png')
            st.session_state.corners = corner_change_2('corners.png')
# Размеры комнаты и выбор способа набора текстуры
with change_col_3:
    if path_photo is not None:
        st.session_state.room_l = st.number_input('Глубина комнаты', min_value=0., max_value=6., value=4., step=0.1, format=None, key='room_l_key', help=None, on_change=None, args=None, kwargs=None, placeholder=None, disabled=False, label_visibility="visible")
        st.session_state.room_w = st.number_input('Ширина комнаты', min_value=0., max_value=6., value=3., step=0.1, format=None, key='room_w_key', help=None, on_change=None, args=None, kwargs=None, placeholder=None, disabled=False, label_visibility="visible")
        st.session_state.room_h = st.number_input('Высота потолка', min_value=0., max_value=4., value=3., step=0.1, format=None, key='room_h_key', help=None, on_change=None, args=None, kwargs=None, placeholder=None, disabled=False, label_visibility="visible")
        # Выбор текстур
if 'corners'  in st.session_state and st.session_state.corners != None:
    st.session_state.pp = 1
col1, col2, col3 = st.columns(3)
# Текстура пола
with col1:
    if path_photo is not None and st.session_state.pp is not None:
        floor = st.checkbox('Текстура пола', value=False, key=None, help=None, on_change=None, args=None, kwargs=None, disabled=False, label_visibility="visible")
        if floor:
            floor_select = st.button("Выбор текстуры пола", key = '1')
            if floor_select:
                    select_floor(img_list_1, cv2.cvtColor(st.session_state.modification_photo[0], cv2.COLOR_BGR2RGB), img_list)
            try:
                st.image(img_list_1[st.session_state.select_floor[0]], caption=None, width=200, use_column_width=None, clamp=False, channels='RGB', output_format="auto")
                st.session_state.json_list[0] = img_list[st.session_state.select_floor[0]]
            except:
                b=1
        else:
            st.session_state.json_list[0] = None
# Текстура стен
with col2:
    if path_photo is not None and st.session_state.pp is not None:
        wall = st.checkbox('Текстура стен', value=False, key=None, help=None, on_change=None, args=None, kwargs=None, disabled=False, label_visibility="visible")
        if wall:
            wall_select = st.button("Выбор текстуры стен", key = '2')
            if wall_select:
                select_wall(img_list_1, cv2.cvtColor(st.session_state.modification_photo[0], cv2.COLOR_BGR2RGB), img_list)
            try:
                st.image(img_list_1[st.session_state.select_wall[0]], caption=None, width=200, use_column_width=None, clamp=False, channels='RGB', output_format="auto")
                st.session_state.json_list[1] = img_list[st.session_state.select_wall[0]]
            except:
                b=1
        else:
            st.session_state.json_list[1] = None
# Текстура потолка
with col3:
    if path_photo is not None and st.session_state.pp is not None:
        ceiling = st.checkbox('Текстура потолка', value=False, key=None, help=None, on_change=None, args=None, kwargs=None, disabled=False, label_visibility="visible")
        if ceiling:
            ceiling_select = st.button("Выбор текстуры потолка", key = '3')
            if ceiling_select:
                    select_ceiling(img_list_1, cv2.cvtColor(st.session_state.modification_photo[0], cv2.COLOR_BGR2RGB), img_list)
            try:
                st.image(img_list_1[st.session_state.select_ceiling[0]], caption=None, width=200, use_column_width=None, clamp=False, channels='RGB', output_format="auto")
                st.session_state.json_list[2] = img_list[st.session_state.select_ceiling[0]]
            except:
                b=1
        else:
            st.session_state.json_list[2] = None
if path_photo is not None and st.session_state.pp is not None:
    st.session_state.surface = [floor, wall, ceiling]

# Выбор углов поворота текстур
col1_1, col2_1 = st.columns([3,1])
try:
    with col2_1:
        if path_photo is not None and st.session_state.pp is not None:
            floor_angel = st.slider(label = 'Угол поворота текстуры пола',min_value=0, max_value=360, value=0, step=1, format=None, key='floor_angle_key',
                    help=None, on_change=save_angel, args=None, kwargs=None, disabled=False, label_visibility= 'visible')
            wall_angel = st.slider(label = 'Угол поворота текстуры стен',min_value=0, max_value=360, value=0, step=1, format=None, key='wall_angle_key',
                    help=None, on_change=save_angel, args=None, kwargs=None, disabled=False, label_visibility= 'visible')
            ceiling_angel = st.slider(label = 'Угол поворота текстуры потолка',min_value=0, max_value=360, value=0, step=1, format=None, key='ceiling_angle_key',
                    help=None, on_change=save_angel, args=None, kwargs=None, disabled=False, label_visibility= 'visible')
            st.write(":orange[Выбор способа набора текстур]")
            st.write(":orange[Если больше 1 то набор ёлочкой]")
            st.session_state.f_h = st.number_input('Пол', min_value=0, max_value=3, value=0, step=1, format=None, key='f_h_key', help=None, on_change=None, args=None, kwargs=None, placeholder=None, disabled=False, label_visibility="visible")
            st.session_state.w_h = st.number_input('Стены', min_value=0, max_value=3, value=0, step=1, format=None, key='w_h_key', help=None, on_change=None, args=None, kwargs=None, placeholder=None, disabled=False, label_visibility="visible")
            st.session_state.c_h = st.number_input('Потолок', min_value=0, max_value=3, value=0, step=1, format=None, key='c_h_key', help=None, on_change=None, args=None, kwargs=None, placeholder=None, disabled=False, label_visibility="visible")

            if 'floor_angel' not in st.session_state:
                st.session_state.floor_angel = 0
            if 'wall_angel' not in st.session_state:
                st.session_state.wall_angel = 0
            if 'ceiling_angel' not in st.session_state:
                st.session_state.ceiling_angel = 0
            try:
                if st.session_state.floor_angel != floor_angel:
                    modificate_image('floor')
            except:
                print('no_change_angel_floor')
            try:
                if st.session_state.wall_angel != wall_angel:
                    modificate_image('floor')
            except:
                print('no_change_angel_wall')
            try:
                if st.session_state.ceiling_angel != ceiling_angel:
                    modificate_image('ceiling')
            except:
                print('no_change_angel_ceiling')

# Отображение фотографии с наложением текстур    
    with col1_1:
    #    scroll = st.slider(label = '123',min_value=0., max_value=1., value=0.5, step=0.01, format=None, key=33, 
    #            help=None, on_change=None, args=None, kwargs=None, disabled=False, label_visibility= 'hidden')
        if path_photo is not None and st.session_state.pp is not None:
            if st.session_state.modification_photo is not None:
                image_1 = st.session_state.modification_photo[0]
            else:
                image_1 = cv2.cvtColor(st.session_state.predict[1], cv2.COLOR_BGR2RGB)
            image = cv2.cvtColor(st.session_state.predict[1], cv2.COLOR_BGR2RGB)
            #st.write(st.get_option())
            image_comparison(
                            img1=image,
                            img2=image_1,
                            width = 1080,
                            starting_position = 1
                            )
            #fix_line = np.zeros((image.shape[0], 5, 3), dtype=np.uint8)
            #fix_line[:,:,1]  = 254
            #image_2 =cv2.hconcat([image[:, 0:int(image.shape[1]*scroll)], fix_line])
            #image_3 =cv2.hconcat([image_2, image_1[: , int(image_1.shape[1]*scroll+5):]])
            #img = st.image(image_3, caption=None, use_column_width='always', clamp=False, channels="RGB", output_format="auto")
            cv2.imwrite('Change_photo.jpg', cv2.cvtColor(st.session_state.modification_photo[0], cv2.COLOR_BGR2RGB))
            with open('Change_photo.jpg', "rb") as file:
                print(type(file))
                btn = st.download_button(
                        label="Скачать изображение с наложенной текстурой",
                        data=file,
                        file_name="Change_photo.png",
                        mime="image/png"
                    )
except:
    b=1