#!/usr/bin/env python3
"""
Test Fixed Input Size

Verify that the CNN model now receives consistent input dimensions
"""

import numpy as np
from datetime import datetime
from core.data_models import BaseDataInput, OHLCVBar
from core.enhanced_cnn_adapter import EnhancedCNNAdapter

def create_test_data(with_cob=True, with_indicators=True):
    """Create test BaseDataInput with varying data completeness"""
    
    # Create basic OHLCV data
    ohlcv_bars = []
    for i in range(100):  # Less than 300 to test padding
        bar = OHLCVBar(
            symbol="ETH/USDT",
            timestamp=datetime.now(),
            open=100.0 + i,
            high=101.0 + i,
            low=99.0 + i,
            close=100.5 + i,
            volume=1000 + i,
            timeframe="1s"
        )
        ohlcv_bars.append(bar)
    
    # Create test data
    base_data = BaseDataInput(
        symbol="ETH/USDT",
        timestamp=datetime.now(),
        ohlcv_1s=ohlcv_bars,
        ohlcv_1m=ohlcv_bars[:50],  # Even less data
        ohlcv_1h=ohlcv_bars[:20],
        ohlcv_1d=ohlcv_bars[:10],
        btc_ohlcv_1s=ohlcv_bars[:80],  # Incomplete BTC data
        technical_indicators={'rsi': 50.0, 'macd': 0.1} if with_indicators else {},
        last_predictions={}
    )
    
    # Add COB data if requested (simplified for testing)
    if with_cob:
        # Create a simple mock COB data object
        class MockCOBData:
            def __init__(self):
                self.price_buckets = {
                    2500.0: {'bid_volume': 100, 'ask_volume': 90, 'total_volume': 190, 'imbalance': 0.05},
                    2501.0: {'bid_volume': 80, 'ask_volume': 120, 'total_volume': 200, 'imbalance': -0.2}
                }
                self.ma_1s_imbalance = {2500.0: 0.1, 2501.0: -0.1}
                self.ma_5s_imbalance = {2500.0: 0.05, 2501.0: -0.05}
        
        base_data.cob_data = MockCOBData()
    
    return base_data

def test_consistent_feature_size():
    """Test that feature vectors are always the same size"""
    print("=== Testing Consistent Feature Size ===")
    
    # Test different data scenarios
    scenarios = [
        ("Full data", True, True),
        ("No COB data", False, True),
        ("No indicators", True, False),
        ("Minimal data", False, False)
    ]
    
    feature_sizes = []
    
    for name, with_cob, with_indicators in scenarios:
        base_data = create_test_data(with_cob, with_indicators)
        features = base_data.get_feature_vector()
        
        print(f"{name}: {len(features)} features")
        feature_sizes.append(len(features))
    
    # Check if all sizes are the same
    if len(set(feature_sizes)) == 1:
        print(f"✅ All feature vectors have consistent size: {feature_sizes[0]}")
        return feature_sizes[0]
    else:
        print(f"❌ Inconsistent feature sizes: {feature_sizes}")
        return None

def test_cnn_adapter():
    """Test that CNN adapter works with fixed input size"""
    print("\n=== Testing CNN Adapter ===")
    
    try:
        # Create CNN adapter
        adapter = EnhancedCNNAdapter()
        print(f"CNN model initialized with feature_dim: {adapter.model.feature_dim}")
        
        # Test with different data scenarios
        scenarios = [
            ("Full data", True, True),
            ("No COB data", False, True),
            ("Minimal data", False, False)
        ]
        
        for name, with_cob, with_indicators in scenarios:
            try:
                base_data = create_test_data(with_cob, with_indicators)
                
                # Make prediction
                result = adapter.predict(base_data)
                
                print(f"✅ {name}: Prediction successful - {result.action} (conf={result.confidence:.3f})")
                
            except Exception as e:
                print(f"❌ {name}: Prediction failed - {e}")
        
        return True
        
    except Exception as e:
        print(f"❌ CNN adapter initialization failed: {e}")
        return False

def test_no_network_rebuilding():
    """Test that network doesn't rebuild during runtime"""
    print("\n=== Testing No Network Rebuilding ===")
    
    try:
        adapter = EnhancedCNNAdapter()
        original_feature_dim = adapter.model.feature_dim
        
        print(f"Original feature_dim: {original_feature_dim}")
        
        # Make multiple predictions with different data
        for i in range(5):
            base_data = create_test_data(with_cob=(i % 2 == 0), with_indicators=(i % 3 == 0))
            
            try:
                result = adapter.predict(base_data)
                current_feature_dim = adapter.model.feature_dim
                
                if current_feature_dim != original_feature_dim:
                    print(f"❌ Network was rebuilt! Original: {original_feature_dim}, Current: {current_feature_dim}")
                    return False
                
                print(f"✅ Prediction {i+1}: No rebuilding, feature_dim stable at {current_feature_dim}")
                
            except Exception as e:
                print(f"❌ Prediction {i+1} failed: {e}")
                return False
        
        print("✅ Network architecture remained stable throughout all predictions")
        return True
        
    except Exception as e:
        print(f"❌ Test failed: {e}")
        return False

def main():
    """Run all tests"""
    print("=== Fixed Input Size Test Suite ===\n")
    
    # Test 1: Consistent feature size
    fixed_size = test_consistent_feature_size()
    
    if fixed_size:
        # Test 2: CNN adapter works
        adapter_works = test_cnn_adapter()
        
        if adapter_works:
            # Test 3: No network rebuilding
            no_rebuilding = test_no_network_rebuilding()
            
            if no_rebuilding:
                print("\n✅ ALL TESTS PASSED!")
                print("✅ Feature vectors have consistent size")
                print("✅ CNN adapter works with fixed input")
                print("✅ No runtime network rebuilding")
                print(f"✅ Fixed feature size: {fixed_size}")
            else:
                print("\n❌ Network rebuilding test failed")
        else:
            print("\n❌ CNN adapter test failed")
    else:
        print("\n❌ Feature size consistency test failed")

if __name__ == "__main__":
    main()