import streamlit as st
import json
import os

# 设置JSON文件的目录
JSON_DIR = 'table_result'

# 加载指定的JSON文件
def load_json(file_name):
    file_path = os.path.join(JSON_DIR, file_name)
    with open(file_path, "r", encoding="utf-8") as f:
        return json.load(f)

# 显示单个字典的信息
def display_dict(data):
    st.write("### 文件信息")
    st.write(f"**Path:** {data['path']}")
    st.write(f"**Table ID:** {data['table_id']}")
    st.write(f"**Section:** {data['section']}")
    
    st.write("### Table")
    # 添加滚动容器
    table_container = f"""
    <div style="max-height: 400px; overflow-y: auto; border: 1px solid #ddd; padding: 10px;">
        {data['table_html']}
    </div>
    """
    st.markdown(table_container, unsafe_allow_html=True)
    
    st.write("### Context")
    # 拼接 all_context 并高亮 target_context_ids 的句子
    all_context = data["all_context"]
    highlighted_context = ""
    for idx, sentence in enumerate(all_context):
        if idx in data["perturb_sentence_id"]:
            highlighted_context += f"<span style='color:red;'>{sentence}</span> "
        elif idx in data["target_context_ids"]:
            highlighted_context += f"**{sentence}** "
        else:
            highlighted_context += sentence + " "
    st.markdown(highlighted_context, unsafe_allow_html=True)

    st.write("### Selected Paragraphs")
    for paragraph in data["selected_paragraphs"]:
        st.write(paragraph)
    st.write("### perturb_sentence_id")
    st.write(data["perturb_sentence_id"][0])    
    st.write(data["perturb_sentence_id"][1]) 

    st.write("### Output")
    st.write("**Perturbed Statement:**")
    st.write(data["output"]["perturbed_statement"])
    st.write("**Perturbed Explanation:**")
    st.write(data["output"]["perturbed_explanation"])


# 主程序
def main():
    st.title("Perturb For Table")

    # 获取文件列表
    file_list = [f for f in os.listdir(JSON_DIR) if f.endswith(".json")]

    # 搜索框选择文件
    selected_file = st.selectbox("选择一个文件", file_list)

    # 加载并展示JSON内容
    if selected_file:
        st.write(f"当前选择文件: **{selected_file}**")
        json_data = load_json(selected_file)

        # 展示每个字典的信息
        for idx, single_dict in enumerate(json_data):
            st.write(f"## 数据 {idx + 1}")
            display_dict(single_dict)

if __name__ == "__main__":
    main()