"""Amostras do Catalogo de Teses e Dissertacoes da CAPES.

Le o CSV oficial em fluxo (sem salvar o arquivo bruto). Dois modos que podem
rodar juntos na mesma leitura:
- amostra estratificada: ate K trabalhos por area de conhecimento, por
  reservoir sampling, nas grandes areas escolhidas;
- amostra dirigida (--policy-map): ate P trabalhos por politica cujo alias
  aparece no titulo ou no resumo.
Nomes de discentes e orientadores sao descartados.

Uso:
  python collect_capes.py --year 2023 --per-area 60 --out ../data/capes_2023.jsonl
  python collect_capes.py --year 2023 --per-area 60 --grandes-areas "CIÊNCIAS DA SAÚDE" --out ...
  python collect_capes.py --year 2023 --per-area 0 --policy-map ../policy_map.csv --per-policy 15 --policy-out ...
"""
import argparse, csv, io, json, random, sys, urllib.request

URLS = {
    2021: "https://dadosabertos.capes.gov.br/dataset/dbb9b83e-7779-4b99-89d2-1ff8099f36f6/resource/068003e4-196c-41f4-8c35-1f7c94b4e55c/download/br-capes-btd-2021-2025-03-31.csv",
    2022: "https://dadosabertos.capes.gov.br/dataset/dbb9b83e-7779-4b99-89d2-1ff8099f36f6/resource/78f73608-6f5e-463c-ba79-0bff4f8a578d/download/br-capes-btd-2022-2025-03-31.csv",
    2023: "https://dadosabertos.capes.gov.br/dataset/dbb9b83e-7779-4b99-89d2-1ff8099f36f6/resource/b69baf26-8d02-4c10-ba39-7e9ab799e6ed/download/br-capes-btd-2023-2025-03-31.csv",
    2024: "https://dadosabertos.capes.gov.br/dataset/dbb9b83e-7779-4b99-89d2-1ff8099f36f6/resource/87133ba7-ac99-4d87-966e-8f580bc96231/download/br-capes-btd-2024-2025-12-01.csv",
}
DEFAULT_AREAS = "CIÊNCIAS HUMANAS;CIÊNCIAS SOCIAIS APLICADAS"
KEEP = ["AN_BASE", "ID_PRODUCAO_INTELECTUAL", "NM_PRODUCAO", "NM_SUBTIPO_PRODUCAO",
        "NM_GRAU_ACADEMICO", "SG_ENTIDADE_ENSINO", "NM_REGIAO", "SG_UF_IES",
        "NM_PROGRAMA", "NM_GRANDE_AREA_CONHECIMENTO", "NM_AREA_CONHECIMENTO",
        "NM_SUBAREA_CONHECIMENTO", "NM_AREA_AVALIACAO", "NM_LINHA_PESQUISA",
        "DS_PALAVRA_CHAVE", "NM_IDIOMA", "DS_RESUMO"]


def load_policies(path):
    out = []
    with open(path, encoding="utf-8") as f:
        for row in csv.DictReader(f):
            aliases = [x.strip().lower() for x in row["aliases"].split("|") if x.strip()]
            out.append((row["politica"], row["cor"], aliases))
    return out


def reservoir(store, seen, key, rec, k, rng):
    seen[key] = seen.get(key, 0) + 1
    bucket = store.setdefault(key, [])
    if len(bucket) < k:
        bucket.append(rec)
    else:
        j = rng.randrange(seen[key])
        if j < k:
            bucket[j] = rec


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--year", type=int, required=True, choices=sorted(URLS))
    ap.add_argument("--per-area", type=int, default=5)
    ap.add_argument("--grandes-areas", default=DEFAULT_AREAS, help="separadas por ;")
    ap.add_argument("--min-chars", type=int, default=400)
    ap.add_argument("--seed", type=int, default=20261013)
    ap.add_argument("--out")
    ap.add_argument("--policy-map")
    ap.add_argument("--per-policy", type=int, default=15)
    ap.add_argument("--policy-out")
    a = ap.parse_args()
    rng = random.Random(a.seed + a.year)
    grandes = {g.strip() for g in a.grandes_areas.split(";") if g.strip()}
    policies = load_policies(a.policy_map) if a.policy_map else []
    csv.field_size_limit(2**31 - 1)
    resp = urllib.request.urlopen(URLS[a.year], timeout=180)
    reader = csv.DictReader(io.TextIOWrapper(resp, encoding="latin-1", newline=""), delimiter=";")
    seen, res, pseen, pres, n = {}, {}, {}, {}, 0
    for row in reader:
        n += 1
        if (row.get("NM_GRANDE_AREA_CONHECIMENTO") or "").strip() not in grandes:
            continue
        resumo = row.get("DS_RESUMO") or ""
        if len(resumo) < a.min_chars:
            continue
        rec = {k: (row.get(k) or "").strip() for k in KEEP}
        rec["fonte"] = "capes"
        rec["id"] = f"capes-{rec['AN_BASE']}-{rec['ID_PRODUCAO_INTELECTUAL']}"
        if a.per_area > 0:
            reservoir(res, seen, rec["NM_AREA_CONHECIMENTO"], rec, a.per_area, rng)
        if policies:
            txt = (rec["NM_PRODUCAO"] + " " + resumo).lower()
            for nome, cor, aliases in policies:
                if any(al in txt for al in aliases):
                    prec = dict(rec, politica_busca=nome, cor_busca=cor)
                    reservoir(pres, pseen, nome, prec, a.per_policy, rng)
    if a.out and a.per_area > 0:
        with open(a.out, "w", encoding="utf-8") as f:
            for area, recs in sorted(res.items()):
                for r in recs:
                    r["populacao_area"] = seen[area]
                    f.write(json.dumps(r, ensure_ascii=False) + "\n")
    if a.policy_out and policies:
        with open(a.policy_out, "w", encoding="utf-8") as f:
            for nome, recs in sorted(pres.items()):
                for r in recs:
                    r["populacao_politica"] = pseen[nome]
                    f.write(json.dumps(r, ensure_ascii=False) + "\n")
    print(json.dumps({"ano": a.year, "linhas_lidas": n, "areas": seen, "politicas": pseen},
                     ensure_ascii=False))


if __name__ == "__main__":
    main()