# Install required packages first:
# pip install langchain huggingface_hub pandas faiss-cpu numpy transformers torch accelerate bitsandbytes gradio

import pandas as pd
import numpy as np
import gradio as gr
import torch
from langchain.vectorstores import FAISS
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
from langchain.llms import HuggingFacePipeline
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from huggingface_hub import InferenceClient

# Configuration for both components
EMBEDDING_MODEL_NAME = "sentence-transformers/all-mpnet-base-v2"
LLM_MODEL_NAME = "EleutherAI/gpt-neo-1.3B"
CHAT_MODEL_NAME = "PygmalionAI/pygmalion-1.3b"

# ------ Chat Component ------
chat_client = InferenceClient(CHAT_MODEL_NAME)

def chat_respond(
    message: str,
    history: list[list[str]],
    system_message: str,
    max_tokens: int,
    temperature: float,
    top_p: float,   
):
    messages = [{"role": "system", "content": system_message}]
    
    for human, assistant in history:
        messages.extend([
            {"role": "user", "content": human},
            {"role": "assistant", "content": assistant}
        ])
    
    messages.append({"role": "user", "content": message})

    response = ""
    for chunk in chat_client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = chunk.choices[0].delta.content
        response += token
        yield response

# ------ RAG Component ------
def load_embeddings(csv_path):
    df = pd.read_csv(csv_path)
    
    if 'Cleaned_Ideas' not in df.columns or 'Embeddings' not in df.columns:
        raise ValueError("CSV must contain 'Cleaned_Ideas' and 'Embeddings' columns")
        
    df['Embeddings'] = df['Embeddings'].apply(
        lambda x: np.fromstring(
            x.strip("[]").replace("\n", ""),
            sep=", ",
            dtype=np.float32
        )
    )
    
    expected_dim = 768
    valid_embeddings = df['Embeddings'].apply(lambda x: len(x) == expected_dim)
    if not valid_embeddings.all():
        invalid_count = len(df) - valid_embeddings.sum()
        raise ValueError(f"{invalid_count} entries have invalid embedding dimensions")
        
    return df['Cleaned_Ideas'].tolist(), np.array(df['Embeddings'].tolist())
#                                                           Fixed here → ↑
def setup_rag_system():
    # Load data and create vector store
    texts, embeddings = load_embeddings("ideas_with_embeddings.csv")
    
    embedding_model = HuggingFaceEmbeddings(
        model_name=EMBEDDING_MODEL_NAME,
        model_kwargs={'device': 'cpu'},
        encode_kwargs={'normalize_embeddings': True}
    )
    
    vector_store = FAISS.from_embeddings(
        text_embeddings=list(zip(texts, embeddings)),
        embedding=embedding_model,
        normalize_L2=True
    )

    # Load LLM
    tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL_NAME)
    model = AutoModelForCausalLM.from_pretrained(
        LLM_MODEL_NAME,
        device_map="auto",
        torch_dtype=torch.float16,
        trust_remote_code=True
    )
    
    pipe = pipeline(
        "text-generation",
        model=model,
        tokenizer=tokenizer,
        max_new_tokens=250,
        temperature=0.7,
        top_p=0.9,
        repetition_penalty=1.1,
        do_sample=True,
        return_full_text=False,
        pad_token_id=tokenizer.eos_token_id
    )

    # Setup QA chain
    template = """<context>
    {context}
    </context>

    <question>
    {question}
    </question>

    Analyze the philosophical concept using ONLY the provided context. Follow these steps:
    1. Identify key elements from the context
    2. Explain relationships between concepts
    3. Highlight any mentioned philosophical implications

    If context is insufficient, state: "This requires deeper analysis beyond current knowledge." 

    Answer:"""
    
    qa_chain = RetrievalQA.from_chain_type(
        llm=HuggingFacePipeline(pipeline=pipe),
        chain_type="stuff",
        retriever=vector_store.as_retriever(
            search_type="similarity_score_threshold",
            search_kwargs={
                "k": 6,
                "score_threshold": 0.55,
                "include_metadata": True
            }
        ),
        return_source_documents=True,
        chain_type_kwargs={
            "prompt": PromptTemplate(
                template=template,
                input_variables=["context", "question"],
                template_format="f-string"
            )
        },
        verbose=True
    )
    return qa_chain

# ------ Unified Interface ------
def create_interface():
    qa_chain = setup_rag_system()
    
    with gr.Blocks(title="Philosophical AI Assistant") as demo:
        gr.Markdown("# 🤖 AI Philosophy Assistant")
        gr.Markdown("Combining conversational AI with philosophical knowledge base")
        
        # Chat Tab
        with gr.Tab("Chat"):
            gr.ChatInterface(
                fn=chat_respond,
                additional_inputs=[
                    gr.Textbox(value="You are a philosophy expert assistant.", label="System Role"),
                    gr.Slider(1, 2048, value=512, label="Max Tokens"),
                    gr.Slider(0.1, 4.0, value=0.7, label="Temperature"),
                    gr.Slider(0.1, 1.0, value=0.95, label="Top-p")
                ],
                chatbot=gr.Chatbot(height=500, avatar_images=("👤", "🤖")),
                textbox=gr.Textbox(placeholder="Ask me anything about philosophy...", container=False, scale=7),
            )
        
        # Research Assistant Tab
        with gr.Tab("Research Assistant"):
            with gr.Row():
                with gr.Column(scale=3):
                    query_input = gr.Textbox(label="Research Query", lines=3)
                    submit_btn = gr.Button("Analyze", variant="primary")
                with gr.Column(scale=2):
                    sources_output = gr.Markdown(label="Source Context")
            response_output = gr.Markdown(label="Analysis Results")
            
            def process_query(query):
                try:
                    result = qa_chain({"query": query})
                    response = result['result'].split("Answer:")[-1].strip()
                    sources = "\n".join(
                        f"**Source {i+1}:** {doc.page_content[:200]}..." 
                        for i, doc in enumerate(result['source_documents'][:3])
                    )
                    return response, sources
                except Exception as e:
                    return f"Error: {str(e)}", ""
            
            submit_btn.click(
                process_query,
                inputs=query_input,
                outputs=[response_output, sources_output]
            )
    
    return demo

if __name__ == "__main__":
    demo = create_interface()
    demo.launch(share=True)