"""follow the loans through a library

Generated from src/content/docs/reference/census.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 = [
    ("Central", "Adult", "Print", 8),
    ("Central", "Adult", "Ebook", 4),
    ("Central", "Children", "Print", 4),
    ("Riverside", "Adult", "Print", 2),
    ("Riverside", "Children", "Print", 6),
]
rows = []
for branch, collection, format_, n in batches:
    for _ in range(n):
        rows.append((branch, collection, format_))

loans = pd.DataFrame(
    rows,
    columns=["branch", "collection", "format"],
)
print(f"Loans: {len(loans)}")

dimensions = ["branch", "collection", "format"]
full = fw.census(loans, dimensions)
print(fw.render_plaintext(full))

figure_data = fw.visualization_data(full)
groups = pd.DataFrame(figure_data["rows"])
central_collections = groups[
    groups["parent_label"] == "branch=Central"
]
print(
    central_collections[
        ["label", "count", "parent_share", "share"]
    ].to_string(index=False)
)

post = fw.census(
    loans, dimensions, top_n=1, top_n_mode="post", top_n_per_parent=True
)
print("Keep the original population:")
print(fw.render_plaintext(post))

pre = fw.census(
    loans, dimensions, top_n=1, top_n_mode="pre", top_n_per_parent=True
)
print("Count only the selected rows:")
print(fw.render_plaintext(pre))

collection_first = fw.census(
    loans, ["collection", "branch"], max_depth=2
)
print(fw.render_plaintext(collection_first))

from pathlib import Path

Path("census.svg").write_text(fw.render_svg(full), encoding="utf-8")