import gradio as gr
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from datetime import datetime

HEIGHT = 400
WIDTH = 1100


def prepare_trades(trades_df: pd.DataFrame) -> pd.DataFrame:
    """Prepares the trades data for analysis."""
    trades_df["creation_timestamp"] = pd.to_datetime(trades_df["creation_timestamp"])
    trades_df["creation_date"] = trades_df["creation_timestamp"].dt.date
    trades_df["creation_timestamp"] = trades_df["creation_timestamp"].dt.tz_convert(
        "UTC"
    )
    trades_df = trades_df.sort_values(by="creation_timestamp", ascending=True)
    trades_df["month_year"] = (
        trades_df["creation_timestamp"].dt.to_period("M").astype(str)
    )
    trades_df["month_year_week"] = (
        trades_df["creation_timestamp"]
        .dt.to_period("W")
        .dt.start_time.dt.strftime("%b-%d-%Y")
    )
    trades_df["winning_trade"] = trades_df["winning_trade"].astype(int)
    return trades_df


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


def get_overall_by_market_trades(trades_df: pd.DataFrame) -> pd.DataFrame:
    """Gets the overall trades data"""
    trades_count = (
        trades_df.groupby(["month_year_week", "market_creator"], 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 get_overall_winning_trades(trades_df: pd.DataFrame) -> pd.DataFrame:
    """Gets the overall winning trades data for the given tools and calculates the winning percentage."""
    winning_trades = (
        trades_df.groupby(["month_year_week"])["winning_trade"].sum()
        / trades_df.groupby(["month_year_week"])["winning_trade"].count()
        * 100
    )
    # winning_trades is a series, give it a dataframe
    winning_trades = winning_trades.reset_index()
    winning_trades.columns = winning_trades.columns.astype(str)
    winning_trades.columns = ["month_year_week", "winning_trade"]
    return winning_trades


def get_overall_winning_by_market_trades(trades_df: pd.DataFrame) -> pd.DataFrame:
    """Gets the overall winning trades data for the given tools and calculates the winning percentage."""
    winning_trades = (
        trades_df.groupby(["month_year_week", "market_creator"], sort=False)[
            "winning_trade"
        ].sum()
        / trades_df.groupby(["month_year_week", "market_creator"], sort=False)[
            "winning_trade"
        ].count()
        * 100
    )

    winning_trades = winning_trades.reset_index()
    winning_trades.columns = winning_trades.columns.astype(str)
    winning_trades.columns = ["month_year_week", "market_creator", "winning_trade"]
    return winning_trades


def get_overall_winning_by_market_and_trader_type(
    trades_df: pd.DataFrame,
) -> pd.DataFrame:
    """Gets the overall winning trades data for the given tools and calculates the winning percentage."""
    # Group by week, market_creator and staking_type
    winning_trades = (
        trades_df.groupby(
            ["month_year_week", "market_creator", "staking_type"], sort=False
        )["winning_trade"].sum()
        / trades_df.groupby(
            ["month_year_week", "market_creator", "staking_type"], sort=False
        )["winning_trade"].count()
        * 100
    )

    winning_trades = winning_trades.reset_index()
    winning_trades.columns = winning_trades.columns.astype(str)
    winning_trades.columns = [
        "month_year_week",
        "market_creator",
        "staking_type",
        "winning_trade",
    ]
    return winning_trades


def plot_trades_by_week(trades_df: pd.DataFrame) -> gr.BarPlot:
    """Plots the weekly trades data ."""
    return gr.BarPlot(
        value=trades_df,
        x="month_year_week",
        y="trades",
        show_label=True,
        interactive=True,
        show_actions_button=True,
        tooltip=["month_year_week", "trades"],
        height=HEIGHT,
        width=WIDTH,
    )


def integrated_plot_trades_per_market_by_week(trades_df: pd.DataFrame) -> gr.Plot:

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

    # 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
    )

    trades = get_overall_by_market_trades(all_filtered_trades)
    fig = px.bar(
        trades,
        x="month_year_week",
        y="trades",
        color="market_creator",
        barmode="group",
        color_discrete_sequence=["purple", "goldenrod", "darkgreen"],
        category_orders={"market_creator": ["pearl", "quickstart", "all"]},
    )

    fig.update_layout(
        xaxis_title="Week",
        yaxis_title="Weekly nr of trades",
        legend=dict(yanchor="top", y=0.5),
    )
    fig.update_layout(width=WIDTH, height=HEIGHT)
    fig.update_xaxes(tickformat="%b %d\n%Y")
    return gr.Plot(value=fig)


def integrated_plot_trades_per_market_by_week_v2(trades_df: pd.DataFrame) -> gr.Plot:
    # adding the total
    trades_all = trades_df.copy(deep=True)
    trades_all["market_creator"] = "all"

    # 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
    )
    # Create binary staking category
    all_filtered_trades["staking_type"] = all_filtered_trades["staking"].apply(
        lambda x: "non_Olas" if x == "non_Olas" else "Olas"
    )

    # Group by week, market_creator and staking_type
    trades = (
        all_filtered_trades.groupby(
            ["month_year_week", "market_creator", "staking_type"], sort=False
        )
        .size()
        .reset_index(name="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]
    # Combine the traces
    final_traces = []
    market_colors = {"pearl": "darkviolet", "quickstart": "goldenrod", "all": "green"}
    market_darker_colors = {
        "pearl": "purple",
        "quickstart": "darkgoldenrod",
        "all": "darkgreen",
    }

    # Process both Olas and non-Olas traces for each market together
    for market in ["pearl", "quickstart", "all"]:
        market_data = trades[trades["market_creator"] == market]
        # Create a dictionary to store the Olas values for each week
        olas_values = dict(
            zip(
                market_data[market_data["staking_type"] == "Olas"]["month_year_week"],
                market_data[market_data["staking_type"] == "Olas"]["trades"],
            )
        )
        # First add 'Olas' trace
        olas_data = market_data[market_data["staking_type"] == "Olas"]
        olas_trace = go.Bar(
            x=olas_data["month_year_week"],
            y=olas_data["trades"],
            name=f"{market}-Olas",
            marker_color=market_colors[market],
            offsetgroup=market,  # Keep the market grouping
            showlegend=True,
        )

        # Then add 'non_Olas' trace with base set to olas values
        non_Olas_data = market_data[market_data["staking_type"] == "non_Olas"]
        non_Olas_trace = go.Bar(
            x=non_Olas_data["month_year_week"],
            y=non_Olas_data["trades"],
            name=f"{market}-non_Olas",
            marker_color=market_darker_colors[market],
            offsetgroup=market,  # Keep the market grouping
            base=[olas_values.get(x, 0) for x in non_Olas_data["month_year_week"]],
            showlegend=True,
        )

        final_traces.extend([olas_trace, non_Olas_trace])

    # Create new figure with the combined traces
    fig = go.Figure(data=final_traces)

    # Update layout
    fig.update_layout(
        xaxis_title="Week",
        yaxis_title="Weekly nr of trades",
        legend=dict(yanchor="top", y=0.5),
        width=WIDTH,
        height=HEIGHT,
        barmode="group",
    )

    # Update x-axis format
    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)


def integrated_plot_winning_trades_per_market_by_week(
    trades_df: pd.DataFrame,
) -> gr.Plot:
    # adding the total
    trades_all = trades_df.copy(deep=True)
    trades_all["market_creator"] = "all"

    # 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
    )
    final_df = get_overall_winning_by_market_trades(all_filtered_trades)
    fig = px.bar(
        final_df,
        x="month_year_week",
        y="winning_trade",
        color="market_creator",
        barmode="group",
        color_discrete_sequence=["purple", "goldenrod", "darkgreen"],
        category_orders={"market_creator": ["pearl", "quickstart", "all"]},
    )
    fig.update_layout(
        xaxis_title="Week",
        yaxis_title="Weekly % of winning trades",
        legend=dict(yanchor="top", y=0.5),
    )
    fig.update_layout(width=WIDTH, height=HEIGHT)
    fig.update_xaxes(tickformat="%b %d\n%Y")
    return gr.Plot(
        value=fig,
    )


def integrated_plot_winning_trades_per_market_by_week_v2(
    trades_df: pd.DataFrame, trader_filter: str = "all"
) -> gr.Plot:
    # adding the total
    trades_all = trades_df.copy(deep=True)
    trades_all["market_creator"] = "all"

    # 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
    )
    # Create binary staking category
    all_filtered_trades["staking_type"] = all_filtered_trades["staking"].apply(
        lambda x: "non_Olas" if x == "non_Olas" else "Olas"
    )
    if trader_filter == "all":
        final_df = get_overall_winning_by_market_trades(all_filtered_trades)
    else:
        final_df = get_overall_winning_by_market_and_trader_type(all_filtered_trades)

    # Convert string dates to datetime and sort them
    all_dates_dt = sorted(
        [
            datetime.strptime(date, "%b-%d-%Y")
            for date in final_df["month_year_week"].unique()
        ]
    )
    # Convert back to string format
    all_dates = [date.strftime("%b-%d-%Y") for date in all_dates_dt]
    color_discrete_sequence = ["darkviolet", "goldenrod", "green"]
    if trader_filter == "Olas":
        final_df = final_df[final_df["staking_type"] == "Olas"]
    elif trader_filter == "non_Olas":
        final_df = final_df[final_df["staking_type"] == "non_Olas"]
        color_discrete_sequence = ["purple", "darkgoldenrod", "darkgreen"]

    fig = px.bar(
        final_df,
        x="month_year_week",
        y="winning_trade",
        color="market_creator",
        barmode="group",
        color_discrete_sequence=color_discrete_sequence,
        category_orders={"market_creator": ["pearl", "quickstart", "all"]},
    )
    fig.update_layout(
        xaxis_title="Week",
        yaxis_title="Weekly % of winning trades",
        legend=dict(yanchor="top", y=0.5),
    )

    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,
    )