#!/usr/bin/env python3
# usage:        uv run dir_treemap.py <dir>
# requirements: plotly
import sys
from pathlib import Path
from collections import defaultdict
import pandas as pd
import plotly.express as px

# visualization depth
depth = 3

# use either . or first arg as start directory
root = Path.cwd()
if len(sys.argv) > 1:
    root = Path(sys.argv[1]).resolve()

# display size human readable
def human(n):
    for u in ("B", "KB", "MB", "GB"):
        if n < 1024: return f"{n:.1f} {u}"
        n /= 1024
    return f"{n:.1f} TB"

# accumulate file sizes 
sizes = defaultdict(int)
for p in root.rglob("*"):  # get (recursive) list of file pathes
    if not p.is_file():
        continue
    try:
        # deep files are rolled up into closest visible ancestor
        # ("a", "b", "c", "d", "file.txt")
        parts = p.relative_to(root).parts
        # ("a", "b", "c")
        target = parts[:depth]             
        sizes[target] += p.stat().st_size
    except (PermissionError, OSError):
        pass

# build node list from the accumulated paths
# hard coded start node
nodes = [{"id": "__root__",
          "parent": "",
          "label": root.name,
          "size": 0}]
seen = set()
for parts, size in sizes.items():
    for level in range(1, len(parts) + 1):
        seg = parts[:level]
        if seg in seen:
            continue
        seen.add(seg)
        parent = "/".join(seg[:-1]) if len(seg) > 1 else "__root__"
        nodes.append({
            "id":     "/".join(seg),
            "parent": parent,
            "label":  seg[-1],
            "size":   size if level == len(parts) else 0,
        })
df = pd.DataFrame(nodes)

# visualize as tree map
fig = px.treemap(df,
                 ids="id",
                 names="label",
                 parents="parent",
                 values="size",
                 title=f"Directory contents: {root}",
                 color_discrete_sequence=px.colors.qualitative.Pastel)
fig.update_traces(
    texttemplate="%{label}<br>%{customdata[0]}",
    customdata=[[human(r)] for r in df["size"]],
    hovertemplate="<b>%{label}</b><br>%{customdata[0]}<extra></extra>",
    marker=dict(line=dict(color="black", width=1)))
fig.update_layout(margin=dict(t=40, l=5, r=5, b=5))
fig.show()
