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


print("Starting the app ...")

# Sanity check on the yfinance_utils
print("Sanity check on the yfinance_utils ...")
# print(fetch_yfinance_daily("MSFT", "2020-01-01", "2020-01-03"))
data = yf.download("MSFT", "2020-01-01", "2020-01-03")
print(data.head())

# Set page config
st.set_page_config(page_title="Asset Class Comparison", layout="wide")

# Title and description
st.title("Asset Class Performance Comparison")
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 Large Cap",
        "NASDAQ 100",
        "S&P 500",
        "Dow Jones",
        "Microsoft",
        "Google",
        "Nvidia",
        "Apple",
        "Amazon",
        "Tesla",
        "Netflix",
        "Meta",
        "Bitcoin",
        "Ethereum",
        "Solana",
        "Dogecoin",
    ],
    default=[
        "Fixed Deposit", 
        "HDB",
        "Gold", 
        "US Treasury Bonds", "SGS Bonds", 
        "S&P 500", "Dow Jones", "NASDAQ Composite", #"NASDAQ Large Cap", "NASDAQ 100",
        "Microsoft", "Google", "Nvidia",
        "Bitcoin"  
        ]
)

# Today's date for reference
today = datetime.now().date()

usd_to_sgd = get_usd_sgd_rate() if currency == "SGD" else 1.0
currency_symbol = "$" if currency == "USD" else "S$"

# Create a dictionary of tickers for yfinance
tickers = {
    "Gold": "GC=F",
    "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": "GOOGL",
    "Nvidia": "NVDA",
    "Apple": "AAPL",
    "Amazon": "AMZN",
    "Tesla": "TSLA",
    "Netflix": "NFLX",
    "Meta": "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_yfinance_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
        
        summary_data.append({
            "Asset": asset,
            f"Final Value ({currency_symbol})": final_value,
            "Annualized Return (%)": annualized_return,
            "Positive Years": f"{positive_years}/{total_years}",
            "Positive Years %": positive_percentage,
        })
    else:
        summary_data.append({
            "Asset": asset,
            f"Final Value ({currency_symbol})": None,
            "Annualized 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["Annualized Return (%)"] = df["Annualized 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"
        ),
        "Annualized 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)}")