import torch
import gradio as gr
import faiss
import numpy as np

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from sentence_transformers import SentenceTransformer

GEN_MODEL_NAME = "google/flan-t5-small"
EMB_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained(GEN_MODEL_NAME)
gen_model = AutoModelForSeq2SeqLM.from_pretrained(GEN_MODEL_NAME).to(DEVICE)
embedder = SentenceTransformer(EMB_MODEL_NAME, device=DEVICE)

DISCLAIMER = (
    "Disclaimer: This pregnancy assistant provides general educational information only. "
    "It is not a substitute for antenatal care or professional medical advice. "
    "If you have warning signs (heavy bleeding, severe headache, severe abdominal pain, reduced baby movements) "
    "or feel unsafe, please go to a clinic or hospital immediately."
)

pregnancy_facts = [
    {
        "id": 1,
        "topic": "Antenatal care",
        "text": (
            "Pregnant women should attend regular antenatal visits to monitor their own health "
            "and the baby's development. Early booking in the first trimester is recommended so "
            "that any health issues can be detected and managed early. During antenatal visits, "
            "healthcare providers may check blood pressure, weight, urine, and the baby's heartbeat, "
            "and may request simple laboratory tests or scans when needed. These visits are also a "
            "good time to ask questions about nutrition, exercise, birth preparation, and any symptoms "
            "that are worrying you. Skipping antenatal care can make it easier to miss important warning "
            "signs, so it is better to keep appointments and follow the advice of qualified health professionals."
        )
    },
    {
        "id": 2,
        "topic": "Warning signs",
        "text": (
            "Some warning signs in pregnancy need urgent medical attention. These include severe or sudden "
            "headache, blurred vision, swelling of the face or hands, heavy vaginal bleeding, severe abdominal "
            "or pelvic pain, or fluid leaking from the vagina before labour starts. Other serious signs are fever, "
            "painful urination, convulsions, difficulty breathing, or a sudden reduction or change in the baby's "
            "movements. If a pregnant woman notices any of these symptoms, she should go to a clinic or hospital "
            "as soon as possible or call a healthcare professional for guidance. It is safer to be checked and told "
            "that everything is fine than to ignore a possible emergency."
        )
    },
    {
        "id": 3,
        "topic": "Nutrition",
        "text": (
            "During pregnancy, eating a balanced diet helps support the health of the mother and the growing baby. "
            "Meals should include vegetables, fruits, whole grains, and sources of protein such as beans, eggs, fish, "
            "or lean meat, depending on what is culturally acceptable and available. Pregnant women are often advised "
            "by healthcare providers to take folic acid and iron supplements to help reduce the risk of anaemia and "
            "support the baby's development. Drinking enough clean water throughout the day is also important. "
            "Highly processed foods, very sugary drinks, or large amounts of junk food should be limited. "
            "For any supplements or herbal products, it is safer to ask a qualified healthcare professional before use."
        )
    },
    {
        "id": 4,
        "topic": "Substances to avoid",
        "text": (
            "Alcohol, tobacco, and illicit or non-prescribed drugs should be avoided in pregnancy because they may "
            "harm the baby and affect the mother's health. Smoking or exposure to second-hand smoke can increase the "
            "risk of complications in pregnancy and after birth. Some medicines that are safe for non-pregnant adults "
            "may not be safe in pregnancy, especially if taken without medical advice. Before starting, stopping, or "
            "changing any medicine, including herbal remedies or over-the-counter drugs, a pregnant woman should speak "
            "to a qualified healthcare professional. It is better to check first than to take a product that could be "
            "unsafe for the baby."
        )
    },
    {
        "id": 5,
        "topic": "Exercise",
        "text": (
            "For many women, moderate physical activity during pregnancy can be safe and beneficial. Simple activities "
            "like walking, gentle stretching, or pregnancy-safe exercises can help improve mood, sleep, and general "
            "fitness. However, very intense exercise, activities with a high risk of falling, or contact sports may not "
            "be suitable in pregnancy. Pregnant women who were not exercising before pregnancy should start slowly and "
            "listen to their bodies, stopping if they feel pain, dizziness, or shortness of breath. It is always a good "
            "idea to discuss exercise plans with a healthcare professional, especially if there are any medical problems "
            "or previous pregnancy complications."
        )
    },
]

corpus = [item["text"] for item in pregnancy_facts]
corpus_embeddings = embedder.encode(corpus, convert_to_numpy=True)

dimension = corpus_embeddings.shape[1]
index = faiss.IndexFlatL2(dimension)
index.add(corpus_embeddings)

def retrieve_pregnancy_facts(query: str, k: int = 3):
    query_embedding = embedder.encode([query], convert_to_numpy=True)
    distances, indices = index.search(query_embedding, k)
    results = []
    for idx in indices[0]:
        results.append(pregnancy_facts[idx]["text"])
    return results

def pregnancy_qa_pipeline(question: str) -> str:
    question = question.strip()
    if not question:
        return "Please type a question about pregnancy."

    contexts = retrieve_pregnancy_facts(question, k=3)
    context_str = "\n\n".join(contexts)

    prompt = f"""You are an educational pregnancy information assistant.

Context:
{context_str}

Instructions:
- Give general educational information only.
- Do not diagnose or confirm pregnancy status.
- Do not provide medication doses, prescriptions, or step-by-step instructions for taking medicines.
- Do not provide instructions for ending a pregnancy.
- If the question sounds urgent or includes warning signs, strongly recommend immediate in-person care.
- If the context is insufficient, say what you can safely say and recommend seeing a healthcare professional.

Question: {question}

Answer:"""

    inputs = tokenizer(
        prompt,
        return_tensors="pt",
        truncation=True,
        max_length=1024
    ).to(DEVICE)

    with torch.no_grad():
        outputs = gen_model.generate(
            **inputs,
            max_new_tokens=256,
            num_beams=4,
            early_stopping=True
        )

    answer = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
    return f"{answer}\n\n{DISCLAIMER}"

def pregnancy_qa_interface(question: str) -> str:
    return pregnancy_qa_pipeline(question)

with gr.Blocks() as demo:
    gr.Markdown("# 🤰 Pregnancy Information Assistant (Educational Use Only)")
    gr.Markdown(
        "Ask general questions about pregnancy, antenatal care, warning signs, and lifestyle. "
        "This tool is for education only and does not replace a healthcare professional."
    )

    with gr.Row():
        with gr.Column():
            q_box = gr.Textbox(
                label="Ask a question about pregnancy",
                placeholder="e.g. Is it safe to exercise during pregnancy?",
                lines=2
            )
            ask_btn = gr.Button("Ask")
        with gr.Column():
            a_box = gr.Textbox(
                label="Answer",
                lines=10
            )

    ask_btn.click(
        pregnancy_qa_interface,
        inputs=q_box,
        outputs=a_box
    )

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