import pandas as pd


# This script is referenced in the 2b_model_train.py as a python_transform(). It
# must contain only the single 'transform()' method with the exact signature
# defined here.
#
def transform(df: pd.DataFrame) -> pd.DataFrame:
    # Create new columns named "height", "weight" and "bmi" based on the
    # existing data in columns "height_cm" and "weight_kg".
    df["height"] = df["height_cm"] / 2.54
    df["weight"] = df["weight_kg"] / 0.45359237
    df["bmi"] = (
        1 / (df["height_cm"].astype(float) ** 2) * 10000 * df["weight_kg"].astype(float)
    )

    # Return a dataframe with only the given columns:
    out_df = df[["pid", "height", "weight", "bmi"]]
    return out_df
