import pandas as pd
import gradio as gr
import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import matplotlib.pyplot as plt
import seaborn as sns


def plot_kl_div_per_market(closed_markets: pd.DataFrame) -> gr.Plot:

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

    # merging both dataframes
    final_markets = pd.concat([closed_markets, all_markets], ignore_index=True)
    final_markets = final_markets.sort_values(by="opening_datetime", ascending=True)

    fig = px.box(
        final_markets,
        x="month_year_week",
        y="kl_divergence",
        color="market_creator",
        color_discrete_sequence=["purple", "goldenrod", "darkgreen"],
        category_orders={"market_creator": ["pearl", "quickstart", "all"]},
    )
    fig.update_traces(boxmean=True)
    fig.update_layout(
        xaxis_title="Markets closing Week",
        yaxis_title="Kullback–Leibler divergence",
        legend=dict(yanchor="top", y=0.5),
        width=800,  # Adjusted for better fit on laptop screens
        height=600,  # Adjusted for better fit on laptop screens
    )

    fig.update_xaxes(tickformat="%b %d\n%Y")

    return gr.Plot(
        value=fig,
    )


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

    # merging both dataframes
    final_markets = pd.concat([closed_markets, all_markets], ignore_index=True)
    final_markets = final_markets.sort_values(by="opening_datetime", ascending=True)

    # Create the main figure and axis
    fig, ax1 = plt.subplots(figsize=(10, 6))

    # Create the boxplot using seaborn
    sns.boxplot(
        data=final_markets,
        x="month_year_week",
        y="kl_divergence",
        ax=ax1,
        hue="market_creator",
    )

    # Set labels and title for the main axis
    ax1.set_xlabel("Week")
    ax1.set_ylabel("KL Divergence")
    ax1.set_title("KL Divergence Boxplot with Off-by Percentage")

    # Create a secondary y-axis
    ax2 = ax1.twinx()

    # Plot the off_by_perc values on the secondary y-axis
    for i, week in enumerate(closed_markets["month_year_week"].unique()):
        off_by_perc = closed_markets[closed_markets["month_year_week"] == week][
            "off_by_perc"
        ]
        ax2.scatter([i] * len(off_by_perc), off_by_perc, color="red", alpha=0.01)

    # Set label for the secondary y-axis
    ax2.set_ylabel("Off-by Percentage")

    # Adjust the layout and display the plot
    plt.tight_layout()
    return gr.Plot(
        value=fig,
    )


def plot_total_bet_amount(trades_df: pd.DataFrame) -> gr.Plot:
    """Plots the trade metrics."""
    # Create binary staking category
    trades_df["trader_type"] = trades_df["staking"].apply(
        lambda x: "non_agent" if x == "non_agent" else "agent"
    )

    total_bet_amount = (
        trades_df.groupby(
            ["month_year_week", "market_creator", "trader_type"], sort=False
        )["collateral_amount"]
        .sum()
        .reset_index(name="total_bet_amount")
    )
    color_mapping = [
        "darkviolet",
        "purple",
        "goldenrod",
        "darkgoldenrod",
        "green",
        "darkgreen",
    ]
    total_bet_amount["trader_market"] = total_bet_amount.apply(
        lambda x: (x["trader_type"], x["market_creator"]), axis=1
    )

    fig = px.bar(
        total_bet_amount,
        x="month_year_week",
        y="total_bet_amount",
        color="trader_market",
        color_discrete_sequence=color_mapping,
        category_orders={
            "market_creator": ["pearl", "quickstart", "all"],
            "trader_market": [
                ("agent", "pearl"),
                ("non_agent", "pearl"),
                ("agent", "quickstart"),
                ("non_agent", "quickstart"),
                ("agent", "all"),
                ("non_agent", "all"),
            ],
        },
        barmode="group",
        facet_col="market_creator",
    )

    fig.update_layout(
        xaxis_title="Week",
        yaxis_title="Weekly total bet amount",
        legend=dict(yanchor="top", y=0.5),
    )
    for axis in fig.layout:
        if axis.startswith("xaxis"):
            fig.layout[axis].update(title="Week")
    fig.update_xaxes(tickformat="%b %d")
    return gr.Plot(
        value=fig,
    )