import gradio as gr
import numpy as np
import inspect
import pandas as pd
import torch
import re
from datasets import load_dataset
from sentence_transformers import SentenceTransformer
import faiss
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM
import spaces  # <--- Import the Hugging Face spaces library for Free GPU

# =========================================================================
# 1. LAZY LOADING ARCHITECTURE (The Ultimate Fix for Smart AI)
# To get a truly "smart" pitch, we MUST use a modern State-of-the-Art LLM.
# We are upgrading from the ancient 'flan-t5' to the brilliant 'Qwen2.5' model.
# By lazy loading, we ensure Hugging Face never crashes during startup!
# =========================================================================

# Load the dataset globally so the UI Dropdown knows what sectors exist
print("Loading Dataset...")
df = pd.DataFrame(load_dataset('Yoel125/synthetic-companies-12k', split='train'))
df['full_text'] = df['sector'] + " - " + df['theme'] + ": " + df['description']
all_sectors = ["All Sectors"] + sorted(list(df['sector'].unique()))

# Global caches for the heavy AI models
embedding_model_cache = None
faiss_index_cache = None
gen_tokenizer_cache = None
gen_model_cache = None

def get_ai_models():
    global embedding_model_cache, faiss_index_cache, gen_tokenizer_cache, gen_model_cache
    
    if embedding_model_cache is None:
        print("Lazy-loading Embedding Model...")
        embedding_model_cache = SentenceTransformer('paraphrase-MiniLM-L3-v2', device='cpu')
        
    if faiss_index_cache is None:
        try:
            print("Loading FAISS index...")
            embeddings = np.load('company_embeddings.npy')
            faiss_index_cache = faiss.IndexFlatL2(embeddings.shape[1])
            faiss_index_cache.add(np.array(embeddings).astype('float32'))
        except FileNotFoundError:
            print("Generating new FAISS index...")
            embeddings = embedding_model_cache.encode(df['full_text'].tolist(), show_progress_bar=False)
            faiss_index_cache = faiss.IndexFlatL2(embeddings.shape[1])
            faiss_index_cache.add(np.array(embeddings).astype('float32'))

    if gen_model_cache is None:
        print("Lazy-loading State-of-the-Art GenAI Model (Qwen2.5-0.5B-Instruct)...")
        # UPGRADED to a massively smarter, modern Causal LLM (ChatGPT equivalent for small models)
        gen_tokenizer_cache = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
        gen_model_cache = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct", torch_dtype="auto")
        gen_model_cache.eval()
        
    return embedding_model_cache, faiss_index_cache, gen_tokenizer_cache, gen_model_cache


def generate_sales_pitch(user_query, company_name, sector, theme, description, tokenizer, model):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    # Modern ChatML format used by state-of-the-art models like Qwen and Llama
    messages = [
        {"role": "system", "content": "You are a brilliant, aggressive Wall Street investment analyst. Your job is to write a single, highly persuasive, creative sentence explaining why a company is a massive investment opportunity."},
        {"role": "user", "content": f"Company: {company_name}\nIndustry: {theme}\nWhat they do: {description[:300]}\n\nWrite a 1-sentence sales pitch explaining why this company is the ultimate strategic investment for someone focused on '{user_query}'. Do not just summarize what they do. Be creative and aggressive."}
    ]
    
    # Apply the exact chat template the model was trained on
    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

    try:
        model.to(device)
        inputs = tokenizer(prompt, return_tensors="pt").to(device)

        # 100% STRICT DETERMINISM: Forces the exact same brilliant output every time!
        torch.manual_seed(42)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(42)
            torch.backends.cudnn.deterministic = True
            torch.backends.cudnn.benchmark = False

        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=75,
                do_sample=False,
                repetition_penalty=1.1,
            )
            
        # Causal LMs output the prompt + generation. We slice off the prompt.
        input_length = inputs.input_ids.shape[1]
        pitch = tokenizer.decode(outputs[0][input_length:], skip_special_tokens=True).strip()

        # Clean up quotes
        if pitch and pitch[0] == '"' and pitch[-1] == '"':
            pitch = pitch[1:-1]
            
        # Just in case the AI gets chatty, force it to 1-2 sentences max
        sentences = pitch.split(". ")
        if len(sentences) > 2:
            pitch = ". ".join(sentences[:2]) + "."

        if pitch:
            pitch = pitch[0].upper() + pitch[1:]
            
        return pitch

    except Exception as e:
        print(f"GenAI Generation Error: {e}")
        return f"By leveraging their advanced {theme} capabilities, {company_name} is perfectly positioned to capture explosive growth and completely dominate the '{user_query}' space."


# --- 2. RECOMMENDATION ENGINE LOGIC ---
@spaces.GPU  # <--- Hugging Face ZeroGPU Decorator!
def recommend_investment(user_query, selected_sector, top_k=3):
    if (not user_query or not str(user_query).strip()) and selected_sector == "All Sectors":
        yield "Please enter an investment thesis or keyword in the text box above, or select a specific industry sector from the dropdown menu."
        return

    if user_query and str(user_query).strip():
        non_english_chars = sum(1 for char in str(user_query) if ord(char) > 127)
        if non_english_chars > 2:
            yield "⚠️ **Language Not Supported:** SectorSync AI is currently optimized exclusively for English data. Please write your investment thesis in English and try again."
            return

    if not user_query or not str(user_query).strip():
        user_query = f"innovative {selected_sector} companies"

    # Streaming a status message so you know why the first click takes a few seconds!
    yield "🚀 Initializing Smart AI Engine... (The very first search takes ~10 seconds to load the heavy AI. Future searches will be instant!)"

    # Boot up the heavy AI models safely!
    embedding_model, faiss_index, gen_tokenizer, gen_model = get_ai_models()
    
    yield "🔍 Searching for matching companies and writing smart pitches..."

    if selected_sector != "All Sectors":
        enriched_query = f"{selected_sector} industry B2B company specializing in: {user_query}"
    else:
        enriched_query = f"B2B investment opportunity specializing in: {user_query}"

    query_vector = embedding_model.encode([enriched_query])
    search_k = 25 if selected_sector != "All Sectors" else top_k
    distances, indices = faiss_index.search(np.array(query_vector).astype('float32'), k=search_k)

    matches = []
    for i in range(search_k):
        idx = indices[0][i]
        sim_score = 1 / (1 + distances[0][i])
        company_row = df.iloc[idx]
        if selected_sector != "All Sectors" and company_row['sector'] != selected_sector:
            continue
        matches.append((company_row, sim_score))
        if len(matches) == top_k:
            break

    if not matches:
        for i in range(min(top_k, len(indices[0]))):
            idx = indices[0][i]
            sim_score = 1 / (1 + distances[0][i])
            matches.append((df.iloc[idx], sim_score))

    output_markdown = f"### Top {len(matches)} AI-Recommended Matches for: *'{user_query}'*\n\n"
    if selected_sector != "All Sectors":
        output_markdown += f"**Filtered by Sector:** `{selected_sector}`\n\n---\n\n"
    else:
        output_markdown += "---\n\n"

    match_labels = ["🥇 Strongest Match", "🥈 Close Match", "🥉 Close Match"]

    for rank, (row, score) in enumerate(matches, 1):
        c_name = row['company_name']
        ticker = row.get('ticker', 'N/A')
        sector = row['sector']
        theme = row['theme']
        desc = row['description']

        try:
            pitch = generate_sales_pitch(user_query, c_name, sector, theme, desc, gen_tokenizer, gen_model)
        except Exception:
            pitch = f"By leveraging their advanced {theme} capabilities, {c_name} is perfectly positioned to capture explosive growth and completely dominate the '{user_query}' space."

        tier_label = match_labels[rank - 1] if rank <= len(match_labels) else "Match"

        output_markdown += f"#### #{rank}. {c_name} (`{ticker}`) — *{sector}*\n\n"
        output_markdown += f"**{tier_label}** (Similarity Score: `{score*100:.1f}%`)\n\n"
        output_markdown += f"**Industry:** {theme}\n\n"
        output_markdown += f"**GenAI Investment Pitch:** *'{pitch}'*\n\n"
        output_markdown += f"**Company Overview:** {desc[:200]}...\n\n---\n\n"

        yield output_markdown


# --- 3. GRADIO USER INTERFACE ---
custom_theme = gr.themes.Base(
    primary_hue="emerald",
    neutral_hue="slate"
).set(
    body_background_fill="#121212",
    body_background_fill_dark="#121212",
    body_text_color="#f0f4f8",
    body_text_color_dark="#f0f4f8",
    background_fill_primary="#1e1e1e",
    background_fill_primary_dark="#1e1e1e",
    background_fill_secondary="#121212",
    background_fill_secondary_dark="#121212",
    border_color_primary="#333333",
    border_color_primary_dark="#333333",
    block_background_fill="#1e1e1e",
    block_background_fill_dark="#1e1e1e",
    block_label_text_color="#f0f4f8",
    block_label_text_color_dark="#f0f4f8",
    input_background_fill="#2a2a2a",
    input_background_fill_dark="#2a2a2a",
    button_primary_background_fill="#10b981",
    button_primary_background_fill_dark="#10b981",
    button_primary_text_color="#121212",
    button_primary_text_color_dark="#121212",
    table_even_background_fill="#1e1e1e",
    table_even_background_fill_dark="#1e1e1e",
    table_odd_background_fill="#121212",
    table_odd_background_fill_dark="#121212",
    table_border_color="#333333",
    table_border_color_dark="#333333"
)

# THE ULTIMATE CSS FIX for the invisible text!
# This ensures that no matter what mode the OS is in, the little boxes 
# holding the Ticker and Similarity Score are forced to have a dark background 
# and a bright emerald green text color so they pop out beautifully.
custom_css = """
table, table.dataset, tbody, thead, tr, td, th {
    background-color: #1e1e1e !important;
    color: #f0f4f8 !important;
    border-color: #333333 !important;
}

code, pre {
    background-color: #2a2a2a !important;
    color: #10b981 !important;
    border: 1px solid #333333 !important;
    padding: 2px 6px !important;
    border-radius: 4px !important;
}
"""

with gr.Blocks(theme=custom_theme, css=custom_css, title="SectorSync AI") as demo:
    gr.Markdown("# SectorSync AI")
    gr.Markdown("Cut Through the Market Noise — Find Your Next Winning Stock in Seconds. Discover high-growth companies matching your investment thesis using FAISS Vector Search and Generative AI.")
    gr.Markdown("*Similarity Score reflects how closely a company profile matches your query in AI-embedding space — a relative ranking signal, not a calibrated financial confidence rating.*")

    with gr.Row():
        with gr.Column(scale=2):
            query_input = gr.Textbox(
                label="What is your investment thesis or topic?",
                placeholder="e.g., autonomous robotics, clean battery storage, gene therapy for rare diseases...",
                lines=2
            )
        with gr.Column(scale=1):
            sector_dropdown = gr.Dropdown(
                choices=all_sectors,
                value="All Sectors",
                label="Sector"
            )

    search_button = gr.Button("Find Investment Matches", variant="primary")

    gr.Examples(
        examples=[
            ["Artificial Intelligence and Machine Learning in Healthcare", "Healthcare"],
            ["Next-generation Renewable Energy and Solar Battery Storage", "Energy"],
            ["Autonomous Robotics and Supply Chain Logistics Automation", "Industrials"]
        ],
        inputs=[query_input, sector_dropdown],
        label="Quick Starters (1-Click Example Searches)"
    )

    results_output = gr.Markdown(label="Recommendation Results")

    search_button.click(fn=recommend_investment, inputs=[query_input, sector_dropdown], outputs=results_output)
    query_input.submit(fn=recommend_investment, inputs=[query_input, sector_dropdown], outputs=results_output)

if __name__ == "__main__":
    print("Launching SectorSync AI Recommender App...")
    demo.launch(share=False)