import gradio as gr
import pandas as pd
import plotly.express as px
from datetime import datetime


def get_overall_by_staking_traders(trades_df: pd.DataFrame) -> pd.DataFrame:
    """Gets the overall trades data"""
    trades_count = (
        trades_df.groupby(["month_year_week", "staking"], sort=False)
        .size()
        .reset_index()
    )
    trades_count.columns = trades_count.columns.astype(str)
    trades_count.rename(columns={"0": "trades"}, inplace=True)
    return trades_count


def plot_staking_trades_per_market_by_week(
    trades_df: pd.DataFrame, market_creator: str
) -> gr.Plot:

    # adding the total
    trades_all = trades_df.copy(deep=True)
    trades_all["market_creator"] = "all"

    # choose colour
    market_colour = "green"
    if market_creator == "pearl":
        market_colour = "darkviolet"
    elif market_creator == "quickstart":
        market_colour = "goldenrod"

    # merging both dataframes
    all_filtered_trades = pd.concat([trades_df, trades_all], ignore_index=True)
    all_filtered_trades = all_filtered_trades.sort_values(
        by="creation_timestamp", ascending=True
    )
    all_filtered_trades = all_filtered_trades.loc[
        all_filtered_trades["market_creator"] == market_creator
    ]
    if market_creator != "all":
        if market_creator == "pearl":
            # remove the staking data from quickstart
            all_filtered_trades = all_filtered_trades.loc[
                all_filtered_trades["staking"] != "quickstart"
            ]
        else:
            # remove the staking data from pearl
            all_filtered_trades = all_filtered_trades.loc[
                all_filtered_trades["staking"] != "pearl"
            ]
        all_filtered_trades["staking"] = all_filtered_trades["staking"].replace(
            {market_creator: "staking_traders", "non_Olas": "non_Olas_traders"}
        )
        colour_sequence = ["gray", market_colour, "black"]
        categories_sorted = {
            "staking": ["non_staking_traders", "staking_traders", "non_Olas_traders"]
        }
    else:
        all_filtered_trades["staking"] = all_filtered_trades["staking"].replace(
            {
                "pearl": "staking_pearl_traders",
                "quickstart": "staking_quickstart_traders",
                "non_Olas": "non_Olas_traders",
            }
        )
        colour_sequence = ["gray", "darkviolet", "goldenrod", "black"]
        categories_sorted = {
            "staking": [
                "non_staking_traders",
                "staking_pearl_traders",
                "staking_quickstart_traders",
                "non_Olas_traders",
            ]
        }
    trades = get_overall_by_staking_traders(all_filtered_trades)
    # Convert string dates to datetime and sort them
    all_dates_dt = sorted(
        [
            datetime.strptime(date, "%b-%d-%Y")
            for date in trades["month_year_week"].unique()
        ]
    )
    # Convert back to string format
    all_dates = [date.strftime("%b-%d-%Y") for date in all_dates_dt]

    fig = px.bar(
        trades,
        x="month_year_week",
        y="trades",
        color="staking",
        barmode="group",
        color_discrete_sequence=colour_sequence,
        category_orders=categories_sorted,
    )

    fig.update_layout(
        xaxis_title="Week",
        yaxis_title="Weekly nr of trades",
        legend=dict(yanchor="top", y=0.5),
        width=1000,  # Adjusted for better fit on laptop screens
        height=600,  # Adjusted for better fit on laptop screens
    )
    fig.update_xaxes(tickformat="%b %d\n%Y")
    # Update layout to force x-axis category order (hotfix for a sorting issue)
    fig.update_layout(xaxis={"categoryorder": "array", "categoryarray": all_dates})
    return gr.Plot(value=fig)