tasks/task_conditional_count.py

"""
Conditional Count task generator for temporal reasoning dataset.

Multi-hop task: First apply a temporal condition (before, after, between),
then count events satisfying that condition.
"""

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

from tasks.multihop_base import MultihopBaseGenerator


class ConditionalCountTaskGenerator(MultihopBaseGenerator):
    """Generates conditional_count task dataset samples."""

    TASK_NAME = "conditional_count"

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

        Pipeline:
        1. Build a sequential scene with 4-8 events (some categories repeated)
        2. Pick a random question type (count_after, count_before, count_between)
        3. Select anchor sound(s) and target sound from the sequence
        4. Compute the correct count
        5. Generate MCQ and open-text questions
        """
        n_events = random.randint(
            self.task_config.get("min_events", 4),
            self.task_config.get("max_events", 8),
        )

        # Build scene — we allow repeated categories for counting
        n_unique = random.randint(max(2, n_events // 2), min(n_events, len(self.dataset.CATEGORIES)))
        categories_pool = self.dataset.sample_categories(n_unique)
        # Fill sequence to n_events by repeating some
        categories = list(categories_pool)
        while len(categories) < n_events:
            categories.append(random.choice(categories_pool))
        random.shuffle(categories)

        # Build audio scene
        final_audio, source_files, events_meta = self._build_scene_from_categories(
            categories, target_duration_seconds
        )

        # Save audio
        output_path = self.audio_output / f"{sample_id}.wav"
        final_audio.export(str(output_path), format="wav")

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

        # Generate question
        mcq_data, open_data, answer, q_meta = self._generate_question(
            question_type, categories, events_meta
        )

        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_path.relative_to(self.output_base.parent)),
            "n_events": n_events,
            "categories": categories,
            "source_files": source_files,
            "question_type": question_type,
            "target_duration_s": target_duration_seconds,
            "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_data["question"],
            "open_text_answer": open_data["correct_answer"],
            **q_meta,
        }

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

    def _build_scene_from_categories(
        self, categories: List[str], target_duration_s: float
    ):
        """Build audio scene from pre-specified category sequence."""
        from pydub import AudioSegment as PydubSegment
        from utils import concatenate_to_target_duration, generate_controlled_gap_durations

        n_events = len(categories)
        num_gaps = n_events - 1

        # Estimate per-event duration
        gap_total_est = num_gaps * 500  # ~500ms avg gap
        avail = max(n_events * self.source_clip_duration, target_duration_s - gap_total_est / 1000)
        per_event_s = max(self.source_clip_duration, avail / n_events)

        source_files = []
        audio_segments = []
        for cat in categories:
            fname, fpath = self.dataset.sample_file_from_category(cat)
            audio = self.audio_processor.load_audio(fpath)
            audio_segments.append(audio)
            source_files.append(fname)

        # Gaps
        if num_gaps > 0:
            gap_durations = generate_controlled_gap_durations(
                num_gaps, min_gap_ms=300, max_gap_ms=1500, gap_multiplier=2.0
            )
        else:
            gap_durations = []

        # Assemble
        result = audio_segments[0]
        current_ms = 0
        events_meta = [
            {
                "index": 0,
                "category": categories[0],
                "start_ms": 0,
                "end_ms": len(audio_segments[0]),
                "duration_ms": len(audio_segments[0]),
            }
        ]
        current_ms = len(audio_segments[0])

        for i in range(1, n_events):
            gap_ms = gap_durations[i - 1]
            result = result + PydubSegment.silent(duration=gap_ms)
            current_ms += gap_ms
            event_start = current_ms
            result = result + audio_segments[i]
            current_ms += len(audio_segments[i])
            events_meta.append(
                {
                    "index": i,
                    "category": categories[i],
                    "start_ms": event_start,
                    "end_ms": current_ms,
                    "duration_ms": len(audio_segments[i]),
                    "gap_before_ms": gap_ms,
                }
            )

        return result, source_files, events_meta

    def _generate_question(
        self,
        question_type: str,
        categories: List[str],
        events_meta: List[Dict],
    ):
        """Generate question and compute answer for conditional count."""
        unique_cats = list(set(categories))
        n = len(categories)

        if question_type in ("count_after", "count_before"):
            # Pick anchor and target
            anchor_idx = random.randint(1, n - 2) if n > 2 else 0
            anchor_sound = categories[anchor_idx]

            # Pick target sound (one that appears in the sequence)
            target_sound = random.choice(unique_cats)

            if question_type == "count_after":
                count = sum(1 for i in range(anchor_idx + 1, n) if categories[i] == target_sound)
            else:
                count = sum(1 for i in range(0, anchor_idx) if categories[i] == target_sound)

            # Format question
            mcq_templates = self.task_config["mcq_questions"][question_type]
            open_templates = self.task_config["open_text_questions"][question_type]
            mcq_text = mcq_templates.format(
                target_sound=target_sound, anchor_sound=anchor_sound
            )
            open_text = open_templates.format(
                target_sound=target_sound, anchor_sound=anchor_sound
            )

            mcq_data = self.question_generator.generate_count_mcq(
                mcq_text, count, self.dataset.CATEGORIES
            )
            open_data = self.question_generator.generate_count_open_text(open_text, count)

            q_meta = {
                "anchor_sound": anchor_sound,
                "anchor_index": anchor_idx,
                "target_sound": target_sound,
                "correct_count": count,
            }
            return mcq_data, open_data, count, q_meta

        elif question_type == "count_between":
            if n < 3:
                return None, None, None, {}
            # Pick two anchor indices
            idx1, idx2 = sorted(random.sample(range(n), 2))
            sound1 = categories[idx1]
            sound2 = categories[idx2]

            # Count events strictly between idx1 and idx2
            between_cats = categories[idx1 + 1 : idx2]
            # Pick a target — or count all
            if random.random() < 0.5 and between_cats:
                target_sound = random.choice(list(set(between_cats)))
                count = between_cats.count(target_sound)
                mcq_text = self.task_config["mcq_questions"]["count_between"].format(target_sound=target_sound, sound1=sound1, sound2=sound2)
                open_text = self.task_config["open_text_questions"]["count_between"].format(target_sound=target_sound, sound1=sound1, sound2=sound2)
            else:
                target_sound = "all"
                count = len(between_cats)
                mcq_text = f"How many sounds occur between {sound1} and {sound2}?"
                open_text = f"How many sounds occur between {sound1} and {sound2}?"

            mcq_data = self.question_generator.generate_count_mcq(
                mcq_text, count, self.dataset.CATEGORIES
            )
            open_data = self.question_generator.generate_count_open_text(open_text, count)

            q_meta = {
                "sound1": sound1,
                "sound2": sound2,
                "target_sound": target_sound,
                "correct_count": count,
            }
            return mcq_data, open_data, count, q_meta

        elif question_type == "count_during":
            # Simplified: pick anchor, count other events (all overlap conceptually
            # since events are sequential — we interpret "during" as overlapping time
            # window centered on the anchor)
            anchor_idx = random.randint(0, n - 1)
            anchor_sound = categories[anchor_idx]
            # For sequential audio, "during" is approximated as the anchor's neighbors
            # We count how many of target_sound appear adjacent to anchor
            target_sound = random.choice(unique_cats)
            count = categories.count(target_sound) - (1 if target_sound == anchor_sound else 0)
            count = max(0, count)

            mcq_text = self.task_config["mcq_questions"]["count_during"].format(target_sound=target_sound, anchor_sound=anchor_sound)
            open_text = self.task_config["open_text_questions"]["count_during"].format(target_sound=target_sound, anchor_sound=anchor_sound)

            mcq_data = self.question_generator.generate_count_mcq(
                mcq_text, count, self.dataset.CATEGORIES
            )
            open_data = self.question_generator.generate_count_open_text(open_text, count)

            q_meta = {
                "anchor_sound": anchor_sound,
                "target_sound": target_sound,
                "correct_count": count,
            }
            return mcq_data, open_data, count, q_meta

        elif question_type == "count_overlap":
            # In sequential audio, overlap count is 0 by construction unless
            # we build overlap scenes. Simplify: count events near anchor.
            anchor_idx = random.randint(0, n - 1)
            anchor_sound = categories[anchor_idx]
            # Count adjacent events as "overlapping" conceptually
            neighbors = set()
            if anchor_idx > 0:
                neighbors.add(anchor_idx - 1)
            if anchor_idx < n - 1:
                neighbors.add(anchor_idx + 1)
            count = len(neighbors)

            mcq_text = self.task_config["mcq_questions"]["count_overlap"].format(anchor_sound=anchor_sound, target_sound=anchor_sound)
            open_text = self.task_config["open_text_questions"]["count_overlap"].format(anchor_sound=anchor_sound, target_sound=anchor_sound)

            mcq_data = self.question_generator.generate_count_mcq(
                mcq_text, count, self.dataset.CATEGORIES
            )
            open_data = self.question_generator.generate_count_open_text(open_text, count)

            q_meta = {"anchor_sound": anchor_sound, "correct_count": count}
            return mcq_data, open_data, count, q_meta

        return None, None, None, {}


def main(config_path: str = None):
    """Main entry point for conditional_count 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(
        "conditional_count_task",
        log_file=str(Path(config["output"]["base_path"]) / config["logging"]["log_file"]),
        level=config["logging"]["level"],
        console_output=config["logging"]["console_output"],
    )

    generator = ConditionalCountTaskGenerator(config, logger)
    generator.generate_dataset()


if __name__ == "__main__":
    main()