import pandas as pd

file = "./data/pnad_trimestral_trimestre_012026.parquet"


regions_ufs = {
    "Sul": [41, 42, 43],
    "Sudeste": [31, 32, 33, 35],
    "Centro-Oeste": [50, 51, 52, 53],
    "Norte": [11, 12, 13, 14, 15, 16, 17],
    "Nordeste": [21, 22, 23, 24, 25, 26, 27, 28, 29],
}
regions_names = list(regions_ufs.keys())
regions: dict[str, dict[str, int]] = {}


carteira_assinada = "V4029"
# carteira_assinada_opcoes = [1, # sim
#                           2] # não

renda = "VD4019"

df = pd.read_parquet(file)


def get_regions():

    for region_name in regions_names:
        region_df = df[df["UF"].isin(regions_ufs[region_name])]
        regions[region_name] = {
            "informal_count": (region_df[carteira_assinada] == 2).shape[0],
            "informal_mean": region_df.loc[
                region_df[carteira_assinada] == 2, renda
            ].mean(),
            "informal_std": region_df.loc[
                region_df[carteira_assinada] == 2, renda
            ].std(),
            "formal_count": (region_df[carteira_assinada] == 1).shape[0],
            "formal_mean": region_df.loc[
                region_df[carteira_assinada] == 1, renda
            ].mean(),
            "formal_std": region_df.loc[region_df[carteira_assinada] == 1, renda].std(),
        }
    return regions


def main():
    get_regions()
    print(f"TOTAL: {df.shape[0]}")
    print(f"QUANTIDADE DE INFORMAIS: {(df[carteira_assinada] == 2).sum()}")
    print(f"QUANTIDADE DE FORMAIS: {(df[carteira_assinada] == 1).sum()}")
    print("\n")
    for region_name in regions_names:
        print(f"REGIÃO: {region_name}")
        for key, value in regions[region_name].items():
            print(f"{key}: {value}")
        print("------------------------------")
        print("")


if __name__ == "__main__":
    main()