import pytest
import pandas as pd
import numpy as np
import copy
import sys
import os

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from backtest import expanding_window_backtest
from core_types import PortfolioState
from solver import build_and_optimize

def mock_expanding_window_backtest(returns_df, spy_rets, capital, rfr, cfg, model, allocation_engine, spread_map, initial_train_days=10, rebalance_freq=5, ff_df=None, yield_df=None):
    return expanding_window_backtest(returns_df, spy_rets, capital, rfr, cfg, model, allocation_engine, spread_map, initial_train_days, rebalance_freq, ff_df, yield_df)

def test_delisting_terminal_returns():
    """
    Simulates a delisting event where a stock drops to 0.0, causing a -100% return
    and verifying the backtest correctly carries the total loss for that asset.
    """
    dates = pd.date_range("2020-01-01", periods=30, freq="B")
    # Ticker A is stable. Ticker B drops to 0.0 at period 20.
    
    # We create raw prices to ensure compounding works exactly
    prices_A = np.linspace(100, 130, 30)
    prices_B = np.linspace(50, 60, 20).tolist() + [0.0] * 10
    
    prices_df = pd.DataFrame({'A': prices_A, 'B': prices_B}, index=dates)
    
    # Calculate returns
    returns_df = prices_df.pct_change()
    
    # In core_engine.py we pad delisted tickers with 0.0 returns after the -1.0
    # Let's manually do that here to simulate core_engine's behavior
    returns_df.loc[dates[20], 'B'] = -1.0 # The drop
    returns_df.loc[dates[21]:, 'B'] = 0.0 # The padding
    
    returns_df = returns_df.dropna() # Drop first day NaN
    
    spy_rets = pd.Series(0.001, index=returns_df.index)
    
    cfg = {"transaction_cost": 0.0, "trading_days_per_year": 252, "tax_enabled": False}
    spread_map = {"A": 0.0, "B": 0.0}
    
    # Run backtest
    # We use a dummy model that allocates 50/50 initially
    
    # Mocking build_and_optimize to return 50/50
    class MockOptRes:
        weights = pd.Series({"A": 0.5, "B": 0.5, "CASH": 0.0})
    
    import solver
    original_build = solver.build_and_optimize
    solver.build_and_optimize = lambda *args, **kwargs: MockOptRes()
    
    try:
        eq, spy_eq = mock_expanding_window_backtest(
            returns_df, spy_rets, capital=100_000.0, rfr=0.0, cfg=cfg, 
            model=1, allocation_engine=1, spread_map=spread_map,
            initial_train_days=10, rebalance_freq=5
        )
    finally:
        solver.build_and_optimize = original_build
        
    assert eq is not None
    assert len(eq) > 0
    
    # The portfolio had a 50% allocation to B at some point.
    # When B drops to 0.0 (return = -1.0), the portfolio should take a massive hit, but not go negative.
    
    # Ensure no NaNs or infs in the equity curve
    assert not eq.isna().any()
    assert not np.isinf(eq).any()
    
    # Portfolio shouldn't drop below 0
    assert eq.min() >= 0.0