"""Use column expressions for derived values and boolean filtering."""

import emzed

peaks = emzed.Table.create_table(
    ["compound", "mz", "rt", "area"],
    [str, emzed.MzType, emzed.RtType, float],
    rows=[
        ["caffeine", 195.08765, 42.0, 2.40e6],
        ["glucose", 181.07066, 60.0, 8.10e5],
        ["arginine", 175.11895, 95.0, 1.20e6],
        ["alanine", 90.0550, 30.0, 4.10e5],
    ],
)

peaks.add_column("area_million", peaks.area / 1e6, float)
peaks.add_column("is_late", peaks.rt > 50.0, bool)
peaks.add_column(
    "intensity_class",
    (peaks.area > 1e6).then_else("high", "mid_or_low"),
    str,
)
peaks.add_column(
    "mz_rounded",
    peaks.apply(lambda mz: round(mz, 2), peaks.mz),
    float,
)

filtered = peaks.filter((peaks.rt > 40.0) & (peaks.area > 7e5))
filtered = filtered.sort_by("area", ascending=False)

print("full table with derived columns:")
print(peaks)
print()
print("filtered result:")
print(filtered.extract_columns("compound", "mz", "rt", "area_million", "is_late"))
