import gradio as gr
import folium
from folium.plugins import MarkerCluster
import tempfile
import math

# =========================
# DATABASE YA HOSPITALS MU RWANDA
# =========================
hospitals_data = [
    {"name": "King Faisal Hospital", "district": "Gasabo",
     "gps": (-1.949, 30.089), "specialty": ["general", "maternity", "surgery"]},

    {"name": "CHUK", "district": "Nyarugenge",
     "gps": (-1.950, 30.058), "specialty": ["general", "pediatrics", "emergency"]},

    {"name": "Butaro Hospital", "district": "Burera",
     "gps": (-1.641, 29.893), "specialty": ["general", "oncology", "pediatrics"]},

    {"name": "Gahini District Hospital", "district": "Kayonza",
     "gps": (-2.019, 30.576), "specialty": ["general", "orthopedic"]},
]

# =========================
# HAVERSINE DISTANCE (NO geopy)
# =========================
def haversine(coord1, coord2):
    lat1, lon1 = coord1
    lat2, lon2 = coord2
    R = 6371  # km

    dlat = math.radians(lat2 - lat1)
    dlon = math.radians(lon2 - lon1)

    a = math.sin(dlat/2)**2 + math.cos(math.radians(lat1)) \
        * math.cos(math.radians(lat2)) * math.sin(dlon/2)**2

    return 2 * R * math.atan2(math.sqrt(a), math.sqrt(1 - a))

# =========================
# SIMPLE NLP (PURE PYTHON)
# =========================
def extract_specialty(text):
    text = text.lower()
    keywords = {
        "child": "pediatrics",
        "baby": "pediatrics",
        "pediatric": "pediatrics",
        "pregnant": "maternity",
        "birth": "maternity",
        "cancer": "oncology",
        "tumor": "oncology",
        "surgery": "surgery",
        "operation": "surgery",
        "emergency": "emergency",
        "urgent": "emergency"
    }

    for key, value in keywords.items():
        if key in text:
            return value
    return ""

# =========================
# MAIN AI FUNCTION
# =========================
def ai_healthcare_assistant(user_text, lat, lon):
    specialty = extract_specialty(user_text)
    user_location = (lat, lon)

    results = []

    for h in hospitals_data:
        if specialty and specialty not in h["specialty"]:
            continue

        distance = round(haversine(user_location, h["gps"]), 2)
        h_copy = h.copy()
        h_copy["distance"] = distance
        results.append(h_copy)

    if not results:
        return "❌ Nta bitaro bibonetse bijyanye n'icyo ushaka.", None

    results = sorted(results, key=lambda x: x["distance"])

    # =========================
    # CREATE MAP
    # =========================
    m = folium.Map(location=[lat, lon], zoom_start=8)
    cluster = MarkerCluster().add_to(m)

    folium.Marker(
        [lat, lon],
        popup="📍 AHO URI",
        icon=folium.Icon(color="blue")
    ).add_to(m)

    for h in results:
        popup = f"""
        <b>{h['name']}</b><br>
        District: {h['district']}<br>
        Services: {', '.join(h['specialty'])}<br>
        Distance: {h['distance']} km
        """
        folium.Marker(h["gps"], popup=popup).add_to(cluster)

    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".html")
    m.save(tmp.name)

    # =========================
    # TEXT RESPONSE
    # =========================
    response = "🏥 Ibitaro byegereye aho uri:\n"
    for h in results:
        response += f"- {h['name']} ({h['distance']} km)\n"

    return response, tmp.name

# =========================
# GRADIO INTERFACE
# =========================
iface = gr.Interface(
    fn=ai_healthcare_assistant,
    inputs=[
        gr.Textbox(label="Sobanura icyo ukeneye (ex: I need cancer hospital near me)"),
        gr.Number(label="Latitude (aho uri)"),
        gr.Number(label="Longitude (aho uri)")
    ],
    outputs=[
        gr.Textbox(label="AI Recommendation"),
        gr.File(label="Interactive Hospital Map")
    ],
    title="AI Healthcare Assistant – Rwanda",
    description="AI ikumva icyo ushaka ikakwereka ibitaro byegereye aho uri hakurikijwe service ukeneye."
)

iface.launch()