tasks/task_temporal_arithmetic.py

"""
Temporal Arithmetic task generator for temporal reasoning dataset.

Multi-hop inter-task: Combines duration, count/repetition, and silence
reasoning by comparing total active time across labels or vs silence.
"""

import random
from pathlib import Path
from typing import Dict, List, Optional

from utils import (
    setup_logger,
    set_random_seed,
    concatenate_to_target_duration,
    generate_controlled_gap_durations,
    create_preprocessed_dataset,
)
from tasks.multihop_base import MultihopBaseGenerator


class TemporalArithmeticTaskGenerator(MultihopBaseGenerator):
    """Generates temporal_arithmetic task dataset samples."""

    TASK_NAME = "temporal_arithmetic"

    def __init__(self, config: dict, logger=None):
        super().__init__(config, logger)
        preprocessed_path = self.task_config.get(
            "preprocessed_data_path",
            config["tasks"].get("duration", {}).get("preprocessed_data_path", ""),
        )
        self.dataset = create_preprocessed_dataset(config, preprocessed_path=preprocessed_path)

    def generate_sample(
        self,
        sample_id: int,
        target_duration_seconds: float = None,
        question_type: str = None,
    ) -> Optional[Dict]:
        """
        Generate a single temporal_arithmetic sample.

        Pipeline:
        1. Build scene with 4-8 events (some repeated categories)
        2. Track total duration per label and total silence
        3. Generate arithmetic comparison questions
        """
        n_events = random.randint(
            self.task_config.get("min_events", 4),
            self.task_config.get("max_events", 8),
        )

        # Allow repeated categories for aggregation
        n_unique = random.randint(2, min(n_events, len(self.dataset.CATEGORIES)))
        categories_pool = self.dataset.sample_categories(n_unique)
        categories = list(categories_pool)
        while len(categories) < n_events:
            categories.append(random.choice(categories_pool))
        random.shuffle(categories)

        from pydub import AudioSegment as PydubSegment

        source_files = []
        audio_segments = []
        effective_durations_ms = []

        for cat in categories:
            fname, fpath, eff_dur = self.dataset.sample_file_from_category_with_duration(cat)
            audio = self.audio_processor.load_audio(fpath)
            audio = concatenate_to_target_duration(audio, max(self.source_clip_duration, eff_dur))
            audio_segments.append(audio)
            source_files.append(fname)
            effective_durations_ms.append(int(eff_dur * 1000))

        # Generate gaps
        num_gaps = n_events - 1
        gap_min = self.task_config.get("min_gap_ms", 300)
        gap_max = self.task_config.get("max_gap_ms", 2000)
        if num_gaps > 0:
            gap_durations = generate_controlled_gap_durations(
                num_gaps, min_gap_ms=gap_min, max_gap_ms=gap_max, gap_multiplier=2.0
            )
        else:
            gap_durations = []

        # Assemble
        result = audio_segments[0]
        for i in range(1, n_events):
            gap_ms = gap_durations[i - 1]
            result = result + PydubSegment.silent(duration=gap_ms)
            result = result + audio_segments[i]

        output_path = self.audio_output / f"{sample_id}.wav"
        result.export(str(output_path), format="wav")

        # Compute aggregates
        # Total duration per label
        label_totals = {}
        for cat, dur in zip(categories, effective_durations_ms):
            label_totals[cat] = label_totals.get(cat, 0) + dur

        total_silence_ms = sum(gap_durations) if gap_durations else 0
        longest_silence_ms = max(gap_durations) if gap_durations else 0

        if question_type is None:
            question_type = random.choice(self.task_config["question_types"])

        mcq_data, open_data, q_meta = self._generate_question(
            question_type, categories, label_totals, total_silence_ms,
            longest_silence_ms, gap_durations
        )

        if mcq_data is None:
            return None

        metadata = {
            "id": sample_id,
            "audio_path": str(output_path.relative_to(self.output_base.parent)),
            "n_events": n_events,
            "categories": categories,
            "source_files": source_files,
            "question_type": question_type,
            "label_totals_ms": label_totals,
            "total_silence_ms": total_silence_ms,
            "gap_durations_ms": gap_durations,
            "target_duration_s": target_duration_seconds,
            "actual_duration_s": len(result) / 1000.0,
            "mcq_question": mcq_data["question"],
            "mcq_options": mcq_data["options"],
            "mcq_correct_answer": mcq_data["correct_answer"],
            "open_text_question": open_data["question"],
            "open_text_answer": open_data["correct_answer"],
            **q_meta,
        }

        self.logger.info(
            f"Generated temporal_arithmetic sample {sample_id}: "
            f"{n_events} events, type={question_type}"
        )
        return metadata

    def _generate_question(
        self, question_type, categories, label_totals,
        total_silence_ms, longest_silence_ms, gap_durations
    ):
        """Generate temporal arithmetic question."""
        unique_cats = list(label_totals.keys())

        if question_type == "total_label_duration":
            if len(unique_cats) < 2:
                return None, None, {}
            sound1, sound2 = random.sample(unique_cats, 2)
            dur1 = label_totals[sound1]
            dur2 = label_totals[sound2]

            correct = sound1 if dur1 >= dur2 else sound2

            mcq_text = self.task_config["mcq_questions"]["total_label_duration"].format(sound1=sound1, sound2=sound2)
            open_text = self.task_config["open_text_questions"]["total_label_duration"].format(sound1=sound1, sound2=sound2)

            mcq_data = self.question_generator.generate_pairwise_comparison_mcq(
                mcq_text, correct, (sound1, sound2), self.dataset.CATEGORIES
            )
            open_data = self.question_generator.generate_category_open_text(
                open_text, correct
            )
            q_meta = {"dur1_ms": dur1, "dur2_ms": dur2}
            return mcq_data, open_data, q_meta

        elif question_type == "label_vs_silence":
            target_sound = random.choice(unique_cats)
            label_dur = label_totals[target_sound]

            if label_dur >= total_silence_ms:
                correct = target_sound
            else:
                correct = "total silence"

            mcq_text = self.task_config["mcq_questions"]["label_vs_silence"].format(target_sound=target_sound)
            open_text = self.task_config["open_text_questions"]["label_vs_silence"].format(target_sound=target_sound)

            options = [target_sound, "total silence"]
            other = [c for c in self.dataset.CATEGORIES if c != target_sound][:2]
            options.extend(other)
            random.shuffle(options)
            options = options[:4]
            if correct not in options:
                options[0] = correct

            option_labels = ["A", "B", "C", "D"]
            correct_label = option_labels[options.index(correct)]
            option_map = {l: v for l, v in zip(option_labels, options)}

            mcq_data = {"question": mcq_text, "options": option_map,
                        "correct_answer": correct_label}
            open_data = {"question": open_text, "correct_answer": correct}
            q_meta = {"label_duration_ms": label_dur, "total_silence_ms": total_silence_ms}
            return mcq_data, open_data, q_meta

        elif question_type == "combined_duration":
            target_sound = random.choice(unique_cats)
            combined_dur = label_totals[target_sound]

            if combined_dur >= longest_silence_ms:
                correct = f"all {target_sound} sounds combined"
            else:
                correct = "the longest silence"

            mcq_text = self.task_config["mcq_questions"]["combined_duration"].format(target_sound=target_sound, sound1=target_sound,
                     sound2=random.choice([c for c in unique_cats if c != target_sound] or unique_cats))
            open_text = self.task_config["open_text_questions"]["combined_duration"].format(target_sound=target_sound, sound1=target_sound,
                     sound2=random.choice([c for c in unique_cats if c != target_sound] or unique_cats))

            options = [f"all {target_sound} sounds combined", "the longest silence"]
            other = [c for c in self.dataset.CATEGORIES if c != target_sound][:2]
            options.extend(other)
            random.shuffle(options)
            options = options[:4]
            if correct not in options:
                options[0] = correct

            option_labels = ["A", "B", "C", "D"]
            correct_label = option_labels[options.index(correct)]
            option_map = {l: v for l, v in zip(option_labels, options)}

            mcq_data = {"question": mcq_text, "options": option_map,
                        "correct_answer": correct_label}
            open_data = {"question": open_text, "correct_answer": correct}
            q_meta = {"combined_duration_ms": combined_dur,
                      "longest_silence_ms": longest_silence_ms}
            return mcq_data, open_data, q_meta

        return None, None, {}


def main(config_path: str = None):
    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(
        "temporal_arithmetic_task",
        log_file=str(Path(config["output"]["base_path"]) / config["logging"]["log_file"]),
        level=config["logging"]["level"],
        console_output=config["logging"]["console_output"],
    )
    generator = TemporalArithmeticTaskGenerator(config, logger)
    generator.generate_dataset()

if __name__ == "__main__":
    main()