JPMC-quant/run_filter_combinations.py

#!/usr/bin/env python
"""
Batch runner for quant report with different filter combinations.

Runs 03_quant_report.script.py for each single-filter combination:
- Each age group (with all others active)
- Each gender (with all others active)
- Each ethnicity (with all others active)
- Each income group (with all others active)
- Each consumer segment (with all others active)

Usage:
    uv run python run_filter_combinations.py
    uv run python run_filter_combinations.py --dry-run  # Preview combinations without running
"""

import subprocess
import sys
import json
from pathlib import Path

from tqdm import tqdm

from utils import QualtricsSurvey


# Default data paths (same as in 03_quant_report.script.py)
RESULTS_FILE = 'data/exports/2-2-26/JPMC_Chase Brand Personality_Quant Round 1_February 2, 2026_Labels.csv'
QSF_FILE = 'data/exports/OneDrive_2026-01-21/Soft Launch Data/JPMC_Chase_Brand_Personality_Quant_Round_1.qsf'

REPORT_SCRIPT = Path(__file__).parent / '03_quant_report.script.py'


def get_filter_combinations(survey: QualtricsSurvey) -> list[dict]:
    """
    Generate all single-filter combinations.
    
    Each combination isolates ONE filter value while keeping all others at "all selected".
    Returns list of dicts with filter kwargs for each run.
    """
    combinations = []
    
    # Add "All Respondents" run (no filters = all options selected)
    combinations.append({
        'name': 'All_Respondents',
        'filters': {}  # Empty = use defaults (all selected)
    })
    
    # Age groups - one at a time
    for age in survey.options_age:
        combinations.append({
            'name': f'Age-{age}',
            'filters': {'age': [age]}
        })
    
    # Gender - one at a time
    for gender in survey.options_gender:
        combinations.append({
            'name': f'Gender-{gender}',
            'filters': {'gender': [gender]}
        })
    
    # Ethnicity - grouped by individual values
    # Ethnicity options are comma-separated (e.g., "White or Caucasian, Hispanic or Latino")
    # Create filters that include ALL options containing each individual ethnicity value
    ethnicity_values = set()
    for ethnicity_option in survey.options_ethnicity:
        # Split by comma and strip whitespace
        values = [v.strip() for v in ethnicity_option.split(',')]
        ethnicity_values.update(values)
    
    for ethnicity_value in sorted(ethnicity_values):
        # Find all options that contain this value
        matching_options = [
            opt for opt in survey.options_ethnicity 
            if ethnicity_value in [v.strip() for v in opt.split(',')]
        ]
        combinations.append({
            'name': f'Ethnicity-{ethnicity_value}',
            'filters': {'ethnicity': matching_options}
        })
    
    # Income - one at a time
    for income in survey.options_income:
        combinations.append({
            'name': f'Income-{income}',
            'filters': {'income': [income]}
        })
    
    # Consumer segments - combine _A and _B options
    # Group options by base name (removing _A/_B suffix)
    consumer_groups = {}
    for consumer in survey.options_consumer:
        # Check if ends with _A or _B
        if consumer.endswith('_A') or consumer.endswith('_B'):
            base_name = consumer[:-2]  # Remove last 2 chars (_A or _B)
            if base_name not in consumer_groups:
                consumer_groups[base_name] = []
            consumer_groups[base_name].append(consumer)
        else:
            # Not an _A/_B option, keep as-is
            consumer_groups[consumer] = [consumer]
    
    for base_name, options in consumer_groups.items():
        combinations.append({
            'name': f'Consumer-{base_name}',
            'filters': {'consumer': options}
        })
    
    return combinations


def run_report(filters: dict, dry_run: bool = False) -> bool:
    """
    Run the report script with given filters.
    
    Args:
        filters: Dict of filter_name -> list of values
        dry_run: If True, just print command without running
        
    Returns:
        True if successful, False otherwise
    """
    cmd = [sys.executable, str(REPORT_SCRIPT)]
    
    for filter_name, values in filters.items():
        if values:
            cmd.extend([f'--{filter_name}', json.dumps(values)])
    
    if dry_run:
        print(f"  Would run: {' '.join(cmd)}")
        return True
    
    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            cwd=Path(__file__).parent
        )
        if result.returncode != 0:
            print(f"\n  ERROR: {result.stderr[:500]}")
            return False
        return True
    except Exception as e:
        print(f"\n  ERROR: {e}")
        return False


def main():
    import argparse
    parser = argparse.ArgumentParser(description='Run quant report for all filter combinations')
    parser.add_argument('--dry-run', action='store_true', help='Preview combinations without running')
    args = parser.parse_args()
    
    # Load survey to get available filter options
    print("Loading survey to get filter options...")
    survey = QualtricsSurvey(RESULTS_FILE, QSF_FILE)
    survey.load_data()  # Populates options_* attributes
    
    # Generate all combinations
    combinations = get_filter_combinations(survey)
    print(f"Generated {len(combinations)} filter combinations")
    
    if args.dry_run:
        print("\nDRY RUN - Commands that would be executed:")
        for combo in combinations:
            print(f"\n{combo['name']}:")
            run_report(combo['filters'], dry_run=True)
        return
    
    # Run each combination with progress bar
    successful = 0
    failed = []
    
    for combo in tqdm(combinations, desc="Running reports", unit="filter"):
        tqdm.write(f"Running: {combo['name']}")
        if run_report(combo['filters']):
            successful += 1
        else:
            failed.append(combo['name'])
    
    # Summary
    print(f"\n{'='*50}")
    print(f"Completed: {successful}/{len(combinations)} successful")
    if failed:
        print(f"Failed: {', '.join(failed)}")


if __name__ == '__main__':
    main()