"""
User subset selection and subscription feature extraction.
"""

import pandas as pd
import numpy as np
from typing import Dict, List, Optional, Tuple, Any
import gc


class UserSelector:
    """Handle user subset selection and basic feature extraction."""

    def __init__(self):
        self.subscription_mapping = {"active": 2, "past_due": 1, "free": 0}

    def select_creators(
        self,
        total_clip_df: pd.DataFrame,
        discord_info_df: pd.DataFrame,
        include_all_users: bool = True,
        verbose: bool = True,
    ) -> pd.DataFrame:
        """
        Select users for analysis, including both creators and non-creators.

        Args:
            total_clip_df: DataFrame containing all clips
            discord_info_df: DataFrame containing user info
            include_all_users: If True, include all users (creators and non-creators)
            verbose: Whether to print progress messages

        Returns:
            DataFrame with selected user IDs and creator flag
        """
        if verbose:
            if include_all_users:
                print("🎯 USER SUBSET SELECTION - ALL USERS")
            else:
                print("🎯 USER SUBSET SELECTION - CREATORS ONLY")
            print("=" * 50)

        # Get users who have created clips
        clip_creators = set(total_clip_df["user_id"].dropna().unique())

        if include_all_users:
            # Start with ALL users from discord_info_df
            users_df = pd.DataFrame({"user_id": discord_info_df["user_id"].unique()})
            users_df["user_id"] = pd.to_numeric(users_df["user_id"], errors="coerce")
            users_df = users_df.dropna(subset=["user_id"]).astype({"user_id": int})
            users_df = users_df.drop_duplicates().reset_index(drop=True)

            # Add creator flag
            users_df["is_creator"] = users_df["user_id"].isin(clip_creators).astype(int)

            if verbose:
                n_creators = users_df["is_creator"].sum()
                n_non_creators = len(users_df) - n_creators
                print(f"📊 Total users selected: {len(users_df):,}")
                print(
                    f"   • Content creators: {n_creators:,} ({n_creators/len(users_df)*100:.1f}%)"
                )
                print(
                    f"   • Non-creators: {n_non_creators:,} ({n_non_creators/len(users_df)*100:.1f}%)"
                )
                print(f"✅ Focus: Analyzing entire user base including free users")
        else:
            # Original behavior - creators only
            users_df = pd.DataFrame({"user_id": list(clip_creators)})
            users_df["user_id"] = pd.to_numeric(users_df["user_id"], errors="coerce")
            users_df = users_df.dropna(subset=["user_id"]).astype({"user_id": int})
            users_df = users_df.drop_duplicates().reset_index(drop=True)

            # Filter to only users that exist in discord_info_df
            valid_users = set(discord_info_df["user_id"].unique())
            users_df = users_df[users_df["user_id"].isin(valid_users)]
            users_df["is_creator"] = 1  # All are creators

            if verbose:
                print(f"📊 Content creators selected: {len(users_df):,} users")
                print(
                    f"📊 Total unique creators in clips dataset: {len(clip_creators):,}"
                )
                print(f"📊 Valid creators (exist in user database): {len(users_df):,}")
                print(
                    f"📊 This represents {len(users_df)/len(discord_info_df)*100:.1f}% of all users in the system"
                )
                print(
                    f"✅ Focus: Analyzing actual content creators from interesting clips dataset"
                )

        return users_df

    def extract_subscription_features(
        self,
        features_df: pd.DataFrame,
        discord_info_df: pd.DataFrame,
        verbose: bool = True,
    ) -> pd.DataFrame:
        """
        Extract subscription-related features.

        Args:
            features_df: DataFrame with user IDs
            discord_info_df: DataFrame containing subscription info
            verbose: Whether to print progress messages

        Returns:
            DataFrame with subscription features added
        """
        if verbose:
            print("💳 SUBSCRIPTION FEATURES")
            print("=" * 50)

        # Get subscription data for selected users
        subscription_data = discord_info_df[
            discord_info_df["user_id"].isin(features_df["user_id"])
        ]

        # Merge subscription status
        features_df = features_df.merge(
            subscription_data[["user_id", "subscription_status"]],
            on="user_id",
            how="left",
        )

        # Map to numeric tiers
        features_df["subscription_tier"] = (
            features_df["subscription_status"]
            .map(self.subscription_mapping)
            .fillna(0)
            .astype(int)
        )

        # Drop the temporary column
        features_df.drop("subscription_status", axis=1, inplace=True)

        if verbose:
            # Summary
            sub_dist = features_df["subscription_tier"].value_counts()
            print(
                f"📊 Distribution: Free={sub_dist.get(0, 0):,} | "
                f"Past Due={sub_dist.get(1, 0):,} | Active={sub_dist.get(2, 0):,}"
            )
            print(f"✅ Shape: {features_df.shape}")

        return features_df

    def create_initial_features(
        self,
        total_clip_df: pd.DataFrame,
        discord_info_df: pd.DataFrame,
        include_all_users: bool = True,
        verbose: bool = True,
    ) -> pd.DataFrame:
        """
        Create initial feature dataframe with user selection and subscription features.

        Args:
            total_clip_df: DataFrame containing all clips
            discord_info_df: DataFrame containing user info
            include_all_users: If True, include all users (creators and non-creators)
            verbose: Whether to print progress messages

        Returns:
            Initial features DataFrame with is_creator flag
        """
        # Select users (now includes all users by default)
        features_df = self.select_creators(
            total_clip_df, discord_info_df, include_all_users, verbose
        )

        # Add subscription features
        features_df = self.extract_subscription_features(
            features_df, discord_info_df, verbose
        )

        return features_df

    def get_user_statistics(self, features_df: pd.DataFrame) -> Dict[str, Any]:
        """
        Get statistics about the selected users.

        Args:
            features_df: DataFrame with selected users

        Returns:
            Dictionary with user statistics
        """
        stats = {
            "total_users": len(features_df),
            "subscription_distribution": features_df["subscription_tier"]
            .value_counts()
            .to_dict(),
            "active_subscribers": (features_df["subscription_tier"] == 2).sum(),
            "active_subscriber_pct": (features_df["subscription_tier"] == 2).mean()
            * 100,
            "paying_users": (features_df["subscription_tier"] > 0).sum(),
            "paying_user_pct": (features_df["subscription_tier"] > 0).mean() * 100,
        }

        return stats