tasks/task_silence_gap.py

"""
Silence Gap task generator for temporal reasoning dataset.

Generates audio samples where sequential sound clips are separated by silences
of varying durations, and asks questions about those silence gaps (longest gap,
shortest gap, which sound follows the longest silence, etc.).

Uses preprocessed ESC-50 clips (trimmed to active audio regions) so that
embedded silences in the source clips don't interfere with the controlled gaps.
"""

import csv
import json
import random
from collections import Counter
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,
    generate_controlled_gap_durations,
    build_silence_gap_task_audio,
    create_preprocessed_dataset
)


class SilenceGapTaskGenerator:
    """Generates silence_gap task dataset samples."""
    
    def __init__(self, config: dict, logger=None):
        """
        Initialize SilenceGapTaskGenerator.
        
        Args:
            config: Full pipeline configuration dictionary
            logger: Logger instance
        """
        self.config = config
        self.logger = logger or setup_logger(__name__)
        
        self.task_config = config['tasks']['silence_gap']
        
        # 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.min_clips = self.task_config.get('min_clips_per_sample', 3)
        self.max_clips = self.task_config.get('max_clips_per_sample', 10)
        self.min_gap_ms = self.task_config.get('min_gap_duration_ms', 500)
        self.max_gap_ms = self.task_config.get('max_gap_duration_ms', 3000)
        self.gap_multiplier = self.task_config.get('gap_multiplier', 2.0)
        
        # Preprocessed data path
        preprocessed_path = self.task_config.get(
            'preprocessed_data_path',
            config['tasks'].get('duration', {}).get('preprocessed_data_path', '')
        )
        
        # Dataset
        self.dataset = create_preprocessed_dataset(config, preprocessed_path=preprocessed_path)
        
        # 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,  # We control silences manually
            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']) / 'silence_gap'
        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 silence_gap task sample.
        
        Pipeline:
        1. Sample N categories (min_clips to max_clips)
        2. Load preprocessed (trimmed) audio for each
        3. Extend each clip to a reasonable segment duration
        4. Generate controlled gap durations with multiplier constraint
        5. Build final audio with controlled gaps
        6. Generate questions about the gaps
        
        Args:
            sample_id: Sample ID
            target_duration_seconds: Pre-generated target duration
            
        Returns:
            Metadata dictionary, or None if failed
        """
        # Step 1: Determine clip duration and number of clips
        if target_duration_seconds is not None:
            clip_duration_s = target_duration_seconds
        else:
            clip_duration_s = generate_single_clip_duration(
                self.min_clip_duration, self.max_clip_duration
            )
        
        n_clips = random.randint(self.min_clips, self.max_clips)
        n_clips = min(n_clips, len(self.dataset.CATEGORIES))  # Can't exceed available categories
        
        # Step 2: Sample categories
        try:
            categories = self.dataset.sample_categories(n_clips)
        except ValueError:
            self.logger.warning(f"Sample {sample_id}: Cannot sample {n_clips} categories")
            return None
        
        # Step 3: Load preprocessed audio and extend clips
        audio_segments = []
        source_files = []
        
        # Calculate approximate duration per clip
        # Total = sum(clip_durations) + sum(gaps)
        # Rough estimate: allocate ~60% to clips, ~40% to gaps
        estimated_gap_total_ms = n_clips * (self.min_gap_ms + self.max_gap_ms) / 2
        available_for_clips_s = clip_duration_s - (estimated_gap_total_ms / 1000.0)
        per_clip_s = max(self.source_clip_duration, available_for_clips_s / n_clips)
        
        for category in categories:
            filename, filepath, eff_dur = self.dataset.sample_file_from_category_with_duration(
                category
            )
            audio = self.audio_processor.load_audio(filepath)
            
            # Extend clip to per_clip_s by repeating
            audio_segments.append(audio)
            source_files.append(filename)
        
        # Step 4: Generate controlled gap durations
        num_gaps = n_clips - 1
        gap_durations = generate_controlled_gap_durations(
            num_gaps,
            min_gap_ms=self.min_gap_ms,
            max_gap_ms=self.max_gap_ms,
            gap_multiplier=self.gap_multiplier
        )
        
        # Step 5: Build final audio
        final_audio, _, build_metadata = build_silence_gap_task_audio(
            audio_segments,
            categories,
            gap_durations,
            target_duration_seconds=clip_duration_s
        )
        
        # 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
        gap_info = build_metadata['gap_info']
        longest_idx = build_metadata['longest_gap_idx']
        shortest_idx = build_metadata['shortest_gap_idx']
        
        longest_gap = gap_info[longest_idx]
        shortest_gap = gap_info[shortest_idx]
        
        # Select a random question type
        question_type = random.choice(self.task_config['question_types'])
        
        mcq_data, open_text_data = self._generate_question(
            question_type, categories, gap_info, longest_idx, shortest_idx
        )
        
        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)),
            'n_clips': n_clips,
            'categories': categories,
            'source_files': source_files,
            'question_type': question_type,
            'gap_durations_ms': gap_durations,
            'longest_gap_idx': longest_idx,
            'shortest_gap_idx': shortest_idx,
            'longest_gap_ms': longest_gap['gap_duration_ms'],
            'shortest_gap_ms': shortest_gap['gap_duration_ms'],
            'longest_gap_pair': (longest_gap['left_category'], longest_gap['right_category']),
            'shortest_gap_pair': (shortest_gap['left_category'], shortest_gap['right_category']),
            'target_duration_s': clip_duration_s,
            '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 silence_gap sample {sample_id}: {n_clips} clips, "
            f"gaps={gap_durations}, type={question_type}"
        )
        
        return metadata
    
    def _generate_question(
        self,
        question_type: str,
        categories: List[str],
        gap_info: List[Dict],
        longest_idx: int,
        shortest_idx: int
    ) -> Tuple[Optional[Dict], Optional[Dict]]:
        """Generate MCQ and open-text question based on type."""
        
        longest_gap = gap_info[longest_idx]
        shortest_gap = gap_info[shortest_idx]
        
        mcq_template = self.task_config['mcq_questions'].get(question_type, '')
        open_template = self.task_config['open_text_questions'].get(question_type, '')
        
        if question_type == 'longest_gap':
            correct_pair = (longest_gap['left_category'], longest_gap['right_category'])
            mcq_data = self.question_generator.generate_pair_mcq(
                mcq_template, correct_pair, categories, self.dataset.CATEGORIES
            )
            open_data = self.question_generator.generate_pair_open_text(
                open_template, correct_pair
            )
        
        elif question_type == 'shortest_gap':
            correct_pair = (shortest_gap['left_category'], shortest_gap['right_category'])
            mcq_data = self.question_generator.generate_pair_mcq(
                mcq_template, correct_pair, categories, self.dataset.CATEGORIES
            )
            open_data = self.question_generator.generate_pair_open_text(
                open_template, correct_pair
            )
        
        elif question_type == 'after_gap':
            correct_category = longest_gap['right_category']
            mcq_data = self.question_generator.generate_category_mcq(
                mcq_template, correct_category, categories, self.dataset.CATEGORIES
            )
            open_data = self.question_generator.generate_category_open_text(
                open_template, correct_category
            )
        
        elif question_type == 'before_gap':
            correct_category = longest_gap['left_category']
            mcq_data = self.question_generator.generate_category_mcq(
                mcq_template, correct_category, categories, self.dataset.CATEGORIES
            )
            open_data = self.question_generator.generate_category_open_text(
                open_template, correct_category
            )
        
        elif question_type == 'compare_gaps':
            # Pick two gaps and compare
            if len(gap_info) < 2:
                return None, None
            
            gap_indices = random.sample(range(len(gap_info)), 2)
            gap_a, gap_b = gap_info[gap_indices[0]], gap_info[gap_indices[1]]
            sound1 = gap_a['left_category']
            sound2 = gap_b['left_category']
            
            # Which gap is longer?
            if gap_a['gap_duration_ms'] >= gap_b['gap_duration_ms']:
                correct_answer = f"after {sound1}"
            else:
                correct_answer = f"after {sound2}"
            
            formatted_mcq = mcq_template.format(sound1=sound1, sound2=sound2)
            formatted_open = open_template.format(sound1=sound1, sound2=sound2)
            
            # For compare_gaps, use category MCQ with "after X" as options
            options = [f"after {sound1}", f"after {sound2}"]
            other_cats = [c for c in categories if c != sound1 and c != sound2]
            for c in other_cats:
                if len(options) < 4:
                    options.append(f"after {c}")
            
            # Fill with distractors from all categories if still < 4
            all_cats = [c for c in self.dataset.CATEGORIES if f"after {c}" not in options]
            random.shuffle(all_cats)
            for c in all_cats:
                if len(options) < 4:
                    options.append(f"after {c}")
                else:
                    break
                    
            random.shuffle(options)
            
            option_labels = ['A', 'B', 'C', 'D']
            correct_label = option_labels[options.index(correct_answer)]
            option_map = {label: val for label, val in zip(option_labels, options)}
            
            mcq_data = {
                'question': formatted_mcq,
                'options': option_map,
                'correct_answer': correct_label,
                'correct_value': correct_answer
            }
            open_data = {
                'question': formatted_open,
                'correct_answer': correct_answer
            }
        
        else:
            return None, None
        
        return mcq_data, open_data
    
    def generate_dataset(self) -> tuple:
        """Generate the complete silence_gap 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} silence_gap 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 / 'silence_gap_mcq.csv'
        self._save_mcq_csv(all_metadata, mcq_csv_path)
        
        open_text_csv_path = self.output_base / 'silence_gap_open_text.csv'
        self._save_open_text_csv(all_metadata, open_text_csv_path)
        
        metadata_csv_path = self.output_base / 'silence_gap_metadata.csv'
        self._save_metadata_csv(all_metadata, metadata_csv_path)
        
        self.logger.info(f"Silence gap 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', 'gap_durations_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']),
                    str(meta['gap_durations_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', 'gap_durations_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']),
                    str(meta['gap_durations_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', 'n_clips',
                'source_files', 'source_categories', 'gap_durations_ms',
                'longest_gap_idx', 'shortest_gap_idx', 'longest_gap_ms', 'shortest_gap_ms',
                'longest_gap_pair', 'shortest_gap_pair',
                'target_duration_s', 'actual_duration_s', 'question_type'
            ])
            
            for meta in metadata_list:
                writer.writerow([
                    meta['id'],
                    meta['audio_path'],
                    meta['n_clips'],
                    str(meta['source_files']),
                    str(meta['categories']),
                    str(meta['gap_durations_ms']),
                    meta['longest_gap_idx'],
                    meta['shortest_gap_idx'],
                    meta['longest_gap_ms'],
                    meta['shortest_gap_ms'],
                    str(meta['longest_gap_pair']),
                    str(meta['shortest_gap_pair']),
                    meta['target_duration_s'],
                    meta['actual_duration_s'],
                    meta['question_type']
                ])


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


if __name__ == '__main__':
    main()