"""Test the correct implementation that handles any input order."""
import polars as pl


def correct_create_ranking_targets(df, horizons=[2], price_col="close", ticker_col="ticker", date_col="date"):
    """
    Correct implementation that:
    1. Explicitly sorts by [ticker, date] ascending
    2. Uses shift(-1) to move returns FORWARD in time
    
    Per docstring: features on day D predict returns from D+1 to D+1+h
    """
    # CRITICAL: Sort to ensure consistent ordering
    result_df = df.sort([ticker_col, date_col])
    
    for h in horizons:
        fwd_temp = f"_fwd_{h}d_temp"
        fwd_return = f"fwd_return_{h}d"
        
        result_df = (
            result_df
            # Step 1: Calculate h-day forward return on day D: price[D+h]/price[D] - 1
            .with_columns(
                (
                    pl.col(price_col).shift(-h).over(ticker_col) / pl.col(price_col) - 1
                ).alias(fwd_temp)
            )
            # Step 2: Shift UP by 1 so the return sits on day D-1
            # This way, features on D-1 predict returns from D to D+h
            # (Equivalently: features on D predict returns from D+1 to D+1+h)
            .with_columns(pl.col(fwd_temp).shift(-1).over(ticker_col).alias(fwd_return))
            .drop(fwd_temp)
        )
    
    return result_df


def main():
    # Create test data
    base_df = pl.DataFrame({
        "date": pl.date_range(pl.date(2024, 1, 1), pl.date(2024, 1, 10), eager=True),
        "ticker": ["BTC"] * 10,
        "close": [100.0, 110.0, 120.0, 115.0, 125.0, 130.0, 140.0, 135.0, 145.0, 150.0]
    })
    
    prices = base_df["close"].to_list()
    h = 2
    
    print("=" * 80)
    print("CORRECT IMPLEMENTATION TEST")
    print("=" * 80)
    
    print("\nExpected values per docstring ('features on D predict returns from D+1 to D+1+h'):")
    print("For h=2, row D should have: price[D+3] / price[D+1] - 1")
    for i in range(len(prices)):
        if i + 3 < len(prices) and i + 1 < len(prices):
            ret = prices[i + 3] / prices[i + 1] - 1
            print(f"  Row {i} ({base_df['date'][i]}): {prices[i+3]}/{prices[i+1]} - 1 = {ret:.6f}")
        else:
            print(f"  Row {i} ({base_df['date'][i]}): null")
    
    print("\n" + "-" * 80)
    print("Test with different input orders:")
    print("-" * 80)
    
    # Test all orderings
    df_asc = base_df.sort("date")
    df_desc = base_df.sort("date", descending=True)
    df_shuffled = base_df.sample(fraction=1.0, seed=42)
    
    result_asc = correct_create_ranking_targets(df_asc, horizons=[2])
    result_desc = correct_create_ranking_targets(df_desc, horizons=[2])
    result_shuffled = correct_create_ranking_targets(df_shuffled, horizons=[2])
    
    # Compare all (sorted by date)
    comparison = (
        result_asc.sort("date").select(["date", "close", pl.col("fwd_return_2d").alias("from_asc")])
        .join(result_desc.sort("date").select(["date", pl.col("fwd_return_2d").alias("from_desc")]), on="date")
        .join(result_shuffled.sort("date").select(["date", pl.col("fwd_return_2d").alias("from_shuffled")]), on="date")
    )
    
    print("\nAll three should be IDENTICAL:")
    print(comparison)
    
    # Verify against expected
    print("\n" + "-" * 80)
    print("Verify against manual expected values:")
    print("-" * 80)
    
    expected = []
    for i in range(len(prices)):
        if i + 3 < len(prices) and i + 1 < len(prices):
            expected.append(round(prices[i + 3] / prices[i + 1] - 1, 6))
        else:
            expected.append(None)
    
    actual = [round(v, 6) if v is not None else None for v in result_asc.sort("date")["fwd_return_2d"].to_list()]
    
    print(f"Expected: {expected}")
    print(f"Actual:   {actual}")
    
    matches = all(
        (e is None and a is None) or (e is not None and a is not None and abs(e - a) < 1e-5)
        for e, a in zip(expected, actual)
    )
    print(f"\nMatches expected: {matches}")


if __name__ == "__main__":
    main()
