#!/usr/bin/env python3
"""Test script to verify extractor consistency."""

import pandas as pd
import numpy as np
from reaction_features import ReactionFeatureExtractor
from content_features import ContentFeatureExtractor
from bot_features import BotFeatureExtractor
from engagement_features import EngagementFeatureCreator

def test_extractors():
    """Test that all extractors return consistent tuple format."""
    print("Testing extractor return patterns...")
    
    # Create dummy data
    user_ids = pd.Series([1, 2, 3, 4, 5])
    
    # Test reaction extractor
    print("\n1. Testing ReactionFeatureExtractor...")
    reaction_df = pd.DataFrame({
        'user_id': [1, 2, 3, 1, 2],
        'reaction_type': ['L', 'D', 'L', 'L', 'D'],
        'flagged': [0, 0, 1, 0, 0]
    })
    
    reaction_extractor = ReactionFeatureExtractor()
    result = reaction_extractor.extract_features(reaction_df, user_ids, verbose=False)
    print(f"   Returns: {type(result)} with {len(result)} elements")
    if isinstance(result, tuple):
        print(f"   Element 1: {type(result[0])} with shape {result[0].shape if hasattr(result[0], 'shape') else 'N/A'}")
        print(f"   Element 2: {type(result[1])}")
    
    # Test content extractor
    print("\n2. Testing ContentFeatureExtractor...")
    total_clip_df = pd.DataFrame({
        'user_id': [1, 2, 3, 1, 2],
        'id': [101, 102, 103, 104, 105],
        'model_name': ['model1', 'model2', 'v4p5', 'model1', 'v4p5'],
        'created_at': pd.date_range('2025-01-01', periods=5),
        'type': ['public', 'private', 'public', 'public', 'private']
    })
    
    content_extractor = ContentFeatureExtractor()
    result = content_extractor.extract_features(total_clip_df, user_ids, verbose=False)
    print(f"   Returns: {type(result)} with {len(result)} elements")
    if isinstance(result, tuple):
        print(f"   Element 1: {type(result[0])} with shape {result[0].shape if hasattr(result[0], 'shape') else 'N/A'}")
        print(f"   Element 2: {type(result[1])}")
    
    # Test bot extractor
    print("\n3. Testing BotFeatureExtractor...")
    bots_action_df = pd.DataFrame({
        'clip_id': [101, 102, 103],
        'download_audio_count': [1, 0, 2],
        'share_count': [0, 1, 1]
    })
    
    features_df = pd.DataFrame({
        'user_id': user_ids,
        'total_clips_created': [10, 20, 30, 40, 50]
    })
    
    bot_extractor = BotFeatureExtractor()
    result = bot_extractor.extract_features(
        bots_action_df, total_clip_df, user_ids, features_df, verbose=False
    )
    print(f"   Returns: {type(result)} with {len(result)} elements")
    if isinstance(result, tuple):
        print(f"   Element 1: {type(result[0])} with shape {result[0].shape if hasattr(result[0], 'shape') else 'N/A'}")
        print(f"   Element 2: {type(result[1])}")
    
    # Test engagement creator
    print("\n4. Testing EngagementFeatureCreator...")
    features_df = pd.DataFrame({
        'user_id': user_ids,
        'total_clips_created': [10, 20, 30, 40, 50],
        'reaction_frequency': [1, 2, 3, 4, 5],
        'total_bot_actions': [5, 10, 15, 20, 25],
        'subscription_tier': [1, 2, 1, 2, 3],
        'is_recent_creator': [1, 1, 0, 1, 1]
    })
    
    engagement_creator = EngagementFeatureCreator()
    result = engagement_creator.create_engagement_features(features_df, verbose=False)
    print(f"   Returns: {type(result)} with {len(result)} elements")
    if isinstance(result, tuple):
        print(f"   Element 1: {type(result[0])} with shape {result[0].shape if hasattr(result[0], 'shape') else 'N/A'}")
        print(f"   Element 2: {type(result[1])}")
    
    print("\n✅ All extractors tested successfully!")
    print("\nExpected pattern: All should return tuple of (DataFrame, dict)")

if __name__ == "__main__":
    test_extractors()