data_processor.py

# MovieMindAI - Data Processor
# Processes and joins IMDb data into a unified format

import pandas as pd
import numpy as np
from pathlib import Path
from typing import Optional
import config
import data_loader


def calculate_popularity(df: pd.DataFrame) -> pd.Series:
    """
    Calculate popularity score based on rating, votes, and recency.
    
    Formula:
    popularity = (RATING_WEIGHT * normalized_rating) + 
                 (VOTES_WEIGHT * normalized_log_votes) + 
                 (RECENCY_WEIGHT * normalized_recency)
    """
    # Normalize rating (0-10 scale, already normalized conceptually)
    # But we'll scale it 0-1 for combination
    norm_rating = (df["averageRating"] - df["averageRating"].min()) / (
        df["averageRating"].max() - df["averageRating"].min()
    )
    
    # Normalize log votes (log scale to reduce impact of extreme values)
    log_votes = np.log10(df["numVotes"] + 1)
    norm_votes = (log_votes - log_votes.min()) / (log_votes.max() - log_votes.min())
    
    # Normalize recency (newer = higher score)
    # Use startYear, with current year as reference
    year_diff = config.CURRENT_YEAR - df["startYear"]
    max_diff = config.CURRENT_YEAR - config.MIN_YEAR  # Max possible difference
    norm_recency = 1 - (year_diff / max_diff)  # Invert so newer = higher
    norm_recency = norm_recency.clip(0, 1)  # Ensure bounds
    
    # Calculate weighted popularity
    popularity = (
        config.RATING_WEIGHT * norm_rating +
        config.VOTES_WEIGHT * norm_votes +
        config.RECENCY_WEIGHT * norm_recency
    )
    
    return popularity


def extract_origin_country(akas_df: pd.DataFrame) -> pd.DataFrame:
    """
    Extract origin country from title.akas.tsv.
    Uses the region with ordering=1 or isOriginalTitle=1.
    """
    # Prioritize isOriginalTitle, then lowest ordering
    original = akas_df[akas_df["isOriginalTitle"] == True].copy()
    
    if len(original) > 0:
        # Take the first region for original titles
        origin = original.groupby("titleId").first()["region"].reset_index()
    else:
        # Fallback to ordering=1
        first_order = akas_df[akas_df["ordering"] == 1].copy()
        origin = first_order.groupby("titleId").first()["region"].reset_index()
    
    origin.columns = ["tconst", "originCountry"]
    return origin


def extract_actors_directors(
    principals_df: pd.DataFrame,
    names_df: pd.DataFrame,
    max_actors: int = 5,
    max_directors: int = 3
) -> pd.DataFrame:
    """
    Extract top actors and directors for each title.
    """
    # Create name lookup
    name_lookup = names_df.set_index("nconst")["primaryName"].to_dict()
    
    # Filter to actors and directors
    actors = principals_df[principals_df["category"].isin(["actor", "actress"])].copy()
    directors = principals_df[principals_df["category"] == "director"].copy()
    
    # Map nconst to names
    actors["name"] = actors["nconst"].map(name_lookup)
    directors["name"] = directors["nconst"].map(name_lookup)
    
    # Get top N actors per title (by ordering)
    actors_sorted = actors.sort_values(["tconst", "ordering"])
    top_actors = (
        actors_sorted.groupby("tconst")["name"]
        .apply(lambda x: list(x.dropna().head(max_actors)))
        .reset_index()
    )
    top_actors.columns = ["tconst", "actors"]
    
    # Get top N directors per title
    directors_sorted = directors.sort_values(["tconst", "ordering"])
    top_directors = (
        directors_sorted.groupby("tconst")["name"]
        .apply(lambda x: list(x.dropna().head(max_directors)))
        .reset_index()
    )
    top_directors.columns = ["tconst", "directors"]
    
    # Merge actors and directors
    result = pd.merge(top_actors, top_directors, on="tconst", how="outer")
    
    return result


def process_data() -> pd.DataFrame:
    """
    Main processing function. Loads, filters, and joins all data.
    """
    print("=" * 60)
    print("MOVIEMINDAI DATA PROCESSING")
    print("=" * 60)
    
    # Step 1: Load title basics with initial filters
    print("\n[1/7] Loading title basics...")
    basics = data_loader.load_title_basics(
        filter_types=config.VALID_TITLE_TYPES,
        min_year=config.MIN_YEAR
    )
    
    # Step 2: Load and join ratings
    print("\n[2/7] Loading and joining ratings...")
    ratings = data_loader.load_title_ratings()
    
    # Merge and apply rating/vote filters
    df = pd.merge(basics, ratings, on="tconst", how="inner")
    df = df[df["numVotes"] >= config.MIN_VOTES]
    df = df[df["averageRating"] >= config.MIN_RATING]
    
    print(f"  After rating filters: {len(df):,} titles")
    
    # Get the filtered title IDs for subsequent loads
    title_ids = set(df["tconst"].unique())
    print(f"  Unique titles to process: {len(title_ids):,}")
    
    # Step 3: Load name basics for actor/director lookup
    print("\n[3/7] Loading name basics...")
    names = data_loader.load_name_basics()
    
    # Step 4: Load principals (filtered to our titles)
    print("\n[4/7] Loading principals...")
    principals = data_loader.load_title_principals_filtered(title_ids)
    
    # Step 5: Extract actors and directors
    print("\n[5/7] Extracting actors and directors...")
    cast_crew = extract_actors_directors(
        principals,
        names,
        max_actors=config.MAX_ACTORS_PER_TITLE,
        max_directors=config.MAX_DIRECTORS_PER_TITLE
    )
    df = pd.merge(df, cast_crew, on="tconst", how="left")
    
    # Fill missing with empty lists
    df["actors"] = df["actors"].apply(lambda x: x if isinstance(x, list) else [])
    df["directors"] = df["directors"].apply(lambda x: x if isinstance(x, list) else [])
    
    # Step 6: Load akas and extract origin country
    print("\n[6/7] Loading akas and extracting origin countries...")
    akas = data_loader.load_title_akas_filtered(title_ids)
    origin = extract_origin_country(akas)
    df = pd.merge(df, origin, on="tconst", how="left")
    df["originCountry"] = df["originCountry"].fillna("Unknown")
    
    # Step 7: Calculate popularity and final processing
    print("\n[7/7] Calculating popularity and finalizing...")
    df["popularity"] = calculate_popularity(df)
    
    # Parse genres into list
    df["genres"] = df["genres"].apply(
        lambda x: x.split(",") if isinstance(x, str) else []
    )
    
    # Select and order final columns
    final_columns = [
        "tconst",
        "titleType",
        "primaryTitle",
        "originalTitle",
        "startYear",
        "endYear",
        "runtimeMinutes",
        "genres",
        "averageRating",
        "numVotes",
        "directors",
        "actors",
        "originCountry",
        "popularity",
    ]
    
    df = df[final_columns].copy()
    
    # Sort by popularity descending
    df = df.sort_values("popularity", ascending=False).reset_index(drop=True)
    
    print("\n" + "=" * 60)
    print("PROCESSING COMPLETE")
    print("=" * 60)
    print(f"Total titles: {len(df):,}")
    print(f"Movies: {len(df[df['titleType'] == 'movie']):,}")
    print(f"TV Series: {len(df[df['titleType'] == 'tvSeries']):,}")
    print(f"TV Mini-Series: {len(df[df['titleType'] == 'tvMiniSeries']):,}")
    print(f"Year range: {df['startYear'].min():.0f} - {df['startYear'].max():.0f}")
    print(f"Rating range: {df['averageRating'].min():.1f} - {df['averageRating'].max():.1f}")
    
    return df


def save_processed_data(df: pd.DataFrame) -> None:
    """Save processed data to parquet file."""
    # Create processed directory if it doesn't exist
    config.PROCESSED_DIR.mkdir(parents=True, exist_ok=True)
    
    # Save to parquet
    df.to_parquet(config.PROCESSED_TITLES_PATH, index=False)
    print(f"\nSaved to: {config.PROCESSED_TITLES_PATH}")


def load_processed_data() -> pd.DataFrame:
    """Load processed data from parquet file."""
    return pd.read_parquet(config.PROCESSED_TITLES_PATH)


if __name__ == "__main__":
    # Process and save data
    df = process_data()
    save_processed_data(df)
    
    # Show sample
    print("\nSample titles:")
    print(df[["primaryTitle", "startYear", "averageRating", "popularity"]].head(10))