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

@st.cache_resource
def get_sp500_list():
    table = pd.read_html('https://en.wikipedia.org/wiki/List_of_S%26P_500_companies')
    return table[0]['Symbol'].tolist()

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', None),
        'P/B Ratio': info.get('priceToBook', None),
        'P/S Ratio': info.get('priceToSalesTrailing12Months', None),
        'Debt to Equity Ratio': info.get('debtToEquity', None),
        'Return on Equity': info.get('returnOnEquity', None),
        'Book-to-Market Ratio': 1 / info.get('priceToBook', float('nan')) 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

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)

def color_combined_score(value):
    if value > 0:
        color = 'green'
    elif value < 0:
        color = 'red'
    else:
        color = 'lightgrey'
    return f'background-color: {color};'

def filter_incomplete_stocks(df, required_columns):
    # Ensure all required columns exist in the DataFrame
    for column in required_columns:
        if column not in df.columns:
            df[column] = pd.NA
    return df.dropna(subset=required_columns)

st.title('S&P 500 Stock Comparison Tool')

sp500_list = get_sp500_list()
sp500_averages = load_sp500_averages('sp500_averages.csv')

scores_df = calculate_combined_scores_for_stocks(sp500_list, sp500_averages)
required_columns = ['P/E Ratio', 'P/B Ratio', 'P/S Ratio', 'Debt to Equity Ratio', 'Return on Equity', 'Book-to-Market Ratio']

scores_df_filtered = filter_incomplete_stocks(scores_df, required_columns)
scores_df_sorted = scores_df_filtered.sort_values(by='Combined Score', ascending=False)

col1, col2 = st.columns([3, 5])

with col1:
    st.subheader("Stock Overview")
    styled_scores_df = scores_df_sorted.style.applymap(color_combined_score, subset=['Combined Score'])
    st.dataframe(styled_scores_df)

with col2:
    st.subheader("Stock Details")
    sorted_tickers = scores_df_sorted['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, info = fetch_stock_data(ticker_symbol)
            comparison, _ = compare_to_index(stock_data, sp500_averages)
            
            st.write(f"**{info.get('longName', 'N/A')}** ({ticker_symbol})")
            st.write(info.get('longBusinessSummary', 'N/A'))
            
            for ratio in required_columns:
                value = stock_data.get(ratio, 'N/A')
                average = sp500_averages.loc[ratio, 'Average'] if ratio in sp500_averages.index else 'N/A'
                status = comparison.get(ratio, 'N/A')
                st.write(f"{ratio}: {value} (Your Ratio) | {average} (S&P 500 Avg) - {status}")