Advanced Project: a data pipeline

import numpy as np
import pandas as pd

HOSPITAIS = [f"Hospital {letra}" for letra in "ABCDEFGHIJKL"]
CONVENIOS = ["Alfa Saúde", "Beta Plano", "Gama Vida", "Delta Med", "Particular", "SUS"]
ESPECIALIDADES = [
    "Cardiologia", "Ortopedia", "Pediatria", "Clínica Geral",
    "Neurologia", "Oncologia", "Obstetrícia", "Cirurgia",
]  # fmt: skip
PROBLEMAS_INJETADOS = (
    "data_invalida",
    "alta_antes",
    "valor_negativo",
    "pago_maior",
    "id_repetido",
    "idade_absurda",
    "convenio_fantasma",
)


def gerar_atendimentos(n: int, semente: int = 11) -> pd.DataFrame:
    """Atendimentos sintéticos, com cerca de 1% de cada problema, em linhas sorteadas e fixas.

    Os convênios também aparecem com caixa e espaços diferentes (isso é sujeira, não erro).
    """
    rng = np.random.default_rng(semente)
    internacao = np.datetime64("2024-01-01") + rng.integers(0, 730, n).astype("timedelta64[D]")
    permanencia = rng.integers(1, 15, n)
    alta = internacao + permanencia.astype("timedelta64[D]")
    cobrado = (permanencia * rng.uniform(800, 2500, n) + rng.uniform(200, 3000, n)).round(2)
    pago = (cobrado * rng.uniform(0.7, 1.0, n)).round(2)
    pago[rng.choice(n, max(1, n // 20), replace=False)] = np.nan  # pagamento ainda pendente
    convenio = np.array(CONVENIOS)[rng.integers(0, len(CONVENIOS), n)].astype(object)
    sujos = rng.choice(n, max(1, n // 10), replace=False)
    convenio[sujos] = [c.lower() + " " for c in convenio[sujos]]

    df = pd.DataFrame(
        {
            "id_atendimento": np.arange(1, n + 1),
            "id_paciente": rng.integers(1, n // 3 + 2, n),
            "hospital": np.array(HOSPITAIS)[rng.integers(0, len(HOSPITAIS), n)],
            "convenio": convenio,
            "especialidade": np.array(ESPECIALIDADES)[rng.integers(0, len(ESPECIALIDADES), n)],
            "data_internacao": internacao.astype(str),
            "data_alta": alta.astype(str),
            "valor_cobrado": cobrado,
            "valor_pago": pago,
            "idade": rng.integers(0, 95, n),
        }
    )
    k = max(1, n // 100)
    sorteadas = rng.permutation(n)
    fatias = {nome: sorteadas[i * k : (i + 1) * k] for i, nome in enumerate(PROBLEMAS_INJETADOS)}
    df.loc[fatias["data_invalida"], "data_alta"] = "sem data"
    df.loc[fatias["alta_antes"], ["data_internacao", "data_alta"]] = df.loc[
        fatias["alta_antes"], ["data_alta", "data_internacao"]
    ].to_numpy()
    df.loc[fatias["valor_negativo"], "valor_cobrado"] = -df.loc[
        fatias["valor_negativo"], "valor_cobrado"
    ]
    df.loc[fatias["pago_maior"], "valor_pago"] = df.loc[fatias["pago_maior"], "valor_cobrado"] * 1.5
    df.loc[fatias["idade_absurda"], "idade"] = 150
    df.loc[fatias["convenio_fantasma"], "convenio"] = "Convênio Fantasma"
    repetidas = fatias["id_repetido"]
    df.loc[repetidas, "id_atendimento"] = (
        df.loc[repetidas, "id_atendimento"].to_numpy() % max(1, n // 2) + 1
    )
    return df

from collections.abc import Collection

import pandas as pd

from saude.gerar import CONVENIOS

_MAPA = {c.lower(): c for c in CONVENIOS}


def normalizar_convenio(serie: pd.Series) -> pd.Series:
    """Nome correto do convênio. O que não está na lista vira ausente (e é detectado)."""
    return serie.str.strip().str.lower().map(_MAPA)


def marcar_problemas(df: pd.DataFrame, ids_vistos: Collection[int] = frozenset()) -> pd.DataFrame:
    """Uma coluna booleana por regra, uma linha por atendimento. Nada é descartado aqui.

    `ids_vistos` são os ids de pedaços anteriores: um id repetido entre pedaços só é detectado
    se o estado for levado de um pedaço para o outro.
    """
    internacao = pd.to_datetime(df["data_internacao"], errors="coerce")
    alta = pd.to_datetime(df["data_alta"], errors="coerce")
    return pd.DataFrame(
        {
            "data_invalida": internacao.isna() | alta.isna(),
            "alta_antes_da_internacao": alta < internacao,
            "valor_negativo": (df["valor_cobrado"] < 0) | (df["valor_pago"] < 0),
            "pago_maior_que_cobrado": df["valor_pago"] > df["valor_cobrado"],
            "id_repetido": df["id_atendimento"].duplicated()
            | df["id_atendimento"].isin(ids_vistos),
            "idade_invalida": ~df["idade"].between(0, 120),
            "convenio_desconhecido": normalizar_convenio(df["convenio"]).isna(),
        },
        index=df.index,
    )


def contar_problemas(problemas: pd.DataFrame) -> dict[str, int]:
    return {motivo: int(n) for motivo, n in problemas.sum().items()}

from collections.abc import Collection
from dataclasses import dataclass

import pandas as pd

from saude.validacao import contar_problemas, marcar_problemas, normalizar_convenio


@dataclass(frozen=True)
class Resultado:
    limpos: pd.DataFrame
    quarentena: pd.DataFrame
    contagem: dict[str, int]


def tipar(df: pd.DataFrame) -> pd.DataFrame:
    """Tipos certos e colunas derivadas. Só vale para linhas já validadas."""
    return (
        df.assign(
            data_internacao=pd.to_datetime(df["data_internacao"]),
            data_alta=pd.to_datetime(df["data_alta"]),
            convenio=normalizar_convenio(df["convenio"]),
        )
        .assign(
            dias_internacao=lambda d: (d["data_alta"] - d["data_internacao"]).dt.days,
            taxa_glosa=lambda d: 1 - d["valor_pago"] / d["valor_cobrado"],
            mes=lambda d: d["data_internacao"].dt.to_period("M").astype(str),
        )
        .astype(
            {
                "hospital": "category",
                "convenio": "category",
                "especialidade": "category",
                "id_atendimento": "int32",
                "id_paciente": "int32",
                "idade": "int16",
                "dias_internacao": "int16",
            }
        )
    )


def separar(df: pd.DataFrame, ids_vistos: Collection[int] = frozenset()) -> Resultado:
    """Divide em limpos e quarentena: nenhuma linha some, cada uma vai para um dos lados."""
    problemas = marcar_problemas(df, ids_vistos)
    invalido = problemas.any(axis=1)
    rotulos = problemas.columns.to_numpy()
    motivos = [",".join(rotulos[linha]) for linha in problemas[invalido].to_numpy()]
    quarentena = df[invalido].assign(motivos=motivos)
    return Resultado(tipar(df[~invalido]), quarentena, contar_problemas(problemas))


def memoria_mb(df: pd.DataFrame) -> float:
    return float(df.memory_usage(deep=True).sum() / 1024 / 1024)

from collections import Counter
from pathlib import Path

import pandas as pd

from saude.limpeza import separar

CHAVES = ["hospital", "mes"]
SOMAS = ["atendimentos", "dias", "cobrado", "pago", "cobrado_com_pagamento"]


def somar(limpos: pd.DataFrame) -> pd.DataFrame:
    """Somas e contagens por hospital e mês: o que pode ser combinado entre pedaços."""
    base = limpos.assign(
        hospital=limpos["hospital"].astype(str),
        cobrado_com_pagamento=limpos["valor_cobrado"].where(limpos["valor_pago"].notna()),
    )
    return (
        base.groupby(CHAVES, observed=True)
        .agg(
            atendimentos=("id_atendimento", "count"),
            dias=("dias_internacao", "sum"),
            cobrado=("valor_cobrado", "sum"),
            pago=("valor_pago", "sum"),
            cobrado_com_pagamento=("cobrado_com_pagamento", "sum"),
        )
        .reset_index()
    )


def finalizar(somas: pd.DataFrame) -> pd.DataFrame:
    """As razões saem dos totais, nunca de médias de médias."""
    return (
        somas.groupby(CHAVES, as_index=False)[SOMAS]
        .sum()
        .assign(
            permanencia_media=lambda d: d["dias"] / d["atendimentos"],
            percentual_pago=lambda d: d["pago"] / d["cobrado_com_pagamento"],
        )
        .sort_values(CHAVES)
        .reset_index(drop=True)
    )


def agregar(limpos: pd.DataFrame) -> pd.DataFrame:
    return finalizar(somar(limpos))


def agregar_em_pedacos(
    caminho: Path, tamanho: int, quarentena_csv: Path | None = None
) -> tuple[pd.DataFrame, dict[str, int], int]:
    """Lê o arquivo em pedaços, levando o conjunto de ids vistos de um pedaço ao outro."""
    vistos: set[int] = set()
    parciais = []
    contagem: Counter[str] = Counter()
    em_quarentena = 0
    primeiro_pedaco = True
    for pedaco in pd.read_csv(caminho, chunksize=tamanho):
        resultado = separar(pedaco, frozenset(vistos))
        vistos.update(pedaco["id_atendimento"].tolist())
        parciais.append(somar(resultado.limpos))
        contagem.update(resultado.contagem)
        em_quarentena += len(resultado.quarentena)
        if quarentena_csv is not None:
            resultado.quarentena.to_csv(
                quarentena_csv,
                mode="w" if primeiro_pedaco else "a",
                header=primeiro_pedaco,
                index=False,
            )
        primeiro_pedaco = False
    return finalizar(pd.concat(parciais, ignore_index=True)), dict(contagem), em_quarentena

import argparse
import time
from pathlib import Path

import pandas as pd

from saude.agregacao import agregar_em_pedacos
from saude.gerar import gerar_atendimentos
from saude.limpeza import memoria_mb, separar


def criar_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="saude", description="Pipeline de atendimentos")
    sub = parser.add_subparsers(dest="comando", required=True)
    gerar = sub.add_parser("gerar", help="gera um arquivo de atendimentos sintéticos")
    gerar.add_argument("--linhas", type=int, default=200_000)
    gerar.add_argument("--semente", type=int, default=11)
    gerar.add_argument("--saida", type=Path, required=True)
    processar = sub.add_parser("processar", help="valida, separa a quarentena e agrega")
    processar.add_argument("arquivo", type=Path)
    processar.add_argument("--saida", type=Path, required=True)
    processar.add_argument("--tamanho", type=int, default=50_000, help="linhas por pedaço")
    return parser


def main(argv: list[str] | None = None) -> int:
    args = criar_parser().parse_args(argv)
    if args.comando == "gerar":
        args.saida.parent.mkdir(parents=True, exist_ok=True)
        gerar_atendimentos(args.linhas, args.semente).to_csv(args.saida, index=False)
        print(f"{args.linhas} atendimentos gravados em {args.saida}")
        return 0

    args.saida.mkdir(parents=True, exist_ok=True)
    inicio = time.perf_counter()
    try:
        resumo, contagem, quarentena = agregar_em_pedacos(
            args.arquivo, args.tamanho, args.saida / "quarentena.csv"
        )
        amostra = pd.read_csv(args.arquivo, nrows=args.tamanho)
    except (OSError, ValueError, KeyError) as erro:
        print(f"Erro: {erro}")
        return 1
    resumo.to_parquet(args.saida / "resumo_hospital_mes.parquet", index=False)
    validas = int(resumo["atendimentos"].sum())
    print(
        f"linhas lidas: {validas + quarentena} | válidas: {validas} | em quarentena: {quarentena}"
    )
    print("linhas com cada problema (uma linha pode ter vários):")
    for motivo, n in contagem.items():
        print(f"  {motivo:<26}{n:>6}")
    print(f"resumo: {len(resumo)} combinações de hospital e mês")
    crus = amostra.head(1000)
    tipados = separar(crus).limpos
    print(
        f"memória ({len(tipados)} linhas válidas): "
        f"{memoria_mb(crus.loc[tipados.index]):.2f} MB crus, {memoria_mb(tipados):.2f} MB tipados"
    )
    print(f"tempo: {time.perf_counter() - inicio:.1f} s")
    return 0

import pandas as pd
import pytest
from pandas.testing import assert_frame_equal

from saude.gerar import PROBLEMAS_INJETADOS, gerar_atendimentos
from saude.limpeza import memoria_mb, separar
from saude.validacao import contar_problemas, marcar_problemas, normalizar_convenio


def linha(**mudancas):
    base = {
        "id_atendimento": 1, "id_paciente": 7, "hospital": "Hospital A", "convenio": "SUS",
        "especialidade": "Cirurgia", "data_internacao": "2025-01-10", "data_alta": "2025-01-15",
        "valor_cobrado": 1000.0, "valor_pago": 900.0, "idade": 40,
    }  # fmt: skip
    return {**base, **mudancas}


def test_a_geracao_e_deterministica_e_tem_cada_problema() -> None:
    a, b = gerar_atendimentos(3000), gerar_atendimentos(3000)
    assert_frame_equal(a, b)
    contagem = contar_problemas(marcar_problemas(a))
    assert all(n > 0 for n in contagem.values()), contagem
    assert len(PROBLEMAS_INJETADOS) == 7


def test_cada_regra_pega_o_seu_caso() -> None:
    df = pd.DataFrame(
        [
            linha(id_atendimento=1),
            linha(id_atendimento=2, data_alta="2025-01-05"),
            linha(id_atendimento=3, valor_cobrado=-5.0),
            linha(id_atendimento=4, valor_pago=2000.0),
            linha(id_atendimento=1),
            linha(id_atendimento=6, idade=150),
            linha(id_atendimento=7, convenio="Inventado"),
            linha(id_atendimento=8, data_internacao="ontem"),
        ]
    )
    p = marcar_problemas(df)
    assert p["alta_antes_da_internacao"].tolist() == [False, True] + [False] * 6
    assert p["valor_negativo"].iloc[2] and p["pago_maior_que_cobrado"].iloc[3]
    assert p["id_repetido"].tolist() == [False, False, False, False, True, False, False, False]
    assert p["idade_invalida"].iloc[5] and p["convenio_desconhecido"].iloc[6]
    assert p["data_invalida"].iloc[7]
    assert p.iloc[0].sum() == 0


def test_id_repetido_entre_pedacos_exige_o_estado() -> None:
    df = pd.DataFrame([linha(id_atendimento=5)])
    assert not marcar_problemas(df)["id_repetido"].iloc[0]
    assert marcar_problemas(df, ids_vistos={5})["id_repetido"].iloc[0]


def test_convenio_normaliza_caixa_e_espacos() -> None:
    assert normalizar_convenio(pd.Series([" sus ", "alfa saúde ", "xyz"])).tolist()[:2] == [
        "SUS",
        "Alfa Saúde",
    ]
    assert normalizar_convenio(pd.Series(["xyz"])).isna().all()


@pytest.fixture(scope="module")
def gerado() -> pd.DataFrame:
    return gerar_atendimentos(20_000)


def test_nenhuma_linha_se_perde(gerado) -> None:
    r = separar(gerado)
    assert len(r.limpos) + len(r.quarentena) == len(gerado)
    assert set(r.limpos.index).isdisjoint(r.quarentena.index)
    assert r.quarentena["motivos"].str.len().gt(0).all()


def test_o_que_sobra_nao_viola_nenhuma_regra(gerado) -> None:
    r = separar(gerado)
    depois = marcar_problemas(
        r.limpos.assign(
            data_internacao=r.limpos["data_internacao"].astype(str),
            data_alta=r.limpos["data_alta"].astype(str),
            convenio=r.limpos["convenio"].astype(str),
        )
    )
    assert not depois.any().any()


def test_separar_nao_altera_a_entrada(gerado) -> None:
    copia = gerado.copy()
    separar(gerado)
    assert_frame_equal(gerado, copia)


def test_tipos_e_colunas_derivadas(gerado) -> None:
    limpos = separar(gerado).limpos
    assert (
        str(limpos["hospital"].dtype) == "category"
        and str(limpos["id_atendimento"].dtype) == "int32"
    )
    assert (limpos["dias_internacao"] >= 0).all()
    assert limpos["taxa_glosa"].dropna().between(0, 0.31).all()


def test_tipar_reduz_a_memoria(gerado) -> None:
    limpos = separar(gerado).limpos
    assert memoria_mb(limpos) < 0.7 * memoria_mb(gerado.loc[limpos.index])

from pathlib import Path

import pandas as pd
import pytest
from pandas.testing import assert_frame_equal

from saude.agregacao import agregar, agregar_em_pedacos
from saude.cli import main
from saude.gerar import gerar_atendimentos
from saude.limpeza import separar


@pytest.fixture(scope="module")
def arquivo(tmp_path_factory) -> Path:
    caminho = tmp_path_factory.mktemp("dados") / "atendimentos.csv"
    gerar_atendimentos(9000).to_csv(caminho, index=False)
    return caminho


@pytest.mark.parametrize("tamanho", [1000, 2500, 20_000])
def test_em_pedacos_e_igual_a_ler_tudo(arquivo: Path, tamanho: int) -> None:
    completo = agregar(separar(pd.read_csv(arquivo)).limpos)
    em_pedacos, _, _ = agregar_em_pedacos(arquivo, tamanho)
    assert_frame_equal(completo, em_pedacos)


def test_a_quarentena_tambem_bate_entre_os_dois_caminhos(arquivo: Path) -> None:
    completo = separar(pd.read_csv(arquivo))
    resumo, contagem, em_quarentena = agregar_em_pedacos(arquivo, 1500)
    assert em_quarentena == len(completo.quarentena)
    assert int(resumo["atendimentos"].sum()) == len(completo.limpos)
    assert contagem == completo.contagem


def test_agregacao_com_resposta_conhecida() -> None:
    base = pd.DataFrame(
        {
            "id_atendimento": [1, 2, 3],
            "hospital": ["Hospital A", "Hospital A", "Hospital B"],
            "convenio": ["SUS", "SUS", "SUS"],
            "especialidade": ["Cirurgia"] * 3,
            "id_paciente": [1, 2, 3],
            "data_internacao": ["2025-01-10", "2025-01-20", "2025-01-12"],
            "data_alta": ["2025-01-12", "2025-01-25", "2025-01-13"],
            "valor_cobrado": [100.0, 200.0, 50.0],
            "valor_pago": [100.0, None, 25.0],
            "idade": [30, 40, 50],
        }
    )
    resumo = agregar(separar(base).limpos).set_index("hospital")
    assert resumo.loc["Hospital A", "atendimentos"] == 2
    assert resumo.loc["Hospital A", "permanencia_media"] == pytest.approx((2 + 5) / 2)
    assert resumo.loc["Hospital A", "percentual_pago"] == pytest.approx(1.0)
    assert resumo.loc["Hospital B", "percentual_pago"] == pytest.approx(0.5)


def test_pipeline_de_ponta_a_ponta(tmp_path: Path, capsys) -> None:
    entrada = tmp_path / "a.csv"
    assert main(["gerar", "--linhas", "4000", "--saida", str(entrada)]) == 0
    destino = tmp_path / "res"
    assert main(["processar", str(entrada), "--saida", str(destino), "--tamanho", "1500"]) == 0
    saida = capsys.readouterr().out
    assert "linhas lidas: 4000" in saida
    assert len(pd.read_parquet(destino / "resumo_hospital_mes.parquet")) > 0
    assert len(pd.read_csv(destino / "quarentena.csv")) > 0


def test_arquivo_inexistente_devolve_1(tmp_path: Path, capsys) -> None:
    assert main(["processar", str(tmp_path / "nada.csv"), "--saida", str(tmp_path / "r")]) == 1
    assert "Erro" in capsys.readouterr().out

cd projetos/saude
uv sync
uv run pytest
uv run saude gerar --linhas 200000 --saida saida/atendimentos.csv
uv run saude processar saida/atendimentos.csv --saida saida/resultado

200000 atendimentos gravados em saida/atendimentos.csv
linhas lidas: 200000 | válidas: 186081 | em quarentena: 13919
linhas com cada problema (uma linha pode ter vários):
  data_invalida               2000
  alta_antes_da_internacao    2000
  valor_negativo              2000
  pago_maior_que_cobrado      3893
  id_repetido                 1979
  idade_invalida              2000
  convenio_desconhecido       2000
resumo: 288 combinações de hospital e mês
memória (936 linhas válidas): 0.12 MB crus, 0.07 MB tipados
tempo: 1.0 s