import streamlit as st
import torch
from transformers import AutoTokenizer
from semviqa.ser.qatc_model import QATCForQuestionAnswering
from semviqa.tvc.model import ClaimModelForClassification
from semviqa.ser.ser_eval import extract_evidence_tfidf_qatc
from semviqa.tvc.tvc_eval import classify_claim
import time
import pandas as pd
import os
import psutil
import gc
import threading
from queue import Queue

# Set environment variables to optimize CPU performance
os.environ["OMP_NUM_THREADS"] = str(psutil.cpu_count(logical=False))
os.environ["MKL_NUM_THREADS"] = str(psutil.cpu_count(logical=False))

# Set device globally
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Load models with caching and CPU optimization
@st.cache_resource()
def load_model(model_name, model_class, is_bc=False, device=None):
    if device is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"
    
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = model_class.from_pretrained(model_name, num_labels=3 if not is_bc else 2)
    model.eval()
    
    # CPU-specific optimizations
    if device == "cpu":
        # Use torch's quantization for CPU inference speed boost
        try:
            import torch.quantization
            # Quantize the model to INT8
            model = torch.quantization.quantize_dynamic(
                model, {torch.nn.Linear}, dtype=torch.qint8
            )
        except Exception as e:
            st.warning(f"Quantization failed, using default model: {e}")
    
    model.to(device)
    return tokenizer, model

# Pre-process text function to avoid doing it multiple times
@st.cache_data
def preprocess_text(text):
    # Add any text cleaning or normalization here
    return text.strip()

# Function to extract evidence in a separate thread for better CPU utilization
def extract_evidence_threaded(queue, claim, context, model_qatc, tokenizer_qatc, device, 
                             tfidf_threshold, length_ratio_threshold):
    start_time = time.time()
    with torch.no_grad():
        evidence = extract_evidence_tfidf_qatc(
            claim, context, model_qatc, tokenizer_qatc,
            device,
            confidence_threshold=tfidf_threshold,
            length_ratio_threshold=length_ratio_threshold
        )
    evidence_time = time.time() - start_time
    queue.put((evidence, evidence_time))

# Function to classify in a separate thread
def classify_claim_threaded(queue, claim, evidence, model, tokenizer, device):
    with torch.no_grad():
        result = classify_claim(claim, evidence, model, tokenizer, device)
    queue.put(result)

# Optimized function for evidence extraction and classification with better CPU performance
def perform_verification(claim, context, model_qatc, tokenizer_qatc, model_tc, tokenizer_tc, 
                          model_bc, tokenizer_bc, tfidf_threshold, length_ratio_threshold):
    # Use thread for evidence extraction to allow garbage collection in between
    evidence_queue = Queue()
    evidence_thread = threading.Thread(
        target=extract_evidence_threaded,
        args=(evidence_queue, claim, context, model_qatc, tokenizer_qatc, DEVICE,
              tfidf_threshold, length_ratio_threshold)
    )
    evidence_thread.start()
    evidence_thread.join()
    evidence, evidence_time = evidence_queue.get()
    
    # Explicit garbage collection after evidence extraction
    gc.collect()
    
    # Classify the claim
    verdict_start_time = time.time()
    with torch.no_grad():
        prob3class, pred_tc = classify_claim(
            claim, evidence, model_tc, tokenizer_tc, DEVICE
        )
        
        # Only run binary classifier if needed
        prob2class, pred_bc = 0, 0
        if pred_tc != 0:
            prob2class, pred_bc = classify_claim(
                claim, evidence, model_bc, tokenizer_bc, DEVICE
            )
            verdict = "SUPPORTED" if pred_bc == 0 else "REFUTED" if prob2class > prob3class else ["NEI", "SUPPORTED", "REFUTED"][pred_tc]
        else:
            verdict = "NEI"
            
    verdict_time = time.time() - verdict_start_time

    return {
        "evidence": evidence,
        "verdict": verdict,
        "evidence_time": evidence_time,
        "verdict_time": verdict_time,
        "prob3class": prob3class.item() if isinstance(prob3class, torch.Tensor) else prob3class,
        "pred_tc": pred_tc,
        "prob2class": prob2class.item() if isinstance(prob2class, torch.Tensor) else prob2class,
        "pred_bc": pred_bc
    }

# Set up page configuration
st.set_page_config(page_title="SemViQA Demo", layout="wide")

# Custom CSS: fixed header and tabs, dynamic height, result box formatting
# Custom CSS: fixed header and tabs, dynamic height, result box formatting
st.markdown(
    """
    <style>
    html, body {
        height: 100%;
        margin: 0;
        overflow: hidden;
    }
    .main-container {
        height: calc(100vh - 55px); /* Browser height - 55px */
        overflow-y: auto;
        padding: 20px;
    }
    .big-title {
        font-size: 36px !important;
        font-weight: bold;
        color: var(--primary-color);
        text-align: center;
        margin-top: 20px;
        position: sticky; /* Pin the header */
        top: 0;
        background-color: var(--background-color); /* Ensure the header covers content when scrolling */
        z-index: 100; /* Ensure it's above other content */
    }
    .sub-title {
        font-size: 20px;
        color: var(--text-color);
        text-align: center;
        margin-bottom: 20px;
    }
    .stButton>button {
        background-color: var(--primary-color);
        color: red; 
        font-size: 16px;
        width: 100%;
        border-radius: 8px;
        padding: 10px;
        border: 2px solid red;
    }
    .stTextArea textarea {
        font-size: 16px;
        min-height: 120px;
        background-color: var(--secondary-background-color);
        color: var(--text-color);
    }
    .result-box {
        background-color: var(--secondary-background-color);
        padding: 20px;
        border-radius: 10px;
        box-shadow: 0px 4px 8px rgba(0, 0, 0, 0.1);
        margin-top: 20px;
        color: var(--text-color);
    }
    .verdict {
        font-size: 24px;
        font-weight: bold;
        margin: 0;
        display: flex;
        align-items: center;
        color: var(--text-color);
    }
    .verdict-icon {
        margin-right: 10px;
    }
    /* Fix the tabs at the top */
    div[data-baseweb="tab-list"] {
        position: sticky;
        top: 55px; /* Height of the header */
        background-color: var(--background-color);
        z-index: 99;
    }
    /* Fix the sidebar content */
    .stSidebar .sidebar-content {
        background-color: var(--secondary-background-color);
        padding: 20px;
        border-radius: 10px;
        position: sticky;
        top: 55px; /* Height of the header */
        height: calc(100vh - 75px); /* Adjust height to fit within the viewport */
        overflow-y: auto; /* Enable scrolling within the sidebar if content is too long */
    }
    .stSidebar .st-expander {
        background-color: var(--background-color);
        border-radius: 10px;
        padding: 10px;
        margin-bottom: 10px;
    }
    .stSidebar .stSlider {
        margin-bottom: 20px;
    }   
    .stSidebar .stSelectbox {
        margin-bottom: 20px;
    }
    .stSidebar .stCheckbox {
        margin-bottom: 20px;
    }

    .code-block {
        background-color: #f5f5f5; /* Màu nền */
        border: 1px solid #ccc; /* Viền */
        padding: 10px; /* Khoảng cách bên trong */
        border-radius: 5px; /* Bo tròn góc */
        font-family: monospace; /* Font chữ đơn cách */
        white-space: pre-wrap; /* Giữ nguyên định dạng xuống dòng */
        color: #333; /* Màu chữ */
    }
    </style>
    """,
    unsafe_allow_html=True,
)

# Container for the whole content with dynamic height
with st.container():
    st.markdown("<p class='big-title'>SemViQA: A Semantic Question Answering System for Vietnamese Information Fact-Checking</p>", unsafe_allow_html=True)
    st.markdown("""
        <div style="text-align: center; margin-bottom: 20px;">
            <p>
                <a href="https://github.com/DAVID-NGUYEN-S16">Nam V. Nguyen</a>*, 
                <a href="https://github.com/xndien2004">Dien X. Tran</a>*, 
                Thanh T. Tran, 
                Anh T. Hoang, 
                Tai V. Duong, 
                Di T. Le, 
                Phuc-Lu Le
            </p> 
        </div>
    """, unsafe_allow_html=True)

    # Sidebar: Global Settings
    with st.sidebar.expander("⚙️ Settings", expanded=True):
        tfidf_threshold = st.slider("TF-IDF Threshold", 0.0, 1.0, 0.5, 0.01)
        length_ratio_threshold = st.slider("Length Ratio Threshold", 0.1, 1.0, 0.5, 0.01)
        qatc_model_name = st.selectbox("QATC Model", [
            "SemViQA/qatc-infoxlm-viwikifc",
            "SemViQA/qatc-infoxlm-isedsc01",
            "SemViQA/qatc-vimrc-viwikifc",
            "SemViQA/qatc-vimrc-isedsc01"
        ])
        bc_model_name = st.selectbox("Binary Classification Model", [
            "SemViQA/bc-xlmr-viwikifc",
            "SemViQA/bc-xlmr-isedsc01",
            "SemViQA/bc-infoxlm-viwikifc",
            "SemViQA/bc-infoxlm-isedsc01",
            "SemViQA/bc-erniem-viwikifc",
            "SemViQA/bc-erniem-isedsc01"
        ])
        tc_model_name = st.selectbox("3-Class Classification Model", [
            "SemViQA/tc-xlmr-viwikifc",
            "SemViQA/tc-xlmr-isedsc01",
            "SemViQA/tc-infoxlm-viwikifc",
            "SemViQA/tc-infoxlm-isedsc01",
            "SemViQA/tc-erniem-viwikifc",
            "SemViQA/tc-erniem-isedsc01"
        ])
        show_details = st.checkbox("Show Probability Details", value=False)
        
        # Add CPU optimization settings
        st.subheader("CPU Performance Settings")
        num_threads = st.slider("Number of CPU Threads", 1, psutil.cpu_count(), 
                               psutil.cpu_count(logical=False))
        os.environ["OMP_NUM_THREADS"] = str(num_threads)
        os.environ["MKL_NUM_THREADS"] = str(num_threads)

    # Load models once and keep them in memory
    tokenizer_qatc, model_qatc = load_model(qatc_model_name, QATCForQuestionAnswering, device=DEVICE)
    tokenizer_bc, model_bc = load_model(bc_model_name, ClaimModelForClassification, is_bc=True, device=DEVICE)
    tokenizer_tc, model_tc = load_model(tc_model_name, ClaimModelForClassification, device=DEVICE)
    st.session_state.models_loaded = True

    # Store verification history
    if 'history' not in st.session_state:
        st.session_state.history = []
    if 'latest_result' not in st.session_state:
        st.session_state.latest_result = None

    # Icons for results
    verdict_icons = {
        "SUPPORTED": "✅",
        "REFUTED": "❌",
        "NEI": "⚠️"
    }

    # Tabs: Verify, History
    tabs = st.tabs(["Verify", "History"])

    # --- Tab Verify ---
    with tabs[0]:
        st.subheader("Verify a Claim")
        # 2-column layout: input on the left, results on the right
        col_input, col_result = st.columns([2, 1])

        with col_input:
            claim = st.text_area("Enter Claim", "Chiến tranh với Campuchia đã kết thúc trước khi Việt Nam thống nhất.")
            context = st.text_area("Enter Context", "Sau khi thống nhất, Việt Nam tiếp tục gặp khó khăn do sự sụp đổ và tan rã của đồng minh Liên Xô cùng Khối phía Đông, các lệnh cấm vận của Hoa Kỳ, chiến tranh với Campuchia, biên giới giáp Trung Quốc và hậu quả của chính sách bao cấp sau nhiều năm áp dụng. Năm 1986, Đảng Cộng sản ban hành cải cách đổi mới, tạo điều kiện hình thành kinh tế thị trường và hội nhập sâu rộng. Cải cách đổi mới kết hợp cùng quy mô dân số lớn đưa Việt Nam trở thành một trong những nước đang phát triển có tốc độ tăng trưởng thuộc nhóm nhanh nhất thế giới, được coi là Hổ mới châu Á dù cho vẫn gặp phải những thách thức như tham nhũng, tội phạm gia tăng, ô nhiễm môi trường và phúc lợi xã hội chưa đầy đủ. Ngoài ra, giới bất đồng chính kiến, chính phủ một số nước phương Tây và các tổ chức theo dõi nhân quyền có quan điểm chỉ trích hồ sơ nhân quyền của Việt Nam liên quan đến các vấn đề tôn giáo, kiểm duyệt truyền thông, hạn chế hoạt động ủng hộ nhân quyền cùng các quyền tự do dân sự.")
            verify_button = st.button("Verify", key="verify_button")

        with col_result:
            st.markdown("<h3>Verification Result</h3>", unsafe_allow_html=True)
            if verify_button:
                # Preprocess texts to improve tokenization speed
                preprocessed_claim = preprocess_text(claim)
                preprocessed_context = preprocess_text(context)
                
                # Placeholder for displaying result/loading
                with st.spinner("Verifying..."): 
                    start_time = time.time()
                    
                    # Clear memory before verification
                    gc.collect()
                    
                    # Use the optimized verification function
                    result = perform_verification(
                        preprocessed_claim, preprocessed_context, 
                        model_qatc, tokenizer_qatc,
                        model_tc, tokenizer_tc,
                        model_bc, tokenizer_bc,
                        tfidf_threshold, length_ratio_threshold
                    )
                    
                    total_time = time.time() - start_time
                    
                    # Format details if needed
                    details = ""
                    if show_details:
                        details = f"""
                            3-Class Probability: {result['prob3class']:.2f}
                            3-Class Predicted Label: {['NEI', 'SUPPORTED', 'REFUTED'][result['pred_tc']]}
                            2-Class Probability: {result['prob2class']:.2f}
                            2-Class Predicted Label: {['SUPPORTED', 'REFUTED'][result['pred_bc']] if isinstance(result['pred_bc'], int) and result['pred_tc'] != 0 else 'Not used'}
                        """
                    
                    st.session_state.latest_result = {
                        "claim": claim,
                        "evidence": result['evidence'],
                        "verdict": result['verdict'],
                        "evidence_time": result['evidence_time'],
                        "verdict_time": result['verdict_time'],
                        "total_time": total_time,
                        "details": details,
                        "qatc_model": qatc_model_name,
                        "bc_model": bc_model_name,
                        "tc_model": tc_model_name
                    }

                    # Add new result to history
                    st.session_state.history.append(st.session_state.latest_result)

                    # Clear memory after processing
                    gc.collect()

                    # Display the result after verification
                    res = st.session_state.latest_result
                    st.markdown(f"""
                        <div class='result-box'>
                            <p><strong>Claim:</strong> {res['claim']}</p>
                            <p><strong>Evidence:</strong> {res['evidence']}</p>
                            <p class='verdict'><span class='verdict-icon'>{verdict_icons.get(res['verdict'], '')}</span>{res['verdict']}</p>
                            <p><strong>Evidence Inference Time:</strong> {res['evidence_time']:.2f} seconds</p>
                            <p><strong>Verdict Inference Time:</strong> {res['verdict_time']:.2f} seconds</p>
                            <p><strong>Total Execution Time:</strong> {res['total_time']:.2f} seconds</p>
                            {f"<div class='code-block'><pre>{res['details']}</pre></div>" if show_details else ""}
                        </div>
                    """, unsafe_allow_html=True)
                    
                    # Download Verification Result Feature
                    result_text = f"Claim: {res['claim']}\nEvidence: {res['evidence']}\nVerdict: {res['verdict']}\nDetails: {res['details']}"
                    st.download_button("Download Result", data=result_text, file_name="verification_result.txt", mime="text/plain")
            else:
                st.info("No verification result yet.")

    # --- Tab History ---
    with tabs[1]:
        st.subheader("Verification History")
        if st.session_state.history:
            # Convert history to DataFrame for easy download
            history_df = pd.DataFrame(st.session_state.history)
            st.download_button(
                label="Download Full History",
                data=history_df.to_csv(index=False).encode('utf-8'),
                file_name="verification_history.csv",
                mime="text/csv",
            )
            for idx, record in enumerate(reversed(st.session_state.history), 1):
                st.markdown(f"**{idx}. Claim:** {record['claim']}  \n**Result:** {verdict_icons.get(record['verdict'], '')} {record['verdict']}")
        else:
            st.write("No verification history yet.")