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 io

# Load models with caching
@st.cache_resource()
def load_model(model_name, model_class, is_bc=False):
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = model_class.from_pretrained(model_name, num_labels=3 if not is_bc else 2)
    model.eval()
    return tokenizer, model

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

# Custom CSS: fixed navigation, header, and adjusted height
st.markdown("""
    <style>
        html, body {
            height: 100%;
            margin: 0;
            overflow: hidden;
        }
        .main-container {
            height: calc(100vh - 55px);
            overflow-y: auto;
            padding: 20px;
        }
        .big-title { 
            font-size: 36px; 
            font-weight: bold; 
            color: #4A90E2; 
            text-align: center; 
            margin-top: 20px;
            position: sticky;
            top: 0;
            background-color: white;
            z-index: 100;
            padding: 10px 0;
        }
        .sub-title { 
            font-size: 20px; 
            color: #666; 
            text-align: center; 
            margin-bottom: 20px;
            position: sticky;
            top: 56px;
            background-color: white;
            z-index: 100;
            padding-bottom: 10px;
        }
        .stButton>button { 
            background-color: #4CAF50; 
            color: white; 
            font-size: 16px; 
            width: 100%; 
            border-radius: 8px; 
            padding: 10px; 
        }
        .stTextArea textarea { 
            font-size: 16px;
            min-height: 120px;
        }
        .result-box { 
            background-color: #f9f9f9; 
            padding: 20px; 
            border-radius: 10px; 
            box-shadow: 0px 4px 8px rgba(0, 0, 0, 0.1); 
            margin-top: 20px;
        }
        .verdict { 
            font-size: 24px; 
            font-weight: bold; 
            margin: 0;
            display: flex;
            align-items: center;
        }
        .verdict-icon { 
            margin-right: 10px;
        }
        /* Fixed sidebar */
        .css-1d391kg, .css-1cypcdb {
            position: sticky;
            top: 0;
            height: calc(100vh - 55px);
            overflow-y: auto;
        }
        /* Tab content area */
        .stTabs [data-baseweb="tab-panel"] {
            height: calc(100vh - 150px);
            overflow-y: auto;
        }
        /* Loading animation */
        .loading-animation {
            text-align: center;
            padding: 20px;
        }
        .loading-animation .dot {
            display: inline-block;
            width: 12px;
            height: 12px;
            border-radius: 50%;
            background-color: #4A90E2;
            margin: 0 5px;
            animation: pulse 1.4s infinite ease-in-out;
        }
        .loading-animation .dot:nth-child(2) {
            animation-delay: 0.2s;
        }
        .loading-animation .dot:nth-child(3) {
            animation-delay: 0.4s;
        }
        @keyframes pulse {
            0%, 100% { transform: scale(0.8); opacity: 0.5; }
            50% { transform: scale(1.2); opacity: 1; }
        }
    </style>
""", unsafe_allow_html=True)

# Main container with fixed height
with st.container():
    st.markdown("<p class='big-title'>SemViQA: Semantic QA Information Verification System for Vietnamese</p>", unsafe_allow_html=True)
    st.markdown("<p class='sub-title'>Enter information to verify and context to check its accuracy</p>", 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)

    # 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
    if 'is_verifying' not in st.session_state:
        st.session_state.is_verifying = False

    # Load selected models
    tokenizer_qatc, model_qatc = load_model(qatc_model_name, QATCForQuestionAnswering)
    tokenizer_bc, model_bc = load_model(bc_model_name, ClaimModelForClassification, is_bc=True)
    tokenizer_tc, model_tc = load_model(tc_model_name, ClaimModelForClassification)

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

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

    # --- Verify Tab ---
    with tabs[0]:
        st.subheader("Verify Information")
        # Use 2-column layout: inputs on left, results on right
        col_input, col_result = st.columns([2, 1])
        
        with col_input:
            claim = st.text_area("Enter Claim", "Vietnam is a country in Southeast Asia.")
            context = st.text_area("Enter Context", "Vietnam is a country located in Southeast Asia, covering an area of over 331,000 km² with a population of more than 98 million people.")
            
            def start_verification():
                st.session_state.is_verifying = True
                st.experimental_rerun()
                
            if st.button("Verify", key="verify_button", on_click=start_verification):
                pass
        
        # Display results in right column
        with col_result:
            st.markdown("<h3>Verification Results</h3>", unsafe_allow_html=True)
            
            if st.session_state.is_verifying:
                # Show loading animation
                st.markdown("""
                    <div class="result-box">
                        <p><strong>Processing verification...</strong></p>
                        <div class="loading-animation">
                            <span class="dot"></span>
                            <span class="dot"></span>
                            <span class="dot"></span>
                        </div>
                        <p>1. Extracting evidence...</p>
                        <p>2. Running binary classification...</p>
                        <p>3. Running 3-class classification...</p>
                        <p>4. Determining final verdict...</p>
                    </div>
                """, unsafe_allow_html=True)
                
                # Perform actual verification
                with torch.no_grad():
                    # Extract evidence and classify information
                    evidence = extract_evidence_tfidf_qatc(
                        claim, context, model_qatc, tokenizer_qatc,
                        "cuda" if torch.cuda.is_available() else "cpu",
                        confidence_threshold=tfidf_threshold,
                        length_ratio_threshold=length_ratio_threshold
                    )
                    verdict = "NEI"
                    details = ""
                    prob3class, pred_tc = classify_claim(
                        claim, evidence, model_tc, tokenizer_tc,
                        "cuda" if torch.cuda.is_available() else "cpu"
                    )
                    if pred_tc != 0:
                        prob2class, pred_bc = classify_claim(
                            claim, evidence, model_bc, tokenizer_bc,
                            "cuda" if torch.cuda.is_available() else "cpu"
                        )
                        if pred_bc == 0:
                            verdict = "SUPPORTED"
                        elif prob2class > prob3class:
                            verdict = "REFUTED"
                        else:
                            verdict = ["NEI", "SUPPORTED", "REFUTED"][pred_tc]
                        if show_details:
                            details = f"<p><strong>3-Class Probability:</strong> {prob3class.item():.2f} - <strong>2-Class Probability:</strong> {prob2class.item():.2f}</p>"
                    
                    # Save verification history and latest result
                    st.session_state.history.append({
                        "claim": claim,
                        "evidence": evidence,
                        "verdict": verdict
                    })
                    st.session_state.latest_result = {
                        "claim": claim,
                        "evidence": evidence,
                        "verdict": verdict,
                        "details": details
                    }
                    
                    if torch.cuda.is_available():
                        torch.cuda.empty_cache()
                
                # Turn off verification flag
                st.session_state.is_verifying = False
                st.experimental_rerun()
                
            elif st.session_state.latest_result is not None:
                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>
                        {res['details']}
                    </div>
                """, unsafe_allow_html=True)
                # Download verification result
                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 results yet.")

    # --- History Tab ---
    with tabs[1]:
        st.subheader("Verification History")
        if st.session_state.history:
            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.")

    # --- About Tab ---
    with tabs[2]:
        st.subheader("About")
        st.markdown("""
            <p align="center">
                <a href="https://arxiv.org/abs/2503.00955">
                    <img src="https://img.shields.io/badge/arXiv-2411.00918-red?style=flat&label=arXiv">
                </a>
                <a href="https://huggingface.co/SemViQA">
                    <img src="https://img.shields.io/badge/Hugging%20Face-Model-yellow?style=flat">
                </a>
                <a href="https://pypi.org/project/SemViQA">
                    <img src="https://img.shields.io/pypi/v/SemViQA?color=blue&label=PyPI">
                </a>
                <a href="https://github.com/DAVID-NGUYEN-S16/SemViQA">
                    <img src="https://img.shields.io/github/stars/DAVID-NGUYEN-S16/SemViQA?style=social">
                </a>
            </p>
        """, unsafe_allow_html=True)
        st.markdown("""
            **Description:**  
            SemViQA is a Semantic QA system designed to verify information in Vietnamese.  
            The system extracts evidence from the provided context and classifies information as **SUPPORTED**, **REFUTED**, or **NEI** (Not Enough Information) based on advanced models.
        """)