tasks/task_temporal_loudness.py

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

Multi-hop inter-task: Combines temporal filtering with loudness comparison.
First applies a before/after condition, then compares volume levels.
"""

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


class TemporalLoudnessTaskGenerator(MultihopBaseGenerator):
    """Generates temporal_loudness task dataset samples."""

    TASK_NAME = "temporal_loudness"

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

        Pipeline:
        1. Build scene with 4-8 events, each at different volume levels
        2. Pick anchor sound
        3. Among events before/after anchor, find loudest/softest
        4. Generate question
        """
        n_events = random.randint(
            self.task_config.get("min_events", 4),
            self.task_config.get("max_events", 8),
        )

        # For repeated_loudness, we need repeated categories
        if question_type == "repeated_loudness":
            n_unique = random.randint(2, max(2, min(n_events - 1, len(self.dataset.CATEGORIES))))
        else:
            n_unique = 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)

        # Generate distinct volume levels for each event
        volume_range = self.task_config.get("volume_range_db", [-12, 6])
        volume_levels = []
        for _ in range(n_events):
            vol = random.uniform(volume_range[0], volume_range[1])
            volume_levels.append(round(vol, 1))

        # Ensure sufficient volume difference
        min_diff_db = self.task_config.get("min_volume_diff_db", 3.0)
        sorted_vols = sorted(volume_levels)
        if len(sorted_vols) > 1 and (sorted_vols[-1] - sorted_vols[0]) < min_diff_db:
            volume_levels[0] = sorted_vols[0] - min_diff_db

        # Build audio
        from pydub import AudioSegment as PydubSegment

        source_files = []
        audio_segments = []
        measured_loudness = []
        per_event_s = max(self.source_clip_duration,
                          (target_duration_seconds or 30) / n_events)

        for i, cat in enumerate(categories):
            fname, fpath = self.dataset.sample_file_from_category(cat)
            audio = self.audio_processor.load_audio(fpath)
            audio = audio.apply_gain(volume_levels[i])
            loudness = get_lufs_loudness(audio)
            audio_segments.append(audio)
            source_files.append(fname)
            measured_loudness.append(loudness)

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

        # Assemble
        result = audio_segments[0]
        events_meta = [
            {"index": 0, "category": categories[0],
             "volume_db": volume_levels[0], "loudness_lufs": measured_loudness[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]
            events_meta.append(
                {"index": i, "category": categories[i],
                 "volume_db": volume_levels[i], "loudness_lufs": measured_loudness[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"])

        # Anchor selection
        anchor_idx = random.randint(1, n_events - 2) if n_events > 2 else 0
        anchor_sound = categories[anchor_idx]

        mcq_data, open_data, q_meta = self._generate_question(
            question_type, categories, volume_levels, measured_loudness,
            anchor_idx, anchor_sound
        )

        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,
            "measured_loudness_lufs": measured_loudness,
            "question_type": question_type,
            "anchor_sound": anchor_sound,
            "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_loudness sample {sample_id}: "
            f"{n_events} events, type={question_type}"
        )
        return metadata

    def _generate_question(
        self, question_type, categories, volume_levels, loudness_levels,
        anchor_idx, anchor_sound
    ):
        """Generate temporal loudness question."""
        n = len(categories)

        if question_type in ("loudest_after", "softest_after"):
            subset = [(i, categories[i], loudness_levels[i])
                       for i in range(anchor_idx + 1, n)]
        elif question_type in ("loudest_before", "softest_before"):
            subset = [(i, categories[i], loudness_levels[i])
                       for i in range(0, anchor_idx)]
        elif question_type == "repeated_loudness":
            # Find categories that appear multiple times
            from collections import Counter
            counts = Counter(categories)
            repeated = [c for c, cnt in counts.items() if cnt >= 2]
            if not repeated:
                return None, None, {}
            target_sound = random.choice(repeated)
            indices = [i for i, c in enumerate(categories) if c == target_sound]

            # Find loudest occurrence
            loudest_idx = max(indices, key=lambda i: loudness_levels[i])
            occurrence_num = indices.index(loudest_idx) + 1  # 1-indexed
            correct = f"occurrence {occurrence_num}"

            # Check trend
            occ_loudness = [loudness_levels[i] for i in indices]
            if all(occ_loudness[i] <= occ_loudness[i + 1]
                   for i in range(len(occ_loudness) - 1)):
                trend = "louder"
            elif all(occ_loudness[i] >= occ_loudness[i + 1]
                     for i in range(len(occ_loudness) - 1)):
                trend = "softer"
            else:
                trend = "neither"

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

            options = [f"occurrence {j + 1}" for j in range(len(indices))]
            while len(options) < 4:
                options.append(f"occurrence {len(options) + 1}")
            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 = {"target_sound": target_sound, "trend": trend}
            return mcq_data, open_data, q_meta
        else:
            return None, None, {}

        if not subset:
            return None, None, {}

        if "loudest" in question_type:
            best = max(subset, key=lambda x: x[2])
        else:
            best = min(subset, key=lambda x: x[2])

        correct_category = best[1]
        present_cats = [s[1] for s in subset]

        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)

        mcq_data = self.question_generator.generate_category_mcq(
            mcq_text, correct_category, present_cats, self.dataset.CATEGORIES
        )
        open_data = self.question_generator.generate_category_open_text(
            open_text, correct_category
        )
        q_meta = {"correct_category": correct_category,
                   "correct_loudness_lufs": best[2]}
        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(
        "temporal_loudness_task",
        log_file=str(Path(config["output"]["base_path"]) / config["logging"]["log_file"]),
        level=config["logging"]["level"],
        console_output=config["logging"]["console_output"],
    )
    generator = TemporalLoudnessTaskGenerator(config, logger)
    generator.generate_dataset()

if __name__ == "__main__":
    main()