import streamlit as st
import pandas as pd
import yfinance as yf
import plotly.graph_objects as go
from datetime import datetime, timedelta
import numpy as np

# Import utility functions
from utils.yfinance_utils import fetch_yfinance_daily
from utils.currency_utils import get_usd_sgd_rate
from utils.fd_utils import calculate_fd_returns
from utils.hdb_utils import calculate_hdb_returns

def fetch_stock_daily(symbol, start_date, end_date):
    print (f"Fetching stock data for {symbol} from {start_date} to {end_date}")
    start_date_dt = pd.Timestamp(start_date)
    end_date_dt = pd.Timestamp(end_date)
    try:
        folder = "notebooks/stock_data/"
        data = pd.read_csv(
            f"{folder}{symbol}.csv", 
            parse_dates=["Date"]
            )
        print (f'Date range: {data["Date"].min()} - {data["Date"].max()}')
        print (f"start_date_dt: {start_date_dt}, end_date_dt: {end_date_dt}")
        data = data[(data["Date"] >= start_date_dt) & (data["Date"] <= end_date_dt)]
        data.set_index("Date", inplace=True)
        data = data["Close"]
        return data
    except FileNotFoundError:
        print(f"File {folder}{symbol}.csv not found.")
        return None
    

print("Starting the app ...")

# Set page config
st.set_page_config(
    page_title="Asset Class Comparison", 
    layout="wide",
    menu_items={
        'Get Help': 'https://github.com/yourusername/asset-class-comparison',
        'Report a bug': 'https://github.com/yourusername/asset-class-comparison/issues',
        'About': 'Compare the performance of different asset classes over time'
    }
)

# Title and description
st.title("Asset Class Performance Comparison")

# Initialize session state to track which tab is active
if 'active_tab' not in st.session_state:
    st.session_state.active_tab = 0

# Get URL parameters and set the active tab accordingly
if "tab" in st.query_params:
    tab_value = st.query_params["tab"].lower()
    if tab_value in ["readme", "2", "calculator", "formulas", "guide"]:
        st.session_state.active_tab = 1
    elif tab_value in ["1", "comparison", "main", "invest", "investment"]:
        st.session_state.active_tab = 0

# Create a tab selector
tab_options = ["Investment Comparison", "README & Formulas"]
selected_tab = st.radio("Select Tab:", tab_options, horizontal=True, index=st.session_state.active_tab, label_visibility="collapsed")

# Update session state and URL when tab changes
if selected_tab == "Investment Comparison" and st.session_state.active_tab != 0:
    st.session_state.active_tab = 0
    st.query_params["tab"] = "comparison"
    st.rerun()
elif selected_tab == "README & Formulas" and st.session_state.active_tab != 1:
    st.session_state.active_tab = 1
    st.query_params["tab"] = "readme"
    st.rerun()
    
# Display URL tips
with st.expander("URL Tips 🔗"):
    st.markdown("""
    You can use URL parameters to directly navigate to specific tabs:
    
    * Main investment comparison: `?tab=1` or `?tab=comparison` 
    * README & Calculator: `?tab=2` or `?tab=readme`
    
    Example: `http://yourapp.com/?tab=readme`
    
    These links can be bookmarked or shared to go directly to a specific tab.
    
    You can also set the tab programmatically in your app:
    ```python
    # To switch to the README tab
    st.query_params["tab"] = "readme"
    
    # To clear all parameters
    st.query_params.clear()
    ```
    """)

# Display content based on the selected tab
if st.session_state.active_tab == 0:
    # INVESTMENT COMPARISON TAB
    st.write("Compare the performance of different asset classes over time")
    # st.write("Note: Cryptocurrencies (BTC, ETH, SOL, DOGE) are highly volatile and should be considered high-risk investments")
    
    # Sidebar for user inputs
    st.sidebar.header("Investment Parameters")
    currency = st.sidebar.selectbox("Display Currency", ["USD", "SGD"], index=0)
    initial_investment = st.sidebar.number_input(f"Initial Investment Amount ({currency})", min_value=1000, value=10000, step=1000)
    start_date = st.sidebar.date_input("Start Date", value=datetime.now() - timedelta(days=365*25))
    user_end_date = st.sidebar.date_input("End Date", value=datetime.now())
    
    fd_rate = st.sidebar.number_input("Fixed Deposit Rate (%)", min_value=0.0, value=2.9, step=0.1) / 100
    use_log_scale = st.sidebar.checkbox("Use Log Scale", value=True)
    
    # Calculate and display investment period
    investment_days = (user_end_date - start_date).days
    investment_years = investment_days / 365
    st.write(f"Investment Period: {investment_days} days ({investment_years:.1f} years)")
    
    # Asset selection
    selected_assets = st.sidebar.multiselect(
        "Select Assets to Compare",
        [
            "Fixed Deposit",
            "HDB",
            "Gold",
            "SGS Bonds",
            "US Treasury Bonds",
            "NASDAQ Composite",
            "NASDAQ-100",
            "Invesco QQQ Trust",
            "Fidelity NASDAQ Composite Index ETF",
            "Invesco NASDAQ 100 ETF",
            "S&P 500",
            "Dow Jones",
            "Microsoft",
            "Google (Alphabet)",
            "Nvidia",
            "Apple",
            "Amazon",
            "Tesla",
            "Netflix",
            "Meta (Facebook)",
            "Bitcoin",
            "Ethereum",
            "Solana",
            "Dogecoin",
        ],
        default=[
            "Fixed Deposit", 
            "HDB",
            "Gold", 
            "US Treasury Bonds", 
            "SGS Bonds", 
            "S&P 500", 
            "Dow Jones", 
            "NASDAQ Composite",
            "Microsoft", 
            "Google (Alphabet)", 
            "Meta (Facebook)",
            "Nvidia",
            "Bitcoin"  
        ]
    )

    # Today's date for reference
    today = datetime.now().date()
    today = today# - timedelta(days=1) # Temporarily set to yesterday

    # usd_to_sgd = get_usd_sgd_rate() if currency == "SGD" else 1.0
    usd_to_sgd = 1.30 # Temporarily set to 1.30 to test the app
    currency_symbol = "$" if currency == "USD" else "S$"

    # Create a dictionary of tickers for yfinance
    tickers = {
        "Gold": "GLD",
        # "HDB": "A12.SI",
        "SGS Bonds": "A35.SI",  # Nikko AM SGD Investment Grade Corporate Bond ETF
        "US Treasury Bonds": "TLT",  # iShares 20+ Year Treasury Bond ETF
        "NASDAQ Composite": "^IXIC",
        "NASDAQ Large Cap": "^NDX",
        "NASDAQ 100": "^NDX",
        "S&P 500": "^GSPC",
        "Dow Jones": "^DJI",
        "Microsoft": "MSFT",
        "Google (Alphabet)": "GOOGL",
        "Nvidia": "NVDA",
        "Apple": "AAPL",
        "Amazon": "AMZN",
        "Tesla": "TSLA",
        "Netflix": "NFLX",
        "Meta (Facebook)": "META",
        "Bitcoin": "BTC-USD",
        "Ethereum": "ETH-USD",
        "Solana": "SOL-USD",
        "Dogecoin": "DOGE-USD",
    }

    # Determine the effective end date for each asset
    asset_end_dates = {}
    for asset in selected_assets:
        if asset == "Fixed Deposit":
            asset_end_dates[asset] = user_end_date
        else:
            if user_end_date > today:
                asset_end_dates[asset] = today
            else:
                asset_end_dates[asset] = user_end_date

    # Warn the user if a future end date is selected for market assets
    if any(user_end_date > today and asset != "Fixed Deposit" for asset in selected_assets):
        st.warning(f"Market data is only available up to today ({today}). For market assets, the end date has been set to today.")

    # Calculate returns for each selected asset
    asset_series = {}
    failed_assets = []
    actual_start_dates = {}

    for asset in selected_assets:
        asset_start = start_date
        asset_end = asset_end_dates[asset]
        if asset == "Fixed Deposit":
            fd_index = pd.date_range(start=asset_start, end=user_end_date)
            daily_rate = (1 + fd_rate) ** (1/365) - 1
            fd_values = initial_investment * (1 + daily_rate) ** np.arange(len(fd_index))
            if currency == "SGD":
                fd_values = fd_values * usd_to_sgd
            asset_series[asset] = pd.Series(fd_values, index=fd_index)
            actual_start_dates[asset] = asset_start
        elif asset == "HDB":
            hdb_values = calculate_hdb_returns(asset_start, asset_end, initial_investment)
            if hdb_values is not None:
                if currency == "SGD":
                    hdb_values = hdb_values * usd_to_sgd
                asset_series[asset] = hdb_values
                actual_start_dates[asset] = asset_start
            else:
                failed_assets.append(asset)
        else:
            price_data = fetch_stock_daily(tickers[asset], asset_start, asset_end)
            if price_data is not None and not price_data.empty:
                price_data = price_data.sort_index()
                actual_start = price_data.index[0]
                actual_start_dates[asset] = actual_start
                aligned_index = pd.date_range(start=actual_start, end=asset_end)
                price_data = price_data.reindex(aligned_index)
                price_data = price_data.ffill()
                asset_values = initial_investment * (price_data / price_data.iloc[0])
                if currency == "SGD":
                    asset_values = asset_values * usd_to_sgd
                asset_series[asset] = asset_values
            else:
                failed_assets.append(asset)

    # Combine all asset series into a single DataFrame
    if asset_series:
        returns_data = pd.DataFrame(asset_series)
    else:
        returns_data = pd.DataFrame()

    # Remove failed assets from selected_assets (except FD)
    selected_assets = [asset for asset in selected_assets if asset not in failed_assets or asset == "Fixed Deposit"]

    if not selected_assets:
        st.error("No assets could be loaded. Please try different assets.")
        st.stop()

    # Create the plot
    fig = go.Figure()

    # Add vertical lines for every 5 years
    start_year = returns_data.index[0].year
    end_year = returns_data.index[-1].year
    for year in range(start_year, end_year + 1, 5):
        fig.add_vline(x=datetime(year, 1, 1), line_dash="dash", line_color="gray", opacity=0.3)

    for asset in selected_assets:
        fig.add_trace(go.Scatter(
            x=returns_data.index,
            y=returns_data[asset],
            name=asset,
            mode='lines'
        ))

    fig.update_layout(
        title="Asset Performance Comparison",
        xaxis_title="Date",
        yaxis_title=f"Investment Value ({currency_symbol})",
        hovermode="x unified",
        height=600,
        yaxis_type="log" if use_log_scale else "linear"
    )

    # Display the plot
    st.plotly_chart(fig, use_container_width=True)

    # Create a summary table
    st.subheader("Investment Summary")
    summary_data = []
    for asset in selected_assets:
        valid_series = returns_data[asset].dropna()
        if not valid_series.empty:
            final_value = valid_series.iloc[-1]
            days = (valid_series.index[-1] - valid_series.index[0]).days
            years = days / 365
            annualized_return = ((final_value / initial_investment) ** (1/years) - 1) * 100
            
            # Calculate yearly return statistics
            yearly_data = valid_series.resample('YE').first()
            yearly_returns = yearly_data.pct_change().dropna()
            positive_years = (yearly_returns > 0).sum()
            total_years = len(yearly_returns)
            positive_percentage = (positive_years / total_years) * 100
            
            # Calculate return multiple
            return_multiple = final_value / initial_investment
            
            # Calculate simple interest based annual return
            simple_annual_return = ((return_multiple - 1) / years) * 100
            
            summary_data.append({
                "Asset": asset,
                f"Final Value ({currency_symbol})": final_value,
                "Return Multiple": return_multiple,
                "Annualized Return (%)": annualized_return,
                "Simple Annual Return (%)": simple_annual_return,
                "Positive Years": f"{positive_years}/{total_years}",
                "Positive Years %": positive_percentage,
            })
        else:
            summary_data.append({
                "Asset": asset,
                f"Final Value ({currency_symbol})": None,
                "Return Multiple": None,
                "Annualized Return (%)": None,
                "Simple Annual Return (%)": None,
                "Positive Years": "N/A",
                "Positive Years %": None,
            })

    # Convert to DataFrame
    df = pd.DataFrame(summary_data)

    # Format the display values
    df[f"Final Value ({currency_symbol})"] = df[f"Final Value ({currency_symbol})"].apply(lambda x: f"{x:,.2f}" if x is not None else "N/A")
    df["Return Multiple"] = df["Return Multiple"].apply(lambda x: f"{x:.2f}" if x is not None else "N/A")
    df["Annualized Return (%)"] = df["Annualized Return (%)"].apply(lambda x: f"{x:.2f}" if x is not None else "N/A")
    df["Simple Annual Return (%)"] = df["Simple Annual Return (%)"].apply(lambda x: f"{x:.2f}" if x is not None else "N/A")
    df["Positive Years %"] = df["Positive Years %"].apply(lambda x: f"{x:.1f}" if x is not None else "N/A")

    # Display the summary table with sorting enabled
    st.dataframe(
        df,
        hide_index=True,
        column_config={
            f"Final Value ({currency_symbol})": st.column_config.NumberColumn(
                format="%.2f"
            ),
            "Return Multiple": st.column_config.NumberColumn(
                format="%.2f"
            ),
            "Annualized Return (%)": st.column_config.NumberColumn(
                format="%.2f"
            ),
            "Simple Annual Return (%)": st.column_config.NumberColumn(
                format="%.2f"
            ),
            "Positive Years %": st.column_config.NumberColumn(
                format="%.1f"
            ),
            "Performance": st.column_config.ImageColumn(
                "Performance"
            )
        }
    )

    # Calculate and display final returns
    st.subheader("Final Investment Values")
    for asset in selected_assets:
        valid_series = returns_data[asset].dropna()
        if not valid_series.empty:
            final_value = valid_series.iloc[-1]
            st.write(f"{asset}: {currency_symbol}{final_value:,.2f}")
        else:
            st.write(f"{asset}: Data unavailable")

    # Calculate and display annualized returns
    st.subheader("Annualized Returns")
    for asset in selected_assets:
        valid_series = returns_data[asset].dropna()
        if len(valid_series) > 1:
            actual_start = actual_start_dates[asset]
            days = (valid_series.index[-1] - valid_series.index[0]).days
            years = days / 365
            final_value = valid_series.iloc[-1]
            annualized_return = ((final_value / initial_investment) ** (1/years) - 1) * 100
            if pd.Timestamp(actual_start).date() > start_date:
                st.write(f"{asset}: {annualized_return:.2f}% (Data available from {actual_start.strftime('%Y-%m-%d')})")
            else:
                st.write(f"{asset}: {annualized_return:.2f}%")
        else:
            st.write(f"{asset}: N/A")

    # Calculate and display yearly return statistics
    st.subheader("Yearly Return Statistics")
    for asset in selected_assets:
        valid_series = returns_data[asset].dropna()
        if len(valid_series) > 1:
            # Resample to yearly data
            yearly_data = valid_series.resample('YE').first()
            
            # Calculate yearly returns
            yearly_returns = yearly_data.pct_change().dropna()
            
            # Count positive and negative years
            positive_years = (yearly_returns > 0).sum()
            total_years = len(yearly_returns)
            positive_percentage = (positive_years / total_years) * 100
            
            st.write(f"{asset}: {positive_years} out of {total_years} years ({positive_percentage:.1f}%) had positive returns")
        else:
            st.write(f"{asset}: Insufficient data for yearly analysis")

    # Show warnings for data availability
    for asset in selected_assets:
        if asset in actual_start_dates and pd.Timestamp(actual_start_dates[asset]).date() > start_date:
            st.warning(f"Data for {asset} is only available from {actual_start_dates[asset].strftime('%Y-%m-%d')}. The analysis starts from this date.")

    # Show warning for failed assets
    if failed_assets:
        st.warning(f"Could not load data for the following assets: {', '.join(failed_assets)}")

else:
    # README & FORMULAS TAB
    st.header("Investment Calculation Guide")
    
    st.markdown("""
    ## Key Investment Formulas
    
    This guide explains the key formulas used in the investment comparison tool and how to interpret them.
    
    ### Return Multiple
    
    The return multiple shows how many times your initial investment has grown over the investment horizon.
    
    ```
    Return Multiple = Final Value / Initial Investment
    ```
    
    **Example:**
    - Initial Investment: $10,000
    - Final Value: $25,000
    - Return Multiple = 25,000 / 10,000 = 2.5x
    
    This means your investment has grown to 2.5 times its original value over the investment period.
    
    ### Annualized Return (CAGR)
    
    The annualized return is the geometric average annual return over the investment period, expressed as a percentage. This is equivalent to the Compound Annual Growth Rate (CAGR).
    
    ```
    Annualized Return (%) = ((Final Value / Initial Investment) ^ (1 / Years) - 1) × 100
    ```
    
    **Example:**
    - Initial Investment: $10,000
    - Final Value: $20,000
    - Investment Period: 10 years
    - Annualized Return = ((20,000 / 10,000) ^ (1/10) - 1) × 100 = 7.18%
    
    This means your investment grew at an average rate of 7.18% per year, compounded annually.
    
    ### Simple Annual Return
    
    Simple annual return calculates the average yearly return by dividing the total return by the number of years.
    
    ```
    Simple Annual Return (%) = ((Return Multiple - 1) / Years) × 100
    ```
    
    **Example:**
    - Return Multiple: 2.0x
    - Investment Period: 10 years
    - Simple Annual Return = ((2.0 - 1) / 10) × 100 = 10.0%
    
    This represents the linear average annual return, which is typically higher than CAGR.
    
    ### CAGR vs. Simple Annual Return
    
    **Compound Annual Growth Rate (CAGR)** and simple annual return represent two different ways to measure annual performance:
    
    - **CAGR** (our "Annualized Return") assumes compounding - each year's gains build upon the previous year's value
    - **Simple Annual Return** assumes linear growth - dividing the total gain by the number of years
    
    For the same investment period and final value, the simple annual return is typically higher than CAGR. The difference increases with longer investment periods and higher returns.
    
    **Mathematical relationship:**
    - CAGR: (1 + r)^n = Final Value / Initial Value
    - Simple: 1 + (r × n) = Final Value / Initial Value
    
    Where r is the rate (as a decimal) and n is the number of years.
    
    ### Calculating Final Value from Annualized Return
    
    If you know the annualized return, you can calculate the expected final value:
    
    ```
    Final Value = Initial Investment × (1 + Annualized Return/100) ^ Years
    ```
    
    **Example:**
    - Initial Investment: $10,000
    - Annualized Return: 7%
    - Investment Period: 10 years
    - Final Value = $10000 × (1 + 0.07)^{10} = 19672$
    
    ### Positive Years Percentage
    
    The percentage of years with positive returns:
    
    ```
    Positive Years (%) = (Number of Years with Positive Returns / Total Years) × 100
    ```
    
    This helps assess the consistency of positive performance.
    """)
    
    st.subheader("How to Use This Information")
    
    st.markdown("""
    - **Return Multiple** is useful for quickly understanding total growth
    - **Annualized Return** allows fair comparison between investments with different time periods
    - **Positive Years %** helps assess the consistency and volatility of returns
    
    When comparing investments:
    1. Higher annualized returns are generally better
    2. More consistent returns (higher positive years %) may indicate lower volatility
    3. Return multiples must be viewed in context of the time period
    """)
    
    # Add interactive investment calculator
    st.subheader("Investment Calculator")
    
    st.markdown("""
    Use this calculator to estimate future investment values based on different parameters.
    """)
    
    # Add explanation about simple vs compound interest
    with st.expander("Simple vs. Compound Interest"):
        st.markdown("""
        ### Simple Interest vs. Compound Interest
        
        **Simple Interest** is calculated only on the initial principal, without considering accumulated interest.
        
        ```
        Final Value = Initial Investment × (1 + (Annual Return/100 × Years))
        ```
        
        **Compound Interest** is calculated on both the initial principal and the accumulated interest.
        
        ```
        Final Value = Initial Investment × (1 + Annual Return/100)^Years
        ```
        
        **Key Differences:**
        - Simple interest grows linearly (straight line on chart)
        - Compound interest grows exponentially (curved line on chart)
        - Over long periods, the difference becomes substantial
        - Most real-world investments compound (stocks, bonds, etc.)
        """)
    
    calc_col1, calc_col2 = st.columns(2)
    
    with calc_col1:
        calc_initial = st.number_input("Initial Investment ($)", min_value=100, value=10000, step=1000)
        calc_years = st.number_input("Investment Period (Years)", min_value=1, value=10, step=1)
        interest_type = st.radio("Interest Type", ["Compound", "Simple"], index=0)
    
    with calc_col2:
        calc_return = st.number_input("Annual Return (%)", min_value=0.0, value=7.0, step=0.5)
        calc_currency = st.selectbox("Currency", ["USD", "SGD"], index=0)
    
    # Calculate results based on interest type
    calc_currency_symbol = "$" if calc_currency == "USD" else "S$"
    
    if interest_type == "Compound":
        calc_final_value = calc_initial * ((1 + calc_return/100) ** calc_years)
        interest_formula = f"${calc_initial:,.0f} × (1 + {calc_return/100:.4f})^{calc_years}$"
    else:  # Simple Interest
        calc_final_value = calc_initial * (1 + (calc_return/100 * calc_years))
        interest_formula = f"${calc_initial:,.0f} × (1 + {calc_return/100:.4f} × {calc_years})$"
    
    calc_return_multiple = calc_final_value / calc_initial
    total_interest = calc_final_value - calc_initial
    
    # Display results in a highlighted box
    st.markdown(f"""
    #### Results
    """)
    
    results_col1, results_col2, results_col3 = st.columns(3)
    
    with results_col1:
        st.metric("Final Investment Value", f"{calc_currency_symbol}{calc_final_value:,.2f}")
        
    with results_col2:
        st.metric("Return Multiple", f"{calc_return_multiple:.2f}x")
        
    with results_col3:
        st.metric("Total Interest Earned", f"{calc_currency_symbol}{total_interest:,.2f}")
    
    st.markdown(f"**Formula Applied:** {interest_formula} = {calc_currency_symbol}{calc_final_value:,.2f}")
    
    # Add a chart to visualize growth and compare both methods
    years_range = list(range(0, calc_years + 1))
    
    # Calculate both types of growth for comparison
    compound_values = [calc_initial * ((1 + calc_return/100) ** year) for year in years_range]
    simple_values = [calc_initial * (1 + (calc_return/100 * year)) for year in years_range]
    
    # Create DataFrame with both growth types
    growth_chart = pd.DataFrame({
        'Year': years_range,
        'Compound Interest': compound_values,
        'Simple Interest': simple_values
    })
    
    st.subheader("Investment Growth Projection")
    
    # Highlight the selected interest type
    if interest_type == "Compound":
        st.line_chart(growth_chart, x='Year', y=['Compound Interest', 'Simple Interest'])
    else:
        st.line_chart(growth_chart, x='Year', y=['Simple Interest', 'Compound Interest'])