#!/usr/bin/env python3
"""Counts the Spanish and Balearic rental fleet from DGT "parque" microdata.

Reads the national yearly ZIPs (`parque_consolidado_<year>.zip`) straight from
the archive, never unpacking them, and writes aggregates only: no row, make,
model or owner attribute of a single vehicle leaves this process. See README.md
for the sources, the field definitions and the method.

Usage:
    python3 -I count_fleet.py <dir-with-zips> <output-dir> [year ...]

Writes `raw_aggregates.json` (working file, includes municipality cells, do not
publish) and `results.csv` (the published table).
"""

import csv
import json
import multiprocessing
import sys
import zipfile
from collections import Counter
from pathlib import Path

YEARS = (2020, 2021, 2022, 2023, 2024, 2025)
RENTAL_SERVICE_CODE = "A01"  # RD 2822/1998 Annex II.D: public service, hire without driver
BALEARICS = "07"
UNALLOCATED = "Unallocated"

# INE municipality codes per island. Mallorca is listed too, so an unexpected
# code raises KeyError instead of silently landing on an island.
ISLAND_MUNICIPALITIES = {
    "Ibiza": "07026 07046 07048 07050 07054",
    "Formentera": "07024",
    "Menorca": "07002 07015 07023 07032 07037 07052 07064 07902",
    "Mallorca": (
        "07001 07003 07004 07005 07006 07007 07008 07009 07010 07011 07012 "
        "07013 07014 07016 07017 07018 07019 07020 07021 07022 07025 07027 "
        "07028 07029 07030 07031 07033 07034 07035 07036 07038 07039 07040 "
        "07041 07042 07043 07044 07045 07047 07049 07051 07053 07055 07056 "
        "07057 07058 07059 07060 07061 07062 07063 07065 07901"
    ),
}
ISLAND_BY_MUNICIPALITY = {
    code: island
    for island, codes in ISLAND_MUNICIPALITIES.items()
    for code in codes.split()
}

assert len(ISLAND_BY_MUNICIPALITY) == 67, "INE lists 67 Balearic municipalities"

CARS = "TURISMOS"
TWO_WHEELERS = ("MOTOCICLETAS", "CICLOMOTORES")


def type_group(tipo_dgt: str) -> str:
    """Folds TIPO_DGT into the three groups the published table uses."""
    if tipo_dgt == CARS:
        return "cars"
    if tipo_dgt in TWO_WHEELERS:
        return "motorcycles_mopeds"
    return "other"


def island_of(municipio: str) -> str:
    """Maps MUNICIPIO to an island; blank means DGT suppressed it (< 10,000 residents)."""
    return ISLAND_BY_MUNICIPALITY[municipio] if municipio else UNALLOCATED


def count_year(args: tuple[Path, int]) -> dict:
    """Streams one year's archive and returns its aggregates."""
    directory, year = args
    rows = malformed = 0
    latest = ""  # yyyymmdd of the newest FECHA_MATR, proves the file's cut-off
    provinces: Counter[str] = Counter()
    services: Counter[str] = Counter()
    types: Counter[str] = Counter()
    renting: Counter[str] = Counter()
    cells: Counter[tuple[str, str, str]] = Counter()
    municipalities: Counter[str] = Counter()
    holders: Counter[str] = Counter()

    with zipfile.ZipFile(directory / f"parque_consolidado_{year}.zip") as bundle:
        members = bundle.namelist()
        if len(members) != 1:
            raise SystemExit(f"{year}: expected one member, found {members}")
        with bundle.open(members[0]) as stream:
            header = next(stream).decode("latin-1").strip().split("|")
            at = {name: i for i, name in enumerate(header)}
            width = len(header)
            for raw in stream:
                rows += 1
                # latin-1 never fails and only ASCII codes are read.
                f = raw.decode("latin-1").rstrip("\r\n").split("|")
                if len(f) != width:
                    malformed += 1
                    continue
                d = f[at["FECHA_MATR"]]
                stamp = d[6:] + d[3:5] + d[:2]
                if stamp > latest and stamp[:4] <= str(year):
                    latest = stamp
                province = f[at["PROVINCIA"]].strip().zfill(2)
                service = f[at["SERVICIO"]].strip()
                if province == BALEARICS:
                    services[service] += 1
                if service != RENTAL_SERVICE_CODE:
                    continue
                provinces[province] += 1
                vehicle = f[at["TIPO_DGT"]].strip()
                flag = f[at["RENTING"]].strip()
                types[vehicle] += 1
                renting[flag] += 1
                if province == BALEARICS:
                    municipio = f[at["MUNICIPIO"]].strip()
                    cells[(island_of(municipio), type_group(vehicle), flag)] += 1
                    municipalities[municipio] += 1
                    holders[f[at["TIPO_TITULAR"]].strip()] += 1

    assert sum(cells.values()) == provinces[BALEARICS], "islands must sum to province"
    return {
        "year": year,
        "member": members[0],
        "rows": rows,
        "malformed_rows": malformed,
        "latest_registration": latest,
        "a01_by_province": dict(sorted(provinces.items())),
        "a01_by_type_national": dict(types.most_common()),
        "a01_renting_national": dict(renting),
        "balearic_rows_by_service": dict(services.most_common()),
        "balearic_cells": [[*key, n] for key, n in sorted(cells.items())],
        "balearic_by_municipality": dict(municipalities.most_common()),
        "balearic_holder_type": dict(holders),
    }


def write_results(path: Path, results: list[dict]) -> None:
    """Writes year × scope × metric counts.

    Renting stays at province level: per island it drops to cells of one vehicle.
    """
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(("year", "scope", "metric", "vehicles"))
        for r in results:
            year, provinces = r["year"], r["a01_by_province"]
            writer.writerow((year, "Spain", "all_types", sum(provinces.values())))
            writer.writerow((year, "Spain", "cars", r["a01_by_type_national"][CARS]))
            writer.writerow((year, "Madrid province", "all_types", provinces["28"]))
            writer.writerow(
                (year, "Sant Josep de sa Talaia", "all_types", r["balearic_by_municipality"]["07048"])
            )
            totals: Counter[tuple[str, str]] = Counter()
            for island, group, flag, n in r["balearic_cells"]:
                for scope in (island, "Balearics"):
                    totals[(scope, "all_types")] += n
                    totals[(scope, group)] += n
                totals[("Balearics", {"S": "renting_yes", "N": "renting_no"}.get(flag, "renting_blank"))] += n
            for (scope, metric), n in sorted(totals.items()):
                writer.writerow((year, scope, metric, n))


def main() -> None:
    if len(sys.argv) < 3:
        raise SystemExit(__doc__)
    directory, output = Path(sys.argv[1]), Path(sys.argv[2])
    output.mkdir(parents=True, exist_ok=True)
    years = tuple(int(y) for y in sys.argv[3:]) or YEARS
    with multiprocessing.Pool(min(len(years), 4)) as pool:
        results = pool.map(count_year, [(directory, y) for y in years])
    (output / "raw_aggregates.json").write_text(
        json.dumps(results, indent=1) + "\n", encoding="utf-8"
    )
    write_results(output / "results.csv", results)


if __name__ == "__main__":
    main()