tasks/task_multi_hop.py

"""
Multi-Hop task generator for temporal reasoning dataset.

Multi-hop inter-task: Explicit two-step or three-step temporal reasoning
over combinations of order, duration, volume, silence, and count.

Example chains:
  - "What sound occurs after the longest sound?" (order + duration)
  - "What sound occurs before the loudest sound?" (order + volume)
  - "How many sounds occur after the longest silence?" (count + silence)
"""

import random
from collections import Counter
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,
    get_lufs_loudness,
    create_preprocessed_dataset,
)
from tasks.multihop_base import MultihopBaseGenerator


class MultiHopTaskGenerator(MultihopBaseGenerator):
    """Generates multi_hop task dataset samples."""

    TASK_NAME = "multi_hop"

    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 multi_hop sample.

        Builds a rich scene with duration, volume, and silence variation,
        then asks questions requiring chained reasoning.
        """
        n_events = random.randint(
            self.task_config.get("min_events", 5),
            self.task_config.get("max_events", 8),
        )

        # Allow some repeats for count_before_loudest / count_after_longest
        n_unique = random.randint(
            max(3, n_events // 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)

        # Volume variation
        volume_range = self.task_config.get("volume_range_db", [-10, 6])
        volume_levels = [
            round(random.uniform(volume_range[0], volume_range[1]), 1)
            for _ in range(n_events)
        ]

        # Ensure min volume difference
        min_diff = self.task_config.get("min_volume_diff_db", 4.0)
        vmin, vmax = min(volume_levels), max(volume_levels)
        if vmax - vmin < min_diff:
            idx_min = volume_levels.index(vmin)
            volume_levels[idx_min] = vmax - min_diff

        from pydub import AudioSegment as PydubSegment

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

        for i, cat in enumerate(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 = audio.apply_gain(volume_levels[i])
            loudness = get_lufs_loudness(audio)
            audio_segments.append(audio)
            source_files.append(fname)
            effective_durations_ms.append(int(eff_dur * 1000))
            measured_loudness.append(loudness)

        # Gaps
        num_gaps = n_events - 1
        gap_min = self.task_config.get("min_gap_ms", 500)
        gap_max = self.task_config.get("max_gap_ms", 3000)
        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.5,
            )
        else:
            gap_durations = []

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

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

        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, effective_durations_ms,
            measured_loudness, volume_levels, 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,
            "volume_levels_db": volume_levels,
            "effective_durations_ms": effective_durations_ms,
            "measured_loudness_lufs": measured_loudness,
            "gap_durations_ms": gap_durations,
            "question_type": question_type,
            "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 multi_hop sample {sample_id}: "
            f"{n_events} events, type={question_type}"
        )
        return metadata

    def _generate_question(
        self, question_type, categories, durations_ms,
        loudness_levels, volume_levels, gap_durations
    ):
        """Generate multi-hop chained reasoning question."""
        n = len(categories)

        # ---- Step 1: Identify the "pivot" event based on first property ----

        if question_type in ("after_longest", "before_longest"):
            pivot_idx = durations_ms.index(max(durations_ms))
            pivot_label = "longest sound"
        elif question_type == "after_shortest":
            pivot_idx = durations_ms.index(min(durations_ms))
            pivot_label = "shortest sound"
        elif question_type == "before_loudest":
            pivot_idx = loudness_levels.index(max(loudness_levels))
            pivot_label = "loudest sound"
        elif question_type == "after_longest_gap":
            if not gap_durations:
                return None, None, {}
            longest_gap_idx = gap_durations.index(max(gap_durations))
            # Events after the longest gap
            after_idx = longest_gap_idx + 1
            events_after = categories[after_idx:]
            count_after = len(events_after)

            # Pick template variant — count or identify
            if random.random() < 0.5 and events_after:
                # Identify first event after longest gap
                correct = categories[after_idx]
                mcq_text = self.task_config["mcq_questions"]["after_longest_gap"]
                open_text = self.task_config["open_text_questions"]["after_longest_gap"]
                # If template asks "how many" use count, else use category
                if "how many" in mcq_text.lower():
                    mcq_data = self.question_generator.generate_count_mcq(
                        mcq_text, count_after, self.dataset.CATEGORIES
                    )
                    open_data = self.question_generator.generate_count_open_text(
                        open_text, count_after
                    )
                else:
                    mcq_data = self.question_generator.generate_category_mcq(
                        mcq_text, correct, categories, self.dataset.CATEGORIES
                    )
                    open_data = self.question_generator.generate_category_open_text(
                        open_text, correct
                    )
                q_meta = {"longest_gap_idx": longest_gap_idx, "correct_value": correct,
                          "count_after_longest_gap": count_after}
                return mcq_data, open_data, q_meta
            else:
                mcq_text = f"How many sounds occur after the longest silence?"
                open_text = f"How many sounds occur after the longest silence?"
                mcq_data = self.question_generator.generate_count_mcq(
                    mcq_text, count_after, self.dataset.CATEGORIES
                )
                open_data = self.question_generator.generate_count_open_text(
                    open_text, count_after
                )
                q_meta = {"longest_gap_idx": longest_gap_idx,
                          "count_after_longest_gap": count_after}
                return mcq_data, open_data, q_meta

        elif question_type in ("count_before_loudest", "count_after_longest"):
            if question_type == "count_before_loudest":
                pivot_idx = loudness_levels.index(max(loudness_levels))
                region = categories[:pivot_idx]
            else:
                pivot_idx = durations_ms.index(max(durations_ms))
                region = categories[pivot_idx + 1:]

            # Count occurrences of a target sound in region
            unique_in_region = list(set(region))
            if unique_in_region:
                target_sound = random.choice(unique_in_region)
            else:
                target_sound = random.choice(list(set(categories)))
            count = region.count(target_sound)

            mcq_text = self.task_config["mcq_questions"][question_type].format(target_sound=target_sound)
            open_text = self.task_config["open_text_questions"][question_type].format(target_sound=target_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 = {"target_sound": target_sound, "correct_count": count,
                      "pivot_index": pivot_idx}
            return mcq_data, open_data, q_meta

        elif question_type in ("overlap_after_anchor", "loudest_after_anchor",
                                "longest_before_anchor"):
            # Pick anchor
            anchor_idx = random.randint(1, n - 2) if n > 2 else 0
            anchor_sound = categories[anchor_idx]

            if question_type == "loudest_after_anchor":
                subset = [(i, categories[i], loudness_levels[i])
                          for i in range(anchor_idx + 1, n)]
                if not subset:
                    return None, None, {}
                best = max(subset, key=lambda x: x[2])
                correct = best[1]
            elif question_type == "longest_before_anchor":
                subset = [(i, categories[i], durations_ms[i])
                          for i in range(0, anchor_idx)]
                if not subset:
                    return None, None, {}
                best = max(subset, key=lambda x: x[2])
                correct = best[1]
            elif question_type == "overlap_after_anchor":
                # In sequential audio, no true overlap — use event right after anchor
                if anchor_idx + 1 < n:
                    correct = categories[anchor_idx + 1]
                else:
                    return None, None, {}
            else:
                return None, None, {}

            mcq_text = self.task_config["mcq_questions"][question_type].format(anchor_sound=anchor_sound)
            open_text = self.task_config["open_text_questions"][question_type].format(anchor_sound=anchor_sound)

            present = categories[anchor_idx + 1:] if "after" in question_type else categories[:anchor_idx]
            mcq_data = self.question_generator.generate_category_mcq(
                mcq_text, correct, present or categories, self.dataset.CATEGORIES
            )
            open_data = self.question_generator.generate_category_open_text(
                open_text, correct
            )
            q_meta = {"anchor_sound": anchor_sound, "correct_category": correct}
            return mcq_data, open_data, q_meta

        else:
            # Default: not one of the special cases above
            return None, None, {}

        # ---- Step 2: Get the event after/before the pivot ----
        if "after" in question_type:
            target_idx = pivot_idx + 1
            if target_idx >= n:
                return None, None, {}
        else:  # "before"
            target_idx = pivot_idx - 1
            if target_idx < 0:
                return None, None, {}

        correct = categories[target_idx]

        mcq_text = self.task_config["mcq_questions"][question_type]
        open_text = self.task_config["open_text_questions"][question_type]

        mcq_data = self.question_generator.generate_category_mcq(
            mcq_text, correct, categories, self.dataset.CATEGORIES
        )
        open_data = self.question_generator.generate_category_open_text(
            open_text, correct
        )
        q_meta = {
            "pivot_index": pivot_idx,
            "pivot_category": categories[pivot_idx],
            "correct_category": correct,
        }
        return mcq_data, open_data, q_meta


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(
        "multi_hop_task",
        log_file=str(Path(config["output"]["base_path"]) / config["logging"]["log_file"]),
        level=config["logging"]["level"],
        console_output=config["logging"]["console_output"],
    )
    generator = MultiHopTaskGenerator(config, logger)
    generator.generate_dataset()

if __name__ == "__main__":
    main()