#!/usr/bin/env python3
"""Rate each FEC election cycle of one public file, cycle by cycle, for four uses.

The file is the FEC "contributions from committees to candidates and other committees" bulk file
(pas2) for the five closed cycles 2014, 2016, 2018, 2020 and 2022. For every cycle the script measures how
complete and usable each field is, then rates the cycle for each use in USES below.

    python3 fec_cycle_readiness.py --data-dir <empty folder> --out fec_cycle_readiness.csv
    python3 fec_cycle_readiness.py --data-dir <another empty folder> --out planted.csv --plant duplicate-rows,unmapped-code

Needs Python 3 and the duckdb module (tested on Python 3.12 and 3.14, DuckDB 1.5.6). It prints each downloaded file's byte size and Last-Modified
header, the result table, and the result of its own checks. It prints counts and shares only: no name,
address or other personal value is ever printed or written. It exits 1, and writes no CSV, if a check fails.

The FEC data stands in for a company's own data. FEC contributor names and addresses may not be sold or
used for solicitation or commercial purposes: https://www.fec.gov/updates/sale-or-use-contributor-information/
"""
import argparse
import csv
import os
import shutil
import sys
import urllib.request
import zipfile

import duckdb

BASE = "https://www.fec.gov/files/bulk-downloads"
CYCLES = [2014, 2016, 2018, 2020, 2022]
HEADER_URL = f"{BASE}/data_dictionaries/pas2_header_file.csv"

# What each TRANSACTION_TP code in this file means (FEC: "Transaction type codes"). A code that is not
# listed here is counted as unmapped, and any unmapped row fails the run.
TRANSACTION_TYPES = {
    "24A": "Independent expenditure opposing election of candidate",
    "24C": "Coordinated party expenditure",
    "24E": "Independent expenditure advocating election of candidate",
    "24F": "Communication cost for candidate (only for Form 7 filer)",
    "24K": "Contribution made to nonaffiliated committee",
    "24N": "Communication cost against candidate (only for Form 7 filer)",
    "24P": "Contribution made to possible federal candidate including in-kind contributions",
    "24R": "Election recount disbursement",
    "24Z": "In-kind contribution made to registered filer",
}

# The four uses a buyer or a product team has for a file like this. Each use needs some measures to be at
# least `ok` percent for a plain "Sell", and at least `floor` percent for "Sell with a caveat"; a cycle
# below the floor on any of them is "Fix first". For a measure in LESS_IS_BETTER (negative amounts),
# `ok` is the most allowed for "Sell" and `floor` the most allowed for "Sell with a caveat".
# These thresholds are our choices: change them to yours.
# measure -> (ok, floor)
USES = {
    "link_to_recipient": ("Link each row to its recipient",
                          {"pct_recipient_id": (99.0, 95.0)}),
    "monthly_trend": ("Month-by-month trend inside the cycle",
                      {"pct_date_in_cycle": (99.0, 95.0), "months_present": (24, 22)}),
    "state_zip_map": ("State and ZIP maps",
                      {"pct_state": (99.0, 95.0), "pct_zip5": (95.0, 90.0)}),
    "dollar_totals": ("Dollar totals a buyer can reconcile",
                      {"pct_amount_parsed": (100.0, 99.9), "pct_sub_id_unique": (100.0, 99.9),
                       "pct_tran_id": (99.0, 95.0), "pct_amount_negative": (10.0, 20.0)}),
}

LESS_IS_BETTER = {"pct_amount_negative"}
PLANTS = {"duplicate-rows", "unmapped-code"}

MEASURES = ["pct_recipient_id", "pct_cand_id", "pct_date_in_cycle", "pct_date_outside_cycle",
            "pct_date_unusable", "months_present", "pct_state", "pct_zip5", "pct_amount_parsed",
            "pct_amount_negative", "pct_sub_id_unique", "pct_tran_id", "pct_entity_tp", "pct_tp_mapped"]


def fetch(url, dest):
    """Download url to dest (skip if the same size is already there). Print size and Last-Modified."""
    head = urllib.request.urlopen(urllib.request.Request(url, method="HEAD"), timeout=60)
    size, modified = head.headers.get("Content-Length"), head.headers.get("Last-Modified")
    if not (os.path.exists(dest) and str(os.path.getsize(dest)) == size):
        with urllib.request.urlopen(url, timeout=60) as src, open(dest, "wb") as dst:
            shutil.copyfileobj(src, dst)
    print(f"source  {url}  bytes={os.path.getsize(dest)}  last-modified={modified}")
    return os.path.getsize(dest)


def load(data_dir, plant):
    os.makedirs(data_dir, exist_ok=True)
    if os.listdir(data_dir):
        sys.exit(f"{data_dir} is not empty: give the script a new, empty folder")
    header_path = os.path.join(data_dir, "pas2_header_file.csv")
    total = fetch(HEADER_URL, header_path)
    with open(header_path, encoding="utf-8") as f:
        cols = f.read().strip().split(",")
    con = duckdb.connect()
    colspec = ", ".join(f"'{c}': 'VARCHAR'" for c in cols)
    source_lines = {}
    first = True
    for cy in CYCLES:
        yy = str(cy)[2:]
        zpath = os.path.join(data_dir, f"pas2{yy}.zip")
        total += fetch(f"{BASE}/{cy}/pas2{yy}.zip", zpath)
        txt = os.path.join(data_dir, f"itpas2_{cy}.txt")
        with zipfile.ZipFile(zpath) as z:
            if z.namelist() != ["itpas2.txt"]:
                sys.exit(f"unexpected zip contents: {z.namelist()}")
            with z.open("itpas2.txt") as src, open(txt, "wb") as dst:
                n_lines = 0
                while chunk := src.read(1 << 20):
                    dst.write(chunk)
                    n_lines += chunk.count(b"\n")
        source_lines[cy] = n_lines
        select = (f"SELECT {cy} AS cycle, row_number() OVER () AS file_row, * FROM read_csv('{txt}', delim='|', "
                  f"header=false, columns={{{colspec}}}, quote='', escape='', encoding='latin-1')")
        con.sql(("CREATE TABLE pas2 AS " if first else "INSERT INTO pas2 ") + select)
        first = False
    print(f"total downloaded: {total} bytes")
    if "duplicate-rows" in plant:
        con.sql("INSERT INTO pas2 SELECT * FROM pas2 WHERE cycle = 2018 ORDER BY file_row LIMIT 1000")
        print("PLANTED: 1,000 rows of the 2018 cycle loaded twice")
    if "unmapped-code" in plant:
        con.sql("UPDATE pas2 SET TRANSACTION_TP = '99' WHERE cycle = 2020 AND file_row <= 500")
        print("PLANTED: 500 rows of the 2020 cycle given the transaction type '99'")
    return con, source_lines


def measure(con):
    con.sql("CREATE TABLE tt (code VARCHAR, meaning VARCHAR)")
    con.executemany("INSERT INTO tt VALUES (?, ?)", list(TRANSACTION_TYPES.items()))
    nb = lambda c: f"NULLIF(TRIM({c}), '')"
    return con.sql(f"""
    WITH d AS (
      SELECT p.*, TRY_STRPTIME({nb('TRANSACTION_DT')}, '%m%d%Y')::DATE AS dt,
             TRY_CAST({nb('TRANSACTION_AMT')} AS DECIMAL(14, 2)) AS amt,
             COUNT(*) OVER (PARTITION BY cycle, {nb('SUB_ID')}) AS sub_id_n
      FROM pas2 p
    ), m AS (
      SELECT cycle, COUNT(*) AS n, COUNT(*) FILTER (WHERE dt >= MAKE_DATE(cycle - 1, 1, 1) AND dt <= MAKE_DATE(cycle, 12, 31)) AS in_cycle,
             COUNT(*) FILTER (WHERE dt IS NOT NULL AND NOT (dt >= MAKE_DATE(cycle - 1, 1, 1) AND dt <= MAKE_DATE(cycle, 12, 31))) AS outside,
             COUNT(*) FILTER (WHERE dt IS NULL) AS unusable,
             COUNT(DISTINCT DATE_TRUNC('month', dt)) FILTER (WHERE dt >= MAKE_DATE(cycle - 1, 1, 1) AND dt <= MAKE_DATE(cycle, 12, 31)) AS months_present,
             COUNT(*) FILTER (WHERE {nb('CAND_ID')} IS NOT NULL OR {nb('OTHER_ID')} IS NOT NULL) AS recipient,
             COUNT(*) FILTER (WHERE {nb('CAND_ID')} IS NOT NULL) AS cand,
             COUNT(*) FILTER (WHERE {nb('STATE')} IS NOT NULL) AS st,
             COUNT(*) FILTER (WHERE REGEXP_MATCHES(TRIM(ZIP_CODE), '^[0-9]{{5}}([0-9]{{4}})?$')) AS zip5,
             COUNT(*) FILTER (WHERE amt IS NOT NULL) AS amt_ok,
             COUNT(*) FILTER (WHERE amt < 0) AS amt_neg,
             COUNT(*) FILTER (WHERE {nb('SUB_ID')} IS NOT NULL AND sub_id_n = 1) AS sub_ok,
             COUNT(*) FILTER (WHERE {nb('TRAN_ID')} IS NOT NULL) AS tran,
             COUNT(*) FILTER (WHERE {nb('ENTITY_TP')} IS NOT NULL) AS ent,
             COUNT(*) FILTER (WHERE TRANSACTION_TP IN (SELECT code FROM tt)) AS mapped
      FROM d GROUP BY cycle
    )
    SELECT cycle, n, in_cycle, outside, unusable, months_present, mapped,
           100.0 * recipient / n, 100.0 * cand / n, 100.0 * in_cycle / n, 100.0 * outside / n, 100.0 * unusable / n,
           100.0 * st / n, 100.0 * zip5 / n, 100.0 * amt_ok / n, 100.0 * amt_neg / n, 100.0 * sub_ok / n,
           100.0 * tran / n, 100.0 * ent / n, 100.0 * mapped / n
    FROM m ORDER BY cycle""").fetchall()


def rate(row):
    """Return {use: rating} for one cycle's measures."""
    out = {}
    for use, (_label, needs) in USES.items():
        rating = "Sell"
        for measure, (ok, floor) in needs.items():
            sign = -1 if measure in LESS_IS_BETTER else 1
            v = sign * row[measure]
            if v < sign * floor:
                rating = "Fix first"
                break
            if v < sign * ok:
                rating = "Sell with a caveat"
        out[use] = rating
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data-dir", required=True, help="a new, empty folder for the downloads")
    ap.add_argument("--out", required=True, help="CSV to write")
    ap.add_argument("--plant", default="", help="comma list: duplicate-rows, unmapped-code (test faults)")
    a = ap.parse_args()
    plant = [x for x in a.plant.split(",") if x]
    if set(plant) - PLANTS:
        ap.error(f"unknown --plant value(s) {sorted(set(plant) - PLANTS)}; use {sorted(PLANTS)}")
    for use, (_label, needs) in USES.items():
        for m, (ok, floor) in needs.items():
            if (ok > floor) if m in LESS_IS_BETTER else (ok < floor):
                sys.exit(f"USES[{use!r}][{m!r}]: Sell threshold {ok} is looser than the caveat threshold {floor}")
    print(f"duckdb {duckdb.__version__}, python {sys.version.split()[0]}")

    con, source_lines = load(a.data_dir, plant)
    raw = measure(con)
    rows = []
    for r in raw:
        (cy, n, in_cycle, outside, unusable, months, mapped, recipient, cand, pin, pout, punus, st, zip5, amt,
         neg, sub, tran, ent, pmapped) = r
        rec = {"cycle": cy, "rows": n, "pct_recipient_id": recipient, "pct_cand_id": cand,
               "pct_date_in_cycle": pin, "pct_date_outside_cycle": pout, "pct_date_unusable": punus,
               "months_present": months, "pct_state": st, "pct_zip5": zip5, "pct_amount_parsed": amt,
               "pct_amount_negative": neg, "pct_sub_id_unique": sub, "pct_tran_id": tran, "pct_entity_tp": ent,
               "pct_tp_mapped": pmapped}
        rec.update({f"rating_{k}": v for k, v in rate(rec).items()})
        rec["_counts"] = (n, in_cycle, outside, unusable, mapped)
        rows.append(rec)

    # The script's own checks.
    failures = []
    for rec in rows:
        cy = rec["cycle"]
        n, in_cycle, outside, unusable, mapped = rec["_counts"]
        if n != source_lines[cy]:
            failures.append(f"{cy}: loaded {n} rows but the source file has {source_lines[cy]} lines "
                            f"({n - source_lines[cy]:+d})")
        if in_cycle + outside + unusable != n:
            failures.append(f"{cy}: date classes add up to {in_cycle + outside + unusable}, not {n}")
        share_sum = rec["pct_date_in_cycle"] + rec["pct_date_outside_cycle"] + rec["pct_date_unusable"]
        if abs(share_sum - 100.0) > 1e-9:
            failures.append(f"{cy}: date shares add up to {share_sum:.6f}%, not 100%")
        for m in MEASURES:
            if m != "months_present" and not (0.0 <= rec[m] <= 100.0):
                failures.append(f"{cy}: {m} is {rec[m]}, outside 0 to 100")
        if mapped != n:
            codes = con.execute("SELECT TRANSACTION_TP, COUNT(*) FROM pas2 WHERE cycle = ? AND TRANSACTION_TP "
                                "NOT IN (SELECT code FROM tt) GROUP BY 1 ORDER BY 1", [cy]).fetchall()
            failures.append(f"{cy}: {n - mapped} rows have a transaction type that is not in TRANSACTION_TYPES: "
                            + ", ".join(f"'{c}' x {k}" for c, k in codes))
        if len(rec) - 1 != 2 + len(MEASURES) + len(USES):
            failures.append(f"{cy}: expected one rating per use, got a different count")
    if [r["cycle"] for r in rows] != CYCLES:
        failures.append(f"cycles measured {[r['cycle'] for r in rows]} differ from {CYCLES}")

    fields = ["cycle", "rows"] + MEASURES + [f"rating_{k}" for k in USES]
    print()
    print("cycle  " + "rows".rjust(8) + "  " + "  ".join(m.replace("pct_", "%") for m in MEASURES))
    for rec in rows:
        print(f"{rec['cycle']}  {rec['rows']:>8}  " + "  ".join(
            (f"{rec[m]:.2f}" if m != "months_present" else f"{rec[m]}").rjust(len(m.replace("pct_", "%")))
            for m in MEASURES))
    print()
    print("cycle  " + "  ".join(f"{k}" for k in USES))
    for rec in rows:
        print(f"{rec['cycle']}  " + "  ".join(rec[f'rating_{k}'].ljust(len(k)) for k in USES))
    print()
    if failures:
        print("SELF-CHECK FAILED: no CSV written")
        for f in failures:
            print("  -", f)
        sys.exit(1)
    print(f"SELF-CHECK PASSED: {len(CYCLES)} cycles, row counts equal the source files, shares add up, "
          f"every transaction type is mapped")
    with open(a.out, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=fields, extrasaction="ignore", lineterminator="\n")
        w.writeheader()
        for rec in rows:
            w.writerow({k: (f"{v:.2f}" if isinstance(v, float) else v) for k, v in rec.items() if k in fields})
    print(f"wrote {a.out}")


if __name__ == "__main__":
    main()