Py.Cafe

############ Imports ##############
import vizro.plotly.express as px
import vizro.tables as vt
import vizro.models as vm
import vizro.figures as vf
from vizro.models.types import capture
from vizro import Vizro
import pandas as pd
from vizro.managers import data_manager


####### Function definitions ######
@capture("graph")
def tests_by_analyst_chart(data_frame: pd.DataFrame):
    color_map = {
        "Pass": "#689f38",
        "OOS": "#ff5267",
        "OOT": "#ff9222",
        "Pending": "#9e9e9e",
    }
    counts = (
        data_frame.groupby(["analyst", "result_status"])
        .size()
        .reset_index(name="count")
    )
    fig = px.bar(
        counts,
        y="analyst",
        x="count",
        color="result_status",
        color_discrete_map=color_map,
        orientation="h",
        category_orders={"result_status": ["Pass", "OOS", "OOT", "Pending"]},
    )
    fig.update_layout(
        xaxis_title="Number of Tests",
        yaxis_title="Analyst",
        legend_title="Result Status",
        barmode="stack",
    )
    return fig


@capture("graph")
def sla_status_bar_chart(data_frame: pd.DataFrame):
    color_map = {
        "Breached": "#ff5267",
        "At Risk": "#ff9222",
        "On Track": "#2196f3",
        "Met": "#689f38",
    }
    counts = data_frame.groupby("sla_status").size().reset_index(name="count")
    fig = px.bar(
        counts,
        x="sla_status",
        y="count",
        color="sla_status",
        color_discrete_map=color_map,
        category_orders={"sla_status": ["Breached", "At Risk", "On Track", "Met"]},
    )
    fig.update_layout(
        xaxis_title="SLA Status", yaxis_title="Number of Tests", showlegend=False
    )
    return fig


@capture("graph")
def result_status_pie_chart(data_frame: pd.DataFrame):
    color_map = {
        "Pass": "#689f38",
        "OOS": "#ff5267",
        "OOT": "#ff9222",
        "Pending": "#9e9e9e",
    }
    counts = data_frame.groupby("result_status").size().reset_index(name="count")
    fig = px.pie(
        counts,
        values="count",
        names="result_status",
        color="result_status",
        color_discrete_map=color_map,
        category_orders={"result_status": ["Pass", "OOS", "OOT", "Pending"]},
    )
    fig.update_traces(textposition="inside", textinfo="percent+label")
    return fig


@capture("graph")
def batch_risk_scatter_chart(data_frame: pd.DataFrame):
    color_map = {"High": "#ff5267", "Medium": "#ff9222", "Low": "#689f38"}
    fig = px.scatter(
        data_frame,
        x="days_to_target",
        y="risk_score",
        color="risk_level",
        color_discrete_map=color_map,
        size="deviation_count",
        size_max=20,
        hover_data=["batch_id", "product", "site", "status", "blocker_type"],
        category_orders={"risk_level": ["High", "Medium", "Low"]},
    )
    fig.update_layout(
        xaxis_title="Days to Target Release",
        yaxis_title="Risk Score",
        legend_title="Risk Level",
    )
    return fig


@capture("graph")
def test_result_heatmap_chart(data_frame: pd.DataFrame):
    pivot = (
        data_frame.groupby(["test_type", "result_status"])
        .size()
        .reset_index(name="count")
    )
    pivot_wide = pivot.pivot(
        index="test_type", columns="result_status", values="count"
    ).fillna(0)
    result_order = ["Pass", "OOS", "OOT", "Pending"]
    existing_cols = [c for c in result_order if c in pivot_wide.columns]
    pivot_wide = pivot_wide[existing_cols]
    fig = px.imshow(
        pivot_wide, text_auto=True, color_continuous_scale="Blues", aspect="auto"
    )
    fig.update_layout(xaxis_title="Result Status", yaxis_title="Test Type")
    return fig


@capture("graph")
def status_by_site_chart(data_frame: pd.DataFrame):
    color_map = {
        "Blocked": "#ff5267",
        "On Hold": "#ff9222",
        "In Progress": "#2196f3",
        "QC Complete": "#4caf50",
        "Ready for Release": "#689f38",
    }
    counts = data_frame.groupby(["site", "status"]).size().reset_index(name="count")
    fig = px.bar(
        counts,
        y="site",
        x="count",
        color="status",
        color_discrete_map=color_map,
        orientation="h",
        category_orders={
            "status": [
                "Blocked",
                "On Hold",
                "In Progress",
                "QC Complete",
                "Ready for Release",
            ]
        },
    )
    fig.update_layout(
        xaxis_title="Number of Batches",
        yaxis_title="Site",
        legend_title="Status",
        barmode="stack",
    )
    return fig


@capture("graph")
def tests_by_instrument_chart(data_frame: pd.DataFrame):
    color_map = {
        "Breached": "#ff5267",
        "At Risk": "#ff9222",
        "On Track": "#2196f3",
        "Met": "#689f38",
    }
    counts = (
        data_frame.groupby(["instrument", "sla_status"])
        .size()
        .reset_index(name="count")
    )
    fig = px.bar(
        counts,
        y="instrument",
        x="count",
        color="sla_status",
        color_discrete_map=color_map,
        orientation="h",
        category_orders={"sla_status": ["Breached", "At Risk", "On Track", "Met"]},
    )
    fig.update_layout(
        xaxis_title="Number of Tests",
        yaxis_title="Instrument",
        legend_title="SLA Status",
        barmode="stack",
    )
    return fig


@capture("graph")
def queue_by_shift_chart(data_frame: pd.DataFrame):
    color_map = {
        "Breached": "#ff5267",
        "At Risk": "#ff9222",
        "On Track": "#2196f3",
        "Met": "#689f38",
    }
    counts = (
        data_frame.groupby(["shift", "sla_status"]).size().reset_index(name="count")
    )
    fig = px.bar(
        counts,
        x="shift",
        y="count",
        color="sla_status",
        color_discrete_map=color_map,
        barmode="group",
        category_orders={
            "shift": ["Day", "Night", "Weekend"],
            "sla_status": ["Breached", "At Risk", "On Track", "Met"],
        },
    )
    fig.update_layout(
        xaxis_title="Shift", yaxis_title="Number of Tests", legend_title="SLA Status"
    )
    return fig


####### Data Loading and Processing #####
# Load raw data
batch_data_raw = pd.read_csv("batch_data.csv")
qc_queue_data_raw = pd.read_csv("qc_queue_data.csv")

# Create filtered datasets for KPIs
batch_data_high_risk = batch_data_raw[batch_data_raw["risk_level"] == "High"].copy()
batch_data_blocked = batch_data_raw[batch_data_raw["status"] == "Blocked"].copy()

# Calculate on-time percentage
on_time_count = (batch_data_raw["days_to_target"] >= 0).sum()
total_count = len(batch_data_raw)
batch_data_on_time_pct = pd.DataFrame({"on_time_pct": [100 * on_time_count / total_count]})

# QC Queue filtered datasets
qc_queue_pending = qc_queue_data_raw[qc_queue_data_raw["result_status"] == "Pending"].copy()
qc_queue_breached = qc_queue_data_raw[qc_queue_data_raw["sla_status"] == "Breached"].copy()
qc_queue_completed = qc_queue_data_raw[qc_queue_data_raw["tat_hours"].notna()].copy()
qc_queue_oos_oot = qc_queue_data_raw[qc_queue_data_raw["result_status"].isin(["OOS", "OOT"])].copy()
qc_queue_retest = qc_queue_data_raw[qc_queue_data_raw["retest_flag"] == True].copy()

####### Data Manager Settings #####
data_manager["batch_data"] = batch_data_raw
data_manager["batch_data_high_risk"] = batch_data_high_risk
data_manager["batch_data_blocked"] = batch_data_blocked
data_manager["batch_data_on_time_pct"] = batch_data_on_time_pct
data_manager["qc_queue_data"] = qc_queue_data_raw
data_manager["qc_queue_pending"] = qc_queue_pending
data_manager["qc_queue_breached"] = qc_queue_breached
data_manager["qc_queue_completed"] = qc_queue_completed
data_manager["qc_queue_oos_oot"] = qc_queue_oos_oot
data_manager["qc_queue_retest"] = qc_queue_retest

########### Model code ############
model = vm.Dashboard(
    pages=[
        vm.Page(
            components=[
                vm.Figure(
                    id="kpi_total_batches",
                    figure=vf.kpi_card(
                        data_frame="batch_data",
                        value_column="batch_id",
                        value_format="{value}",
                        agg_func="count",
                        title="Total Batches",
                        icon="inventory_2",
                    ),
                ),
                vm.Figure(
                    id="kpi_high_risk",
                    figure=vf.kpi_card(
                        data_frame="batch_data_high_risk",
                        value_column="batch_id",
                        value_format="{value}",
                        agg_func="count",
                        title="High Risk",
                        icon="warning",
                    ),
                ),
                vm.Figure(
                    id="kpi_blocked",
                    figure=vf.kpi_card(
                        data_frame="batch_data_blocked",
                        value_column="batch_id",
                        value_format="{value}",
                        agg_func="count",
                        title="Blocked",
                        icon="block",
                    ),
                ),
                vm.Figure(
                    id="kpi_days_to_target",
                    figure=vf.kpi_card(
                        data_frame="batch_data",
                        value_column="days_to_target",
                        value_format="{value:.1f} days",
                        agg_func="mean",
                        title="Avg Days to Target",
                        icon="schedule",
                    ),
                ),
                vm.Figure(
                    id="kpi_on_time",
                    figure=vf.kpi_card(
                        data_frame="batch_data_on_time_pct",
                        value_column="on_time_pct",
                        value_format="{value:.1f}%",
                        agg_func="mean",
                        title="On-Time Rate",
                        icon="check_circle",
                    ),
                ),
                vm.Graph(
                    id="batch_risk_scatter",
                    figure=batch_risk_scatter_chart(data_frame="batch_data"),
                    title="Batch Risk Overview",
                ),
                vm.Graph(
                    id="status_by_site",
                    figure=status_by_site_chart(data_frame="batch_data"),
                    title="Status by Site",
                ),
                vm.AgGrid(
                    id="batch_table",
                    figure=vt.dash_ag_grid(data_frame="batch_data"),
                    title="Batch Details",
                ),
            ],
            title="At-Risk Batches",
            layout=vm.Grid(
                grid=[
                    [0, 0, 0, 1, 1, 1, 2, 2, 3, 3, 4, 4],
                    [5, 5, 5, 5, 5, 5, 6, 6, 6, 6, 6, 6],
                    [5, 5, 5, 5, 5, 5, 6, 6, 6, 6, 6, 6],
                    [5, 5, 5, 5, 5, 5, 6, 6, 6, 6, 6, 6],
                    [7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7],
                    [7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7],
                    [7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7],
                ],
                row_min_height="140px",
            ),
            controls=[
                vm.Filter(
                    column="site",
                    selector=vm.Dropdown(multi=True, title="Site"),
                ),
                vm.Filter(
                    column="product",
                    selector=vm.Dropdown(multi=True, title="Product"),
                ),
                vm.Filter(
                    column="risk_level",
                    selector=vm.Dropdown(multi=True, title="Risk Level"),
                ),
                vm.Filter(
                    column="status",
                    selector=vm.Dropdown(multi=True, title="Status"),
                ),
            ],
        ),
        vm.Page(
            components=[
                vm.Figure(
                    id="kpi_queue_total",
                    figure=vf.kpi_card(
                        data_frame="qc_queue_pending",
                        value_column="test_id",
                        value_format="{value}",
                        agg_func="count",
                        title="Tests in Queue",
                        icon="science",
                    ),
                ),
                vm.Figure(
                    id="kpi_sla_breaches",
                    figure=vf.kpi_card(
                        data_frame="qc_queue_breached",
                        value_column="test_id",
                        value_format="{value}",
                        agg_func="count",
                        title="SLA Breaches",
                        icon="error",
                    ),
                ),
                vm.Figure(
                    id="kpi_avg_tat",
                    figure=vf.kpi_card(
                        data_frame="qc_queue_completed",
                        value_column="tat_hours",
                        value_format="{value:.1f}h",
                        agg_func="mean",
                        title="Avg TAT",
                        icon="timer",
                    ),
                ),
                vm.Figure(
                    id="kpi_oos_oot",
                    figure=vf.kpi_card(
                        data_frame="qc_queue_oos_oot",
                        value_column="test_id",
                        value_format="{value}",
                        agg_func="count",
                        title="OOS/OOT",
                        icon="report_problem",
                    ),
                ),
                vm.Figure(
                    id="kpi_retest",
                    figure=vf.kpi_card(
                        data_frame="qc_queue_retest",
                        value_column="test_id",
                        value_format="{value}",
                        agg_func="count",
                        title="Retests",
                        icon="replay",
                    ),
                ),
                vm.Graph(
                    id="tat_by_test_type",
                    figure=px.box(
                        data_frame="qc_queue_completed",
                        x="test_type",
                        y="tat_hours",
                        points="outliers",
                    ),
                    title="TAT by Test Type",
                ),
                vm.Graph(
                    id="queue_by_sla",
                    figure=sla_status_bar_chart(data_frame="qc_queue_data"),
                    title="Queue by SLA Status",
                ),
                vm.Graph(
                    id="tests_by_instrument",
                    figure=tests_by_instrument_chart(data_frame="qc_queue_data"),
                    title="Tests by Instrument",
                ),
                vm.Graph(
                    id="result_status_pie",
                    figure=result_status_pie_chart(data_frame="qc_queue_data"),
                    title="Result Status",
                ),
                vm.Graph(
                    id="tests_by_analyst",
                    figure=tests_by_analyst_chart(data_frame="qc_queue_data"),
                    title="Tests by Analyst",
                ),
                vm.Graph(
                    id="queue_by_shift",
                    figure=queue_by_shift_chart(data_frame="qc_queue_data"),
                    title="Queue by Shift",
                ),
                vm.Graph(
                    id="test_result_heatmap",
                    figure=test_result_heatmap_chart(data_frame="qc_queue_data"),
                    title="Test Type vs Result",
                ),
                vm.AgGrid(
                    id="qc_table",
                    figure=vt.dash_ag_grid(data_frame="qc_queue_data"),
                    title="Test Queue Details",
                ),
            ],
            title="QC Lab Queue",
            layout=vm.Grid(
                grid=[
                    [0, 0, 0, 1, 1, 1, 2, 2, 3, 3, 4, 4],
                    [5, 5, 5, 5, 5, 5, 6, 6, 6, 6, 6, 6],
                    [5, 5, 5, 5, 5, 5, 6, 6, 6, 6, 6, 6],
                    [5, 5, 5, 5, 5, 5, 6, 6, 6, 6, 6, 6],
                    [7, 7, 7, 7, 8, 8, 8, 8, 9, 9, 9, 9],
                    [7, 7, 7, 7, 8, 8, 8, 8, 9, 9, 9, 9],
                    [7, 7, 7, 7, 8, 8, 8, 8, 9, 9, 9, 9],
                    [10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11],
                    [10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11],
                    [10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11],
                    [12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12],
                    [12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12],
                    [12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12],
                ],
                row_min_height="140px",
            ),
            controls=[
                vm.Filter(
                    column="site",
                    selector=vm.Dropdown(multi=True, title="Site"),
                ),
                vm.Filter(
                    column="test_type",
                    selector=vm.Dropdown(multi=True, title="Test Type"),
                ),
                vm.Filter(
                    column="sla_status",
                    selector=vm.Dropdown(multi=True, title="SLA Status"),
                ),
                vm.Filter(
                    column="result_status",
                    selector=vm.Dropdown(multi=True, title="Result Status"),
                ),
                vm.Filter(
                    column="instrument",
                    selector=vm.Dropdown(multi=True, title="Instrument"),
                ),
                vm.Filter(
                    column="analyst",
                    selector=vm.Dropdown(multi=True, title="Analyst"),
                ),
            ],
        ),
    ],
    theme="vizro_dark",
    title="Batch Release & Quality Ops Command Center",
)

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