"""Clean and summarize loans.csv using only the Python standard library.

Run it:
    python wrangle.py
    python wrangle.py --top 3

What it does, step by step:
    1. Load the raw CSV.
    2. Clean it: trim spaces, fix names, parse numbers, drop duplicates and unusable rows.
    3. Summarize it: totals, a rate, and an average, grouped by branch or category.
    4. Save loans_clean.csv next to this script. The raw file is never changed.

Dedicated to the public domain under CC0 1.0.
"""
from __future__ import annotations

import argparse
import csv
from collections import defaultdict
from pathlib import Path

HERE = Path(__file__).parent
RAW = HERE / "loans.csv"
CLEAN = HERE / "loans_clean.csv"

# Known spellings. Anything else is reported instead of silently kept.
BRANCHES = {"central": "Central", "harbor": "Harbor", "hillside": "Hillside", "riverside": "Riverside"}
CATEGORIES = {
    "fiction": "Fiction",
    "non-fiction": "Non-fiction",
    "non fiction": "Non-fiction",
    "children": "Children",
    "audio": "Audio",
    "reference": "Reference",
}
FIELDS = ["date", "month", "branch", "category", "loans", "late_returns", "avg_wait_days", "notes"]


def parse_number(text: str, as_int: bool = True):
    """Turn text such as '1,204' or '3.5' into a number. Empty text becomes None."""
    text = text.strip().replace(",", "")
    if text == "":
        return None
    return int(float(text)) if as_int else float(text)


def clean_rows(raw_rows: list[dict]) -> tuple[list[dict], dict]:
    """Return (clean_rows, report). The report counts what was changed."""
    report = {"rows_read": len(raw_rows), "names_fixed": 0, "thousands_separators": 0,
              "duplicates_removed": 0, "dropped_missing_loans": 0, "unknown_names": 0}
    seen = set()
    clean = []
    for raw in raw_rows:
        branch = BRANCHES.get(raw["branch"].strip().lower())
        category = CATEGORIES.get(raw["category"].strip().lower())
        if branch is None or category is None:
            report["unknown_names"] += 1
            continue
        report["names_fixed"] += (raw["branch"] != branch) + (raw["category"] != category)
        if "," in raw["loans"]:
            report["thousands_separators"] += 1
        row = {
            "date": raw["date"].strip(),
            "month": raw["date"].strip()[:7],
            "branch": branch,
            "category": category,
            "loans": parse_number(raw["loans"]),
            "late_returns": parse_number(raw["late_returns"]),
            "avg_wait_days": parse_number(raw["avg_wait_days"], as_int=False),
            "notes": raw["notes"].strip(),
        }
        key = tuple(row.values())
        if key in seen:
            report["duplicates_removed"] += 1
            continue
        seen.add(key)
        if row["loans"] is None:
            report["dropped_missing_loans"] += 1
            continue
        clean.append(row)
    report["rows_clean"] = len(clean)
    return clean, report


def load_and_clean(path: Path = RAW) -> tuple[list[dict], dict]:
    with path.open(newline="", encoding="utf-8") as f:
        return clean_rows(list(csv.DictReader(f)))


def sum_by(rows: list[dict], group_by: str, metric: str) -> dict:
    """Total of a numeric column for each group. Missing values are skipped."""
    totals: dict = defaultdict(int)
    for row in rows:
        if row[metric] is not None:
            totals[row[group_by]] += row[metric]
    return dict(totals)


def late_rate_by(rows: list[dict], group_by: str) -> dict:
    """Late returns divided by loans, using only rows where late_returns is known."""
    late: dict = defaultdict(int)
    loans: dict = defaultdict(int)
    for row in rows:
        if row["late_returns"] is not None:
            late[row[group_by]] += row["late_returns"]
            loans[row[group_by]] += row["loans"]
    return {key: late[key] / loans[key] for key in loans}


def mean_by(rows: list[dict], group_by: str, metric: str) -> dict:
    """Average of a numeric column for each group. Missing values are skipped."""
    values: dict = defaultdict(list)
    for row in rows:
        if row[metric] is not None:
            values[row[group_by]].append(row[metric])
    return {key: sum(v) / len(v) for key, v in values.items()}


def save_clean(rows: list[dict], path: Path = CLEAN) -> None:
    with path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=FIELDS, lineterminator="\n")
        writer.writeheader()
        for row in rows:
            writer.writerow({k: ("" if v is None else v) for k, v in row.items()})


def main() -> None:
    parser = argparse.ArgumentParser(description="Clean and summarize loans.csv")
    parser.add_argument("--top", type=int, default=5, help="how many rows of each table to show")
    args = parser.parse_args()

    rows, report = load_and_clean()
    save_clean(rows)

    print("CLEANING REPORT")
    for key, value in report.items():
        print(f"  {key.replace('_', ' '):<24}{value}")

    print("\nTOTAL LOANS BY BRANCH")
    for branch, total in sorted(sum_by(rows, "branch", "loans").items(), key=lambda kv: -kv[1])[: args.top]:
        print(f"  {branch:<10}{total:>8,}")

    print("\nLATE-RETURN RATE BY BRANCH")
    for branch, rate in sorted(late_rate_by(rows, "branch").items(), key=lambda kv: -kv[1])[: args.top]:
        print(f"  {branch:<10}{rate:>7.1%}")

    print("\nAVERAGE WAIT (DAYS) BY CATEGORY")
    for category, wait in sorted(mean_by(rows, "category", "avg_wait_days").items(), key=lambda kv: kv[1])[: args.top]:
        print(f"  {category:<12}{wait:>6.2f}")

    print(f"\nSaved {len(rows)} clean rows to {CLEAN.name}")


if __name__ == "__main__":
    main()