import streamlit as st
import yfinance as yf
from langchain_groq import ChatGroq
import spacy
import pandas as pd
import os
# Load spaCy model
nlp = spacy.load("en_core_web_sm")
from dotenv import load_dotenv
load_dotenv()
# Replace with your actual API key
LLM = ChatGroq(
    temperature=0.7,
    model="llama-3.1-70b-versatile",
    api_key=os.environ['GROQ_API_KEY'],
    verbose=True
)

# Function to extract key investment parameters
def extract_investment_params(text):
    doc = nlp(text)
    investment_params = {
        "goal": None,
        "risk_tolerance": None,
        "amount": None,
        "horizon": None,
        "sectors": [],
        "volatility_tolerance": None
    }

    # Example extraction logic
    for ent in doc.ents:
        if ent.label_ == "MONEY":
            investment_params["amount"] = ent.text
        elif ent.label_ == "DATE":
            investment_params["horizon"] = ent.text
    
    # Extract sectors (this is a simple approach and might need refinement)
    sectors = ["technology", "healthcare", "finance", "energy", "consumer"]
    for sector in sectors:
        if sector in text.lower():
            investment_params["sectors"].append(sector)

    return investment_params

# Set default parameters if not provided by user
def set_default_params(params):
    if not params["goal"]:
        params["goal"] = "medium-term"
    if not params["risk_tolerance"]:
        params["risk_tolerance"] = "medium"
    if not params["amount"]:
        params["amount"] = "10000"
    if not params["horizon"]:
        params["horizon"] = "1 year"
    if not params["sectors"]:
        params["sectors"] = ["technology", "healthcare"]
    if not params["volatility_tolerance"]:
        params["volatility_tolerance"] = 0.2  # Assume 20% as medium volatility tolerance
    return params

# Fetch top stocks using yfinance library
def fetch_top_stocks(sectors, num_stocks=5):
    sector_tickers = {
        "technology": ["AAPL", "MSFT", "GOOGL", "AMZN", "NVDA", "TSLA", "META", "ADBE", "CRM", "INTC"],
        "healthcare": ["JNJ", "UNH", "PFE", "ABT", "TMO", "MRK", "ABBV", "DHR", "BMY", "AMGN"],
        "finance": ["JPM", "BAC", "WFC", "C", "GS", "MS", "BLK", "SPGI", "AXP", "V"],
        "energy": ["XOM", "CVX", "COP", "SLB", "EOG", "MPC", "PSX", "VLO", "OXY", "KMI"],
        "consumer": ["PG", "KO", "PEP", "COST", "WMT", "HD", "MCD", "NKE", "SBUX", "DIS"]
    }

    all_tickers = []
    for sector in sectors:
        all_tickers.extend(sector_tickers.get(sector, []))

    if not all_tickers:
        return []

    data = yf.download(all_tickers, period="1mo")
    market_caps = {}

    for ticker in all_tickers:
        try:
            info = yf.Ticker(ticker).info
            market_cap = info.get('marketCap', 0)
            market_caps[ticker] = market_cap
        except:
            print(f"Error fetching data for {ticker}")

    sorted_tickers = sorted(market_caps, key=market_caps.get, reverse=True)
    return sorted_tickers[:num_stocks]

# Predict stock prices (using simple moving average)
def predict_stock_prices(stocks, horizon):
    predictions = {}
    for stock in stocks:
        ticker = yf.Ticker(stock)
        hist = ticker.history(period="1y")
        current_price = hist['Close'].iloc[-1]
        sma = hist['Close'].rolling(window=30).mean().iloc[-1]
        predicted_price = sma * (1 + 0.1)  # Assuming 10% growth, adjust as needed
        predictions[stock] = f"Current: ${current_price:.2f}, Predicted: ${predicted_price:.2f}"
    return predictions

# Calculate volatility for each stock
def calculate_volatility(stocks):
    volatilities = {}
    for stock in stocks:
        ticker = yf.Ticker(stock)
        hist = ticker.history(period="1y")
        volatilities[stock] = hist['Close'].pct_change().std()
    return volatilities

# Analyze stocks based on volatility
def analyze_stocks(stocks, predictions, volatilities, volatility_tolerance):
    filtered_stocks = []
    for stock in stocks:
        if volatilities[stock] < volatility_tolerance:
            filtered_stocks.append(stock)
    return filtered_stocks

# Summarize the stock analysis
def summarize_analysis(stocks, predictions, volatilities):
    summary = "Based on your preferences, here are the recommended stocks:\n\n"
    for stock in stocks:
        summary += f"{stock}: {predictions[stock]}, Volatility: {volatilities[stock]:.2%}\n"
    return summary

# Generate advice using ChatGroq model
def generate_advice(summary):
    prompt = f'''Hi you are an helpful assistant, you have to greet the user before answering and be polite and provide me investment
    advise based on : {summary}'''
    response = LLM.invoke(prompt)
    return response.content

# Main app layout
st.title("Investment Strategy Chatbot")
st.write("Start by entering your investment preferences, then ask follow-up questions about the advice.")

# User input
message = st.text_input("Enter your investment preferences:")
if message:
    params = extract_investment_params(message)
    params = set_default_params(params)
    stocks = fetch_top_stocks(params["sectors"])
    predictions = predict_stock_prices(stocks, params["horizon"])
    volatilities = calculate_volatility(stocks)
    analyzed_stocks = analyze_stocks(stocks, predictions, volatilities, params["volatility_tolerance"])
    summary = summarize_analysis(analyzed_stocks, predictions, volatilities)
    advice = generate_advice(summary)

    # Display results
    st.write(advice)

# Handle follow-up questions
follow_up = st.text_input("Ask any follow-up questions:")
if follow_up:
    follow_up_response = LLM.invoke(follow_up)
    st.write(follow_up_response.content)