import pandas as pd
import gradio as gr
import gc
import plotly.express as px


def plot_rolling_average_dune(
    daa_df: pd.DataFrame,
) -> gr.Plot:
    """Function to plot the rolling average of daily active agents"""

    fig = px.bar(
        daa_df,
        x="tx_date",
        y="seven_day_trailing_avg",
    )
    fig.update_layout(
        xaxis_title="Date",
        yaxis_title="7-day rolling average of DAA",
    )

    return gr.Plot(
        value=fig,
    )


def plot_rolling_average(
    daa_df: pd.DataFrame,
    market_creator: str = None,
) -> gr.Plot:
    """Function to plot the rolling average of daily active agents"""
    if market_creator is not None:
        filtered_traders_df = daa_df.loc[daa_df["market_creator"] == market_creator]
        rolling_avg_df = get_sevenday_rolling_average(filtered_traders_df)
    else:
        rolling_avg_df = get_sevenday_rolling_average(daa_df)
    print(rolling_avg_df.head())

    # Ensure 'creation_date' is a column, not an index
    if "tx_date" not in rolling_avg_df.columns:
        rolling_avg_df = rolling_avg_df.reset_index()

    fig = px.bar(
        rolling_avg_df,
        x="tx_date",
        y="rolling_avg_traders",
    )
    fig.update_layout(
        xaxis_title="Date",
        yaxis_title="7-day rolling average of DAA",
    )

    return gr.Plot(
        value=fig,
    )


def get_sevenday_rolling_average(daa_df: pd.DataFrame) -> pd.DataFrame:
    """Function to get the 7-day rolling average of the number of unique
    trader_address"""
    # Create a local copy of the dataframe
    local_df = daa_df.copy()

    # Sort the dataframe by date
    local_df = local_df.sort_values(by="tx_date").set_index("tx_date")

    # Group by market_creator and calculate rolling average of unique trader_address
    rolling_avg = (
        local_df.resample("D")["trader_address"]
        .nunique()
        .rolling(window=7)
        .mean()
        .reset_index()
    )
    rolling_avg.rename(columns={"trader_address": "rolling_avg_traders"}, inplace=True)
    return rolling_avg


def plot_rolling_average_roi(
    traders_data: pd.DataFrame, pearl_agents: pd.DataFrame
) -> gr.Plot:
    """Function to plot the rolling average of ROI for pearl agents"""
    # Get the list of unique addresses from the daa_pearl_df
    unique_addresses = pearl_agents["safe_address"].unique()
    # Filter the weekly_roi_df to include only those addresses
    filtered_traders_data = traders_data[
        traders_data["trader_address"].isin(unique_addresses)
    ]
    # create the date column
    filtered_traders_data["creation_timestamp"] = pd.to_datetime(
        filtered_traders_data["creation_timestamp"], timezone="UTC"
    )
    filtered_traders_data["creation_date"] = filtered_traders_data[
        "creation_timestamp"
    ].dt.date

    # Get the 2-week rolling average of ROI
    rolling_avg_roi_df = get_twoweeks_rolling_average_roi(filtered_traders_data)
    print(rolling_avg_roi_df.head())
    # Ensure 'month_year_week' is a column, not an index
    if "month_year_week" not in rolling_avg_roi_df.columns:
        rolling_avg_roi_df = rolling_avg_roi_df.reset_index()
    fig = px.bar(
        rolling_avg_roi_df,
        x="month_year_week",
        y="rolling_avg_roi",
    )
    fig.update_layout(
        xaxis_title="Week",
        yaxis_title="2-week rolling average of ROI at the trader level",
    )

    return gr.Plot(
        value=fig,
    )


def get_twoweeks_rolling_average_roi(traders_data: pd.DataFrame) -> pd.DataFrame:
    """Function to get the 2-week rolling average of the ROI by market_creator and total"""

    # Aggregate ROI at the date level
    daily_avg = traders_data.groupby("creation_date")["roi"].mean().reset_index()
    # Ensure creation_date is datetime64[ns]
    daily_avg["creation_date"] = pd.to_datetime(daily_avg["creation_date"])
    daily_avg = daily_avg.set_index("creation_date")
    # Now resample and rolling
    weekly_avg = daily_avg.resample("W").mean()
    rolling_avg = weekly_avg.rolling(window=2).mean().reset_index()

    rolling_avg.rename(columns={"roi": "rolling_avg_roi"}, inplace=True)
    # Optionally, rename the date column for clarity
    rolling_avg.rename(columns={"creation_date": "month_year_week"}, inplace=True)
    return rolling_avg


def get_weekly_average_roi(weekly_roi_df: pd.DataFrame) -> pd.DataFrame:
    """Function to get the weekly average ROI for pearl agents"""
    # Create a local copy of the dataframe
    local_df = weekly_roi_df.copy()

    # Convert string dates to datetime
    local_df["month_year_week"] = pd.to_datetime(
        local_df["month_year_week"], format="%b-%d-%Y"
    )

    # Group by month_year_week and market_creator, then calculate the mean ROI
    weekly_avg_roi = (
        local_df.groupby(["month_year_week"], sort=False)["roi"].mean().reset_index()
    )
    return weekly_avg_roi


def plot_weekly_average_roi(
    weekly_roi_df: pd.DataFrame, pearl_agents: pd.DataFrame
) -> gr.Plot:
    """Function to plot the weekly average of ROI for pearl agents"""
    # Get the list of unique addresses from the daa_pearl_df
    unique_addresses = pearl_agents["safe_address"].unique()
    # Filter the weekly_roi_df to include only those addresses
    filtered_weekly_roi_df = weekly_roi_df[
        weekly_roi_df["trader_address"].isin(unique_addresses)
    ]
    # Select only the columns: "roi", "month_year_week", "trader_address"
    filtered_weekly_roi_df = filtered_weekly_roi_df[
        ["roi", "month_year_week", "trader_address"]
    ].copy()
    # Remove duplicates
    filtered_weekly_roi_df = filtered_weekly_roi_df.drop_duplicates(
        subset=["month_year_week", "trader_address"]
    )
    # Get the weekly average ROI
    weekly_avg_roi_df = get_weekly_average_roi(filtered_weekly_roi_df)
    # plot the weekly average ROI
    print(weekly_avg_roi_df.head())
    # Ensure 'month_year_week' is a column, not an index
    if "month_year_week" not in weekly_avg_roi_df.columns:
        weekly_avg_roi_df = weekly_avg_roi_df.reset_index()
    fig = px.line(
        weekly_avg_roi_df,
        x="month_year_week",
        y="roi",
    )
    fig.update_layout(
        xaxis_title="Week",
        yaxis_title="Weekly average ROI for pearl agents",
    )
    return gr.Plot(
        value=fig,
    )