tasks/task_overlap.py

"""
Overlap task generator for temporal reasoning dataset.

Generates audio samples where two sound events partially overlap in time,
and asks questions about which sounds overlap, identification of overlapping
pairs, and yes/no overlap 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_overlap_task_audio,
    create_dataset
)


class OverlapTaskGenerator:
    """Generates overlap task dataset samples."""
    
    def __init__(self, config: dict, logger=None):
        """
        Initialize OverlapTaskGenerator.
        
        Args:
            config: Full pipeline configuration dictionary
            logger: Logger instance
        """
        self.config = config
        self.logger = logger or setup_logger(__name__)
        
        self.task_config = config['tasks']['overlap']
        
        # 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_overlap_ratio = self.task_config.get('min_overlap_ratio', 0.2)
        self.max_overlap_ratio = self.task_config.get('max_overlap_ratio', 0.5)
        
        # 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']) / 'overlap'
        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 overlap task sample.
        
        Pipeline:
        1. Sample 2 categories
        2. Load audio for each
        3. Extend clips to reasonable duration
        4. Build overlapping audio using overlay()
        5. Generate questions about the overlap
        
        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
        
        category_a, category_b = categories[0], categories[1]
        
        # Step 2: Load audio
        filename_a, filepath_a = self.dataset.sample_file_from_category(category_a)
        filename_b, filepath_b = self.dataset.sample_file_from_category(category_b)
        
        audio_a = self.audio_processor.load_audio(filepath_a)
        audio_b = self.audio_processor.load_audio(filepath_b)
        
        # Step 3: Extend clips to a reasonable duration for overlap
        # Each clip should be at least source_clip_duration
        clip_target_s = max(self.source_clip_duration, 5.0)
        if target_duration_seconds:
            # Each clip takes roughly half the total minus overlap
            clip_target_s = max(clip_target_s, target_duration_seconds * 0.6)
        # Step 4: Build overlapping audio
        try:
            final_audio, build_metadata = build_overlap_task_audio(
                audio_a, audio_b,
                category_a, category_b,
                min_overlap_ratio=self.min_overlap_ratio,
                max_overlap_ratio=self.max_overlap_ratio
            )
        except Exception as e:
            self.logger.warning(f"Sample {sample_id}: Failed to build overlap 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 5: Generate questions
        question_type = random.choice(self.task_config['question_types'])
        
        mcq_data, open_text_data = self._generate_question(
            question_type, category_a, category_b, 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_a, filename_b],
            'category_a': category_a,
            'category_b': category_b,
            'categories': categories,
            'question_type': question_type,
            'a_start_ms': build_metadata['a_start_ms'],
            'a_end_ms': build_metadata['a_end_ms'],
            'b_start_ms': build_metadata['b_start_ms'],
            'b_end_ms': build_metadata['b_end_ms'],
            'overlap_duration_ms': build_metadata['overlap_duration_ms'],
            'overlap_ratio': build_metadata['overlap_ratio'],
            '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 overlap sample {sample_id}: {category_a} + {category_b}, "
            f"overlap={build_metadata['overlap_duration_ms']}ms "
            f"({build_metadata['overlap_ratio']*100:.1f}%), type={question_type}"
        )
        
        return metadata
    
    def _generate_question(
        self,
        question_type: str,
        category_a: str,
        category_b: 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 = [category_a, category_b]
        
        if question_type == 'identify_overlap':
            # "Which sound overlaps with {anchor_sound}?"
            anchor = random.choice([category_a, category_b])
            correct = category_b if anchor == category_a else category_a
            
            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 == 'overlap_pair':
            # "Which two sounds occur at the same time?"
            correct_pair = (category_a, category_b)
            mcq_data = self.question_generator.generate_pair_mcq(
                mcq_template, correct_pair, present_categories, self.dataset.CATEGORIES
            )
            open_data = self.question_generator.generate_pair_open_text(
                open_template, correct_pair
            )
        
        elif question_type == 'yes_no_overlap':
            # "Does {sound1} overlap with {sound2}?"
            # Always True for the overlapping pair
            formatted_mcq = mcq_template.format(sound1=category_a, sound2=category_b)
            formatted_open = open_template.format(sound1=category_a, sound2=category_b)
            
            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
            )
        
        elif question_type == 'starts_before_end':
            # "Which sound starts before {anchor_sound} ends?"
            # Sound A starts first, Sound B starts before A ends (by construction)
            anchor = category_a  # A starts first
            correct = category_b  # B starts during A
            
            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
            )
        
        else:
            return None, None
        
        return mcq_data, open_data
    
    def generate_dataset(self) -> tuple:
        """Generate the complete overlap 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} overlap 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 / 'overlap_mcq.csv'
        self._save_mcq_csv(all_metadata, mcq_csv_path)
        
        open_text_csv_path = self.output_base / 'overlap_open_text.csv'
        self._save_open_text_csv(all_metadata, open_text_csv_path)
        
        metadata_csv_path = self.output_base / 'overlap_metadata.csv'
        self._save_metadata_csv(all_metadata, metadata_csv_path)
        
        self.logger.info(f"Overlap 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',
                'overlap_duration_ms', 'overlap_ratio'
            ])
            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['overlap_duration_ms'],
                    round(meta['overlap_ratio'], 3)
                ])
    
    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',
                'overlap_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['overlap_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',
                'category_a', 'category_b',
                'a_start_ms', 'a_end_ms', 'b_start_ms', 'b_end_ms',
                'overlap_duration_ms', 'overlap_ratio', '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['category_a'],
                    meta['category_b'],
                    meta['a_start_ms'],
                    meta['a_end_ms'],
                    meta['b_start_ms'],
                    meta['b_end_ms'],
                    meta['overlap_duration_ms'],
                    round(meta['overlap_ratio'], 3),
                    meta['total_duration_ms'],
                    meta['actual_duration_s'],
                    meta['question_type']
                ])


def main(config_path: str = None):
    """Main entry point for overlap 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(
        'overlap_task',
        log_file=str(Path(config['output']['base_path']) / config['logging']['log_file']),
        level=config['logging']['level'],
        console_output=config['logging']['console_output']
    )
    
    generator = OverlapTaskGenerator(config, logger)
    generator.generate_dataset()


if __name__ == '__main__':
    main()