"""Test script to investigate potential bug in create_ranking_targets."""
import polars as pl
from crowdcent_challenge.scoring import create_ranking_targets


def trace_simple_example():
    """Create a simple example to trace through the logic step by step."""
    print("=" * 80)
    print("SIMPLE EXAMPLE - Single Ticker")
    print("=" * 80)
    
    # Create simple price data for one ticker
    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]
    })
    
    print("\nInput data:")
    print(df)
    
    # Calculate manually what the forward returns should be for h=2
    print("\n" + "-" * 80)
    print("MANUAL CALCULATION (h=2)")
    print("-" * 80)
    
    prices = df["close"].to_list()
    
    print("\nStep 1: h-day forward return sitting on day D")
    print("Formula: price[D+h] / price[D] - 1")
    for i in range(len(prices)):
        if i + 2 < len(prices):
            ret = prices[i + 2] / prices[i] - 1
            print(f"  Day {i}: price[{i+2}]/price[{i}] - 1 = {prices[i+2]}/{prices[i]} - 1 = {ret:.4f}")
        else:
            print(f"  Day {i}: null (no price[{i+2}])")
    
    print("\nStep 2: shift(1) - shift values DOWN (result[i] = value[i-1])")
    print("After shift(1), the return that was on row X moves to row X+1")
    
    # What the CURRENT code does:
    # - fwd_temp[i] = price[i+h]/price[i] - 1
    # - fwd_return[i] = fwd_temp[i-1] = price[i-1+h]/price[i-1] - 1
    
    print("\nWhat code produces (fwd_return[i] = fwd_temp[i-1]):")
    for i in range(len(prices)):
        if i == 0:
            print(f"  Day {i}: null (fwd_temp[-1] doesn't exist)")
        elif (i - 1) + 2 < len(prices):
            ret = prices[i - 1 + 2] / prices[i - 1] - 1
            print(f"  Day {i}: price[{i+1}]/price[{i-1}] - 1 = {prices[i+1]}/{prices[i-1]} - 1 = {ret:.4f}")
            print(f"         This is return from day {i-1} to day {i+1}")
        else:
            print(f"  Day {i}: null")
    
    print("\n" + "-" * 80)
    print("WHAT THE DOCSTRING SAYS WE WANT")
    print("-" * 80)
    print("""
    'Shift returns back by 1 day (so features on day D predict returns from D+1 to D+h+1)'
    
    So on row D, we should have: return from D+1 to D+1+h = price[D+1+h]/price[D+1] - 1
    """)
    
    print("What we SHOULD have (price[D+1+h]/price[D+1] - 1):")
    for i in range(len(prices)):
        if (i + 1 + 2) < len(prices) and (i + 1) < len(prices):
            ret = prices[i + 1 + 2] / prices[i + 1] - 1
            print(f"  Day {i}: price[{i+3}]/price[{i+1}] - 1 = {prices[i+3]}/{prices[i+1]} - 1 = {ret:.4f}")
            print(f"         This is return from day {i+1} to day {i+3}")
        else:
            print(f"  Day {i}: null")
    
    # Run actual function
    print("\n" + "-" * 80)
    print("ACTUAL FUNCTION OUTPUT")
    print("-" * 80)
    result = create_ranking_targets(
        df, 
        horizons=[2], 
        return_raw_returns=True,
        drop_incomplete=False
    )
    print(result.select(["date", "close", "fwd_return_2d", "target_2d"]))
    
    print("\n" + "=" * 80)
    print("BUG ANALYSIS")
    print("=" * 80)
    print("""
    On Day 1 (2024-01-02), features are calculated from Day 1 close.
    According to docstring, we should predict returns from Day 2 to Day 4.
    
    Expected fwd_return: price[4]/price[2] - 1 = 125/120 - 1 = 0.0417
    
    But looking at the actual output above, what does Day 1 actually show?
    """)


def trace_with_two_tickers():
    """Test cross-sectional ranking with two tickers."""
    print("\n" + "=" * 80)
    print("TWO TICKERS - Testing Cross-Sectional Ranking")
    print("=" * 80)
    
    # Create data for two tickers with known relative performance
    dates = pl.date_range(pl.date(2024, 1, 1), pl.date(2024, 1, 10), eager=True)
    
    df = pl.DataFrame({
        "date": list(dates) * 2,
        "ticker": ["BTC"] * 10 + ["ETH"] * 10,
        "close": (
            [100.0, 110.0, 120.0, 130.0, 140.0, 150.0, 160.0, 170.0, 180.0, 190.0] +  # BTC: steady 10% growth
            [100.0, 105.0, 108.0, 110.0, 112.0, 114.0, 116.0, 118.0, 120.0, 122.0]    # ETH: slower growth
        )
    }).sort(["date", "ticker"])
    
    print("\nInput data (first 6 rows per ticker):")
    print(df.head(12))
    
    result = create_ranking_targets(
        df, 
        horizons=[2], 
        return_raw_returns=True,
        drop_incomplete=False
    )
    
    print("\nResults with forward returns:")
    print(result.sort(["date", "ticker"]).select(["date", "ticker", "close", "fwd_return_2d", "target_2d"]))
    
    # Check a specific date
    print("\n" + "-" * 80)
    date_to_check = pl.date(2024, 1, 3)
    print(f"On {date_to_check}:")
    subset = result.filter(pl.col("date") == date_to_check)
    print(subset)
    
    print("""
    If BTC consistently outperforms ETH:
    - BTC should have higher fwd_return_2d
    - BTC should have higher target_2d (closer to 1)
    - ETH should have lower target_2d (closer to 0)
    """)


def verify_alignment():
    """Verify the return actually aligns with the right trading day."""
    print("\n" + "=" * 80)
    print("RETURN ALIGNMENT VERIFICATION")
    print("=" * 80)
    
    # Simple increasing prices where we can verify returns exactly
    df = pl.DataFrame({
        "date": pl.date_range(pl.date(2024, 1, 1), pl.date(2024, 1, 8), eager=True),
        "ticker": ["TEST"] * 8,
        # Day 0=100, Day 1=102 (2%), Day 2=104 (2%), Day 3=110 (5.77%), etc.
        "close": [100.0, 102.0, 104.0, 110.0, 115.0, 120.0, 130.0, 140.0]
    })
    
    result = create_ranking_targets(
        df,
        horizons=[2],
        return_raw_returns=True,
        drop_incomplete=False
    )
    
    print("Data with calculated returns:")
    result_with_dates = result.with_columns([
        pl.lit("").alias("interpretation")
    ])
    
    print(result.select(["date", "close", "fwd_return_2d"]))
    
    print("\nLet's verify each row:")
    for row in result.iter_rows(named=True):
        date = row["date"]
        ret = row["fwd_return_2d"]
        if ret is not None:
            # What return SHOULD be there per the docstring?
            # Features on day D predict returns from D+1 to D+1+h
            print(f"\nDate {date}:")
            print(f"  fwd_return_2d = {ret:.6f}")
            # Find what this return actually represents
            # If shift(1) moves values DOWN, then the return on row i is fwd_temp[i-1]
            # which is price[i-1+h]/price[i-1] - 1


if __name__ == "__main__":
    trace_simple_example()
    trace_with_two_tickers()
    verify_alignment()
