import spaces
import os
import torch
import random
import numpy as np
import pandas as pd
import gradio as gr
from pathlib import Path

from transformers import pipeline, set_seed
from sentence_transformers import SentenceTransformer
import transformers.utils.import_utils as transformers_import_utils
from datasets import load_dataset

# Transformers 5 may query the CUDA stream while preparing CPU attention
# masks (used internally when we call embedding_model.encode(...) on CPU).
# In the GPU-less ZeroGPU parent process, that stray CUDA runtime call
# leaves CUDA in a state that then breaks the next forked @spaces.GPU
# worker's attempt to initialize its dynamically assigned GPU, causing
# "RuntimeError: No CUDA GPUs are available" even though the code never
# explicitly touches CUDA outside the decorated function.
# Fix: skip only that query in the ZeroGPU parent; real GPU workers keep
# the original Transformers behavior once CUDA_VISIBLE_DEVICES is set.
# See: https://huggingface.co/spaces/zero-gpu-explorers/README/discussions/181
if hasattr(transformers_import_utils, "is_cuda_stream_capturing"):
    _original_is_cuda_stream_capturing = transformers_import_utils.is_cuda_stream_capturing

    def _zerogpu_safe_is_cuda_stream_capturing(*args, **kwargs):
        if os.getenv("SPACES_ZERO_GPU") == "1" and not os.getenv("CUDA_VISIBLE_DEVICES"):
            return False
        return _original_is_cuda_stream_capturing(*args, **kwargs)

    transformers_import_utils.is_cuda_stream_capturing = _zerogpu_safe_is_cuda_stream_capturing

# 1. Initialize LLM
print("Loading BioMistral LLM...")
llm = pipeline('text-generation', model='BioMistral/BioMistral-7B', torch_dtype=torch.bfloat16, device_map="cuda")

# 2. Load Datasets
print("Loading datasets...")
dataset = load_dataset("cbasu/Med-EASi", split="train")
dataset2 = load_dataset("cbasu/Med-EASi", split="validation")
dataset3 = load_dataset("cbasu/Med-EASi", split="test")
dataset4 = load_dataset("csv", data_files="custom_terms.csv", split="train")

# 3. Format dataset rows into strings
texts = []
for row in dataset:
    content = f"Medical Context: {row['Expert']}\nSimplified Explanation: {row['Simple']}\nTerms: {row['expert_terms']}, {row['layman_terms']}"
    texts.append(content)
for row in dataset2:
    content = f"Medical Context: {row['Expert']}\nSimplified Explanation: {row['Simple']}\nTerms: {row['expert_terms']}, {row['layman_terms']}"
    texts.append(content)
for row in dataset3:
    content = f"Medical Context: {row['Expert']}\nSimplified Explanation: {row['Simple']}\nTerms: {row['expert_terms']}, {row['layman_terms']}"
    texts.append(content)
for row in dataset4:
    content = f"Medical Context: {row['Expert']}\nSimplified Explanation: {row['Simple']}\nTerms: {row['expert_terms']}, {row['layman_terms']}"
    texts.append(content)

print(f"Total processed documents: {len(texts)}")

# 4. Generate Embeddings & Index
print("Indexing dataset embeddings...")
embedding_model = SentenceTransformer('all-MiniLM-L6-v2', device='cpu')
embeddings = embedding_model.encode(texts)

df = pd.DataFrame({
    "Document": texts,
    "Embedding": list(embeddings)
})

# 5. Retrieval Helper Function
def retrieve_with_pandas(query, top_k=3):
    query_embedding = embedding_model.encode([query])[0]
    df['Similarity'] = df['Embedding'].apply(
        lambda x: np.dot(query_embedding, x) / (np.linalg.norm(query_embedding) * np.linalg.norm(x))
    )
    results = df.sort_values(by="Similarity", ascending=False).head(top_k)
    return results["Document"].tolist()

# 6. Prompt Builder Helper Function
def build_prompt(context_docs: list, user_query: str) -> str:
    context = "\n\n".join(context_docs)
    prompt = f"""<s>[INST] You are a helpful medical assistant.
Use the following retrieved medical context to answer the user's question clearly and concisely.
Explain it in very simple terms as if the user has the vocabulary of an elementary school student.
If the context does not contain enough information to answer, state that you don't know based on the provided text.
Answer in 2 to 3 sentences maximum.

Context:
{context}

Question:
{user_query} [/INST]"""
    return prompt

# 7. Answer Generation Helper Function
@spaces.GPU(duration=60)
def generate_answer(prompt: str, max_new_tokens: int = 80) -> str:
    set_seed(42)
    if llm.model.config.pad_token_id is None:
        llm.model.config.pad_token_id = llm.model.config.eos_token_id

    output = llm(
        prompt,
        max_new_tokens=max_new_tokens,
        num_return_sequences=1,
        do_sample=False,
    )[0]["generated_text"]

    # Strip out the prompt text to leave only the model's new response
    answer = output[len(prompt):].strip()
    return answer

# 8. RAG Handler Function
def respond_to_query(user_query: str) -> str:
    if not user_query.strip():
        return "Please enter a valid medical question or term."
    
    # Step A: Retrieve context documents based on user input
    context_docs = retrieve_with_pandas(user_query, top_k=3)
    
    # Step B: Build prompt with retrieved context
    prompt = build_prompt(context_docs, user_query)
    
    # Step C: Generate model answer
    answer = generate_answer(prompt, max_new_tokens=80)
    
    return answer

# 9. Gradio Web Interface
demo = gr.Interface(
    fn=respond_to_query,
    inputs=gr.Textbox(lines=2, label="User Question", placeholder="e.g., What is necrosis? or What is cholera?"),
    outputs=gr.Textbox(label="BioMistral RAG Response", lines=4),
    title="Medical Terminology RAG Assistant",
    description="Ask a medical question to retrieve context from Med-EASi + custom terms and generate simplified explanations.",
    examples=[
        ["What causes necrosis?"],
        ["What is cholera?"],
        ["What is asbestosis?"],
        ["How can syphilis be treated?"]
    ]
)

if __name__ == "__main__":
    demo.launch()