import chromadb
import gradio as gr
import matplotlib.pyplot as plt
import io
from PIL import Image
import re

# 初始化集合
client = chromadb.PersistentClient(path="chroma_db2")
output_collection = client.get_collection("data_collection")

def remove1(latex_str):
    # 定义需要移除的命令的正则表达式
    unsupported_commands = [
        r'\\ref\{([^}]+)\}',  # 匹配 \ref{...}，保留括号内容
        r'\\label\{([^}]+)\}',  # 匹配 \label{...}，保留括号内容
        r'\\Ob\(([^)]+)\)',
    ]
    latex_str=latex_str.replace("\\\\","\\")
    # 逐个移除不支持的命令
    for pattern in unsupported_commands:
        latex_str = re.sub(pattern, r'\1', latex_str)

    return latex_str

plt.rcParams.update({
    "text.usetex": True,
    "font.family": "serif",
    "axes.spines.right": False,
    "axes.spines.top": False,
    "text.latex.preamble": r"\usepackage{amsmath,amssymb}"
})

def search(query,enable_latex):
    output_results = output_collection.get(
        where={"$or": [{"tag": query}, {"label": query}]},
        limit=1
    )
    
    if not output_results["ids"]:
        return ["No matching results found",None,'']
    
    # 获取匹配的output信息
    tag= output_results["metadatas"][0]["tag"]
    label= output_results["metadatas"][0]["label"]
    latex= output_results["documents"][0]
    latex_output = f"""
    <div style="margin-bottom: 25px; padding: 15px; border: 1px solid #e0e0e0; border-radius: 8px;">
        <div style="margin: 15px 0;">
                <p style="font-size: 1.1em; margin: 8px 0;">
                    <strong>✍ Tag:</strong> {tag}
                </p>
                <p style="font-size: 1.1em; margin: 8px 0;">
                    <strong>💬 Label:</strong> {label}
                </p>
                <p style="font-size: 1.1em; margin: 8px 0; color: #e67e22;">
                    <strong>⭐ LaTeX:</strong> {latex}
                </p>
        </div>
        <div style="text-align: center; margin-top: 20px;">
            <a href="https://stacks.math.columbia.edu/tag/{tag}" 
                target="_blank" 
                style="display: inline-block;
                    background: #3498db;
                    color: white !important;
                    padding: 12px 25px;
                    border-radius: 5px;
                    text-decoration: none;
                    font-size: 1.1em;
                    transition: all 0.3s ease;
                    box-shadow: 0 2px 5px rgba(0,0,0,0.1);">
                📖 Proceed To The Docs Link
            </a>
        </div>
    </div>
    """
    answer = '<div style="padding: 10px;">'
    for l in range(10):
        kind = output_results["metadatas"][0][str(l)+'kind'].strip()
        fn = output_results["metadatas"][0][str(l)+'formal name'].strip("'")
        fs = output_results["metadatas"][0][str(l)+'formal statement']
        inn = output_results["metadatas"][0][str(l)+'informal name'].strip("'")
        ins = output_results["metadatas"][0][str(l)+'informal statement']
    
        answer += f"""
        <div style="margin-bottom: 25px; padding: 15px; border: 1px solid #e0e0e0; border-radius: 8px;">
            <!-- 放大并居中的formal name -->
            <div style="text-align: center; margin: 15px 0;">
                <h2 style="font-size: 1.8em; color: #2c3e50; 
                         margin: 0;
                          padding-bottom: 10px;
                          border-bottom: 3px solid #3498db;
                          display: inline-block;">
                    {fn}
                </h2>
            </div>

            <!-- 强调的formal statement -->
            <div style="background: #f8f9fa;
                       padding: 15px;
                       border-radius: 6px;
                       margin: 15px 0;
                       border-left: 4px solid #3498db;">
                <p style="font-size: 1.2em; line-height: 1.6; color: #2c3e50; margin: 0;font-family: 'Courier New', Consolas, Monaco, monospace;">
                    <strong style="font-size: 1.1em; color: #3498db;">📜 Formal Statement:</strong><br>
                    {fs}
                </p>
            </div>

            <!-- 其他信息 -->
            <div style="margin: 15px 0;">
                <p style="font-size: 1.1em; margin: 8px 0;">
                    <strong>⭐ Kind:</strong> {kind}
                </p>
                <p style="font-size: 1.1em; margin: 8px 0;">
                    <strong>✍ Informal Name:</strong> {inn}
                </p>
                <p style="font-size: 1.1em; margin: 8px 0;">
                    <strong>💬 Informal Statement:</strong> {ins}
                </p>
                <p style="font-size: 1.1em; margin: 8px 0; color: #e67e22;">
                    <strong>✌ Similarity Ranking:</strong> {l+1}
                </p>
            </div>

            <!-- 文档链接 -->
            <div style="text-align: center; margin-top: 20px;">
                <a href="https://leanprover-community.github.io/mathlib4_docs/find/?pattern={fn}#doc" 
                   target="_blank" 
                   style="display: inline-block;
                         background: #3498db;
                          color: white !important;
                          padding: 12px 25px;
                          border-radius: 5px;
                          text-decoration: none;
                          font-size: 1.1em;
                          transition: all 0.3s ease;
                          box-shadow: 0 2px 5px rgba(0,0,0,0.1);">
                    📖 View Full Documentation
                </a>
            </div>
        </div>
        """
    answer += '</div>'
    
    return [latex_output,answer]

with gr.Blocks(theme=gr.themes.Soft(), title="Stack-Mathlib4 Search") as demo:
    gr.Markdown("""<center><font size=8> Stack-Mathlib4 Search</font></center>""")
        
    with gr.Tab("Theorem Search"):
        label_input = gr.Textbox(
            placeholder="", 
            label="Please enter a label or a tag"
        )
        search_button = gr.Button("Search", variant="primary")
        
        information = gr.HTML(label="Information of the Theorem", elem_classes="result-box")
        answer = gr.HTML(label="Top Matching Theorems", elem_classes="result-box")

        search_button.click(
            fn=search,
            inputs=[label_input],
            outputs=[information, answer]
        )

demo.launch(share=True)