#!/usr/bin/env python3
"""
AZR Integration Test Script

AZR REINFORCE++ 학습 통합이 제대로 작동하는지 테스트
"""

import os
import sys
import pandas as pd
import torch

# 경로 설정
sys.path.append('/home/ubuntu/RLVR/TestTime-RLVR-v2')
sys.path.append('/home/ubuntu/RLVR/TestTime-RLVR-v2/test')

from utils.azr_trainer_integration import AZRTrainerIntegration
from absolute_zero_reasoner.testtime.logger import TestTimeLogger


def test_data_loading():
    """Parquet 데이터 로딩 테스트"""
    print("=== Testing Data Loading ===")
    
    # 실제 데이터 경로
    data_path = "/home/ubuntu/RLVR/TestTime-RLVR-v2/tmp/batch_results/ttrlvr_azr_20250729_141828/humaneval/HumanEval_1/round_5/azr_training_data"
    
    # 각 파일 확인
    for task_type in ['induction', 'deduction', 'abduction']:
        file_path = os.path.join(data_path, f"{task_type}.parquet")
        if os.path.exists(file_path):
            df = pd.read_parquet(file_path)
            print(f"✅ {task_type}: {len(df)} rows")
            
            # 첫 번째 행의 reward 확인
            if len(df) > 0:
                first_reward = df.iloc[0]['basic_accuracy']
                print(f"   First row reward: {first_reward}")
        else:
            print(f"❌ {task_type}: File not found")
    
    print()


def test_azr_trainer_init():
    """AZR Trainer 초기화 테스트"""
    print("=== Testing AZR Trainer Initialization ===")
    
    try:
        # 더미 토크나이저
        from transformers import AutoTokenizer
        tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B")
        
        # 로거
        logger = TestTimeLogger(log_dir="/tmp/test_logs")
        
        # AZR Trainer 초기화
        trainer = AZRTrainerIntegration(
            model_path="Qwen/Qwen2.5-7B",
            tokenizer=tokenizer,
            logger=logger,
            gpu_id=0,
            num_cpus=8
        )
        
        print("✅ AZR Trainer initialized successfully")
        
        # 설정 생성 테스트
        config = trainer.create_azr_config(round_num=1, batch_size=24)
        print(f"✅ Config created: {config.trainer.experiment_name}")
        
        # 정리
        trainer.cleanup()
        print("✅ Cleanup completed")
        
    except Exception as e:
        print(f"❌ Error: {e}")
        import traceback
        traceback.print_exc()
    
    print()


def test_data_conversion():
    """데이터 변환 테스트"""
    print("=== Testing Data Conversion ===")
    
    try:
        from transformers import AutoTokenizer
        tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B")
        logger = TestTimeLogger(log_dir="/tmp/test_logs")
        
        trainer = AZRTrainerIntegration(
            model_path="Qwen/Qwen2.5-7B",
            tokenizer=tokenizer,
            logger=logger,
            gpu_id=0,
            num_cpus=8
        )
        
        # 데이터 로드 및 변환
        data_path = "/home/ubuntu/RLVR/TestTime-RLVR-v2/tmp/batch_results/ttrlvr_azr_20250729_141828/humaneval/HumanEval_1/round_5/azr_training_data"
        data_protos, stats = trainer.load_and_prepare_data(data_path)
        
        print(f"✅ Loaded {len(data_protos)} data protos")
        print(f"   Stats: {dict(stats)}")
        
        # 배치 생성 테스트
        if data_protos:
            batches = trainer.prepare_batches(data_protos, batch_size=24)
            print(f"✅ Created {len(batches)} batches")
            
            if batches:
                first_batch_size = len(batches[0])
                print(f"   First batch size: {first_batch_size}")
        
        trainer.cleanup()
        
    except Exception as e:
        print(f"❌ Error: {e}")
        import traceback
        traceback.print_exc()
    
    print()


def main():
    """메인 테스트 함수"""
    print("🧪 Starting AZR Integration Tests")
    print("=" * 60)
    
    # 환경 설정
    os.environ['CUDA_VISIBLE_DEVICES'] = '4'
    
    # 테스트 실행
    test_data_loading()
    test_azr_trainer_init()
    test_data_conversion()
    
    print("✅ Tests completed!")


if __name__ == '__main__':
    main()