import gradio as gr
from huggingface_hub import InferenceClient
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

client = InferenceClient(model="Qwen/Qwen2.5-7B-Instruct")

with open("knowledge.txt", "r", encoding="utf-8") as file:
    knowledge_text = file.read()

def preprocess_text(text):
    chunks = text.strip().split("\n")
    return [chunk.strip() for chunk in chunks if chunk.strip()]

chunks = preprocess_text(knowledge_text)

embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
chunk_embeddings = embedding_model.encode(chunks)

def get_relevant_context(user_query, top_k=5):
    query_embedding = embedding_model.encode([user_query])
    similarity_scores = cosine_similarity(query_embedding, chunk_embeddings)[0]
    top_indices = np.argsort(similarity_scores)[-top_k:][::-1]
    return "\n".join([chunks[i] for i in top_indices])

def respond(message, history):
    context = get_relevant_context(message)

    system_prompt = f"""
You are Femory, a friendly reproductive health education chatbot.

Use the context below to answer the user's question.

Context:
{context}

Rules:
- Explain in simple, supportive language.
- Do not diagnose the user.
- Do not claim to replace a doctor.
- Only answer reproductive health education questions.
- Always end with: "This is for educational purposes only. Please talk to a healthcare professional for medical advice."
"""

    messages = [{"role": "system", "content": system_prompt}]

    if history:
        messages.extend(history)

    messages.append({"role": "user", "content": message})

    response = client.chat_completion(
        messages=messages,
        max_tokens=600
    )

    return response.choices[0].message.content.strip()

def chat(message, history):
    if history is None:
        history = []

    bot_response = respond(message, history)

    history.append({"role": "user", "content": message})
    history.append({"role": "assistant", "content": bot_response})

    return history, ""

def fill_pmos():
    return "What is PMOS?"

def fill_endometriosis():
    return "What are symptoms of endometriosis?"

def fill_adenomyosis():
    return "What is adenomyosis?"

def fill_fibroids():
    return "What are fibroids?"

def fill_painful_periods():
    return "What can cause painful periods?"

def fill_heavy_bleeding():
    return "What can cause heavy bleeding?"

def fill_irregular_cycle():
    return "What can cause an irregular cycle?"

custom_css = """
.gradio-container {
    background: #FDE7EC !important;
}

.header-card {
    background: white;
    border-radius: 22px;
    padding: 22px;
    margin-bottom: 20px;
    text-align: center;
}

.main-panel {
    background: white;
    border-radius: 20px;
    padding: 20px;
}

.condition-card {
    background: white;
    border-radius: 20px;
    padding: 15px;
    margin-bottom: 12px;
}

button {
    border-radius: 20px !important;
}

footer {
    display: none !important;
}
"""

with gr.Blocks() as demo:

    gr.HTML("""
    <div class="header-card">
        <h1 style="color:#E96A8D;">🌸 Femory</h1>
        <p>Your reproductive health companion</p>
    </div>
    """)

    with gr.Row():

        with gr.Column(scale=1):
            gr.Markdown("### 💗 Common Topics")

            pmos_btn = gr.Button("PMOS")
            endo_btn = gr.Button("Endometriosis")
            adeno_btn = gr.Button("Adenomyosis")
            fibroid_btn = gr.Button("Fibroids")

            gr.Markdown("""
### 🌷 Safe Space
- Symptoms
- Treatments
- Diagnoses
- Women's health conditions
            """)

        with gr.Column(scale=3):

            gr.HTML("""
            <div class="main-panel">
                <h2 style="color:#E96A8D;">Chat with Femory</h2>
                <p>Ask questions about symptoms, treatments, and reproductive health.</p>
            </div>
            """)

            chatbot = gr.Chatbot(height=500)

            with gr.Row():
                painful_btn = gr.Button("Painful Periods")
                heavy_btn = gr.Button("Heavy Bleeding")
                irregular_btn = gr.Button("Irregular Cycle")

            msg = gr.Textbox(
                placeholder="Ask a reproductive health question...",
                show_label=False
            )

            send = gr.Button("Send")

        with gr.Column(scale=1):
            gr.HTML("""
            <div class="condition-card">
                <h4>💗 PMOS</h4>
                Hormonal condition that can affect periods, skin, hair growth, and fertility.
            </div>

            <div class="condition-card">
                <h4>💜 Endometriosis</h4>
                Tissue similar to the uterine lining grows outside the uterus.
            </div>

            <div class="condition-card">
                <h4>🌸 Adenomyosis</h4>
                Tissue grows into the muscle wall of the uterus.
            </div>

            <div class="condition-card">
                <h4>🩷 Fibroids</h4>
                Noncancerous growths that can develop in or around the uterus.
            </div>
            """)

    send.click(chat, inputs=[msg, chatbot], outputs=[chatbot, msg])
    msg.submit(chat, inputs=[msg, chatbot], outputs=[chatbot, msg])

    pmos_btn.click(fill_pmos, outputs=msg)
    endo_btn.click(fill_endometriosis, outputs=msg)
    adeno_btn.click(fill_adenomyosis, outputs=msg)
    fibroid_btn.click(fill_fibroids, outputs=msg)
    painful_btn.click(fill_painful_periods, outputs=msg)
    heavy_btn.click(fill_heavy_bleeding, outputs=msg)
    irregular_btn.click(fill_irregular_cycle, outputs=msg)

demo.launch(css=custom_css, theme=gr.themes.Soft())