import streamlit as st
import leafmap.foliumap as leafmap
import ee
import os
import datetime
from tools import split_tif, analyze_stydy_area, maketrain, modelTrain, export_classification, Make_result, \
    analyze_result
import pandas as pd

os.environ['PROJ_LIB'] = r'B:\anaconda3\Lib\site-packages\pyproj\proj_dir\share\proj'
from utils import *
from glob import glob
from datetime import datetime

st.set_page_config(page_title='TSSCD土地覆盖变化监测系统', layout="wide")


def get_file_path(data):
    import tempfile
    import os
    import uuid
    _, file_extension = os.path.splitext(data.name)
    file_id = str(uuid.uuid4())
    file_path = os.path.join(tempfile.gettempdir(), f"{file_id}{file_extension}")
    return file_path


def save_uploaded_file(uploadedfile):
    with open(os.path.join("tempDir", uploadedfile.name), "wb") as f:
        f.write(uploadedfile.getbuffer())


def disable():
    st.session_state.disabled = True


# Initialize disabled for form_submit_button to False
if "disabled" not in st.session_state:
    st.session_state.disabled = False

if "analyze1_button" not in st.session_state:
    st.session_state.analyze1_button = False

if "sourcemap" not in st.session_state:
    st.session_state.sourcemap = True

temp_dir = r'tempDir'  # 临时文件存放区


def update_map(Map, Data):
    file_list = glob(os.path.join(Data, '数据', '*.tif'))
    start_tif_path = file_list[0]
    end_tif_path = file_list[-1]

    s, geo, pro = readGeoTIFF(start_tif_path)
    e, _, _ = readGeoTIFF(end_tif_path)

    CreateGeoTiff(os.path.join(temp_dir, 'start.tif'), s[:3],
                  geo, pro)
    CreateGeoTiff(os.path.join(temp_dir, 'end.tif'), e[:3],
                  geo, pro)
    Map.add_raster(os.path.join(temp_dir, 'start.tif'),
                   bands=[3, 2, 1], vmin=0, vmax=4000,
                   layer_name=os.path.basename(start_tif_path))
    Map.add_raster(os.path.join(temp_dir, 'end.tif'),
                   bands=[3, 2, 1], vmin=0, vmax=4000,
                   layer_name=os.path.basename(end_tif_path))
    Map.add_raster(glob(os.path.join(Data, '产品', '*2020*.tif'))[0], layer_name='ESA 2020')
    Map.add_raster(glob(os.path.join(Data, '产品', '*2021*.tif'))[0], layer_name='ESA 2021')
    Map.add_legend(title='ESA Land Cover', builtin_legend='ESA_WorldCover')
    # Map.to_streamlit()


def main(Map):
    with st.sidebar:
        st.title("土地覆盖监测系统")
        Data = st.text_input("请输入工程文件夹: ", 'D:\雄安新区')
        data_path = os.path.join(Data, '数据')
        datelist = []
        for file in os.listdir(data_path):
            if file.endswith('tif'):
                year = int(file[:4])
                month = int(file[4:6])
                datelist.append([year, month])
        submitted = st.button("确 定", on_click=disable, disabled=st.session_state.disabled)
    if submitted:
        if st.session_state.sourcemap:
            st.session_state.sourcemap = False
            update_map(Map, Data)

    with st.sidebar:
        with st.form("submit_split_zone"):
            st.header('Step1：数据切片')
            split_value = st.slider("请选择分块尺寸", min_value=100, max_value=1000, value=300, step=100)
            split_tif_button = st.form_submit_button("数据切片")  # 数据切片
            if split_tif_button:
                print('开始切片!')
                split_tif.gen(Data, split_value)
        with st.form("dataset"):
            st.header('Step2：变化频次定义')
            change1 = st.slider('不变样本占比', min_value=0, max_value=100, value=60, step=5)
            change2 = st.slider('变化一次占比', min_value=0, max_value=100, value=20, step=5)
            change3 = st.slider('变化两次占比', min_value=0, max_value=100, value=15, step=5)
            change4 = st.slider('变化三次占比', min_value=0, max_value=100, value=5, step=5)
            st.form_submit_button('变化次数比例计算')
            weight = np.array([change1, change2, change3, change4])
            weight_p_change_count = weight / np.sum(weight)
            p = st.empty()
            p.text(np.round(weight_p_change_count, 2))
        with st.form("change_type"):
            st.header('Step3：土地覆盖类型占比定义')
            water = st.slider('水体', min_value=0, max_value=100, value=60, step=5)
            forest = st.slider('林地', min_value=0, max_value=100, value=20, step=5)
            grass = st.slider('草地', min_value=0, max_value=100, value=15, step=5)
            bare = st.slider('裸地', min_value=0, max_value=100, value=5, step=5)
            impervious = st.slider('不透水面', min_value=0, max_value=100, value=15, step=5)
            crop = st.slider('耕地', min_value=0, max_value=100, value=5, step=5)

            st.form_submit_button('土地覆盖比例计算')
            weight = np.array([water, forest, grass, bare, impervious, crop])
            weight_p_lc = weight / np.sum(weight)
            p = st.empty()
            p.text(np.round(weight_p_lc, 2))

        with st.form("change_pro"):
            st.header('Step4：土地覆盖类型占比定义')
            # 定义行和列的名称
            rows = ['水体', '林地', '草地', '裸地', '不透水面', '耕地']
            columns = ['水体', '林地', '草地', '裸地', '不透水面', '耕地']

            # 创建一个空的DataFrame
            data = pd.DataFrame(index=range(6), columns=columns)
            data.insert(0, '土地覆盖类型', rows)
            transfer_matrix = np.array([[0, 0.2, 0.2, 0.2, 0.2, 0.2],
                                        [0.1, 0, 0.35, 0.2, 0.1, 0.25],
                                        [0.2, 0.2, 0, 0.35, 0.25, 0],
                                        [0.15, 0.05, 0.3, 0, 0.4, 0.1],
                                        [0.1, 0.1, 0.1, 0.6, 0, 0.1],
                                        [0.1, 0.3, 0, 0.3, 0.3, 0]]) * 100
            # 将NumPy矩阵赋给DataFrame
            data.iloc[:, 1:] = transfer_matrix
            # 在Streamlit应用中显示表格
            table = st.data_editor(data, use_container_width=True, disabled=['土地覆盖类型'], hide_index=True)

            temp = st.empty()
            temp.dataframe(data, use_container_width=True, hide_index=True)
            change_type_pro = st.form_submit_button('转移矩阵计算')
            if change_type_pro:
                change_type = table.values[:, 1:]
                change_type = (change_type * 100 / np.tile(np.sum(change_type, axis=1), 6).reshape(6, 6).T).astype(
                    'int')
                data.iloc[:, 1:] = change_type.astype('int')
                table = temp.dataframe(data, use_container_width=True, hide_index=True)

        with st.form("create_sample"):
            st.header('Step5：训练样本生成')
            maketrain_button = st.form_submit_button('开始生成样本')
            if maketrain_button:
                maketrain.choose_points(Data)
                maketrain.get_series(Data, split_value)
                maketrain.makedataset(Data, weight_p_lc, weight_p_change_count)
        with st.form("train_model"):
            st.header('Step6：模型参数设置')
            batch_size = st.slider('batch_size', min_value=8, max_value=512, value=128, step=1)
            c1 = st.slider('c1', min_value=0, max_value=2000, value=128, step=5)
            c2 = st.slider('c2', min_value=0, max_value=2000, value=256, step=5)
            c3 = st.slider('c3', min_value=0, max_value=2000, value=512, step=5)
            c4 = st.slider('c4', min_value=0, max_value=2000, value=1024, step=5)
            train = st.form_submit_button('开始训练！')
            if train:
                model_config = [c1, c2, c3, c4]
                modelTrain.train_moodel(Data, batch_size, model_config)
        with st.form('predict_true'):
            st.header('Step7：数据推理')
            pre_button = st.form_submit_button('开始推理')
            if pre_button:
                export_classification.predict_true_tif(Data, [c1, c2, c3, c4])
                Make_result.make_result(Data, split_value)
            analyze1_button = st.form_submit_button('结果分析')
            if analyze1_button:
                st.session_state.analyze1_button = True
                analyze_result.get_trend(Data)
                out_path = analyze_result.get_gif(Data)
                st.image(out_path)
    if st.session_state.analyze1_button:

        start_time = st.slider(
            "日期",
            min_value=datetime(datelist[0][0], datelist[0][1], 1),
            max_value=datetime(datelist[-1][0], datelist[-1][1], 1),
            value=datetime(datelist[0][0], datelist[0][1], 1),
            format="YYYY年MM月")
        choose_date = [int(start_time.year), int(start_time.month)]
        index = datelist.index(choose_date)
        data_path = glob(os.path.join(Data, r'数据\*.tif'))[index]
        tif_path = glob(os.path.join(Data, r'split_zone\OutputResult\逐时相的结果\tif_rgb\*.tif'))[index]
        if index == 0:
            index_change = 0
        else:
            index_change = index - 1
        change_tif_path = glob(os.path.join(Data, rf'split_zone\OutputResult\逐时相的结果\change_tif\*.tif'))[
            index_change]
        arr, geo, pro = readGeoTIFF(data_path)
        CreateGeoTiff(os.path.join(temp_dir, 'temp.tif'), arr[:3], geo, pro)

        Map.add_raster(os.path.join(temp_dir, 'temp.tif'), bands=[3, 2, 1], vmin=0, vmax=4000,
                       layer_name=f'{datelist[index][0]}年{datelist[index][1]}月Sentinel-2影像')
        Map.add_raster(change_tif_path, palette='jet',
                       layer_name=f'{datelist[index][0]}年{datelist[index][1]}月累积变化图')
        Map.add_raster(tif_path, bands=[1, 2, 3], vmin=0, vmax=255,
                       layer_name=f'{datelist[index][0]}年{datelist[index][1]}月土地覆盖制图')

        labels = ["Water", "Woodland", "Grassland", "Bare soil", "Impervious surface", "Cropland"]

        color_list = [(88, 184, 255), (25, 70, 31), (138, 208, 27), (222, 168, 128), (212, 67, 56), (255, 214, 156)]

        Map.add_legend(title="Legend", labels=labels, colors=color_list)
        Map.add_colorbar(colors=Create_color_list(48), vim=0, vmax=48)

    Map.to_streamlit()


def rgb_to_hex(rgb):
    r, g, b = rgb
    return '#' + ('{:02X}' * 3).format(r, g, b)


def Create_color_list(num, color_prompt='jet'):
    # 创建渐变色带
    gradient = np.linspace(0, 1, num)
    gradient_colors = [plt.colormaps.get_cmap(color_prompt)(g) for g in gradient]
    color_list = [rgb_to_hex(i) for i in (np.array(gradient_colors)[:, :-1] * 255).astype('int')]
    return color_list


def app():
    Map1 = leafmap.Map()
    main(Map1)


app()