tests/test_p2_feature_engineering.py

"""Tests for the P2 feature & target engineering tools (owner: Aditya).

Fully hermetic — in-memory DuckDB only, no PKDD DuckDB file, no Lexsi
SDK, no LLM. Reuses the offline-synthetic fixture family from
`conftest.py` (`ctx_offline`, `loan_context_df`, `seed_sql_result`) and
adds one temporal fixture local to this module.

Covers, per the sprint definition of done:
  * happy path + error paths for each tool;
  * the cache-key contracts (`last_split`, `last_target_definition`,
    `last_imbalance_report`, in-place `sql_result:<label>` mutation);
  * the feature-lineage audit log;
  * registry wiring (all four tools present and Tool-protocol-shaped).
"""

from __future__ import annotations

import pandas as pd
import pytest

from lexsi_ds.agent.tools import REGISTRY, Tool
from lexsi_ds.agent.tools.define_target import TOOL as define_target_tool
from lexsi_ds.agent.tools.define_target import DefineTargetArgs
from lexsi_ds.agent.tools.derive_feature import TOOL as derive_feature_tool
from lexsi_ds.agent.tools.derive_feature import DeriveFeatureArgs
from lexsi_ds.agent.tools.handle_imbalance import TOOL as handle_imbalance_tool
from lexsi_ds.agent.tools.handle_imbalance import HandleImbalanceArgs
from lexsi_ds.agent.tools.time_split import TOOL as time_split_tool
from lexsi_ds.agent.tools.time_split import TimeSplitArgs

P2_TOOLS = ["derive_feature", "define_target", "time_split", "handle_imbalance"]


@pytest.fixture
def temporal_loan_df() -> pd.DataFrame:
    """120 loans across 2016-2019 with a label that drifts over time —
    the shape the leakage scenario needs."""
    import random

    rng = random.Random(13)
    rows = []
    for i in range(120):
        year = 2016 + i // 30
        month = 1 + (i * 7) % 12
        amount = rng.randint(5_000, 200_000)
        duration = rng.choice([12, 24, 36, 48, 60])
        # default rate drifts upward over the years (temporal signal)
        p_default = 0.05 + 0.08 * (year - 2016) + 0.2 * (amount / 200_000)
        status = "B" if rng.random() < p_default else "A"
        rows.append({
            "loan_id": 2000 + i,
            "granted_date": f"{year}-{month:02d}-15",
            "amount": amount,
            "duration": duration,
            "payments": round(amount / duration, 2),
            "status": status,
        })
    return pd.DataFrame(rows)


# ===========================================================================
# registry wiring
# ===========================================================================


def test_p2_tools_registered_and_protocol_shaped():
    for name in P2_TOOLS:
        assert name in REGISTRY, f"{name} missing from REGISTRY"
        tool = REGISTRY[name]
        assert isinstance(tool, Tool)
        assert tool.spec.name == name
        assert tool.spec.description and tool.spec.returns


# ===========================================================================
# derive_feature
# ===========================================================================


def test_derive_feature_duckdb_happy(ctx_offline, loan_context_df, seed_sql_result):
    seed_sql_result(ctx_offline, "loans", loan_context_df)
    res = derive_feature_tool.run(
        DeriveFeatureArgs(df_label="loans", name="payment_burden",
                          expression="payments / amount"),
        ctx_offline,
    )
    assert res.ok, res.summary
    df = ctx_offline.cache["sql_result:loans"]
    assert "payment_burden" in df.columns
    assert len(df) == len(loan_context_df)
    assert "payment_burden" in res.summary
    # lineage recorded
    lineage = ctx_offline.cache["feature_lineage"]
    assert lineage[-1]["tool"] == "derive_feature"
    assert lineage[-1]["column"] == "payment_burden"


def test_derive_feature_pandas_engine(ctx_offline, loan_context_df, seed_sql_result):
    seed_sql_result(ctx_offline, "loans", loan_context_df)
    res = derive_feature_tool.run(
        DeriveFeatureArgs(df_label="loans", name="amt_per_month",
                          expression="amount / duration", engine="pandas"),
        ctx_offline,
    )
    assert res.ok, res.summary
    assert "amt_per_month" in ctx_offline.cache["sql_result:loans"].columns


def test_derive_feature_unknown_label(ctx_offline):
    res = derive_feature_tool.run(
        DeriveFeatureArgs(df_label="nope", name="x", expression="1 + 1"),
        ctx_offline,
    )
    assert not res.ok
    assert res.error == "unknown_label"


def test_derive_feature_unknown_column_lists_available(
    ctx_offline, loan_context_df, seed_sql_result
):
    seed_sql_result(ctx_offline, "loans", loan_context_df)
    res = derive_feature_tool.run(
        DeriveFeatureArgs(df_label="loans", name="risk_score",
                          expression="loan_amount / monthly_income"),
        ctx_offline,
    )
    assert not res.ok
    assert "Available columns" in res.summary
    assert "risk_score" not in ctx_offline.cache["sql_result:loans"].columns


def test_derive_feature_overwrite_guard(ctx_offline, loan_context_df, seed_sql_result):
    seed_sql_result(ctx_offline, "loans", loan_context_df)
    res = derive_feature_tool.run(
        DeriveFeatureArgs(df_label="loans", name="amount", expression="amount * 2"),
        ctx_offline,
    )
    assert not res.ok and res.error == "column_exists"
    res2 = derive_feature_tool.run(
        DeriveFeatureArgs(df_label="loans", name="amount",
                          expression="amount * 2", overwrite=True),
        ctx_offline,
    )
    assert res2.ok, res2.summary


def test_derive_feature_blocks_statement_sql(ctx_offline, loan_context_df, seed_sql_result):
    seed_sql_result(ctx_offline, "loans", loan_context_df)
    res = derive_feature_tool.run(
        DeriveFeatureArgs(df_label="loans", name="x",
                          expression="1; DROP TABLE loans"),
        ctx_offline,
    )
    assert not res.ok and res.error == "forbidden_expression"


def test_derive_feature_rejects_aggregate(ctx_offline, loan_context_df, seed_sql_result):
    seed_sql_result(ctx_offline, "loans", loan_context_df)
    res = derive_feature_tool.run(
        DeriveFeatureArgs(df_label="loans", name="total", expression="SUM(amount)"),
        ctx_offline,
    )
    assert not res.ok  # 1 value for 60 rows → non_rowwise_expression
    assert res.error == "non_rowwise_expression"


# ===========================================================================
# time_split
# ===========================================================================


def test_time_split_cutoff_happy(ctx_offline, temporal_loan_df, seed_sql_result):
    seed_sql_result(ctx_offline, "loans", temporal_loan_df)
    res = time_split_tool.run(
        TimeSplitArgs(df_label="loans", time_col="granted_date",
                      cutoff="2018-01-01"),
        ctx_offline,
    )
    assert res.ok, res.summary
    train = ctx_offline.cache["sql_result:train"]
    test = ctx_offline.cache["sql_result:test"]
    assert len(train) + len(test) == len(temporal_loan_df)
    # no temporal overlap — the anti-leakage invariant
    assert pd.to_datetime(train["granted_date"]).max() < pd.to_datetime(
        test["granted_date"]
    ).min()
    split = ctx_offline.cache["last_split"]
    assert split["kind"] == "time"
    assert split["n_train"] == len(train) and split["n_test"] == len(test)


def test_time_split_fraction_mode(ctx_offline, temporal_loan_df, seed_sql_result):
    seed_sql_result(ctx_offline, "loans", temporal_loan_df)
    res = time_split_tool.run(
        TimeSplitArgs(df_label="loans", time_col="granted_date",
                      test_fraction=0.25, train_label="tr", test_label="te"),
        ctx_offline,
    )
    assert res.ok, res.summary
    te = ctx_offline.cache["sql_result:te"]
    assert 0 < len(te) <= int(0.35 * len(temporal_loan_df))  # ties allowed


def test_time_split_requires_exactly_one_mode(ctx_offline, temporal_loan_df, seed_sql_result):
    seed_sql_result(ctx_offline, "loans", temporal_loan_df)
    both = time_split_tool.run(
        TimeSplitArgs(df_label="loans", time_col="granted_date",
                      cutoff="2018-01-01", test_fraction=0.2),
        ctx_offline,
    )
    neither = time_split_tool.run(
        TimeSplitArgs(df_label="loans", time_col="granted_date"),
        ctx_offline,
    )
    assert not both.ok and both.error == "bad_args"
    assert not neither.ok and neither.error == "bad_args"


def test_time_split_degenerate_cutoff(ctx_offline, temporal_loan_df, seed_sql_result):
    seed_sql_result(ctx_offline, "loans", temporal_loan_df)
    res = time_split_tool.run(
        TimeSplitArgs(df_label="loans", time_col="granted_date",
                      cutoff="1990-01-01"),
        ctx_offline,
    )
    assert not res.ok and res.error == "degenerate_split"
    assert "spans" in res.summary  # reports the real range for recovery


def test_time_split_non_time_column(ctx_offline, loan_context_df, seed_sql_result):
    seed_sql_result(ctx_offline, "loans", loan_context_df)
    res = time_split_tool.run(
        TimeSplitArgs(df_label="loans", time_col="status", cutoff="2018-01-01"),
        ctx_offline,
    )
    assert not res.ok and res.error == "bad_time_column"


def test_time_split_small_side_warning(ctx_offline, temporal_loan_df, seed_sql_result):
    seed_sql_result(ctx_offline, "loans", temporal_loan_df)
    res = time_split_tool.run(
        TimeSplitArgs(df_label="loans", time_col="granted_date",
                      cutoff="2019-11-01"),
        ctx_offline,
    )
    assert res.ok
    assert "Warnings" in res.summary  # tiny test side


# ===========================================================================
# define_target
# ===========================================================================


def test_define_target_rule_with_scope(ctx_offline, temporal_loan_df, seed_sql_result):
    seed_sql_result(ctx_offline, "loans", temporal_loan_df)
    res = define_target_tool.run(
        DefineTargetArgs(df_label="loans", target_column="y_default",
                         positive_definition="status = 'B'",
                         scope_filter="status IN ('A','B')",
                         entity_column="loan_id"),
        ctx_offline,
    )
    assert res.ok, res.summary
    df = ctx_offline.cache["sql_result:loans"]
    assert set(df["y_default"].unique()) == {0, 1}
    definition = ctx_offline.cache["last_target_definition"]
    assert definition["target_column"] == "y_default"
    assert definition["entity_column"] == "loan_id"
    assert definition["class_balance"] is not None
    assert "Class balance" in res.summary


def test_define_target_existing_numeric_regression(
    ctx_offline, loan_context_df, seed_sql_result
):
    seed_sql_result(ctx_offline, "loans", loan_context_df)
    res = define_target_tool.run(
        DefineTargetArgs(df_label="loans", target_column="amount",
                         task_type="regression"),
        ctx_offline,
    )
    assert res.ok, res.summary
    assert "range" in res.summary.lower()


def test_define_target_single_class_error(ctx_offline, temporal_loan_df, seed_sql_result):
    seed_sql_result(ctx_offline, "loans", temporal_loan_df)
    res = define_target_tool.run(
        DefineTargetArgs(df_label="loans", target_column="y",
                         positive_definition="status = 'X'",
                         scope_filter="status = 'A'"),
        ctx_offline,
    )
    assert not res.ok and res.error == "single_class_target"


def test_define_target_unknown_existing_column(ctx_offline, loan_context_df, seed_sql_result):
    seed_sql_result(ctx_offline, "loans", loan_context_df)
    res = define_target_tool.run(
        DefineTargetArgs(df_label="loans", target_column="not_there"),
        ctx_offline,
    )
    assert not res.ok and res.error == "unknown_column"


def test_define_target_regression_non_numeric(ctx_offline, loan_context_df, seed_sql_result):
    seed_sql_result(ctx_offline, "loans", loan_context_df)
    res = define_target_tool.run(
        DefineTargetArgs(df_label="loans", target_column="status",
                         task_type="regression"),
        ctx_offline,
    )
    assert not res.ok and res.error == "non_numeric_target"


# ===========================================================================
# handle_imbalance
# ===========================================================================


def test_handle_imbalance_severe(ctx_offline, seed_sql_result):
    df = pd.DataFrame({
        "y": [1] * 8 + [0] * 192,
        "x": range(200),
    })
    seed_sql_result(ctx_offline, "loans", df)
    res = handle_imbalance_tool.run(
        HandleImbalanceArgs(df_label="loans", target_column="y"),
        ctx_offline,
    )
    assert res.ok, res.summary
    report = ctx_offline.cache["last_imbalance_report"]
    assert report["severity"] == "severe"
    assert report["imbalance_ratio"] == 24.0
    assert report["strategies"][0]["strategy"] == "class_weights"
    # offline ⇒ no Lexsi synthetic option in the proposal
    assert all(
        s["strategy"] != "lexsi_synthetic_augmentation"
        for s in report["strategies"]
    )
    assert "AUROC" in res.summary


def test_handle_imbalance_balanced(ctx_offline, seed_sql_result):
    df = pd.DataFrame({"y": [0, 1] * 50})
    seed_sql_result(ctx_offline, "d", df)
    res = handle_imbalance_tool.run(
        HandleImbalanceArgs(df_label="d", target_column="y"), ctx_offline
    )
    assert res.ok
    assert ctx_offline.cache["last_imbalance_report"]["severity"] == "balanced"
    assert res.payload["strategies"][0]["strategy"] == "none"


def test_handle_imbalance_defaults_target_from_definition(
    ctx_offline, temporal_loan_df, seed_sql_result
):
    seed_sql_result(ctx_offline, "loans", temporal_loan_df)
    define_target_tool.run(
        DefineTargetArgs(df_label="loans", target_column="y_default",
                         positive_definition="status = 'B'"),
        ctx_offline,
    )
    res = handle_imbalance_tool.run(
        HandleImbalanceArgs(df_label="loans"), ctx_offline  # no target arg
    )
    assert res.ok, res.summary
    assert res.payload["target_column"] == "y_default"


def test_handle_imbalance_continuous_target_error(
    ctx_offline, loan_context_df, seed_sql_result
):
    seed_sql_result(ctx_offline, "loans", loan_context_df)
    res = handle_imbalance_tool.run(
        HandleImbalanceArgs(df_label="loans", target_column="amount"),
        ctx_offline,
    )
    assert not res.ok and res.error == "non_categorical_target"


def test_handle_imbalance_no_target_anywhere(ctx_offline, loan_context_df, seed_sql_result):
    seed_sql_result(ctx_offline, "loans", loan_context_df)
    res = handle_imbalance_tool.run(
        HandleImbalanceArgs(df_label="loans"), ctx_offline
    )
    assert not res.ok and res.error == "no_target"


# ===========================================================================
# end-to-end P2 chain + lineage (the demo-step-3 shape, offline)
# ===========================================================================


def test_p2_chain_lineage(ctx_offline, temporal_loan_df, seed_sql_result):
    """run_sql → derive_feature → define_target → time_split →
    handle_imbalance leaves a complete, ordered audit trail."""
    seed_sql_result(ctx_offline, "loans", temporal_loan_df)

    assert derive_feature_tool.run(
        DeriveFeatureArgs(df_label="loans", name="payment_burden",
                          expression="payments / amount"),
        ctx_offline,
    ).ok
    assert define_target_tool.run(
        DefineTargetArgs(df_label="loans", target_column="y_default",
                         positive_definition="status = 'B'",
                         entity_column="loan_id"),
        ctx_offline,
    ).ok
    assert time_split_tool.run(
        TimeSplitArgs(df_label="loans", time_col="granted_date",
                      cutoff="2018-01-01"),
        ctx_offline,
    ).ok
    assert handle_imbalance_tool.run(
        HandleImbalanceArgs(df_label="train"), ctx_offline
    ).ok

    lineage = ctx_offline.cache["feature_lineage"]
    assert [r["tool"] for r in lineage] == [
        "derive_feature", "define_target", "time_split", "handle_imbalance",
    ]
    assert [r["step"] for r in lineage] == [1, 2, 3, 4]
    # the train slice carries both the derived feature and the target
    train = ctx_offline.cache["sql_result:train"]
    assert {"payment_burden", "y_default"} <= set(train.columns)