#!/usr/bin/env python3
"""
Test script to verify pta_reload indicators work with NumPy 2+ and Pandas 2+
"""

import numpy as np
import pandas as pd
import sys
import traceback
from pathlib import Path

# Add pta_reload to path
sys.path.insert(0, str(Path(__file__).parent))

def create_test_data():
    """Create sample OHLCV data for testing"""
    np.random.seed(42)
    n = 100
    
    # Generate realistic price data
    close_prices = np.random.normal(100, 5, n).cumsum() + 100
    close_prices = np.maximum(close_prices, 50)  # Minimum price of 50
    
    # Generate OHLC data
    high = close_prices + np.random.uniform(0, 5, n)
    low = close_prices - np.random.uniform(0, 5, n)
    open_prices = close_prices + np.random.uniform(-2, 2, n)
    volume = np.random.randint(1000, 10000, n)
    
    # Create DataFrame
    dates = pd.date_range('2024-01-01', periods=n, freq='D')
    
    df = pd.DataFrame({
        'Open': open_prices,
        'High': high,
        'Low': low,
        'Close': close_prices,
        'Volume': volume
    }, index=dates)
    
    return df

def test_indicators():
    """Test all available indicators"""
    print(f"NumPy version: {np.__version__}")
    print(f"Pandas version: {pd.__version__}")
    print("=" * 50)
    
    try:
        from pta_reload import ta
        print("✓ Successfully imported pta_reload")
    except Exception as e:
        print(f"✗ Failed to import pta_reload: {e}")
        return False
    
    # Create test data
    df = create_test_data()
    print(f"✓ Created test data: {len(df)} rows")
    
    # Get all available indicators
    indicators = ta.list_indicators()
    print(f"✓ Found {len(indicators)} indicators to test")
    
    success_count = 0
    failure_count = 0
    failed_indicators = []
    
    for indicator_name in indicators:
        try:
            print(f"\nTesting {indicator_name}...", end="")
            
            # Get the indicator function
            indicator_func = getattr(ta, indicator_name)
            
            # Try to get some metadata about the function
            try:
                info = ta.get_indicator_info(indicator_name)
                category = info.get('category', 'unknown')
            except:
                category = 'unknown'
            
            # Test with different input combinations based on category
            result = None
            
            if category == 'overlap':
                # Most overlap indicators need just close price
                result = indicator_func(df['Close'])
            elif category == 'momentum':
                # Try close price first
                try:
                    result = indicator_func(df['Close'])
                except:
                    # Some might need high/low
                    try:
                        result = indicator_func(df['High'], df['Low'], df['Close'])
                    except:
                        result = indicator_func(df['Close'], df['Volume'])
            elif category == 'volatility':
                # Usually need high, low, close
                try:
                    result = indicator_func(df['High'], df['Low'], df['Close'])
                except:
                    result = indicator_func(df['Close'])
            elif category == 'volume':
                # Usually need close and volume
                try:
                    result = indicator_func(df['Close'], df['Volume'])
                except:
                    result = indicator_func(df['Volume'])
            elif category == 'trend':
                # Various combinations
                try:
                    result = indicator_func(df['High'], df['Low'], df['Close'])
                except:
                    try:
                        result = indicator_func(df['Close'])
                    except:
                        result = indicator_func(df['High'], df['Low'])
            else:
                # Default: try close price
                result = indicator_func(df['Close'])
            
            # Verify result
            if result is not None:
                if isinstance(result, pd.Series):
                    if len(result) > 0:
                        print(f" ✓ ({category})")
                        success_count += 1
                    else:
                        print(f" ⚠ Empty result ({category})")
                        failure_count += 1
                        failed_indicators.append(indicator_name)
                elif isinstance(result, pd.DataFrame):
                    if len(result) > 0 and len(result.columns) > 0:
                        print(f" ✓ ({category}) - {len(result.columns)} columns")
                        success_count += 1
                    else:
                        print(f" ⚠ Empty DataFrame ({category})")
                        failure_count += 1
                        failed_indicators.append(indicator_name)
                else:
                    print(f" ✓ ({category}) - {type(result)}")
                    success_count += 1
            else:
                print(f" ✗ Returned None ({category})")
                failure_count += 1
                failed_indicators.append(indicator_name)
                
        except Exception as e:
            print(f" ✗ Error: {str(e)[:50]}...")
            failure_count += 1
            failed_indicators.append(indicator_name)
    
    # Summary
    print("\n" + "=" * 50)
    print("TEST SUMMARY")
    print("=" * 50)
    print(f"Total indicators tested: {len(indicators)}")
    print(f"Successful: {success_count}")
    print(f"Failed: {failure_count}")
    print(f"Success rate: {(success_count / len(indicators) * 100):.1f}%")
    
    if failed_indicators:
        print("\nFailed indicators:")
        for indicator in failed_indicators:
            print(f"  - {indicator}")
    
    return failure_count == 0

def test_specific_compatibility():
    """Test specific NumPy/Pandas compatibility scenarios"""
    print("\n" + "=" * 50)
    print("COMPATIBILITY TESTS")
    print("=" * 50)
    
    df = create_test_data()
    
    try:
        from pta_reload import ta
        
        # Test 1: SMA with different pandas/numpy interactions
        print("Test 1: SMA indicator...")
        sma_result = ta.sma(df['Close'], length=10)
        assert isinstance(sma_result, pd.Series)
        assert not sma_result.empty
        print("✓ SMA works correctly")
        
        # Test 2: RSI with fillna handling
        print("Test 2: RSI with fill options...")
        rsi_result = ta.rsi(df['Close'], length=14, fillna=0)
        assert isinstance(rsi_result, pd.Series)
        print("✓ RSI with fillna works correctly")
        
        # Test 3: MACD with complex returns
        print("Test 3: MACD indicator...")
        macd_result = ta.macd(df['Close'])
        assert isinstance(macd_result, pd.DataFrame)
        assert len(macd_result.columns) >= 2  # Should have multiple columns
        print("✓ MACD works correctly")
        
        # Test 4: Bollinger Bands
        print("Test 4: Bollinger Bands...")
        bb_result = ta.bbands(df['Close'])
        assert isinstance(bb_result, pd.DataFrame)
        print("✓ Bollinger Bands work correctly")
        
        print("\n✓ All compatibility tests passed!")
        return True
        
    except Exception as e:
        print(f"\n✗ Compatibility test failed: {e}")
        traceback.print_exc()
        return False

if __name__ == "__main__":
    print("Testing pta_reload compatibility with NumPy 2+ and Pandas 2+")
    print("=" * 60)
    
    # Run basic indicator tests
    basic_success = test_indicators()
    
    # Run specific compatibility tests  
    compat_success = test_specific_compatibility()
    
    print("\n" + "=" * 60)
    if basic_success and compat_success:
        print("🎉 ALL TESTS PASSED! pta_reload is compatible with NumPy 2+ and Pandas 2+")
        sys.exit(0)
    else:
        print("❌ SOME TESTS FAILED! Review the output above for details")
        sys.exit(1)