app.py

import streamlit as st
import pandas as pd
import os
import json
import yfinance as yf
from langchain_core.output_parsers import JsonOutputParser
from pydantic import BaseModel, Field, ValidationError
from typing import List, Optional, Dict
from langchain_groq import ChatGroq
from dataclasses import dataclass, field
from dotenv import load_dotenv
import pickle

import requests
from bs4 import BeautifulSoup
import re
import google.generativeai as genai
import numpy as np
import logging  # Added logging
import time # added time for loading animation


# Load environment variables
load_dotenv()

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# Configuration
class Config:
    ALPHA_VANTAGE_API_KEY = os.getenv("ALPHA_VANTAGE_API_KEY")
    GROQ_API_KEY = os.getenv("GROQ_API_KEY")
    STOCK_DATA_DIR = "stock_data_NSE"
    OUTPUT_FILE = "output_files/portfolio.json"
    SECTORS = [
        "Communication Services",
        "Consumer Discretionary",
        "Consumer Staples",
        "Energy",
        "Financials",
        "Health Care",
        "Industrials",
        "Information Technology",
        "Materials",
        "Real Estate",
        "Utilities"
    ]
    GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")


# Create output directories if they don't exist
if not os.path.exists(Config.STOCK_DATA_DIR):
    os.makedirs(Config.STOCK_DATA_DIR)
if not os.path.exists("output_files"):
    os.makedirs("output_files")


# --------------------- Function from portfolio.py ---------------------

def fetch_stock_data(symbols: List[str]) -> Dict[str, pd.DataFrame | None]:
    """Fetches stock data for multiple symbols from Yahoo Finance."""
    stock_dataframes = {}
    for symbol in symbols:
        try:
            logging.info(f"Fetching stock data for {symbol}")
            ticker = yf.Ticker(symbol)
            data = ticker.history(period="max")
            if data.empty:
                logging.warning(f"No data found for symbol '{symbol}'.")
                stock_dataframes[symbol] = None
                continue
            stock_dataframes[symbol] = data
            logging.info(f"Successfully fetched stock data for {symbol}")
        except Exception as e:
            logging.error(f"Error fetching data for symbol '{symbol}': {e}")
            stock_dataframes[symbol] = None
    return stock_dataframes


def store_stock_data(stock_dataframes: Dict[str, pd.DataFrame | None],
                     output_path: str = Config.STOCK_DATA_DIR) -> None:
    """Stores stock data to local CSV files."""
    for symbol, data in stock_dataframes.items():
        if data is not None:
            file_name = f"{symbol}_daily_data.csv"
            file_path = os.path.join(output_path, file_name)
            try:
                logging.info(f"Saving data for '{symbol}' to {file_path}")
                data.to_csv(file_path)
                logging.info(f"Data for '{symbol}' saved to {file_path}")
            except Exception as e:
                logging.error(f"Error saving data for '{symbol}' to {file_path}: {e}")
        else:
            logging.warning(f"No data available for '{symbol}', skipping storage.")


def load_stock_data_and_extract_price(output_path_dir: str) -> Dict[str, Dict[str, float]]:
    """Loads stock data from CSV files and extracts the most recent (last) day's closing price."""
    all_stock_data = {}
    for filename in os.listdir(output_path_dir):
        if filename.endswith("_daily_data.csv"):
            symbol = filename.replace("_daily_data.csv", "")
            file_path = os.path.join(output_path_dir, filename)
            try:
                logging.info(f"Loading data from {file_path} for symbol {symbol}")
                df = pd.read_csv(file_path, index_col=0)
                if not df.empty:
                    initial_price = df.iloc[-1]['Close']
                    all_stock_data[symbol] = {"initial_price": initial_price}
                    logging.info(f"Initial price extracted for {symbol}: {initial_price}")
                else:
                    logging.warning(f"Empty dataframe for symbol '{symbol}'. Setting initial price to 0")
                    all_stock_data[symbol] = {"initial_price": 0.0}
            except (IndexError, KeyError, FileNotFoundError) as e:
                logging.error(f"Error occurred for reading {symbol}, due to: {e}")
                all_stock_data[symbol] = {"initial_price": 0.0}  # default initial price is 0.0

    return all_stock_data


def merge_stock_data_with_price(stock_data: Dict, extracted_data: Dict) -> Dict:
    """Merges the extracted price data with the main stock data."""
    merged_stock_data = stock_data.copy()
    for key, value in stock_data.items():
        symbol = value["symbol"]
        if symbol in extracted_data:
            merged_stock_data[key]["initial_price"] = extracted_data[symbol]["initial_price"]
            logging.info(f"Merged initial price for {symbol} in main stock data")
        else:
            merged_stock_data[key]["initial_price"] = 0.0  # default value if it cannot be extracted
            logging.warning(f"Could not extract price for {symbol}. Setting default value to 0")
    return merged_stock_data


def generate_prompt(stock_data: Dict) -> str:
    """Generates a prompt for the language model with all the stock data"""
    prompt_template_with_price = """
    You are a financial analysis expert.
    Please provide a summary of the following stock data, including the company name, stock symbol, and initial purchase price.

    Stock Data:
    {stock_data}

    Summary:
    """
    stock_json_str = json.dumps(stock_data)
    formatted_prompt_with_price = prompt_template_with_price.format(stock_data=stock_json_str)
    logging.info(f"Generated LLM prompt: {formatted_prompt_with_price}")
    return formatted_prompt_with_price


class Asset(BaseModel):
    """Represents an asset within a portfolio."""
    quantity: int = Field(..., description="The number of shares or units held for this specific asset.")
    initial_price: float = Field(..., description="The initial purchase price per share or unit of this asset.")
    sector: str = Field(..., description=f"""The economic sector of the asset, based on the stock symbol or company name.
                                        For example, use this {Config.SECTORS}'Financials' for HDFC or JPM, 'consumer' for PG, 'Information Technology' for GOOG. This categorization
                                        should be done based on the business nature of the company whose stock is traded. For instance,
                                         if the stock symbol is 'HDFCBANK', the sector is expected to be 'Financials'.""")


class Portfolio(BaseModel):
    """Represents an individual portfolio."""
    name: str = Field(...,
                      description="The name given to this portfolio, for example 'Diversified Portfolio'. 'Aggressive Tech Portfolio' ")
    assets: Dict[str, Asset] = Field(..., description="""A dictionary containing the assets within this portfolio. The keys of the dictionary
                                   are the ticker symbols of the stocks (e.g., 'JPM', 'PG'), and the values are the corresponding
                                   'Asset' objects, which define the quantity, initial price, and sector for each asset.
                                   Example: {'JPM': {'quantity': 150, 'initial_price': 140, 'sector': 'finance'},
                                            'PG': {'quantity': 200, 'initial_price': 160, 'sector': 'consumer'}}"""
                                     )


def invoke_llm_for_portfolio(formatted_prompt: str) -> Portfolio:
    """Invokes the LLM for structured output of the portfolio"""
    llm = ChatGroq(groq_api_key=Config.GROQ_API_KEY, model_name="llama-3.1-8b-instant")
    structured_llm = llm.with_structured_output(Portfolio)
    try:
        logging.info(f"Invoking LLM for portfolio generation")
        output = structured_llm.invoke(formatted_prompt)
        logging.info(f"LLM returned Portfolio data {output}")
        return output
    except ValidationError as e:
        logging.error(f"Error during LLM invocation: {e}")
        raise
    except Exception as e:
        logging.error(f"Unexpected error during LLM invocation {e}")
        raise


def portfolio_to_json(portfolio: Portfolio, output_file: str = Config.OUTPUT_FILE) -> None:
    """Converts a Portfolio object to a JSON string and saves it to a file."""
    try:
        logging.info(f"Saving portfolio to JSON file: {output_file}")
        json_str = portfolio.model_dump_json(indent=4)
        with open(output_file, "w") as f:
            f.write(json_str)
        logging.info(f"Portfolio saved to '{output_file}'")
    except Exception as e:
        logging.error(f"Error saving JSON file {e}")


# --------------------- Function from scenario.py ---------------------
def scrape_website(url):
    headers = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
    }
    logging.info(f"Scraping website: {url}")
    try:
        response = requests.get(url, headers=headers)
        response.raise_for_status()  # Raise an exception for bad status codes
        soup = BeautifulSoup(response.text, "html.parser")
        logging.info(f"Successfully scraped website: {url}")
        return soup.prettify()
    except requests.exceptions.RequestException as e:
        logging.error(f"Failed to retrieve page. Status code: {e}")
        return f"Failed to retrieve page. Status code: {e}"


genai.configure(api_key=Config.GOOGLE_API_KEY)  # Replace with your API key
generation_config = {
    "temperature": 1,
    "top_p": 0.95,
    "top_k": 40,
    "max_output_tokens": 8192,
    "response_mime_type": "text/plain",
}

model = genai.GenerativeModel(
    model_name="gemini-2.0-flash-exp",
    generation_config=generation_config,
)

chat_session = model.start_chat()


def get_response(llm, prompt):
    logging.info(f"Sending prompt to LLM for scenario: {prompt}")
    response = llm.send_message(prompt)
    logging.info(f"LLM returned a response for scenario")
    return response


def extract_json_content(text):
    match = re.search(r"```json\n(.*?)```", text, re.DOTALL)
    if match:
        logging.info("Extracted JSON content from LLM response")
        return match.group(1).strip()
    else:
        logging.warning("Could not extract JSON content from LLM response")
        return None


def invoke_llm_for_scenario(context_data):
    sectors = Config.SECTORS

    prompt = f"""
        # TASK: Analyze market context and identify potential market scenarios.

        # CONTEXT:
        {context_data}
        # END CONTEXT

        # INSTRUCTION: Based on the provided market context, analyze and identify up to three plausible market scenarios.
        # For each scenario, determine its name (e.g., "Moderate Downturn"), the general market direction ("up" or "down"), a major trigger point that could cause the scenario to unfold, and a list of sectors that would be significantly impacted. Each 'sector_impact' list should have less than or equal to 4 sectors.

        # OUTPUT FORMAT: Provide the analysis in JSON format with the following structure.
        # Use the sector names provided:
        {sectors}

        # EXAMPLE:
        ```json
        {{
            "market_scenarios": {{
                "scenario1": {{
                    "name": "Moderate Downturn",
                    "direction": "down",
                    "trigger": "Interest rate hike",
                    "sector_impact": [
                        "Financials",
                        "Energy"
                    ]
                }},
                "scenario2": {{
                    "name": "Bullish Growth",
                    "direction": "up",
                    "trigger": "Successful vaccine rollout",
                    "sector_impact": [
                        "Health Care",
                        "Information Technology"
                    ]
                }}
            }}
        }}
        """

    answer = get_response(chat_session, prompt)
    json_output = extract_json_content(answer.text)
    output_file = "output_files/scenario.json"
    try:
        analysis_json = json.loads(json_output)
        os.makedirs(os.path.dirname(output_file), exist_ok=True)
        with open(output_file, "w") as f:
            json.dump(analysis_json, f, indent=4)
        logging.info(f"Analysis saved to '{output_file}'")
        return analysis_json
    except json.JSONDecodeError:
        logging.error("Could not decode the output from the model into JSON format.")
        return None
    except Exception as e:
        logging.error(f"Error: {e}")
        return None


# --------------------- Function from simulation_data.py ---------------------
def monte_carlo_simulation(portfolio_data, scenario_data, num_simulations=10000):
    """Performs a Monte Carlo simulation on a portfolio based on market scenarios."""
    portfolio = portfolio_data
    scenarios = scenario_data["market_scenarios"]

    results = {}

    for scenario_key, scenario_details in scenarios.items():
        scenario_name = scenario_details["name"]
        sector_impacts = scenario_details.get("sector_impact", {})
        results[scenario_name] = {
            "portfolio_values": [],
            "average_return": 0,
            "std_dev_return": 0,
            "percentiles": {},
        }

        for _ in range(num_simulations):
            portfolio_value = 0
            for asset_name, asset_details in portfolio["assets"].items():
                sector = asset_details["sector"]
                quantity = asset_details["quantity"]
                initial_price = asset_details["initial_price"]

                price_change_percentage = 0
                if isinstance(sector_impacts, dict) and sector in sector_impacts:
                    price_change_percentage = np.random.normal(loc=sector_impacts[sector] / 100, scale=0.1)
                # Calculate the new price
                new_price = initial_price * (1 + price_change_percentage)

                portfolio_value += new_price * quantity
            results[scenario_name]["portfolio_values"].append(portfolio_value)

        # Calculate Results
        portfolio_values = results[scenario_name]["portfolio_values"]
        initial_portfolio_value = sum(
            asset["quantity"] * asset["initial_price"] for asset in portfolio["assets"].values())
        returns = [(value - initial_portfolio_value) / initial_portfolio_value for value in portfolio_values]

        results[scenario_name]["average_return"] = np.mean(returns)
        results[scenario_name]["std_dev_return"] = np.std(returns)
        results[scenario_name]["percentiles"] = {
            5: np.percentile(returns, 5),
            25: np.percentile(returns, 25),
            50: np.percentile(returns, 50),
            75: np.percentile(returns, 75),
            95: np.percentile(returns, 95),
        }
    logging.info(f"Monte Carlo simulation completed for scenario {scenario_name}")
    return results


def load_dataframes(filename="output_files/saved_dataframes.pkl"):
    try:
        logging.info(f"Loading dataframes from file: {filename}")
        with open(filename, 'rb') as file:
            saved_dataframes = pickle.load(file)
            logging.info(f"DataFrames successfully loaded from {filename}.")
            return saved_dataframes
    except FileNotFoundError:
        logging.error(f"File {filename} not found.")
        return None


def calculate_scenario_magnitudes(portfolio_data, scenario_data, saved_dataframes):
    scenario_results = {}

    for scenario_name, scenario_details in scenario_data["market_scenarios"].items():
        impacted_sectors = scenario_details["sector_impact"]

        # Filter assets in the impacted sectors
        relevant_assets = [
            symbol
            for symbol, details in portfolio_data["assets"].items()
            if details["sector"] in impacted_sectors
        ]

        # Calculate magnitudes for the scenario
        sector_magnitudes = {}
        for symbol in relevant_assets:
            df = saved_dataframes[symbol]
            sector = portfolio_data["assets"][symbol]["sector"]

            # Calculate magnitude as the absolute difference between first and last Close price
            magnitude = abs(df["Close"].iloc[-2] - df["Close"].iloc[-1])

            # Aggregate by sector
            if sector not in sector_magnitudes:
                sector_magnitudes[sector] = 0
            sector_magnitudes[sector] += magnitude

        # Calculate aggregated magnitude for the scenario
        aggregated_magnitude = sum(sector_magnitudes.values())

        # Store results
        scenario_results[scenario_name] = {
            "individual_magnitudes": sector_magnitudes,
            "aggregated_magnitude": aggregated_magnitude,
        }
    logging.info(f"Magnitudes calculated for scenario {scenario_name}")
    return scenario_results


def update_scenario_data(scenario_data, scenario_results):
    for scenario_id, results in scenario_results.items():
        # Update the sector impacts to include individual magnitudes
        scenario_data["market_scenarios"][scenario_id]["sector_impact"] = results["individual_magnitudes"]
        # Update aggregated magnitude
        scenario_data["market_scenarios"][scenario_id]["aggregated_magnitude"] = results["aggregated_magnitude"]

    logging.info(f"Scenario data updated with calculated magnitudes")
    return scenario_data


# --------------------- Streamlit App ---------------------
def main():
    st.title("Portfolio Analysis and Simulation App")

    # Initialize session state for data storage
    if 'stock_data' not in st.session_state:
        st.session_state['stock_data'] = {}

    if 'saved_dataframes' not in st.session_state:
        st.session_state['saved_dataframes'] = None

    if 'portfolio_data' not in st.session_state:
        st.session_state['portfolio_data'] = {}

    if 'scenario_data' not in st.session_state:
        st.session_state['scenario_data'] = {}

    if 'simulation_results' not in st.session_state:
        st.session_state['simulation_results'] = {}

    # Input for stock data as a JSON file
    st.header("1. Upload Portfolio Data (JSON)")
    uploaded_file = st.file_uploader("Upload your stock_data.json file", type=["json"])

    if uploaded_file:
        with st.spinner("Processing stock data..."):
          try:
              stock_data = json.load(uploaded_file)
              st.session_state['uploaded_stock_data'] = stock_data
              st.success("Stock data file uploaded successfully!")

              # Fetch stock data
              stock_symbols = [value["symbol"] for value in stock_data.values()]
              stock_dfs = fetch_stock_data(stock_symbols)

              # Save DataFrames in a dictionary for future use
              saved_dataframes = {}
              if stock_dfs:
                for symbol, df in stock_dfs.items():
                    if df is not None:
                        # Save DataFrame in the variable
                        saved_dataframes[symbol] = df
                        logging.info(f"Data for '{symbol}' loaded into variable.")
                    else:
                        logging.warning(f"No data found for '{symbol}'")
              else:
                  logging.error("Error occurred during fetching data. DataFrames are not returned.")

              # Save the dictionary to a local file
              with open('output_files/saved_dataframes.pkl', 'wb') as file:
                pickle.dump(saved_dataframes, file)
                logging.info(f"DataFrames successfully saved to output_files/saved_dataframes.pkl.")
              st.session_state['saved_dataframes'] = saved_dataframes
              # Store Data
              store_stock_data(stock_dfs)

              # Load the last price
              extracted_data = load_stock_data_and_extract_price(Config.STOCK_DATA_DIR)

              # Merge extracted price with the main dictionary
              merged_stock_data = merge_stock_data_with_price(stock_data, extracted_data)
              st.session_state['stock_data'] = merged_stock_data

              # Generate prompt for LLM
              formatted_prompt = generate_prompt(merged_stock_data)

              # Invoke LLM
              try:
                portfolio_output = invoke_llm_for_portfolio(formatted_prompt)
                portfolio_to_json(portfolio_output)
                st.session_state['portfolio_data'] = portfolio_output.model_dump()
                st.success("Stock data processed successfully. Portfolio data generated!")
              except Exception as e:
                st.error(f"An unexpected error occurred during the LLM invocation: {e}")
          except json.JSONDecodeError:
              st.error("Invalid JSON format. Please upload a valid JSON file.")
          except Exception as e:
              st.error(f"An error occurred while processing the uploaded file: {e}")


    st.header("2. Fetch Market Scenario")
    # Input for market analysis URL
    url = st.text_input("Enter Livemint URL (e.g. https://www.livemint.com/market/stock-market-news/page-7)",
                        value="https://www.livemint.com/market/stock-market-news/page-7")
    fetch_market_scenario = st.button("Fetch Market Scenario")

    if fetch_market_scenario:
        with st.spinner("Fetching market scenario..."):
          # Market Analysis
          context_data = scrape_website(url)  # Changed here
          scenario_data = invoke_llm_for_scenario(context_data)
          if scenario_data:
              st.session_state['scenario_data'] = scenario_data
              st.success("Market scenario data generated")
          else:
              st.error("Error occurred while generating market scenarios")

    st.header("3. Run Simulation")

    run_simulation = st.button("Run Monte Carlo Simulation")
    if run_simulation:
        with st.spinner("Running Monte Carlo Simulation..."):
          if st.session_state['portfolio_data'] and st.session_state['scenario_data']:
              saved_dataframes = st.session_state['saved_dataframes']
              # Update scenario data with magnitudes
              scenario_results = calculate_scenario_magnitudes(st.session_state['portfolio_data'],
                                                               st.session_state['scenario_data'], saved_dataframes)
              updated_scenario_data = update_scenario_data(st.session_state['scenario_data'], scenario_results)
              # Run Monte Carlo simulation
              simulation_results = monte_carlo_simulation(st.session_state['portfolio_data'], updated_scenario_data)
              st.session_state['simulation_results'] = simulation_results
              # Display simulation results
              st.subheader("Simulation Results")
              for scenario_name, results in simulation_results.items():
                  st.write(f"**Scenario:** {scenario_name}")
                  st.write(f"  **Average Return:** {results['average_return']:.4f}")
                  st.write(f"  **Std Dev Return:** {results['std_dev_return']:.4f}")
                  st.write("  **Return Percentiles:**")
                  for percentile, value in results["percentiles"].items():
                      st.write(f"    {percentile}th: {value:.4f}")
                  st.write("-" * 40)
              st.success("Monte Carlo simulation completed.")
          else:
              st.error("Please ensure both portfolio and scenario data are available.")


if __name__ == "__main__":
    main()