import os
import numpy as np
import pandas as pd
import gradio as gr
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity

# =========================
# 1. ASSETS & CONFIG
# =========================
MODEL_ID = "sentence-transformers/all-MiniLM-L6-v2"
DATA_FILE = "synthetic_analyst_notes_10k.csv"
EMBED_FILE = "embeddings.npy"

# =========================
# 2. WALL STREET STYLING (CSS) - FINAL POLISH
# =========================
CSS = r"""
:root{
  --bg:#0b1220;
  --panel:#0f1a2f;
  --text:#e7eefc;
  --muted:#9db0d0;
  --accent:#2b76ff;
  --accent2:#00c2ff;
  --good:#18d38a;
  --bad:#ff4d5e;
  --border:rgba(255,255,255,.10);
  --shadow:0 22px 70px rgba(0,0,0,.55);
}
*{box-sizing:border-box;}
html,body{background:var(--bg); color:var(--text);}
.gradio-container{max-width:1200px !important; background:var(--bg) !important; color:var(--text) !important;}
h1,h2,h3,h4, label, .label, .wrap, .prose, .markdown, .md{color:var(--text) !important;}

/* --- FINAL NAV BAR FIX (ALL WHITE) --- */
.nav{
    position:sticky; top:0; z-index:50; 
    display:flex; align-items:center; gap:16px; 
    padding:14px 18px; 
    background:rgba(11,18,32,.95); 
    backdrop-filter: blur(10px); 
    border-bottom:1px solid rgba(255,255,255,.15);
}
/* Force all text inside nav to be white */
.nav * { color: #ffffff !important; }

.brand{display:flex; align-items:center; gap:10px; font-weight:900; letter-spacing:.4px;}
.brand .logo{width:22px; height:22px; border-radius:7px; background:linear-gradient(135deg, var(--accent), var(--accent2)); box-shadow:0 16px 55px rgba(43,118,255,.28);}

.links{display:flex; gap:14px; font-size:13px; font-weight: 500;}
.links div { opacity: 0.9; cursor: pointer; }
.links div:hover { opacity: 1; text-decoration: underline; }

.spacer{flex:1;}
.pill{border:1px solid rgba(255,255,255,.30); border-radius:999px; padding:7px 10px; font-size:12px; background:rgba(255,255,255,.15);}
.kpi{display:flex; align-items:center; gap:10px; font-size:12px;}
.dot{width:7px; height:7px; border-radius:999px; background:rgba(24,211,138,.95); box-shadow:0 0 18px rgba(24,211,138,.30);}

/* HERO & SECTIONS */
.hero{margin:18px; border-radius:22px; overflow:hidden; position:relative; min-height:260px; background:radial-gradient(1100px 420px at 20% 10%, rgba(43,118,255,.34), transparent 62%), radial-gradient(900px 380px at 80% 30%, rgba(0,194,255,.16), transparent 60%), linear-gradient(180deg, rgba(255,255,255,.06), rgba(255,255,255,0)); border:1px solid rgba(255,255,255,.08); box-shadow:var(--shadow);}
.hero-inner{position:relative; padding:30px 32px 22px 32px; display:grid; grid-template-columns: 1.25fr .75fr; gap:18px; align-items:end;}
.hero h1{margin:0; font-size:38px; line-height:1.06;}
.hero p{margin:10px 0 0 0; color:var(--muted); max-width:62ch; font-size:14px;}
.tagrow{display:flex; gap:10px; flex-wrap:wrap; justify-content:flex-end; padding-bottom:6px;}
.tag{font-size:12px; color:var(--muted); border:1px solid rgba(255,255,255,.10); background:rgba(255,255,255,.04); padding:7px 10px; border-radius:999px;}

.section{margin:18px; border:1px solid rgba(255,255,255,.08); background:rgba(255,255,255,.03); border-radius:18px; box-shadow:0 18px 55px rgba(0,0,0,.35); overflow:hidden;}
.section-head{padding:14px 16px; border-bottom:1px solid rgba(255,255,255,.06); display:flex; align-items:baseline; justify-content:space-between; gap:10px;}
.pad{padding:14px 16px;}
.whitepanel{background:#ffffff !important; color:#111111 !important; border-radius:16px !important; border:1px solid rgba(0,0,0,.08) !important;}
.whitepanel *{color:#111111 !important;}
.gr-textbox textarea{background:#ffffff !important; color:#111111 !important; border:1px solid rgba(0,0,0,.10) !important; border-radius:16px !important;}

/* BUTTONS UNIFORMITY */
button { height: 100% !important; min-height: 50px !important; }
.runbtn{border:1px solid rgba(43,118,255,.55) !important; background:linear-gradient(135deg, rgba(43,118,255,.95), rgba(0,194,255,.70)) !important; color:#ffffff !important; box-shadow:0 18px 45px rgba(43,118,255,.22) !important; font-weight: bold !important;}
.quick{border:1px solid rgba(0,0,0,.12) !important; background:#f0f2f5 !important; color:#111111 !important;}
.quick:hover{border-color:rgba(43,118,255,.35) !important; background:#ffffff !important;}

/* CARDS */
.cards{display:grid; grid-template-columns: repeat(3, minmax(0, 1fr)); gap:12px;}
@media (max-width: 980px){.cards{grid-template-columns: 1fr;}}
.card{border-radius:16px; border:1px solid rgba(0,0,0,.10); background:#ffffff; padding:12px; box-shadow:0 14px 40px rgba(0,0,0,.10);}
.card .top{display:flex; align-items:center; justify-content:space-between; gap:10px; margin-bottom:8px;}
.badge{font-size:11px; font-weight:800; padding:6px 8px; border-radius:999px; background:rgba(43,118,255,.10); border:1px solid rgba(43,118,255,.35); color:#111111;}
.sym{font-weight:900; letter-spacing:.2px; font-size:18px;}
.meta{color:#333; font-size:12px; display:flex; gap:10px; flex-wrap:wrap; margin-bottom:10px;}
.note{font-size:13px; line-height:1.35; color:#111; white-space:pre-wrap;}
"""

NAV_HTML = """
<div class="nav">
  <div class="brand"><div class="logo"></div><div>WALLSTREET AI</div></div>
  <div class="links">
    <div>Semantic Search</div><div>Signals</div><div>Grounding</div><div>Terminal</div>
  </div>
  <div class="spacer"></div>
  <div class="kpi"><div class="dot"></div><div>System Live</div></div>
  <div class="pill">10k Synthetic Data</div>
</div>
"""

HERO_HTML = """
<div class="hero">
  <div class="hero-inner">
    <div>
      <h1>Semantic Strategy Engine</h1>
      <p>
        Describe your investment thesis (e.g., "Aggressive tech growth but with debt risks"). 
        The AI performs <b>Semantic Search</b> to find grounded historical cases and generates a top pick.
      </p>
    </div>
    <div class="tagrow">
      <div class="tag">Transformer Models</div>
      <div class="tag">Contextual Grounding</div>
      <div class="tag">Real-time Inference</div>
    </div>
  </div>
</div>
"""

# =========================
# 3. BACKEND LOGIC
# =========================
print("⏳ Loading Models & Data...")
# Load Model
model = SentenceTransformer(MODEL_ID)

# Load Data
try:
    df = pd.read_csv(DATA_FILE)
    df['analyst_note'] = df['analyst_note'].astype(str)
    embeddings = np.load(EMBED_FILE)
    print("✅ Assets Loaded.")
except Exception as e:
    print(f"❌ Error loading assets: {e}")
    df = pd.DataFrame()
    embeddings = np.array([])

def search(query):
    if not query or df.empty:
        return pd.DataFrame(), "", ""

    # 1. Embed Query
    query_vec = model.encode([query])
    
    # 2. Similarity Search
    sims = cosine_similarity(query_vec, embeddings)[0]
    
    # 3. Get Top 3
    top_indices = sims.argsort()[-3:][::-1]
    results = df.iloc[top_indices].copy()
    results['score'] = sims[top_indices]
    
    # Prepare Data Table
    display_cols = ['company_name', 'sector', 'recommendation', 'revenue_musd', 'revenue_growth', 'score']
    table_data = results[[c for c in display_cols if c in results.columns]]

    # 4. Generate "Cards" HTML (For Top 3)
    cards_html = "<div class='cards'>"
    for i, (_, row) in enumerate(results.iterrows()):
        badge = "TOP MATCH" if i == 0 else f"#{i+1}"
        note_snippet = row['analyst_note'][:180] + "..."
        
        cards_html += f"""
        <div class="card">
            <div class="top">
                <div class="sym">{row.get('company_name', 'UNK')}</div>
                <div class="badge">{badge}</div>
            </div>
            <div class="meta">
                {row.get('sector', '')} • Rec: <b>{row.get('recommendation', '')}</b> • Score: {row['score']:.3f}
            </div>
            <div class="note">"{note_snippet}"</div>
        </div>
        """
    cards_html += "</div>"

    # 5. Generate "AI PICK" HTML (Enriched/Faker Analysis)
    best = results.iloc[0]
    
    # Constructing a "Smart" Fake Rationale based on real data
    rev_growth = float(best.get('revenue_growth', 0)) * 100
    rev_abs = best.get('revenue_musd', 0)
    sector = best.get('sector', 'Unknown')
    rec = best.get('recommendation', 'HOLD')
    
    # Richer text generation (template)
    analysis_text = f"""
    <b>Strategic Rationale:</b><br>
    Our semantic algorithm identified <b>{best.get('company_name')}</b> as the optimal historical precedent. 
    Operating within the <b>{sector}</b> sector, this entity demonstrates a classic <b>{rec}</b> signal configuration, 
    aligning with your request for specific risk/reward parameters.
    <br><br>
    <b>Key Financial Indicators:</b><br>
    • <b>Revenue Trajectory:</b> The company reported ${rev_abs}M with a growth rate of <b>{rev_growth:.1f}%</b>, suggesting significant market movement.<br>
    • <b>Analyst Consensus:</b> The underlying sentiment extraction highlights correlation to your query keywords.
    <br><br>
    <b>Contextual Grounding (Source Note):</b><br>
    <i>"{best['analyst_note']}"</i>
    """

    pick_html = f"""
    <div class='card' style='margin-top:12px; border:2px solid #4f8cff; background:#f0f7ff'>
      <div class='top'>
        <div class='sym'>{best.get('company_name', 'UNK')}</div>
        <div class='badge' style='background:#2b76ff; color:white'>🏆 AI TOP PICK</div>
      </div>
      <div class='meta'>
         <b>Strategy: {rec}</b> • Confidence Score: {best['score']:.4f}
      </div>
      <div class='note'>
        {analysis_text}
      </div>
    </div>
    """
    
    return table_data, cards_html, pick_html

# =========================
# 4. QUICK STARTERS (STRATEGY TYPES)
# =========================
def q1(): return "Companies with aggressive revenue growth >20% despite high burn rate and volatility"
def q2(): return "Distressed assets with high debt-to-equity ratios and liquidity concerns"
def q3(): return "Established firms with compressing gross margins and operational inefficiencies"

# =========================
# 5. UI CONSTRUCTION
# =========================
with gr.Blocks(css=CSS, theme=gr.themes.Base()) as demo:
    gr.HTML(NAV_HTML)
    gr.HTML(HERO_HTML)

    with gr.Group(elem_classes=["section", "whitepanel"]):
        gr.HTML("<div class='section-head'><h2>Investment Query</h2><div class='hint'>Describe what you are looking for</div></div>")
        
        with gr.Row(elem_classes=["pad"]):
            query_in = gr.Textbox(label="Your Thesis", placeholder="e.g. 'Aggressive growth but burning cash...'", lines=2)
            
        # BUTTONS ROW - All equal size now
        with gr.Row(elem_classes=["pad"]):
            run_btn = gr.Button("🔍 SCAN MARKET", elem_classes=["runbtn"])
            b1 = gr.Button("🔥 High Growth", elem_classes=["quick"])
            b2 = gr.Button("⚠️ High Risk", elem_classes=["quick"])
            b3 = gr.Button("📉 Low Margins", elem_classes=["quick"])

        gr.HTML("<div class='section-head'><h2>Market Intelligence</h2><div class='hint'>Similar Historical Cases & AI Pick</div></div>")
        
        # ORDER FIXED: TABS (Visual Cards) FIRST, THEN AI PICK
        with gr.Tabs(elem_classes=["tabs"]):
            with gr.TabItem("Visual Cards"):
                out_cards = gr.HTML("<div class='smallmuted'>Results will appear here...</div>")
            with gr.TabItem("Data Table"):
                out_table = gr.Dataframe()

        # AI Pick is now below the tabs
        out_pick = gr.HTML()

    # Interactions
    run_btn.click(fn=search, inputs=query_in, outputs=[out_table, out_cards, out_pick])
    b1.click(fn=q1, outputs=query_in).then(fn=search, inputs=query_in, outputs=[out_table, out_cards, out_pick])
    b2.click(fn=q2, outputs=query_in).then(fn=search, inputs=query_in, outputs=[out_table, out_cards, out_pick])
    b3.click(fn=q3, outputs=query_in).then(fn=search, inputs=query_in, outputs=[out_table, out_cards, out_pick])

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