"""Demonstrate the bug and the fix in create_ranking_targets."""
import polars as pl
import narwhals as nw


def current_implementation(df, h=2):
    """Current (buggy) implementation."""
    result_df = df.clone()
    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))  # BUG: shift(1)
        .drop(fwd_temp)
    )
    return result_df


def fixed_implementation(df, h=2):
    """Fixed implementation - shift(-1) instead of shift(1)."""
    result_df = df.clone()
    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))  # FIX: shift(-1)
        .drop(fwd_temp)
    )
    return result_df


def expected_return_from_d_to_d_plus_h(df, h=2):
    """Direct calculation: return from D to D+h, placed on row D-1."""
    # This is what the docstring says we want:
    # Features on day D-1 predict returns from D to D+h
    # Equivalently: Features on day D predict returns from D+1 to D+1+h
    
    result_df = df.clone()
    fwd_return = f"fwd_return_{h}d"
    
    # price[D+h+1] / price[D+1] - 1, on row D
    # = price shifted by -(h+1) / price shifted by -1 - 1
    result_df = result_df.with_columns(
        (
            pl.col("close").shift(-(h+1)).over("ticker") 
            / pl.col("close").shift(-1).over("ticker") 
            - 1
        ).alias(fwd_return)
    )
    return result_df


def main():
    # Create test data
    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 = df["close"].to_list()
    h = 2
    
    print("=" * 90)
    print("DEMONSTRATION: Bug in create_ranking_targets")
    print("=" * 90)
    
    print("\nTest data:")
    print(df)
    
    print("\n" + "-" * 90)
    print("Manual calculation of EXPECTED returns")
    print("-" * 90)
    print("""
    Per docstring: 'features on day D predict returns from D+1 to D+h+1'
    For h=2, on row D we want: price[D+3] / price[D+1] - 1
    """)
    
    print("Expected fwd_return_2d for each row:")
    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} ({df['date'][i]}): price[{i+3}]/price[{i+1}] - 1 = {prices[i+3]}/{prices[i+1]} - 1 = {ret:.6f}")
        else:
            print(f"  Row {i} ({df['date'][i]}): null")
    
    print("\n" + "-" * 90)
    print("COMPARISON: Current vs Fixed vs Expected")
    print("-" * 90)
    
    current = current_implementation(df, h)
    fixed = fixed_implementation(df, h)
    expected = expected_return_from_d_to_d_plus_h(df, h)
    
    comparison = df.select(["date", "close"]).with_columns([
        current["fwd_return_2d"].alias("current_buggy"),
        fixed["fwd_return_2d"].alias("fixed"),
        expected["fwd_return_2d"].alias("expected"),
    ])
    
    print("\n", comparison)
    
    print("\n" + "-" * 90)
    print("VERIFICATION")
    print("-" * 90)
    
    # Check if fixed matches expected
    fixed_vals = [v for v in fixed["fwd_return_2d"].to_list() if v is not None]
    expected_vals = [v for v in expected["fwd_return_2d"].to_list() if v is not None]
    
    print(f"\nFixed values: {[round(v, 6) for v in fixed_vals]}")
    print(f"Expected values: {[round(v, 6) for v in expected_vals]}")
    
    # They should match
    matches = all(
        abs(f - e) < 1e-10 
        for f, e in zip(fixed_vals, expected_vals)
    )
    print(f"\nFixed matches expected: {matches}")
    
    print("\n" + "=" * 90)
    print("CONCLUSION")
    print("=" * 90)
    print("""
    The bug is in line 556 of scoring.py:
    
        .with_columns(nw.col(fwd_temp).shift(1).over(ticker_col).alias(fwd_return))
                                       ^^^^^^^^
    Should be:
    
        .with_columns(nw.col(fwd_temp).shift(-1).over(ticker_col).alias(fwd_return))
                                       ^^^^^^^^^
    
    shift(1) moves values DOWN (result[i] = value[i-1]), giving us old returns.
    shift(-1) moves values UP (result[i] = value[i+1]), giving us future returns.
    
    The current code attaches the return from day D-1 to D-1+h on row D.
    The fix attaches the return from day D+1 to D+1+h on row D (as documented).
    """)


if __name__ == "__main__":
    main()
