tasks/task_during_contains.py

"""
During/Contains task generator for temporal reasoning dataset.

Generates audio samples where one sound event is fully contained within
another (temporally). The "container" sound plays for a longer duration,
and the "contained" sound starts and ends entirely within the container.

Questions test whether a model can identify containment relationships:
- Which sound occurs during another?
- Which sound contains another?
- Yes/No containment verification
"""

import csv
import json
import random
from pathlib import Path
from typing import Dict, List, Optional, Tuple

from utils import (
    AudioProcessor,
    QuestionGenerator,
    setup_logger,
    set_random_seed,
    generate_sample_durations_for_task,
    generate_single_clip_duration,
    concatenate_to_target_duration,
    build_during_contains_task_audio,
    create_dataset
)


class DuringContainsTaskGenerator:
    """Generates during_contains task dataset samples."""
    
    def __init__(self, config: dict, logger=None):
        """
        Initialize DuringContainsTaskGenerator.
        
        Args:
            config: Full pipeline configuration dictionary
            logger: Logger instance
        """
        self.config = config
        self.logger = logger or setup_logger(__name__)
        
        self.task_config = config['tasks']['during_contains']
        
        # Audio parameters
        audio_config = config['audio']
        self.min_clip_duration = audio_config['min_clip_duration']
        self.max_clip_duration = audio_config['max_clip_duration']
        self.source_clip_duration = audio_config.get('source_clip_duration', 5.0)
        
        # Task-specific parameters
        self.task_duration_hours = self.task_config['task_duration_size']
        self.num_sounds = self.task_config.get('num_sounds', 2)
        self.min_container_duration_s = self.task_config.get('min_container_duration_s', 8.0)
        self.max_container_duration_s = self.task_config.get('max_container_duration_s', 15.0)
        self.min_margin_s = self.task_config.get('min_margin_s', 0.5)
        self.min_margin_ms = int(self.min_margin_s * 1000)
        
        # Dataset
        self.dataset = create_dataset(config)
        
        # Audio processor
        self.audio_processor = AudioProcessor(
            crossfade_duration=audio_config.get('crossfade_duration', 500),
            silence_duration=audio_config.get('silence_duration', 1000),
            with_silence=False,
            normalize=audio_config.get('normalize', False),
            normalize_target_dBFS=audio_config.get('normalize_target_dBFS', -20.0)
        )
        
        # Question generator
        mcq_config = config.get('mcq', {})
        self.question_generator = QuestionGenerator(
            num_options=mcq_config.get('num_options', 4),
            option_labels=mcq_config.get('option_labels', ['A', 'B', 'C', 'D']),
            distractor_strategy=mcq_config.get('distractor_strategy', 'balanced')
        )
        
        # Output paths
        self.output_base = Path(config['output']['base_path']) / 'during_contains'
        self.audio_output = self.output_base / 'audio'
        self.output_base.mkdir(parents=True, exist_ok=True)
        self.audio_output.mkdir(parents=True, exist_ok=True)
    
    def generate_sample(
        self,
        sample_id: int,
        target_duration_seconds: float = None
    ) -> Optional[Dict]:
        """
        Generate a single during_contains task sample.
        
        Pipeline:
        1. Sample 2 categories (container and contained)
        2. Load audio for each
        3. Extend container to min_container_duration_s..max_container_duration_s
        4. Keep contained at source clip duration (shorter)
        5. Build containment audio using overlay()
        6. Generate questions about the containment
        
        Args:
            sample_id: Sample ID
            target_duration_seconds: Pre-generated target duration
            
        Returns:
            Metadata dictionary, or None if failed
        """
        # Step 1: Sample categories
        try:
            categories = self.dataset.sample_categories(self.num_sounds)
        except ValueError:
            self.logger.warning(f"Sample {sample_id}: Cannot sample {self.num_sounds} categories")
            return None
        
        container_category = categories[0]
        contained_category = categories[1]
        
        # Step 2: Load audio
        filename_container, filepath_container = self.dataset.sample_file_from_category(container_category)
        filename_contained, filepath_contained = self.dataset.sample_file_from_category(contained_category)
        
        container_audio = self.audio_processor.load_audio(filepath_container)
        contained_audio = self.audio_processor.load_audio(filepath_contained)
        
        # Step 3: Extend container to target duration
        container_duration_s = random.uniform(
            self.min_container_duration_s,
            self.max_container_duration_s
        )
        # Step 4: Contained stays at source clip duration (5s typically)
        # Make sure contained is shorter than container minus margins
        max_contained_duration_ms = len(container_audio) - 2 * self.min_margin_ms
        if len(contained_audio) > max_contained_duration_ms:
            contained_audio = contained_audio[:max_contained_duration_ms]
        
        # Step 5: Build containment audio
        try:
            final_audio, build_metadata = build_during_contains_task_audio(
                container_audio,
                contained_audio,
                container_category,
                contained_category,
                min_margin_ms=self.min_margin_ms
            )
        except ValueError as e:
            self.logger.warning(f"Sample {sample_id}: Failed to build containment audio: {e}")
            return None
        
        # Save audio
        output_audio_path = self.audio_output / f"{sample_id}.wav"
        final_audio.export(str(output_audio_path), format="wav")
        
        # Step 6: Generate questions
        question_type = random.choice(self.task_config['question_types'])
        
        mcq_data, open_text_data = self._generate_question(
            question_type, container_category, contained_category, build_metadata
        )
        
        if mcq_data is None:
            self.logger.warning(f"Sample {sample_id}: Failed to generate question")
            return None
        
        metadata = {
            'id': sample_id,
            'audio_path': str(output_audio_path.relative_to(self.output_base.parent)),
            'num_sounds': self.num_sounds,
            'source_files': [filename_container, filename_contained],
            'container_category': container_category,
            'contained_category': contained_category,
            'categories': categories,
            'question_type': question_type,
            'container_duration_ms': build_metadata['container_duration_ms'],
            'contained_duration_ms': build_metadata['contained_duration_ms'],
            'contained_start_ms': build_metadata['contained_start_ms'],
            'contained_end_ms': build_metadata['contained_end_ms'],
            'margin_before_ms': build_metadata['margin_before_ms'],
            'margin_after_ms': build_metadata['margin_after_ms'],
            'total_duration_ms': build_metadata['total_duration_ms'],
            'actual_duration_s': len(final_audio) / 1000.0,
            'mcq_question': mcq_data['question'],
            'mcq_options': mcq_data['options'],
            'mcq_correct_answer': mcq_data['correct_answer'],
            'open_text_question': open_text_data['question'],
            'open_text_answer': open_text_data['correct_answer']
        }
        
        self.logger.info(
            f"Generated during_contains sample {sample_id}: "
            f"{container_category} contains {contained_category}, "
            f"contained at {build_metadata['contained_start_ms']}-"
            f"{build_metadata['contained_end_ms']}ms, type={question_type}"
        )
        
        return metadata
    
    def _generate_question(
        self,
        question_type: str,
        container_category: str,
        contained_category: str,
        build_metadata: Dict
    ) -> Tuple[Optional[Dict], Optional[Dict]]:
        """Generate MCQ and open-text question based on type."""
        
        mcq_template = self.task_config['mcq_questions'].get(question_type, '')
        open_template = self.task_config['open_text_questions'].get(question_type, '')
        
        present_categories = [container_category, contained_category]
        
        if question_type == 'during':
            # "Which sound occurs during {anchor_sound}?"
            anchor = container_category  # The container is the anchor
            correct = contained_category  # The contained occurs "during" the container
            
            formatted_mcq = mcq_template.format(anchor_sound=anchor)
            formatted_open = open_template.format(anchor_sound=anchor)
            
            mcq_data = self.question_generator.generate_category_mcq(
                formatted_mcq, correct, present_categories, self.dataset.CATEGORIES
            )
            open_data = self.question_generator.generate_category_open_text(
                formatted_open, correct
            )
        
        elif question_type == 'contains':
            # "Which longer sound contains {target_sound}?"
            target = contained_category  # The shorter contained sound
            correct = container_category  # The container "contains" the target
            
            formatted_mcq = mcq_template.format(target_sound=target)
            formatted_open = open_template.format(target_sound=target)
            
            mcq_data = self.question_generator.generate_category_mcq(
                formatted_mcq, correct, present_categories, self.dataset.CATEGORIES
            )
            open_data = self.question_generator.generate_category_open_text(
                formatted_open, correct
            )
        
        elif question_type == 'yes_no_during':
            # "Does {target_sound} happen entirely during {anchor_sound}?"
            # Always True by construction
            formatted_mcq = mcq_template.format(
                target_sound=contained_category,
                anchor_sound=container_category
            )
            formatted_open = open_template.format(
                target_sound=contained_category,
                anchor_sound=container_category
            )
            
            mcq_data = self.question_generator.generate_yes_no_mcq(
                formatted_mcq, correct_answer=True
            )
            open_data = self.question_generator.generate_yes_no_open_text(
                formatted_open, correct_answer=True
            )
        
        else:
            return None, None
        
        return mcq_data, open_data
    
    def generate_dataset(self) -> tuple:
        """Generate the complete during_contains dataset."""
        sample_durations = generate_sample_durations_for_task(
            self.task_duration_hours,
            self.min_clip_duration,
            self.max_clip_duration
        )
        num_samples = len(sample_durations)
        
        self.logger.info(
            f"Generating {num_samples} during_contains samples "
            f"(target: {self.task_duration_hours}h)..."
        )
        
        all_metadata = []
        for i, duration in enumerate(sample_durations):
            metadata = self.generate_sample(i, target_duration_seconds=duration)
            if metadata is not None:
                all_metadata.append(metadata)
        
        self.logger.info(f"Generated {len(all_metadata)}/{num_samples} samples successfully")
        
        # Save CSVs
        mcq_csv_path = self.output_base / 'during_contains_mcq.csv'
        self._save_mcq_csv(all_metadata, mcq_csv_path)
        
        open_text_csv_path = self.output_base / 'during_contains_open_text.csv'
        self._save_open_text_csv(all_metadata, open_text_csv_path)
        
        metadata_csv_path = self.output_base / 'during_contains_metadata.csv'
        self._save_metadata_csv(all_metadata, metadata_csv_path)
        
        self.logger.info(f"During/Contains task complete!")
        self.logger.info(f"  - MCQ CSV: {mcq_csv_path}")
        self.logger.info(f"  - Open-text CSV: {open_text_csv_path}")
        self.logger.info(f"  - Metadata CSV: {metadata_csv_path}")
        
        return mcq_csv_path, open_text_csv_path
    
    def _save_mcq_csv(self, metadata_list: List[Dict], output_path: Path):
        """Save MCQ format CSV."""
        with open(output_path, 'w', newline='') as f:
            writer = csv.writer(f)
            writer.writerow([
                'question', 'id', 'audio_path',
                'optionA', 'optionB', 'optionC', 'optionD',
                'correct', 'question_type', 'source_wavs', 'source_categories',
                'container_duration_ms', 'contained_duration_ms'
            ])
            for meta in metadata_list:
                writer.writerow([
                    meta['mcq_question'],
                    meta['id'],
                    meta['audio_path'],
                    meta['mcq_options']['A'],
                    meta['mcq_options']['B'],
                    meta['mcq_options']['C'],
                    meta['mcq_options']['D'],
                    meta['mcq_correct_answer'],
                    meta['question_type'],
                    str(meta['source_files']),
                    str(meta['categories']),
                    meta['container_duration_ms'],
                    meta['contained_duration_ms']
                ])
    
    def _save_open_text_csv(self, metadata_list: List[Dict], output_path: Path):
        """Save open-text format CSV."""
        with open(output_path, 'w', newline='') as f:
            writer = csv.writer(f)
            writer.writerow([
                'question', 'id', 'audio_path', 'answer',
                'question_type', 'source_wavs', 'source_categories',
                'container_duration_ms'
            ])
            for meta in metadata_list:
                writer.writerow([
                    meta['open_text_question'],
                    meta['id'],
                    meta['audio_path'],
                    meta['open_text_answer'],
                    meta['question_type'],
                    str(meta['source_files']),
                    str(meta['categories']),
                    meta['container_duration_ms']
                ])
                
    def _save_metadata_csv(self, metadata_list: List[Dict], output_path: Path):
        """Save detailed metadata CSV."""
        with open(output_path, 'w', newline='') as f:
            writer = csv.writer(f)
            writer.writerow([
                'id', 'audio_path', 'num_sounds',
                'source_files', 'source_categories',
                'container_category', 'contained_category',
                'container_duration_ms', 'contained_duration_ms',
                'contained_start_ms', 'contained_end_ms',
                'margin_before_ms', 'margin_after_ms',
                'total_duration_ms', 'actual_duration_s', 'question_type'
            ])
            
            for meta in metadata_list:
                writer.writerow([
                    meta['id'],
                    meta['audio_path'],
                    meta['num_sounds'],
                    str(meta['source_files']),
                    str(meta['categories']),
                    meta['container_category'],
                    meta['contained_category'],
                    meta['container_duration_ms'],
                    meta['contained_duration_ms'],
                    meta['contained_start_ms'],
                    meta['contained_end_ms'],
                    meta['margin_before_ms'],
                    meta['margin_after_ms'],
                    meta['total_duration_ms'],
                    meta['actual_duration_s'],
                    meta['question_type']
                ])


def main(config_path: str = None):
    """Main entry point for during_contains task generation."""
    import yaml
    
    if config_path is None:
        config_path = Path(__file__).parent.parent / 'config.yaml'
    
    with open(config_path, 'r') as f:
        config = yaml.safe_load(f)
    
    set_random_seed(config['random_seed'])
    
    logger = setup_logger(
        'during_contains_task',
        log_file=str(Path(config['output']['base_path']) / config['logging']['log_file']),
        level=config['logging']['level'],
        console_output=config['logging']['console_output']
    )
    
    generator = DuringContainsTaskGenerator(config, logger)
    generator.generate_dataset()


if __name__ == '__main__':
    main()