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"]
    )
    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())

    fig = px.bar(
        rolling_avg_roi_df,
        x="creation_date",
        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"""

    # Create a copy to avoid SettingWithCopyWarning
    local_df = traders_data.copy()

    # Ensure creation_date is datetime64[ns]
    # Since creation_date comes from .dt.date, it's a date object, not datetime
    local_df["creation_date"] = pd.to_datetime(local_df["creation_date"])

    # Aggregate ROI at the date level
    daily_avg = local_df.groupby("creation_date")["roi"].mean().reset_index()

    # Set the datetime index
    daily_avg = daily_avg.set_index("creation_date")

    # Now resample and rolling average
    weekly_avg = daily_avg.resample("W").mean()
    rolling_avg = weekly_avg.rolling(window=2).mean().reset_index()

    # Rename columns
    rolling_avg.rename(
        columns={"roi": "rolling_avg_roi", "date": "creation_date"},
        inplace=True,
    )

    return rolling_avg


def get_weekly_average_roi(traders_data: pd.DataFrame) -> pd.DataFrame:
    """Function to get the weekly average ROI by market_creator and total"""

    # Create a copy to avoid SettingWithCopyWarning
    local_df = traders_data.copy()

    # Ensure creation_date is datetime64[ns]
    # Since creation_date comes from .dt.date, it's a date object, not datetime
    local_df["creation_date"] = pd.to_datetime(local_df["creation_date"])

    # Aggregate ROI at the date level first
    daily_avg = local_df.groupby("creation_date")["roi"].mean().reset_index()

    # Set the datetime index
    daily_avg = daily_avg.set_index("creation_date")

    # Resample to weekly frequency and calculate mean
    weekly_avg = daily_avg.resample("W").mean().reset_index()

    return weekly_avg


def plot_weekly_average_roi(
    traders_data: 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
    # 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"]
    )
    filtered_traders_data["creation_date"] = filtered_traders_data[
        "creation_timestamp"
    ].dt.date

    # Get the weekly average ROI
    weekly_avg_roi_df = get_weekly_average_roi(filtered_traders_data)
    # plot the weekly average ROI
    print(weekly_avg_roi_df.head())

    fig = px.line(
        weekly_avg_roi_df,
        x="creation_date",
        y="weekly_avg_roi",
    )
    fig.update_layout(
        xaxis_title="Week",
        yaxis_title="Weekly average ROI for pearl agents",
    )
    return gr.Plot(
        value=fig,
    )