main.py

"""
Main pipeline runner for temporal reasoning audio dataset generation.

This script orchestrates the generation of all task datasets.
"""

import argparse
import sys
import yaml
from pathlib import Path
from typing import List, Optional

# Add project root to path
sys.path.append(str(Path(__file__).parent))

from utils import setup_logger, set_random_seed
from tasks.task_count import CountTaskGenerator
from tasks.task_duration import DurationTaskGenerator
from tasks.task_order import OrderTaskGenerator
from tasks.task_volume import VolumeTaskGenerator
from tasks.task_silence_gap import SilenceGapTaskGenerator
from tasks.task_overlap import OverlapTaskGenerator
from tasks.task_during_contains import DuringContainsTaskGenerator

# Multihop temporal reasoning — singular tasks
from tasks.task_conditional_count import ConditionalCountTaskGenerator
from tasks.task_conditional_duration import ConditionalDurationTaskGenerator
from tasks.task_between_events import BetweenEventsTaskGenerator
from tasks.task_event_density import EventDensityTaskGenerator

# Multihop temporal reasoning — inter-task
from tasks.task_duration_gap import DurationGapTaskGenerator
from tasks.task_temporal_arithmetic import TemporalArithmeticTaskGenerator
from tasks.task_temporal_loudness import TemporalLoudnessTaskGenerator
from tasks.task_multi_hop import MultiHopTaskGenerator


def load_config(config_path: str) -> dict:
    """Load configuration from YAML file."""
    with open(config_path, 'r') as f:
        config = yaml.safe_load(f)
    return config


def run_count_task(config: dict, logger):
    """Run the count task generation."""
    if not config['tasks']['count']['enabled']:
        logger.info("Count task is disabled, skipping...")
        return
    
    logger.info("=" * 80)
    logger.info("STARTING COUNT TASK GENERATION")
    logger.info("=" * 80)
    
    generator = CountTaskGenerator(config, logger)
    generator.dataset.reset_category_usage()  # Reset counter for this task
    generator.generate_dataset()
    
    # Log category usage statistics
    usage_stats = generator.dataset.get_category_usage_stats()
    sorted_stats = sorted(usage_stats.items(), key=lambda x: x[1], reverse=True)
    logger.info("Category usage statistics (as answers):")
    logger.info(f"  Min usage: {sorted_stats[-1][1]} (category: {sorted_stats[-1][0]})")
    logger.info(f"  Max usage: {sorted_stats[0][1]} (category: {sorted_stats[0][0]})")
    logger.info(f"  Mean usage: {sum(usage_stats.values()) / len(usage_stats):.2f}")
    
    logger.info("Count task completed successfully!")


def run_duration_task(config: dict, logger):
    """Run the duration task generation."""
    if not config['tasks']['duration']['enabled']:
        logger.info("Duration task is disabled, skipping...")
        return
    
    logger.info("=" * 80)
    logger.info("STARTING DURATION TASK GENERATION")
    logger.info("=" * 80)
    
    generator = DurationTaskGenerator(config, logger)
    generator.dataset.reset_category_usage()  # Reset counter for this task
    generator.generate_dataset()
    
    # Log category usage statistics
    usage_stats = generator.dataset.get_category_usage_stats()
    sorted_stats = sorted(usage_stats.items(), key=lambda x: x[1], reverse=True)
    logger.info("Category usage statistics (as longest/shortest answers):")
    logger.info(f"  Min usage: {sorted_stats[-1][1]} (category: {sorted_stats[-1][0]})")
    logger.info(f"  Max usage: {sorted_stats[0][1]} (category: {sorted_stats[0][0]})")
    logger.info(f"  Mean usage: {sum(usage_stats.values()) / len(usage_stats):.2f}")
    
    logger.info("Duration task completed successfully!")


def run_order_task(config: dict, logger):
    """Run the order task generation."""
    if not config['tasks']['order']['enabled']:
        logger.info("Order task is disabled, skipping...")
        return
    
    logger.info("=" * 80)
    logger.info("STARTING ORDER TASK GENERATION")
    logger.info("=" * 80)
    
    generator = OrderTaskGenerator(config, logger)
    generator.dataset.reset_category_usage()  # Reset counter for this task
    generator.generate_dataset()
    
    # Log category usage statistics
    usage_stats = generator.dataset.get_category_usage_stats()
    sorted_stats = sorted(usage_stats.items(), key=lambda x: x[1], reverse=True)
    logger.info("Category usage statistics (as first/last/after/before answers):")
    logger.info(f"  Min usage: {sorted_stats[-1][1]} (category: {sorted_stats[-1][0]})")
    logger.info(f"  Max usage: {sorted_stats[0][1]} (category: {sorted_stats[0][0]})")
    logger.info(f"  Mean usage: {sum(usage_stats.values()) / len(usage_stats):.2f}")
    
    logger.info("Order task completed successfully!")


def run_volume_task(config: dict, logger):
    """Run the volume task generation."""
    if not config['tasks']['volume']['enabled']:
        logger.info("Volume task is disabled, skipping...")
        return
    
    logger.info("=" * 80)
    logger.info("STARTING VOLUME TASK GENERATION")
    logger.info("=" * 80)
    
    generator = VolumeTaskGenerator(config, logger)
    generator.dataset.reset_category_usage()  # Reset counter for this task
    generator.generate_dataset()
    
    # Log category usage statistics
    usage_stats = generator.dataset.get_category_usage_stats()
    sorted_stats = sorted(usage_stats.items(), key=lambda x: x[1], reverse=True)
    logger.info("Category usage statistics (as loudest/softest answers):")
    logger.info(f"  Min usage: {sorted_stats[-1][1]} (category: {sorted_stats[-1][0]})")
    logger.info(f"  Max usage: {sorted_stats[0][1]} (category: {sorted_stats[0][0]})")
    logger.info(f"  Mean usage: {sum(usage_stats.values()) / len(usage_stats):.2f}")
    
    logger.info("Volume task completed successfully!")


def run_silence_gap_task(config: dict, logger):
    """Run the silence gap task generation."""
    if not config['tasks'].get('silence_gap', {}).get('enabled', False):
        logger.info("Silence gap task is disabled, skipping...")
        return
    
    logger.info("=" * 80)
    logger.info("STARTING SILENCE GAP TASK GENERATION")
    logger.info("=" * 80)
    
    generator = SilenceGapTaskGenerator(config, logger)
    generator.generate_dataset()
    
    logger.info("Silence gap task completed successfully!")


def run_overlap_task(config: dict, logger):
    """Run the overlap task generation."""
    if not config['tasks'].get('overlap', {}).get('enabled', False):
        logger.info("Overlap task is disabled, skipping...")
        return
    
    logger.info("=" * 80)
    logger.info("STARTING OVERLAP TASK GENERATION")
    logger.info("=" * 80)
    
    generator = OverlapTaskGenerator(config, logger)
    generator.dataset.reset_category_usage()
    generator.generate_dataset()
    
    logger.info("Overlap task completed successfully!")


def run_during_contains_task(config: dict, logger):
    """Run the during/contains task generation."""
    if not config['tasks'].get('during_contains', {}).get('enabled', False):
        logger.info("During/Contains task is disabled, skipping...")
        return
    
    logger.info("=" * 80)
    logger.info("STARTING DURING/CONTAINS TASK GENERATION")
    logger.info("=" * 80)
    
    generator = DuringContainsTaskGenerator(config, logger)
    generator.dataset.reset_category_usage()
    generator.generate_dataset()
    
    logger.info("During/Contains task completed successfully!")


# =====================================================================
# MULTIHOP TEMPORAL REASONING — SINGULAR TASKS
# =====================================================================

def run_conditional_count_task(config: dict, logger):
    """Run the conditional count task generation."""
    if not config['tasks'].get('conditional_count', {}).get('enabled', False):
        logger.info("Conditional count task is disabled, skipping...")
        return
    
    logger.info("=" * 80)
    logger.info("STARTING CONDITIONAL COUNT TASK GENERATION")
    logger.info("=" * 80)
    
    generator = ConditionalCountTaskGenerator(config, logger)
    generator.generate_dataset()
    
    logger.info("Conditional count task completed successfully!")


def run_conditional_duration_task(config: dict, logger):
    """Run the conditional duration task generation."""
    if not config['tasks'].get('conditional_duration', {}).get('enabled', False):
        logger.info("Conditional duration task is disabled, skipping...")
        return
    
    logger.info("=" * 80)
    logger.info("STARTING CONDITIONAL DURATION TASK GENERATION")
    logger.info("=" * 80)
    
    generator = ConditionalDurationTaskGenerator(config, logger)
    generator.generate_dataset()
    
    logger.info("Conditional duration task completed successfully!")


def run_between_events_task(config: dict, logger):
    """Run the between events task generation."""
    if not config['tasks'].get('between_events', {}).get('enabled', False):
        logger.info("Between events task is disabled, skipping...")
        return
    
    logger.info("=" * 80)
    logger.info("STARTING BETWEEN EVENTS TASK GENERATION")
    logger.info("=" * 80)
    
    generator = BetweenEventsTaskGenerator(config, logger)
    generator.generate_dataset()
    
    logger.info("Between events task completed successfully!")


def run_event_density_task(config: dict, logger):
    """Run the event density task generation."""
    if not config['tasks'].get('event_density', {}).get('enabled', False):
        logger.info("Event density task is disabled, skipping...")
        return
    
    logger.info("=" * 80)
    logger.info("STARTING EVENT DENSITY TASK GENERATION")
    logger.info("=" * 80)
    
    generator = EventDensityTaskGenerator(config, logger)
    generator.generate_dataset()
    
    logger.info("Event density task completed successfully!")


# =====================================================================
# MULTIHOP TEMPORAL REASONING — INTER-TASK
# =====================================================================

def run_duration_gap_task(config: dict, logger):
    """Run the duration gap task generation."""
    if not config['tasks'].get('duration_gap', {}).get('enabled', False):
        logger.info("Duration gap task is disabled, skipping...")
        return
    
    logger.info("=" * 80)
    logger.info("STARTING DURATION GAP TASK GENERATION")
    logger.info("=" * 80)
    
    generator = DurationGapTaskGenerator(config, logger)
    generator.generate_dataset()
    
    logger.info("Duration gap task completed successfully!")


def run_temporal_arithmetic_task(config: dict, logger):
    """Run the temporal arithmetic task generation."""
    if not config['tasks'].get('temporal_arithmetic', {}).get('enabled', False):
        logger.info("Temporal arithmetic task is disabled, skipping...")
        return
    
    logger.info("=" * 80)
    logger.info("STARTING TEMPORAL ARITHMETIC TASK GENERATION")
    logger.info("=" * 80)
    
    generator = TemporalArithmeticTaskGenerator(config, logger)
    generator.generate_dataset()
    
    logger.info("Temporal arithmetic task completed successfully!")


def run_temporal_loudness_task(config: dict, logger):
    """Run the temporal loudness task generation."""
    if not config['tasks'].get('temporal_loudness', {}).get('enabled', False):
        logger.info("Temporal loudness task is disabled, skipping...")
        return
    
    logger.info("=" * 80)
    logger.info("STARTING TEMPORAL LOUDNESS TASK GENERATION")
    logger.info("=" * 80)
    
    generator = TemporalLoudnessTaskGenerator(config, logger)
    generator.generate_dataset()
    
    logger.info("Temporal loudness task completed successfully!")


def run_multi_hop_task(config: dict, logger):
    """Run the multi-hop task generation."""
    if not config['tasks'].get('multi_hop', {}).get('enabled', False):
        logger.info("Multi-hop task is disabled, skipping...")
        return
    
    logger.info("=" * 80)
    logger.info("STARTING MULTI-HOP TASK GENERATION")
    logger.info("=" * 80)
    
    generator = MultiHopTaskGenerator(config, logger)
    generator.generate_dataset()
    
    logger.info("Multi-hop task completed successfully!")


def run_pipeline(
    config_path: str,
    tasks: Optional[List[str]] = None,
    output_path: Optional[str] = None
):
    """
    Run the complete dataset generation pipeline.
    
    Args:
        config_path: Path to configuration YAML file
        tasks: Optional list of specific tasks to run (default: all enabled tasks)
        output_path: Optional custom output path (overrides config)
    """
    # Load configuration
    config = load_config(config_path)
    
    # Override output path if provided
    if output_path:
        config['output']['base_path'] = output_path
    
    # Create output directory
    output_base = Path(config['output']['base_path'])
    output_base.mkdir(parents=True, exist_ok=True)
    
    # Set random seed
    set_random_seed(config['random_seed'])
    
    # Setup main logger
    logger = setup_logger(
        'pipeline',
        log_file=str(output_base / config['logging']['log_file']),
        level=config['logging']['level'],
        console_output=config['logging']['console_output']
    )
    
    logger.info("=" * 80)
    logger.info("TEMPORAL REASONING AUDIO DATASET GENERATION PIPELINE")
    logger.info("=" * 80)
    logger.info(f"Configuration: {config_path}")
    logger.info(f"Output directory: {output_base}")
    logger.info(f"Random seed: {config['random_seed']}")
    if 'dataset_source' in config:
        ds = config['dataset_source']
        logger.info(f"Dataset source: {ds.get('name', 'Custom')}")
        logger.info(f"Dataset audio path: {ds['audio_path']}")
        logger.info(f"Dataset metadata path: {ds['metadata_path']}")
    else:
        logger.info(f"ESC-50 audio path: {config['esc50']['audio_path']}")
        logger.info(f"ESC-50 metadata path: {config['esc50']['metadata_path']}")
    
    # Determine which tasks to run
    task_map = {
        # Simple singular temporal reasoning
        'count': run_count_task,
        'duration': run_duration_task,
        'order': run_order_task,
        'volume': run_volume_task,
        'silence_gap': run_silence_gap_task,
        'overlap': run_overlap_task,
        'during_contains': run_during_contains_task,
        # Multihop temporal reasoning — singular
        'conditional_count': run_conditional_count_task,
        'conditional_duration': run_conditional_duration_task,
        'between_events': run_between_events_task,
        'event_density': run_event_density_task,
        # Multihop temporal reasoning — inter-task
        'duration_gap': run_duration_gap_task,
        'temporal_arithmetic': run_temporal_arithmetic_task,
        'temporal_loudness': run_temporal_loudness_task,
        'multi_hop': run_multi_hop_task,
    }
    
    if tasks:
        tasks_to_run = {k: v for k, v in task_map.items() if k in tasks}
        logger.info(f"Running specific tasks: {', '.join(tasks)}")
    else:
        tasks_to_run = task_map
        logger.info("Running all enabled tasks")
    
    # Run tasks
    for task_name, task_func in tasks_to_run.items():
        try:
            task_dir = output_base / task_name
            metadata_file = task_dir / f"{task_name}_metadata.csv"
            if metadata_file.exists():
                logger.info(f"Task '{task_name}' output already exists at {metadata_file}. Skipping...")
                continue
            task_func(config, logger)
        except Exception as e:
            logger.error(f"Error running {task_name} task: {e}", exc_info=True)
            raise
    
    logger.info("=" * 80)
    logger.info("PIPELINE COMPLETED SUCCESSFULLY!")
    logger.info("=" * 80)
    logger.info(f"All outputs saved to: {output_base}")


def main():
    """Main entry point with argument parsing."""
    parser = argparse.ArgumentParser(
        description="Temporal Reasoning Audio Dataset Generation Pipeline",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Run all tasks with default config
  python main.py
  
  # Run with custom config
  python main.py --config my_config.yaml
  
  # Run specific tasks only
  python main.py --tasks count duration
  
  # Use custom output directory
  python main.py --output /path/to/output
  
  # Combine options
  python main.py --config custom.yaml --tasks count order --output ./my_dataset
        """
    )
    
    parser.add_argument(
        '--config', '-c',
        type=str,
        default='config.yaml',
        help='Path to configuration YAML file (default: config.yaml)'
    )
    
    parser.add_argument(
        '--tasks', '-t',
        nargs='+',
        choices=[
            'count', 'duration', 'order', 'volume', 'silence_gap', 'overlap', 'during_contains',
            'conditional_count', 'conditional_duration', 'between_events', 'event_density',
            'duration_gap', 'temporal_arithmetic', 'temporal_loudness', 'multi_hop'
        ],
        help='Specific tasks to run (default: all enabled tasks)'
    )
    
    parser.add_argument(
        '--output', '-o',
        type=str,
        help='Custom output directory (overrides config)'
    )
    
    args = parser.parse_args()
    
    # Check if config file exists
    config_path = Path(args.config)
    if not config_path.exists():
        # Try relative to script directory
        script_dir = Path(__file__).parent
        config_path = script_dir / args.config
        if not config_path.exists():
            print(f"Error: Config file not found: {args.config}")
            sys.exit(1)
    
    # Run pipeline
    try:
        run_pipeline(
            config_path=str(config_path),
            tasks=args.tasks,
            output_path=args.output
        )
    except Exception as e:
        print(f"Pipeline failed with error: {e}")
        sys.exit(1)


if __name__ == '__main__':
    main()