import streamlit as st
import requests
import json

# 创建 Streamlit 页面布局
st.title('RWKV6 Query/Context相关性测试')

# 初始化 session state 中的 query 和 contexts 数组
if 'query' not in st.session_state:
    st.session_state['query'] = ''
if 'contexts' not in st.session_state:
    st.session_state['contexts'] = []

# 输入框用于输入 Query
query = st.text_input('Query', value=st.session_state['query'])

# 添加按钮，按下后增加一个输入框到 contexts 数组中
if st.button('Add'):
    st.session_state['contexts'].append('')

# 显示所有的 contexts 输入框
for i, context in enumerate(st.session_state['contexts']):
    st.session_state['contexts'][i] = st.text_input(f'Context {i+1}', value=context, key=f'context_{i}')

# 发送 HTTP 请求的函数
def send_post_request(query, contexts, url):
    headers = {
        'accept': 'application/json',
        'Content-Type': 'application/json'
    }
    data = {
        'input_text': query,
        'compared_texts': contexts
    }
    response = requests.post(url, headers=headers, data=json.dumps(data))
    return response.json()

import pandas as pd

def process_and_display_response(query, contexts, url, score_key='similarities'):
    st.write('Query:', query)
    
    # 发送 POST 请求
    response = send_post_request(query, contexts, url)
    
    # 显示返回的内容
    if response:
        scores = response.get(score_key, [])
        elapsed_time = response.get('elapsed_time', '未知')
        st.write(f'计算花费时间为：{elapsed_time}秒')
        
        # 创建一个包含上下文和分数的数据框
        df = pd.DataFrame({
            'Context': contexts,
            'Score': scores
        })
        
        # 找出最高分数的索引
        max_score_index = df['Score'].idxmax()
        
        # 创建 HTML 表格
        table_html = "<table style='width:100%; border-collapse: collapse;'>"
        table_html += "<tr><th style='border: 1px solid black; padding: 8px;'>Context</th><th style='border: 1px solid black; padding: 8px;'>Score</th></tr>"
        
        for index, row in df.iterrows():
            bg_color = "yellow" if index == max_score_index else "white"
            table_html += f"<tr style='background-color: {bg_color};'>"
            table_html += f"<td style='border: 1px solid black; padding: 8px; word-wrap: break-word; max-width: 500px;'>{row['Context']}</td>"
            table_html += f"<td style='border: 1px solid black; padding: 8px;'>{row['Score']:.4f}</td>"
            table_html += "</tr>"
        
        table_html += "</table>"
        
        # 使用 st.write 显示 HTML
        st.write(table_html, unsafe_allow_html=True) 
# Commit 按钮
if st.button('Commit'):
    process_and_display_response(
        query, 
        st.session_state['contexts'], 
        'http://47.113.113.230:6969/compute_sentence_similarities',
        'similarities'
    )

# Rerank 按钮
if st.button('Rerank'):
    process_and_display_response(
        query, 
        st.session_state['contexts'], 
        'http://47.113.113.230:6969/compute_scores',
        'similarities'
    )
    
# Clear 按钮
if st.button('Clear'):
    st.session_state['query'] = ''
    st.session_state['contexts'] = []
    st.rerun()

# 示例数据
demos = [
    {
        'query': '小米汽车价格区间是多少？',
        'contexts': [
            '小米汽车首款作品Xiaomi SU7 Logo三年打磨，满怀敬畏之心十倍投入，自研关键技术融汇优雅设计、澎湃性能与生态科技C 级豪华科技轿车人车合一，我心澎湃',
            '遵循「符合直觉」的美学设计理念，造就 Xiaomi SU7 经典的流畅车身线条。',
            '细节设计与整体外观相得益彰。水滴形大灯、全车低风阻设计，配合四种轮毂设计，与整车的优雅外观浑然一体。',
            '3000mm 超长轴距 带来宽敞的车内乘坐空间充裕的头部和腿部空间 前后排都有自在的乘坐体验',
            '座椅采用多层三明治结构坐垫 针对性优化腰部和脊椎承托 Nappa 真皮包裹，触感柔软舒适',
            '小米工况算法，在 CLTC 工况的基础上，针对气温差异和人的驾驶习惯差异，加入高温工况、低温工况、激烈驾驶工况，续航预测更准确。',
            '小米汽车统一零售价SU7：¥ 215,900 SU7 Pro：¥ 245,900 SU7 Max：¥ 299,900'
        ]
    },
    {
        'query': '成年人右下腹很痛怎么办？',
        'contexts': [
            '肚子疼，学名腹痛，是临床常见的症状。绝大部分人都会有过腹痛的经历，感觉忍忍就过去了，其实没有那么简单。腹痛其实是人体的自我保护反应，一般而言，无法及时引起腹痛示警的只有肿瘤。从常见的消化道疾病，到急性心肌梗死、妇科疾病或者一些全身性疾病，都可能引起腹痛',
            '上腹部疼痛 可以定位于胃、胆囊或者小肠上段，可能发生的疾病有烧心、胃炎、胆石症以及消化性溃疡等。',
            '中腹部疼痛 可以定位于小肠或者大肠，可能发生的疾病有阑尾炎、结肠炎、肠胀气、肠胃炎、肠易激综合征以及由精神刺激引起的“胃痛”等。',
            '下腹部疼痛 可定位于结肠下段或者盆腔器官，可能发生的疾病有阑尾炎晚期、结肠炎、肠胀气、肠易激综合征及一些膀胱疾病，女性还可发生子宫、输卵管或者卵巢问题。',
            '分不出来，感觉哪里都疼 可能是肠梗阻（呕吐腹胀，不排便不排气），阑尾炎初期（转移性疼痛）。也有可能病因不在腹部，比方说急性心肌梗死、脊柱增生性骨关节炎、肿瘤等，只有去医院才能查出来。',
            '肚子疼几乎每个人都经历过，有些疼痛忍一忍就过去了，那么什么样的腹痛，需要去医院呢？ 气温变化、大量饮酒、摄入过多油腻或不洁食物、外伤、过度劳累、睡眠不足等因素都可能诱发腹痛， 多数情况并不严重，部分可自行缓解。 若腹痛严重、不能缓解或反复出现，应及时前往医院就诊。 尤其是在剧烈腹痛的同时伴有发热、黑便或血便、无法控制的呕吐和呕血、无尿、晕厥或昏迷，或其他部位的疼痛延续到腹部，需紧急就诊。如果不知道该去哪个科室，可以直接挂急诊。',
            '切记，这篇文章并非要教你如何自己给自己看病，这样的事情还是交给专业的人来做；而是想要告诉你，因为腹痛就医，怎么样才能准确地向医生描述自己的症状，方便医生诊断。良好的医患沟通能减小误诊风险，患者也能早些康复。 个别人抱着考验医生的心态，故意隐瞒病症，只会害人害己。难倒了医生，自己也受罪，何必呢？'
        ]
    },
    {
        'query':'现在几点了？',
        'contexts':[
            '现在下雨了',
            '现在是晚上',
            '现在下午8点15分',
            '我去踢球了',
            '现在是夏天',
        ]
    },
]

# 显示示例数据
st.write('示例数据:')
for demo in demos:
    if st.button(f"Apply: {demo['query']}"):
        st.session_state['query'] = demo['query']
        st.session_state['contexts'] = demo['contexts']
        st.rerun()
    st.write('Query:', demo['query'])
    st.write('Contexts:', demo['contexts'])