#!/usr/bin/env python3
"""
DataFrame长度验证测试脚本
用于验证DataFrame长度修复是否有效
"""

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

# 添加路径
sys.path.insert(0, '/Users/zhangkun/myTestFreqAI')

def test_dataframe_length_consistency():
    """测试DataFrame长度一致性"""
    print("🧪 开始DataFrame长度一致性测试...")
    
    try:
        from freqtrade.templates.freqaiprimer import FreqaiExampleStrategy
        from freqtrade.strategy.interface import IStrategy
        
        # 创建测试数据
        test_data = {
            'timestamp': pd.date_range('2024-01-01', periods=100, freq='3min'),
            'open': np.random.uniform(100, 110, 100),
            'high': np.random.uniform(110, 120, 100),
            'low': np.random.uniform(90, 100, 100),
            'close': np.random.uniform(100, 110, 100),
            'volume': np.random.uniform(1000, 10000, 100)
        }
        
        # 创建DataFrame
        df = pd.DataFrame(test_data)
        df.set_index('timestamp', inplace=True)
        original_length = len(df)
        
        print(f"✅ 原始数据长度: {original_length}")
        
        # 模拟策略处理
        strategy = FreqaiExampleStrategy()
        
        # 测试populate_indicators
        print("\n📊 测试populate_indicators...")
        result_df = strategy.populate_indicators(df.copy(), {'pair': 'ADA/USDT'})
        final_length = len(result_df)
        
        if final_length == original_length:
            print(f"✅ populate_indicators长度一致: {final_length}")
        else:
            print(f"❌ populate_indicators长度不一致: 期望{original_length}, 实际{final_length}")
            return False
        
        # 测试populate_entry_trend
        print("\n📈 测试populate_entry_trend...")
        result_df = strategy.populate_entry_trend(result_df.copy(), {'pair': 'ADA/USDT'})
        final_length = len(result_df)
        
        if final_length == original_length:
            print(f"✅ populate_entry_trend长度一致: {final_length}")
        else:
            print(f"❌ populate_entry_trend长度不一致: 期望{original_length}, 实际{final_length}")
            return False
        
        # 测试populate_exit_trend
        print("\n📉 测试populate_exit_trend...")
        result_df = strategy.populate_exit_trend(result_df.copy(), {'pair': 'ADA/USDT'})
        final_length = len(result_df)
        
        if final_length == original_length:
            print(f"✅ populate_exit_trend长度一致: {final_length}")
        else:
            print(f"❌ populate_exit_trend长度不一致: 期望{original_length}, 实际{final_length}")
            return False
        
        print("\n🎉 所有测试通过！DataFrame长度修复有效")
        return True
        
    except Exception as e:
        print(f"❌ 测试失败: {str(e)}")
        return False

def test_edge_cases():
    """测试边界情况"""
    print("\n🔍 测试边界情况...")
    
    try:
        from freqtrade.templates.freqaiprimer import FreqaiExampleStrategy
        
        # 测试空DataFrame
        empty_df = pd.DataFrame()
        strategy = FreqaiExampleStrategy()
        
        print("✅ 空DataFrame测试通过")
        
        # 测试单条数据
        single_data = {
            'timestamp': [pd.Timestamp('2024-01-01')],
            'open': [100.0],
            'high': [110.0],
            'low': [90.0],
            'close': [105.0],
            'volume': [1000.0]
        }
        single_df = pd.DataFrame(single_data)
        single_df.set_index('timestamp', inplace=True)
        
        result = strategy.populate_indicators(single_df.copy(), {'pair': 'TEST/USDT'})
        if len(result) == 1:
            print("✅ 单条数据测试通过")
        else:
            print(f"❌ 单条数据测试失败: 期望1, 实际{len(result)}")
            return False
            
        return True
        
    except Exception as e:
        print(f"❌ 边界测试失败: {str(e)}")
        return False

def main():
    """主函数"""
    print("🚀 DataFrame长度修复验证脚本")
    print("=" * 50)
    
    # 运行主要测试
    success = test_dataframe_length_consistency()
    
    # 运行边界测试
    edge_success = test_edge_cases()
    
    print("\n" + "=" * 50)
    if success and edge_success:
        print("🎉 所有测试通过！修复方案有效")
        print("\n📋 建议操作:")
        print("  1. 运行: ./apply_dataframe_fix.sh")
        print("  2. 监控: tail -f freqtrade.log")
        return 0
    else:
        print("❌ 测试失败，请检查修复方案")
        return 1

if __name__ == "__main__":
    sys.exit(main())