"""spk_preprocessor_lora.py — Custom SPK preprocessor for LoRA training.

Problem: ESPnet's SpkPreprocessor._speech_process() unconditionally accesses
data["speech2"] in eval mode (self.train=False). During collect_stats,
ESPnet uses train=False for ALL data (including train data), which only has
"speech" (no "speech2"). This causes KeyError: 'speech2'.

Fix: Override _speech_process to check for "speech2" presence:
  - If speech2 present (trial pairs / actual validation): normal multi-crop eval.
  - If speech2 absent (train data during collect_stats): single-crop like train mode.

This survives ESPnet reinstalls (unlike patching the installed package).
"""
import random
import numpy as np
from espnet2.train.preprocessor import SpkPreprocessor


class SpkLoRAPreprocessor(SpkPreprocessor):
    """SpkPreprocessor that handles missing speech2 in collect_stats mode."""

    def _speech_process(self, data):
        if self.train:
            # Normal train mode: single speech, random crop
            return super()._speech_process(data)

        # Eval mode (self.train=False) — used by collect_stats for ALL data
        if "speech2" not in data:
            # Train data during collect_stats: no speech2, do single crop
            audio = data["speech"]
            if len(audio) < self.target_duration:
                shortage = self.target_duration - len(audio) + 1
                audio = np.pad(audio, (0, shortage), "wrap")
            startframe = np.int64(
                random.random() * max(0, len(audio) - self.target_duration)
            )
            data["speech"] = audio[
                int(startframe) : int(startframe) + self.target_duration
            ]
            return data
        else:
            # Trial pair (actual validation): normal multi-crop for both speech+speech2
            return super()._speech_process(data)