import streamlit as st
import yfinance as yf
import pandas as pd

# Correctly using st.cache_data as per Streamlit's new caching mechanism
@st.cache_data
def get_sp500_list():
    table = pd.read_html('https://en.wikipedia.org/wiki/List_of_S%26P_500_companies')
    return table[0]['Symbol'].tolist()

# Define the path to your CSV file with S&P 500 averages
sp500_averages_path = 'sp500_averages.csv'

def load_sp500_averages(filepath):
    return pd.read_csv(filepath, header=0, names=['Ratio', 'Average']).set_index('Ratio')

def fetch_stock_data(ticker_symbol):
    ticker = yf.Ticker(ticker_symbol)
    info = ticker.info
    financials = {
        'P/E Ratio': info.get('forwardPE'),
        'P/B Ratio': info.get('priceToBook'),
        'P/S Ratio': info.get('priceToSalesTrailing12Months'),
        'Debt to Equity Ratio': info.get('debtToEquity'),
        'Return on Equity': info.get('returnOnEquity'),
        'Book-to-Market Ratio': 1 / info.get('priceToBook') if info.get('priceToBook') else None
    }
    return financials, info

def compare_to_index(stock_ratios, index_averages):
    comparison = {}
    score = 0
    for ratio, value in stock_ratios.items():
        if ratio in index_averages.index and pd.notna(value):
            average = index_averages.loc[ratio, 'Average']
            comparison[ratio] = 'Undervalued' if value < average else 'Overvalued'
            score += 1 if value < average else -1
    return comparison, score

# Ensure this function is defined before it's called in the script
def calculate_combined_scores_for_stocks(stocks, index_averages):
    scores = []
    for ticker_symbol in stocks:
        stock_data, _ = fetch_stock_data(ticker_symbol)
        comparison, score = compare_to_index(stock_data, index_averages)
        scores.append({'Stock': ticker_symbol, 'Combined Score': score})
    return pd.DataFrame(scores)

# Define the color-coding function for the 'Combined Score' column
def color_combined_score(value):
    """Colors the combined score cell based on its value."""
    if value > 0:
        color = 'green'
    elif value < 0:
        color = 'red'
    else:
        color = 'none'
    return f'background-color: {color};'

def filter_incomplete_stocks(df):
    # Assuming that 'N/A' represents missing data, replace it with NaN
    df.replace('N/A', pd.NA, inplace=True)
    # Drop rows with any NaN values in the specified columns
    return df.dropna(subset=['P/E Ratio', 'P/B Ratio', 'P/S Ratio', 'Debt to Equity Ratio', 'Return on Equity', 'Book-to-Market Ratio'])
    
# User interface in Streamlit
st.title('S&P 500 Stock Comparison Tool')


# Load the current S&P 500 list and averages
sp500_list = get_sp500_list()
sp500_averages = load_sp500_averages(sp500_averages_path)

# Calculate combined scores for all S&P 500 stocks
scores_df = calculate_combined_scores_for_stocks(sp500_list, sp500_averages)
scores_df_sorted = scores_df.sort_values(by='Combined Score', ascending=False)

# Filter out stocks with incomplete data
scores_df_filtered = filter_incomplete_stocks(scores_df_sorted)

# Layout for displaying overview and details
col1, col2 = st.columns([3, 5])  # For example, this will give the first column 3/8 of the width

with col1:
    st.subheader("Stock Overview")
    # Apply color based on 'Combined Score' value and display the DataFrame
    styled_scores_df = scores_df_filtered.style.applymap(color_combined_score, subset=['Combined Score'])
    st.dataframe(styled_scores_df)

with col2:
    st.subheader("Stock Details")
    # Use the filtered DataFrame to create the dropdown for stock selection
    sorted_tickers = scores_df_filtered['Stock'].tolist()
    ticker_symbol = st.selectbox('Select a stock for details', options=sorted_tickers)
    if ticker_symbol:
        with st.spinner(f'Fetching data for {ticker_symbol}...'):
            stock_data, _ = fetch_stock_data(ticker_symbol)
            comparison, _ = compare_to_index(stock_data, sp500_averages)
            
            # Display the company name and ticker symbol
            st.write(f"**{ticker_symbol}**")
            
            # Display each financial ratio, its value, and the S&P 500 average
            for ratio in stock_data.keys():
                value = stock_data[ratio]
                average = sp500_averages.loc[ratio, 'Average'] if ratio in sp500_averages.index else 'N/A'
                comparison_result = comparison[ratio] if ratio in comparison else 'N/A'
                st.write(f"{ratio}: {value} (Your Ratio) | {average} (S&P 500 Avg) - {comparison_result}")