"""three sourdough orders before lunch

Generated from src/content/docs/reference/joint-counts.md.
Run with the documented Fieldwork release in a writable directory.
Exports are saved there using the filenames shown in the walkthrough.
"""

import pandas as pd
import fieldwork as fw

batches = [
    ("Saturday", "Quay", "early", "sourdough", 3),
    ("Saturday", "Quay", "early", "rye", 1),
    ("Saturday", "Quay", "late", "sourdough", 2),
    ("Saturday", "Quay", "late", "focaccia", 2),
    ("Saturday", "Hill", "early", "rye", 2),
    ("Saturday", "Hill", "noon", "focaccia", 2),
    ("Sunday", "Quay", "early", "sourdough", 1),
    ("Sunday", "Quay", "noon", "rye", 1),
    ("Sunday", "Quay", "late", "focaccia", 2),
]
rows = []
for day, stall, pickup, bread, count in batches:
    for _ in range(count):
        rows.append((day, stall, pickup, bread))

orders = pd.DataFrame(
    rows,
    columns=["day", "stall", "pickup", "bread"],
)
print(f"Orders: {len(orders)}")

quay = fw.joint_counts(
    orders,
    ["pickup", "bread"],
    context={"day": "Saturday", "stall": "Quay"},
)
data = quay.to_dict()
pickups, breads = data["a"], data["b"]
observed = {
    (pickups[cell["a"]], breads[cell["b"]]): cell["count"]
    for cell in data["cells"]
}
print("Early sourdough orders:", observed[("early", "sourdough")])
print(
    "Input / outside context / evaluated:",
    data["scope"]["input_rows"],
    data["restriction_excluded_rows"],
    data["evaluated_rows"],
)

# Fill unobserved combinations only across the returned supported axes.
table = pd.DataFrame(0, index=pickups, columns=breads)
for (pickup, bread), count in observed.items():
    table.loc[pickup, bread] = count
print(table.to_string())
print("Stored cells / grid cells:", len(data["cells"]), table.size)
svg = fw.render_svg(quay, view="heatmap")

all_orders = fw.joint_counts(orders, ["pickup", "bread"])
print("Global pickup domain:", all_orders["a"])
print("Saturday Quay pickup domain:", pickups)
print(
    "Early sourdough share: {:.1%}".format(
        observed[("early", "sourdough")] / data["evaluated_rows"]
    )
)

try:
    fw.joint_counts(
        orders,
        ["pickup", "bread"],
        context={"day": "Saturday", "stall": "Quay"},
        max_cells=4,
    )
except ValueError as error:
    print(type(error).__name__ + ": " + str(error))

sunday_hill = fw.joint_counts(
    orders,
    ["pickup", "bread"],
    context={"day": "Sunday", "stall": "Hill"},
)
print(
    "Sunday Hill:",
    sunday_hill["status"],
    sunday_hill["a"],
    sunday_hill["b"],
    sunday_hill["cells"],
)

from pathlib import Path

Path("joint_counts.svg").write_text(svg, encoding="utf-8")