import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
#from langchain.prompts import PromptTemplate
from langchain_core.prompts import PromptTemplate
from huggingface_hub import hf_hub_download
import os
from huggingface_hub import login

hf_token = os.environ["HF_TOKEN"]
login(token=hf_token)

model_name = "SelmaNajih001/GRPORagMinstral2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
pipe = pipeline(
    "text-generation",
    model=model_name,
    tokenizer=model_name,
    max_new_tokens=600,
    temperature=0.0,
    do_sample=False,
    num_beams=6,
    repetition_penalty=1.5,
    return_full_text=False,
    eos_token_id=tokenizer.eos_token_id
)

# Prompt template
prompt_template = """
You are a financial market analyst.
Before making a prediction, you must analyze the past, provided in the Context below.
Your goal is to identify similar historical situations and use them to infer what may happen next.
Your analysis must be comprehensive, covering macroeconomic, sectoral, and corporate-specific factors.
You should identify past periods that closely resemble the current environment (e.g., "high inflation + geopolitical conflict" or "rate hikes + tech earnings slump").
Base your reasoning on actual market reactions from those periods — specify which companies or sectors moved and how.
If multiple scenarios are possible, explain each one and why the market may react differently under varying conditions.
Explicitly name the historical reference period(s) used (e.g., "2008 financial crisis," "2020 pandemic crash and recovery," etc.).

Response Format:

-Chosen Stock or List of Stocks:
(name/names)

-Prediction(s):
(expected price change or direction)

-Explanation:
A concise, factual analysis linking the historical precedent to the current conditions.
Mention the relevant macroeconomic, sector, and corporate factors, explaining how they interacted in the past and why similar outcomes may occur again. "

Here the Context:
{context}

Question:
What could happen after: {question}
"""
prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/multi-qa-mpnet-base-dot-v1")


import tempfile
tmp_dir = tempfile.mkdtemp()



vectorstore = FAISS.load_local(
    folder_path=".",  
    embeddings=embeddings,
    allow_dangerous_deserialization=True
)

def analisi_finanziaria(query, k=3):

    docs_found = vectorstore.similarity_search(query, k=k)

    context = "\n".join([doc.page_content for doc in docs_found])

    final_prompt = prompt.format(context=context, question=query)

    # Genera risposta
    result = pipe(final_prompt)[0]['generated_text']
    return result
    
import gradio as gr

examples = [
    "Trump imposes tariffs",
    "Tesla announces a new affordable electric vehicle model",
    "Nvidia releases new GPU technology",
    "Apple launches Apple TV+ subscription service",
    "Elon Musk created a new political party to run against Trump"
]

description_md = """
This tool analyzes financial events using a retrieval-augmented **language model**.

**How it works:**
- The model leverages historical events and news to provide predictions.
- For each input, similar past events are retrieved to give context.
- The output includes:
  - **Chosen Stocks**: stocks likely impacted
  - **Prediction**: expected price change
  - **Explanation**: brief reasoning based on historical context

**Example use cases:**
- Market reactions to political events
- Corporate announcements and earnings reports
- Technological product launches

Click an example below to quickly test the model.
"""

# Gradio interface
iface = gr.Interface(
    fn=analisi_finanziaria, 
    inputs=gr.Textbox(label="Enter financial event", placeholder="Type an event here..."),
    outputs=gr.Textbox(label="Prediction"),
    title="GRPO Financial Analyst",
    description=description_md,
    examples=[[e] for e in examples],
    allow_flagging="never"
)

iface.launch()