app.py

from typing import Dict, Any, List, Tuple
import streamlit as st
import yfinance as yf
import pandas as pd
import plotly.graph_objects as go
import plotly.express as px
from datetime import datetime, timedelta
from newsapi.newsapi_client import NewsApiClient
import requests
import os
import sqlite3
from sqlite3 import Error
import json
import openai
from prophet import Prophet
from prophet.plot import plot_plotly
from sklearn.preprocessing import StandardScaler, MinMaxScaler
import asyncio
import aiohttp
from functools import partial
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch_geometric.nn import GCNConv
from torch_geometric.data import Data
import skfuzzy as fuzz
import skfuzzy.control as ctrl
import numpy as np
import networkx as nx
import random
import matplotlib.pyplot as plt
from streamlit_autorefresh import st_autorefresh
import cvxpy as cp  # For portfolio optimization
from sklearn.ensemble import IsolationForest  # For anomaly detection

# ----------------------------
# Configuration and Constants
# ----------------------------

# Load environment variables
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
FMP_API_KEY = os.getenv("FMP_API_KEY")
NEWS_API_KEY = os.getenv("NEWS_API_KEY")

if not OPENAI_API_KEY or not FMP_API_KEY or not NEWS_API_KEY:
    st.error(
        "API keys for OpenAI, Financial Modeling Prep, and NewsAPI are not set. Please set them in the `.streamlit/secrets.toml` file."
    )
    st.stop()

# Initialize OpenAI
openai.api_key = OPENAI_API_KEY

# Initialize NewsApiClient
newsapi = NewsApiClient(api_key=NEWS_API_KEY)

# Database Configuration
DATABASE = "stock_dashboard.db"

# ----------------------------
# API Rate Limits Configuration
# ----------------------------

# Define API rate limits (example limits; adjust based on your subscription)
API_RATE_LIMITS = {
    "FMP": {
        "max_requests_per_day": 500,
        "current_count": 0,
        "last_reset": datetime.utcnow().date()
    },
    "NewsAPI": {
        "max_requests_per_day": 500,
        "current_count": 0,
        "last_reset": datetime.utcnow().date()
    },
    "OpenAI": {
        "max_requests_per_day": 1000,
        "current_count": 0,
        "last_reset": datetime.utcnow().date()
    }
}

# ----------------------------
# Helper Functions
# ----------------------------

def local_css():
    """Injects custom CSS for enhanced styling."""
    st.markdown(
        """
        <style>
        /* Sidebar Styling */
        .css-1d391kg { 
            background-color: #f0f2f6;
        }
        /* Header Styling */
        .title {
            font-size: 3rem;
            text-align: center;
            color: #2e86de;
            margin-bottom: 0;
        }
        .description {
            text-align: center;
            color: #555555;
            margin-top: 0;
            margin-bottom: 2rem;
            font-size: 1.2rem;
        }
        /* DataFrame Styling */
        .dataframe th, .dataframe td {
            text-align: center;
            padding: 10px;
        }
        /* Footer Styling */
        .footer {
            text-align: center;
            color: #888888;
            margin-top: 3rem;
            font-size: 0.9rem;
        }
        /* Dark Mode Styling */
        .dark-mode {
            background-color: #1e1e1e;
            color: #ffffff;
        }
        /* Chat Interface Styling */
        .chat-container {
            max-height: 500px;
            overflow-y: auto;
            padding: 10px;
            border: 1px solid #ccc;
            border-radius: 5px;
            background-color: #f9f9f9;
            margin-bottom: 1rem;
        }
        .user-message {
            text-align: right;
            margin: 5px 0;
            color: #2e86de;
        }
        .assistant-message {
            text-align: left;
            margin: 5px 0;
            color: #e74c3c;
        }
        /* Tooltip Styling */
        .tooltip {
            position: relative;
            display: inline-block;
            border-bottom: 1px dotted black;
        }

        .tooltip .tooltiptext {
            visibility: hidden;
            width: 220px;
            background-color: #555;
            color: #fff;
            text-align: left;
            border-radius: 6px;
            padding: 5px;
            position: absolute;
            z-index: 1;
            bottom: 125%; /* Position above */
            left: 50%;
            margin-left: -110px;
            opacity: 0;
            transition: opacity 0.3s;
        }

        .tooltip:hover .tooltiptext {
            visibility: visible;
            opacity: 1;
        }

        /* Button Styling */
        .css-1aumxhk {
            background-color: #2e86de;
            color: white;
        }
        </style>
        """,
        unsafe_allow_html=True,
    )

def initialize_database():
    """
    Initializes the SQLite database and creates necessary tables if they don't exist.
    """
    try:
        conn = sqlite3.connect(DATABASE)
        cursor = conn.cursor()
        # Create interactions table
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS interactions (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                timestamp TEXT NOT NULL,
                user_input TEXT NOT NULL,
                assistant_response TEXT NOT NULL
            )
        """)
        # Create stock_cache table
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS stock_cache (
                ticker TEXT PRIMARY KEY,
                fetched_at TEXT NOT NULL,
                data TEXT NOT NULL
            )
        """)
        # Create portfolio table
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS portfolio (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                ticker TEXT NOT NULL,
                added_at TEXT NOT NULL
            )
        """)
        # Create api_usage table
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS api_usage (
                api_name TEXT PRIMARY KEY,
                request_count INTEGER NOT NULL,
                last_reset TEXT NOT NULL
            )
        """)
        # Initialize api_usage records if they don't exist
        for api in API_RATE_LIMITS.keys():
            cursor.execute("""
                INSERT OR IGNORE INTO api_usage (api_name, request_count, last_reset)
                VALUES (?, ?, ?)
            """, (api, 0, datetime.utcnow().date().isoformat()))
        conn.commit()
        conn.close()
    except Error as e:
        st.error(f"Error initializing database: {e}")
        st.stop()

def insert_interaction(user_input, assistant_response):
    """
    Inserts a user interaction into the interactions table.
    """
    try:
        conn = sqlite3.connect(DATABASE)
        cursor = conn.cursor()
        cursor.execute("""
            INSERT INTO interactions (timestamp, user_input, assistant_response)
            VALUES (?, ?, ?)
        """, (datetime.utcnow().isoformat(), user_input, assistant_response))
        conn.commit()
        conn.close()
    except Error as e:
        st.error(f"Error inserting interaction: {e}")

def fetch_interactions(limit=50):
    """
    Fetches the most recent user interactions.
    """
    try:
        conn = sqlite3.connect(DATABASE)
        cursor = conn.cursor()
        cursor.execute("""
            SELECT timestamp, user_input, assistant_response
            FROM interactions
            ORDER BY id DESC
            LIMIT ?
        """, (limit,))
        rows = cursor.fetchall()
        conn.close()
        return rows
    except Error as e:
        st.error(f"Error fetching interactions: {e}")
        return []

def clear_interactions():
    """
    Clears all records from the interactions table.
    """
    try:
        conn = sqlite3.connect(DATABASE)
        cursor = conn.cursor()
        cursor.execute("DELETE FROM interactions")
        conn.commit()
        conn.close()
        st.success("All interactions have been cleared.")
    except Error as e:
        st.error(f"Error clearing interactions: {e}")

def insert_stock_cache(ticker, data):
    """
    Inserts or updates stock data in the stock_cache table.
    """
    try:
        conn = sqlite3.connect(DATABASE)
        cursor = conn.cursor()
        cursor.execute("""
            INSERT INTO stock_cache (ticker, fetched_at, data)
            VALUES (?, ?, ?)
            ON CONFLICT(ticker) DO UPDATE SET
                fetched_at=excluded.fetched_at,
                data=excluded.data
        """, (ticker.upper(), datetime.utcnow().isoformat(), json.dumps(data, default=str)))
        conn.commit()
        conn.close()
    except Error as e:
        st.error(f"Error inserting stock cache: {e}")

def fetch_stock_cache(ticker):
    """
    Fetches cached stock data for a given ticker.
    """
    try:
        conn = sqlite3.connect(DATABASE)
        cursor = conn.cursor()
        cursor.execute("""
            SELECT fetched_at, data
            FROM stock_cache
            WHERE ticker = ?
        """, (ticker.upper(),))
        row = cursor.fetchone()
        conn.close()
        if row:
            fetched_at, data = row
            return json.loads(data), fetched_at
        return None, None
    except Error as e:
        st.error(f"Error fetching stock cache: {e}")
        return None, None

def clear_stock_cache():
    """
    Clears all records from the stock_cache table.
    """
    try:
        conn = sqlite3.connect(DATABASE)
        cursor = conn.cursor()
        cursor.execute("DELETE FROM stock_cache")
        conn.commit()
        conn.close()
        st.success("All cached stock data has been cleared.")
    except Error as e:
        st.error(f"Error clearing stock cache: {e}")

def add_to_portfolio(ticker):
    """
    Adds a ticker to the user's portfolio.
    """
    try:
        conn = sqlite3.connect(DATABASE)
        cursor = conn.cursor()
        cursor.execute("""
            INSERT INTO portfolio (ticker, added_at)
            VALUES (?, ?)
        """, (ticker.upper(), datetime.utcnow().isoformat()))
        conn.commit()
        conn.close()
        st.success(f"{ticker.upper()} has been added to your portfolio.")
    except Error as e:
        st.error(f"Error adding to portfolio: {e}")

def remove_from_portfolio(ticker):
    """
    Removes a ticker from the user's portfolio.
    """
    try:
        conn = sqlite3.connect(DATABASE)
        cursor = conn.cursor()
        cursor.execute("""
            DELETE FROM portfolio
            WHERE ticker = ?
        """, (ticker.upper(),))
        conn.commit()
        conn.close()
        st.success(f"{ticker.upper()} has been removed from your portfolio.")
    except Error as e:
        st.error(f"Error removing from portfolio: {e}")

def fetch_portfolio():
    """
    Fetches the user's portfolio.
    """
    try:
        conn = sqlite3.connect(DATABASE)
        cursor = conn.cursor()
        cursor.execute("""
            SELECT ticker
            FROM portfolio
        """)
        rows = cursor.fetchall()
        conn.close()
        return [row[0] for row in rows]
    except Error as e:
        st.error(f"Error fetching portfolio: {e}")
        return []

def update_api_usage(api_name):
    """
    Increments the API usage count for the specified API and checks rate limits.
    Returns True if the API call is allowed, False otherwise.
    """
    try:
        conn = sqlite3.connect(DATABASE)
        cursor = conn.cursor()
        cursor.execute("""
            SELECT request_count, last_reset
            FROM api_usage
            WHERE api_name = ?
        """, (api_name,))
        row = cursor.fetchone()
        if row:
            request_count, last_reset = row
            last_reset_date = datetime.fromisoformat(last_reset).date()
            today = datetime.utcnow().date()
            if last_reset_date < today:
                # Reset the count
                cursor.execute("""
                    UPDATE api_usage
                    SET request_count = 1, last_reset = ?
                    WHERE api_name = ?
                """, (today.isoformat(), api_name))
                conn.commit()
                conn.close()
                API_RATE_LIMITS[api_name]["current_count"] = 1
                API_RATE_LIMITS[api_name]["last_reset"] = today
                return True
            else:
                if request_count < API_RATE_LIMITS[api_name]["max_requests_per_day"]:
                    # Increment the count
                    cursor.execute("""
                        UPDATE api_usage
                        SET request_count = request_count + 1
                        WHERE api_name = ?
                    """, (api_name,))
                    conn.commit()
                    conn.close()
                    API_RATE_LIMITS[api_name]["current_count"] += 1
                    return True
                else:
                    # Rate limit exceeded
                    conn.close()
                    st.warning(f"{api_name} API rate limit exceeded for today.")
                    return False
        else:
            # API not found in usage table
            conn.close()
            st.error(f"API usage record for {api_name} not found.")
            return False
    except Error as e:
        st.error(f"Error updating API usage: {e}")
        return False

async def fetch_single_stock_data(session, ticker):
    """
    Asynchronously fetches stock data for a single ticker, respecting API rate limits.
    """
    if not update_api_usage("FMP"):
        st.warning(f"Financial Modeling Prep API rate limit exceeded. Skipping {ticker.upper()}.")
        return ticker, {}
    cached_data, fetched_at = fetch_stock_cache(ticker)
    if cached_data:
        return ticker, cached_data
    else:
        try:
            stock = yf.Ticker(ticker)
            info = stock.info
            insert_stock_cache(ticker, info)
            return ticker, info
        except Exception as e:
            st.error(f"Error fetching data for {ticker}: {e}")
            return ticker, {}

async def fetch_all_stock_data(tickers):
    """
    Asynchronously fetches stock data for all tickers.
    """
    async with aiohttp.ClientSession() as session:
        tasks = []
        for ticker in tickers:
            task = asyncio.ensure_future(fetch_single_stock_data(session, ticker))
            tasks.append(task)
        responses = await asyncio.gather(*tasks)
        return responses

# -----------------------------
# Fuzzy Logic for Mutation Rate
# -----------------------------

def prepare_graph(num_nodes=10):
    """
    Prepares a graph with edge weights.
    """
    graph = nx.complete_graph(num_nodes)
    for u, v in graph.edges:
        graph[u][v]['weight'] = random.randint(1, 100)  # Assign random weights
    return graph

def visualize_fuzzy_logic(controller):
    """
    Visualizes fuzzy membership functions.
    """
    x = np.arange(0, 1.01, 0.01)
    fig, ax = plt.subplots(figsize=(8, 5))
    ax.plot(x, fuzz.trapmf(x, [0, 0, 0.3, 0.5]), label="Low Uncertainty")
    ax.plot(x, fuzz.trimf(x, [0.3, 0.5, 0.7]), label="Medium Uncertainty")
    ax.plot(x, fuzz.trapmf(x, [0.5, 0.7, 1, 1]), label="High Uncertainty")
    ax.legend()
    ax.set_title("Fuzzy Membership Functions")
    ax.set_xlabel("Uncertainty")
    ax.set_ylabel("Membership Degree")
    st.pyplot(fig)

class FuzzyMutationController:
    def __init__(self):
        # Define fuzzy variables
        self.uncertainty = ctrl.Antecedent(np.arange(0, 1.01, 0.01), 'uncertainty')
        self.mutation_rate = ctrl.Consequent(np.arange(0, 1.01, 0.01), 'mutation_rate')

        # Define membership functions
        self.uncertainty['low'] = fuzz.trapmf(self.uncertainty.universe, [0, 0, 0.3, 0.5])
        self.uncertainty['medium'] = fuzz.trimf(self.uncertainty.universe, [0.3, 0.5, 0.7])
        self.uncertainty['high'] = fuzz.trapmf(self.uncertainty.universe, [0.5, 0.7, 1, 1])

        self.mutation_rate['low'] = fuzz.trapmf(self.mutation_rate.universe, [0, 0, 0.1, 0.3])
        self.mutation_rate['medium'] = fuzz.trimf(self.mutation_rate.universe, [0.1, 0.3, 0.5])
        self.mutation_rate['high'] = fuzz.trapmf(self.mutation_rate.universe, [0.3, 0.5, 1, 1])

        # Define fuzzy rules
        rule1 = ctrl.Rule(self.uncertainty['low'], self.mutation_rate['low'])
        rule2 = ctrl.Rule(self.uncertainty['medium'], self.mutation_rate['medium'])
        rule3 = ctrl.Rule(self.uncertainty['high'], self.mutation_rate['high'])

        # Create control system and simulation
        self.mutation_ctrl = ctrl.ControlSystem([rule1, rule2, rule3])
        self.mutation_sim = ctrl.ControlSystemSimulation(self.mutation_ctrl)

    def compute_mutation_rate(self, uncertainty_value):
        # Ensure uncertainty_value is within [0, 1]
        uncertainty_value = min(max(uncertainty_value, 0), 1)
        self.mutation_sim.input['uncertainty'] = uncertainty_value
        self.mutation_sim.compute()
        mutation_rate = self.mutation_sim.output['mutation_rate']
        # Ensure mutation_rate is within [0, 1]
        mutation_rate = min(max(mutation_rate, 0), 1)
        return mutation_rate

# -----------------------------
# Genetic Algorithm Definition
# -----------------------------

class GeneticAlgorithm:
    def __init__(self, graph, fuzzy_controller, population_size=50, generations=100):
        self.graph = graph
        self.nodes = list(graph.nodes)
        self.fuzzy_controller = fuzzy_controller
        self.population_size = population_size
        self.generations = generations

    def fitness(self, path):
        """
        Calculates the fitness of a path.
        """
        score = 0
        for i in range(len(path)):
            u = path[i]
            v = path[(i + 1) % len(path)]  # Wrap around to the start
            if self.graph.has_edge(u, v):
                score += self.graph[u][v].get('weight', 0)  # Use weight attribute
            else:
                score -= 1000  # Penalize missing edges heavily
        return score

    def run(self) -> Tuple[List[Any], List[float], List[float]]:
        # Initialize random population of paths
        population = [random.sample(self.nodes, len(self.nodes)) for _ in range(self.population_size)]
        best_fitness_history = []
        mutation_rate_history = []

        for generation in range(self.generations):
            # Evaluate fitness of the population
            fitness_scores = [self.fitness(path) for path in population]

            # Normalize fitness scores
            fitness_array = np.array(fitness_scores)
            fitness_mean = np.mean(fitness_array)
            fitness_std = np.std(fitness_array)
            uncertainty_value = fitness_std / (abs(fitness_mean) + 1e-6)
            # Normalize uncertainty_value to [0, 1]
            uncertainty_value = uncertainty_value / (uncertainty_value + 1)

            # Compute mutation rate using fuzzy logic
            mutation_rate = self.fuzzy_controller.compute_mutation_rate(uncertainty_value)

            # Select top candidates (elitism)
            sorted_population = [path for _, path in sorted(zip(fitness_scores, population), reverse=True)]
            population = sorted_population[:self.population_size // 2]

            # Generate new population through crossover and mutation
            new_population = population.copy()
            while len(new_population) < self.population_size:
                parent1, parent2 = random.sample(population, 2)
                child = self.crossover(parent1, parent2)
                child = self.mutate(child, mutation_rate)
                new_population.append(child)

            population = new_population

            # Record best fitness and mutation rate
            best_fitness = max(fitness_scores)
            best_fitness_history.append(best_fitness)
            mutation_rate_history.append(mutation_rate)

            # Update progress bar in Streamlit
            if 'ga_progress' in st.session_state:
                st.session_state.ga_progress.progress((generation + 1) / self.generations)
                st.session_state.ga_status.text(f"Generation {generation + 1}/{self.generations}")
                st.session_state.ga_chart.add_rows({"Best Fitness": best_fitness})

        # Return the best path found
        fitness_scores = [self.fitness(path) for path in population]
        best_index = fitness_scores.index(max(fitness_scores))
        best_path = population[best_index]
        return best_path, best_fitness_history, mutation_rate_history

    def crossover(self, parent1, parent2):
        # Ordered Crossover (OX)
        size = len(parent1)
        start, end = sorted(random.sample(range(size), 2))
        child = [None] * size

        # Copy a slice from parent1 to child
        child[start:end] = parent1[start:end]

        # Fill the remaining positions with genes from parent2
        ptr = end
        for gene in parent2:
            if gene not in child:
                if ptr >= size:
                    ptr = 0
                child[ptr] = gene
                ptr += 1
        return child

    def mutate(self, individual, mutation_rate):
        if random.random() < mutation_rate:
            idx1, idx2 = random.sample(range(len(individual)), 2)
            individual[idx1], individual[idx2] = individual[idx2], individual[idx1]
        return individual

# ----------------------------
# AI Assistant Integration
# ----------------------------

class RealAgent:
    """Main agent logic handling interactions with the OpenAI Assistant."""

    def __init__(self):
        self.agent_state = AgentState()

    def process(self, user_input: str, selected_tickers: List[str], stock_df: pd.DataFrame,
                historical_dfs: Dict[str, pd.DataFrame], news_articles: Dict[str, List[Dict[str, Any]]],
                fsirdm_data: Dict[str, Any]) -> str:
        """
        Processes user input and generates a response using the AI Assistant.
        """
        # Retrieve conversation history for context
        conversation_history = self.agent_state.get_conversation_history()

        # Generate additional context from the data
        additional_context = self.generate_additional_context(selected_tickers, stock_df, historical_dfs, news_articles, fsirdm_data)

        # Generate response using OpenAI
        response = generate_openai_response(conversation_history, user_input, additional_context)

        # Store interaction
        self.agent_state.store_interaction(user_input, response)
        insert_interaction(user_input, response)
        return response

    def generate_additional_context(self, selected_tickers: List[str], stock_df: pd.DataFrame,
                                    historical_dfs: Dict[str, pd.DataFrame], news_articles: Dict[str, List[Dict[str, Any]]],
                                    fsirdm_data: Dict[str, Any]) -> str:
        """
        Generates a concise summary of the selected stocks' data, historical data, news, and FSIRDM analysis to provide context to the assistant.
        """
        try:
            # Get stock data
            stock_data = stock_df[stock_df['Ticker'].isin(selected_tickers)].to_dict(orient='records')

            # Get historical data metrics
            historical_metrics: Dict[str, Any] = {}
            for ticker in selected_tickers:
                hist_df = historical_dfs.get(ticker, pd.DataFrame())
                if not hist_df.empty:
                    metrics = {
                        "Latest Close": hist_df['Close'].iloc[-1],
                        "52 Week High": hist_df['Close'].max(),
                        "52 Week Low": hist_df['Close'].min(),
                        "SMA 20": hist_df['SMA_20'].iloc[-1],
                        "EMA 20": hist_df['EMA_20'].iloc[-1],
                        "Volatility (Std Dev)": hist_df['Close'].std(),
                        "RSI": hist_df['RSI'].iloc[-1]
                    }
                else:
                    metrics = "No historical data available."
                historical_metrics[ticker] = metrics

            # Get FSIRDM analysis
            fsirdm_summary = fsirdm_data

            # Get latest news headlines
            news_summary: Dict[str, Any] = {}
            if news_articles and isinstance(news_articles, dict):
                for ticker, articles in news_articles.items():
                    if articles and isinstance(articles, list):
                        news_summary[ticker] = [
                            {
                                "Title": article['Title'],
                                "Description": article['Description'],
                                "URL": article['URL'],
                                "Published At": article['Published At']
                            }
                            for article in articles
                        ]
                    else:
                        news_summary[ticker] = "No recent news articles found."
            else:
                news_summary = "No recent news articles found."

            # Combine all summaries into a structured JSON format
            additional_context = {
                "Stock Data": stock_data,
                "Historical Metrics": historical_metrics,
                "FSIRDM Analysis": fsirdm_summary,
                "Latest News": news_summary
            }

            # Convert to JSON string with proper date formatting
            additional_context_json = json.dumps(additional_context, default=str, indent=4)

            # Truncate if necessary to stay within token limits
            max_length = 32000  # Adjust based on token count estimates
            if len(additional_context_json) > max_length:
                additional_context_json = additional_context_json[:max_length] + "..."

            return additional_context_json
        except Exception as e:
            st.error(f"Error generating additional context: {e}")
            return ""

class AgentState:
    """Manages the agent's memory and state."""

    def __init__(self):
        self.short_term_memory: List[Dict[str, str]] = self.load_memory()

    def load_memory(self) -> List[Dict[str, str]]:
        """
        Loads recent interactions from the database to provide context.
        """
        interactions = fetch_interactions(limit=50)
        memory: List[Dict[str, str]] = []
        for interaction in reversed(interactions):  # oldest first
            timestamp, user_input, assistant_response = interaction
            memory.append({
                "user_input": user_input,
                "assistant_response": assistant_response
            })
        return memory

    def store_interaction(self, user_input: str, response: str) -> None:
        """
        Stores a new interaction in the memory.
        """
        self.short_term_memory.append({"user_input": user_input, "assistant_response": response})
        if len(self.short_term_memory) > 10:
            self.short_term_memory.pop(0)

    def get_conversation_history(self) -> List[Dict[str, str]]:
        """
        Returns the current conversation history.
        """
        return self.short_term_memory

    def reset_memory(self) -> None:
        """
        Resets the short-term memory.
        """
        self.short_term_memory = []

def generate_openai_response(conversation_history: List[Dict[str, str]], user_input: str, additional_context: str) -> str:
    """
    Generates a response from OpenAI's GPT-4 model based on the conversation history, user input, and additional context.
    """
    try:
        # Prepare the messages for the model
        messages = [
            {
                "role": "system",
                "content": (
                    "You are a financial AI assistant leveraging the FS-IRDM framework to analyze real-time stock data and news sentiment. "
                    "You compare multiple stocks, quantify uncertainties with fuzzy memberships, and classify stocks into high-growth, stable, and risky categories. "
                    "Utilize transformation matrices and continuous utility functions to optimize portfolio decisions while conserving expected utility. "
                    "Dynamically adapt through sensitivity analysis and stochastic modeling of market volatility, ensuring decision integrity with homeomorphic mappings. "
                    "Refine utility functions based on investor preferences and risk aversion, employ advanced non-linear optimization and probabilistic risk analysis, "
                    "and ensure secure data handling with fuzzy logic-enhanced memory compression and diffeomorphic encryption. Additionally, provide robust, explainable recommendations "
                    "and comprehensive reporting, maintaining consistency, adaptability, and efficiency in dynamic financial environments. "
                    "Use and learn from every single interaction to better tune your strategy and optimize the portfolio."
                )
            }
        ]

        # Add additional context (stock data, historical data, news, FSIRDM analysis)
        if additional_context:
            messages.append({"role": "system", "content": f"Here is the relevant data:\n{additional_context}"})

        # Add past interactions
        for interaction in conversation_history:
            messages.append({"role": "user", "content": interaction['user_input']})
            messages.append({"role": "assistant", "content": interaction['assistant_response']})

        # Add the current user input
        messages.append({"role": "user", "content": user_input})

        # Call OpenAI API
        response = openai.ChatCompletion.create(
            model="gpt-4",
            messages=messages,
            max_tokens=1500,  # Adjust based on requirements
            n=1,
            stop=None,
            temperature=0.7,
        )

        assistant_response: str = response.choices[0].message['content'].strip()
        return assistant_response
    except Exception as e:
        st.error(f"Error generating response from OpenAI: {e}")
        return "I'm sorry, I couldn't process your request at the moment."

# ----------------------------
# Swarm of AI Agents Definitions
# ----------------------------

class DataAnalysisAgent:
    """Agent responsible for in-depth data analysis."""

    def analyze_market_trends(self, stock_df: pd.DataFrame) -> Dict[str, Any]:
        """
        Analyzes market trends based on current stock data.
        """
        try:
            analysis = {}
            # Example: Calculate average PE ratio
            avg_pe = stock_df['PE Ratio'].mean()
            analysis['Average PE Ratio'] = avg_pe

            # Example: Identify top-performing sectors
            sector_performance = stock_df.groupby('Sector')['Market Cap'].sum().reset_index()
            top_sectors = sector_performance.sort_values(by='Market Cap', ascending=False).head(3)
            analysis['Top Sectors by Market Cap'] = top_sectors.to_dict(orient='records')

            # Add more sophisticated analyses as needed
            return analysis
        except Exception as e:
            st.error(f"Error in Data Analysis Agent: {e}")
            return {}

class PredictionAgent:
    """Agent responsible for forecasting and predictions."""

    def forecast_stock_trends(self, historical_dfs: Dict[str, pd.DataFrame], forecast_period: int = 90) -> Dict[str, Any]:
        """
        Forecasts future stock trends using Prophet.
        """
        try:
            forecasts = {}
            for ticker, hist_df in historical_dfs.items():
                if hist_df.empty:
                    forecasts[ticker] = "No historical data available."
                    continue
                model, forecast_df = forecast_stock_price(hist_df, forecast_period)
                if model is not None and not forecast_df.empty:
                    forecasts[ticker] = forecast_df[['ds', 'yhat', 'yhat_lower', 'yhat_upper']].tail(30).to_dict(orient='records')
                else:
                    forecasts[ticker] = "Forecasting model could not generate predictions."
            return forecasts
        except Exception as e:
            st.error(f"Error in Prediction Agent: {e}")
            return {}

class SentimentAnalysisAgent:
    """Agent responsible for analyzing news sentiment."""

    def analyze_sentiment(self, news_articles: Dict[str, List[Dict[str, Any]]]) -> Dict[str, float]:
        """
        Analyzes sentiment of news articles using OpenAI's sentiment analysis.
        Returns sentiment scores per ticker.
        """
        sentiment_scores = {}
        try:
            for ticker, articles in news_articles.items():
                if not articles:
                    sentiment_scores[ticker] = 0.0  # Neutral sentiment
                    continue
                sentiments = []
                for article in articles:
                    prompt = f"Analyze the sentiment of the following news article and rate it from -1 (very negative) to 1 (very positive):\n\nTitle: {article['Title']}\nDescription: {article['Description']}"
                    sentiment = get_sentiment_from_openai(prompt)
                    sentiments.append(sentiment)
                # Average sentiment score
                if sentiments:
                    avg_sentiment = sum(sentiments) / len(sentiments)
                    sentiment_scores[ticker] = avg_sentiment
                else:
                    sentiment_scores[ticker] = 0.0
            return sentiment_scores
        except Exception as e:
            st.error(f"Error in Sentiment Analysis Agent: {e}")
            return sentiment_scores

class AnomalyDetectionAgent:
    """Agent responsible for detecting anomalies in stock data."""

    def detect_anomalies(self, stock_df: pd.DataFrame, historical_dfs: Dict[str, pd.DataFrame]) -> Dict[str, List[str]]:
        """
        Detects anomalies in stock data using Isolation Forest.
        Returns a dictionary with tickers as keys and list of anomaly dates as values.
        """
        anomalies = {}
        try:
            for ticker in stock_df['Ticker']:
                hist_df = historical_dfs.get(ticker, pd.DataFrame())
                if hist_df.empty:
                    anomalies[ticker] = ["No historical data available."]
                    continue
                # Use Isolation Forest on Close prices
                model = IsolationForest(contamination=0.05, random_state=42)
                hist_df = hist_df.sort_values(by='Date')
                hist_df['Close_Log'] = np.log(hist_df['Close'] + 1)  # Log transform to stabilize variance
                model.fit(hist_df[['Close_Log']])
                hist_df['Anomaly'] = model.predict(hist_df[['Close_Log']])
                anomaly_dates = hist_df[hist_df['Anomaly'] == -1]['Date'].tolist()
                anomalies[ticker] = anomaly_dates if len(anomaly_dates) > 0 else ["No anomalies detected."]
            return anomalies
        except Exception as e:
            st.error(f"Error in Anomaly Detection Agent: {e}")
            return anomalies

class PortfolioOptimizationAgent:
    """Agent responsible for optimizing the user's portfolio."""

    def optimize_portfolio(self, stock_df: pd.DataFrame, historical_dfs: Dict[str, pd.DataFrame]) -> Dict[str, Any]:
        """
        Optimizes the portfolio using mean-variance optimization.
        Returns the optimal weights for each stock.
        """
        try:
            tickers = stock_df['Ticker'].tolist()
            returns_data = []
            valid_tickers = []

            # Calculate daily returns for each ticker
            for ticker in tickers:
                hist_df = historical_dfs.get(ticker, pd.DataFrame())
                if hist_df.empty or len(hist_df) < 2:
                    st.warning(f"Not enough historical data for {ticker} to calculate returns.")
                    continue
                hist_df = hist_df.sort_values(by='Date')
                hist_df['Return'] = hist_df['Close'].pct_change()
                returns = hist_df['Return'].dropna()
                returns_data.append(returns)
                valid_tickers.append(ticker)

            # Check if there are enough tickers to optimize
            if len(valid_tickers) < 2:
                st.error("Not enough tickers with sufficient historical data for portfolio optimization.")
                return {"Optimal Weights": "Insufficient data."}

            # Create a DataFrame of returns
            returns_df = pd.concat(returns_data, axis=1, join='inner')
            returns_df.columns = valid_tickers

            # Calculate expected returns and covariance matrix
            expected_returns = returns_df.mean() * 252  # Annualize
            cov_matrix = returns_df.cov() * 252  # Annualize

            # Check if covariance matrix is positive semi-definite
            if not self.is_positive_semi_definite(cov_matrix.values):
                st.error("Covariance matrix is not positive semi-definite. Adjusting...")
                cov_matrix_values = self.nearest_positive_semi_definite(cov_matrix.values)
            else:
                cov_matrix_values = cov_matrix.values

            n = len(valid_tickers)
            weights = cp.Variable(n)
            portfolio_return = expected_returns.values @ weights
            portfolio_risk = cp.quad_form(weights, cov_matrix_values)
            risk_aversion = 0.5  # Adjust based on preference

            # Define the optimization problem
            problem = cp.Problem(cp.Maximize(portfolio_return - risk_aversion * portfolio_risk),
                                 [cp.sum(weights) == 1, weights >= 0])
            problem.solve()

            if weights.value is not None:
                optimal_weights = {ticker: round(weight, 4) for ticker, weight in zip(valid_tickers, weights.value)}
                return {"Optimal Weights": optimal_weights}
            else:
                return {"Optimal Weights": "Optimization failed."}
        except Exception as e:
            st.error(f"Error in Portfolio Optimization Agent: {e}")
            return {"Optimal Weights": "Optimization failed."}

    def nearest_positive_semi_definite(self, A):
        """
        Finds the nearest positive semi-definite matrix to A.
        """
        try:
            B = (A + A.T) / 2
            _, s, V = np.linalg.svd(B)
            H = np.dot(V.T, np.dot(np.diag(s), V))
            A2 = (B + H) / 2
            A3 = (A2 + A2.T) / 2

            if self.is_positive_semi_definite(A3):
                return A3
            else:
                return np.eye(A.shape[0])  # Fallback to identity matrix
        except Exception as e:
            st.error(f"Error in making matrix positive semi-definite: {e}")
            return np.eye(A.shape[0])

    def is_positive_semi_definite(self, A):
        """
        Checks if matrix A is positive semi-definite.
        """
        try:
            return np.all(np.linalg.eigvals(A) >= -1e-10)
        except Exception as e:
            st.error(f"Error checking positive semi-definiteness: {e}")
            return False

class RealTimeAlertingAgent:
    """Agent responsible for real-time alerting based on stock data."""

    def generate_alerts(self, stock_df: pd.DataFrame, sentiment_scores: Dict[str, float],
                        anomaly_data: Dict[str, List[str]]) -> Dict[str, List[str]]:
        """
        Generates alerts based on specific conditions such as high volatility, negative sentiment, or detected anomalies.
        Returns a dictionary with tickers as keys and list of alerts as values.
        """
        alerts = {}
        try:
            for index, row in stock_df.iterrows():
                ticker = row['Ticker']
                alerts[ticker] = []
                # Example Alert 1: High PE Ratio
                if row['PE Ratio'] > 25:
                    alerts[ticker].append("High PE Ratio detected.")

                # Example Alert 2: Negative Sentiment
                sentiment = sentiment_scores.get(ticker, 0)
                if sentiment < -0.5:
                    alerts[ticker].append("Negative sentiment in recent news.")

                # Example Alert 3: Anomalies detected
                anomaly_dates = anomaly_data.get(ticker, [])
                if anomaly_dates and anomaly_dates != ["No anomalies detected."]:
                    alerts[ticker].append(f"Anomalies detected on dates: {', '.join(map(str, anomaly_dates))}")

                # Add more alert conditions as needed
            return alerts
        except Exception as e:
            st.error(f"Error in Real-Time Alerting Agent: {e}")
            return alerts

# Initialize agents
if 'real_agent' not in st.session_state:
    st.session_state.real_agent = RealAgent()

if 'data_analysis_agent' not in st.session_state:
    st.session_state.data_analysis_agent = DataAnalysisAgent()

if 'prediction_agent' not in st.session_state:
    st.session_state.prediction_agent = PredictionAgent()

if 'sentiment_analysis_agent' not in st.session_state:
    st.session_state.sentiment_analysis_agent = SentimentAnalysisAgent()

if 'anomaly_detection_agent' not in st.session_state:
    st.session_state.anomaly_detection_agent = AnomalyDetectionAgent()

if 'portfolio_optimization_agent' not in st.session_state:
    st.session_state.portfolio_optimization_agent = PortfolioOptimizationAgent()

if 'real_time_alerting_agent' not in st.session_state:
    st.session_state.real_time_alerting_agent = RealTimeAlertingAgent()

# ----------------------------
# AI Assistant Tools
# ----------------------------

def get_sentiment_from_openai(prompt: str) -> float:
    """
    Uses OpenAI's GPT-4 to analyze sentiment and return a numerical score.
    """
    try:
        response = openai.ChatCompletion.create(
            model="gpt-4",
            messages=[
                {"role": "system", "content": "You are an assistant that analyzes the sentiment of news articles. Rate the sentiment from -1 (very negative) to 1 (very positive)."},
                {"role": "user", "content": prompt}
            ],
            max_tokens=10,
            temperature=0.0,
        )
        sentiment_text = response.choices[0].message['content'].strip()
        # Extract numerical value from response
        try:
            sentiment_score = float(sentiment_text)
        except ValueError:
            sentiment_score = 0.0  # Default to neutral if parsing fails
        # Clamp the score between -1 and 1
        sentiment_score = max(min(sentiment_score, 1.0), -1.0)
        return sentiment_score
    except Exception as e:
        st.error(f"Error in Sentiment Analysis: {e}")
        return 0.0

# ----------------------------
# FSIRDM Framework Integration
# ----------------------------

def generate_fsirdm_analysis(stock_df: pd.DataFrame, historical_dfs: Dict[str, pd.DataFrame],
                             forecast_period: int) -> Dict[str, Any]:
    """
    Generates FSIRDM analysis based on stock data and historical data.
    Returns a dictionary summarizing the analysis.
    """
    try:
        # Example FSIRDM Analysis Components:
        # Fuzzy Set Membership for Risk
        risk_levels = {}
        for index, row in stock_df.iterrows():
            rsi = row['RSI']
            if rsi < 30:
                risk = "High Risk"
            elif 30 <= rsi <= 70:
                risk = "Moderate Risk"
            else:
                risk = "Low Risk"
            risk_levels[row['Ticker']] = risk

        # Quantitative Analysis
        quantitative_analysis = {}
        for ticker in stock_df['Ticker']:
            hist_df = historical_dfs.get(ticker, pd.DataFrame())
            if not hist_df.empty:
                latest_close = hist_df['Close'].iloc[-1]
                sma = hist_df['SMA_20'].iloc[-1]
                ema = hist_df['EMA_20'].iloc[-1]
                volatility = hist_df['Close'].std()
                quantitative_analysis[ticker] = {
                    "Latest Close": latest_close,
                    "SMA 20": sma,
                    "EMA 20": ema,
                    "Volatility": volatility
                }
            else:
                quantitative_analysis[ticker] = "No historical data available."

        # Risk Assessment
        risk_assessment = risk_levels

        # Portfolio Optimization Suggestions
        optimization_suggestions = {}
        for ticker in stock_df['Ticker']:
            # Simple heuristic: If RSI < 30 and high volatility, suggest to hold or buy
            # If RSI > 70 and low volatility, suggest to sell
            if risk_assessment[ticker] == "High Risk" and quantitative_analysis[ticker] != "No historical data available.":
                optimization_suggestions[ticker] = "Consider holding or buying."
            elif risk_assessment[ticker] == "Low Risk" and quantitative_analysis[ticker] != "No historical data available.":
                optimization_suggestions[ticker] = "Consider selling or reducing position."
            else:
                optimization_suggestions[ticker] = "Hold current position."

        # Combine all FSIRDM components
        fsirdm_summary = {
            "Risk Assessment": risk_assessment,
            "Quantitative Analysis": quantitative_analysis,
            "Optimization Suggestions": optimization_suggestions
        }

        return fsirdm_summary
    except Exception as e:
        st.error(f"Error generating FSIRDM analysis: {e}")
        return {}

# ----------------------------
# AI Orchestrator Definition
# ----------------------------

class AIOrchestrator:
    """Orchestrates the swarm of AI agents to provide comprehensive insights."""

    def __init__(self):
        self.data_analysis_agent = st.session_state.data_analysis_agent
        self.prediction_agent = st.session_state.prediction_agent
        self.sentiment_analysis_agent = st.session_state.sentiment_analysis_agent
        self.anomaly_detection_agent = st.session_state.anomaly_detection_agent
        self.portfolio_optimization_agent = st.session_state.portfolio_optimization_agent
        self.real_time_alerting_agent = st.session_state.real_time_alerting_agent

    def generate_insights(self, stock_df: pd.DataFrame, historical_dfs: Dict[str, pd.DataFrame],
                         news_articles: Dict[str, List[Dict[str, Any]]], forecast_period: int) -> Dict[str, Any]:
        """
        Runs all agents and aggregates their insights.
        """
        insights = {}
        # Data Analysis
        data_trends = self.data_analysis_agent.analyze_market_trends(stock_df)
        insights['Data Trends'] = data_trends

        # Predictions
        forecasts = self.prediction_agent.forecast_stock_trends(historical_dfs, forecast_period)
        insights['Forecasts'] = forecasts

        # Sentiment Analysis
        sentiments = self.sentiment_analysis_agent.analyze_sentiment(news_articles)
        insights['Sentiment Scores'] = sentiments

        # Anomaly Detection
        anomalies = self.anomaly_detection_agent.detect_anomalies(stock_df, historical_dfs)
        insights['Anomalies'] = anomalies

        # Portfolio Optimization
        portfolio_opt = self.portfolio_optimization_agent.optimize_portfolio(stock_df, historical_dfs)
        insights['Portfolio Optimization'] = portfolio_opt

        # Real-Time Alerting
        alerts = self.real_time_alerting_agent.generate_alerts(stock_df, sentiments, anomalies)
        insights['Alerts'] = alerts

        return insights

# ----------------------------
# Main Application
# ----------------------------

def run_and_visualize(ga):
    """
    Runs the GA and visualizes results.
    """
    with st.spinner("Running Genetic Algorithm..."):
        best_path, fitness_history, mutation_rate_history = ga.run()

    st.success("Genetic Algorithm completed!")

    # Best Path Visualization
    st.write("### Best Path Found:")
    st.write(best_path)

    # Fitness and Mutation Rate Trends
    st.write("### Fitness and Mutation Rate Trends")
    st.line_chart({
        "Best Fitness": fitness_history,
        "Mutation Rate": mutation_rate_history
    })

def main():
    # Initialize Database
    initialize_database()

    # Apply local CSS
    local_css()

    # Auto-refresh the app every 60 seconds for real-time data
    count = st_autorefresh(interval=60 * 1000, limit=100, key="autorefreshcounter")

    # Title and Description with enhanced styling
    st.markdown("<h1 class='title'>📈 Stock Dashboard</h1>", unsafe_allow_html=True)
    st.markdown("""
    <p class='description'>
    Explore and manage your favorite stocks. View comprehensive financial metrics, analyze historical performance with predictive insights, compare multiple stocks, stay updated with the latest news, and interact with our AI Assistant.
    </p>
    """, unsafe_allow_html=True)

    # Sidebar Configuration
    st.sidebar.header("🔧 Settings")

    # Dark/Light Mode Toggle
    theme = st.sidebar.radio("Theme", ("Light", "Dark"))
    if theme == "Dark":
        st.markdown(
            """
            <style>
            body {
                background-color: #1e1e1e;
                color: #ffffff;
            }
            .chat-container {
                background-color: #2e2e2e;
                color: #ffffff;
            }
            </style>
            """,
            unsafe_allow_html=True
        )
    else:
        st.markdown(
            """
            <style>
            body {
                background-color: #ffffff;
                color: #000000;
            }
            .chat-container {
                background-color: #f9f9f9;
                color: #000000;
            }
            </style>
            """,
            unsafe_allow_html=True
        )

    # Sidebar Options for Database Management
    st.sidebar.header("🗂️ Database Management")
    db_option = st.sidebar.selectbox(
        "Select an option",
        ("None", "View Interactions", "Clear Interactions", "Clear Cache", "Manage Portfolio")
    )

    if db_option == "View Interactions":
        st.sidebar.markdown("### Recent Interactions")
        interactions = fetch_interactions()
        if interactions:
            for interaction in interactions:
                timestamp, user_input, assistant_response = interaction
                st.sidebar.markdown(
                    f"- **{timestamp}**\n  **You:** {user_input}\n  **Assistant:** {assistant_response}\n")
        else:
            st.sidebar.info("No interactions to display.")
    elif db_option == "Clear Interactions":
        if st.sidebar.button("🗑️ Clear All Interactions"):
            clear_interactions()
            # Also reset AgentState's memory
            st.session_state.real_agent.agent_state.reset_memory()
    elif db_option == "Clear Cache":
        if st.sidebar.button("🗑️ Clear All Cached Stock Data"):
            clear_stock_cache()
    elif db_option == "Manage Portfolio":
        manage_portfolio()

    # Fetch the top 10 stock tickers
    with st.spinner("Fetching top 10 stocks by market capitalization..."):
        top_10_tickers = get_top_10_stocks()

    # User can add more stocks
    st.sidebar.header("📈 Add Stocks to Dashboard")
    user_tickers = st.sidebar.text_input("Enter stock tickers separated by commas (e.g., AAPL, MSFT, GOOGL):")
    add_button = st.sidebar.button("➕ Add Stocks", key="add_button")
    remove_button = st.sidebar.button("- Remove Stocks", key="remove_button")

    if add_button and user_tickers:
        user_tickers_list = [ticker.strip().upper() for ticker in user_tickers.split(",")]
        for ticker in user_tickers_list:
            if ticker and ticker not in top_10_tickers and ticker not in fetch_portfolio():
                add_to_portfolio(ticker)
        st.sidebar.success("Selected stocks have been added to your portfolio.")
    if remove_button and user_tickers:
        user_tickers_list = [ticker.strip().upper() for ticker in user_tickers.split(",")]
        for ticker in user_tickers_list:
            if ticker and ticker in top_10_tickers:
                remove_from_portfolio(ticker)
        st.sidebar.success("Selected stocks have been removed from your portfolio.")

    # Sidebar Configuration
    st.sidebar.title("🧬 Genetic Algorithm with Fuzzy Logic")
    st.sidebar.header("Configuration")

    # Parameters
    population_size = st.sidebar.slider("Population Size", 10, 200, 50)
    generations = st.sidebar.slider("Generations", 10, 500, 100)
    uncertainty_input = st.sidebar.slider("Uncertainty Level", 0.0, 1.0, 0.5)

    # Graph Preparation
    st.sidebar.header("Graph Configuration")
    num_nodes = st.sidebar.slider("Number of Nodes", 5, 20, 10)
    graph = prepare_graph(num_nodes)

    # Fuzzy Logic
    fuzzy_controller = FuzzyMutationController()
    mutation_rate = fuzzy_controller.compute_mutation_rate(uncertainty_input)
    st.sidebar.write(f"Computed Mutation Rate: {mutation_rate:.2f}")

    # Visualize Fuzzy Logic
    st.sidebar.header("Fuzzy Logic Visualization")
    visualize_fuzzy_logic(fuzzy_controller)

    # Run Genetic Algorithm
    st.header("🧬 Genetic Algorithm Results")
    ga = GeneticAlgorithm(graph, fuzzy_controller, population_size, generations)

    if st.button("Run Genetic Algorithm"):
        run_and_visualize(ga)
        visualize_graph(graph, best_path=ga.run()[0])

    # Fetch portfolio tickers
    portfolio_tickers = fetch_portfolio()

    # Combine top 10 and portfolio tickers, ensuring uniqueness
    all_tickers = list(set(top_10_tickers + portfolio_tickers))

    if not all_tickers:
        st.warning("No stocks to display. Please add stocks to your portfolio.")
        st.stop()

    # Sidebar Options
    refresh_data = st.sidebar.button("🔄 Refresh Data")
    download_format = st.sidebar.selectbox("Download Format", ("CSV", "JSON"))
    time_period = st.sidebar.selectbox(
        "Select Time Period for Historical Data",
        ("1mo", "3mo", "6mo", "1y", "2y", "5y", "10y", "ytd", "max")
    )

    # Refresh data if button is clicked
    if refresh_data:
        with st.spinner("Refreshing stock data..."):
            stock_df = get_stock_data(all_tickers)
        st.success("Data refreshed successfully!")
    else:
        with st.spinner("Fetching stock data..."):
            stock_df = get_stock_data(all_tickers)

    # Ensure relevant columns are numeric
    numeric_cols = ["Market Cap", "Current Price (USD)", "52 Week High", "52 Week Low",
                    "PE Ratio", "Dividend Yield", "EPS", "Beta", "Revenue", "Net Income", "RSI"]

    for col in numeric_cols:
        stock_df[col] = pd.to_numeric(stock_df[col], errors='coerce').fillna(0)

    # Verify data types
    st.write("**Data Types After Conversion:**")
    st.write(stock_df.dtypes)

    # Now, compute sector_performance from numeric 'Market Cap'
    sector_performance = stock_df.groupby('Sector')['Market Cap'].sum().reset_index()

    # Filter out sectors with zero or negative Market Cap
    sector_performance = sector_performance[sector_performance['Market Cap'] > 0]

    if sector_performance.empty:
        st.warning("No valid sector performance data available for visualization.")
    else:
        # Continue with formatting and plotting
        stock_df_formatted = format_data(stock_df)

        # Display the DataFrame with enhanced styling
        st.subheader("📊 Stocks by Market Capitalization")
        st.data_editor(
            stock_df_formatted.style.set_table_styles([
                {'selector': 'th',
                 'props': [('background-color', '#2e86de'), ('color', 'white'), ('font-size', '14px'),
                           ('text-align', 'center')]},
                {'selector': 'td', 'props': [('text-align', 'center'), ('font-size', '13px')]},
                {'selector': 'tr:nth-child(even)', 'props': [('background-color', '#f2f2f2')]},
            ]).hide(axis='index'),
            height=600,
            use_container_width=True
        )

        # Sector Performance Treemap
        st.markdown("### 🔥 Sector Performance Treemap")
        try:
            fig_heatmap = px.treemap(
                sector_performance,
                path=['Sector'],
                values='Market Cap',
                title='Market Capitalization by Sector',
                color='Market Cap',
                color_continuous_scale='RdBu',
                hover_data={'Market Cap': ':.2f'}
            )
            st.plotly_chart(fig_heatmap, use_container_width=True)
        except ZeroDivisionError as zde:
            st.error(f"Error creating treemap: {zde}")
        except Exception as e:
            st.error(f"An unexpected error occurred while creating treemap: {e}")

        # Comparative Metrics Visualization
        st.markdown("### 📊 Comparative Metrics")

        comparative_metrics = ["Market Cap", "Current Price (USD)", "52 Week High", "52 Week Low",
                               "PE Ratio", "Dividend Yield", "EPS", "Beta", "Revenue", "Net Income", "RSI"]

        for metric in comparative_metrics:
            fig = px.bar(
                stock_df,
                x='Ticker',
                y=metric,
                color='Sector',
                title=f'{metric} Comparison',
                labels={metric: metric, 'Ticker': 'Stock Ticker'},
                text_auto='.2s' if 'Cap' in metric or 'Revenue' in metric or 'Income' in metric else '.2f'
            )
            st.plotly_chart(fig, use_container_width=True)

        # Download Buttons
        st.markdown("### 📥 Download Data")
        col1, col2 = st.columns(2)

        with col1:
            if download_format == "CSV":
                csv_data = convert_df_to_csv(stock_df)
                st.download_button(
                    label="📥 Download CSV",
                    data=csv_data,
                    file_name='stocks_data.csv',
                    mime='text/csv',
                )
        with col2:
            if download_format == "JSON":
                json_data = convert_df_to_json(stock_df)
                st.download_button(
                    label="📥 Download JSON",
                    data=json_data,
                    file_name='stocks_data.json',
                    mime='application/json',
                )

        # Additional Features: Historical Performance
        st.markdown("---")
        st.markdown("### 📈 Stock Performance Over Time")

        # Let user select multiple tickers to view historical data
        selected_tickers = st.multiselect("🔍 Select Stocks to Compare Historical Performance", all_tickers, default=[all_tickers[0]])

        if not selected_tickers:
            st.warning("Please select at least one stock to view historical performance.")
        else:
            historical_dfs = {}
            for ticker in selected_tickers:
                historical_df = get_historical_data(ticker, period=time_period)
                if not historical_df.empty:
                    # Calculate additional metrics for better insights
                    historical_df['SMA_20'] = historical_df['Close'].rolling(window=20).mean()
                    historical_df['EMA_20'] = historical_df['Close'].ewm(span=20, adjust=False).mean()
                    historical_df['BB_upper'] = historical_df['SMA_20'] + 2 * historical_df['Close'].rolling(window=20).std()
                    historical_df['BB_lower'] = historical_df['SMA_20'] - 2 * historical_df['Close'].rolling(window=20).std()
                    delta = historical_df['Close'].diff()
                    up = delta.clip(lower=0)
                    down = -1 * delta.clip(upper=0)
                    roll_up = up.rolling(window=14).mean()
                    roll_down = down.rolling(window=14).mean()
                    RS = roll_up / roll_down
                    historical_df['RSI'] = 100.0 - (100.0 / (1.0 + RS))
                    historical_dfs[ticker] = historical_df
                else:
                    historical_dfs[ticker] = pd.DataFrame()

            # Interactive Comparative Candlestick Chart using Plotly
            fig_candlestick = go.Figure()

            for ticker in selected_tickers:
                hist_df = historical_dfs.get(ticker, pd.DataFrame())
                if hist_df.empty:
                    continue
                fig_candlestick.add_trace(go.Candlestick(
                    x=hist_df['Date'],
                    open=hist_df['Open'],
                    high=hist_df['High'],
                    low=hist_df['Low'],
                    close=hist_df['Close'],
                    name=f'{ticker} Candlestick'
                ))
                fig_candlestick.add_trace(go.Scatter(
                    x=hist_df['Date'],
                    y=hist_df['SMA_20'],
                    mode='lines',
                    line=dict(width=1),
                    name=f'{ticker} SMA 20'
                ))
                fig_candlestick.add_trace(go.Scatter(
                    x=hist_df['Date'],
                    y=hist_df['EMA_20'],
                    mode='lines',
                    line=dict(width=1),
                    name=f'{ticker} EMA 20'
                ))

            fig_candlestick.update_layout(
                title="Comparative Candlestick Chart with Moving Averages",
                xaxis_title="Date",
                yaxis_title="Price (USD)",
                xaxis_rangeslider_visible=False
            )

            st.plotly_chart(fig_candlestick, use_container_width=True)

            # Comparative RSI
            st.markdown("#### Comparative Relative Strength Index (RSI)")
            fig_rsi = go.Figure()

            for ticker in selected_tickers:
                hist_df = historical_dfs.get(ticker, pd.DataFrame())
                if hist_df.empty:
                    continue
                fig_rsi.add_trace(go.Scatter(
                    x=hist_df['Date'],
                    y=hist_df['RSI'],
                    mode='lines',
                    name=f'{ticker} RSI'
                ))

            fig_rsi.add_hline(y=70, line_dash="dash", line_color="red")
            fig_rsi.add_hline(y=30, line_dash="dash", line_color="green")
            fig_rsi.update_layout(
                title="Relative Strength Index (RSI) Comparison",
                xaxis_title="Date",
                yaxis_title="RSI",
                yaxis=dict(range=[0, 100])
            )
            st.plotly_chart(fig_rsi, use_container_width=True)

            # Forecasting with Prophet for each selected ticker
            st.markdown("#### 🔮 Future Price Prediction")
            forecast_period = st.slider("Select Forecast Period (Days):", min_value=30, max_value=365, value=90,
                                        step=30)

            forecast_figs = []
            # Initialize AI Orchestrator
            ai_orchestrator = AIOrchestrator()
            # Generate insights from agents
            fsirdm_data = ai_orchestrator.generate_insights(stock_df, historical_dfs, {}, forecast_period)

            for ticker in selected_tickers:
                model, forecast_df = forecast_stock_price(historical_dfs.get(ticker, pd.DataFrame()), forecast_period)
                if model is not None and not forecast_df.empty:
                    fig_forecast = plot_plotly(model, forecast_df)
                    fig_forecast.update_layout(
                        title=f"{ticker} Price Forecast for Next {forecast_period} Days",
                        xaxis_title="Date",
                        yaxis_title="Price (USD)"
                    )
                    forecast_figs.append(fig_forecast)
                    st.plotly_chart(fig_forecast, use_container_width=True)
                else:
                    st.warning(f"Forecasting model could not generate predictions for {ticker}.")

            # Real-Time Notifications
            # Implement notifications for significant changes here if needed

        # News and Analysis
            st.markdown("---")
            st.markdown("### 📰 Latest News")

            news_articles = get_stock_news(selected_tickers)

            if news_articles:
                for ticker, articles in news_articles.items():
                    st.markdown(f"#### {ticker} News")
                    if articles and isinstance(articles, list):
                        for article in articles:
                            st.markdown(f"**[{article['Title']}]({article['URL']})**")
                            try:
                                published_date = datetime.strptime(article['Published At'], "%Y-%m-%dT%H:%M:%SZ")
                                formatted_date = published_date.strftime("%B %d, %Y %H:%M")
                            except ValueError:
                                formatted_date = article['Published At']
                            st.markdown(f"*{formatted_date}*")
                            st.markdown(f"{article['Description']}")
                            st.markdown("---")
                    else:
                        st.info("No recent news articles found.")
            else:
                st.info("No recent news articles found.")

        # AI Assistant Interaction
            st.markdown("---")
            st.markdown("### 🤖 Ask the Generis AI")
            st.empty()
            user_input = st.text_input("Ask a question about stocks or market trends:", type="default")

            # ----------------------------
            # Added "Clear Chat" Button
            # ----------------------------
            if st.button("🗑️ Clear Chat"):
                clear_chat()
                st.success("Chat history has been cleared.")

            if user_input:
                with st.spinner("Processing your request..."):
                    response = st.session_state.real_agent.process(
                        user_input, selected_tickers, stock_df, historical_dfs, news_articles, fsirdm_data
                    )
                # Display the conversation in a chat-like format
                st.markdown(f"<div class='user-message'><strong>You:</strong> {user_input}</div>",
                            unsafe_allow_html=True)
                st.markdown(f"<div class='assistant-message'><strong>Assistant:</strong> {response}</div>",
                            unsafe_allow_html=True)
                st.markdown("</div>", unsafe_allow_html=True)

        # Display AI Generated Insights
            st.markdown("---")
            st.markdown("### 🤖 AI Generated Insights")
            ai_orchestrator = AIOrchestrator()
            insights = ai_orchestrator.generate_insights(stock_df, historical_dfs, news_articles, forecast_period)

            # Display Data Trends
            st.subheader("📊 Data Trends")
            data_trends = insights.get('Data Trends', {})
            if data_trends:
                st.json(data_trends)
            else:
                st.info("No data trends available.")

            # Display Forecasts
            st.subheader("🔮 Forecasts")
            forecasts = insights.get('Forecasts', {})
            if forecasts:
                for ticker, forecast in forecasts.items():
                    st.markdown(f"**{ticker} Forecast:**")
                    if isinstance(forecast, list):
                        forecast_df = pd.DataFrame(forecast)
                        st.dataframe(forecast_df)
                    else:
                        st.write(forecast)
            else:
                st.info("No forecasts available.")

            # Display Sentiment Scores
            st.subheader("😊 Sentiment Scores")
            sentiments = insights.get('Sentiment Scores', {})
            if sentiments:
                st.json(sentiments)
            else:
                st.info("No sentiment scores available.")

            # Display Anomalies
            st.subheader("⚠️ Anomalies Detected")
            anomalies = insights.get('Anomalies', {})
            if anomalies:
                for ticker, anomaly_dates in anomalies.items():
                    st.markdown(f"**{ticker}:** {', '.join(map(str, anomaly_dates))}")
            else:
                st.info("No anomalies detected.")

            # Display Portfolio Optimization
            st.subheader("💼 Portfolio Optimization")
            portfolio_opt = insights.get('Portfolio Optimization', {})
            if portfolio_opt:
                st.json(portfolio_opt)
            else:
                st.info("No portfolio optimization available.")

            # Display Alerts
            st.subheader("🚨 Alerts")
            alerts = insights.get('Alerts', {})
            if alerts:
                for ticker, alert_list in alerts.items():
                    if alert_list:
                        for alert in alert_list:
                            st.warning(f"{ticker}: {alert}")
            else:
                st.info("No alerts generated.")

        # Footer
            st.markdown("---")
            st.markdown("<p class='footer'>© 2024 Your Company Name. All rights reserved.</p>", unsafe_allow_html=True)

def forecast_stock_price(historical_df: pd.DataFrame, periods: int = 90) -> Tuple[Any, pd.DataFrame]:
    """
    Uses Facebook's Prophet to forecast future stock prices.
    Returns both the fitted model and the forecast DataFrame.
    """
    try:
        if historical_df.empty:
            return None, pd.DataFrame()
        df_prophet = historical_df[['Date', 'Close']].rename(columns={'Date': 'ds', 'Close': 'y'})
        model = Prophet(daily_seasonality=False, yearly_seasonality=True, weekly_seasonality=True)
        model.fit(df_prophet)
        future = model.make_future_dataframe(periods=periods)
        forecast = model.predict(future)
        return model, forecast  # Return both model and forecast
    except Exception as e:
        st.error(f"Error in forecasting: {e}")
        return None, pd.DataFrame()

def clear_chat():
    """
    Clears the conversation history by deleting interactions from the database
    and resetting the AgentState's memory.
    """
    clear_interactions()
    st.session_state.real_agent.agent_state.reset_memory()

def manage_portfolio():
    """
    Manages the user's portfolio by allowing addition and removal of stocks.
    """
    st.header("📁 Manage Your Portfolio")
    portfolio = fetch_portfolio()
    if portfolio:
        st.subheader("Your Portfolio:")
        portfolio_df = pd.DataFrame(portfolio, columns=["Ticker"])
        st.data_editor(portfolio_df, height=500, use_container_width=True, key="portfolio_editor")

        remove_tickers = st.multiselect(label="Select stocks to remove from your portfolio:", options=portfolio)
        if st.button("🗑️ Remove Selected Stocks") and remove_tickers:
            for ticker in remove_tickers:
                remove_from_portfolio(ticker)
            st.success("Selected stocks have been removed from your portfolio.")
    else:
        st.info("Your portfolio is empty. Add stocks to get started.")

    st.markdown("---")
    st.subheader("Add New Stocks to Portfolio")
    new_ticker = st.text_input("Enter a stock ticker to add:")
    if st.button("➕ Add to Portfolio"):
        if new_ticker:
            add_to_portfolio(new_ticker)
        else:
            st.warning("Please enter a valid stock ticker.")

def get_top_10_stocks() -> List[str]:
    """
    Fetches the top 10 stocks by market capitalization using Financial Modeling Prep API.
    Utilizes the database cache if available and not expired.
    """
    try:
        # Check if cache exists for 'TOP10'
        conn = sqlite3.connect(DATABASE)
        cursor = conn.cursor()
        cursor.execute("""
            SELECT fetched_at, data
            FROM stock_cache
            WHERE ticker = 'TOP10'
        """)
        row = cursor.fetchone()
        conn.close()
        use_cache = False
        if row:
            fetched_at, data = row
            fetched_time = datetime.fromisoformat(fetched_at).date()
            today = datetime.utcnow().date()
            if (today - fetched_time).days < 1:  # 1 day TTL
                use_cache = True
                tickers = json.loads(data)
        if not use_cache:
            if not update_api_usage("FMP"):
                st.warning("Financial Modeling Prep API rate limit exceeded. Cannot fetch top 10 stocks.")
                return []
            url = f"https://financialmodelingprep.com/api/v3/stock-screener?marketCapMoreThan=1000000000&limit=100&apikey={FMP_API_KEY}"
            response = requests.get(url)
            response.raise_for_status()
            data = response.json()
            if not isinstance(data, list):
                st.error("Unexpected response format from Financial Modeling Prep API.")
                return []
            sorted_data = sorted(data, key=lambda x: x.get('marketCap', 0), reverse=True)
            top_10 = sorted_data[:10]
            tickers = [stock['symbol'] for stock in top_10]
            # Cache the top 10 tickers
            insert_stock_cache('TOP10', tickers)
        return tickers
    except Exception as e:
        st.error(f"Error fetching top 10 stocks: {e}")
        return []

@st.cache_data(ttl=600)
def get_stock_data(tickers: List[str]) -> pd.DataFrame:
    """
    Fetches detailed stock data for the given list of tickers using yfinance.
    Utilizes the database cache if available.
    """
    # Use asyncio to fetch data concurrently
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    responses = loop.run_until_complete(fetch_all_stock_data(tickers))
    loop.close()

    stock_info = []
    for ticker, info in responses:
        stock_info.append({
            "Ticker": ticker,
            "Name": info.get("shortName", "N/A"),
            "Sector": info.get("sector", "N/A"),
            "Industry": info.get("industry", "N/A"),
            "Market Cap": info.get("marketCap", 0),  # Default to 0 if missing
            "Current Price (USD)": info.get("currentPrice", 0),
            "52 Week High": info.get("fiftyTwoWeekHigh", 0),
            "52 Week Low": info.get("fiftyTwoWeekLow", 0),
            "PE Ratio": info.get("trailingPE", 0),
            "Dividend Yield": info.get("dividendYield", 0),
            "EPS": info.get("trailingEps", 0),
            "Beta": info.get("beta", 0),
            "Revenue": info.get("totalRevenue", 0),
            "Net Income": info.get("netIncomeToCommon", 0),
            "RSI": calculate_rsi(info.get("symbol", "N/A"), period=14)  # Calculate RSI
        })
    df = pd.DataFrame(stock_info)

    # Data Cleaning: Ensure all numeric columns are indeed numeric
    numeric_cols = ["Market Cap", "Current Price (USD)", "52 Week High", "52 Week Low",
                    "PE Ratio", "Dividend Yield", "EPS", "Beta", "Revenue", "Net Income", "RSI"]

    for col in numeric_cols:
        df[col] = pd.to_numeric(df[col], errors='coerce').fillna(0)

    return df

def calculate_rsi(ticker: str, period: int = 14) -> float:
    """
    Calculates the Relative Strength Index (RSI) for a given ticker.
    """
    try:
        stock = yf.Ticker(ticker)
        hist = stock.history(period="1y")
        if hist.empty:
            return 0.0
        delta = hist['Close'].diff()
        up = delta.clip(lower=0)
        down = -1 * delta.clip(upper=0)
        roll_up = up.rolling(window=period).mean()
        roll_down = down.rolling(window=period).mean()
        RS = roll_up / roll_down
        RSI = 100.0 - (100.0 / (1.0 + RS))
        return RSI.iloc[-1] if not RSI.empty else 0.0
    except Exception as e:
        st.error(f"Error calculating RSI for {ticker}: {e}")
        return 0.0

@st.cache_data(ttl=600)
def get_historical_data(ticker: str, period: str = "1mo") -> pd.DataFrame:
    """
    Fetches historical stock data.
    """
    try:
        stock = yf.Ticker(ticker)
        hist = stock.history(period=period)
        if hist.empty:
            st.warning(f"No historical data available for {ticker} in the selected period.")
            return pd.DataFrame()
        hist.reset_index(inplace=True)
        hist['Date'] = hist['Date'].dt.date
        return hist
    except Exception as e:
        st.error(f"Error fetching historical data for {ticker}: {e}")
        return pd.DataFrame()

@st.cache_data(ttl=600)
def get_stock_news(tickers: List[str]) -> Dict[str, List[Dict[str, Any]]]:
    """
    Fetches latest news for the given tickers using NewsAPI.
    Returns a dictionary with tickers as keys and list of articles as values.
    """
    news_dict = {}
    for ticker in tickers:
        if not update_api_usage("NewsAPI"):
            st.warning("NewsAPI rate limit exceeded. Cannot fetch latest news.")
            news_dict[ticker] = []
            continue
        try:
            stock = yf.Ticker(ticker)
            company_name = stock.info.get('shortName', ticker)
            query = f"{ticker} OR \"{company_name}\""

            articles = newsapi.get_everything(
                q=query,
                language='en',
                sort_by='publishedAt',
                page_size=5
            )

            news = []
            for article in articles.get('articles', []):
                news.append({
                    "Title": article['title'],
                    "Description": article['description'],
                    "URL": article['url'],
                    "Published At": article['publishedAt']
                })
            news_dict[ticker] = news
        except Exception as e:
            st.error(f"Error fetching news for {ticker}: {e}")
            news_dict[ticker] = []
    return news_dict

def convert_df_to_csv(df: pd.DataFrame) -> bytes:
    """Converts DataFrame to CSV."""
    return df.to_csv(index=False).encode('utf-8')

def convert_df_to_json(df: pd.DataFrame) -> bytes:
    """Converts DataFrame to JSON."""
    return df.to_json(orient='records', indent=4).encode('utf-8')

def format_data(df: pd.DataFrame) -> pd.DataFrame:
    """
    Formats numerical columns for better readability.
    """
    df_formatted = df.copy()
    df_formatted["Market Cap"] = df_formatted["Market Cap"].apply(
        lambda x: f"${x:,.2f}" if isinstance(x, (int, float)) else "N/A"
    )
    df_formatted["Current Price (USD)"] = df_formatted["Current Price (USD)"].apply(
        lambda x: f"${x:,.2f}" if isinstance(x, (int, float)) else "N/A"
    )
    df_formatted["52 Week High"] = df_formatted["52 Week High"].apply(
        lambda x: f"${x:,.2f}" if isinstance(x, (int, float)) else "N/A"
    )
    df_formatted["52 Week Low"] = df_formatted["52 Week Low"].apply(
        lambda x: f"${x:,.2f}" if isinstance(x, (int, float)) else "N/A"
    )
    df_formatted["PE Ratio"] = df_formatted["PE Ratio"].apply(
        lambda x: f"{x:.2f}" if isinstance(x, (int, float)) else "N/A"
    )
    df_formatted["Dividend Yield"] = df_formatted["Dividend Yield"].apply(
        lambda x: f"{x:.2%}" if isinstance(x, (int, float)) else "N/A"
    )
    df_formatted["EPS"] = df_formatted["EPS"].apply(
        lambda x: f"{x:.2f}" if isinstance(x, (int, float)) else "N/A"
    )
    df_formatted["Beta"] = df_formatted["Beta"].apply(
        lambda x: f"{x:.2f}" if isinstance(x, (int, float)) else "N/A"
    )
    df_formatted["Revenue"] = df_formatted["Revenue"].apply(
        lambda x: f"${x:,.2f}" if isinstance(x, (int, float)) else "N/A"
    )
    df_formatted["Net Income"] = df_formatted["Net Income"].apply(
        lambda x: f"${x:,.2f}" if isinstance(x, (int, float)) else "N/A"
    )
    df_formatted["RSI"] = df_formatted["RSI"].apply(
        lambda x: f"{x:.2f}" if isinstance(x, (int, float)) else "N/A"
    )
    return df_formatted

def get_stock_price(ticker: str) -> str:
    """
    Fetches the current price of the given ticker.
    """
    try:
        stock = yf.Ticker(ticker)
        price = stock.info.get('currentPrice', None)
        if price is None:
            return f"Could not fetch the current price for {ticker.upper()}."
        return f"The current price of {ticker.upper()} is ${price:.2f}"
    except Exception as e:
        return f"Error fetching stock price for {ticker.upper()}: {e}"

def get_stock_summary(ticker: str) -> str:
    """
    Fetches a summary of the given ticker.
    """
    try:
        stock = yf.Ticker(ticker)
        info = stock.info
        summary = {
            "Name": info.get("shortName", "N/A"),
            "Sector": info.get("sector", "N/A"),
            "Industry": info.get("industry", "N/A"),
            "Current Price (USD)": info.get("currentPrice", "N/A"),
            "52 Week High": info.get("fiftyTwoWeekHigh", "N/A"),
            "52 Week Low": info.get("fiftyTwoWeekLow", "N/A"),
            "Market Cap": info.get("marketCap", "N/A"),
        }
        response = "\n".join([f"{key}: {value}" for key, value in summary.items()])
        return response
    except Exception as e:
        return f"Error fetching summary for {ticker.upper()}: {e}"

def get_latest_news(query: str) -> List[Dict[str, Any]]:
    """
    Fetches the latest news articles based on the query.
    """
    if not update_api_usage("NewsAPI"):
        st.warning("NewsAPI rate limit exceeded. Cannot fetch latest news.")
        return []
    try:
        articles = newsapi.get_everything(q=query, language='en', sort_by='publishedAt', page_size=3)
        if not articles['articles']:
            return []
        news_list: List[Dict[str, Any]] = []
        for article in articles['articles']:
            news_list.append({
                "Title": article['title'],
                "Description": article['description'],
                "URL": article['url'],
                "Published At": article['publishedAt']
            })
        return news_list
    except Exception as e:
        st.error(f"Error fetching news for {query}: {e}")
        return []

# ----------------------------
# Run the Application
# ----------------------------

if __name__ == "__main__":
    main()