import polars as pl
from datetime import datetime

print("polars", pl.__version__)

df = pl.DataFrame(
    {
        "id_store": [1],
        "id_product_group": [10],
        "target_date": [datetime(2025, 1, 1)],
        "id_product": [2],
        "prediction": [1.0],
        "price": [5.0],
    }
)

lf = (
    df.lazy()
    # Force empty with a filter that removes everything
    .filter(pl.lit(False))
    .sort(
        "id_store", "id_product_group", "target_date",
        (pl.col("prediction") * pl.col("price")),
        descending=[False, False, False, True],
    )
    .with_columns(
        id_products = pl.col("id_product")
            .cumulative_eval(pl.element().explode().sort().implode())
            .over("id_store", "id_product_group", "target_date"),
    )
    .with_columns(
        n_products = pl.col("id_products").list.len().cast(pl.Int16),
        id_products = pl.col("id_products").cast(pl.List(pl.Int64)),
    )
)

print("collect_schema:", lf.collect_schema()["id_products"])  # planner's view
out = lf.collect()
print("post_collect schema:", out.collect_schema()["id_products"])  # check if dtype flipped (e.g., List[List[Int64]])

if lf.collect_schema()["id_products"] != out.collect_schema()["id_products"]:
    raise ValueError("Schema changed after collect!")