import pandas as pd
import numpy as np
import plotly.express as px
from sklearn.decomposition import PCA
from sklearn.metrics.pairwise import cosine_similarity
from sentence_transformers import SentenceTransformer
import warnings
warnings.filterwarnings('ignore')

# ==========================================
# 1. KHỞI TẠO MODELS (Load sẵn khi import file)
# ==========================================
print("[Utils] Đang khởi tạo các mô hình...")
model_sbert = SentenceTransformer('keepitreal/vietnamese-sbert')
model_e5 = SentenceTransformer('intfloat/multilingual-e5-base')
model_bge = SentenceTransformer('BAAI/bge-m3')
print("[Utils] Đã load xong 3 models!")

# ==========================================
# 2. DỮ LIỆU NỀN (KNOWLEDGE BASE)
# ==========================================
base_sentences = [
    "Tôi rất thích nghiên cứu về Trí tuệ nhân tạo.",
    "I love studying Artificial Intelligence.",
    "Con mèo đang ngủ trên ghế sofa.",
    "The cat is taking a nap on the couch.",
    "Công thức làm bánh pizza ngon nhất thế giới."
]

# ==========================================
# 3. HÀM XỬ LÝ LOGIC LÕI
# ==========================================
def process_embedding_analysis(user_text, base_text_input):
    """
    Nhận text từ UI, tính toán vector, trả về Biểu đồ Plotly và Dataframe kết quả.
    """
    if not user_text.strip():
        return None, None, None, pd.DataFrame()
    
    base_sentences = [s.strip() for s in base_text_input.split('\n') if s.strip()]
    
    if len(base_sentences) == 0:
        return None, None, None, pd.DataFrame({"Lỗi": ["Vui lòng nhập ít nhất 1 câu vào Dữ liệu nền!"]})
        
    # Tui viết một cái hàm nhỏ (helper function) bên trong để tái sử dụng code cho gọn nè
    def get_plot_and_sim(model, model_name, is_e5=False):
        sentences_to_encode = base_sentences.copy()
        query_to_encode = user_text
        
        # Xử lý riêng cho E5 (cái vụ gắn prefix á)
        if is_e5:
            sentences_to_encode = [f"passage: {text}" for text in base_sentences]
            query_to_encode = f"query: {user_text}"
            
        base_embeddings = model.encode(sentences_to_encode)
        user_embedding = model.encode([query_to_encode])
        similarities = cosine_similarity(user_embedding, base_embeddings)[0]
        
        # Tính PCA
        all_texts = sentences_to_encode + [query_to_encode]
        all_emb = model.encode(all_texts)
        
        n_samples = len(all_texts)
        n_comp = min(2, n_samples)
        pca = PCA(n_components=n_comp)
        emb_2d = pca.fit_transform(all_emb)
        # pca = PCA(n_components=2)
        # emb_2d = pca.fit_transform(all_emb)
        
        if n_comp == 1:
            emb_2d = np.hstack((emb_2d, np.zeros((n_samples, 1))))
        
        df_plot = pd.DataFrame({
            'X': emb_2d[:, 0], 'Y': emb_2d[:, 1],
            'Text': base_sentences + [user_text],
            'Loại': ["Dữ liệu nền"] * len(base_sentences) + ["Câu bạn nhập"]
        })
        
        # Vẽ hình
        fig = px.scatter(
            df_plot, x='X', y='Y', color='Loại', hover_name='Text',
            title=f"{model_name}",
            color_discrete_map={"Dữ liệu nền": "#636EFA", "Câu bạn nhập": "#EF553B"}
        )
        
        # Vẽ đường nối
        best_match_idx = similarities.argmax()
        fig.add_shape(
            type="line",
            x0=emb_2d[-1, 0], y0=emb_2d[-1, 1],
            x1=emb_2d[best_match_idx, 0], y1=emb_2d[best_match_idx, 1],
            line=dict(color="gray", width=2, dash="dot"), layer="below"
        )
        
        # Thu nhỏ margin xíu cho 3 biểu đồ đứng cạnh nhau không bị chật
        fig.update_traces(marker=dict(size=10, line=dict(width=1, color='white')))
        fig.update_layout(template="plotly_white", margin=dict(l=10, r=10, t=35, b=10))
        
        return fig, similarities

    # Bắt đầu chạy cả 3 con AI nè
    fig_sbert, sim_sbert = get_plot_and_sim(model_sbert, "1. Vietnamese SBERT")
    fig_e5, sim_e5 = get_plot_and_sim(model_e5, "2. Microsoft E5", is_e5=True)
    fig_bge, sim_bge = get_plot_and_sim(model_bge, "3. BAAI BGE-M3")
    
    # Gộp điểm của 3 con vào chung 1 cái bảng cho dễ nhìn á
    results_df = pd.DataFrame({
        "Câu trong Database": base_sentences,
        "SBERT (%)": [round(s * 100, 1) for s in sim_sbert],
        "E5 (%)": [round(s * 100, 1) for s in sim_e5],
        "BGE-M3 (%)": [round(s * 100, 1) for s in sim_bge]
    })
    
    return fig_sbert, fig_e5, fig_bge, results_df