import pandas as pd
import pytest

from psx_signal.labels import add_forward_labels


def test_forward_returns_enter_at_next_open_and_use_trading_rows() -> None:
    dates = pd.to_datetime(["2026-08-13", "2026-08-17", "2026-08-18", "2026-08-19", "2026-08-20", "2026-08-21"])
    frame = pd.DataFrame({
        "symbol": "AAA", "date": dates,
        "open": [90, 100, 110, 120, 130, 140],
        "close": [95, 105, 115, 125, 135, 150],
    })
    result = add_forward_labels(frame, horizons=(1, 5), buy_threshold=0.01, sell_threshold=-0.01)
    first = result.iloc[0]
    assert first["forward_return_1d"] == pytest.approx(0.05)
    assert first["forward_return_5d"] == pytest.approx(0.50)
    assert first["execution_date_1d"] == pd.Timestamp("2026-08-17")
    assert first["exit_date_5d"] == pd.Timestamp("2026-08-21")
    assert first["label_1d"] == 1
    assert pd.isna(result.iloc[-1]["label_1d"])


def test_labels_never_cross_symbols() -> None:
    frame = pd.DataFrame({
        "symbol": ["AAA", "AAA", "BBB", "BBB"],
        "date": pd.to_datetime(["2026-01-01", "2026-01-02"] * 2),
        "open": [10, 20, 100, 200], "close": [10, 22, 100, 180],
    })
    result = add_forward_labels(frame, horizons=(1,), buy_threshold=0.05, sell_threshold=-0.05)
    labels = result[result["date"] == pd.Timestamp("2026-01-01")].set_index("symbol")["label_1d"]
    assert labels.to_dict() == {"AAA": 1, "BBB": -1}