#!/usr/bin/env python3
"""
Test script to verify NaN handling in feature extractors.
Tests edge cases that previously produced NaN values.
"""

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

print("🧪 TESTING NaN HANDLING IN FEATURE EXTRACTORS")
print("=" * 50)

# Create test data with edge cases
test_users = pd.DataFrame({
    'user_id': [1, 2, 3, 4, 5],  # Various edge case users
})

# User 1: Normal user with activity
# User 2: User with no clips
# User 3: User with no reactions
# User 4: User with no bot actions
# User 5: Brand new user (no activity)

# Create test dataframes
total_clip_df = pd.DataFrame({
    'user_id': [1, 1, 1, 3, 3],  # Users 2, 4, 5 have no clips
    'id': [101, 102, 103, 301, 302],
    'created_at': pd.to_datetime(['2024-01-01', '2024-01-15', '2024-02-01', '2024-01-10', '2024-01-20']),
    'model_name': ['v4', 'v4p5', 'v4', 'v3', 'v4'],
    'task': ['create', 'create', 'remix', 'create', 'create'],
    'type': ['public', 'public', 'private', 'public', 'public']
})

reaction_df = pd.DataFrame({
    'user_id': [1, 1, 2, 2, 4],  # Users 3, 5 have no reactions
    'reaction_type': ['L', 'D', 'L', 'L', 'L'],
    'play_count': [10, 5, 20, 15, 8]
})

bots_action_df = pd.DataFrame({
    'clip_id': [101, 101, 102, 301],  # No actions for user 4, 5's clips
    'download_audio_count': [1, 0, 1, 0],
    'download_video_count': [0, 1, 0, 1],
    'share_count': [1, 0, 2, 0]
})

discord_info_df = pd.DataFrame({
    'user_id': [1, 2, 3, 4, 5],
    'subscription_status': ['active', 'free', 'past_due', 'free', 'free']
})

# Initialize extractors
user_selector = UserSelector()
content_extractor = ContentFeatureExtractor()
reaction_extractor = ReactionFeatureExtractor()
bot_extractor = BotFeatureExtractor()
engagement_creator = EngagementFeatureCreator()

print("\n1️⃣ Testing UserSelector...")
features_df = user_selector.create_initial_features(
    total_clip_df, discord_info_df, verbose=False
)
print(f"   ✅ Created features for {len(features_df)} users")
print(f"   ✅ No NaN in subscription_tier: {not features_df['subscription_tier'].isna().any()}")

print("\n2️⃣ Testing ContentFeatureExtractor...")
content_features, _ = content_extractor.extract_features(
    total_clip_df, features_df['user_id'], verbose=False
)
features_df = content_extractor.merge_features(features_df, content_features)
print(f"   ✅ Content features shape: {content_features.shape}")
print(f"   ✅ No NaN in clip_creation_rate: {not features_df['clip_creation_rate'].isna().any()}")
print(f"   ✅ No NaN in daily_generation_rate: {not features_df['daily_generation_rate'].isna().any()}")

print("\n3️⃣ Testing ReactionFeatureExtractor...")
reaction_features, _ = reaction_extractor.extract_features(
    reaction_df, features_df['user_id'], verbose=False
)
features_df = reaction_extractor.merge_features(features_df, reaction_features)
print(f"   ✅ Reaction features shape: {reaction_features.shape}")
print(f"   ✅ No NaN in reaction_like_ratio: {not features_df['reaction_like_ratio'].isna().any()}")
print(f"   ✅ No NaN in engagement_ratio: {not features_df['engagement_ratio'].isna().any()}")

print("\n4️⃣ Testing BotFeatureExtractor...")
bot_features, _ = bot_extractor.extract_features(
    bots_action_df, total_clip_df, features_df['user_id'], features_df, verbose=False
)
features_df = bot_extractor.merge_features(features_df, bot_features)
print(f"   ✅ Bot features shape: {bot_features.shape}")
print(f"   ✅ No NaN in bot_action_rate: {not features_df['bot_action_rate'].isna().any()}")
print(f"   ✅ No NaN in audio_preference_ratio: {not features_df['audio_preference_ratio'].isna().any()}")

print("\n5️⃣ Testing EngagementFeatureCreator...")
features_df, _ = engagement_creator.create_engagement_features(
    features_df, verbose=False
)
print(f"   ✅ No NaN in creator_consumer_ratio: {not features_df['creator_consumer_ratio'].isna().any()}")
print(f"   ✅ No NaN in engagement_efficiency: {not features_df['engagement_efficiency'].isna().any()}")
print(f"   ✅ No NaN in viral_potential: {not features_df['viral_potential'].isna().any()}")

print("\n📊 FINAL CHECK: Scanning all features for NaN values...")
nan_columns = features_df.columns[features_df.isna().any()].tolist()
if nan_columns:
    print(f"   ❌ Found NaN in columns: {nan_columns}")
    for col in nan_columns:
        nan_count = features_df[col].isna().sum()
        print(f"      - {col}: {nan_count} NaN values")
else:
    print(f"   ✅ NO NaN VALUES FOUND in any of the {len(features_df.columns)} columns!")

print("\n🎯 Edge Case Results:")
print("User Types Tested:")
print("- User 1: Normal active user")
print("- User 2: User with no clips created") 
print("- User 3: User with clips but no reactions")
print("- User 4: User with reactions but no bot actions")
print("- User 5: Brand new user with no activity")

print(f"\n✅ All {len(features_df)} users processed without NaN values!")
print(f"✅ Total features: {len(features_df.columns) - 1}")  # Exclude user_id