Py.Cafe

############ Imports ##############
import vizro.models as vm
from vizro.models.types import capture
from vizro import Vizro
import pandas as pd
from vizro.managers import data_manager
import plotly.graph_objects as go
import numpy as np
from vizro.models.types import capture


####### Function definitions ######
@capture("graph")
def horizontal_category_subcategory_sales(data_frame):
    # Define low volume stationery subcategories to combine
    low_volume_items = [
        "Envelopes",
        "Pens & Art Supplies",
        "Scissors, Rulers and Trimmers",
        "Labels",
        "Rubber Bands",
    ]

    # Create a copy of the dataframe for processing
    df_processed = data_frame.copy()

    # Replace the low volume subcategories with 'Low volume stationery'
    df_processed.loc[
        df_processed["Product Sub-Category"].isin(low_volume_items),
        "Product Sub-Category",
    ] = "Low volume stationery"

    # Group by Product Category and Product Sub-Category to get total sales and profit
    grouped = (
        df_processed.groupby(["Product Category", "Product Sub-Category"])
        .agg({"Sales": "sum", "Profit": "sum"})
        .reset_index()
    )

    # Calculate Profit Margin as sum(profit)/sum(sales) for each row
    grouped["Profit Margin"] = grouped["Profit"] / grouped["Sales"]

    # Sort by Product Category, then by Sales DESCENDING within each category (highest first)
    grouped = grouped.sort_values(
        ["Product Category", "Sales"], ascending=[True, False]
    )

    # Get positive margins only for blue gradient calculation
    positive_margins = grouped[grouped["Profit Margin"] > 0]["Profit Margin"]

    if len(positive_margins) > 0:
        min_positive_margin = positive_margins.min()
        max_positive_margin = positive_margins.max()
        positive_range = max_positive_margin - min_positive_margin
    else:
        min_positive_margin = 0
        max_positive_margin = 1
        positive_range = 1

    # Create color scale: orange for negative margins, pale blue to dark blue for positive margins
    def get_color(margin_value):
        if margin_value <= 0:
            return "#ff7f0e"  # Orange for negative margins (losses)
        else:
            # Blue gradient for positive margins
            if positive_range == 0:
                return (
                    "#1f4e79"  # Default dark blue if no variation in positive margins
                )

            # Normalize positive margin to 0-1 scale
            normalized = (margin_value - min_positive_margin) / positive_range

            # Interpolate between pale blue and dark blue for positive margins
            # Pale Blue: #add8e6, Dark Blue: #1f4e79
            pale_r, pale_g, pale_b = 173, 216, 230
            dark_r, dark_g, dark_b = 31, 78, 121

            r = int(pale_r + (dark_r - pale_r) * normalized)
            g = int(pale_g + (dark_g - pale_g) * normalized)
            b = int(pale_b + (dark_b - pale_b) * normalized)

            return f"rgb({r},{g},{b})"

    # Create colors for each bar based on profit margin
    colors = [get_color(margin) for margin in grouped["Profit Margin"]]

    # Create y-axis labels that maintain the sorted order (reverse for plotly display)
    y_labels = grouped["Product Sub-Category"].tolist()
    y_labels.reverse()  # Reverse so highest sales appear at top

    fig = go.Figure()

    # Reverse the data order for display (highest sales at top)
    grouped_reversed = grouped.iloc[::-1].reset_index(drop=True)
    colors_reversed = colors[::-1]

    fig.add_trace(
        go.Bar(
            y=grouped_reversed["Product Sub-Category"],
            x=grouped_reversed["Sales"],
            orientation="h",
            marker=dict(
                color=colors_reversed,
                line=dict(width=0.5, color="rgba(255,255,255,0.4)"),
            ),
            hovertemplate="<b>%{y}</b><br>Sales: $%{x:,.0f}<br>Profit: $%{customdata[0]:,.0f}<br>Profit Margin: %{customdata[1]:.1%}<br><extra></extra>",
            customdata=list(
                zip(grouped_reversed["Profit"], grouped_reversed["Profit Margin"])
            ),
            name="Sales by Sub-Category",
        )
    )

    # Add dividing lines between product categories (adjust for reversed order)
    categories = grouped["Product Category"].unique()
    total_items = len(grouped)
    y_position = 0

    for i, category in enumerate(categories):
        category_count = len(grouped[grouped["Product Category"] == category])

        # Calculate position from bottom for reversed display
        if i > 0:
            line_position = total_items - y_position - 0.5
            fig.add_hline(
                y=line_position,
                line=dict(color="rgba(255,255,255,0.6)", width=2, dash="solid"),
            )

        y_position += category_count

    # Create gradient bar for legend (only for positive margins)
    gradient_x = np.linspace(0, 1, 100)
    gradient_colors = []

    for x in gradient_x:
        if len(positive_margins) > 0:
            margin_val = (
                min_positive_margin + (max_positive_margin - min_positive_margin) * x
            )
            gradient_colors.append(get_color(margin_val))
        else:
            gradient_colors.append("#1f4e79")

    # Add gradient bar as a separate subplot area
    fig.add_trace(
        go.Bar(
            x=gradient_x,
            y=["Profit Margin"] * 100,
            orientation="h",
            marker=dict(color=gradient_colors, line=dict(width=0)),
            showlegend=False,
            hoverinfo="skip",
            yaxis="y2",
            xaxis="x2",
        )
    )

    # Calculate dynamic height based on number of items (minimum 400px, ~30px per item)
    chart_height = max(400, len(grouped) * 30 + 100)

    fig.update_layout(
        xaxis_title="Sales ($)",
        yaxis_title="Product Sub-Category",
        height=chart_height,
        hovermode="closest",
        paper_bgcolor="rgba(0,0,0,0)",
        plot_bgcolor="rgba(0,0,0,0)",
        showlegend=False,
        yaxis=dict(categoryorder="array", categoryarray=y_labels, domain=[0, 0.85]),
        # Second y-axis for gradient legend
        yaxis2=dict(
            domain=[0.9, 0.95],
            anchor="x2",
            showticklabels=False,
            showgrid=False,
            zeroline=False,
        ),
        # Second x-axis for gradient legend
        xaxis2=dict(
            domain=[0.7, 0.98],
            anchor="y2",
            showticklabels=False,
            showgrid=False,
            zeroline=False,
        ),
        annotations=[
            # Gradient legend labels
            dict(
                x=0.68,
                y=0.925,
                xref="paper",
                yref="paper",
                text="Low Margin",
                showarrow=False,
                font=dict(size=10, color="white"),
                xanchor="right",
            ),
            dict(
                x=1.0,
                y=0.925,
                xref="paper",
                yref="paper",
                text="High Margin",
                showarrow=False,
                font=dict(size=10, color="white"),
                xanchor="left",
            ),
            dict(
                x=0.84,
                y=0.97,
                xref="paper",
                yref="paper",
                text="<b>Profit Margin</b>",
                showarrow=False,
                font=dict(size=11, color="white"),
                xanchor="center",
            ),
            # Add legend for orange (negative margins)
            dict(
                x=0.68,
                y=0.875,
                xref="paper",
                yref="paper",
                text="🟠 = Loss",
                showarrow=False,
                font=dict(size=10, color="white"),
                xanchor="right",
            ),
        ],
    )

    return fig


####### Data Manager Settings #####
data_manager["megastore_data"] = pd.read_csv(
    "https://raw.githubusercontent.com/stichbury/vizro_projects/main/Megastore/MegastoreData.csv"
)

########### Model code ############
model = vm.Dashboard(
    pages=[
        vm.Page(
            components=[
                vm.Graph(
                    id="sales_chart",
                    type="graph",
                    figure=horizontal_category_subcategory_sales(
                        data_frame="megastore_data"
                    ),
                    title="Sales by Product Sub-Category (Color = Profit Margin)",
                )
            ],
            title="Megastore Sales Analysis",
            controls=[
                vm.Filter(
                    type="filter",
                    column="Continent",
                    targets=["sales_chart"],
                    selector=vm.Dropdown(
                        type="dropdown",
                        value=["Asia", "Europe", "North America", "South America"],
                        multi=True,
                    ),
                )
            ],
        )
    ],
    theme="vizro_dark",
    title="Megastore Dashboard",
)

Vizro().build(model).run()