from autark import *
from autark.uprogram import dist_normalized


def main() -> None:
    url = "https://vincentarelbundock.github.io/Rdatasets/csv/AER/Fertility.csv"
    frame = OnceFrame(CsvReader(url))

    program = Program()

    program.columns(None).slice(0, 5).alias("head", frame.schema())
    program.columns("work").reduce("count").alias("row_count")
    program.columns("work").reduce("mean").alias("work_mean")
    program.columns("age").reduce("mean").alias("age_mean")

    morekids_pct = dist_normalized(program.columns("morekids"))
    morekids_pct = morekids_pct.columns(0).concat([morekids_pct.columns(1) * 100.0]).order_by(
        morekids_pct.columns(1), ascending=False
    )
    morekids_pct.alias(
        "morekids_pct", frame.schema_of_columns(["morekids"]).column_f64("percent")
    )

    afam_pct = dist_normalized(program.columns("afam"))
    afam_pct = afam_pct.columns(0).concat([afam_pct.columns(1) * 100.0]).order_by(
        afam_pct.columns(1), ascending=False
    )
    afam_pct.alias("afam_pct", frame.schema_of_columns(["afam"]).column_f64("percent"))

    realized = frame.realize(program)

    print("KEYS:", realized.keys())
    print("\nrow_count:\n", realized["row_count"])
    print("\nwork_mean:\n", realized["work_mean"])
    print("\nage_mean:\n", realized["age_mean"])
    print("\nmorekids_pct:\n", realized["morekids_pct"])
    print("\nafam_pct:\n", realized["afam_pct"])
    print("\nhead:\n", realized["head"])


if __name__ == "__main__":
    main()
