"""Test: sort descending + shift(1) - does it work?"""
import polars as pl


def sort_desc_then_shift1(df, h=2):
    """Sort descending, then use shift(1)."""
    result_df = df.sort(["ticker", "date"], descending=[False, True])  # date descending
    
    fwd_temp = f"_fwd_{h}d_temp"
    fwd_return = f"fwd_return_{h}d"
    
    result_df = (
        result_df
        .with_columns(
            (pl.col("close").shift(-h).over("ticker") / pl.col("close") - 1).alias(fwd_temp)
        )
        .with_columns(pl.col(fwd_temp).shift(1).over("ticker").alias(fwd_return))
        .drop(fwd_temp)
    )
    return result_df


def sort_asc_then_shift_neg1(df, h=2):
    """Sort ascending, then use shift(-1) - the "correct" way."""
    result_df = df.sort(["ticker", "date"])  # date ascending
    
    fwd_temp = f"_fwd_{h}d_temp"
    fwd_return = f"fwd_return_{h}d"
    
    result_df = (
        result_df
        .with_columns(
            (pl.col("close").shift(-h).over("ticker") / pl.col("close") - 1).alias(fwd_temp)
        )
        .with_columns(pl.col(fwd_temp).shift(-1).over("ticker").alias(fwd_return))
        .drop(fwd_temp)
    )
    return result_df


def main():
    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("TEST: Sort DESCENDING + shift(1) vs Sort ASCENDING + shift(-1)")
    print("=" * 80)
    
    print("\nExpected (features on D predict returns D+1 to D+3):")
    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"  {base_df['date'][i]}: {ret:.6f}")
        else:
            print(f"  {base_df['date'][i]}: null")
    
    result_desc = sort_desc_then_shift1(base_df, h)
    result_asc = sort_asc_then_shift_neg1(base_df, h)
    
    print("\n" + "-" * 80)
    print("Sort DESC + shift(1):")
    print(result_desc.sort("date").select(["date", "close", "fwd_return_2d"]))
    
    print("\nSort ASC + shift(-1):")
    print(result_asc.sort("date").select(["date", "close", "fwd_return_2d"]))
    
    # Are they the same?
    desc_vals = result_desc.sort("date")["fwd_return_2d"].to_list()
    asc_vals = result_asc.sort("date")["fwd_return_2d"].to_list()
    
    def vals_equal(a, b):
        for x, y in zip(a, b):
            if x is None and y is None:
                continue
            if x is None or y is None:
                return False
            if abs(x - y) > 1e-10:
                return False
        return True
    
    print(f"\nBoth methods equal? {vals_equal(desc_vals, asc_vals)}")


if __name__ == "__main__":
    main()
