import yfinance as yf
import pandas as pd
import gradio as gr
import plotly.graph_objs as go

def fetch_and_rebalance(etfs_str, period, rebalance_threshold_percentage):
    etfs = [etf.strip() for etf in etfs_str.split(",")]
    # Fetch historical data for the given period
    etf_data = {etf: yf.Ticker(etf).history(period=period) for etf in etfs}
    
    # Ensure the index is a DatetimeIndex
    for etf in etfs:
        etf_data[etf].index = pd.to_datetime(etf_data[etf].index)
    
    # Fetch dividend data and ensure the index is a DatetimeIndex
    etf_dividends = {etf: yf.Ticker(etf).dividends for etf in etfs}
    
    # Ensure both data series have the same dates
    common_dates = etf_data[etfs[0]].index
    for etf in etfs[1:]:
        common_dates = common_dates.intersection(etf_data[etf].index)
    
    for etf in etfs:
        etf_data[etf] = etf_data[etf].reindex(common_dates)
        etf_dividends[etf] = etf_dividends[etf].reindex(common_dates, fill_value=0)

    # Initial investment
    initial_investment = 10000
    investments = {etf: initial_investment / len(etfs) for etf in etfs}
    shares = {etf: investments[etf] / etf_data[etf]['Close'].iloc[0] for etf in etfs}

    # DataFrame to track the investments
    df_rebalance = pd.DataFrame(columns=["Date"] + [f"{etf} Value" for etf in etfs] + ["Total Value"])
    rebalance_dates = []

    # Perform the rebalancing based on the percentage difference condition
    for date in common_dates[1:]:  # Skip the first date for the initial investment
        total_value = 0
        for etf in etfs:
            # Update the values based on daily returns
            previous_date = etf_data[etf].index[etf_data[etf].index.get_loc(date) - 1]
            investments[etf] *= (etf_data[etf]['Close'].loc[date] / etf_data[etf]['Close'].loc[previous_date])
            
            # Add dividends
            investments[etf] += etf_dividends[etf].loc[date] * shares[etf]

            total_value += investments[etf]
        
        # Check if rebalancing is needed
        value_difference = abs(investments[etfs[0]] - investments[etfs[1]])
        percentage_difference = value_difference / total_value
        rebalance_needed = percentage_difference > rebalance_threshold_percentage
        
        if rebalance_needed:
            # Rebalance to equal weight
            investments = {etf: total_value / len(etfs) for etf in etfs}
            shares = {etf: investments[etf] / etf_data[etf]['Close'].loc[date] for etf in etfs}
            rebalance_dates.append(date)
        
        # Append to DataFrame using pd.concat
        df_rebalance = pd.concat([
            df_rebalance,
            pd.DataFrame({
                "Date": [date], 
                **{f"{etf} Value": [investments[etf]] for etf in etfs}, 
                "Total Value": [total_value]
            })
        ], ignore_index=True)

    # Check if df_rebalance is empty
    if df_rebalance.empty:
        return "No data was appended."
    final_value = df_rebalance["Total Value"].iloc[-1]
    total_return = (final_value - initial_investment) / initial_investment * 100

    # Create the plot
    fig = go.Figure()
    for etf in etfs:
        fig.add_trace(go.Scatter(x=df_rebalance["Date"], y=df_rebalance[f"{etf} Value"], mode='lines', name=etf))
    
    # Add rebalancing markers
    for date in rebalance_dates:
        fig.add_trace(go.Scatter(x=[date], y=[df_rebalance.loc[df_rebalance['Date'] == date, f"{etf} Value"].values[0]],
                                 mode='markers', name='Rebalance', marker=dict(size=10, color='red')))

    fig.update_layout(title='ETF Valuations Over Time', xaxis_title='Date', yaxis_title='Value')

    return f"Final Value: {final_value}, Total Return: {total_return}%", fig

interface = gr.Interface(
    fn=fetch_and_rebalance, 
    inputs=[
        gr.Textbox(lines=1, placeholder="ETFs (comma separated)", value="TQQQ,JEPI"), 
        gr.Textbox(lines=1, placeholder="Period", value="ytd"), 
        gr.Slider(minimum=0.0, maximum=1.0, value=0.05, label="Rebalance Threshold Percentage")
    ],
    outputs=["text", "plot"]
)

interface.launch()