import polars as pl

from scripts import steps_to_parquet


def _frame() -> pl.DataFrame:
    # event 0: step (1,0) spawns track 2 (first-step pre_E=15) → e_sec=15.
    # event 1: step (1,0) spawns tracks 2 & 3 (20 + 30)        → e_sec=50.
    return pl.DataFrame(
        {
            "event_id": [0, 0, 0, 1, 1, 1],
            "track_id": [1, 1, 2, 1, 2, 3],
            "step_no": [0, 1, 0, 0, 0, 0],
            "pre_E": [100.0, 80.0, 15.0, 200.0, 20.0, 30.0],
            "child_track_ids": [[2], [], [], [2, 3], [], []],
        }
    )


def test_e_sec_sums_child_first_step_energy():
    out = steps_to_parquet._add_secondary_energy(_frame())
    e_sec = dict(
        zip(zip(out["track_id"], out["step_no"], out["event_id"]), out["e_sec"])
    )
    assert e_sec[(1, 0, 0)] == 15.0  # one child, first-step pre_E 15
    assert e_sec[(1, 0, 1)] == 50.0  # two children, 20 + 30


def test_e_sec_zero_when_no_children():
    out = steps_to_parquet._add_secondary_energy(_frame())
    childless = out.filter(
        (pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 1)
    )
    assert childless["e_sec"].item() == 0.0


def test_e_sec_preserves_row_count_and_order():
    df = _frame()
    out = steps_to_parquet._add_secondary_energy(df)
    assert out.height == df.height
    assert out["pre_E"].to_list() == df["pre_E"].to_list()