import numpy as np
import polars as pl

from wp39.models import Preprocessor


def test_preprocessor_statistics_come_from_train_only():
    tr = pl.DataFrame({"x": [1.0, 2.0, 3.0, 4.0, 100.0], "z": [0.0, 0.0, 1.0, 1.0, 1.0]})
    te = pl.DataFrame({"x": [1000.0, None, -5.0], "z": [5.0, 5.0, 5.0]})
    pp = Preprocessor(["x", "z"]).fit(tr)
    Xte = pp.transform(te)
    assert pp.fitted_on_rows == 5
    assert Xte[0, 0] <= (pp.hi["x"] - pp.mu["x"]) / pp.sd["x"] + 1e-9        # clipped at the train 99th percentile
    assert abs(Xte[1, 0] - (pp.med["x"] - pp.mu["x"]) / pp.sd["x"]) < 1e-9  # imputed with the train median
    assert np.all(np.isfinite(Xte))