import gradio as gr
import pandas as pd
import numpy as np
import plotly.express as px
import plotly.graph_objects as go
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import joblib

# ==========================================================
# PHARMABRIDGE — PREMIUM UI EDITION
# Optimized for HuggingFace Spaces
# ==========================================================

# -----------------------------
# LOAD DATA
# -----------------------------

try:
    DF = pd.read_csv("drug_database.csv")
except:
    DF = pd.read_csv("merged_pharma_dataset.csv")

DF = DF.fillna("")

TEXT_COL = "drug_text" if "drug_text" in DF.columns else "generic_name"

# -----------------------------
# LOAD TFIDF MODEL
# -----------------------------

try:
    VEC = joblib.load("tfidf_vectorizer.pkl")
    MAT = joblib.load("tfidf_matrix.pkl")
except:
    VEC = TfidfVectorizer(max_features=15000, ngram_range=(1,2))
    MAT = VEC.fit_transform(DF[TEXT_COL])

# -----------------------------
# COLOR SYSTEM
# -----------------------------

COLORS = {
    "Allopathic":"#2563EB",
    "Ayurvedic":"#16A34A",
    "Unani":"#F59E0B",
    "Homeopathic":"#7C3AED",
    "Herbal":"#EF4444"
}

SYSTEMS = ["All Systems","Allopathic","Ayurvedic","Unani","Homeopathic","Herbal"]

# ==========================================================
# SMART SEARCH
# ==========================================================

def search_drug(query, system, topn):

    if query.strip()=="":
        return "<div class='ph'>🔍 Enter a drug name, symptom, or compound</div>", None

    q_vec = VEC.transform([query])
    sims = cosine_similarity(q_vec, MAT).flatten()

    if system!="All Systems":
        mask = DF["medical_system"]==system
        sims[~mask.values]=0

    idx = sims.argsort()[::-1][:topn]

    results = DF.iloc[idx].copy()
    results["score"] = sims[idx]

    cards = "<div class='grid'>"

    for _,r in results.iterrows():

        sys = r.get("medical_system","Unknown")
        col = COLORS.get(sys,"#64748B")

        cards += f"""
        <div class='card'>
            <div class='badge' style='background:{col}20;color:{col}'>{sys}</div>
            <div class='bn'>{r.get('brand_name','Unknown')}</div>
            <div class='gn'>{r.get('generic_name','')}</div>
            <div class='meta'>💊 {r.get('dosage_form','')} • 🏭 {r.get('manufacturer','')}</div>
            <div class='score'>{round(r['score']*100,1)}% similarity</div>
        </div>
        """

    cards += "</div>"

    fig = px.bar(
        results.head(15),
        x="score",
        y="brand_name",
        color="medical_system",
        orientation="h",
        title="Similarity Scores"
    )

    fig.update_layout(height=420)

    return cards, fig

# ==========================================================
# CROSS SYSTEM COMPARISON
# ==========================================================

def cross_compare(query):

    if query.strip()=="":
        return None

    q_vec = VEC.transform([query])
    sims = cosine_similarity(q_vec, MAT).flatten()

    data=[]

    for sys in COLORS:

        mask = DF["medical_system"]==sys
        scores = sims.copy()
        scores[~mask.values]=0

        idx = scores.argsort()[::-1][:3]

        for i in idx:
            data.append({
                "system":sys,
                "drug":DF.iloc[i].get("brand_name",""),
                "score":scores[i]
            })

    df=pd.DataFrame(data)

    fig=px.bar(
        df,
        x="system",
        y="score",
        color="system",
        hover_name="drug",
        title="Cross Medical System Comparison"
    )

    return fig

# ==========================================================
# DRUG FINGERPRINT
# ==========================================================

def fingerprint(drug):

    if drug.strip()=="":
        return "<div class='ph'>Enter a drug name</div>",None

    q_vec = VEC.transform([drug])
    sims = cosine_similarity(q_vec, MAT).flatten()

    idx = sims.argmax()

    row = DF.iloc[idx]

    html=f"""
    <div class='profile'>
        <h2>{row.get('brand_name','')}</h2>
        <p><b>Medical System:</b> {row.get('medical_system','')}</p>
        <p><b>Generic:</b> {row.get('generic_name','')}</p>
        <p><b>Dosage Form:</b> {row.get('dosage_form','')}</p>
        <p><b>Manufacturer:</b> {row.get('manufacturer','')}</p>
    </div>
    """

    vec = MAT[idx].toarray().flatten()
    top = vec.argsort()[-20:]

    fig = go.Figure()

    fig.add_bar(
        x=vec[top],
        y=VEC.get_feature_names_out()[top],
        orientation="h"
    )

    fig.update_layout(title="Drug Feature Fingerprint",height=420)

    return html, fig

# ==========================================================
# DATASET ANALYTICS
# ==========================================================

def dataset_dashboard():

    sys = DF["medical_system"].value_counts()

    fig = px.pie(
        names=sys.index,
        values=sys.values,
        title="Drug Distribution Across Medical Systems"
    )

    fig.update_layout(height=500)

    return fig

# ==========================================================
# DRUG EXPLORER
# ==========================================================

def explorer(system, search):

    sub = DF.copy()

    if system!="All":
        sub=sub[sub["medical_system"]==system]

    if search:
        s=search.lower()
        sub=sub[sub["brand_name"].str.lower().str.contains(s)]

    rows=""

    for _,r in sub.head(100).iterrows():

        rows+=f"""
        <tr>
        <td>{r.get('brand_name','')}</td>
        <td>{r.get('generic_name','')}</td>
        <td>{r.get('dosage_form','')}</td>
        <td>{r.get('medical_system','')}</td>
        </tr>
        """

    table=f"""
    <table class='tbl'>
    <tr>
    <th>Brand</th>
    <th>Generic</th>
    <th>Dosage</th>
    <th>System</th>
    </tr>
    {rows}
    </table>
    """

    return table

# ==========================================================
# PREMIUM GLASS UI CSS
# ==========================================================

CSS="""
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700;800&display=swap');

body{
background:linear-gradient(135deg,#eef2ff,#f8fafc);
font-family:'Inter';
}

.grid{
display:grid;
grid-template-columns:repeat(auto-fill,minmax(260px,1fr));
gap:18px
}

.card{
background:rgba(255,255,255,0.85);
backdrop-filter:blur(10px);
padding:20px;
border-radius:16px;
box-shadow:0 10px 30px rgba(0,0,0,.08);
transition:all .25s
}

.card:hover{
transform:translateY(-4px);
box-shadow:0 20px 40px rgba(0,0,0,.12)
}

.badge{
font-size:12px;
padding:4px 10px;
border-radius:20px;
display:inline-block;
margin-bottom:6px
}

.bn{
font-size:18px;
font-weight:700
}

.gn{
color:#475569;
font-size:14px;
margin-top:4px
}

.meta{
font-size:13px;
color:#64748B;
margin-top:8px
}

.score{
color:#2563EB;
font-weight:600;
margin-top:10px
}

.ph{
text-align:center;
padding:60px;
background:white;
border-radius:16px;
color:#64748B
}

.profile{
background:white;
padding:25px;
border-radius:16px;
box-shadow:0 8px 24px rgba(0,0,0,.08)
}

.tbl{
width:100%;
border-collapse:collapse
}

.tbl th{
background:#0f172a;
color:white;
padding:10px
}

.tbl td{
padding:10px;
border-bottom:1px solid #e2e8f0
}

"""

# ==========================================================
# UI
# ==========================================================

with gr.Blocks(css=CSS,title="PharmaBridge") as app:

    gr.Markdown("""
# 💊 PharmaBridge AI
### Cross‑Medical‑System Drug Intelligence Platform
""")

    with gr.Tabs():

        with gr.Tab("🔍 Smart Search"):

            q=gr.Textbox(label="Search Drug")
            sys=gr.Dropdown(SYSTEMS,value="All Systems")
            n=gr.Slider(5,30,value=10)
            btn=gr.Button("Search")

            cards=gr.HTML()
            chart=gr.Plot()

            btn.click(search_drug,[q,sys,n],[cards,chart])

        with gr.Tab("⚖️ Cross System Compare"):

            q2=gr.Textbox(label="Query")
            btn2=gr.Button("Compare")

            fig=gr.Plot()

            btn2.click(cross_compare,q2,fig)

        with gr.Tab("📊 Dataset Analytics"):

            btn3=gr.Button("Load Dashboard")
            dash=gr.Plot()

            btn3.click(dataset_dashboard,outputs=dash)

        with gr.Tab("🧬 Drug Fingerprint"):

            q3=gr.Textbox(label="Drug Name")
            btn4=gr.Button("Analyze")

            card=gr.HTML()
            fig2=gr.Plot()

            btn4.click(fingerprint,q3,[card,fig2])

        with gr.Tab("📋 Drug Explorer"):

            sys2=gr.Dropdown(["All","Allopathic","Ayurvedic","Unani","Homeopathic","Herbal"],value="All")
            search=gr.Textbox(label="Search")
            btn5=gr.Button("Browse")

            table=gr.HTML()

            btn5.click(explorer,[sys2,search],table)

        with gr.Tab("ℹ️ About"):

            gr.Markdown("""
PharmaBridge is a pharmaceutical AI research system designed to explore drug datasets across multiple medical traditions.

Features:

• Smart Drug Search
• Cross System Medicine Comparison
• Dataset Analytics Dashboard
• Drug Fingerprint Analysis
• Drug Explorer

Built for research and educational use.
""")

app.launch()