main_app.py

import sys
import os
from pathlib import Path
import shutil
import time
import numpy as np
import librosa
from pydub import AudioSegment, silence # Keep for merging/exporting, not for VC pre-chunking
from moviepy.editor import VideoFileClip # AudioFileClip, CompositeAudioClip, concatenate_audioclips # Not used directly for now
import soundfile as sf
import yaml
import torch
import torchaudio
from tqdm import tqdm # For console progress during dev, GUI will use progress bar
import ffmpeg 
from huggingface_hub import hf_hub_download
import gc #For garbage collection
import subprocess # For direct ffmpeg calls if needed, though moviepy/ffmpeg-python is preferred

from demucs.pretrained import get_model as get_demucs_model
from demucs.apply import apply_model as apply_demucs_model
from demucs.audio import AudioFile as DemucsAudioFile

# --- Gdown Import for Google Drive downloads ---
try:
    import gdown
    GDOWN_AVAILABLE = True
except ImportError:
    GDOWN_AVAILABLE = False
    print("Warning: gdown library not found. Automatic download of Speech2RIR model will be disabled.")
    print("Install with: pip install gdown")

try:
    from scenedetect import open_video, SceneManager
    # Import the detectors
    from scenedetect.detectors import ContentDetector, HashDetector, AdaptiveDetector, ThresholdDetector, HistogramDetector 
    SCENEDETECT_AVAILABLE = True
except ImportError:
    SCENEDETECT_AVAILABLE = False
    print("Warning: PySceneDetect not found. Scene cut detection will be disabled.")
    print("Install with: pip install scenedetect[opencv]")

try:
    from modules.commons import build_model, load_checkpoint, recursive_munch
    from modules.campplus.DTDNN import CAMPPlus
    from modules.bigvgan import bigvgan
    from modules.rmvpe import RMVPE
    from modules.audio import mel_spectrogram
except ImportError as e:
    print(f"Error importing from 'modules': {e}. Ensure 'modules' directory is accessible.")
    raise

from transformers import AutoFeatureExtractor, WhisperModel

# --- HF_HOME Setup ---
APP_ROOT = Path(__file__).resolve().parent
HF_MODELS_DIR = APP_ROOT / "checkpoints"
HF_MODELS_DIR.mkdir(parents=True, exist_ok=True)
os.environ['HF_HOME'] = str(HF_MODELS_DIR)
print(f"Hugging Face home directory set to: {os.environ['HF_HOME']}")

# --- Speech2RIR Setup ---
S2R_PROJECT_ROOT = HF_MODELS_DIR / "speech2rir"
# Ensure the Speech2RIR project directory exists
S2R_PROJECT_ROOT.mkdir(parents=True, exist_ok=True)
S2R_MODEL_CACHE = None 
if S2R_PROJECT_ROOT.is_dir():
    sys.path.insert(0, str(S2R_PROJECT_ROOT))
    print(f"Added Speech2RIR project root to sys.path: {S2R_PROJECT_ROOT}")
    try:
        from models.autoencoder.AudioDec import Generator as S2R_Generator
        print("Successfully imported Speech2RIR_Generator.")
    except ImportError as e:
        S2R_Generator = None
        print(f"Failed to import s2r_Generator from speech2rir: {e}. Speech2RIR features will be unavailable.")
else:
    S2R_Generator = None
    print(f"Warning: Speech2RIR project directory not found at {S2R_PROJECT_ROOT}. Speech2RIR features will be unavailable.")

# --- Imports for align-voice-rt60.py ---
try:
    from blind_rt60 import BlindRT60
except ImportError:
    print("Warning: 'blind_rt60' not found. Please install it (`pip install blind-rt60`) or place 'blind_rt60.py' in the path.")
    class BlindRT60: 
        def __init__(self, *args, **kwargs): pass
        def estimate(self, *args, **kwargs): return 0.5

import pyroomacoustics as pra
from pedalboard import Convolution, Pedalboard, Gain

# --- PySide6 Imports ---
from PySide6.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
    QPushButton, QLabel, QLineEdit, QFileDialog, QSlider, QCheckBox, QGridLayout, QTabWidget, QSizePolicy,
    QProgressBar, QTextEdit, QGroupBox, QScrollArea, QComboBox, QListWidget, 
    QListWidgetItem, QMessageBox, QSplitter
)
from PySide6.QtCore import Qt, QThread, Signal, Slot, QUrl, QTimer, QRectF, QPointF, QMetaObject
from PySide6.QtGui import QIcon, QPainter, QPen, QColor, QBrush, QPolygonF, QDoubleValidator
from PySide6.QtMultimedia import QMediaPlayer, QAudioOutput
from PySide6.QtWidgets import QWidget

try:
    from pyannote.audio import Pipeline as PyannotePipeline
    PYANNOTE_AVAILABLE = True
except ImportError:
    PYANNOTE_AVAILABLE = False
    print("Warning: pyannote.audio not found. Speaker diarization will be disabled.")
    print("Install with: pip install pyannote.audio")

# --- ClearVoice Imports ---
import modules.clearvoice.networks
try:
    from modules.clearvoice.clearvoice import ClearVoice
    CLEARVOICE_AVAILABLE = True
except ImportError as e:
    CLEARVOICE_AVAILABLE = False
    print(f"Warning: ClearVoice library not found or import error: {e}. Target voice cleaning will be disabled.")

# --- Global Variables & Configuration ---
TMP_DIR = Path("./tmp_processing")
VOICE_MODELS_CACHE = None
CLEARVOICE_MODEL_CACHE = None
HUGGING_FACE_TOKEN_PLACEHOLDER = "YOUR_HUGGING_FACE_ACCESS_TOKEN_HERE"

def load_custom_model_from_hf(repo_id, model_filename="pytorch_model.bin", config_filename="config.yml"):

    os.makedirs("./checkpoints", exist_ok=True)
    model_path = hf_hub_download(repo_id=repo_id, filename=model_filename, cache_dir="./checkpoints")
    if config_filename is None:
        return model_path
    config_path = hf_hub_download(repo_id=repo_id, filename=config_filename, cache_dir="./checkpoints")
    return model_path, config_path

# --- Demucs & Mixing Helper Functions ---
def _save_demucs_audio_file(tensor: torch.Tensor, path: str, sample_rate: int, log_callback):
    """Save audio file from a tensor, typically Demucs output."""
    try:
        audio_numpy = tensor.cpu().numpy()
        # Demucs output can be [channels, samples] or [batch, channels, samples]
        if audio_numpy.ndim == 3: # [batch, channels, samples]
            audio_numpy = audio_numpy.squeeze(0) # Remove batch dim
        
        # Soundfile expects [frames, channels] or [frames,] for mono
        # If [channels, samples], transpose it.
        if audio_numpy.ndim == 2 and audio_numpy.shape[0] < audio_numpy.shape[1]: # Likely [channels, samples]
            audio_numpy = audio_numpy.T # Transpose to [frames, channels]
        elif audio_numpy.ndim == 1: # Mono [frames,]
            pass # Already in correct shape for soundfile
        else: # Should be [frames, channels]
            pass

        sf.write(path, audio_numpy, sample_rate)
        log_callback(f"Saved Demucs component to {path}")
        return path
    except Exception as e:
        log_callback(f"Error saving Demucs audio file {path}: {e}", is_error=True)
        raise

def _mix_audio_torchaudio(audio1_path: str, audio2_path: str, output_path: str, volume_ratio_audio1: float, log_callback):
    """Mix two audio files using torchaudio."""
    try:
        log_callback(f"Mixing {Path(audio1_path).name} ({(volume_ratio_audio1*100):.1f}%) and {Path(audio2_path).name} ({((1.0-volume_ratio_audio1)*100):.1f}%)")
        audio1, sr1 = torchaudio.load(audio1_path)
        audio2, sr2 = torchaudio.load(audio2_path)

        if sr1 != sr2:
            log_callback(f"Sample rates differ. Resampling {Path(audio2_path).name} from {sr2}Hz to {sr1}Hz.")
            resampler = torchaudio.transforms.Resample(sr2, sr1)
            audio2 = resampler(audio2)

        if audio1.shape[0] == 1 and audio2.shape[0] == 2: # audio1 is mono, audio2 is stereo
            audio1 = audio1.repeat(2, 1)
        elif audio2.shape[0] == 1 and audio1.shape[0] == 2: # audio2 is mono, audio1 is stereo
            audio2 = audio2.repeat(2, 1)

        min_len = min(audio1.shape[1], audio2.shape[1])
        if audio1.shape[1] != audio2.shape[1]:
            log_callback(f"Audio lengths differ. Trimming to {min_len} samples.")
        audio1 = audio1[:, :min_len]
        audio2 = audio2[:, :min_len]

        mixed_audio = (audio1 * volume_ratio_audio1) + (audio2 * (1.0 - volume_ratio_audio1))
        max_val = torch.max(torch.abs(mixed_audio))
        if max_val > 1.0: mixed_audio = mixed_audio / max_val # Normalize
        torchaudio.save(output_path, mixed_audio, sr1)
        log_callback(f"Mixed audio saved to {output_path}")
        return output_path
    except Exception as e:
        log_callback(f"Error during torchaudio mixing: {e}", is_error=True)
        raise

# --- Robust Whisper Feature Extractor (Standalone Function) ---
def _process_whisper_features_robust(audio_16k_tensor, whisper_model, whisper_feature_extractor, device):
    """
    Process audio through Whisper model to extract features, handling long audio.
    audio_16k_tensor: [1, num_samples] tensor on the correct device.
    """
    d_model = whisper_model.config.d_model if hasattr(whisper_model, 'config') else whisper_model.encoder.conv1.out_channels # Use out_channels for conv1
    audio_16k_tensor = audio_16k_tensor.to(device) # Ensure input tensor is on the correct device

    if audio_16k_tensor.size(-1) <= 16000 * 30:  # Max 30s for single pass
        audio_list_for_extractor = [audio_16k_tensor.squeeze(0).cpu().numpy()]
        inputs = whisper_feature_extractor(
            audio_list_for_extractor,
            return_tensors="pt", return_attention_mask=True, sampling_rate=16000
        )
        input_features_masked = whisper_model._mask_input_features(
            inputs.input_features, attention_mask=inputs.attention_mask
        ).to(device)
        
        with torch.no_grad(): # Ensure no_grad for inference parts
            outputs = whisper_model.encoder(
                input_features_masked.to(whisper_model.encoder.dtype),
                head_mask=None, output_attentions=False, output_hidden_states=False, return_dict=True,
            )
        features = outputs.last_hidden_state.to(torch.float32)
        features = features[:, :audio_16k_tensor.size(-1) // 320 + 1]
    else:
        overlapping_time_sec = 5
        chunk_duration_sec = 30
        features_list = []
        
        chunk_duration_samples = 16000 * chunk_duration_sec
        overlap_samples = 16000 * overlapping_time_sec

        current_original_start_sample = 0
        idx = 0
        while current_original_start_sample < audio_16k_tensor.size(-1):
            idx += 1
            
            conceptual_chunk_start_original = current_original_start_sample
            conceptual_chunk_end_original = min(current_original_start_sample + chunk_duration_samples, audio_16k_tensor.size(-1))

            whisper_feed_start = max(0, conceptual_chunk_start_original - (overlap_samples if idx > 1 else 0))
            whisper_feed_end = conceptual_chunk_end_original
            
            current_whisper_input_chunk = audio_16k_tensor[:, whisper_feed_start:whisper_feed_end]

            if current_whisper_input_chunk.size(-1) < (16000 * 0.1): # Min 100ms
                if current_original_start_sample >= audio_16k_tensor.size(-1): break
                current_original_start_sample = conceptual_chunk_end_original
                continue

            current_whisper_input_list = [current_whisper_input_chunk.squeeze(0).cpu().numpy()]
            inputs = whisper_feature_extractor(
                current_whisper_input_list,
                return_tensors="pt", return_attention_mask=True, sampling_rate=16000
            )
            input_features_masked = whisper_model._mask_input_features(
                inputs.input_features, attention_mask=inputs.attention_mask
            ).to(device)
            with torch.no_grad(): # Ensure no_grad for inference parts
                outputs = whisper_model.encoder(
                    input_features_masked.to(whisper_model.encoder.dtype),
                    head_mask=None, output_attentions=False, output_hidden_states=False, return_dict=True,
                )
            chunk_features_full = outputs.last_hidden_state.to(torch.float32)
            chunk_features_full = chunk_features_full[:, :current_whisper_input_chunk.size(-1) // 320 + 1]

            if idx == 1:
                features_to_append = chunk_features_full
            else:
                overlap_feature_frames_to_discard = overlap_samples // 320
                if chunk_features_full.size(1) > overlap_feature_frames_to_discard:
                    features_to_append = chunk_features_full[:, overlap_feature_frames_to_discard:, :]
                else:
                    features_to_append = torch.empty((1, 0, d_model), device=device, dtype=torch.float32)
            
            if features_to_append.size(1) > 0:
                features_list.append(features_to_append)
            
            current_original_start_sample = conceptual_chunk_end_original

        if not features_list:
            print("Warning: Whisper feature list is empty after chunking. Input audio might be too short or all chunks were skipped.")
            return torch.empty((1,0,d_model), device=device, dtype=torch.float32)

        features = torch.cat(features_list, dim=1)
    return features

# --- Voice Conversion (Adapted from main_app.py) ---
def adjust_f0_semitones_vp2(f0_sequence, n_semitones):

    factor = 2 ** (n_semitones / 12)
    return f0_sequence * factor

def crossfade_vp2(chunk1, chunk2, overlap):

    fade_out = np.cos(np.linspace(0, np.pi / 2, overlap)) ** 2
    fade_in = np.cos(np.linspace(np.pi / 2, 0, overlap)) ** 2
    if len(chunk2) < overlap: # Handle cases where chunk2 is shorter than overlap
        overlap_actual = len(chunk2)
        chunk2[:overlap_actual] = chunk2[:overlap_actual] * fade_in[:overlap_actual] + \
                                  chunk1[-overlap_actual:] * fade_out[:overlap_actual]
    else:
        chunk2[:overlap] = chunk2[:overlap] * fade_in + chunk1[-overlap:] * fade_out
    return chunk2


def setup_models_vp2(f0_condition, progress_callback):

    progress_callback(0, "Setting up models...")
    print("Setting up voice conversion models...")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")
    
    progress_callback(10, "Loading main DiT model...")
    dit_checkpoint_path, dit_config_path = load_custom_model_from_hf(
        "Plachta/Seed-VC",
        "DiT_seed_v2_uvit_whisper_small_wavenet_bigvgan_pruned.pth",
        "config_dit_mel_seed_uvit_whisper_small_wavenet.yml"
    )
    config = yaml.safe_load(open(dit_config_path, 'r'))
    model_params = recursive_munch(config['model_params'])
    model = build_model(model_params, stage='DiT')
    hop_length = config['preprocess_params']['spect_params']['hop_length']
    sr = config['preprocess_params']['sr']

    model, _, _, _ = load_checkpoint(model, None, dit_checkpoint_path,
                                    load_only_params=True, ignore_modules=[], is_distributed=False)
    for key in model:
        model[key].eval()
        model[key].to(device)
    model.cfm.estimator.setup_caches(max_batch_size=1, max_seq_length=8192)

    progress_callback(30, "Loading CAMPPlus model...")
    campplus_ckpt_path = load_custom_model_from_hf("funasr/campplus", "campplus_cn_common.bin", config_filename=None)
    campplus_model = CAMPPlus(feat_dim=80, embedding_size=192)
    campplus_model.load_state_dict(torch.load(campplus_ckpt_path, map_location="cpu"))
    campplus_model.eval()
    campplus_model.to(device)

    progress_callback(50, "Loading BigVGAN model...")
    bigvgan_model_instance = bigvgan.BigVGAN.from_pretrained('nvidia/bigvgan_v2_22khz_80band_256x', use_cuda_kernel=False)
    bigvgan_model_instance.remove_weight_norm()
    bigvgan_model_instance = bigvgan_model_instance.eval().to(device)
    
    progress_callback(70, "Loading Whisper model...")
    whisper_name = model_params.speech_tokenizer.whisper_name if hasattr(model_params.speech_tokenizer,
                                                                        'whisper_name') else "openai/whisper-small"
    whisper_model_instance = WhisperModel.from_pretrained(whisper_name, torch_dtype=torch.float16 if device.type == 'cuda' else torch.float32)
    # Move to device *after* potential deletion of decoder
    if hasattr(whisper_model_instance, 'decoder'): del whisper_model_instance.decoder
    whisper_model_instance = whisper_model_instance.to(device) # Ensure it's on device

    whisper_feature_extractor_instance = AutoFeatureExtractor.from_pretrained(whisper_name)

    mel_fn_args = {
        "n_fft": config['preprocess_params']['spect_params']['n_fft'],
        "win_size": config['preprocess_params']['spect_params']['win_length'],
        "hop_size": config['preprocess_params']['spect_params']['hop_length'],
        "num_mels": config['preprocess_params']['spect_params']['n_mels'],
        "sampling_rate": sr, "fmin": 0, "fmax": None, "center": False
    }
    to_mel_instance = lambda x: mel_spectrogram(x, **mel_fn_args)
    
    model_f0_instance, rmvpe_instance, to_mel_f0_instance, bigvgan_44k_model_instance, sr_f0, hop_length_f0 = None, None, None, None, None, None
    if f0_condition:
        progress_callback(80, "Loading F0 models...")
        dit_f0_ckpt_path, dit_f0_config_path = load_custom_model_from_hf(
            "Plachta/Seed-VC",
            "DiT_seed_v2_uvit_whisper_base_f0_44k_bigvgan_pruned_ft_ema.pth",
            "config_dit_mel_seed_uvit_whisper_base_f0_44k.yml"
        )
        config_f0 = yaml.safe_load(open(dit_f0_config_path, 'r'))
        model_params_f0 = recursive_munch(config_f0['model_params'])
        model_f0_instance = build_model(model_params_f0, stage='DiT')
        hop_length_f0 = config_f0['preprocess_params']['spect_params']['hop_length']
        sr_f0 = config_f0['preprocess_params']['sr']

        model_f0_instance, _, _, _ = load_checkpoint(model_f0_instance, None, dit_f0_ckpt_path,
                                        load_only_params=True, ignore_modules=[], is_distributed=False)
        for key in model_f0_instance: model_f0_instance[key].eval(); model_f0_instance[key].to(device)
        model_f0_instance.cfm.estimator.setup_caches(max_batch_size=1, max_seq_length=8192)

        rmvpe_model_path = load_custom_model_from_hf("lj1995/VoiceConversionWebUI", "rmvpe.pt", config_filename=None)
        rmvpe_instance = RMVPE(rmvpe_model_path, is_half=False, device=device)

        mel_fn_args_f0 = {
            "n_fft": config_f0['preprocess_params']['spect_params']['n_fft'],
            "win_size": config_f0['preprocess_params']['spect_params']['win_length'],
            "hop_size": hop_length_f0, "num_mels": config_f0['preprocess_params']['spect_params']['n_mels'],
            "sampling_rate": sr_f0, "fmin": 0, "fmax": None, "center": False
        }
        to_mel_f0_instance = lambda x: mel_spectrogram(x, **mel_fn_args_f0)
        
        bigvgan_44k_model_instance = bigvgan.BigVGAN.from_pretrained('nvidia/bigvgan_v2_44khz_128band_512x', use_cuda_kernel=False)
        bigvgan_44k_model_instance.remove_weight_norm()
        bigvgan_44k_model_instance = bigvgan_44k_model_instance.eval().to(device)
    
    progress_callback(100, "Models ready.")
    models_pack = {
        'device': device, 'model': model, 'model_f0': model_f0_instance, 'campplus_model': campplus_model,
        'bigvgan_model': bigvgan_model_instance, 'bigvgan_44k_model': bigvgan_44k_model_instance,
        'whisper_model': whisper_model_instance, 'whisper_feature_extractor': whisper_feature_extractor_instance,
        'rmvpe': rmvpe_instance, 'to_mel': to_mel_instance, 'to_mel_f0': to_mel_f0_instance,
        'sr': sr if not f0_condition else sr_f0,
        'hop_length': hop_length if not f0_condition else hop_length_f0,
        'f0_condition_state': f0_condition 
    }
    print("Voice conversion models set up successfully.")
    return models_pack


@torch.no_grad()
@torch.inference_mode()
def voice_conversion_vp2(source_path, target_path, diffusion_steps, length_adjust, 
                         inference_cfg_rate, f0_condition, auto_f0_adjust, pitch_shift, models,
                         progress_callback=None): # Removed chunk_idx, total_chunks as they are now internal

    device = models['device']
    inference_module = models['model'] if not f0_condition else models['model_f0']
    mel_fn = models['to_mel'] if not f0_condition else models['to_mel_f0']
    bigvgan_fn = models['bigvgan_model'] if not f0_condition else models['bigvgan_44k_model']
    # Whisper models are passed directly to _process_whisper_features_robust
    campplus_model = models['campplus_model']
    rmvpe = models['rmvpe']
    
    sr_vc = models['sr'] 
    hop_length = models['hop_length']
    
    max_context_window = sr_vc // hop_length * 30 # For DiT/CFM processing
    overlap_frame_len = 16 
    overlap_wave_len = overlap_frame_len * hop_length
    
    source_audio_orig, sr_orig = librosa.load(source_path, sr=None)
    if sr_orig != sr_vc:
        source_audio_orig = librosa.resample(source_audio_orig, orig_sr=sr_orig, target_sr=sr_vc)
    
    ref_audio_orig, sr_ref_orig = librosa.load(target_path, sr=None)
    if sr_ref_orig != sr_vc:
        ref_audio_orig = librosa.resample(ref_audio_orig, orig_sr=sr_ref_orig, target_sr=sr_vc)

    if len(ref_audio_orig) > sr_vc * 25:
        ref_audio_orig = ref_audio_orig[:sr_vc * 25]

    source_audio = torch.tensor(source_audio_orig).unsqueeze(0).float().to(device)
    ref_audio = torch.tensor(ref_audio_orig).unsqueeze(0).float().to(device)

    # --- Resample for Whisper ---
    converted_waves_16k = torchaudio.functional.resample(source_audio, sr_vc, 16000)
    ori_waves_16k = torchaudio.functional.resample(ref_audio, sr_vc, 16000)
    
    # --- ROBUST Whisper Feature Extraction ---
    models['whisper_model'].to(device) # Ensure on correct device
    S_alt = _process_whisper_features_robust(converted_waves_16k, models['whisper_model'], models['whisper_feature_extractor'], device)
    S_ori = _process_whisper_features_robust(ori_waves_16k, models['whisper_model'], models['whisper_feature_extractor'], device)

    # --- Mel Spectrograms and Style ---
    mel = mel_fn(source_audio.float())
    mel2 = mel_fn(ref_audio.float())

    target_lengths = torch.LongTensor([int(mel.size(2) * length_adjust)]).to(mel.device)
    target2_lengths = torch.LongTensor([mel2.size(2)]).to(mel2.device)

    feat2 = torchaudio.compliance.kaldi.fbank(ori_waves_16k, num_mel_bins=80, dither=0, sample_frequency=16000)
    feat2 = feat2 - feat2.mean(dim=0, keepdim=True)
    style2 = campplus_model(feat2.unsqueeze(0))

    # --- F0 Processing ---
    F0_ori, F0_alt, shifted_f0_alt = None, None, None
    if f0_condition and rmvpe:
        F0_ori_np = rmvpe.infer_from_audio(ori_waves_16k[0].cpu().numpy(), thred=0.03) # rmvpe expects numpy
        F0_alt_np = rmvpe.infer_from_audio(converted_waves_16k[0].cpu().numpy(), thred=0.03)
        F0_ori = torch.from_numpy(F0_ori_np).to(device)[None]
        F0_alt = torch.from_numpy(F0_alt_np).to(device)[None]
        
        log_f0_alt = torch.log(F0_alt + 1e-5)
        if auto_f0_adjust:
            voiced_F0_ori = F0_ori[F0_ori > 1]; voiced_F0_alt = F0_alt[F0_alt > 1]
            if len(voiced_F0_ori) > 0 and len(voiced_F0_alt) > 0:
                median_log_f0_ori = torch.median(torch.log(voiced_F0_ori + 1e-5))
                median_log_f0_alt = torch.median(torch.log(voiced_F0_alt + 1e-5))
                log_f0_alt[F0_alt > 1] = log_f0_alt[F0_alt > 1] - median_log_f0_alt + median_log_f0_ori
        shifted_f0_alt = torch.exp(log_f0_alt)
        if pitch_shift != 0:
            shifted_f0_alt[F0_alt > 1] = adjust_f0_semitones_vp2(shifted_f0_alt[F0_alt > 1], pitch_shift)
    
    # --- Length Regulation ---
    cond, _, _, _, _ = inference_module.length_regulator(S_alt, ylens=target_lengths, n_quantizers=3, f0=shifted_f0_alt)
    prompt_condition, _, _, _, _ = inference_module.length_regulator(S_ori, ylens=target2_lengths, n_quantizers=3, f0=F0_ori)
    
    # --- DiT/CFM Processing Loop (Handles its own chunking for acoustic model features) ---
    max_source_window_dit = max_context_window - mel2.size(2) 
    if max_source_window_dit <= 0: 
        max_source_window_dit = cond.size(1) 
        print(f"Warning: Reference audio (mel length {mel2.size(2)}) is long. Processing source in one go for DiT/CFM.")

    processed_frames_dit = 0 
    generated_wave_chunks_for_segment = [] 
    previous_vocoder_output_overlap_chunk = None 
    
    total_dit_frames = cond.size(1)
    num_dit_chunks_for_progress = (total_dit_frames + max_source_window_dit -1) // max_source_window_dit if max_source_window_dit > 0 else 1
    
    current_dit_chunk_idx = 0
    while processed_frames_dit < total_dit_frames:
        current_dit_chunk_idx += 1
        chunk_cond_dit = cond[:, processed_frames_dit : processed_frames_dit + max_source_window_dit] 
        is_last_dit_chunk = (processed_frames_dit + max_source_window_dit) >= total_dit_frames
        
        cat_condition = torch.cat([prompt_condition, chunk_cond_dit], dim=1)
        
        with torch.autocast(device_type=device.type, dtype=torch.float16 if device.type == 'cuda' else torch.float32):
            vc_target_mel = inference_module.cfm.inference(cat_condition,
                                                    torch.LongTensor([cat_condition.size(1)]).to(mel2.device),
                                                    mel2, style2, None, diffusion_steps,
                                                    inference_cfg_rate=inference_cfg_rate)
            vc_target_mel = vc_target_mel[:, :, mel2.size(-1):] 
            
        vc_wave_chunk = bigvgan_fn(vc_target_mel.float())[0]
        
        if progress_callback:
            # Progress is now relative to the DiT loop within this segment
            progress_within_segment = (current_dit_chunk_idx / num_dit_chunks_for_progress) * 100
            progress_callback(int(progress_within_segment), 
                              f"VC DiT {current_dit_chunk_idx}/{num_dit_chunks_for_progress}: {min(processed_frames_dit + vc_target_mel.size(2), total_dit_frames)}/{total_dit_frames} frames")

        if not generated_wave_chunks_for_segment: 
            if is_last_dit_chunk or overlap_wave_len == 0:
                output_wave_for_stitch = vc_wave_chunk[0].cpu().numpy()
            else:
                output_wave_for_stitch = vc_wave_chunk[0, :-overlap_wave_len].cpu().numpy()
            generated_wave_chunks_for_segment.append(output_wave_for_stitch)
            previous_vocoder_output_overlap_chunk = vc_wave_chunk[0, -overlap_wave_len:] if overlap_wave_len > 0 else None
        elif is_last_dit_chunk or overlap_wave_len == 0:
            if previous_vocoder_output_overlap_chunk is not None and overlap_wave_len > 0:
                 output_wave_for_stitch = crossfade_vp2(previous_vocoder_output_overlap_chunk.cpu().numpy(), vc_wave_chunk[0].cpu().numpy(), overlap_wave_len)
            else:
                output_wave_for_stitch = vc_wave_chunk[0].cpu().numpy()
            generated_wave_chunks_for_segment.append(output_wave_for_stitch)
        else: 
            output_wave_for_stitch = crossfade_vp2(previous_vocoder_output_overlap_chunk.cpu().numpy(), vc_wave_chunk[0, :-overlap_wave_len].cpu().numpy(), overlap_wave_len)
            generated_wave_chunks_for_segment.append(output_wave_for_stitch)
            previous_vocoder_output_overlap_chunk = vc_wave_chunk[0, -overlap_wave_len:]

        if is_last_dit_chunk: break
        processed_frames_dit += vc_target_mel.size(2) - (overlap_frame_len if overlap_wave_len > 0 else 0)
        if processed_frames_dit >= total_dit_frames: break 

    final_output_for_segment = np.concatenate(generated_wave_chunks_for_segment) if generated_wave_chunks_for_segment else np.array([], dtype=np.float32)
    return sr_vc, final_output_for_segment

# --- RIR Application (Original and New for S2R) ---
def apply_reverb_to_audio_segment( # For simulated RIR
    input_audio_path_str: str, 
    output_audio_path_str: str, 
    target_sr: int, 
    room_dim_list: list, 
    rt60_target: float, 
    progress_callback=None):

    if progress_callback: progress_callback(0, f"Applying reverb: Loading segment {Path(input_audio_path_str).name}")
    studio_audio, studio_fs_orig = librosa.load(input_audio_path_str, sr=None, mono=True)
    if studio_fs_orig != target_sr:
        studio_audio = librosa.resample(studio_audio, orig_sr=studio_fs_orig, target_sr=target_sr)
    
    if studio_audio.ndim == 1:
        studio_audio = studio_audio[:, np.newaxis] 
        studio_audio_pb = studio_audio 
    elif studio_audio.shape[0] < studio_audio.shape[1]: 
        studio_audio_pb = studio_audio.T 
    else: 
        studio_audio_pb = studio_audio

    fs = target_sr 

    V = np.prod(room_dim_list)
    Sx, Sy, Sz = room_dim_list
    S_total = 2 * (Sx*Sy + Sx*Sz + Sy*Sz)
    
    if rt60_target <= 0: rt60_target = 0.01 
    alpha = 0.161 * V / (rt60_target * S_total) if S_total > 0 else 0.5
    alpha = np.clip(alpha, 0.01, 0.99) 

    if progress_callback: progress_callback(30, f"Simulating room (dim: {room_dim_list}, RT60: {rt60_target:.2f}s, alpha: {alpha:.2f})")
    try:
        room = pra.ShoeBox(room_dim_list, fs=fs, materials=pra.Material(alpha), max_order=15)
    except Exception as e:
        print(f"Error creating ShoeBox room (dims: {room_dim_list}, alpha: {alpha:.2f}): {e}")
        sf.write(output_audio_path_str, studio_audio.squeeze(), fs)
        if progress_callback: progress_callback(100, "Room sim failed.")
        return
    
    source_pos = [room_dim_list[0]/2, room_dim_list[1]/2, room_dim_list[2]/2]
    mic_pos = [source_pos[0] + 1, source_pos[1], source_pos[2]] 
    mic_pos = [min(max(mic_pos[i], 0.1), room_dim_list[i]-0.1) for i in range(3)] 

    room.add_source(source_pos)
    room.add_microphone_array(pra.MicrophoneArray(np.array([mic_pos]).T, fs=fs)) 
    
    if progress_callback: progress_callback(60, "Computing RIR...")
    try:
        room.compute_rir()
        if not room.rir or not room.rir[0] or room.rir[0][0] is None or len(room.rir[0][0]) == 0:
            raise ValueError("RIR computation failed or resulted in empty RIR.")
        rir = room.rir[0][0].astype(np.float32)
    except Exception as e:
        print(f"RIR computation error: {e}. Skipping reverb.")
        sf.write(output_audio_path_str, studio_audio.squeeze(), fs)
        if progress_callback: progress_callback(100, "RIR computation failed.")
        return
        
    if progress_callback: progress_callback(80, "Applying RIR via Pedalboard...")
    convolution = Convolution(rir, sample_rate=fs) 
    board = Pedalboard([convolution])
    output_audio_aligned = board(studio_audio_pb, sample_rate=fs) 
    sf.write(output_audio_path_str, output_audio_aligned, fs)
    if progress_callback: progress_callback(100, "Reverb application complete.")
    print(f"Reverberated audio saved to {output_audio_path_str}")

# --- RT60 Estimation Helper ---
def estimate_rt60_from_audio(audio_path_str, progress_callback=None):

    if progress_callback: progress_callback(10, "Loading audio for RT60 estimation...")
    audio, fs = librosa.load(audio_path_str, sr=None, mono=True)
    if progress_callback: progress_callback(50, "Estimating RT60...")
    estimator = BlindRT60(fs=fs)
    rt60 = estimator.estimate(audio, fs) 
    if progress_callback: progress_callback(100, f"Estimated RT60: {rt60:.3f}s")
    return rt60

# --- Speaker Diarization Helper ---
def extract_speaker_times(rttm_file, target_speaker):

    times = []
    try:
        with open(rttm_file, "r") as f:
            for line in f:
                if "SPEAKER" in line: 
                    parts = line.split()
                    if len(parts) >= 8:
                        try:
                            speaker_id_from_file = parts[7] 
                            if speaker_id_from_file == target_speaker:
                                start = float(parts[3])
                                duration = float(parts[4])
                                end = start + duration
                                times.append((start, end))
                        except (ValueError, IndexError):
                            print(f"Could not parse RTTM line: {line.strip()}")
                            continue
    except FileNotFoundError:
        print(f"RTTM file not found: {rttm_file}")
    return times

# --- WaveformWidget (QPainter version, from previous diff) ---
class WaveformWidget(QWidget): 
    markerClicked = Signal(float, float) # Emit start_sec, end_sec of a clicked visual marker

    clickedPosition = Signal(float) 
    selectionFinalized = Signal(float, float)

    def __init__(self, parent_window, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.parent_window = parent_window
        self.setMinimumHeight(70) # Reduced height to 70% of original 100
        self.setMouseTracking(True) 

        self.y_data_normalized = None 
        self.duration_sec = 0
        
        self.playback_cursor_pos_sec = 0.0
        self.show_playback_cursor = False

        self.selection_start_sec = None
        self.selection_end_sec = None
        self.is_selecting = False
        self.is_dragging_playback_cursor = False
        self.playback_cursor_handle_width = 10 # pixels, for easier grabbing

        self.visual_markers = [] 

        self.waveform_pen = QPen(QColor(0, 0, 255, 180), 1.5)
        self.cursor_pen = QPen(Qt.red, 2)
        self.selection_brush = QBrush(QColor(0, 0, 255, 50))

    def load_audio(self, y_for_plot_normalized, original_sr, original_duration_sec):
        self.y_data_normalized = y_for_plot_normalized
        self.duration_sec = original_duration_sec 
        self.playback_cursor_pos_sec = 0.0
        self.show_playback_cursor = False 
        self.selection_start_sec = None
        self.selection_end_sec = None
        self.is_selecting = False
        self.is_dragging_playback_cursor = False
        self.visual_markers.clear()
        self.update() 

    def set_playback_cursor(self, time_sec):
        if self.duration_sec > 0:
            self.playback_cursor_pos_sec = max(0, min(time_sec, self.duration_sec))
            self.show_playback_cursor = True
            self.update()

    def hide_playback_cursor(self):
        self.show_playback_cursor = False
        self.update()

    def add_visual_marker(self, start_sec, end_sec, color=QColor(255, 165, 0, 100)): 
        if self.duration_sec == 0 or end_sec <= start_sec:
            return
        self.visual_markers.append((max(0, start_sec), min(self.duration_sec, end_sec), QBrush(color)))
        self.visual_markers.sort() 
        self.update()

    def clear_all_visual_markers(self):
        self.visual_markers.clear()
        self.selection_start_sec = None
        self.selection_end_sec = None
        self.is_selecting = False
        self.is_dragging_playback_cursor = False
        self.update()

    def get_current_selection(self):
        if self.selection_start_sec is not None and self.selection_end_sec is not None:
            start = min(self.selection_start_sec, self.selection_end_sec)
            end = max(self.selection_start_sec, self.selection_end_sec)
            if end > start + 0.01: 
                return start, end
        return None, None

    def _time_to_x(self, time_sec):
        if self.duration_sec == 0: return 0
        return (time_sec / self.duration_sec) * self.width()

    def _x_to_time(self, x_pos):
        if self.width() == 0: return 0
        return (x_pos / self.width()) * self.duration_sec

    def paintEvent(self, event):
        super().paintEvent(event)
        painter = QPainter(self)
        painter.setRenderHint(QPainter.Antialiasing)
        
        width = self.width()
        height = self.height()
        
        # painter.fillRect(self.rect(), Qt.black) # Removed to make background transparent
        if self.y_data_normalized is not None and len(self.y_data_normalized) > 0 and self.duration_sec > 0:
            for start_s, end_s, brush in self.visual_markers:
                x_start = self._time_to_x(start_s)
                x_end = self._time_to_x(end_s)
                painter.fillRect(QRectF(x_start, 0, x_end - x_start, height), brush)

            current_sel = self.get_current_selection()
            if current_sel[0] is not None:
                sel_x_start = self._time_to_x(current_sel[0])
                sel_x_end = self._time_to_x(current_sel[1])
                painter.fillRect(QRectF(sel_x_start, 0, sel_x_end - sel_x_start, height), self.selection_brush)

            painter.setPen(self.waveform_pen)
            num_points = len(self.y_data_normalized)
            points_per_pixel = num_points / width if width > 0 else num_points
            
            if points_per_pixel < 1 and num_points > 1: 
                path = QPolygonF()
                for i in range(num_points):
                    t = (i / (num_points -1)) * self.duration_sec if num_points > 1 else 0 
                    x = self._time_to_x(t)
                    y_val = self.y_data_normalized[i]
                    y = height / 2 - (y_val * height / 2.2) 
                    path.append(QPointF(x,y))
                painter.drawPolyline(path)
            elif num_points > 0 : 
                for x_pixel in range(width):
                    start_idx = int(x_pixel * points_per_pixel)
                    end_idx = int((x_pixel + 1) * points_per_pixel)
                    if start_idx >= end_idx: end_idx = start_idx + 1 
                    if end_idx > num_points: end_idx = num_points
                    if start_idx >= num_points: break

                    segment_data = self.y_data_normalized[start_idx:end_idx]
                    if len(segment_data) == 0: continue
                    
                    min_val = np.min(segment_data)
                    max_val = np.max(segment_data)
                    
                    y_min = height / 2 - (min_val * height / 2.2)
                    y_max = height / 2 - (max_val * height / 2.2)
                    painter.drawLine(QPointF(x_pixel, y_min), QPointF(x_pixel, y_max))
            
            if self.show_playback_cursor:
                painter.setPen(self.cursor_pen)
                cursor_x = self._time_to_x(self.playback_cursor_pos_sec)
                painter.drawLine(QPointF(cursor_x, 0), QPointF(cursor_x, height))
        else:
            painter.setPen(Qt.white) # Or another contrasting color for your theme
            painter.drawText(self.rect(), Qt.AlignCenter, "No audio loaded")


    def mousePressEvent(self, event):
        if event.button() == Qt.LeftButton and self.y_data_normalized is not None and self.duration_sec > 0:
            time_sec = self._x_to_time(event.pos().x())
            time_sec = max(0, min(time_sec, self.duration_sec)) # Clamp time

            # Check if clicking on the playback cursor handle
            if self.show_playback_cursor:
                cursor_center_x = self._time_to_x(self.playback_cursor_pos_sec)
                handle_half_width = self.playback_cursor_handle_width / 2.0
                if cursor_center_x - handle_half_width <= event.pos().x() <= cursor_center_x + handle_half_width:
                    self.is_dragging_playback_cursor = True
                    self.is_selecting = False  # Prevent selection when dragging cursor
                    self.set_playback_cursor(time_sec)  # Update cursor position immediately
                    self.clickedPosition.emit(time_sec)  # Emit signal to seek audio
                    self.update()
                    return  # Consume event

            # Check if clicking on a visual marker
            # This check should happen BEFORE starting a new selection or general seek,
            # if marker click is meant to be a distinct action.
            clicked_on_marker = False
            for i, (m_start_s, m_end_s, _) in enumerate(self.visual_markers):
                marker_x_start = self._time_to_x(m_start_s)
                marker_x_end = self._time_to_x(m_end_s)
                if marker_x_start <= event.pos().x() <= marker_x_end:
                    self.markerClicked.emit(m_start_s, m_end_s)
                    clicked_on_marker = True
                    # If a marker is clicked, we might not want to immediately start a new selection
                    # or seek, depending on desired UX. For now, we let it proceed.
                    break 

            # If not dragging cursor, proceed with normal click (seek and/or start selection)
            self.clickedPosition.emit(time_sec)
            self.is_selecting = True # Start selection process
            self.selection_start_sec = time_sec # For selection
            self.selection_end_sec = time_sec # For selection
            self.update()
        # super().mousePressEvent(event) # Don't call super if we handle it fully

    def mouseMoveEvent(self, event):
        if self.is_dragging_playback_cursor and self.y_data_normalized is not None and self.duration_sec > 0:
            time_sec = self._x_to_time(event.pos().x())
            time_sec = max(0, min(time_sec, self.duration_sec)) # Clamp time
            self.set_playback_cursor(time_sec)  # Updates visual position and calls self.update()
            self.clickedPosition.emit(time_sec)  # Emit for continuous seeking
        elif self.is_selecting and self.y_data_normalized is not None and self.duration_sec > 0:
            self.selection_end_sec = self._x_to_time(event.pos().x()) # For selection
            self.update() # For selection
        # super().mouseMoveEvent(event)

    def mouseReleaseEvent(self, event):
        if event.button() == Qt.LeftButton and self.y_data_normalized is not None and self.duration_sec > 0:
            if self.is_dragging_playback_cursor:
                self.is_dragging_playback_cursor = False
                # Final position already set by set_playback_cursor in mouseMove or mousePress
                # Emit signal one last time to ensure audio seeks to the final dragged position
                time_sec = self._x_to_time(event.pos().x())
                time_sec = max(0, min(time_sec, self.duration_sec)) # Clamp time
                self.clickedPosition.emit(time_sec)
                self.update()
                return # Consume event

            if self.is_selecting: # Handle selection finalization
                self.is_selecting = False
                self.selection_end_sec = self._x_to_time(event.pos().x())
                start, end = self.get_current_selection()
                if start is not None and end is not None:
                    self.selectionFinalized.emit(start, end)
                self.update() # For selection
        # super().mouseReleaseEvent(event)

    def clear_selection(self):
        self.selection_start_sec = None
        self.selection_end_sec = None
        self.is_selecting = False
        self.update()

# --- Worker Thread ---
class WorkerThread(QThread):
    progress = Signal(int, str)
    finished = Signal(object, str)
    log_message = Signal(str, bool) # message, is_error

    def __init__(self, func, *args, **kwargs):
        super().__init__()
        self.func = func
        self.args = args
        self.kwargs = kwargs

    def _progress_callback_adapter(self, value, desc):
        self.progress.emit(int(value), desc)

    def _log_callback_adapter(self, message, is_error=False):
        self.log_message.emit(message, is_error)

    def run(self):
        try:
            self.kwargs['progress_callback'] = self._progress_callback_adapter
            # Add log_callback if the target function expects it
            if 'log_callback' in self.func.__code__.co_varnames:
                 self.kwargs['log_callback'] = self._log_callback_adapter
            result = self.func(*self.args, **self.kwargs)
            self.finished.emit(result, "Task completed successfully.")
        except Exception as e:
            import traceback
            tb_str = traceback.format_exc()
            print(f"Error in worker thread: {e}\n{tb_str}")
            self.log_message.emit(f"Worker Error: {str(e)}\n{tb_str}", True)
            self.finished.emit(None, f"Error: {str(e)}")


# --- Main Window ---
class MainWindow(QMainWindow):
    def __init__(self):
        # ... (constructor and _init_ui, _connect_signals as before, with WaveformWidget(self)) ...
        super().__init__()
        self.setWindowTitle("VideoVoiceSwap") # Updated Title
        self.setGeometry(50, 50, 1600, 950) # Slightly larger

        self.central_widget = QWidget()
        self.setCentralWidget(self.central_widget)
        self.main_layout = QVBoxLayout(self.central_widget)

        self.input_mp4_path = None
        self.target_voice_path = None
        self.output_dir_path = None 
        
        self.source_audio_full_path = None 
        self.source_audio_sr = None
        self.source_audio_duration_sec = 0
        self.source_video_fps = None # To store video FPS
        self.source_time_markers = [] 


        self.swapped_audio_full_path = None
        self.swapped_audio_markers = [] 

        self.estimated_rir_path = None # Path to the RIR estimated by Speech2RIR

        self.diarization_results = None 
        self.auto_rir_segment_duration_sec = 20 # Default duration for auto RIR segments
        self.scene_cuts_sec = [] # To store detected scene cut timestamps in seconds

        self.player = QMediaPlayer()
        self.audio_output = QAudioOutput()
        self.player.setAudioOutput(self.audio_output)
        self.player.positionChanged.connect(self._update_waveform_cursor_on_play)
        self.player.playbackStateChanged.connect(self._on_playback_state_changed)
        
        self._init_ui()
        self._connect_signals()

        TMP_DIR.mkdir(parents=True, exist_ok=True)
        
        self.hf_token = HUGGING_FACE_TOKEN_PLACEHOLDER 
        self.hf_token_edit.setText(self.hf_token) 
        self.clearvoice_se_model_name = 'MossFormer2_SE_48K' # Default SE model
        self.swapped_audio_waveform_widget.hide()
        self.current_worker = None # Explicitly initialize after UI setup

    def _init_ui(self):
        # --- Top Row Container for Height Adjustment ---
        top_row_container = QWidget()
        # Estimate current height or set a fixed one. Let's try 220px for a +30px feel.
        # Original might be around 150-190px depending on font and spacing.
        # Minimum height of input files group + config group + margins.
        # An input file row is (LineEdit + Button) ~30-35px. 3 rows = ~100px. Groupbox title ~20px. Total ~120px.
        # Config group has multiple elements, let's say ~120-150px. Tab bar itself.
        top_row_container.setMinimumHeight(210) # Increased height for the top row
        top_row_layout = QHBoxLayout(top_row_container)
        # ... (UI setup as before, ensuring WaveformWidget(self) is used) ...
        # --- Top Row: File Inputs & Room/HF Token ---
        
        # --- Left Pane for Top Row (Input Files & Global Config) ---
        left_top_pane_widget = QWidget()
        left_top_pane_layout = QHBoxLayout(left_top_pane_widget) # Changed to QHBoxLayout
        left_top_pane_layout.setContentsMargins(0,0,0,0) # Optional: remove inner margins
        # left_top_pane_widget.setStyleSheet("border: 1px solid red;") # For debugging layout

        input_files_group = QGroupBox("Input Files")
        input_files_layout = QVBoxLayout()
        mp4_file_layout = QHBoxLayout()
        self.mp4_label = QLineEdit("No MP4 file selected")
        self.mp4_label.setReadOnly(True)
        self.mp4_button = QPushButton("Select Input MP4")
        mp4_file_layout.addWidget(self.mp4_label)
        mp4_file_layout.addWidget(self.mp4_button)
        input_files_layout.addLayout(mp4_file_layout)
        
        wav_file_layout = QHBoxLayout()
        self.wav_label = QLineEdit("No Target voice audio file selected")
        self.wav_label.setReadOnly(True)
        self.wav_button = QPushButton("Select Target Audio")
        wav_file_layout.addWidget(self.wav_label)
        wav_file_layout.addWidget(self.wav_button)

        self.clean_target_voice_button = QPushButton("Clean Target")
        self.clean_target_voice_button.setToolTip("Clean Target Voice (Speech Enhancement).\nApply Speech Enhancement to the selected target voice file.\nRequires ClearVoice models and configs.")
        self.clean_target_voice_button.setEnabled(False) # Enabled when target wav is selected
        wav_file_layout.addWidget(self.clean_target_voice_button)
        input_files_layout.addLayout(wav_file_layout)

        output_dir_layout = QHBoxLayout()
        self.output_dir_label = QLineEdit("No output directory selected")
        self.output_dir_label.setReadOnly(True)
        self.output_dir_button = QPushButton("Select Output Directory")
        output_dir_layout.addWidget(self.output_dir_label)
        output_dir_layout.addWidget(self.output_dir_button)
        input_files_layout.addLayout(output_dir_layout)
        input_files_group.setLayout(input_files_layout)
        left_top_pane_layout.addWidget(input_files_group) # Add to left pane
        config_group = QGroupBox("Global Configuration")
        config_layout = QVBoxLayout()
        top_columns_layout = QHBoxLayout()

        left_column_v_layout = QVBoxLayout()
        left_column_v_layout.addWidget(QLabel("BlindRT60 - Environment:"))
        self.environment_combo = QComboBox()
        self.environment_combo.addItems(["Living Room", "Outside"]) # Consistent casing
        left_column_v_layout.addWidget(self.environment_combo)
        self.calculate_rt60_button = QPushButton("Calc RT60 (Source)")
        left_column_v_layout.addWidget(self.calculate_rt60_button)
        left_column_v_layout.addStretch(1) 
        top_columns_layout.addLayout(left_column_v_layout)

        right_column_v_layout = QVBoxLayout()
        right_column_v_layout.addWidget(QLabel("Room Dimensions (X,Y,Z m):"))
        self.room_dim_edit = QLineEdit("8,7,3.2") 
        right_column_v_layout.addWidget(self.room_dim_edit)
        right_column_v_layout.addWidget(QLabel("BlindRT60 - Estimated Source RT60 (s):"))
        self.rt60_display_edit = QLineEdit()
        self.rt60_display_edit.setReadOnly(True); self.rt60_display_edit.setPlaceholderText("N/A")
        right_column_v_layout.addWidget(self.rt60_display_edit)
        right_column_v_layout.addStretch(1) 
        top_columns_layout.addLayout(right_column_v_layout)

        config_layout.addLayout(top_columns_layout)

        hf_scene_layout = QHBoxLayout()
        hf_layout_v = QVBoxLayout()
        hf_layout_v.addWidget(QLabel("HF Token (For Speaker Diarization):"))
        self.hf_token_edit = QLineEdit()
        hf_layout_v.addWidget(self.hf_token_edit)
        hf_scene_layout.addLayout(hf_layout_v)

        config_layout.addLayout(hf_scene_layout) # Keep adding to config_layout
        config_group.setLayout(config_layout) # Set layout for the group
        config_group.setCheckable(True) # Make it collapsible
        config_group.setChecked(False)  # Start collapsed

        top_row_layout.addWidget(left_top_pane_widget, 0) # Stretch factor 0 for left_top_pane
        
        # --- VC Parameters Scroll Area (moved to top row, right side) ---
        vc_params_scroll_area = QScrollArea()
        vc_params_scroll_area.setMinimumWidth(300) # Adjusted width
        vc_params_scroll_area.setWidgetResizable(True)

        # --- Create Tab Widget and Individual Tabs ---
        self.settings_tab_widget = QTabWidget()

        # --- Tab 0: Global Configuration (Moved Here) ---
        global_config_tab = QWidget()
        global_config_layout = QVBoxLayout(global_config_tab)
        # config_group is already created and populated from above
        global_config_layout.addWidget(config_group)
        global_config_layout.addStretch(1)
        self.settings_tab_widget.addTab(global_config_tab, "Global Config")

        # --- Tab 1: Source Pre-processing ---
        source_preproc_tab = QWidget()
        source_preproc_layout = QVBoxLayout(source_preproc_tab)
        echo_reduction_group = QGroupBox("Source Pre-processing")
        echo_reduction_layout = QVBoxLayout()
        self.reduce_source_echo_checkbox = QCheckBox("Reduce Source Echo (via ClearVoice SE before Voice Conversion)")
        echo_reduction_layout.addWidget(self.reduce_source_echo_checkbox)
        self.source_echo_rt60_threshold_edit = QLineEdit("0.7")
        echo_reduction_layout.addWidget(QLabel("RT60 Threshold for Echo Reduction (s):"))
        echo_reduction_layout.addWidget(self.source_echo_rt60_threshold_edit)
        echo_reduction_group.setLayout(echo_reduction_layout)
        source_preproc_layout.addWidget(echo_reduction_group)
        source_preproc_layout.addStretch(1)
        self.settings_tab_widget.addTab(source_preproc_tab, "Source Pre-process")

        # --- Tab 2: Voice Conversion Model Parameters ---
        vc_params_tab = QWidget()
        vc_params_tab_layout = QVBoxLayout(vc_params_tab)
        vc_params_group = QGroupBox("Voice Conversion Model Parameters")
        vc_params_sliders_layout = QGridLayout() # Use QGridLayout for better alignment
        # Diffusion Steps
        vc_params_sliders_layout.addWidget(QLabel("Diffusion Steps:"), 0, 0)
        self.diffusion_steps_slider = QSlider(Qt.Horizontal)
        self.diffusion_steps_slider.setRange(10, 100); self.diffusion_steps_slider.setValue(25)
        vc_params_sliders_layout.addWidget(self.diffusion_steps_slider, 0, 1)
        self.diffusion_steps_label = QLabel("25")
        vc_params_sliders_layout.addWidget(self.diffusion_steps_label, 0, 2)

        # Length Adjustment
        vc_params_sliders_layout.addWidget(QLabel("Length Adjustment:"), 1, 0)
        self.length_adjust_slider = QSlider(Qt.Horizontal)
        self.length_adjust_slider.setRange(50, 150); self.length_adjust_slider.setValue(100) 
        vc_params_sliders_layout.addWidget(self.length_adjust_slider, 1, 1)
        self.length_adjust_label = QLabel("1.00")
        vc_params_sliders_layout.addWidget(self.length_adjust_label, 1, 2)
        
        # Inference CFG Rate
        vc_params_sliders_layout.addWidget(QLabel("Inference CFG Rate:"), 2, 0)
        self.cfg_rate_slider = QSlider(Qt.Horizontal)
        self.cfg_rate_slider.setRange(10, 200); self.cfg_rate_slider.setValue(70) 
        vc_params_sliders_layout.addWidget(self.cfg_rate_slider, 2, 1)
        self.cfg_rate_label = QLabel("0.70")
        vc_params_sliders_layout.addWidget(self.cfg_rate_label, 2, 2)
        
        # Pitch Shift
        vc_params_sliders_layout.addWidget(QLabel("Pitch Shift (Semitones):"), 3, 0)
        self.pitch_shift_slider = QSlider(Qt.Horizontal)
        self.pitch_shift_slider.setRange(-12, 12); self.pitch_shift_slider.setValue(0)
        vc_params_sliders_layout.addWidget(self.pitch_shift_slider, 3, 1)
        self.pitch_shift_label = QLabel("0")
        vc_params_sliders_layout.addWidget(self.pitch_shift_label, 3, 2)
        
        vc_params_sliders_layout.setColumnStretch(1, 1) # Make slider column stretch
        
        # Layout for inside the VC Params GroupBox
        vc_params_group_internal_layout = QVBoxLayout()
        vc_params_group_internal_layout.addLayout(vc_params_sliders_layout)
        
        f0_checkboxes_layout = QHBoxLayout()
        self.f0_checkbox = QCheckBox("Use F0 Conditioned Model")
        self.auto_f0_checkbox = QCheckBox("Auto Adjust F0"); self.auto_f0_checkbox.setChecked(True)
        f0_checkboxes_layout.addWidget(self.f0_checkbox)
        f0_checkboxes_layout.addWidget(self.auto_f0_checkbox)
        vc_params_group_internal_layout.addLayout(f0_checkboxes_layout)
        vc_params_group.setLayout(vc_params_group_internal_layout)
        vc_params_tab_layout.addWidget(vc_params_group)
        vc_params_tab_layout.addStretch(1)
        self.settings_tab_widget.addTab(vc_params_tab, "Voice Conversion Params")

        # --- Tab 3: Scene Detection Speed Options ---
        scene_detection_tab = QWidget()
        scene_detection_layout = QVBoxLayout(scene_detection_tab)
        scene_detection_options_group = QGroupBox("Scene Detection Options")
        scene_detection_options_layout = QVBoxLayout()

        scene_detection_options_layout.addWidget(QLabel("Detector:"))
        self.sd_detector_combo = QComboBox()
        self.sd_detector_combo.addItems([
            "ContentDetector (Default)", "HashDetector (Faster)", 
            "AdaptiveDetector", "HistogramDetector (Faster, Experimental)"
        ])
        self.sd_detector_combo.setToolTip("Choose the algorithm for scene detection.\nFaster options might be less accurate.")
        scene_detection_options_layout.addWidget(self.sd_detector_combo)

        self.sd_auto_downscale_checkbox = QCheckBox("Enable Auto Downscaling (Recommended)")
        self.sd_auto_downscale_checkbox.setChecked(True)
        self.sd_auto_downscale_checkbox.setToolTip("Automatically reduces video resolution during detection for speed.\nRecommended for faster processing, especially on high-res videos.")
        scene_detection_options_layout.addWidget(self.sd_auto_downscale_checkbox)
        
        downscale_layout = QHBoxLayout()
        downscale_layout.addWidget(QLabel("Manual Downscale Factor:")) # Text updated
        self.sd_manual_downscale_edit = QLineEdit("1") 
        self.sd_manual_downscale_edit.setToolTip("Integer >= 1. Higher values are faster but less accurate (e.g., 2 means half resolution).\nOnly used if Auto Downscaling is disabled.")
        downscale_layout.addWidget(self.sd_manual_downscale_edit)
        scene_detection_options_layout.addLayout(downscale_layout)

        # New UI for "Check 1 frame every X seconds"
        skip_layout_seconds = QHBoxLayout()
        skip_layout_seconds.addWidget(QLabel("Process 1 frame every (seconds):"))
        self.sd_check_every_seconds_edit = QLineEdit("0.5") # Default to 0.5 second
        self.sd_check_every_seconds_edit.setToolTip("e.g., 0.5, 1, 2.0. Determines Frame Skip based on video FPS.\nProcessing fewer frames significantly speeds up detection but reduces accuracy.")
        double_validator_positive = QDoubleValidator(0.01, 3600.0, 2) # Min 0.01s, Max 1hr, 2 decimals
        double_validator_positive.setNotation(QDoubleValidator.StandardNotation)
        self.sd_check_every_seconds_edit.setValidator(double_validator_positive)
        skip_layout_seconds.addWidget(self.sd_check_every_seconds_edit)
        
        self.sd_calculated_frame_skip_label = QLabel("Approx. Frame Skip: N/A")
        skip_layout_seconds.addWidget(self.sd_calculated_frame_skip_label)
        scene_detection_options_layout.addLayout(skip_layout_seconds)

        skip_layout = QHBoxLayout()
        skip_layout.addWidget(QLabel("Manual Frame Skip:")) # Text updated
        self.sd_frame_skip_edit = QLineEdit("0") 
        self.sd_frame_skip_edit.setToolTip("Integer >= 0. Higher values skip more frames (faster, less accurate).\nThis is automatically calculated if 'Process 1 frame every X seconds' is used.")
        self.sd_frame_skip_edit.setReadOnly(True) # Make it read-only as it's auto-calculated now
        skip_layout.addWidget(self.sd_frame_skip_edit)
        scene_detection_options_layout.addLayout(skip_layout)
        
        threshold_layout = QHBoxLayout()
        threshold_layout.addWidget(QLabel("ContentDetector Threshold:"))
        self.sd_content_threshold_edit = QLineEdit("27.0")
        self.sd_content_threshold_edit.setToolTip("Threshold for ContentDetector. Lower is more sensitive. Default: 27.0")
        threshold_layout.addWidget(self.sd_content_threshold_edit)
        scene_detection_options_layout.addLayout(threshold_layout)
        
        scene_detection_options_group.setLayout(scene_detection_options_layout)
        scene_detection_layout.addWidget(scene_detection_options_group)
        scene_detection_layout.addStretch(1)
        self.settings_tab_widget.addTab(scene_detection_tab, "Scene Detection")

        vc_params_scroll_area.setWidget(self.settings_tab_widget) # Set TabWidget to ScrollArea
        top_row_layout.addWidget(vc_params_scroll_area, 1) # Add to top row, stretch factor 1 for settings

        self.main_layout.addWidget(top_row_container) # Add the sized top_row_container
        
        # --- Middle Area: Tab Widget for Steps ---
        self.step_tab_widget = QTabWidget()
        
        # --- Row 1: Source Audio & VC Marking ---
        row1_widget = QWidget()
        row1_layout = QHBoxLayout(row1_widget)

        # Row 1 - Left: Source Audio Pane
        source_audio_pane = QWidget()
        source_audio_layout = QVBoxLayout(source_audio_pane)
        source_audio_pane.setMinimumWidth(200)
        source_waveform_group = QGroupBox("Source Audio (Voice Conversion Marking)")
        source_waveform_layout = QVBoxLayout()
        self.source_waveform_widget = WaveformWidget(self) 
        source_waveform_layout.addWidget(self.source_waveform_widget, stretch=1)
        source_waveform_group.setLayout(source_waveform_layout) 
        source_audio_layout.addWidget(source_waveform_group, stretch=0) 

        audio_ctrl_layout = QHBoxLayout()
        self.play_source_button = QPushButton("Play Source") 
        self.pause_source_button = QPushButton("Pause")
        self.stop_source_button = QPushButton("Stop")
        audio_ctrl_layout.addWidget(self.play_source_button)
        audio_ctrl_layout.addWidget(self.pause_source_button)
        audio_ctrl_layout.addWidget(self.stop_source_button)
        source_audio_layout.addLayout(audio_ctrl_layout)
        row1_layout.addWidget(source_audio_pane, 1) # Stretch factor 1

        # Row 1 - Right: Time Markers / Speaker Diarization
        source_controls_group = QGroupBox("Time Markers / Speaker Diarization")
        source_controls_main_layout = QVBoxLayout(source_controls_group) # Main layout for the group

        self.source_controls_tab_widget = QTabWidget()

        # --- Tab 1: Manual Marking ---
        manual_marking_tab = QWidget()
        manual_marking_layout = QVBoxLayout(manual_marking_tab)

        self.mark_full_audio_button = QPushButton("Mark Full Audio for swap")
        self.mark_full_audio_button.setToolTip("Clears existing markers and adds a single marker spanning the entire source audio duration.")
        manual_marking_layout.addWidget(self.mark_full_audio_button)

        manual_marking_layout.addWidget(QLabel("Current Time Markers (Start - End seconds):"))
        self.source_marker_list_widget = QListWidget()
        self.source_marker_list_widget.setSelectionMode(QListWidget.ExtendedSelection)
        manual_marking_layout.addWidget(self.source_marker_list_widget, 1) # Stretch list

        marker_buttons_layout = QHBoxLayout()
        self.remove_selected_source_markers_button = QPushButton("Remove Selection Markers")
        self.clear_source_markers_button = QPushButton("Clear All Markers")
        marker_buttons_layout.addWidget(self.remove_selected_source_markers_button)
        marker_buttons_layout.addWidget(self.clear_source_markers_button)
        manual_marking_layout.addLayout(marker_buttons_layout)
        self.source_controls_tab_widget.addTab(manual_marking_tab, "Manual Marking")

        # --- Tab 2: Speaker Diarization ---
        speaker_diarization_tab = QWidget()
        speaker_diarization_layout = QVBoxLayout(speaker_diarization_tab)

        self.detect_speakers_button = QPushButton("Detect Speakers")
        if not PYANNOTE_AVAILABLE:
            self.detect_speakers_button.setEnabled(False)
            self.detect_speakers_button.setToolTip("pyannote.audio not installed")
        speaker_diarization_layout.addWidget(self.detect_speakers_button)

        self.speaker_list_combo = QComboBox()
        self.speaker_list_combo.setPlaceholderText("Detected Speakers")
        speaker_diarization_layout.addWidget(self.speaker_list_combo)

        self.apply_speaker_markers_button = QPushButton("Apply Speaker Markers")
        speaker_diarization_layout.addWidget(self.apply_speaker_markers_button)
        speaker_diarization_layout.addStretch(1)
        self.source_controls_tab_widget.addTab(speaker_diarization_tab, "Speaker Diarization")

        source_controls_main_layout.addWidget(self.source_controls_tab_widget)

        self.swap_voice_button = QPushButton("Swap Voice") # Create button instance here
        self.swap_voice_button.setToolTip("Process Marked Segments for Voice Conversion")
        self.swap_voice_button.setStyleSheet("QPushButton { background-color: #2196F3; color: white; padding: 5px; border-radius: 4px; }")
        source_controls_main_layout.addWidget(self.swap_voice_button)

        row1_layout.addWidget(source_controls_group, 1) # Stretch factor 1
        self.step_tab_widget.addTab(row1_widget, "1. Mark Swap Sections")

        # --- Row 2: Swapped Audio & RIR Configuration ---
        row2_widget = QWidget()
        row2_layout = QHBoxLayout(row2_widget)

        # Row 2 - Left: Swapped Audio Pane
        swapped_audio_pane = QWidget()
        swapped_audio_layout = QVBoxLayout(swapped_audio_pane)
        swapped_audio_pane.setMinimumWidth(200) # Consistent minimum width
        swapped_waveform_group = QGroupBox("Swapped Audio (RIR / Preview)")
        swapped_waveform_layout = QVBoxLayout()
        self.swapped_audio_waveform_widget = WaveformWidget(self)
        swapped_waveform_layout.addWidget(self.swapped_audio_waveform_widget, stretch=1)
        swapped_waveform_group.setLayout(swapped_waveform_layout)
        swapped_audio_layout.addWidget(swapped_waveform_group, stretch=0) 
        swapped_audio_ctrl_layout = QHBoxLayout()
        self.play_swapped_button = QPushButton("Play Swapped") 
        self.pause_swapped_button = QPushButton("Pause Swapped") 
        self.stop_swapped_button = QPushButton("Stop Swapped")   
        swapped_audio_ctrl_layout.addWidget(self.play_swapped_button)
        swapped_audio_ctrl_layout.addWidget(self.pause_swapped_button)
        swapped_audio_ctrl_layout.addWidget(self.stop_swapped_button)
        swapped_audio_layout.addLayout(swapped_audio_ctrl_layout)
        row2_layout.addWidget(swapped_audio_pane, 1) # Stretch factor 1

        # Row 2 - Right: RIR Configuration Mode
        rir_config_group = QGroupBox("RIR Configuration Mode")
        rir_config_main_layout = QVBoxLayout(rir_config_group) # Main layout for the group

        # Speech2RIR Global Checkbox (remains at the top)
        self.use_speech2rir_checkbox = QCheckBox("Auto RIR from Source (Speech2RIR) - static environment")
        self.use_speech2rir_checkbox.setToolTip(
            "If checked, RIR will be estimated from the original video's audio and applied globally.\n"
            "Manual RIR markers below will be disabled."
        )
        if S2R_Generator is None:
             self.use_speech2rir_checkbox.setEnabled(False)
             self.use_speech2rir_checkbox.setToolTip("Speech2RIR model components not found or failed to import.")
             self.use_speech2rir_checkbox.setText("Auto RIR from Source (Speech2RIR) - N/A")
        else:
             self.use_speech2rir_checkbox.setEnabled(True)
             self.use_speech2rir_checkbox.setText("Auto RIR from Source (Speech2RIR)- static environment")
             
        self.use_speech2rir_checkbox.setChecked(False) 
        rir_config_main_layout.addWidget(self.use_speech2rir_checkbox)

        self.rir_config_tab_widget = QTabWidget()

        # --- Tab 1: Manual & Auto-Segmented RIR ---
        # This tab was previously self.manual_rir_group
        manual_auto_rir_tab = QWidget()
        # manual_rir_layout will be the layout for this tab
        manual_rir_layout = QVBoxLayout()
        manual_auto_rir_tab.setLayout(manual_rir_layout) # Set layout for the tab early

        # --- Collapsible Group: Auto-Segmentation Settings ---
        auto_segment_group = QGroupBox("Auto-Segmentation Settings")
        auto_segment_group.setCheckable(True)
        auto_segment_group.setChecked(False) # Start collapsed
        auto_segment_layout = QVBoxLayout(auto_segment_group)
        
        auto_rir_config_layout = QHBoxLayout()
        auto_rir_config_layout.addWidget(QLabel("Auto Speech2RIR Segment (s):"))
        self.auto_rir_segment_duration_edit = QLineEdit(str(self.auto_rir_segment_duration_sec))
        self.auto_rir_segment_duration_edit.setToolTip("Duration in seconds for segments when using 'Generate Auto RIR Markers'.")
        double_validator = QDoubleValidator(1.0, 300.0, 1) 
        double_validator.setNotation(QDoubleValidator.StandardNotation)
        self.auto_rir_segment_duration_edit.setValidator(double_validator)
        auto_rir_config_layout.addWidget(self.auto_rir_segment_duration_edit)
        auto_segment_layout.addLayout(auto_rir_config_layout)
        
        # --- Moved Detect Scene Cuts Button Here ---
        self.detect_scene_cuts_button = QPushButton("Detect Scene Changes in Video")
        if not SCENEDETECT_AVAILABLE:
            self.detect_scene_cuts_button.setEnabled(False)
            self.detect_scene_cuts_button.setToolTip("PySceneDetect library not installed.")
        auto_segment_layout.addWidget(self.detect_scene_cuts_button)

        gen_auto_rir_scene_layout = QHBoxLayout()
        self.generate_auto_rir_markers_button = QPushButton("Auto Generate RIR Markers (Speech2RIR/segment)")
        self.generate_auto_rir_markers_button.setToolTip("Automatically segment original source audio and estimate an RIR for each segment using Speech2RIR.\nThese RIRs are then applied to the corresponding segments of the swapped audio.")
        gen_auto_rir_scene_layout.addWidget(self.generate_auto_rir_markers_button)
        self.use_scene_cuts_for_auto_rir_checkbox = QCheckBox("Use detected scene changes for Auto RIR") # Simplified
        self.use_scene_cuts_for_auto_rir_checkbox.setToolTip("If scene cuts are detected and this is checked, Auto RIR will use scenes as segments.")
        self.use_scene_cuts_for_auto_rir_checkbox.setEnabled(False) 
        gen_auto_rir_scene_layout.addWidget(self.use_scene_cuts_for_auto_rir_checkbox)
        auto_segment_layout.addLayout(gen_auto_rir_scene_layout)
        manual_rir_layout.addWidget(auto_segment_group)

        # --- Collapsible Group: Manual RIR Marker Settings ---
        manual_marker_group = QGroupBox("Manual RIR Marker Settings")
        manual_marker_group.setCheckable(True)
        manual_marker_group.setChecked(False) # Start collapsed
        manual_marker_layout = QVBoxLayout(manual_marker_group)
        
        add_rir_rt60_layout = QHBoxLayout()
        self.add_rir_marker_button = QPushButton("Add RIR Marker for Selection")
        add_rir_rt60_layout.addWidget(self.add_rir_marker_button)
        self.use_estimated_rt60_checkbox = QCheckBox("Use Estimated BlindRT60 value from Global Config") 
        self.use_estimated_rt60_checkbox.setToolTip("Use Estimated BlindRT60 value for RIR markers")
        self.use_estimated_rt60_checkbox.setChecked(False) 
        add_rir_rt60_layout.addWidget(self.use_estimated_rt60_checkbox)
        manual_marker_layout.addLayout(add_rir_rt60_layout)

        rir_param_layout = QHBoxLayout()
        rir_param_layout.addWidget(QLabel("Target RT60 (s) for new RIR marker:"))
        self.rir_target_rt60_edit = QLineEdit("0.5") 
        self.rir_target_rt60_edit.setPlaceholderText("e.g., 0.5")
        rir_param_layout.addWidget(self.rir_target_rt60_edit)
        manual_marker_layout.addLayout(rir_param_layout)
        manual_rir_layout.addWidget(manual_marker_group)

        # --- RIR Marker List (remains directly in the tab layout) ---
        manual_rir_layout.addWidget(QLabel("Current RIR Markers (Start - End, Params):"))
        self.rir_marker_list_widget = QListWidget()
        manual_rir_layout.addWidget(self.rir_marker_list_widget, 1) # Add stretch here
        rir_buttons_layout = QHBoxLayout()
        self.remove_selected_rir_markers_button = QPushButton("Remove Selected RIR Markers")
        self.clear_rir_markers_button = QPushButton("Clear All RIR Markers")
        rir_buttons_layout.addWidget(self.remove_selected_rir_markers_button)
        rir_buttons_layout.addWidget(self.clear_rir_markers_button)
        manual_rir_layout.addLayout(rir_buttons_layout)
        self.rir_config_tab_widget.addTab(manual_auto_rir_tab, "Manual & Auto-Segmented RIR")

        # --- Tab 2: Global Speech2RIR Settings ---
        global_s2r_tab = QWidget()
        global_s2r_layout = QVBoxLayout(global_s2r_tab)
        
        self.s2r_mix_layout = QHBoxLayout() # Keep this as QHBoxLayout
        self.s2r_mix_label = QLabel("Speech2RIR Dry/Wet Mix (Wet %):") # This label is part of s2r_mix_layout
        self.s2r_mix_slider = QSlider(Qt.Horizontal)
        self.s2r_mix_slider.setRange(0, 100); self.s2r_mix_slider.setValue(30) 
        self.s2r_mix_value_label = QLabel("30%")
        self.s2r_mix_layout.addWidget(self.s2r_mix_label)
        self.s2r_mix_layout.addWidget(self.s2r_mix_slider)
        self.s2r_mix_layout.addWidget(self.s2r_mix_value_label)
        global_s2r_layout.addLayout(self.s2r_mix_layout)
        global_s2r_layout.addWidget(QLabel("Note: When 'Auto RIR from Source' is checked,\n RIR is estimated globally and applied with this mix."))
        global_s2r_layout.addStretch(1)
        global_s2r_tab.setLayout(global_s2r_layout)
        self.rir_config_tab_widget.addTab(global_s2r_tab, "Global Speech2RIR")

        rir_config_main_layout.addWidget(self.rir_config_tab_widget)

        # New button for generating full preview based on Row 2 RIR settings
        self.generate_full_preview_button = QPushButton("Apply RIRs for Final mix")
        self.generate_full_preview_button.setToolTip(
            "Generates an audio preview using the current RIR markers or global Speech2RIR settings from this section."
        )
        rir_config_main_layout.addWidget(self.generate_full_preview_button)

        row2_layout.addWidget(rir_config_group, 1) # Stretch factor 1
        self.step_tab_widget.addTab(row2_widget, "2. Apply RIR Values")

        # --- Row 3: Final Mix Preview & RIR Adjustment ---
        self.final_mix_rir_settings = [] # For storing adjusted RIRs for final mix
        row3_widget = QWidget()
        row3_layout = QHBoxLayout(row3_widget)

        # Row 3 - Left: Final Mix Waveform Pane
        final_mix_waveform_pane = QWidget()
        final_mix_waveform_layout = QVBoxLayout(final_mix_waveform_pane)
        final_mix_waveform_group = QGroupBox("Final Mix Preview (with RIR adjustments)")
        final_mix_waveform_group_layout = QVBoxLayout(final_mix_waveform_group)
        self.final_mix_waveform_widget = WaveformWidget(self)
        self.final_mix_waveform_widget.hide() # Initially hidden
        final_mix_waveform_group_layout.addWidget(self.final_mix_waveform_widget, 1)
        
        final_mix_audio_ctrl_layout = QHBoxLayout()
        self.play_final_mix_button = QPushButton("Play Preview Mix")
        self.pause_final_mix_button = QPushButton("Pause Preview Mix")
        self.stop_final_mix_button = QPushButton("Stop Preview Mix")
        final_mix_audio_ctrl_layout.addWidget(self.play_final_mix_button)
        final_mix_audio_ctrl_layout.addWidget(self.pause_final_mix_button)
        final_mix_audio_ctrl_layout.addWidget(self.stop_final_mix_button)
        final_mix_waveform_group_layout.addLayout(final_mix_audio_ctrl_layout)
        final_mix_waveform_layout.addWidget(final_mix_waveform_group)
        row3_layout.addWidget(final_mix_waveform_pane, 1)

        # Row 3 - Right: RIR Adjustment Controls
        final_mix_controls_group = QGroupBox("RIR Adjustments for Final Mix")
        final_mix_controls_layout = QGridLayout(final_mix_controls_group) # Use QGridLayout
        self.final_mix_segment_list = QListWidget() # To show segments for adjustment
        self.final_mix_segment_list.itemSelectionChanged.connect(self._on_final_mix_segment_selected) # Keep this

        final_mix_controls_layout.addWidget(QLabel("Segments for Effect Adjustment:"), 0, 0, 1, 2) # Span 2 columns
        final_mix_controls_layout.addWidget(self.final_mix_segment_list, 1, 0, 1, 2) 

        # --- Effect Parameters Tab Widget ---
        self.effects_settings_tab_widget = QTabWidget()

        # --- Tab 1: Basic RIR & Volume ---
        basic_effects_tab = QWidget()
        basic_effects_layout = QGridLayout(basic_effects_tab)

        basic_effects_layout.addWidget(QLabel("Simulated RT60 (s):"), 0, 0)
        self.final_mix_rt60_edit = QLineEdit("0.5")
        basic_effects_layout.addWidget(self.final_mix_rt60_edit, 0, 1)
        basic_effects_layout.addWidget(QLabel("Sim. Room Dim (X,Y,Z m):"), 1, 0)
        self.final_mix_room_dim_edit = QLineEdit("6,4,2.8") # Default
        basic_effects_layout.addWidget(self.final_mix_room_dim_edit, 1, 1)
        basic_effects_layout.addWidget(QLabel("Volume (dB):"), 2, 0)
        self.final_mix_volume_edit = QLineEdit("0.0")
        basic_effects_layout.addWidget(self.final_mix_volume_edit, 2, 1)
        basic_effects_layout.setRowStretch(3, 1) # Add stretch
        self.effects_settings_tab_widget.addTab(basic_effects_tab, "RIR & Volume")

        # --- Tab 2: Gain Effect ---
        gain_effect_tab = QWidget()
        gain_effect_layout = QGridLayout(gain_effect_tab)
        self.final_fx_gain_enable_checkbox = QCheckBox("Enable Gain")
        gain_effect_layout.addWidget(self.final_fx_gain_enable_checkbox, 0, 0, 1, 2)
        gain_effect_layout.addWidget(QLabel("Gain Amount (dB):"), 1, 0)
        self.final_fx_gain_slider = QSlider(Qt.Horizontal)
        self.final_fx_gain_slider.setRange(-24, 24); self.final_fx_gain_slider.setValue(0)
        self.final_fx_gain_label = QLabel("0 dB")
        gain_effect_layout.addWidget(self.final_fx_gain_slider, 1, 1)
        gain_effect_layout.addWidget(self.final_fx_gain_label, 1, 2)
        gain_effect_layout.setRowStretch(2,1)
        self.effects_settings_tab_widget.addTab(gain_effect_tab, "Gain")
        
        # --- Add more effect tabs here (Compressor, EQ, Reverb etc.) ---
        # Example:
        # reverb_effect_tab = QWidget()
        # ... setup layout and controls for pedalboard.Reverb ...
        # self.effects_settings_tab_widget.addTab(reverb_effect_tab, "Reverb")

        final_mix_controls_layout.addWidget(self.effects_settings_tab_widget, 2, 0, 1, 2) # Span 2 columns

        self.apply_segment_effects_button = QPushButton("Apply & Preview Segment Effects")
        final_mix_controls_layout.addWidget(self.apply_segment_effects_button, 3, 0, 1, 2) # Adjusted row

        # --- Demucs Background Mixing Controls ---
        self.demucs_mix_group = QGroupBox("Mix with Original Background") # Assign to self
        self.demucs_mix_group.setCheckable(True) # Allow enabling/disabling the whole group
        self.demucs_mix_group.setChecked(False) # Start disabled
        demucs_mix_layout = QGridLayout(self.demucs_mix_group)

        self.enable_demucs_mix_checkbox = QCheckBox("Enable Demucs Background Mix")
        self.enable_demucs_mix_checkbox.setToolTip("If checked, separates audio from original video, extracts background,\nand mixes it with the swapped voice before final video creation.")
        demucs_mix_layout.addWidget(self.enable_demucs_mix_checkbox, 0, 0, 1, 3) # Span 3 columns

        demucs_mix_layout.addWidget(QLabel("Original Background Volume Ratio:"), 1, 0)
        self.demucs_bg_mix_slider = QSlider(Qt.Horizontal)
        self.demucs_bg_mix_slider.setRange(0, 100); self.demucs_bg_mix_slider.setValue(60) # Default 60%
        demucs_mix_layout.addWidget(self.demucs_bg_mix_slider, 1, 1)
        self.demucs_bg_mix_label = QLabel("0.60")
        demucs_mix_layout.addWidget(self.demucs_bg_mix_label, 1, 2)
        final_mix_controls_layout.addWidget(self.demucs_mix_group, 4, 0, 1, 2) # Add Demucs group

        # --- Final Process Button ---
        self.final_process_button = QPushButton("Final Process") # Create button instance here
        self.final_process_button.setToolTip("Apply RIR & Create Final Video using adjustments from this section.")
        self.final_process_button.setStyleSheet("QPushButton { background-color: #4CAF50; color: white; padding: 5px; border-radius: 4px; font-weight: bold; }")

        final_mix_controls_layout.addWidget(self.final_process_button, 6, 0, 1, 2) 
        
        final_mix_controls_layout.setRowStretch(1, 1) # Segment list stretches
        row3_layout.addWidget(final_mix_controls_group, 1)
        self.step_tab_widget.addTab(row3_widget, "3. Final Mix and Preview")

        # --- Bottom Section: Processing & Status ---
        bottom_section_layout = QHBoxLayout()
        process_status_group = QGroupBox("Processing Controls & Status")
        process_status_layout = QVBoxLayout()
        # process_buttons_layout is removed as buttons are relocated
        self.progress_bar = QProgressBar()
        process_status_layout.addWidget(self.progress_bar)
        self.status_box = QTextEdit()
        self.status_box.setReadOnly(True); self.status_box.setFixedHeight(96) 
        process_status_layout.addWidget(self.status_box)
        
        output_video_path_layout = QHBoxLayout() 
        output_video_path_layout.addWidget(QLabel("Output Video Path:"))
        self.output_video_label = QLineEdit("Output will appear here.")
        self.output_video_label.setReadOnly(True)
        output_video_path_layout.addWidget(self.output_video_label) 
        process_status_layout.addLayout(output_video_path_layout) 
        process_status_group.setLayout(process_status_layout)
        bottom_section_layout.addWidget(process_status_group, stretch=1) 
        
        self.main_layout.addWidget(self.step_tab_widget, 1) # Add the tab widget with stretch
        self.main_layout.addLayout(bottom_section_layout) # Add the bottom section to the main layout
        self._update_ui_for_model_availability() # Initial check
        self._on_environment_changed() 

    def _update_ui_for_model_availability(self):
        """Updates UI elements based on whether S2R or ClearVoice models are available."""
        s2r_available = S2R_Generator is not None
        self.use_speech2rir_checkbox.setEnabled(s2r_available)
        self.generate_auto_rir_markers_button.setEnabled(s2r_available and self.swapped_audio_full_path is not None)
        if not s2r_available:
            self.use_speech2rir_checkbox.setToolTip("Speech2RIR model components not found or failed to import.")
            self.use_speech2rir_checkbox.setText("Auto RIR from Source (Speech2RIR) - N/A")
            self.generate_auto_rir_markers_button.setToolTip("Speech2RIR model components not found.")

        self.clean_target_voice_button.setEnabled(CLEARVOICE_AVAILABLE and self.target_voice_path is not None)
        if not CLEARVOICE_AVAILABLE:
            self.clean_target_voice_button.setToolTip("ClearVoice library not found or import error.")


    def _connect_signals(self):
        self.mp4_button.clicked.connect(self.select_mp4_file)
        self.wav_button.clicked.connect(self.select_wav_file)
        self.output_dir_button.clicked.connect(self.select_output_dir)
        self.calculate_rt60_button.clicked.connect(self.run_rt60_estimation_thread)

        self.detect_scene_cuts_button.clicked.connect(self.run_scene_cut_detection_thread)       
        self.sd_check_every_seconds_edit.textChanged.connect(self._update_frame_skip_from_seconds_input)

        self.diffusion_steps_slider.valueChanged.connect(lambda v: self.diffusion_steps_label.setText(str(v)))
        self.length_adjust_slider.valueChanged.connect(lambda v: self.length_adjust_label.setText(f"{v/100.0:.2f}"))
        self.cfg_rate_slider.valueChanged.connect(lambda v: self.cfg_rate_label.setText(f"{v/100.0:.2f}"))
        self.pitch_shift_slider.valueChanged.connect(lambda v: self.pitch_shift_label.setText(str(v)))

        self.play_source_button.clicked.connect(lambda: self.play_audio_segment(is_source=True))
        self.pause_source_button.clicked.connect(self.pause_audio) 
        self.stop_source_button.clicked.connect(self.stop_audio)   
        self.source_waveform_widget.clickedPosition.connect(lambda t: self.seek_audio(t, is_source=True)) 

        self.play_swapped_button.clicked.connect(lambda: self.play_audio_segment(is_source=False))
        self.pause_swapped_button.clicked.connect(self.pause_audio) 
        self.stop_swapped_button.clicked.connect(self.stop_audio)   
        self.swapped_audio_waveform_widget.clickedPosition.connect(lambda t: self.seek_audio(t, is_source=False))

        self.source_waveform_widget.selectionFinalized.connect(self.add_current_selection_as_marker) 
        self.mark_full_audio_button.clicked.connect(self.mark_full_source_audio_action)
        
        self.detect_speakers_button.clicked.connect(self.run_speaker_diarization_thread)
        self.apply_speaker_markers_button.clicked.connect(self.apply_selected_speaker_to_markers)
        self.remove_selected_source_markers_button.clicked.connect(self.remove_selected_source_markers_from_list)
        self.clear_source_markers_button.clicked.connect(self.clear_all_source_markers_action)
        
        self.swapped_audio_waveform_widget.selectionFinalized.connect(self.add_current_rir_selection_as_marker) 
        # Connect new buttons for final mix preview
        self.play_final_mix_button.clicked.connect(lambda: self.play_audio_segment(is_source=False, is_final_mix_preview=True))
        self.pause_final_mix_button.clicked.connect(self.pause_audio)
        self.stop_final_mix_button.clicked.connect(self.stop_audio)
        self.final_mix_waveform_widget.clickedPosition.connect(lambda t: self.seek_audio(t, is_source=False, is_final_mix_preview=True))
        self.final_mix_waveform_widget.markerClicked.connect(self._on_final_mix_waveform_marker_clicked)
        self.add_rir_marker_button.clicked.connect(self.add_current_rir_selection_as_marker) 
        self.remove_selected_rir_markers_button.clicked.connect(self.remove_selected_rir_markers_from_list)
        self.clear_rir_markers_button.clicked.connect(self.clear_all_rir_markers_action)
        self.use_estimated_rt60_checkbox.toggled.connect(self._update_rir_rt60_input_state)
        self.generate_auto_rir_markers_button.clicked.connect(self.run_auto_rir_generation_thread)
        self.auto_rir_segment_duration_edit.textChanged.connect(self._update_auto_rir_segment_duration)
        self.generate_full_preview_button.clicked.connect(self.initiate_full_rir_preview_generation)
        
        self.use_speech2rir_checkbox.toggled.connect(self._on_speech2rir_toggled)
        self.s2r_mix_slider.valueChanged.connect(lambda v: self.s2r_mix_value_label.setText(f"{v}%"))
        
        self.swap_voice_button.clicked.connect(self.initiate_voice_swap_pipeline)
        
        self.clean_target_voice_button.clicked.connect(self.run_target_voice_cleaning_thread)        

        self.apply_segment_effects_button.clicked.connect(self.apply_and_preview_segment_effects)

        self.final_process_button.clicked.connect(self.initiate_final_processing_pipeline)
        
        # Connect signals for Gain effect example
        self.final_fx_gain_enable_checkbox.toggled.connect(self._update_current_segment_effects_params)
        self.final_fx_gain_slider.valueChanged.connect(lambda v: [self.final_fx_gain_label.setText(f"{v} dB"), self._update_current_segment_effects_params()])

        # Connect Demucs UI
        self.enable_demucs_mix_checkbox.toggled.connect(lambda checked: self.demucs_bg_mix_slider.setEnabled(checked)) # Enable slider only if checkbox is checked
        self.demucs_bg_mix_slider.valueChanged.connect(lambda v: self.demucs_bg_mix_label.setText(f"{v/100.0:.2f}"))
        
        self.environment_combo.currentTextChanged.connect(self._on_environment_changed)

    def _on_environment_changed(self):
        current_environment = self.environment_combo.currentText()
        is_outside = (current_environment == "Outside")
        self.log_status(f"Environment preset changed to: {current_environment}")
        
        # Update default room dimensions based on preset
        if current_environment == "Living Room": # Match exact string from combo box
            self.room_dim_edit.setText("6,4,2.8")
        elif current_environment == "Outside":
            # For "Outside", RIR simulation might use large, open parameters
            self.room_dim_edit.setText("100,100,20") # Example large dimensions

        s2r_features_conditionally_available = (S2R_Generator is not None) and (not is_outside)

        self.use_speech2rir_checkbox.setEnabled(s2r_features_conditionally_available)
        
        if not s2r_features_conditionally_available:
            if self.use_speech2rir_checkbox.isChecked():
                self.use_speech2rir_checkbox.setChecked(False) 
            
            self.use_speech2rir_checkbox.setToolTip(
                "Speech2RIR features are unavailable for the 'Outside' environment or S2R model components are not loaded."
            )
            self.use_speech2rir_checkbox.setText("Auto RIR from Source (Speech2RIR) - N/A (Unavailable for Outside/No Model)")
        else: 
             self.use_speech2rir_checkbox.setToolTip(
                "If checked, RIR will be estimated from the original video's audio and applied globally.\n"
                "Manual RIR markers below will be disabled."
            )
             self.use_speech2rir_checkbox.setText("Auto RIR from Source (Speech2RIR)- static environment")
        
        self._on_speech2rir_toggled(self.use_speech2rir_checkbox.isChecked())

    def _on_speech2rir_toggled(self, checked):
        self.log_status(f"Speech2RIR mode toggled: {'Enabled' if checked else 'Disabled'}")
        is_outside = self.environment_combo.currentText() == "Outside"
        
        # Manual RIR group is enabled if S2R is NOT checked AND it's not "Outside"
        manual_rir_active = not checked and not is_outside
        
        # Enable/disable the "Manual & Auto-Segmented RIR" tab
        manual_auto_tab_index = self.rir_config_tab_widget.indexOf(self.rir_config_tab_widget.findChild(QWidget, "manual_auto_rir_tab_widget_name")) # Find tab by object name if set, or by index
        # Assuming "Manual & Auto-Segmented RIR" is tab 0 and "Global Speech2RIR" is tab 1
        manual_auto_tab_index = 0 
        global_s2r_tab_index = 1

        self.rir_config_tab_widget.setTabEnabled(manual_auto_tab_index, manual_rir_active)
        
        # Enable/disable controls within the "Manual & Auto-Segmented RIR" tab
        # These controls are children of manual_auto_rir_tab (formerly manual_rir_group)
        self.generate_auto_rir_markers_button.setEnabled(
            manual_rir_active and S2R_Generator is not None and self.swapped_audio_full_path is not None
        )
        self.use_scene_cuts_for_auto_rir_checkbox.setEnabled(
            manual_rir_active and S2R_Generator is not None and len(self.scene_cuts_sec) > 0
        )
        # Other controls within this tab (like add_rir_marker_button, rir_target_rt60_edit etc.)
        # will also effectively be controlled by the tab's enabled state.

        # S2R Mix slider (in "Global Speech2RIR" tab) is enabled if S2R IS checked AND it's not "Outside"
        s2r_mix_active = checked and not is_outside and S2R_Generator is not None
        self.s2r_mix_label.setEnabled(s2r_mix_active)
        self.s2r_mix_slider.setEnabled(s2r_mix_active)
        self.s2r_mix_value_label.setEnabled(s2r_mix_active)
        self.rir_config_tab_widget.setTabEnabled(global_s2r_tab_index, s2r_mix_active)

        if checked and S2R_Generator is not None and not is_outside: # S2R mode enabled and possible
            self.rir_config_tab_widget.setCurrentIndex(global_s2r_tab_index) # Switch to Global S2R tab
            self.estimated_rir_path = None 
            if self.source_audio_full_path:
                self.log_status("Speech2RIR enabled, source audio loaded. Triggering RIR estimation.")
                self.run_s2r_rir_estimation_thread()
            else:
                self.log_status("Speech2RIR enabled, but no source audio loaded. RIR will be estimated when audio is loaded.", is_error=False)
            if self.swapped_audio_markers:
                self.log_status("Speech2RIR mode enabled, clearing existing manual RIR markers.")
                self.clear_all_rir_markers_action()
        elif checked and (S2R_Generator is None or is_outside): # S2R mode selected but not possible
            self.log_status("Speech2RIR cannot be enabled: S2R components missing or 'Outside' environment.", is_error=True)
            self.use_speech2rir_checkbox.setChecked(False) # Force uncheck
            return 
        else: # Speech2RIR mode is OFF
            self.log_status("Speech2RIR mode disabled. Manual RIR marking is now active (if not 'Outside').")
            if manual_rir_active: self.rir_config_tab_widget.setCurrentIndex(manual_auto_tab_index) # Switch to Manual tab

    @Slot()
    def _update_frame_skip_from_seconds_input(self):
        if self.source_video_fps is None or self.source_video_fps <= 0:
            self.sd_calculated_frame_skip_label.setText("Approx. Frame Skip: (Video FPS N/A)")
            self.sd_frame_skip_edit.setText("0") # Default to 0 if FPS unknown
            return
        try:
            seconds_interval_str = self.sd_check_every_seconds_edit.text().strip()
            if not seconds_interval_str: # Empty input
                 self.sd_calculated_frame_skip_label.setText("Approx. Frame Skip: (Enter interval)")
                 self.sd_frame_skip_edit.setText("0")
                 return

            seconds_interval = float(seconds_interval_str)
            if seconds_interval <= 0:
                self.sd_frame_skip_edit.setText("0")
                self.sd_calculated_frame_skip_label.setText("Approx. Frame Skip: 0 (Invalid interval)")
                return
            
            frame_skip_value = max(0, round(self.source_video_fps * seconds_interval) - 1)
            self.sd_frame_skip_edit.setText(str(frame_skip_value))
            self.sd_calculated_frame_skip_label.setText(f"Approx. Frame Skip: {frame_skip_value}")
        except ValueError:
            self.sd_calculated_frame_skip_label.setText("Approx. Frame Skip: (Invalid input)")
            self.sd_frame_skip_edit.setText("0") 

    def _update_auto_rir_segment_duration(self, text):
        try:
            val = float(text)
            if 1.0 <= val <= 300.0: 
                self.auto_rir_segment_duration_sec = val
        except ValueError:
            pass 

    def _update_rir_rt60_input_state(self, checked):
        if checked:
            self.rir_target_rt60_edit.setEnabled(False)
            self.rir_target_rt60_edit.setToolTip("RT60 will be taken from the estimated source RT60.")
        else:
            self.rir_target_rt60_edit.setEnabled(True)
            self.rir_target_rt60_edit.setToolTip("Manually enter the target RT60 in seconds for new RIR markers.")

    def log_status(self, message, is_error=False):
        timestamp = time.strftime("%H:%M:%S")
        prefix = "ERROR: " if is_error else "INFO: "        
        try:
            # Check if the widget and its C++ part still exist
            if self.status_box and self.status_box.parent(): # Basic check
                self.status_box.append(f"[{timestamp}] {prefix}{message}")
                if is_error: print(f"ERROR LOGGED: {message}") # Also print errors to console
                QApplication.processEvents()
            else:
                print(f"Skipped logging to GUI (status_box likely deleted): [{timestamp}] {prefix}{message}")
        except RuntimeError: # Catch the specific error if the object is already deleted
            print(f"Failed to log to GUI (RuntimeError): [{timestamp}] {prefix}{message}")

    @Slot(str, bool)
    def _handle_worker_log_message(self, message, is_error):
        self.log_status(message, is_error=is_error)


    def select_output_dir(self):
        dir_path = QFileDialog.getExistingDirectory(self, "Select Output Directory", str(self.output_dir_path) if self.output_dir_path else str(Path.cwd()))
        if dir_path:
            self.output_dir_path = Path(dir_path)
            self.output_dir_label.setText(str(self.output_dir_path))
            self.log_status(f"Selected Output Directory: {self.output_dir_path}")
            try:
                self.output_dir_path.mkdir(parents=True, exist_ok=True)
                self.log_status(f"Ensured output directory exists: {self.output_dir_path}")
            except Exception as e:
                self.log_status(f"Error creating output directory {self.output_dir_path}: {e}", is_error=True)
                QMessageBox.critical(self, "Directory Error", f"Could not create or access output directory: {e}")
                self.output_dir_path = None 
                self.output_dir_label.setText("Error selecting directory!")

    def select_mp4_file(self):
        filepath, _ = QFileDialog.getOpenFileName(self, "Select Input MP4 File", "", "MP4 Files (*.mp4)")
        if filepath:
            self.input_mp4_path = Path(filepath)
            self.mp4_label.setText(self.input_mp4_path.name)
            self.log_status(f"Selected MP4: {filepath}")
            self._extract_and_load_audio()

    def _update_waveform_cursor_on_play(self, position_ms): 
        if self.player.source().isValid():
            source_file = Path(self.player.source().toLocalFile())
            if self.source_audio_full_path and source_file == self.source_audio_full_path and self.source_audio_duration_sec > 0:
                time_sec = position_ms / 1000.0
                self.source_waveform_widget.set_playback_cursor(time_sec)
            elif self.swapped_audio_full_path and source_file == self.swapped_audio_full_path and self.swapped_audio_waveform_widget.duration_sec > 0:
                time_sec = position_ms / 1000.0
                self.swapped_audio_waveform_widget.set_playback_cursor(time_sec)


    def _on_playback_state_changed(self, state): 
        active_waveform = None
        duration_sec = 0
        if self.player.source().isValid():
            source_file = Path(self.player.source().toLocalFile())
            if self.source_audio_full_path and source_file == self.source_audio_full_path:
                active_waveform = self.source_waveform_widget
                duration_sec = self.source_audio_duration_sec
            elif self.swapped_audio_full_path and source_file == self.swapped_audio_full_path:
                active_waveform = self.swapped_audio_waveform_widget
                duration_sec = self.swapped_audio_waveform_widget.duration_sec
            
        if active_waveform:
            if state == QMediaPlayer.PlaybackState.StoppedState:
                # If playback stopped near the end, reset cursor to beginning
                if self.player.duration() > 0 and self.player.position() >= self.player.duration() - 150: 
                    active_waveform.set_playback_cursor(0)
                # Otherwise, keep cursor at current stopped position if not explicitly reset by stop_audio
                # (stop_audio already handles resetting to 0 or selection start)
            elif self.player.source().isValid() and hasattr(self, 'final_mix_preview_audio_path') and self.final_mix_preview_audio_path and source_file == self.final_mix_preview_audio_path:
                active_waveform = self.final_mix_waveform_widget
                # duration_sec = active_waveform.duration_sec # Already handled by waveform widget


    def play_audio_segment(self, is_source=True, is_final_mix_preview=False):
        if is_final_mix_preview:
            audio_path_to_play = getattr(self, 'final_mix_preview_audio_path', None)
            waveform_widget_to_use = self.final_mix_waveform_widget
        else:
            audio_path_to_play = self.source_audio_full_path if is_source else self.swapped_audio_full_path
            waveform_widget_to_use = self.source_waveform_widget if is_source else self.swapped_audio_waveform_widget

        if not audio_path_to_play:
            self.log_status(f"Audio path for {'source' if is_source else ('swapped' if not is_final_mix_preview else 'final mix preview')} not set.", is_error=True)
            return
        if not Path(audio_path_to_play).exists():
            self.log_status(f"Audio file not found: {audio_path_to_play}", is_error=True)
            return

        current_player_source_path = Path(self.player.source().toLocalFile()) if self.player.source().isValid() else None

        if self.player.playbackState() == QMediaPlayer.PlaybackState.PlayingState and \
           current_player_source_path == audio_path_to_play:
            self.pause_audio()
            return
        
        if self.player.playbackState() == QMediaPlayer.PlaybackState.PausedState and \
           current_player_source_path == audio_path_to_play:
            self.player.play()
            current_pos_s = self.player.position() / 1000.0
            # If segment timer was set, re-arm it for remaining duration
            if hasattr(self, '_segment_end_s') and self._segment_end_s is not None and self._segment_end_s > current_pos_s:
                if hasattr(self, '_play_segment_timer'): self._play_segment_timer.stop()
                remaining_duration_ms = int((self._segment_end_s - current_pos_s) * 1000)
                if remaining_duration_ms > 0:
                    self._play_segment_timer.start(remaining_duration_ms)
            self.log_status("Audio resumed.")
            return

        if self.player.playbackState() != QMediaPlayer.PlaybackState.StoppedState:
            self.player.stop() # Stop if playing/paused with a different source

        start_s, end_s = waveform_widget_to_use.get_current_selection()

        if start_s is None or end_s is None: 
            start_s = waveform_widget_to_use.playback_cursor_pos_sec
            if start_s >= waveform_widget_to_use.duration_sec - 0.05: # If at end, play from start
                start_s = 0
            end_s = waveform_widget_to_use.duration_sec 
            self.log_status(f"Playing from cursor at {start_s:.2f}s to end.")
        else: 
            self.log_status(f"Playing selection from {start_s:.2f}s to {end_s:.2f}s.")

        self.player.setSource(QUrl.fromLocalFile(str(audio_path_to_play)))
        self.player.setPosition(int(start_s * 1000))

        if hasattr(self, '_play_segment_timer'):
            self._play_segment_timer.stop()
            self._play_segment_timer.deleteLater()

        self._play_segment_timer = QTimer(self)
        self._play_segment_timer.setSingleShot(True)
        self._play_segment_timer.timeout.connect(self.pause_audio) 

        duration_to_play_ms = int((end_s - start_s) * 1000)
        
        self._segment_end_s = end_s # Store segment end time for pause/resume logic
        if duration_to_play_ms > 50: # Min duration to play
            self._play_segment_timer.start(duration_to_play_ms)
            # Short delay before play() can help in some cases with QMediaPlayer state
            QTimer.singleShot(30, self.player.play) 
        else:
            self.log_status("Selected segment is too short to play or invalid.")
            self._segment_end_s = None 


    def pause_audio(self):
        if self.player.playbackState() == QMediaPlayer.PlaybackState.PlayingState:
            self.player.pause()
            if hasattr(self, '_play_segment_timer'): self._play_segment_timer.stop()
            self.log_status("Audio paused.")

    def stop_audio(self):
        if self.player.playbackState() != QMediaPlayer.PlaybackState.StoppedState:
            self.player.stop() 

        if hasattr(self, '_play_segment_timer'):
            self._play_segment_timer.stop()
        self._segment_end_s = None 

        active_waveform = None
        if self.player.source().isValid(): 
            source_file = Path(self.player.source().toLocalFile())
            if self.source_audio_full_path and source_file == self.source_audio_full_path:
                active_waveform = self.source_waveform_widget
            elif self.swapped_audio_full_path and source_file == self.swapped_audio_full_path:
                active_waveform = self.swapped_audio_waveform_widget
            elif hasattr(self, 'final_mix_preview_audio_path') and self.final_mix_preview_audio_path and source_file == self.final_mix_preview_audio_path:
                active_waveform = self.final_mix_waveform_widget
        
        if active_waveform:
             current_cursor_pos = active_waveform.playback_cursor_pos_sec
             # Reset to 0 only if no selection or if cursor is at the very end.
             # Otherwise, keep cursor at the stopped position for context.
             start_sel, end_sel = active_waveform.get_current_selection()
             if start_sel is None or (active_waveform.duration_sec > 0 and current_cursor_pos >= active_waveform.duration_sec - 0.05) :
                 active_waveform.set_playback_cursor(0)
             # Ensure cursor is visible
             active_waveform.show_playback_cursor = True 
             active_waveform.update()

        self.log_status("Audio stopped.")


    def select_wav_file(self):
        filepath, _ = QFileDialog.getOpenFileName(self, "Select Target Voice File", "", "Audio Files (*.wav *.mp3)")
        if filepath:
            p_filepath = Path(filepath)
            if p_filepath.suffix.lower() == ".mp3":
                self._handle_mp3_target_voice(p_filepath)
            else:
                self.target_voice_path = p_filepath
            self.wav_label.setText(self.target_voice_path.name)
            self.clean_target_voice_button.setEnabled(CLEARVOICE_AVAILABLE and True)
            self.log_status(f"Selected Target WAV: {filepath}")

    def _extract_and_load_audio(self):
        if not self.input_mp4_path: return
        if self.current_worker and self.current_worker.isRunning():
            QMessageBox.warning(self, "Busy", "A process (likely audio extraction) is already running.")
            return
        self.log_status("Starting audio extraction from video...")
        self.progress_bar.setRange(0, 100) 
        self.progress_bar.setValue(0)
        self.set_ui_enabled_for_processing(False)
        self.current_worker = WorkerThread(self._perform_audio_extraction_task, self.input_mp4_path)
        self.current_worker.progress.connect(self.handle_gui_progress_update) 
        self.current_worker.finished.connect(self._on_audio_extraction_finished)
        self.current_worker.log_message.connect(self._handle_worker_log_message)
        self.current_worker.start()

    def _handle_mp3_target_voice(self, mp3_path: Path):
        self.log_status(f"MP3 target voice selected: {mp3_path}. Converting to WAV...")
        try:
            audio = AudioSegment.from_mp3(mp3_path)
            wav_filename = f"{mp3_path.stem}_converted.wav"
            temp_wav_path = TMP_DIR / wav_filename
            audio.export(temp_wav_path, format="wav")
            self.target_voice_path = temp_wav_path
            self.log_status(f"MP3 converted to temporary WAV: {temp_wav_path}")
        except Exception as e:
            self.log_status(f"Error converting MP3 target voice: {e}", is_error=True)
            QMessageBox.critical(self, "MP3 Conversion Error", f"Could not convert MP3 {mp3_path.name} to WAV: {e}")
            self.target_voice_path = None

    def _perform_audio_extraction_task(self, input_mp4_path, progress_callback, log_callback):
        progress_callback(10, "Opening video file...")
        video_clip = VideoFileClip(str(input_mp4_path))
        fps = video_clip.fps # Get FPS
        progress_callback(30, "Extracting audio data...")
        extracted_audio_path = TMP_DIR / f"{input_mp4_path.stem}_extracted_audio.wav"
        video_clip.audio.write_audiofile(str(extracted_audio_path), logger=None) # Use logger=None for moviepy
        video_clip.close()
        progress_callback(60, "Audio extracted. Loading waveform data...")
        y, sr = librosa.load(str(extracted_audio_path), sr=None, mono=True)
        duration_sec = len(y) / sr
        progress_callback(80, "Normalizing audio...")
        y_norm = y / np.max(np.abs(y)) if np.max(np.abs(y)) > 0 else y
        MAX_WAVEFORM_POINTS = 150000 
        y_display = y_norm
        if len(y_norm) > MAX_WAVEFORM_POINTS:
            progress_callback(90, "Downsampling audio for waveform display...")
            downsample_factor = int(np.ceil(len(y_norm) / MAX_WAVEFORM_POINTS))
            y_display = y_norm[::downsample_factor]
        else:
            progress_callback(90, "Preparing audio for display...")
        progress_callback(100, "Audio data ready.")
        return {"audio_path": extracted_audio_path, "y_display": y_display, "sr": sr, "duration_sec": duration_sec, "fps": fps}


    def _on_audio_extraction_finished(self, result, message):
        self.set_ui_enabled_for_processing(True)
        if result:
            self.source_audio_full_path = result["audio_path"]
            self.source_audio_sr = result["sr"]
            self.source_audio_duration_sec = result["duration_sec"]
            self.source_video_fps = result.get("fps") # Store FPS
            self._update_frame_skip_from_seconds_input() # Update UI that depends on FPS

            self.source_waveform_widget.load_audio(result["y_display"], result["sr"], result["duration_sec"])
            self.log_status(f"Source audio loaded. Duration: {self.source_audio_duration_sec:.2f}s, SR: {self.source_audio_sr}Hz, Video FPS: {self.source_video_fps or 'N/A'}")
            self.clear_all_source_markers_action()
            is_outside = self.environment_combo.currentText() == "Outside"
            self.scene_cuts_sec = [] 
            self.use_scene_cuts_for_auto_rir_checkbox.setEnabled(False) 
            self.generate_auto_rir_markers_button.setEnabled(
                S2R_Generator is not None and self.swapped_audio_full_path is not None and not self.use_speech2rir_checkbox.isChecked() and not is_outside)
            self.speaker_list_combo.clear()
            self.diarization_results = None
            self.swapped_audio_waveform_widget.load_audio(np.array([0.0]), 1, 0) 
            self.swapped_audio_waveform_widget.hide()
            self.swapped_audio_full_path = None
            self.clear_all_rir_markers_action()
            self.final_mix_waveform_widget.load_audio(np.array([0.0]), 1, 0) # Clear final mix waveform
            self.final_mix_waveform_widget.hide()
            self.final_mix_preview_audio_path = None
            self.final_mix_segment_list.clear()
            self.final_mix_rir_settings.clear()
            self.estimated_rir_path = None 
            self._on_speech2rir_toggled(self.use_speech2rir_checkbox.isChecked())
        else:
            self.log_status(f"Audio extraction failed: {message}", is_error=True)
            QMessageBox.critical(self, "Audio Extraction Error", f"Could not extract or load audio: {message}")
            self.source_audio_full_path = None
            self.source_video_fps = None
            self.final_mix_preview_audio_path = None
            self.final_mix_waveform_widget.hide()
            self._update_frame_skip_from_seconds_input()
        self.current_worker = None 

    def seek_audio(self, time_sec, is_source=True):
        audio_path_to_seek = self.source_audio_full_path if is_source else self.swapped_audio_full_path
        waveform_widget_to_update = self.source_waveform_widget if is_source else self.swapped_audio_waveform_widget
        duration_to_check = self.source_audio_duration_sec if is_source else waveform_widget_to_update.duration_sec

    def seek_audio(self, time_sec, is_source=True, is_final_mix_preview=False):
        if is_final_mix_preview:
            audio_path_to_seek = getattr(self, 'final_mix_preview_audio_path', None)
            waveform_widget_to_update = self.final_mix_waveform_widget
            duration_to_check = waveform_widget_to_update.duration_sec
        else:
            audio_path_to_seek = self.source_audio_full_path if is_source else self.swapped_audio_full_path
            waveform_widget_to_update = self.source_waveform_widget if is_source else self.swapped_audio_waveform_widget
            duration_to_check = self.source_audio_duration_sec if is_source else waveform_widget_to_update.duration_sec

        if audio_path_to_seek and 0 <= time_sec <= duration_to_check:
            current_player_source_path = Path(self.player.source().toLocalFile()) if self.player.source().isValid() else None
            
            is_playing = self.player.playbackState() == QMediaPlayer.PlaybackState.PlayingState
            
            if current_player_source_path != audio_path_to_seek:
                self.player.stop() 
                self.player.setSource(QUrl.fromLocalFile(str(audio_path_to_seek)))
            
            self.player.setPosition(int(time_sec * 1000))
            waveform_widget_to_update.set_playback_cursor(time_sec)    
            
            if is_playing and current_player_source_path != audio_path_to_seek:
                # If it was playing a different source, restart play for the new source
                # Use a small delay if setSource was just called.
                if is_final_mix_preview:
                    QTimer.singleShot(50, lambda: self.play_audio_segment(is_final_mix_preview=True))
                else:
                    QTimer.singleShot(50, lambda: self.play_audio_segment(is_source=is_source))
            elif is_playing: # If it was playing the same source, it will continue from new position
                pass 
            # If it was paused or stopped, it remains so at the new position.


    def add_current_selection_as_marker(self):
        if not self.source_audio_full_path:
            QMessageBox.warning(self, "No Audio", "Please load an MP4 file first.")
            return
        start_s, end_s = self.source_waveform_widget.get_current_selection()
        if start_s is not None and end_s is not None: 
            self._add_source_marker_to_list(start_s, end_s, "Manual") 
            self.source_waveform_widget.clear_selection() 
        else:
            QMessageBox.information(self, "No Selection", "Please select a region on the waveform by dragging the mouse, then click this button.")

    def mark_full_source_audio_action(self):
        if not self.source_audio_full_path or self.source_audio_duration_sec <= 0:
            QMessageBox.warning(self, "No Audio", "Please load an MP4 file first.")
            return

        self.log_status("Clearing existing markers and marking full audio.")
        self.clear_all_source_markers_action() # Clear existing markers

        self._add_source_marker_to_list(0, self.source_audio_duration_sec, "Full Audio")
        self.source_waveform_widget.clear_selection() # Clear any visual selection
        self.log_status(f"Marked full audio: 0.00s - {self.source_audio_duration_sec:.2f}s")

    def _add_source_marker_to_list(self, start_s, end_s, source_text):
        for s, e, _ in self.source_time_markers:
            if abs(s - start_s) < 0.01 and abs(e - end_s) < 0.01:
                self.log_status(f"Marker {start_s:.2f}s - {end_s:.2f}s already exists.")
                return
        self.source_time_markers.append((start_s, end_s, source_text))
        self.source_time_markers.sort() 
        self._update_source_marker_display()
        self.log_status(f"Added marker: {start_s:.2f}s - {end_s:.2f}s ({source_text})")

    def _update_source_marker_display(self):
        self.source_marker_list_widget.clear()
        self.source_waveform_widget.clear_all_visual_markers() 
        for start_s, end_s, source_text in self.source_time_markers:
            # ... (rest of the method)
            item_text = f"{source_text}: {start_s:.2f}s - {end_s:.2f}s"
            self.source_marker_list_widget.addItem(QListWidgetItem(item_text))
            color = QColor(0, 255, 0, 90) if "Speaker" in source_text else QColor(255, 165, 0, 90) 
            self.source_waveform_widget.add_visual_marker(start_s, end_s, color=color)

    def remove_selected_source_markers_from_list(self):
        selected_items = self.source_marker_list_widget.selectedItems()
        if not selected_items:
            QMessageBox.information(self, "No Selection", "Please select markers from the list to remove.")
            return
        indices_to_remove = sorted([self.source_marker_list_widget.row(item) for item in selected_items], reverse=True)
        for index in indices_to_remove:
            removed_marker = self.source_time_markers.pop(index)
            self.log_status(f"Removed marker: {removed_marker[0]:.2f}s - {removed_marker[1]:.2f}s")
        self._update_source_marker_display()

    def clear_all_source_markers_action(self):
        self.source_time_markers.clear()
        self._update_source_marker_display()
        self.log_status("All time markers cleared.")

    def run_speaker_diarization_thread(self):
        if self.current_worker and self.current_worker.isRunning():
            QMessageBox.warning(self, "Busy", "A process is already running.")
            return
        if not self.source_audio_full_path:
            QMessageBox.warning(self, "No Audio", "Please load an MP4 and extract audio first.")
            return
        if not PYANNOTE_AVAILABLE:
            QMessageBox.critical(self, "Diarization Unavailable", "pyannote.audio library is not installed.")
            return
        self.hf_token = self.hf_token_edit.text().strip()
        if not self.hf_token or self.hf_token == HUGGING_FACE_TOKEN_PLACEHOLDER:
            QMessageBox.warning(self, "HuggingFace Token", "Please enter a valid HuggingFace User Access Token for speaker diarization.")
            self.hf_token_edit.setFocus()
            return
        self.log_status("Starting speaker diarization...")
        self.set_ui_enabled_for_processing(False) 
        self.progress_bar.setRange(0,0) 
        self.current_worker = WorkerThread(self._perform_diarization_task, self.source_audio_full_path, self.hf_token)
        self.current_worker.finished.connect(self._on_diarization_finished)
        self.current_worker.log_message.connect(self._handle_worker_log_message)
        self.current_worker.start()


    def run_rt60_estimation_thread(self):
        if self.current_worker and self.current_worker.isRunning():
            QMessageBox.warning(self, "Busy", "A process is already running.")
            return
        if not self.source_audio_full_path:
            QMessageBox.warning(self, "No Source Audio", "Please load an MP4 and extract audio first to estimate RT60.")
            return
        self.log_status("Starting RT60 estimation from source audio...")
        self.calculate_rt60_button.setEnabled(False)
        self.progress_bar.setRange(0,0) 
        self.current_worker = WorkerThread(self._perform_rt60_estimation_task, str(self.source_audio_full_path))
        self.current_worker.progress.connect(self.handle_gui_progress_update) 
        self.current_worker.finished.connect(self._on_rt60_estimation_finished)
        self.current_worker.log_message.connect(self._handle_worker_log_message)
        self.current_worker.start()

    def _perform_rt60_estimation_task(self, audio_path_str, progress_callback, log_callback=None, is_pre_segmented=False):
            # log_callback is optional for internal calls not needing full GUI logging
            def _log(msg, is_error=False):
                if log_callback: log_callback(msg, is_error)
                else: print(f"RT60 Task: {'ERROR: ' if is_error else ''}{msg}")

            rt60_value = None
            try:
                if progress_callback: progress_callback(5, "RT60: Loading audio...")
                full_audio, fs = librosa.load(audio_path_str, sr=None, mono=True)
                duration_sec = len(full_audio) / fs
                
                MIN_DURATION_FOR_RT60_SEC = 1.0 # Minimum duration for blind_rt60 to work reasonably
                if duration_sec < MIN_DURATION_FOR_RT60_SEC:
                    _log(f"RT60: Audio segment is too short ({duration_sec:.2f}s, min {MIN_DURATION_FOR_RT60_SEC}s). Estimation skipped.", is_error=True)
                    if progress_callback: progress_callback(100, "RT60: Skipped (too short).")
                    return None # Indicate failure or unreliable
                
                audio_to_process = full_audio

                if not is_pre_segmented: # Apply internal segmentation only if not already pre-segmented
                    MAX_DURATION_FOR_RT60_SEC = 60.0 
                    TARGET_SEGMENT_DURATION_SEC = 30.0 
                    
                    if duration_sec > MAX_DURATION_FOR_RT60_SEC:
                        _log(f"RT60: Audio is long ({duration_sec:.1f}s), selecting a segment...")
                        start_sample = int((duration_sec / 2 - TARGET_SEGMENT_DURATION_SEC / 2) * fs)
                        end_sample = int(start_sample + TARGET_SEGMENT_DURATION_SEC * fs)
                        start_sample = max(0, start_sample)
                        end_sample = min(len(full_audio), end_sample)
                        if end_sample > start_sample + fs * 5 : 
                            audio_to_process = full_audio[start_sample:end_sample]
                            _log(f"RT60: Using a {len(audio_to_process)/fs:.1f}s segment from the middle for estimation.")
                        else: 
                            audio_to_process = full_audio[:int(MAX_DURATION_FOR_RT60_SEC * fs)]
                            _log(f"RT60: Using the first {len(audio_to_process)/fs:.1f}s for estimation.")
                    else:
                         _log(f"RT60: Processing full audio ({duration_sec:.1f}s) for estimation (not pre-segmented).")
                else:
                     _log(f"RT60: Processing pre-segmented audio ({duration_sec:.1f}s) for estimation.")


                if progress_callback: progress_callback(20, "RT60: Initializing estimator...")
                estimator = BlindRT60(fs=fs) 
                if progress_callback: progress_callback(30, "RT60: Estimating RT60 from segment...")
                rt60_value = estimator.estimate(audio_to_process, fs) 

                if rt60_value is not None:
                    # Sanity check the estimated RT60
                    if rt60_value <= 0.01 or rt60_value > 10.0: # Example: RT60 > 10s or <=0.01s is likely an error
                        _log(f"RT60: Estimated value ({rt60_value:.3f}s) is outside reasonable bounds. Treating as failed.", True)
                        rt60_value = None
                
            except Exception as e:
                _log(f"RT60 estimation error: {e}", is_error=True)
                rt60_value = None
            
            if rt60_value is not None:
                if progress_callback: progress_callback(100, f"RT60: Estimation complete. Value: {rt60_value:.3f}s")
                else: _log(f"RT60: Estimation complete. Value: {rt60_value:.3f}s") # Log if no progress_callback
            else:
                if progress_callback: progress_callback(100, "RT60: Estimation failed or returned no value.")
                else: _log("RT60: Estimation failed or returned no value.") # Log if no progress_callback
            return rt60_value

    def _on_rt60_estimation_finished(self, result_rt60, message):
        self.calculate_rt60_button.setEnabled(True)
        self.progress_bar.setRange(0,100) 
        self.progress_bar.setValue(100 if result_rt60 is not None else 0)
        if result_rt60 is not None:
            self.rt60_display_edit.setText(f"{result_rt60:.3f}")
            self.log_status(f"RT60 estimation successful: {result_rt60:.3f}s")
            QMessageBox.information(self, "RT60 Estimation Complete", f"Estimated RT60 from source audio: {result_rt60:.3f}s")
        else:
            self.rt60_display_edit.setText("N/A") # Changed from "Error" to "N/A" for clarity
            self.log_status(f"RT60 estimation failed: {message}", is_error=True)
            QMessageBox.critical(self, "RT60 Estimation Error", f"Could not estimate RT60: {message}")
        self.current_worker = None

    def _perform_diarization_task(self, audio_path, hf_token, progress_callback, log_callback): 
        progress_callback(10, "Initializing diarization pipeline...")
        try:
            pipeline = PyannotePipeline.from_pretrained("pyannote/speaker-diarization-3.1", use_auth_token=hf_token)
            pipeline.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
            progress_callback(30, "Applying pipeline to audio...")
            diarization = pipeline(str(audio_path))
            rttm_file_path = TMP_DIR / "diarization_output.rttm"
            with open(rttm_file_path, "w") as rttm_f:
                diarization.write_rttm(rttm_f)
            progress_callback(70, "Parsing RTTM results...")
            all_speaker_segments = {} 
            for speech_turn, _, speaker_id in diarization.itertracks(yield_label=True):
                if speaker_id not in all_speaker_segments: all_speaker_segments[speaker_id] = []
                all_speaker_segments[speaker_id].append((speech_turn.start, speech_turn.end))
            progress_callback(100, "Diarization complete.")
            return all_speaker_segments 
        except Exception as e:
            log_callback(f"Diarization error: {e}", is_error=True) 
            raise Exception(f"Diarization failed: {e}")

    def _on_diarization_finished(self, result, message):
        self.progress_bar.setRange(0,100) 
        self.progress_bar.setValue(100 if result else 0)
        self.set_ui_enabled_for_processing(True) 
        if result:
            self.diarization_results = result 
            self.speaker_list_combo.clear()
            if self.diarization_results:
                for speaker_id in sorted(self.diarization_results.keys()):
                    num_segments = len(self.diarization_results[speaker_id])
                    self.speaker_list_combo.addItem(f"{speaker_id} ({num_segments} segments)", speaker_id)
                self.log_status(f"Speaker diarization successful. Found {len(self.diarization_results)} speakers.")
                QMessageBox.information(self, "Diarization Complete", f"Found {len(self.diarization_results)} speakers. Select a speaker and apply markers.")
            else:
                self.log_status("Diarization complete, but no speakers found.", is_error=True)
                QMessageBox.information(self, "Diarization Complete", "No speakers were detected in the audio.")
        else:
            self.log_status(f"Speaker diarization failed: {message}", is_error=True)
            QMessageBox.critical(self, "Diarization Error", f"Speaker diarization failed: {message}")
        self.current_worker = None

    def apply_selected_speaker_to_markers(self):
        if not self.diarization_results:
            QMessageBox.warning(self, "No Diarization Data", "Please run speaker detection first.")
            return
        selected_speaker_id = self.speaker_list_combo.currentData() 
        if not selected_speaker_id:
            QMessageBox.warning(self, "No Speaker Selected", "Please select a speaker from the dropdown list.")
            return
        speaker_segments = self.diarization_results.get(selected_speaker_id, [])
        if not speaker_segments:
            self.log_status(f"No segments found for speaker {selected_speaker_id}.")
            return
        for start_s, end_s in speaker_segments:
            self._add_source_marker_to_list(start_s, end_s, f"Speaker: {selected_speaker_id}")
        self.log_status(f"Applied {len(speaker_segments)} markers for {selected_speaker_id}.")

    def add_current_rir_selection_as_marker(self):
        if not self.swapped_audio_full_path:
            QMessageBox.warning(self, "No Swapped Audio", "Please perform 'Swap Voice' first.")
            return
        start_s, end_s = self.swapped_audio_waveform_widget.get_current_selection()
        if start_s is not None and end_s is not None: 
            try:
                current_room_dim = [float(d.strip()) for d in self.room_dim_edit.text().split(',')]
                if len(current_room_dim) != 3 or not all(d > 0 for d in current_room_dim):
                    raise ValueError("Room dimensions must be three positive comma-separated numbers.")
            except ValueError as e:
                QMessageBox.critical(self, "Invalid Default Room Dimensions", f"Cannot add RIR marker. Fix default room dimensions in Global Configuration: {e}")
                return
            target_rt60_val = None
            if self.use_estimated_rt60_checkbox.isChecked():
                try:
                    estimated_rt60_str = self.rt60_display_edit.text()
                    if estimated_rt60_str and estimated_rt60_str != "N/A" and estimated_rt60_str != "Error":
                         target_rt60_val = float(estimated_rt60_str)
                except ValueError: pass 
                if target_rt60_val is None:
                     QMessageBox.warning(self, "Estimated RT60 Unavailable", "Could not use estimated RT60. Please calculate it first or uncheck 'Use Estimated Source RT60'.")
                     return
            else:
                try: target_rt60_val = float(self.rir_target_rt60_edit.text())
                except ValueError:
                    QMessageBox.warning(self, "Invalid RT60", "Please enter a valid number for Target RT60 for the RIR marker.")
                    self.rir_target_rt60_edit.setFocus(); return
            rir_params = {'rt60': target_rt60_val, 'room_dim': current_room_dim}
            self._add_rir_marker_to_list(start_s, end_s, rir_params)
            self.swapped_audio_waveform_widget.clear_selection() 
        else:
            QMessageBox.information(self, "No Selection", "Please select a region on the swapped audio waveform by dragging the mouse, then click this button.")


    def _add_rir_marker_to_list(self, start_s, end_s, rir_params):
        self.swapped_audio_markers.append((start_s, end_s, rir_params))
        self.swapped_audio_markers.sort()
        self._update_rir_marker_display()
        self._update_final_mix_segment_list() # Update the list in the third row
        self.log_status(f"Added RIR marker: {start_s:.2f}s - {end_s:.2f}s, Params: {rir_params}")

    def _update_rir_marker_display(self):
        self.rir_marker_list_widget.clear()
        self.swapped_audio_waveform_widget.clear_all_visual_markers()
        for start_s, end_s, rir_params in self.swapped_audio_markers:
            if 'rir_path' in rir_params:
                rt60_val_str = f", Estimated RT60: {rir_params['estimated_rt60']:.3f}s" if rir_params.get('estimated_rt60') is not None else ""
                item_text = f"{start_s:.2f}s - {end_s:.2f}s (Auto Speech2RIR: {Path(rir_params['rir_path']).name}{rt60_val_str})"
                self.swapped_audio_waveform_widget.add_visual_marker(start_s, end_s, color=QColor(0, 200, 200, 90)) 
            else: 
                item_text = f"{start_s:.2f}s - {end_s:.2f}s (Sim: RT60 {rir_params.get('rt60', 'N/A')}, Dim: {rir_params.get('room_dim', 'N/A')})"
                self.swapped_audio_waveform_widget.add_visual_marker(start_s, end_s, color=QColor(100, 100, 255, 90)) 
            self.rir_marker_list_widget.addItem(QListWidgetItem(item_text))

    def remove_selected_rir_markers_from_list(self):
        selected_items = self.rir_marker_list_widget.selectedItems()
        if not selected_items:
            QMessageBox.information(self, "No Selection", "Please select RIR markers from the list to remove.")
            return
        indices_to_remove = sorted([self.rir_marker_list_widget.row(item) for item in selected_items], reverse=True)
        for index in indices_to_remove:
            removed_marker = self.swapped_audio_markers.pop(index)
            self.log_status(f"Removed RIR marker: {removed_marker[0]:.2f}s - {removed_marker[1]:.2f}s")
        self._update_rir_marker_display()
        self._update_final_mix_segment_list()

    def clear_all_rir_markers_action(self):
        self.swapped_audio_markers.clear()
        self._update_rir_marker_display()
        self._update_final_mix_segment_list()
        self.log_status("All RIR markers cleared.")

    def _update_final_mix_segment_list(self):
        self.final_mix_segment_list.clear()
        self.final_mix_rir_settings = [] # Reset adjustments when base markers change
        for i, (start_s, end_s, rir_params) in enumerate(self.swapped_audio_markers):
            # Create a copy of rir_params for final_mix_rir_settings
            self.final_mix_rir_settings.append({'start_s': start_s, 'end_s': end_s, 'params': dict(rir_params)})
            item_text = f"Segment {i+1}: {start_s:.2f}s - {end_s:.2f}s"
            if 'rir_path' in rir_params:
                est_rt60_str = ""
                if rir_params.get('estimated_rt60') is not None:
                    est_rt60_str = f", Est. RT60: {rir_params['estimated_rt60']:.3f}s"
                item_text += f" (S2R: {Path(rir_params['rir_path']).name}{est_rt60_str})"
            elif 'rt60' in rir_params:
                item_text += f" (Sim: RT60 {rir_params.get('rt60', 'N/A')})"
            else:
                item_text += " (No RIR)"
            
            # Add default volume for display, actual params are in self.final_mix_rir_settings
            # This part will be refined when _on_full_rir_preview_finished is updated
            # item_text += f", Vol: {rir_params.get('volume_db', 0.0):.1f}dB"
            self.final_mix_segment_list.addItem(QListWidgetItem(item_text))

    def initiate_voice_swap_pipeline(self):
        if self.current_worker and self.current_worker.isRunning():
            QMessageBox.warning(self, "Busy", "A process is already running. Please wait.")
            return
        if not self.input_mp4_path:
            QMessageBox.warning(self, "Missing Input", "Please select an input MP4 file.")
            return
        if not self.target_voice_path:
            QMessageBox.warning(self, "Missing Input", "Please select a target voice WAV file.")
            return
        if not self.source_time_markers:
            QMessageBox.warning(self, "No Markers", "Please add time markers to indicate which audio segments to process.")
            return

        vc_params = {
            'diffusion_steps': self.diffusion_steps_slider.value(),
            'length_adjust': self.length_adjust_slider.value() / 100.0,
            'inference_cfg_rate': self.cfg_rate_slider.value() / 100.0,
            'f0_condition': self.f0_checkbox.isChecked(),
            'auto_f0_adjust': self.auto_f0_checkbox.isChecked(),
            'pitch_shift': self.pitch_shift_slider.value()
        }

        self.log_status("Starting voice swap process...")
        self.progress_bar.setValue(0)
        self.set_ui_enabled_for_processing(False)
        self.output_video_label.setText("Swapping voices...") 

        self.current_worker = WorkerThread(self._perform_voice_swap_task,
                                           str(self.source_audio_full_path), 
                                           self.source_audio_sr,
                                           str(self.target_voice_path),
                                           vc_params,
                                           list(self.source_time_markers) 
                                          )
        self.current_worker.progress.connect(self.handle_gui_progress_update)
        self.current_worker.finished.connect(self.on_voice_swap_pipeline_finished)
        self.current_worker.log_message.connect(self._handle_worker_log_message)
        self.current_worker.start()

    def on_voice_swap_pipeline_finished(self, result_swapped_audio_path, message):
        self.set_ui_enabled_for_processing(True)
        if result_swapped_audio_path:
            self.swapped_audio_full_path = Path(result_swapped_audio_path)
            self.log_status(f"Voice swap successful! Swapped audio: {self.swapped_audio_full_path}")
            y_swap, sr_swap = librosa.load(self.swapped_audio_full_path, sr=None, mono=True)
            dur_swap = len(y_swap) / sr_swap
            y_swap_norm = y_swap / np.max(np.abs(y_swap)) if np.max(np.abs(y_swap)) > 0 else y_swap
            MAX_WAVEFORM_POINTS = 150000
            y_display_swap = y_swap_norm
            if len(y_swap_norm) > MAX_WAVEFORM_POINTS:
                downsample_factor = int(np.ceil(len(y_swap_norm) / MAX_WAVEFORM_POINTS))
                y_display_swap = y_swap_norm[::downsample_factor]
            self.swapped_audio_waveform_widget.load_audio(y_display_swap, sr_swap, dur_swap)
            self.swapped_audio_waveform_widget.show()
            is_outside = self.environment_combo.currentText() == "Outside"
            self.generate_auto_rir_markers_button.setEnabled(
                S2R_Generator is not None and self.swapped_audio_full_path is not None and not self.use_speech2rir_checkbox.isChecked() and not is_outside
            )
            self.clear_all_rir_markers_action() 
            self._update_final_mix_segment_list() # Populate final mix segment list
            self.final_mix_waveform_widget.hide() # Hide final mix until a preview is generated
            QMessageBox.information(self, "Voice Swap Complete", f"Voice swapping finished.\nSwapped audio ready for RIR marking: {self.swapped_audio_full_path}")
            self.output_video_label.setText(f"Swapped audio ready: {self.swapped_audio_full_path.name}")
            self.step_tab_widget.setCurrentIndex(1) # Switch to "2. Apply RIR Values" tab
        else:
            self.log_status(f"Voice swap failed: {message}", is_error=True)
            self.output_video_label.setText("Voice swap failed.")
            QMessageBox.critical(self, "Voice Swap Error", f"An error occurred during voice swapping:\n{message}")
        self.current_worker = None
        self.unload_vc_models() 

    def initiate_final_processing_pipeline(self):
        if self.current_worker and self.current_worker.isRunning():
            QMessageBox.warning(self, "Busy", "A process is already running.")
            return
        if not self.input_mp4_path: 
            QMessageBox.warning(self, "Missing Input", "Original MP4 not found.")
            return
        if not self.swapped_audio_full_path:
            QMessageBox.warning(self, "Missing Swapped Audio", "Please perform 'Swap Voice' first.")
            return
        if not self.output_dir_path:
            QMessageBox.warning(self, "Missing Output Directory", "Please select an output directory first using the 'Select Output Directory' button.")
            self.output_dir_button.setFocus() 
            return
        use_s2r_mode = self.use_speech2rir_checkbox.isChecked()
        s2r_wet_mix_percentage = self.s2r_mix_slider.value() 
        if use_s2r_mode and (not self.estimated_rir_path or not self.estimated_rir_path.exists()):
            QMessageBox.warning(self, "S2R RIR Missing", "Speech2RIR mode is selected, but no estimated RIR found. Please ensure RIR estimation was run (e.g., by toggling the checkbox after loading source audio, or if it failed).")
            return
        if use_s2r_mode and S2R_Generator is None:
             QMessageBox.critical(self, "S2R Error", "Speech2RIR model components are not loaded. Cannot proceed in S2R mode.")
             return
        try:
            room_dim_list = [float(d.strip()) for d in self.room_dim_edit.text().split(',')]
            if len(room_dim_list) != 3 or not all(d > 0 for d in room_dim_list):
                raise ValueError("Room dimensions must be three positive comma-separated numbers.")
        except ValueError as e:
            QMessageBox.critical(self, "Invalid Default Room Dimensions", str(e))
            return

        # Determine which RIR settings to use
        final_rir_config_to_use = []
        if self.final_mix_rir_settings and any(item['params'] for item in self.final_mix_rir_settings): # Check if any adjustments were made
            # Ensure the structure is a list of dicts, as expected by the processing loop
            final_rir_config_to_use = [dict(item) for item in self.final_mix_rir_settings]
            self.log_status("Using effect settings from the Final Mix Adjustment section for final processing.")
        else: # Should not happen if preview generation populates final_mix_rir_settings
            # Convert from self.swapped_audio_markers (list of tuples) to list of dicts
            final_rir_config_to_use = [{'start_s': s_time, 'end_s': e_time, 'params': dict(r_params)} for s_time, e_time, r_params in self.swapped_audio_markers]
            self.log_status("Warning: Using RIR settings from Row 2 for final processing (Row 3 settings empty).", is_error=True)

        self.log_status("Starting final processing (RIR application and video merge)...")
        self.progress_bar.setValue(0)
        self.set_ui_enabled_for_processing(False)
        self.output_video_label.setText("Finalizing video...")

        enable_demucs = self.enable_demucs_mix_checkbox.isChecked() and self.demucs_mix_group.isChecked()
        demucs_ratio = self.demucs_bg_mix_slider.value() / 100.0

        self.current_worker = WorkerThread(self._perform_final_mixdown_task, 
                                           str(self.input_mp4_path), 
                                           str(self.swapped_audio_full_path), 
                                           self.output_dir_path, 
                                           final_rir_config_to_use, # This is now a list of dicts
                                           use_s2r_mode=use_s2r_mode, 
                                           s2r_wet_mix_percentage=s2r_wet_mix_percentage,
                                           estimated_s2r_rir_path_global=self.estimated_rir_path,
                                           enable_demucs_mix=enable_demucs,
                                           demucs_bg_volume_ratio=demucs_ratio
                                           )
        self.current_worker.progress.connect(self.handle_gui_progress_update)
        self.current_worker.finished.connect(self.on_final_processing_pipeline_finished)
        self.current_worker.log_message.connect(self._handle_worker_log_message)
        self.current_worker.start()

    def handle_gui_progress_update(self, value, description):
        self.progress_bar.setValue(value)
        self.log_status(f"{description} ({value}%)")

    def on_final_processing_pipeline_finished(self, result_path, message):
        self.set_ui_enabled_for_processing(True)
        self.progress_bar.setValue(100 if result_path else 0)
        if result_path:
            self.log_status(f"Processing successful! Output: {result_path}")
            self.output_video_label.setText(str(result_path))
            QMessageBox.information(self, "Processing Complete", f"Video processing finished successfully.\nOutput: {result_path}")
        else:
            self.log_status(f"Final processing failed: {message}", is_error=True)
            self.output_video_label.setText("Processing failed.")
            QMessageBox.critical(self, "Processing Error", f"An error occurred during processing:\n{message}")
        self.current_worker = None

    def clear_cuda_cache_and_gc(self):
        self.log_status("Clearing CUDA cache and running garbage collection...")
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        gc.collect()
        self.log_status("CUDA cache cleared and GC run.")

    def unload_vc_models(self):
        global VOICE_MODELS_CACHE
        if VOICE_MODELS_CACHE is not None:
            self.log_status("Unloading Voice Conversion models from memory...")
            for model_dict_key in ['model', 'model_f0']:
                if model_dict_key in VOICE_MODELS_CACHE and VOICE_MODELS_CACHE[model_dict_key] is not None:
                    actual_model_dict = VOICE_MODELS_CACHE[model_dict_key]
                    if isinstance(actual_model_dict, dict):
                        for key_inner in list(actual_model_dict.keys()): 
                            module = actual_model_dict[key_inner]
                            if hasattr(module, 'to') and callable(module.to) and hasattr(module, 'device') and module.device.type == 'cuda':
                                try: module.to('cpu')
                                except Exception as e: self.log_status(f"Could not move module {model_dict_key}.{key_inner} to CPU: {e}", is_error=True)
                            del actual_model_dict[key_inner] 
            direct_model_keys = ['campplus_model', 'bigvgan_model', 'bigvgan_44k_model', 'whisper_model', 'rmvpe']
            for key in direct_model_keys:
                if key in VOICE_MODELS_CACHE and VOICE_MODELS_CACHE[key] is not None:
                    module = VOICE_MODELS_CACHE[key]
                    if hasattr(module, 'to') and callable(module.to) and hasattr(module, 'device') and module.device.type == 'cuda':
                        try: module.to('cpu')
                        except Exception as e: self.log_status(f"Could not move model {key} to CPU: {e}", is_error=True)
            VOICE_MODELS_CACHE.clear(); VOICE_MODELS_CACHE = None
            self.clear_cuda_cache_and_gc()
            self.log_status("Voice Conversion models unloaded.")
        else:
            self.log_status("Voice Conversion models already unloaded or not loaded.", is_error=False)

    def load_clearvoice_model(self, progress_callback, log_callback): # Added log_callback
        global CLEARVOICE_MODEL_CACHE
        if CLEARVOICE_MODEL_CACHE is not None:
            progress_callback(100, "ClearVoice SE model already loaded.")
            return True
        if not CLEARVOICE_AVAILABLE:
            log_callback("ClearVoice library not available. Cannot load SE model.", True)
            progress_callback(100, "ClearVoice Load Failed (Not Available).")
            return False

        progress_callback(10, f"Loading ClearVoice SE model: {self.clearvoice_se_model_name}...")
        try:
            cv_instance = ClearVoice(task='speech_enhancement', model_names=[self.clearvoice_se_model_name])
            CLEARVOICE_MODEL_CACHE = {'model': cv_instance, 'model_name': self.clearvoice_se_model_name}
            progress_callback(100, "ClearVoice SE model loaded successfully.")
            log_callback(f"ClearVoice SE model '{self.clearvoice_se_model_name}' loaded.")
            return True
        except Exception as e:
            import traceback
            log_callback(f"Error loading ClearVoice SE model: {e}\n{traceback.format_exc()}", True)
            CLEARVOICE_MODEL_CACHE = None
            progress_callback(100, "ClearVoice SE Load Error.")
            return False

    def unload_clearvoice_model(self):
        global CLEARVOICE_MODEL_CACHE
        if CLEARVOICE_MODEL_CACHE is not None:
            self.log_status(f"Unloading ClearVoice SE model '{CLEARVOICE_MODEL_CACHE.get('model_name', '')}' from memory...")
            try:
                del CLEARVOICE_MODEL_CACHE['model'] 
            except KeyError:
                pass 
            except Exception as e:
                self.log_status(f"Error during ClearVoice model deletion: {e}", is_error=True)

            CLEARVOICE_MODEL_CACHE.clear()
            CLEARVOICE_MODEL_CACHE = None
            self.clear_cuda_cache_and_gc()
            self.log_status("ClearVoice SE model unloaded.")
        else:
            self.log_status("ClearVoice SE model already unloaded or not loaded.", is_error=False)

    def _perform_voice_swap_task(self, 
                                 source_audio_full_path_str, 
                                 source_audio_sr,
                                 target_voice_path_str, vc_params,
                                 source_time_segments_for_vc, progress_callback, log_callback):
        global VOICE_MODELS_CACHE
        global CLEARVOICE_MODEL_CACHE
        
        progress_callback(0, "Initializing voice swap (model loading if needed)...")
        f0_selected = vc_params['f0_condition']
        
        if VOICE_MODELS_CACHE is None or VOICE_MODELS_CACHE.get('f0_condition_state') != f0_selected:
            if VOICE_MODELS_CACHE is not None:
                log_callback("F0 condition changed, unloading previous VC models...")
                # Call unload directly as it uses self.log_status which is GUI thread safe
                QMetaObject.invokeMethod(self, "unload_vc_models", Qt.QueuedConnection)


            log_callback("Loading new VC models...")
            VOICE_MODELS_CACHE = setup_models_vp2(f0_selected, 
                                                  progress_callback=lambda v, d: progress_callback(int(v * 0.10), f"Model Load: {d}"))
        else:
            progress_callback(10, "VC Models already loaded.") 

        if torch.cuda.is_available(): torch.cuda.empty_cache()
        gc.collect()

        sr_vc_models = VOICE_MODELS_CACHE['sr'] 

        reduce_echo_enabled = self.reduce_source_echo_checkbox.isChecked() # Access GUI element here (MainThread)
        rt60_threshold_text = self.source_echo_rt60_threshold_edit.text() # Access GUI element here (MainThread)
        rt60_threshold_for_se = 0.0

        if reduce_echo_enabled:
            try:
                rt60_threshold_for_se = float(rt60_threshold_text)
                # Pass log_callback to load_clearvoice_model
                if not self.load_clearvoice_model(lambda v, d: progress_callback(12 + int(v * 0.03), f"SE Model Load: {d}"), log_callback) : 
                    log_callback("Failed to load ClearVoice SE model. Echo reduction will be skipped.", True)
                    reduce_echo_enabled = False 
            except ValueError:
                log_callback(f"Invalid RT60 threshold: {rt60_threshold_text}. Echo reduction skipped.", True)
                reduce_echo_enabled = False

        current_sr_vc = VOICE_MODELS_CACHE['sr']
        total_user_segments = len(source_time_segments_for_vc)
        processed_vc_segment_audio_paths = []
        
        progress_after_model_loads = 10
        if reduce_echo_enabled and CLEARVOICE_MODEL_CACHE is not None: 
            progress_after_model_loads = 15

        segment_processing_total_span_pc = 85.0 - progress_after_model_loads

        for i, (user_segment_start_s, user_segment_end_s, _) in enumerate(source_time_segments_for_vc):
            true_user_segment_duration_s = user_segment_end_s - user_segment_start_s
            if true_user_segment_duration_s <= 0.01: 
                processed_vc_segment_audio_paths.append(None)
                base_progress_empty_seg = progress_after_model_loads + (i / total_user_segments) * segment_processing_total_span_pc
                span_empty_seg = (1.0 / total_user_segments) * segment_processing_total_span_pc
                progress_callback(int(base_progress_empty_seg + span_empty_seg), f"User Segment {i+1}: Was empty or too short.")
                continue
            original_total_user_segment_duration_samples_at_vc_sr = int(true_user_segment_duration_s * sr_vc_models)

            max_sub_chunk_len_ms = 25000  
            num_sub_chunks = int(np.ceil(true_user_segment_duration_s * 1000 / max_sub_chunk_len_ms))
            if num_sub_chunks == 0 and true_user_segment_duration_s > 0.01: num_sub_chunks = 1

            processed_sub_chunks_data_for_this_user_segment = []
            
            base_progress_for_this_user_segment_group = progress_after_model_loads + (i / total_user_segments) * segment_processing_total_span_pc
            total_span_for_this_user_segment_group = (1.0 / total_user_segments) * segment_processing_total_span_pc

            progress_callback(int(base_progress_for_this_user_segment_group),
                              f"User Segment {i+1}/{total_user_segments} ({user_segment_start_s:.2f}s - {user_segment_end_s:.2f}s): Preparing...")

            for k in range(num_sub_chunks):
                current_sub_chunk_original_start_s_in_source = user_segment_start_s + k * (max_sub_chunk_len_ms / 1000.0)
                current_sub_chunk_duration_s = min(max_sub_chunk_len_ms / 1000.0, 
                                                   (user_segment_start_s + true_user_segment_duration_s) - current_sub_chunk_original_start_s_in_source)

                if current_sub_chunk_duration_s <= 0.01: continue

                base_progress_for_sub_chunk_processing = base_progress_for_this_user_segment_group + \
                                                         (k / num_sub_chunks) * total_span_for_this_user_segment_group
                total_span_for_this_sub_chunk = (1.0 / num_sub_chunks) * total_span_for_this_user_segment_group

                span_for_sub_chunk_se_rt60 = total_span_for_this_sub_chunk * 0.05 
                span_for_sub_chunk_se_proc = total_span_for_this_sub_chunk * 0.15 
                span_for_sub_chunk_vc_op = total_span_for_this_sub_chunk * 0.80   
                
                current_progress_within_sub_chunk = base_progress_for_sub_chunk_processing

                temp_sub_chunk_raw_path = TMP_DIR / f"user_segment_{i}_subchunk_{k}_raw_extract.wav"
                try:
                    y_sub_chunk_raw, sr_sub_chunk_raw = librosa.load(
                        source_audio_full_path_str, 
                        sr=None, 
                        offset=current_sub_chunk_original_start_s_in_source, 
                        duration=current_sub_chunk_duration_s
                    )
                    sf.write(str(temp_sub_chunk_raw_path), y_sub_chunk_raw, sr_sub_chunk_raw)
                except Exception as e_extract_sub:
                    log_callback(f"Error extracting sub-chunk {k+1} for user segment {i+1}: {e_extract_sub}", True)
                    silent_sub_data = np.zeros(int(current_sub_chunk_duration_s * sr_vc_models), dtype=np.float32)
                    processed_sub_chunks_data_for_this_user_segment.append(silent_sub_data)
                    progress_callback(int(current_progress_within_sub_chunk + total_span_for_this_sub_chunk), f"SubChunk {k+1} Error, using silence.")
                    if temp_sub_chunk_raw_path.exists(): temp_sub_chunk_raw_path.unlink(missing_ok=True)
                    continue 
                
                path_for_next_step = temp_sub_chunk_raw_path

                if reduce_echo_enabled and CLEARVOICE_MODEL_CACHE is not None:
                    try:
                        rt60_est_progress_val_sub = 0
                        def rt60_sub_progress_adapter(val, desc):
                            nonlocal rt60_est_progress_val_sub
                            rt60_est_progress_val_sub = val
                            actual_gui_progress = int(current_progress_within_sub_chunk + (val/100.0) * span_for_sub_chunk_se_rt60)
                            progress_callback(actual_gui_progress, f"UserSeg {i+1} Sub {k+1} RT60Est: {desc}")
                        
                        # Pass log_callback to _perform_rt60_estimation_task
                        estimated_sub_chunk_rt60 = self._perform_rt60_estimation_task(str(path_for_next_step), progress_callback=rt60_sub_progress_adapter, log_callback=log_callback)
                        current_progress_within_sub_chunk += span_for_sub_chunk_se_rt60

                        if estimated_sub_chunk_rt60 > rt60_threshold_for_se:
                            log_callback(f"Sub-chunk {k+1} (Seg {i+1}) RT60 ({estimated_sub_chunk_rt60:.2f}s) > threshold. Applying SE.")
                            cleaned_sub_chunk_path_temp = TMP_DIR / f"user_segment_{i}_subchunk_{k}_cleaned.wav"
                            
                            def se_sub_progress_adapter(val, desc):
                                actual_gui_progress = int(current_progress_within_sub_chunk + (val/100.0) * span_for_sub_chunk_se_proc)
                                progress_callback(actual_gui_progress, f"UserSeg {i+1} Sub {k+1} SE: {desc}")

                            # Pass log_callback to _apply_speech_enhancement_to_segment
                            self._apply_speech_enhancement_to_segment(str(path_for_next_step), str(cleaned_sub_chunk_path_temp), progress_callback=se_sub_progress_adapter, log_callback=log_callback)
                            if path_for_next_step.exists() and path_for_next_step != cleaned_sub_chunk_path_temp : path_for_next_step.unlink(missing_ok=True) 
                            path_for_next_step = cleaned_sub_chunk_path_temp
                        current_progress_within_sub_chunk += span_for_sub_chunk_se_proc
                    except Exception as e_se_sub:
                        log_callback(f"Error during SE for sub-chunk {k+1} (Seg {i+1}): {e_se_sub}", True)
                        current_progress_within_sub_chunk += span_for_sub_chunk_se_rt60 * (rt60_est_progress_val_sub < 100)
                        current_progress_within_sub_chunk += span_for_sub_chunk_se_proc 
                    
                def sub_chunk_vc_progress_adapter(sub_val_vc, sub_desc_vc):
                    actual_progress_for_gui = int(current_progress_within_sub_chunk + (sub_val_vc / 100.0) * span_for_sub_chunk_vc_op)
                    progress_callback(actual_progress_for_gui, f"UserSeg {i+1}, SubChunk {k+1}/{num_sub_chunks} - {sub_desc_vc}")

                try:
                    vc_sr_output, vc_audio_data_for_sub_chunk = voice_conversion_vp2(
                        source_path=str(path_for_next_step), 
                        target_path=target_voice_path_str, 
                        models=VOICE_MODELS_CACHE,
                        progress_callback=sub_chunk_vc_progress_adapter,
                        **vc_params
                    )
                except Exception as e_vc_sub:
                    log_callback(f"Error during VC for sub-chunk {k+1} (Seg {i+1}): {e_vc_sub}", True)
                    y_fallback, sr_fallback = librosa.load(str(path_for_next_step), sr=None)
                    if sr_fallback != sr_vc_models: y_fallback = librosa.resample(y_fallback, orig_sr=sr_fallback, target_sr=sr_vc_models)
                    vc_audio_data_for_sub_chunk = y_fallback
                    vc_sr_output = sr_vc_models
                
                processed_sub_chunks_data_for_this_user_segment.append(vc_audio_data_for_sub_chunk)
                
                if path_for_next_step.exists(): path_for_next_step.unlink(missing_ok=True)
                if temp_sub_chunk_raw_path.exists() and temp_sub_chunk_raw_path != path_for_next_step:
                     temp_sub_chunk_raw_path.unlink(missing_ok=True)

            if processed_sub_chunks_data_for_this_user_segment:
                full_processed_user_segment_data_np = np.concatenate(processed_sub_chunks_data_for_this_user_segment)
            else: 
                log_callback(f"Warning: No sub-chunks processed for User Segment {i+1}. Using silence.", True)
                full_processed_user_segment_data_np = np.zeros(original_total_user_segment_duration_samples_at_vc_sr, dtype=np.float32)

            current_len_samples = len(full_processed_user_segment_data_np)
            if current_len_samples < original_total_user_segment_duration_samples_at_vc_sr:
                padding = np.zeros(original_total_user_segment_duration_samples_at_vc_sr - current_len_samples, dtype=full_processed_user_segment_data_np.dtype)
                full_processed_user_segment_data_np = np.concatenate((full_processed_user_segment_data_np, padding))
            elif current_len_samples > original_total_user_segment_duration_samples_at_vc_sr:
                full_processed_user_segment_data_np = full_processed_user_segment_data_np[:original_total_user_segment_duration_samples_at_vc_sr]

            final_path_for_this_user_segment = TMP_DIR / f"segment_{i}_vc_processed_final.wav"
            sf.write(str(final_path_for_this_user_segment), full_processed_user_segment_data_np, sr_vc_models)
            processed_vc_segment_audio_paths.append(final_path_for_this_user_segment)
            
            progress_callback(int(base_progress_for_this_user_segment_group + total_span_for_this_user_segment_group),
                            f"User Segment {i+1} voice conversion complete.")

        progress_callback(85, "All segments voice-converted. Merging into full swapped audio track...")
    
        output_segments_pydub = []
        last_original_audio_end_time_ms = 0
    
        for i, (start_s, end_s, _) in enumerate(source_time_segments_for_vc): 
            start_ms = int(start_s * 1000)
            end_ms = int(end_s * 1000)
   
            if start_ms > last_original_audio_end_time_ms:
                y_interlude, sr_interlude = librosa.load(source_audio_full_path_str, sr=source_audio_sr, offset=last_original_audio_end_time_ms/1000.0, duration=(start_ms - last_original_audio_end_time_ms)/1000.0)
                interlude_path_temp = TMP_DIR / f"interlude_before_seg_{i}.wav"; sf.write(str(interlude_path_temp), y_interlude, sr_interlude)
                original_interlude_pydub = AudioSegment.from_wav(interlude_path_temp); interlude_path_temp.unlink(missing_ok=True)
                output_segments_pydub.append(original_interlude_pydub)
    
            processed_segment_path_for_merge = processed_vc_segment_audio_paths[i]
            if processed_segment_path_for_merge is not None and Path(processed_segment_path_for_merge).exists():
                try:
                    swapped_segment_pydub = AudioSegment.from_wav(processed_segment_path_for_merge)
                    output_segments_pydub.append(swapped_segment_pydub)
                except Exception as e_load_swap:
                    log_callback(f"Error loading processed segment {processed_segment_path_for_merge} for merge: {e_load_swap}. Skipping.", True)
            elif processed_segment_path_for_merge is None: 
                pass 
            else: 
                 log_callback(f"Warning: Processed segment path {processed_segment_path_for_merge} for user segment {i+1} not found or invalid. Skipping in merge.", True)

            last_original_audio_end_time_ms = end_ms

        source_full_duration_ms = int(librosa.get_duration(path=source_audio_full_path_str) * 1000)        
        if last_original_audio_end_time_ms < source_full_duration_ms:
            y_tail, sr_tail = librosa.load(source_audio_full_path_str, sr=source_audio_sr, offset=last_original_audio_end_time_ms/1000.0)
            tail_path_temp = TMP_DIR / "final_tail.wav"; sf.write(str(tail_path_temp), y_tail, sr_tail)
            remaining_original_tail_pydub = AudioSegment.from_wav(tail_path_temp); tail_path_temp.unlink(missing_ok=True)
            output_segments_pydub.append(remaining_original_tail_pydub)
    
        final_swapped_audio_file_path = TMP_DIR / f"{Path(source_audio_full_path_str).stem}_swapped.wav"
        if not output_segments_pydub:
            log_callback("Warning: No voice-swapped segments were produced or no segments marked. Resulting audio is the original source audio.", True)
            shutil.copy(source_audio_full_path_str, final_swapped_audio_file_path)
            progress_callback(100, "Voice swap resulted in original audio (no segments processed).")
            # Unload ClearVoice model if it was loaded
            if reduce_echo_enabled: QMetaObject.invokeMethod(self, "unload_clearvoice_model", Qt.QueuedConnection)
            return str(final_swapped_audio_file_path)
    
        final_swapped_audio_pydub = sum(output_segments_pydub, AudioSegment.empty()) 
        final_swapped_audio_pydub = final_swapped_audio_pydub.set_frame_rate(source_audio_sr) 
        final_swapped_audio_pydub.export(str(final_swapped_audio_file_path), format="wav")
        
        progress_callback(100, "Voice swap merging complete.")
        
        if reduce_echo_enabled: QMetaObject.invokeMethod(self, "unload_clearvoice_model", Qt.QueuedConnection)
        return str(final_swapped_audio_file_path)


    def _perform_final_mixdown_task(self, original_mp4_path_str, swapped_audio_path_str,
                                    output_directory_path_obj, rir_markers,
                                    use_s2r_mode, s2r_wet_mix_percentage,
                                    estimated_s2r_rir_path_global,
                                    enable_demucs_mix: bool, demucs_bg_volume_ratio: float, # New params
                                    progress_callback, log_callback):

        original_mp4_path = Path(original_mp4_path_str)
        swapped_audio_path = Path(swapped_audio_path_str) 
        audio_to_mix_with_video = swapped_audio_path # Initially, this is the audio to be muxed

        progress_callback(0, "Initializing final mixdown...")
        final_audio_for_video_path = None

        try:
            swapped_audio_info = sf.info(str(swapped_audio_path))
            swapped_audio_sr = swapped_audio_info.samplerate
            swapped_audio_duration_ms = int(swapped_audio_info.duration * 1000)
        except Exception as e:
            raise Exception(f"Failed to get info for swapped audio {swapped_audio_path}: {e}")

        estimated_s2r_rir_waveform_data = None 
        estimated_s2r_rir_sr = None 

        if use_s2r_mode:
            # Use the passed global S2R RIR path
            if estimated_s2r_rir_path_global and estimated_s2r_rir_path_global.exists():
                estimated_s2r_rir_waveform_data, estimated_s2r_rir_sr = librosa.load(str(estimated_s2r_rir_path_global), sr=None, mono=True)
                log_callback(f"Loaded estimated RIR for application. SR: {estimated_s2r_rir_sr}")
                if estimated_s2r_rir_sr is not None and estimated_s2r_rir_sr != swapped_audio_sr:
                    log_callback(f"Warning: Estimated RIR SR ({estimated_s2r_rir_sr}) differs from swapped audio SR ({swapped_audio_sr}). Pedalboard will resample RIR.")
            else: 
                error_msg = "Speech2RIR mode selected, but no estimated RIR found. Please ensure RIR estimation was run and successful."
                log_callback(error_msg, True)
                progress_callback(100, "RIR processing failed.")
                raise Exception(error_msg)
        
        if use_s2r_mode and estimated_s2r_rir_waveform_data is not None:
            progress_callback(10, "Applying global Speech2RIR estimated RIR...")
            reverbed_audio_path_s2r = TMP_DIR / "final_audio_with_s2r_rir.wav"
            audio_to_mix_with_video = TMP_DIR / "final_audio_mixed_s2r.wav" # This will be the output of S2R + mix
            self.apply_reverb_with_rir_waveform_and_mix( # This method uses self.log_status, should be fine as it's called by worker.
                str(swapped_audio_path), str(reverbed_audio_path_s2r),  
                str(audio_to_mix_with_video), 
                target_sr=swapped_audio_sr,
                rir_waveform=estimated_s2r_rir_waveform_data, rir_sr=estimated_s2r_rir_sr,
                wet_mix_percentage=s2r_wet_mix_percentage,  
                progress_callback=lambda v,d: progress_callback(10 + int(v*0.3), d) # S2R RIR app is 30% of this stage
            )
            progress_callback(80, "Global Speech2RIR RIR application complete.")

        elif rir_markers: 
            progress_callback(10, "Applying RIR/effects based on marker settings...")
            final_output_samples_list = []
            last_processed_end_time_ms = 0
            total_segments = len(rir_markers)

            for i, segment_info in enumerate(rir_markers): # Iterate over list of dicts
                start_s = segment_info['start_s']
                end_s = segment_info['end_s']
                segment_params = segment_info['params'] # This is the dict with all effect settings

                segment_progress_span_pc = (30.0 / total_segments) if total_segments > 0 else 30.0 # Marker processing is 30%
                segment_progress_base = 10 + int(segment_progress_span_pc * i)
                progress_callback(segment_progress_base, f"Processing segment {i+1}/{total_segments} with effects")

                try:
                    start_ms = int(float(start_s) * 1000) # Ensure start_s is float before int conversion
                    end_ms = int(float(end_s) * 1000)     # Ensure end_s is float before int conversion
                except ValueError as e:
                    log_callback(f"Error converting segment times for segment {i+1} (start: {start_s}, end: {end_s}): {e}", is_error=True)
                    # Decide how to handle: skip segment, use defaults, or raise error
                    # For now, let's skip this segment if times are invalid
                    if last_processed_end_time_ms < swapped_audio_duration_ms: # If there's a gap, fill with original
                         y_interlude_error, _ = librosa.load(str(swapped_audio_path), sr=swapped_audio_sr, offset=last_processed_end_time_ms/1000.0, duration=(swapped_audio_duration_ms - last_processed_end_time_ms)/1000.0)
                         final_output_samples_list.append(y_interlude_error)
                         last_processed_end_time_ms = swapped_audio_duration_ms # Mark as processed up to end
                    continue # Skip to next segment

                if start_ms > last_processed_end_time_ms:
                    y_interlude, _ = librosa.load(str(swapped_audio_path), sr=swapped_audio_sr,
                                                  offset=last_processed_end_time_ms / 1000.0,
                                                  duration=(start_ms - last_processed_end_time_ms) / 1000.0)
                    final_output_samples_list.append(y_interlude)

                # Load original audio for the current segment
                y_segment_original, _ = librosa.load(str(swapped_audio_path), sr=swapped_audio_sr,
                                                     offset=start_s, duration=(end_s - start_s))
                
                processed_segment_data = self._apply_effects_to_segment_audio_data(
                    y_segment_original,
                    swapped_audio_sr,
                    segment_params
                )
                final_output_samples_list.append(processed_segment_data)
                last_processed_end_time_ms = end_ms

            if last_processed_end_time_ms < swapped_audio_duration_ms:
                y_tail, _ = librosa.load(str(swapped_audio_path), sr=swapped_audio_sr, offset=last_processed_end_time_ms / 1000.0)
                final_output_samples_list.append(y_tail)
            
            if not final_output_samples_list:
                # This case means rir_markers was empty, or all segments were zero length.
                # If rir_markers was empty, it implies no effects were intended.
                # So, copy the original swapped audio.
                log_callback("No RIR markers or segments to process for final mix. Using original swapped audio.", True)
                audio_to_mix_with_video = swapped_audio_path
            else:
                final_audio_stitched_np = np.concatenate(final_output_samples_list)
                final_processed_audio_path = TMP_DIR / "final_audio_with_all_effects.wav"
                sf.write(str(final_processed_audio_path), final_audio_stitched_np, swapped_audio_sr)
                audio_to_mix_with_video = final_processed_audio_path

            progress_callback(40, "Segment RIR/effect application complete.")

        else: 
            # No S2R mode and no RIR markers, so audio_to_mix_with_video remains swapped_audio_path
            progress_callback(40, "No RIR/effects to apply. Using swapped audio as is for next steps.")

        # --- Optional Demucs Background Mixing ---
        if enable_demucs_mix:
            progress_callback(45, "Demucs Mix: Extracting original video audio...")
            original_full_audio_temp_path = TMP_DIR / f"{original_mp4_path.stem}_original_full_audio.wav"
            try:
                # Use moviepy to extract audio, similar to _perform_audio_extraction_task
                video_clip_orig = VideoFileClip(str(original_mp4_path))
                video_clip_orig.audio.write_audiofile(str(original_full_audio_temp_path), logger=None)
                video_clip_orig.close()
                log_callback(f"Demucs Mix: Extracted original audio to {original_full_audio_temp_path}")
                progress_callback(50, "Demucs Mix: Separating original background...")

                device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
                log_callback(f"CUDA available: {torch.cuda.is_available()}")
                if torch.cuda.is_available():
                    log_callback(f"Using CUDA device: {torch.cuda.get_device_name()}")

                demucs_model = get_demucs_model('htdemucs')
                demucs_model.to(device)
                demucs_model.eval()

                original_audio_demucs_in = DemucsAudioFile(original_full_audio_temp_path).read()
                if original_audio_demucs_in.dim() == 3: original_audio_demucs_in = original_audio_demucs_in.squeeze(0)
                if original_audio_demucs_in.dim() == 1: original_audio_demucs_in = original_audio_demucs_in.unsqueeze(0)
                
                ref_orig = original_audio_demucs_in.mean(1)
                original_audio_demucs_in = (original_audio_demucs_in - ref_orig.mean()) / ref_orig.std()
                original_audio_demucs_in = original_audio_demucs_in.to(device)

                with torch.no_grad():
                    sources_orig = apply_demucs_model(demucs_model, original_audio_demucs_in.unsqueeze(0), split=True)[0]
                
                background_tensor = torch.zeros_like(sources_orig[demucs_model.sources.index('vocals')])
                for i_src, source_name in enumerate(demucs_model.sources):
                    if source_name != 'vocals':
                        background_tensor += sources_orig[i_src]
                
                background_tensor = background_tensor * ref_orig.std() + ref_orig.mean()
                background_tensor = background_tensor.cpu()
                
                original_background_temp_path = TMP_DIR / f"{original_mp4_path.stem}_original_background.wav"
                _save_demucs_audio_file(background_tensor, str(original_background_temp_path), demucs_model.samplerate, log_callback)
                progress_callback(70, "Demucs Mix: Original background separated.")

                # Mix this background with the (RIR-processed) swapped audio
                final_mixed_demucs_path = TMP_DIR / f"{original_mp4_path.stem}_demucs_mixed_audio.wav"
                _mix_audio_torchaudio(str(original_background_temp_path), 
                                      str(audio_to_mix_with_video), # This is the swapped audio + its RIR
                                      str(final_mixed_demucs_path), 
                                      demucs_bg_volume_ratio, log_callback)
                audio_to_mix_with_video = final_mixed_demucs_path # Update the audio to be muxed
                progress_callback(80, "Demucs Mix: Mixed original background with new vocals.")

                # Clean up demucs model and temp files
                del demucs_model, sources_orig, background_tensor, original_audio_demucs_in, ref_orig
                if original_full_audio_temp_path.exists(): original_full_audio_temp_path.unlink()
                if original_background_temp_path.exists(): original_background_temp_path.unlink()
                if torch.cuda.is_available(): torch.cuda.empty_cache()
                gc.collect()

            except Exception as e_demucs:
                log_callback(f"Error during Demucs background mixing: {e_demucs}. Using non-Demucs mixed audio for video.", is_error=True)
                progress_callback(80, "Demucs Mix: Error, skipping.")
        else:
            progress_callback(80, "Demucs Mix: Skipped.") # Skip Demucs part, go straight to ffmpeg merge

        progress_callback(85, "Preparing to merge final audio into video (using ffmpeg)...")

        timestamp = time.strftime("%Y%m%d-%H%M%S")
        final_video_filename = f"output_video_final_{original_mp4_path.stem}_{timestamp}.mp4"
        final_video_path = output_directory_path_obj / final_video_filename
        try:
            input_video_stream = ffmpeg.input(str(original_mp4_path_str)) 
            input_audio_stream = ffmpeg.input(str(audio_to_mix_with_video)) # Use the potentially Demucs-mixed audio
            probe = ffmpeg.probe(str(original_mp4_path_str))
            video_duration = float(probe['format']['duration'])
            (ffmpeg.output(input_video_stream['v'], input_audio_stream['a'],
                           str(final_video_path), vcodec='copy', acodec='aac', audio_bitrate='192k', t=video_duration)
             .overwrite_output().run(capture_stdout=True, capture_stderr=True))
            progress_callback(100, "Video processing complete!")
            return str(final_video_path)
        except ffmpeg.Error as e:
            error_message = f"ffmpeg error during video creation: {e.stderr.decode('utf8')}"
            log_callback(error_message, True); progress_callback(100, "ffmpeg processing failed.") 
            raise Exception(error_message)
        except Exception as e_gen:
            error_message_gen = f"General error during video creation: {str(e_gen)}"
            log_callback(error_message_gen, True); progress_callback(100, "Video creation failed.")
            raise Exception(error_message_gen)

    def set_ui_enabled_for_processing(self, enabled: bool):
        self.mp4_button.setEnabled(enabled)
        self.wav_button.setEnabled(enabled)
        self.output_dir_button.setEnabled(enabled)
        self.detect_speakers_button.setEnabled(enabled and PYANNOTE_AVAILABLE)
        self.calculate_rt60_button.setEnabled(enabled)
        self.detect_scene_cuts_button.setEnabled(enabled and SCENEDETECT_AVAILABLE)
        self.clean_target_voice_button.setEnabled(enabled and CLEARVOICE_AVAILABLE and self.target_voice_path is not None)
        self.swap_voice_button.setEnabled(enabled)
        self.final_process_button.setEnabled(enabled)
        
        is_outside = self.environment_combo.currentText() == "Outside"
        s2r_globally_available_and_not_outside = S2R_Generator is not None and not is_outside
        
        self.use_speech2rir_checkbox.setEnabled(enabled and s2r_globally_available_and_not_outside)
        
        manual_mode_active_and_possible = enabled and not self.use_speech2rir_checkbox.isChecked() and not is_outside
        #self.manual_rir_group.setEnabled(manual_mode_active_and_possible)
        
        self.use_scene_cuts_for_auto_rir_checkbox.setEnabled(
            manual_mode_active_and_possible and S2R_Generator is not None and len(self.scene_cuts_sec) > 0
        )
        self.generate_auto_rir_markers_button.setEnabled(
            manual_mode_active_and_possible and S2R_Generator is not None and self.swapped_audio_full_path is not None
        )
        
        s2r_mix_slider_active = enabled and self.use_speech2rir_checkbox.isChecked() and s2r_globally_available_and_not_outside
        self.s2r_mix_label.setEnabled(s2r_mix_slider_active)
        self.s2r_mix_slider.setEnabled(s2r_mix_slider_active)
        self.s2r_mix_value_label.setEnabled(s2r_mix_slider_active)        

    def run_target_voice_cleaning_thread(self):
        if not CLEARVOICE_AVAILABLE:
            QMessageBox.critical(self, "ClearVoice Unavailable", "ClearVoice library is not loaded. Cannot clean target voice.")
            return
        if self.current_worker and self.current_worker.isRunning():
            QMessageBox.warning(self, "Busy", "A process is already running.")
            return
        if not self.target_voice_path or not self.target_voice_path.exists():
            QMessageBox.warning(self, "No Target Voice", "Please select a target voice audio file first.")
            return
        self.log_status("Starting Speech Enhancement for target voice...")
        self.set_ui_enabled_for_processing(False)
        self.progress_bar.setRange(0, 100) 
        self.current_worker = WorkerThread(self._perform_target_voice_cleaning_task,
                                           str(self.target_voice_path), 'MossFormer2_SE_48K', 48000)
        self.current_worker.progress.connect(self.handle_gui_progress_update)
        self.current_worker.finished.connect(self._on_target_voice_cleaning_finished)
        self.current_worker.log_message.connect(self._handle_worker_log_message)
        self.current_worker.start()

    def _perform_target_voice_cleaning_task(self, original_target_wav_path_str, model_name_se, target_se_sr, progress_callback, log_callback):
        progress_callback(0, "Initializing Speech Enhancement for target voice...")
        try:
            progress_callback(10, f"Loading SE model: {model_name_se}...")
            cv_instance = ClearVoice(task='speech_enhancement', model_names=[model_name_se])
            progress_callback(30, "Performing speech enhancement on target voice...")
            output_data = cv_instance(input_path=original_target_wav_path_str, online_write=False)
            del cv_instance
            if torch.cuda.is_available(): torch.cuda.empty_cache()
            progress_callback(80, "Speech Enhancement model processed and unloaded.")
            if isinstance(output_data, dict):
                key = next(iter(output_data)); cleaned_wav_data = output_data[key]
            else: cleaned_wav_data = output_data
            if cleaned_wav_data.ndim > 1: cleaned_wav_data = cleaned_wav_data[0, :] 
            cleaned_target_voice_path = TMP_DIR / f"{Path(original_target_wav_path_str).stem}_cleaned_se.wav"
            sf.write(str(cleaned_target_voice_path), cleaned_wav_data, target_se_sr)
            progress_callback(90, f"Cleaned target voice saved to {cleaned_target_voice_path.name}")
            progress_callback(100, "Target voice cleaning complete.")
            return str(cleaned_target_voice_path)
        except Exception as e:
            import traceback
            log_callback(f"Error during target voice cleaning: {e}\n{traceback.format_exc()}", True)
            progress_callback(100, "Target voice cleaning failed.")
            raise e 

    def _on_target_voice_cleaning_finished(self, cleaned_path_str, message):
        self.set_ui_enabled_for_processing(True)
        if cleaned_path_str:
            self.target_voice_path = Path(cleaned_path_str)
            self.wav_label.setText(f"{self.target_voice_path.name} (Cleaned SE)")
            self.log_status(f"Target voice cleaned successfully. Using: {self.target_voice_path}")
            QMessageBox.information(self, "Cleaning Complete", f"Target voice enhanced and updated to:\n{self.target_voice_path}")
        else:
            self.log_status(f"Target voice cleaning failed: {message}", is_error=True)
            QMessageBox.critical(self, "Cleaning Error", f"Could not clean target voice: {message}")
        self.current_worker = None

    def _apply_speech_enhancement_to_segment(self, input_segment_path_str, output_segment_path_str, progress_callback, log_callback):
        if CLEARVOICE_MODEL_CACHE is None or CLEARVOICE_MODEL_CACHE.get('model_name') != self.clearvoice_se_model_name:
            log_callback(f"SE model {self.clearvoice_se_model_name} not loaded or different. Loading for segment processing.")
            if not self.load_clearvoice_model(lambda v,d: progress_callback(int(v*0.2), f"SE Model Load: {d}"), log_callback):
                log_callback("Failed to load ClearVoice SE model for segment. Skipping SE.", True)
                shutil.copy(input_segment_path_str, output_segment_path_str)
                progress_callback(100, "SE skipped (model load fail).")
                return
        
        cv_instance = CLEARVOICE_MODEL_CACHE['model']
        progress_callback(20, "Applying SE to segment...")
        try:
            output_data_se = cv_instance(input_path=input_segment_path_str, online_write=False)
            progress_callback(80, "SE processing complete by model.")

            if isinstance(output_data_se, dict):
                key = next(iter(output_data_se)); cleaned_wav_data = output_data_se[key]
            else: cleaned_wav_data = output_data_se
            if cleaned_wav_data.ndim > 1: cleaned_wav_data = cleaned_wav_data[0, :] 
            
            sf.write(output_segment_path_str, cleaned_wav_data, CLEARVOICE_MODEL_CACHE.get('model_sr', 48000)) 
            progress_callback(100, "Cleaned segment saved.")
        except Exception as e:
            log_callback(f"Error applying SE to segment {input_segment_path_str}: {e}", True)
            shutil.copy(input_segment_path_str, output_segment_path_str) 
            progress_callback(100, "SE failed, used original.")

    @Slot()
    def unload_s2r_model(self):
        global S2R_MODEL_CACHE
        if S2R_MODEL_CACHE is not None:
            self.log_status("Unloading Speech2RIR model from memory...")
            if 'model' in S2R_MODEL_CACHE and hasattr(S2R_MODEL_CACHE['model'], 'to') and 'device' in S2R_MODEL_CACHE and S2R_MODEL_CACHE['device'].type == 'cuda':
                try: S2R_MODEL_CACHE['model'].to('cpu')
                except Exception as e: self.log_status(f"Could not move Speech2RIR model to CPU: {e}", is_error=True)
            if 'model' in S2R_MODEL_CACHE: del S2R_MODEL_CACHE['model']
            S2R_MODEL_CACHE.clear(); S2R_MODEL_CACHE = None
            self.clear_cuda_cache_and_gc()
            self.log_status("Speech2RIR model unloaded.")
        else:
            self.log_status("Speech2RIR model already unloaded or not loaded.", is_error=False)

    def load_s2r_model(self, progress_callback):

        global S2R_MODEL_CACHE
        if S2R_MODEL_CACHE is not None:
            progress_callback(100, "Speech2RIR model already loaded.")
            return True
        if S2R_Generator is None: 
            self.log_status("Speech2RIR Generator class not loaded (Import Error). Cannot load model.", is_error=True)
            progress_callback(100, "Speech2RIR Load Failed (Import).")
            return False
        progress_callback(10, "Loading Speech2RIR model...")
        checkpoint_path = S2R_PROJECT_ROOT / "checkpoint-1040000steps.pkl"
        config_path = "configs/config_speech2rir.yml"
        s2r_model_gdrive_id = "1CcF1c9i76-MVPJ-PGoBwvaVtOUBPD57y"
        
        if not checkpoint_path.exists():
            self.log_status(f"Speech2RIR checkpoint not found at {checkpoint_path}. Attempting download...")
            progress_callback(15, "Speech2RIR Ckpt not found. Attempting download...")
            if GDOWN_AVAILABLE:
                try:
                    S2R_PROJECT_ROOT.mkdir(parents=True, exist_ok=True) # Ensure directory exists
                    gdown.download(id=s2r_model_gdrive_id, output=str(checkpoint_path), quiet=False)
                    if checkpoint_path.exists():
                        self.log_status("Speech2RIR checkpoint downloaded successfully.")
                        progress_callback(20, "Speech2RIR Ckpt downloaded.")
                    else:
                        self.log_status("Speech2RIR checkpoint download failed. File still not found.", is_error=True)
                        progress_callback(100, "Speech2RIR Ckpt Download Failed.")
                        return False
                except Exception as e:
                    self.log_status(f"Error downloading Speech2RIR checkpoint: {e}", is_error=True)
                    progress_callback(100, "Speech2RIR Ckpt Download Error.")
                    return False
            else:
                self.log_status("gdown library not available. Cannot download Speech2RIR checkpoint.", is_error=True)
                progress_callback(100, "Speech2RIR Ckpt Download Skipped (gdown missing).")
                return False

        try:
            with open(config_path, 'r') as f: s2r_exp_config = yaml.load(f, Loader=yaml.FullLoader)
            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
            s2r_model = S2R_Generator(**s2r_exp_config['generator_params'])
            state_dict = torch.load(checkpoint_path, map_location='cpu')
            s2r_model.load_state_dict(state_dict['model']['generator'])
            s2r_model.eval().to(device)
            S2R_MODEL_CACHE = {'model': s2r_model, 'config': s2r_exp_config, 'device': device}
            progress_callback(100, "Speech2RIR model loaded successfully.")
            self.log_status(f"Speech2RIR model loaded from {checkpoint_path}.")
            return True
        except Exception as e:
            import traceback
            self.log_status(f"Error loading Speech2RIR model: {e}\n{traceback.format_exc()}", is_error=True)
            S2R_MODEL_CACHE = None
            progress_callback(100, "Speech2RIR Load Error.")
            return False

    def run_s2r_rir_estimation_thread(self):

        if self.current_worker and self.current_worker.isRunning():
            QMessageBox.warning(self, "Busy", "A process is already running.")
            return
        if not self.source_audio_full_path:
            QMessageBox.warning(self, "No Source Audio", "Please load an MP4 and extract audio first.")
            return
        if S2R_Generator is None:
            QMessageBox.critical(self, "Speech2RIR Error", "Speech2RIR model components are not available due to an import error.")
            return
        self.log_status("Starting Speech2RIR estimation from source audio...")
        self.set_ui_enabled_for_processing(False) 
        self.progress_bar.setRange(0,100) 
        self.current_worker = WorkerThread(self._perform_s2r_estimation_task, str(self.source_audio_full_path))
        self.current_worker.progress.connect(self.handle_gui_progress_update)
        self.current_worker.finished.connect(self._on_s2r_estimation_finished)
        self.current_worker.start()

    def _perform_s2r_estimation_task(self, source_audio_path_str, progress_callback, log_callback, output_rir_basename=None):

        print(f"DEBUG: _perform_s2r_estimation_task started for: {source_audio_path_str}, output_basename: {output_rir_basename}")
        if not self.load_s2r_model(lambda v,d: progress_callback(int(v*0.2), f"Speech2RIR Model Load: {d}")): 
            raise Exception("Failed to load Speech2RIR model.")
        s2r_model_pack = S2R_MODEL_CACHE
        s2r_model = s2r_model_pack['model']; s2r_exp_config = s2r_model_pack['config']; device = s2r_model_pack['device']
        target_s2r_sr = s2r_exp_config.get('sampling_rate', 16000) 
        print(f"DEBUG: Speech2RIR Target SR: {target_s2r_sr} Hz")
        progress_callback(25, "Speech2RIR: Loading and preprocessing source audio...")
        
        print(f"DEBUG_PRINT_S2R_PRE_LIBROSA_LOAD: Attempting librosa.load for chunk '{output_rir_basename}'. Path: {source_audio_path_str}"); sys.stdout.flush()
        try:
            if not Path(source_audio_path_str).exists(): raise FileNotFoundError(f"Audio file for chunk '{output_rir_basename}' not found at: {source_audio_path_str}")
            full_audio, full_sr = librosa.load(source_audio_path_str, sr=None, mono=True)
            print(f"DEBUG_PRINT_S2R_POST_LIBROSA_LOAD: librosa.load successful for chunk '{output_rir_basename}'. SR={full_sr}, Audio length: {len(full_audio) if full_audio is not None else 'None'}"); sys.stdout.flush()
        except Exception as e_load: print(f"CRITICAL_ERROR_LIBROSA_LOAD for chunk '{output_rir_basename}': {e_load}"); sys.stdout.flush(); raise 
        
        full_duration_sec = len(full_audio) / full_sr
        print(f"DEBUG: Speech2RIR: Loaded audio for RIR estimation. Path: {source_audio_path_str}, Original Duration: {full_duration_sec:.2f}s, SR: {full_sr}Hz")

        audio_segment_to_process = full_audio; segment_sr = full_sr

        # New logic: Limit input to S2R model to at most the first 20 seconds of the provided chunk
        MAX_INPUT_DURATION_FOR_S2R_MODEL_SEC = 20.0
        if full_duration_sec > MAX_INPUT_DURATION_FOR_S2R_MODEL_SEC:
            progress_callback(30, f"Speech2RIR: Input chunk for RIR estimation is long ({full_duration_sec:.1f}s), using first {MAX_INPUT_DURATION_FOR_S2R_MODEL_SEC:.1f}s.")
            num_samples_to_take = int(MAX_INPUT_DURATION_FOR_S2R_MODEL_SEC * full_sr)
            audio_segment_to_process = full_audio[:num_samples_to_take]
            log_callback(f"Speech2RIR: Using first {len(audio_segment_to_process)/segment_sr:.1f}s of input chunk '{Path(source_audio_path_str).name}' for RIR estimation.")
            print(f"DEBUG: Speech2RIR: Selected first {MAX_INPUT_DURATION_FOR_S2R_MODEL_SEC:.1f}s. New segment length: {len(audio_segment_to_process)/segment_sr:.1f}s")
        else:
            progress_callback(30, f"Speech2RIR: Input chunk for RIR estimation ({full_duration_sec:.1f}s) is within {MAX_INPUT_DURATION_FOR_S2R_MODEL_SEC:.1f}s limit.")
            log_callback(f"Speech2RIR: Using full input chunk '{Path(source_audio_path_str).name}' ({full_duration_sec:.1f}s) for RIR estimation.")
            print(f"DEBUG: Speech2RIR: Using full segment. Length: {len(audio_segment_to_process)/segment_sr:.1f}s")

        print(f"DEBUG_PRINT_PRE_RESAMPLE_CONDITION_CHECK: For chunk '{output_rir_basename}', segment_sr={segment_sr}, target_s2r_sr={target_s2r_sr}, full_duration_sec={full_duration_sec}"); sys.stdout.flush()
        if segment_sr != target_s2r_sr: 
             print(f"DEBUG_PRINT_ENTERED_RESAMPLE_IF: Condition (segment_sr != target_s2r_sr) is TRUE for chunk '{output_rir_basename}'."); sys.stdout.flush()
             current_chunk_id = output_rir_basename if output_rir_basename else "UnknownChunk"
             print(f"DEBUG_PRINT_A: ENTERING RESAMPLE BLOCK for chunk '{current_chunk_id}'. SR: {segment_sr}Hz -> {target_s2r_sr}Hz."); sys.stdout.flush() # type: ignore
             log_callback(f"Speech2RIR: Resampling segment from {segment_sr}Hz to {target_s2r_sr}Hz...")
             print(f"DEBUG_PRINT_B: AFTER self.log_status for chunk '{current_chunk_id}'."); sys.stdout.flush()
             print(f"DEBUG_PRINT_C: Before torch.from_numpy for chunk '{current_chunk_id}'. audio_segment_to_process is None: {audio_segment_to_process is None}")
             if audio_segment_to_process is not None: print(f"DEBUG_PRINT_D: audio_segment_to_process shape: {audio_segment_to_process.shape}, dtype: {audio_segment_to_process.dtype} for chunk '{current_chunk_id}'"); sys.stdout.flush()
             try:
                 audio_segment_tensor = torch.from_numpy(audio_segment_to_process).float().unsqueeze(0) 
                 print(f"DEBUG_PRINT_E: After torch.from_numpy for chunk '{current_chunk_id}'. Shape: {audio_segment_tensor.shape}"); sys.stdout.flush()
                 print(f"DEBUG_PRINT_F: Before torchaudio.transforms.Resample instantiation for chunk '{current_chunk_id}'."); sys.stdout.flush()
                 resampler = torchaudio.transforms.Resample(orig_freq=segment_sr, new_freq=target_s2r_sr, dtype=audio_segment_tensor.dtype)
                 print(f"DEBUG_PRINT_G: After torchaudio.transforms.Resample instantiation for chunk '{current_chunk_id}'."); sys.stdout.flush()
                 print(f"DEBUG_PRINT_H: Before resampler(audio_segment_tensor) call for chunk '{current_chunk_id}'."); sys.stdout.flush()
                 resampled_audio_segment_tensor = resampler(audio_segment_tensor)
                 print(f"DEBUG_PRINT_I: After resampler(audio_segment_tensor) call for chunk '{current_chunk_id}'. Shape: {resampled_audio_segment_tensor.shape}"); sys.stdout.flush()
                 audio_for_s2r = resampled_audio_segment_tensor.squeeze(0).cpu().numpy() 
                 print(f"DEBUG_PRINT_J: After .squeeze().cpu().numpy() for chunk '{current_chunk_id}'. Shape: {audio_for_s2r.shape}"); sys.stdout.flush()
                 print(f"DEBUG: Speech2RIR: Resampling complete for chunk '{current_chunk_id}'. New length: {len(audio_for_s2r)/target_s2r_sr:.2f}s")
                 print(f"DEBUG: Speech2RIR: Explicitly deleting intermediate resampling tensors for chunk '{current_chunk_id}'."); del audio_segment_tensor; del resampled_audio_segment_tensor
                 if device.type == 'cuda': print(f"DEBUG: Speech2RIR: Emptying CUDA cache (device: {device}) for chunk '{current_chunk_id}'."); torch.cuda.empty_cache()
                 print(f"DEBUG: Speech2RIR: Explicit del and gc.collect after resampling block DONE for chunk '{current_chunk_id}'.")
             except Exception as e_resample: print(f"CRITICAL_ERROR_RESAMPLING for chunk '{current_chunk_id}': {e_resample}"); sys.stdout.flush(); import traceback; traceback.print_exc(); raise
        else:
             audio_for_s2r = audio_segment_to_process 
             print(f"DEBUG: Speech2RIR: No resampling needed. Using segment as is. Length: {len(audio_for_s2r)/target_s2r_sr:.2f}s")
        print("DEBUG: Speech2RIR: Converting audio_for_s2r to tensor...")
        if not audio_for_s2r.flags['C_CONTIGUOUS']:
            print(f"DEBUG: Speech2RIR: audio_for_s2r for chunk '{current_chunk_id if 'current_chunk_id' in locals() else final_rir_filename_stem}' is not C-contiguous. Making a copy.")
            audio_for_s2r = np.ascontiguousarray(audio_for_s2r)
        audio_tensor_for_s2r = torch.tensor(audio_for_s2r, dtype=torch.float32).unsqueeze(0).unsqueeze(0).to(device) 
        print(f"DEBUG: Speech2RIR: Converted to tensor. Shape: {audio_tensor_for_s2r.shape}")
        progress_callback(50, "Speech2RIR: Estimating RIR from segment...")
        final_rir_filename_stem = output_rir_basename if output_rir_basename else "estimated_s2r_rir" 
        print(f"DEBUG: Speech2RIR: final_rir_filename_stem: {final_rir_filename_stem}")
        try:
            print("DEBUG: Speech2RIR: Starting model inference (s2r_model(audio_tensor_for_s2r))..."); 
            with torch.no_grad(): estimated_rir_tensor = s2r_model(audio_tensor_for_s2r)
            estimated_rir_waveform = estimated_rir_tensor.squeeze().cpu().numpy()
            rir_sr = target_s2r_sr; initial_meaningful_duration_sec = 0.08; meaningful_samples = int(initial_meaningful_duration_sec * float(rir_sr))
            if len(estimated_rir_waveform) > meaningful_samples:
                decay_window = np.exp(-np.linspace(0, 15, len(estimated_rir_waveform) - meaningful_samples)) # Reduced aggressiveness from 25 to 15
                processed_rir = np.copy(estimated_rir_waveform); processed_rir[meaningful_samples:] *= decay_window # type: ignore
                estimated_rir_waveform = processed_rir; log_callback(f"Speech2RIR: Applied aggressive tail attenuation to RIR.")
            final_peak_val = np.max(np.abs(estimated_rir_waveform)); 
            if final_peak_val > 1e-7: estimated_rir_waveform = estimated_rir_waveform / final_peak_val            
            print(f"DEBUG: Speech2RIR: Model inference complete. Estimated RIR waveform shape: {estimated_rir_waveform.shape if isinstance(estimated_rir_waveform, np.ndarray) else 'None'}")
            del audio_tensor_for_s2r; del estimated_rir_tensor
            if device.type == 'cuda': print(f"DEBUG: Speech2RIR: Emptying CUDA cache after Speech2RIR inference (device: {device})."); torch.cuda.empty_cache()
            print(f"DEBUG: Speech2RIR: Skipped gc.collect() post-inference for chunk '{final_rir_filename_stem}'."); sys.stdout.flush() # type: ignore
        except Exception as e:
            log_callback(f"Speech2RIR: Error during model inference for '{final_rir_filename_stem}': {e}. Creating a dummy impulse.", True)
            print(f"DEBUG: Speech2RIR: Exception during model inference: {e}. Creating dummy RIR.")
            estimated_rir_waveform = np.zeros(1, dtype=np.float32); estimated_rir_waveform[0] = 1.0
        print(f"DEBUG_PRINT_S2R_PRE_PROGRESS_CB_70: Before progress_callback(70) for chunk '{final_rir_filename_stem}'."); sys.stdout.flush()
        progress_callback(70, "Speech2RIR: Post-processing RIR...")
        print(f"DEBUG_PRINT_S2R_POST_PROGRESS_CB_70: After progress_callback(70) for chunk '{final_rir_filename_stem}'."); sys.stdout.flush()
        print(f"DEBUG_PRINT_S2R_PRE_NAN_CHECK: Before NaN/Inf check for chunk '{final_rir_filename_stem}'. estimated_rir_waveform is None: {estimated_rir_waveform is None}")
        if estimated_rir_waveform is not None: print(f"DEBUG_PRINT_S2R_PRE_NAN_CHECK_DETAILS: Waveform type: {type(estimated_rir_waveform)}, size: {estimated_rir_waveform.size if isinstance(estimated_rir_waveform, np.ndarray) else 'N/A'} for chunk '{final_rir_filename_stem}'"); sys.stdout.flush()
        is_invalid_rir = False
        if not isinstance(estimated_rir_waveform, np.ndarray): print(f"DEBUG_PRINT_S2R_NAN_CHECK_FAIL_TYPE: Not ndarray for chunk '{final_rir_filename_stem}'."); is_invalid_rir = True
        elif estimated_rir_waveform.size == 0: print(f"DEBUG_PRINT_S2R_NAN_CHECK_FAIL_EMPTY: Empty ndarray for chunk '{final_rir_filename_stem}'."); is_invalid_rir = True
        elif np.any(np.isnan(estimated_rir_waveform)): print(f"DEBUG_PRINT_S2R_NAN_CHECK_FAIL_NAN: NaN found for chunk '{final_rir_filename_stem}'."); is_invalid_rir = True
        elif np.any(np.isinf(estimated_rir_waveform)): print(f"DEBUG_PRINT_S2R_NAN_CHECK_FAIL_INF: Inf found for chunk '{final_rir_filename_stem}'."); is_invalid_rir = True
        if is_invalid_rir:
            print(f"DEBUG: Speech2RIR: RIR is invalid (NaN/Inf/empty/wrong type). Original shape: {estimated_rir_waveform.shape if isinstance(estimated_rir_waveform, np.ndarray) else 'None'}. Creating dummy for chunk '{final_rir_filename_stem}'.")
            log_callback(f"Speech2RIR: Estimated RIR for '{final_rir_filename_stem}' is invalid (NaN/Inf/empty). Creating a dummy impulse.", True)
            estimated_rir_waveform = np.zeros(1, dtype=np.float32); estimated_rir_waveform[0] = 1.0 
        print(f"DEBUG_PRINT_S2R_POST_NAN_CHECK: After NaN/Inf check block for chunk '{final_rir_filename_stem}'."); sys.stdout.flush()
        print(f"DEBUG: Speech2RIR: After NaN/Inf check, RIR waveform shape: {estimated_rir_waveform.shape if isinstance(estimated_rir_waveform, np.ndarray) and estimated_rir_waveform is not None else 'N/A'}, max_abs: {np.max(np.abs(estimated_rir_waveform)) if isinstance(estimated_rir_waveform, np.ndarray) and estimated_rir_waveform is not None and estimated_rir_waveform.size > 0 else 'N/A'}"); sys.stdout.flush()
        if not (len(estimated_rir_waveform) > 0 and np.max(np.abs(estimated_rir_waveform)) > 1e-7):
            if not (len(estimated_rir_waveform) == 1 and estimated_rir_waveform[0] == 1.0): 
                log_callback(f"Speech2RIR: Estimated RIR for '{final_rir_filename_stem}' is silent or too short. Creating a dummy impulse.", True)
                print(f"DEBUG: Speech2RIR: RIR is silent or too short (and not already the dummy). Creating dummy.")
                estimated_rir_waveform = np.zeros(1, dtype=np.float32); estimated_rir_waveform[0] = 1.0
            output_rir_path = TMP_DIR / f"{final_rir_filename_stem}_dummy.wav"
            print(f"DEBUG: Speech2RIR: Saving dummy RIR to {output_rir_path}")
            sf.write(str(output_rir_path), estimated_rir_waveform, target_s2r_sr) # type: ignore
            progress_callback(100, f"Speech2RIR: Dummy RIR created as {output_rir_path.name}.")
            print(f"DEBUG: _perform_s2r_estimation_task returning dummy RIR path: {output_rir_path}")
            return str(output_rir_path)
        print("DEBUG: Speech2RIR: RIR seems valid for post-processing.")
        peak_val = np.max(np.abs(estimated_rir_waveform))
        if peak_val > 1e-7:
            estimated_rir_waveform = estimated_rir_waveform / peak_val # type: ignore
            log_callback(f"Speech2RIR: RIR for '{final_rir_filename_stem}' initially peak-normalized."); print(f"DEBUG: Speech2RIR: Initial peak normalization done. Peak was {peak_val}.")
        else:
            print(f"DEBUG: Speech2RIR: RIR is effectively silent before normalization (peak_val <= 1e-7). This should have been caught earlier.")
            log_callback(f"Speech2RIR: RIR for '{final_rir_filename_stem}' is effectively silent before normalization. Using as is (might be problematic).", True)

        peak_index = np.argmax(np.abs(estimated_rir_waveform)).item() # Ensure it's a Python int
        
        allow_leading_ms = 0; allowed_leading_samples = int(allow_leading_ms / 1000.0 * target_s2r_sr); start_index = max(0, peak_index - allowed_leading_samples)
        if start_index > 0:
            estimated_rir_waveform = estimated_rir_waveform[start_index:]
            print(f"DEBUG: Speech2RIR: Trimmed {start_index} leading samples. New length: {len(estimated_rir_waveform)}")
            log_callback(f"Speech2RIR: Aligned RIR for '{final_rir_filename_stem}' by trimming {start_index} leading samples.")
        else:
            log_callback(f"Speech2RIR: RIR peak for '{final_rir_filename_stem}' is already near the beginning. No leading trim applied."); print("DEBUG: Speech2RIR: No leading trim applied.")
        max_rir_len_sec_config = 1.5 ; max_rir_samples_config = int(max_rir_len_sec_config * target_s2r_sr)
        if len(estimated_rir_waveform) > max_rir_samples_config:
            log_callback(f"Speech2RIR: Trimming RIR tail for '{final_rir_filename_stem}' from {len(estimated_rir_waveform)/target_s2r_sr:.2f}s to {max_rir_len_sec_config:.2f}s.")
            print(f"DEBUG: Speech2RIR: Trimming RIR tail. Original length: {len(estimated_rir_waveform)}, Max samples: {max_rir_samples_config}")
            estimated_rir_waveform = estimated_rir_waveform[:max_rir_samples_config]; print(f"DEBUG: Speech2RIR: Tail trimmed. New length: {len(estimated_rir_waveform)}")
        if len(estimated_rir_waveform) > 0: 
            print("DEBUG: Speech2RIR: RIR not empty after trimming. Proceeding with final normalization.")
            final_peak_val = np.max(np.abs(estimated_rir_waveform))
            if final_peak_val > 1e-7:
                estimated_rir_waveform = estimated_rir_waveform / final_peak_val # type: ignore
                log_callback(f"Speech2RIR: RIR for '{final_rir_filename_stem}' re-normalized after alignment/trimming."); print(f"DEBUG: Speech2RIR: Final peak normalization done. Peak was {final_peak_val}.")
            else:
                log_callback(f"Speech2RIR: RIR for '{final_rir_filename_stem}' became silent after trimming. Creating dummy.", True)
                print("DEBUG: Speech2RIR: RIR became silent after trimming. Creating dummy."); estimated_rir_waveform = np.zeros(1, dtype=np.float32); estimated_rir_waveform[0] = 1.0
        else: 
             log_callback(f"Speech2RIR: RIR for '{final_rir_filename_stem}' became empty after processing. Creating dummy impulse.", True)
             print("DEBUG: Speech2RIR: RIR became empty after trimming. Creating dummy."); estimated_rir_waveform = np.zeros(1, dtype=np.float32); estimated_rir_waveform[0] = 1.0
        output_rir_path = TMP_DIR / f"{final_rir_filename_stem}.wav"
        print(f"DEBUG: Speech2RIR: Saving final RIR to {output_rir_path}")
        sf.write(str(output_rir_path), estimated_rir_waveform, target_s2r_sr)
        print("DEBUG: Speech2RIR: Final RIR saved.")
        progress_callback(100, "Speech2RIR: RIR estimation and post-processing complete.")
        print(f"DEBUG: _perform_s2r_estimation_task returning final RIR path: {output_rir_path}")
        return str(output_rir_path)

    def _on_s2r_estimation_finished(self, result_rir_path, message):

        self.set_ui_enabled_for_processing(True) 
        if result_rir_path:
            self.estimated_rir_path = Path(result_rir_path)
            self.log_status(f"Speech2RIR estimation successful. RIR saved to: {self.estimated_rir_path}")
            QMessageBox.information(self, "RIR Estimation Complete", f"RIR estimated from source and saved to {self.estimated_rir_path}")
        else:
            self.log_status(f"Speech2RIR estimation failed: {message}", is_error=True)
            self.estimated_rir_path = None
            self.use_speech2rir_checkbox.setChecked(False) 
            QMessageBox.critical(self, "RIR Estimation Error", f"Could not estimate RIR using Speech2RIR: {message}")
        self.current_worker = None
        self.unload_s2r_model() 

    def run_scene_cut_detection_thread(self):

        if not SCENEDETECT_AVAILABLE: QMessageBox.critical(self, "Scene Detection Unavailable", "PySceneDetect library is not installed."); return
        if self.current_worker and self.current_worker.isRunning(): QMessageBox.warning(self, "Busy", "A process is already running."); return
        if not self.input_mp4_path: QMessageBox.warning(self, "No Video", "Please select an input MP4 file first."); return
        self.log_status("Starting scene cut detection...")
        self.set_ui_enabled_for_processing(False)
        self.progress_bar.setRange(0, 0)

        # Get speedup parameters from UI
        detector_name = self.sd_detector_combo.currentText().split(" ")[0] # Get "ContentDetector" from "ContentDetector (Default)"
        use_auto_downscale = self.sd_auto_downscale_checkbox.isChecked()
        try:
            manual_downscale = int(self.sd_manual_downscale_edit.text())
            if manual_downscale < 1: manual_downscale = 1
        except ValueError:
            manual_downscale = 1
            self.log_status("Invalid manual downscale factor, using 1.", is_error=True)
        
        try:
            frame_skip_value = int(self.sd_frame_skip_edit.text())
            if frame_skip_value < 0: frame_skip_value = 0
        except ValueError:
            frame_skip_value = 0
            self.log_status("Invalid frame skip value, using 0.", is_error=True)
            
        try:
            content_threshold = float(self.sd_content_threshold_edit.text())
        except ValueError:
            content_threshold = 27.0 # Default from PySceneDetect
            self.log_status(f"Invalid ContentDetector threshold, using {content_threshold}.", is_error=True)
        
        # Add similar logic for other detector thresholds if you make them selectable

        self.current_worker = WorkerThread(
            self._perform_scene_cut_detection_task, 
            str(self.input_mp4_path),
            detector_name=detector_name,
            content_detector_threshold=content_threshold, # Pass threshold
            # Pass other detector-specific params as needed
            use_auto_downscale=use_auto_downscale,
            manual_downscale_factor=manual_downscale,
            frame_skip=frame_skip_value
        )
        self.current_worker.finished.connect(self._on_scene_cut_detection_finished)
        self.current_worker.log_message.connect(self._handle_worker_log_message) # Connect log message
        self.current_worker.start()

    def _perform_scene_cut_detection_task(self, video_path_str, progress_callback, 
                                          log_callback, detector_name="ContentDetector", # Added log_callback
                                          content_detector_threshold=27.0,
                                          # Add other detector params here, e.g., hash_detector_threshold=0.25
                                          use_auto_downscale=True,
                                          manual_downscale_factor=1,
                                          frame_skip=0):
        progress_callback(0, "Initializing Scene Detector...")
        video = open_video(video_path_str)
        scene_manager = SceneManager()

        # --- Apply Speedup Strategies ---
        if use_auto_downscale:
            scene_manager.auto_downscale = True
            log_callback("Scene Detection: Auto downscaling enabled.")
        elif manual_downscale_factor > 1:
            scene_manager.auto_downscale = False
            scene_manager.downscale = manual_downscale_factor
            log_callback(f"Scene Detection: Manual downscale factor set to {manual_downscale_factor}.")
        else:
            scene_manager.auto_downscale = False # Explicitly disable if factor is 1
            scene_manager.downscale = 1 
            log_callback("Scene Detection: No downscaling applied (factor 1 or auto disabled).")

        if frame_skip > 0:
            log_callback(f"Scene Detection: Frame skip set to {frame_skip} (processing 1 in every {frame_skip+1} frames).")
        
        # --- Detector Selection ---
        detector = None
        if detector_name == "ContentDetector":
            detector = ContentDetector(threshold=content_detector_threshold)
            log_callback(f"Using ContentDetector with threshold: {content_detector_threshold}")
        elif detector_name == "HashDetector":
            # You might want to expose HashDetector's parameters (threshold, size, lowpass)
            # to the UI as well. Using defaults for now.
            detector = HashDetector(threshold=0.25) # Example threshold, PySceneDetect CLI default is 0.395
            log_callback("Using HashDetector (Faster).")
        elif detector_name == "AdaptiveDetector":
            detector = AdaptiveDetector() # Uses ContentDetector logic internally, threshold is adaptive
            log_callback("Using AdaptiveDetector.")
        elif detector_name == "HistogramDetector":
            detector = HistogramDetector() # Example threshold
            log_callback("Using HistogramDetector (Faster).")
        else: # Fallback
            log_callback(f"Unknown detector '{detector_name}', defaulting to ContentDetector.", True)
            detector = ContentDetector(threshold=content_detector_threshold)
        
        scene_manager.add_detector(detector)
        progress_callback(20, "Detecting scenes...")
        # Pass frame_skip to detect_scenes
        scene_manager.detect_scenes(video=video, show_progress=False, frame_skip=frame_skip)
        
        scene_list_msec = scene_manager.get_scene_list()
        scene_cuts_seconds = []
        if scene_list_msec:
            # The first element of scene_list_msec is the start of the first scene (usually 0).
            # Cuts occur at the *start* of subsequent scenes.
            for i, scene_pair in enumerate(scene_list_msec):
                if i > 0: # Skip the very first scene's start time as it's not a "cut"
                    scene_cuts_seconds.append(scene_pair[0].get_seconds())
        
        progress_callback(100, "Scene detection complete.")
        # Clean up
        del video
        del scene_manager
        del detector
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        gc.collect()
        return sorted(list(set(scene_cuts_seconds)))

    def _on_scene_cut_detection_finished(self, scene_cuts_list_sec, message):

        self.set_ui_enabled_for_processing(True)
        self.progress_bar.setRange(0, 100); self.progress_bar.setValue(100 if scene_cuts_list_sec is not None else 0)
        if scene_cuts_list_sec is not None:
            self.scene_cuts_sec = scene_cuts_list_sec
            self.log_status(f"Scene cut detection successful. Found {len(self.scene_cuts_sec)} cuts. {message}")
            QMessageBox.information(self, "Scene Detection Complete", f"Found {len(self.scene_cuts_sec)} scene cuts.")
            self.use_scene_cuts_for_auto_rir_checkbox.setEnabled(len(self.scene_cuts_sec) > 0)
        else:
            self.log_status(f"Scene cut detection failed: {message}", is_error=True)
            QMessageBox.critical(self, "Scene Detection Error", f"Could not detect scene cuts: {message}")
            self.scene_cuts_sec = []; self.use_scene_cuts_for_auto_rir_checkbox.setEnabled(False)
        self.current_worker = None

    def run_auto_rir_generation_thread(self):

        if not S2R_Generator: QMessageBox.critical(self, "Speech2RIR Unavailable", "Speech2RIR model components not loaded."); return
        if not self.swapped_audio_full_path or not self.swapped_audio_full_path.exists():
            QMessageBox.warning(self, "No Swapped Audio", "Please perform 'Swap Voice' first to generate swapped audio."); return
        if self.current_worker and self.current_worker.isRunning(): QMessageBox.warning(self, "Busy", "A process is already running."); return
        if self.use_speech2rir_checkbox.isChecked():
            self.use_speech2rir_checkbox.setChecked(False); # log_status will be called by _on_speech2rir_toggled
        if self.swapped_audio_markers:
            reply = QMessageBox.question(self, "Clear Existing Markers?", "Auto Generating RIR markers will clear all current RIR markers. Continue?", QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
            if reply == QMessageBox.No: return
            self.clear_all_rir_markers_action()
        self.log_status(f"Starting automatic RIR marker generation (segment duration: {self.auto_rir_segment_duration_sec}s)...")
        self.set_ui_enabled_for_processing(False); self.progress_bar.setRange(0, 100)
        self.current_worker = WorkerThread(self._perform_auto_rir_marking_task, str(self.swapped_audio_full_path), self.auto_rir_segment_duration_sec)
        self.current_worker.progress.connect(self.handle_gui_progress_update)
        self.current_worker.finished.connect(self._on_auto_rir_marking_finished)
        self.current_worker.start()

    def _perform_auto_rir_marking_task(self, swapped_audio_path_str, segment_duration_sec, progress_callback, log_callback):

        print(f"DEBUG: _perform_auto_rir_marking_task started for: {swapped_audio_path_str}, segment_duration: {segment_duration_sec}s")
        print("DEBUG: AutoRIR: Loading Speech2RIR model...")
        if not self.load_s2r_model(lambda v, d: progress_callback(int(v * 0.1), f"Speech2RIR Model Load: {d}")) :
            print("DEBUG: AutoRIR: Failed to load Speech2RIR model."); raise Exception("Failed to load Speech2RIR model for auto RIR marking.")
        
        print("DEBUG: AutoRIR: Speech2RIR model loaded (or was already loaded).")
        print(f"DEBUG: AutoRIR: Loading swapped audio from {swapped_audio_path_str}")
        
        if not self.source_audio_full_path or not self.source_audio_full_path.exists():
            log_callback("Original source audio not available for auto RIR marking.", True)
            raise Exception("Original source audio not available for auto RIR marking.")

        source_info = sf.info(str(self.source_audio_full_path)) # Get info without loading full audio
        duration_ms = int(source_info.duration * 1000)
        # self.source_audio_sr is used for loading segments
        segment_duration_ms = int(segment_duration_sec * 1000)
        print(f"DEBUG: AutoRIR: Original source audio loaded. Duration: {duration_ms}ms. Configured segment duration: {segment_duration_ms}ms")
        segment_definitions = [] 
        use_scene_cuts = self.use_scene_cuts_for_auto_rir_checkbox.isChecked() and self.scene_cuts_sec
        if use_scene_cuts:
            log_callback("Using detected scene cuts for RIR segmentation.")
            current_start_ms = 0
            for scene_cut_time_sec in self.scene_cuts_sec:
                scene_cut_ms = int(scene_cut_time_sec * 1000)
                if scene_cut_ms > current_start_ms and scene_cut_ms <= duration_ms:
                    segment_definitions.append((current_start_ms, scene_cut_ms)); current_start_ms = scene_cut_ms
            if current_start_ms < duration_ms: segment_definitions.append((current_start_ms, duration_ms))
        else:
            log_callback(f"Using fixed duration ({segment_duration_sec}s) for RIR segmentation.")
            num_fixed_segments = (duration_ms + segment_duration_ms - 1) // segment_duration_ms
            for i in range(num_fixed_segments):
                start_ms_fixed = i * segment_duration_ms; end_ms_fixed = min((i + 1) * segment_duration_ms, duration_ms)
                if end_ms_fixed > start_ms_fixed: segment_definitions.append((start_ms_fixed, end_ms_fixed))
        if not segment_definitions:
            log_callback("No segments defined for RIR marking (video might be too short or no scene cuts found/used).", True)
            progress_callback(100, "Auto RIR: No segments."); return []
        num_segments_to_process = len(segment_definitions) # type: ignore
        print(f"DEBUG: AutoRIR: Number of segments to process: {num_segments_to_process}")
        generated_markers = [] # Initialize generated_markers

        for i, (start_ms, end_ms) in enumerate(segment_definitions):
            print(f"DEBUG: AutoRIR: Processing segment {i+1}/{num_segments_to_process}: {start_ms}ms - {end_ms}ms")            
            if end_ms <= start_ms: continue 
            segment_progress_base = 10 + int(85 * (i / num_segments_to_process))
            progress_callback(segment_progress_base, f"AutoRIR: Processing segment {i+1}/{num_segments_to_process} ({start_ms/1000:.2f}s - {end_ms/1000:.2f}s)")
            
            chunk_filename = f"auto_s2r_chunk_{i}.wav" 
            chunk_path = TMP_DIR / chunk_filename
            
            log_callback(f"AutoRIR: Extracting segment {start_ms/1000.0:.2f}s - {end_ms/1000.0:.2f}s from source for RIR estimation.")
            audio_chunk_data, sr_chunk = librosa.load(
                str(self.source_audio_full_path),
                sr=self.source_audio_sr, # Load at the source audio's sample rate
                offset=start_ms / 1000.0,
                duration=(end_ms - start_ms) / 1000.0,
                mono=True # Speech2RIR generally expects mono
            )
            sf.write(str(chunk_path), audio_chunk_data, sr_chunk)
            print(f"DEBUG: AutoRIR: Extracted and saved chunk {i+1} (from original source) to {chunk_path}")

            # Estimate RT60 from the original source audio chunk first
            estimated_source_rt60 = None
            try:
                def rt60_source_chunk_progress_adapter(val, desc):
                    # Progress for this sub-step can be minor, or integrated into the overall segment progress
                    # For simplicity, we might not show detailed progress for this specific RT60 estimation here
                    pass 
                log_callback(f"AutoRIR Seg {i+1}: Estimating RT60 from original audio chunk {chunk_path.name}...")
                estimated_source_rt60 = self._perform_rt60_estimation_task(
                    str(chunk_path), progress_callback=rt60_source_chunk_progress_adapter, # Can be None if no GUI update needed
                    log_callback=log_callback, is_pre_segmented=True) # Crucial: is_pre_segmented=True
            except Exception as e_rt60_src:
                log_callback(f"Error estimating RT60 from source chunk {chunk_path.name}: {e_rt60_src}", True)

            try:
                def sub_progress_adapter(val, desc):
                    current_segment_progress_span = (85 / num_segments_to_process) * 0.9 
                    actual_progress = int(segment_progress_base + (val / 100.0 * current_segment_progress_span))
                    progress_callback(actual_progress, f"AutoRIR Seg {i+1}: {desc}")
                print(f"DEBUG: AutoRIR: Calling _perform_s2r_estimation_task for chunk {chunk_path}")
                estimated_chunk_rir_path = self._perform_s2r_estimation_task( # type: ignore
                    source_audio_path_str=str(chunk_path), progress_callback=sub_progress_adapter,
                    log_callback=log_callback, # Pass log_callback here
                    output_rir_basename=f"auto_s2r_rir_segment_{i}")
                print(f"DEBUG: AutoRIR: _perform_s2r_estimation_task returned RIR path: {estimated_chunk_rir_path} for chunk {i+1}")
                # Store the RIR path for application, and the RT60 estimated from the *source* chunk
                generated_markers.append((start_ms / 1000.0, end_ms / 1000.0, {'rir_path': estimated_chunk_rir_path, 'estimated_rt60': estimated_source_rt60}))
                print(f"DEBUG: AutoRIR: Added marker for chunk {i+1}")
            except Exception as e:
                log_callback(f"Error estimating RIR for segment {i+1} ({chunk_path.name}): {e}", True)
                print(f"DEBUG: AutoRIR: Exception estimating RIR for chunk {i+1}: {e}")
            print(f"DEBUG: AutoRIR: Chunk {i+1} (file: {chunk_path.name}) processing attempt finished. Cleaning up chunk file.")
            try:
                if chunk_path.exists(): chunk_path.unlink(); print(f"DEBUG: AutoRIR: Deleted temporary chunk file {chunk_path}")
            except Exception as e_del: print(f"DEBUG: AutoRIR: Error deleting chunk file {chunk_path}: {e_del}")
            if (i + 1) % 1 == 0: print(f"DEBUG: AutoRIR: Chunk {i+1} processing finished. Continuing loop."); sys.stdout.flush()
        print("DEBUG: AutoRIR: Finished processing all segments. Unloading Speech2RIR model...")
        QMetaObject.invokeMethod(self, "unload_s2r_model", Qt.QueuedConnection) # Ensure it runs on main thread if it touches GUI
        print("DEBUG: AutoRIR: Speech2RIR model unloaded.")
        progress_callback(100, "Auto RIR marker generation complete.")
        print(f"DEBUG: _perform_auto_rir_marking_task returning {len(generated_markers)} markers.")
        return generated_markers

    def _on_auto_rir_marking_finished(self, new_markers, message):

        self.set_ui_enabled_for_processing(True)
        if new_markers is not None: 
            self.swapped_audio_markers.clear() 
            self.swapped_audio_markers.extend(new_markers)
            self.swapped_audio_markers.sort(); self._update_rir_marker_display()
            self.log_status(f"Generated {len(new_markers)} RIR markers automatically. {message}")
            QMessageBox.information(self, "Auto RIR Markers Generated", f"Successfully generated {len(new_markers)} RIR markers.")
        else:
            self.log_status(f"Auto RIR marker generation failed: {message}", is_error=True)
            QMessageBox.critical(self, "Auto RIR Error", f"Could not generate RIR markers: {message}")
        self.current_worker = None

    def initiate_full_rir_preview_generation(self):
        if self.current_worker and self.current_worker.isRunning():
            QMessageBox.warning(self, "Busy", "A process is already running.")
            return
        if not self.swapped_audio_full_path or not self.swapped_audio_full_path.exists():
            QMessageBox.warning(self, "No Swapped Audio", "Please perform 'Swap Voice' first.")
            return

        use_s2r_mode_row2 = self.use_speech2rir_checkbox.isChecked()
        s2r_wet_mix_row2 = self.s2r_mix_slider.value()
        estimated_rir_path_row2 = self.estimated_rir_path
        rir_markers_row2 = list(self.swapped_audio_markers) # Use a copy

        if use_s2r_mode_row2 and (not estimated_rir_path_row2 or not estimated_rir_path_row2.exists()):
            QMessageBox.warning(self, "S2R RIR Missing", "Speech2RIR mode is selected, but no estimated RIR found. Please ensure RIR estimation was run.")
            return
        if not use_s2r_mode_row2 and not rir_markers_row2:
            QMessageBox.information(self, "No RIR Defined", "No manual RIR markers are defined, and Speech2RIR mode is not active. Preview will be the dry swapped audio.")
            # Load dry swapped audio into preview
            self.final_mix_preview_audio_path = self.swapped_audio_full_path
            y_preview, sr_preview = librosa.load(str(self.final_mix_preview_audio_path), sr=None, mono=True)
            self.final_mix_waveform_widget.load_audio(y_preview / np.max(np.abs(y_preview)) if np.max(np.abs(y_preview)) > 0 else y_preview, sr_preview, len(y_preview)/sr_preview)
            self.final_mix_waveform_widget.show()
            self.log_status("Loaded dry swapped audio into final mix preview as no RIRs from RIR marking step were active.")
            return

        self.log_status("Starting generation of full RIR preview...")
        self.progress_bar.setValue(0)
        self.set_ui_enabled_for_processing(False)
        self.output_video_label.setText("Generating Preview...")

        self.current_worker = WorkerThread(
            self._perform_full_rir_preview_task,
            str(self.swapped_audio_full_path),
            rir_markers_row2,
            use_s2r_mode_row2,
            s2r_wet_mix_row2,
            estimated_rir_path_row2
        )
        self.current_worker.progress.connect(self.handle_gui_progress_update)
        self.current_worker.finished.connect(self._on_full_rir_preview_finished)
        self.current_worker.log_message.connect(self._handle_worker_log_message)
        self.current_worker.start()

    def _on_full_rir_preview_finished(self, result_preview_audio_path, message):
        self.set_ui_enabled_for_processing(True)
        self.progress_bar.setValue(100 if result_preview_audio_path else 0)
        if result_preview_audio_path:
            self.final_mix_preview_audio_path = Path(result_preview_audio_path)
            self.log_status(f"Full RIR preview generated: {self.final_mix_preview_audio_path}")
            try:
                y_preview, sr_preview = librosa.load(str(self.final_mix_preview_audio_path), sr=None, mono=True)
                self.final_mix_waveform_widget.load_audio(y_preview / np.max(np.abs(y_preview)) if np.max(np.abs(y_preview)) > 0 else y_preview, sr_preview, len(y_preview)/sr_preview)
                self.final_mix_waveform_widget.show()
                self.output_video_label.setText(f"Final Mix Preview Ready: {self.final_mix_preview_audio_path.name}")
                QMessageBox.information(self, "Preview Ready", f"Full RIR preview loaded into final mix display.")
                self.step_tab_widget.setCurrentIndex(2) # Switch to "3. Final Mix and Preview" tab
            except Exception as e:
                self.log_status(f"Error loading generated preview audio: {e}", is_error=True)
                QMessageBox.critical(self, "Preview Load Error", f"Could not load generated preview audio: {e}")
        else:
            self.log_status(f"Full RIR preview generation failed or no RIRs applied: {message}", is_error=True)
            self.output_video_label.setText("Preview Generation Failed.")
            QMessageBox.critical(self, "Preview Generation Error", f"Could not generate full RIR preview: {message}")

        # Populate Row 3 segment list based on the RIRs used for the preview
        self.final_mix_segment_list.clear()
        self.final_mix_rir_settings.clear()
        
        # rir_markers_row2 was passed to the worker, it holds the config used for preview
        # If use_s2r_mode_row2 was true, rir_markers_row2 might be empty,
        # in which case we'd create a single segment for the whole audio with the global S2R RIR.
        # For simplicity, we'll assume rir_markers_row2 reflects the segmentation.
        # A more robust solution would pass back the exact RIR config applied from the worker.
        
        active_rir_config_for_preview = []
        if self.current_worker and len(self.current_worker.args) > 1:
            # args[1] is rir_markers_from_row2
            # args[2] is use_s2r_mode_from_row2
            # args[4] is estimated_s2r_rir_path_from_row2
            row2_markers = self.current_worker.args[1]
            row2_s2r_mode = self.current_worker.args[2]
            row2_s2r_path = self.current_worker.args[4]

            if row2_s2r_mode and row2_s2r_path:
                 # Create a single segment for the entire duration with S2R RIR info
                swapped_duration = librosa.get_duration(path=str(self.swapped_audio_full_path))
                active_rir_config_for_preview.append(
                    (0.0, swapped_duration, {'s2r_rir_path': str(row2_s2r_path), 'estimated_rt60': None}) # RT60 can be estimated later if needed
                )
            elif row2_markers:
                active_rir_config_for_preview = row2_markers
        
        for i, (start_s, end_s, rir_params) in enumerate(active_rir_config_for_preview):
            # Add default volume and pan
            adjusted_params = {'volume_db': 0.0, 'effects': {'Gain': {'enabled': False, 'gain_db': 0.0}}} # Add default for Gain
            if 'rir_path' in rir_params: # S2R from Row 2
                adjusted_params['s2r_rir_path'] = rir_params['rir_path']
                adjusted_params['estimated_rt60'] = rir_params.get('estimated_rt60')
            elif 'rt60' in rir_params and 'room_dim' in rir_params: # Simulated from Row 2
                adjusted_params['simulated_rir'] = {'rt60': rir_params['rt60'], 'room_dim': list(rir_params['room_dim'])}
            
            self.final_mix_rir_settings.append({'start_s': start_s, 'end_s': end_s, 'params': dict(adjusted_params)})

            # Construct item_text for the list, similar to _update_final_mix_segment_list
            item_text = f"Segment {i+1}: {start_s:.2f}s - {end_s:.2f}s"
            if 's2r_rir_path' in adjusted_params: # Check adjusted_params which now holds the S2R info
                est_rt60_str = ""
                if adjusted_params.get('estimated_rt60') is not None:
                    est_rt60_str = f", Est. RT60: {adjusted_params['estimated_rt60']:.3f}s"
                item_text += f" (S2R: {Path(adjusted_params['s2r_rir_path']).name}{est_rt60_str})"
            elif 'simulated_rir' in adjusted_params:
                sim_params = adjusted_params['simulated_rir']
                item_text += f" (Sim: RT60 {sim_params.get('rt60', 'N/A')})"
            else:
                item_text += " (No RIR)"
            item_text += f", Vol: {adjusted_params.get('volume_db', 0.0):.1f}dB" # Add volume info

            self.final_mix_segment_list.addItem(QListWidgetItem(item_text))

        self.current_worker = None
        # Auto-select the first segment in the final mix list if available
        if self.final_mix_segment_list.count() > 0:
            self.final_mix_segment_list.setCurrentRow(0)
            # _on_final_mix_segment_selected will be triggered, highlighting the segment

    def apply_and_preview_segment_effects(self):
        """
        Applies all configured effects (RIR, Volume, Pedalboard FX) from Row 3 
        to the selected segment of the original swapped audio and plays it.
        Updates the internal settings for this segment.
        """
        # Determine which audio to use as the "original" for this segment's processing.
        # If a full preview (self.final_mix_preview_audio_path) exists, it means Row 2 RIRs were applied.
        # However, for segment-specific adjustments in Row 3, we should always start from the "dry" swapped audio
        # from Row 1, and Row 3 controls will define *all* effects for that segment.
        base_audio_for_segment_processing = self.swapped_audio_full_path
        if not base_audio_for_segment_processing or not base_audio_for_segment_processing.exists():
            QMessageBox.warning(self, "No Swapped Audio", "Please perform 'Swap Voice' first.")
            return

        selected_items = self.final_mix_segment_list.selectedItems()
        if not selected_items:
            QMessageBox.warning(self, "No Segment Selected", "Please select a segment from the 'Segments for RIR Adjustment' list.")
            return
        
        selected_index = self.final_mix_segment_list.row(selected_items[0])
        if not (0 <= selected_index < len(self.final_mix_rir_settings)):
            self.log_status("Selected segment index out of bounds for final_mix_rir_settings.", is_error=True)
            return

        segment_info = self.final_mix_rir_settings[selected_index]
        start_s = segment_info['start_s']
        end_s = segment_info['end_s']
        current_segment_params = segment_info['params'] # This is a reference, so updates are direct

        # --- Update params from UI ---
        try: # Basic RIR & Volume Tab
            current_segment_params['simulated_rir'] = {
                'rt60': float(self.final_mix_rt60_edit.text()),
                'room_dim': [float(d.strip()) for d in self.final_mix_room_dim_edit.text().split(',')]
            }
            if len(current_segment_params['simulated_rir']['room_dim']) != 3 or \
               not all(d > 0 for d in current_segment_params['simulated_rir']['room_dim']):
                raise ValueError("Room dimensions must be three positive comma-separated numbers.")
            current_segment_params['volume_db'] = float(self.final_mix_volume_edit.text())
        except ValueError as e:
            QMessageBox.critical(self, "Invalid Basic Parameters", f"Please enter valid RIR/Volume parameters: {e}")
            return

        # Gain Effect Tab
        if 'effects' not in current_segment_params: current_segment_params['effects'] = {}
        current_segment_params['effects']['Gain'] = {
            'enabled': self.final_fx_gain_enable_checkbox.isChecked(),
            'gain_db': float(self.final_fx_gain_slider.value())
        }
        # ... Update for other effects from their UI elements ...

        self.log_status(f"Updated effects for segment {selected_index+1}. Previewing...")

        # --- Generate Audio for this Segment with all effects ---
        try:
            # Load original swapped audio for this segment
            y_segment_original, sr = librosa.load(
                str(base_audio_for_segment_processing), 
                sr=None, 
                offset=start_s, 
                duration=(end_s - start_s)
            )

            processed_segment_data = self._apply_effects_to_segment_audio_data(
                y_segment_original, sr, current_segment_params
            )

            temp_preview_segment_path = TMP_DIR / f"preview_effect_segment_{selected_index}.wav"
            sf.write(str(temp_preview_segment_path), processed_segment_data, sr)

            # Play just this segment
            self.player.stop()
            self.player.setSource(QUrl.fromLocalFile(str(temp_preview_segment_path)))
            self.player.play()
            self.log_status(f"Playing preview of segment {selected_index+1} with applied effects.")

            # Update list item text (optional, can get long)
            # self.final_mix_segment_list.item(selected_index).setText(...)

        except Exception as e:
            self.log_status(f"Error generating/playing effect preview for segment: {e}", is_error=True)
            QMessageBox.critical(self, "Preview Error", f"Could not generate/play effect preview: {e}")

    def _update_current_segment_effects_params(self):
        """Called when an effect's UI control changes, to update the stored params."""
        selected_items = self.final_mix_segment_list.selectedItems()
        if not selected_items: return
        selected_index = self.final_mix_segment_list.row(selected_items[0])
        if not (0 <= selected_index < len(self.final_mix_rir_settings)): return

        params_ref = self.final_mix_rir_settings[selected_index]['params']
        if 'effects' not in params_ref: params_ref['effects'] = {}

        # Update Gain from UI
        params_ref['effects']['Gain'] = {
            'enabled': self.final_fx_gain_enable_checkbox.isChecked(),
            'gain_db': float(self.final_fx_gain_slider.value())
        }
        # Update other effects similarly
        # self.log_status(f"Effect param changed for segment {selected_index + 1}", is_error=False) # Optional: for debugging

    def _on_final_mix_segment_selected(self):
        # This method is called when a segment in the Row 3 list is selected.
        # It populates the RT60, Room Dimension, Volume, and other effect controls
        # based on the stored settings for that segment.
        selected_items = self.final_mix_segment_list.selectedItems()
        if not selected_items:
            # No item selected, clear selection on waveform
            self.final_mix_waveform_widget.clear_selection()
            return

        selected_index = self.final_mix_segment_list.row(selected_items[0])
        if 0 <= selected_index < len(self.final_mix_rir_settings):
            segment_data = self.final_mix_rir_settings[selected_index]
            start_s = segment_data['start_s']
            end_s = segment_data['end_s']

            # Update waveform selection to highlight the segment
            self.final_mix_waveform_widget.selection_start_sec = start_s
            self.final_mix_waveform_widget.selection_end_sec = end_s
            self.final_mix_waveform_widget.update() # Trigger repaint
            params = segment_data.get('params', {})
            

            # RT60 and Room Dim are always editable in Row 3.
            # If the user edits them, it implies they want a simulated RIR for this segment.
            self.final_mix_rt60_edit.setReadOnly(False)
            self.final_mix_room_dim_edit.setReadOnly(False)

            is_s2r_segment = 's2r_rir_path' in params # Check if it *was* an S2R segment

            if is_s2r_segment:
                # If it was S2R, populate RT60 with its estimated value.
                # Populate Room Dim with global defaults, as editing these implies switching to simulated.
                self.final_mix_rt60_edit.setText(str(params.get('estimated_rt60', '0.5'))) # Default to 0.5 if not found
                self.final_mix_room_dim_edit.setText(self.room_dim_edit.text()) # Use global room dims as default
            else: # Simulated RIR or no RIR
                sim_rir_params = params.get('simulated_rir', {'rt60': 0.5, 'room_dim': [6,4,2.8]})
                self.final_mix_rt60_edit.setText(str(sim_rir_params.get('rt60', '0.5')))
                room_dim = sim_rir_params.get('room_dim', [6,4,2.8])
                self.final_mix_room_dim_edit.setText(','.join(map(str, room_dim)))
            
            self.final_mix_volume_edit.setText(str(params.get('volume_db', '0.0')))

            # Effects
            effects_params = params.get('effects', {})
            # Gain
            gain_p = effects_params.get('Gain', {'enabled': False, 'gain_db': 0.0})
            # Block signals temporarily to prevent _update_current_segment_effects_params from firing prematurely
            self.final_fx_gain_enable_checkbox.blockSignals(True)
            self.final_fx_gain_slider.blockSignals(True)
            self.final_fx_gain_enable_checkbox.setChecked(gain_p.get('enabled', False)) # SetValue may trigger if not blocked
            self.final_fx_gain_slider.setValue(int(float(gain_p.get('gain_db', 0.0)))) # Ensure float then int
            self.final_fx_gain_label.setText(f"{int(float(gain_p.get('gain_db', 0.0)))} dB")
            self.final_fx_gain_enable_checkbox.blockSignals(False)
            self.final_fx_gain_slider.blockSignals(False)
            # ... Load other effect UI elements ...

        else: # Should not happen if list and data are in sync
            self.log_status("Error: Selected segment in final mix list out of sync with data.", is_error=True)
            self.final_mix_waveform_widget.clear_selection()

    def _perform_full_rir_preview_task(self, swapped_audio_path_str,
                                       # rir_markers_from_row2 are the markers active in Row 2 when "Generate Full Preview" was clicked
                                       # These define the *initial* state of the segments for Row 3.
                                       # Row 3 then allows further per-segment modification of these initial RIRs/effects.
                                       # The preview generated by this task reflects the Row 2 settings.
                                       rir_markers_from_row2, use_s2r_mode_from_row2,
                                       s2r_wet_mix_percentage_from_row2, estimated_s2r_rir_path_from_row2,
                                       progress_callback, log_callback):
        
        swapped_audio_path = Path(swapped_audio_path_str)
        preview_output_audio_path = TMP_DIR / "row3_full_preview_temp.wav"

        progress_callback(0, "Initializing RIR preview generation...")

        try:
            swapped_audio_info = sf.info(str(swapped_audio_path))
            swapped_audio_sr = swapped_audio_info.samplerate
            swapped_audio_duration_ms = int(swapped_audio_info.duration * 1000)
        except Exception as e:
            raise Exception(f"Failed to get info for swapped audio {swapped_audio_path}: {e}")

        if use_s2r_mode_from_row2:
            if estimated_s2r_rir_path_from_row2 and estimated_s2r_rir_path_from_row2.exists():
                rir_waveform_data, rir_sr_data = librosa.load(str(estimated_s2r_rir_path_from_row2), sr=None, mono=True)
                log_callback(f"Preview: Loaded estimated RIR. SR: {rir_sr_data}")
                
                progress_callback(10, "Preview: Applying global Speech2RIR estimated RIR...")
                temp_reverbed_for_preview = TMP_DIR / "preview_reverbed_s2r.wav"
                self.apply_reverb_with_rir_waveform_and_mix(
                    str(swapped_audio_path), str(temp_reverbed_for_preview),
                    str(preview_output_audio_path),
                    target_sr=swapped_audio_sr,
                    rir_waveform=rir_waveform_data, rir_sr=rir_sr_data,
                    wet_mix_percentage=s2r_wet_mix_percentage_from_row2,
                    progress_callback=lambda v,d: progress_callback(10 + int(v*0.8), f"Preview S2R: {d}")
                )
                progress_callback(90, "Preview: Global Speech2RIR RIR application complete.")
            else:
                raise Exception("Preview: S2R mode selected, but estimated RIR path is invalid.")
        elif rir_markers_from_row2:
            progress_callback(10, "Preview: Applying RIRs based on Row 2 markers for full preview...")
            # This part is similar to _perform_final_mixdown_task's marker processing
            # We can simplify by calling a shared helper or duplicating the logic carefully
            # For now, let's assume a simplified version of the marker logic from _perform_final_mixdown_task
            # is applied here, saving to preview_output_audio_path.
            # This would involve iterating rir_markers_from_row2, applying reverb to segments, and stitching.
            # This is a placeholder for the detailed segment processing logic:
            log_callback("Preview: Marker-based RIR application for preview (simplified for now).")
            # In a full implementation, you'd iterate through rir_markers_from_row2,
            # call apply_specific_rir_to_segment or apply_reverb_to_audio_segment for each,
            # and then stitch them together.
            # For this example, we'll just copy the swapped audio if markers exist but processing is complex.
            # A more complete solution would reuse the logic from _perform_final_mixdown_task's RIR loop.
            
            # --- Actual marker processing for preview ---
            final_output_samples_list = []
            last_rir_processed_end_time_ms = 0
            total_rir_segments = len(rir_markers_from_row2)

            for i, (start_s, end_s, rir_params) in enumerate(rir_markers_from_row2):
                segment_progress_percentage = (80.0 / total_rir_segments) if total_rir_segments > 0 else 80.0
                segment_progress_base = 10 + int(segment_progress_percentage * i)
                progress_callback(segment_progress_base, f"Preview RIR Seg {i+1}/{total_rir_segments}")

                start_ms = int(start_s * 1000); end_ms = int(end_s * 1000)
                if start_ms > last_rir_processed_end_time_ms:
                    y_interlude, _ = librosa.load(str(swapped_audio_path), sr=swapped_audio_sr, offset=last_rir_processed_end_time_ms/1000.0, duration=(start_ms - last_rir_processed_end_time_ms)/1000.0)
                    final_output_samples_list.append(y_interlude)

                segment_to_reverb_path = TMP_DIR / f"preview_marker_segment_{i}_input.wav"
                y_segment_to_reverb, _ = librosa.load(str(swapped_audio_path), sr=swapped_audio_sr, offset=start_s, duration=(end_s-start_s))
                sf.write(str(segment_to_reverb_path), y_segment_to_reverb, swapped_audio_sr)
                reverbed_segment_output_path = TMP_DIR / f"preview_marker_segment_{i}_output.wav"

                if 'rir_path' in rir_params: # Auto S2R per segment
                    # Simplified: Assume rir_path is valid and apply it (full logic in final mixdown)
                    # For preview, we might just copy or apply a generic effect if too complex
                    shutil.copy(str(segment_to_reverb_path), str(reverbed_segment_output_path)) # Placeholder
                elif 'rt60' in rir_params and 'room_dim' in rir_params: # Simulated
                    apply_reverb_to_audio_segment(str(segment_to_reverb_path), str(reverbed_segment_output_path), target_sr=swapped_audio_sr, room_dim_list=rir_params['room_dim'], rt60_target=rir_params['rt60'], progress_callback=lambda v,d: progress_callback(int(segment_progress_base + (v/100.0 * segment_progress_percentage * 0.9)),d))
                else: # No RIR for this segment
                    shutil.copy(str(segment_to_reverb_path), str(reverbed_segment_output_path))
                
                y_reverbed_segment, _ = librosa.load(str(reverbed_segment_output_path), sr=swapped_audio_sr)
                final_output_samples_list.append(y_reverbed_segment)
                last_rir_processed_end_time_ms = end_ms
            
            if last_rir_processed_end_time_ms < swapped_audio_duration_ms:
                y_tail, _ = librosa.load(str(swapped_audio_path), sr=swapped_audio_sr, offset=last_rir_processed_end_time_ms/1000.0)
                final_output_samples_list.append(y_tail)
            final_audio_stitched_np = np.concatenate(final_output_samples_list) if final_output_samples_list else np.array([])
            sf.write(str(preview_output_audio_path), final_audio_stitched_np, swapped_audio_sr)
            
            progress_callback(90, "Preview: Marker-based RIR (placeholder) complete.")
        else:
            # Should have been caught by the initiator, but as a fallback:
            log_callback("Preview: No RIR to apply. Using swapped audio as is for preview.")
            shutil.copy(str(swapped_audio_path), str(preview_output_audio_path))
            progress_callback(90, "Preview: No RIR to apply.")

        progress_callback(100, "RIR Preview generation complete.")
        return str(preview_output_audio_path)

    def _apply_effects_to_segment_audio_data(self, audio_data_np, sr, segment_params_dict):
        """
        Applies all configured effects from segment_params_dict to audio_data_np (which should be the original/dry audio for the segment).
        This method is responsible for applying RIR (Convolution), Volume (as Gain), and other pedalboard effects.
        segment_params_dict is an item from self.final_mix_rir_settings[...]['params']
        """
        board_fx = []

        # 1. Volume (applied as a Gain pedalboard effect for consistency)
        volume_db = segment_params_dict.get('volume_db', 0.0)
        # We can choose to apply this as the first or last Gain, or integrate it.
        # For simplicity, let's assume it's an initial gain stage if non-zero.
        if volume_db != 0.0: # Only add if non-zero to avoid unnecessary processing
             board_fx.append(Gain(gain_db=volume_db)) # Use direct Gain import

        # 2. RIR (Convolution)
        sim_rir_config = segment_params_dict.get('simulated_rir')
        s2r_rir_path_str = segment_params_dict.get('s2r_rir_path')

        if s2r_rir_path_str: # Prefer S2R RIR if path is provided
            s2r_rir_p = Path(s2r_rir_path_str)
            if s2r_rir_p.exists():
                try:
                    rir_wf, rir_s_original = librosa.load(str(s2r_rir_p), sr=None, mono=True)
                    board_fx.append(Convolution(rir_wf, sample_rate=rir_s_original)) # Pedalboard handles resampling
                except Exception as e:
                    self.log_status(f"Error loading S2R RIR {s2r_rir_p} for effect application: {e}", is_error=True)
            else:
                self.log_status(f"S2R RIR path not found: {s2r_rir_p}", is_error=True)
        elif sim_rir_config: # Fallback to simulated RIR
            rt60 = sim_rir_config.get('rt60', 0.5)
            room_dim = sim_rir_config.get('room_dim', [6,4,2.8])
            try: # Simplified RIR generation from apply_reverb_to_audio_segment
                V = np.prod(room_dim); Sx, Sy, Sz = room_dim; S_total = 2 * (Sx*Sy + Sx*Sz + Sy*Sz)
                alpha = np.clip(0.161 * V / (rt60 * S_total) if S_total > 0 and rt60 > 0 else 0.5, 0.01, 0.99)
                room = pra.ShoeBox(room_dim, fs=sr, materials=pra.Material(alpha), max_order=3) # max_order=3 for speed
                room.add_source([room_dim[0]/2, room_dim[1]/2, room_dim[2]/2])
                mic_pos = [min(max(room.source_locs[0][i] + (0.1 if i==0 else 0), 0.01), room_dim[i]-0.01) for i in range(3)]
                room.add_microphone_array(pra.MicrophoneArray(np.array([mic_pos]).T, fs=sr))
                room.compute_rir()
                if room.rir and room.rir[0] and room.rir[0][0] is not None and len(room.rir[0][0]) > 0:
                    board_fx.append(Convolution(room.rir[0][0].astype(np.float32), sample_rate=sr))
            except Exception as e:
                self.log_status(f"Error generating simulated RIR for effect application: {e}", is_error=True)

        # 3. Other Pedalboard Effects from the 'effects' sub-dictionary
        effects_config = segment_params_dict.get('effects', {})
        
        gain_effect_config = effects_config.get('Gain', {'enabled': False, 'gain_db': 0.0})
        if gain_effect_config.get('enabled', False):
            board_fx.append(Gain(gain_db=float(gain_effect_config.get('gain_db', 0.0)))) # Use direct Gain import
        
        # ... Add logic for Compressor, EQ, Reverb, etc. based on effects_config ...

        if not board_fx: return audio_data_np # No effects to apply

        board = Pedalboard(board_fx) # Constructor does not take sample_rate

        processed_audio = board(audio_data_np.T if audio_data_np.ndim > 1 else audio_data_np, sample_rate=sr) # Pass sample_rate to process method
        return processed_audio.T if audio_data_np.ndim > 1 and processed_audio.ndim > 1 else processed_audio

    def apply_reverb_with_rir_waveform_and_mix(self,
       input_audio_path_str: str, temp_reverbed_audio_path_str: str, 
       final_output_audio_path_str: str, target_sr: int, 
       rir_waveform: np.ndarray, rir_sr: int,    
       wet_mix_percentage: int, progress_callback=None):

        if progress_callback: progress_callback(0, f"Applying estimated RIR: Loading audio {Path(input_audio_path_str).name}")
        segment_audio, segment_fs_orig = librosa.load(input_audio_path_str, sr=None, mono=False) 
        if segment_fs_orig != target_sr: segment_audio = librosa.resample(segment_audio, orig_sr=segment_fs_orig, target_sr=target_sr, res_type='kaiser_fast')
        if segment_audio.ndim == 1: segment_audio_pb = segment_audio 
        elif segment_audio.shape[0] == 2 and segment_audio.shape[1] > 2: segment_audio_pb = segment_audio.T 
        else: segment_audio_pb = segment_audio
        rir_waveform_pb = rir_waveform.astype(np.float32)
        if progress_callback: progress_callback(50, "Applying estimated RIR convolution via Pedalboard...")
        if rir_sr != target_sr:
            self.log_status(f"Manually resampling RIR from {rir_sr}Hz to {target_sr}Hz before Pedalboard.")
            rir_waveform_resampled = librosa.resample(rir_waveform.astype(np.float32), orig_sr=rir_sr, target_sr=target_sr, res_type='kaiser_best') 
            rir_waveform_for_pedalboard = rir_waveform_resampled; convolution_sr_param = target_sr
        else: rir_waveform_for_pedalboard = rir_waveform.astype(np.float32); convolution_sr_param = rir_sr
        convolution = Convolution(rir_waveform_for_pedalboard, sample_rate=convolution_sr_param); board = Pedalboard([convolution])       
        output_audio_convolved = board(segment_audio_pb, sample_rate=target_sr) 
        sf.write(temp_reverbed_audio_path_str, output_audio_convolved, target_sr)
        progress_callback(80, "Mixing dry and wet signals...")
        wet_audio, _ = librosa.load(temp_reverbed_audio_path_str, sr=target_sr, mono=False)
        dry_audio, _ = librosa.load(input_audio_path_str, sr=target_sr, mono=False)
        if dry_audio.shape[0] == 1 and wet_audio.shape[0] > 1: 
            dry_audio = np.repeat(dry_audio, wet_audio.shape[0], axis=0) 
        elif dry_audio.shape[0] > 1 and wet_audio.shape[0] == 1: 
            wet_audio = np.repeat(wet_audio, dry_audio.shape[0], axis=0) 
        elif dry_audio.shape[0] != wet_audio.shape[0] and dry_audio.shape[0] > 0 and wet_audio.shape[0] > 0:
            self.log_status(f"Warning: Dry audio has {dry_audio.shape[0]} channels, Wet audio has {wet_audio.shape[0]} channels. Converting both to mono for mixing.", is_error=True)
            dry_audio = librosa.to_mono(dry_audio); wet_audio = librosa.to_mono(wet_audio) 
        len_dry = dry_audio.shape[1] if dry_audio.ndim == 2 else dry_audio.shape[0]
        len_wet = wet_audio.shape[1] if wet_audio.ndim == 2 else wet_audio.shape[0]
        min_len = min(len_dry, len_wet)
        if dry_audio.ndim == 2: 
            dry_audio = dry_audio[:, :min_len]; 
        else: 
            dry_audio = dry_audio[:min_len]
        
        if wet_audio.ndim == 2: 
            wet_audio = wet_audio[:, :min_len]; 
        else: 
            wet_audio = wet_audio[:min_len]

        wet_gain = wet_mix_percentage / 100.0; dry_gain = 1.0 - wet_gain
        mixed_audio = dry_gain * dry_audio + wet_gain * wet_audio
        if mixed_audio.ndim == 2 and mixed_audio.shape[0] < mixed_audio.shape[1]: 
            mixed_audio_to_write = mixed_audio.T
        else: 
            mixed_audio_to_write = mixed_audio
        sf.write(final_output_audio_path_str, mixed_audio_to_write, target_sr)
        if progress_callback: 
            progress_callback(100, "Estimated RIR application and mixing complete.")
        print(f"Audio with estimated RIR saved to {final_output_audio_path_str}")

    def apply_specific_rir_to_segment(self,
                                      input_audio_path_str: str, output_audio_path_str: str,
                                      target_sr: int, rir_waveform: np.ndarray, rir_sr: int,
                                      progress_callback=None):

        if progress_callback: 
            progress_callback(0, f"Applying specific RIR to {Path(input_audio_path_str).name}")
        segment_audio, segment_fs_orig = librosa.load(input_audio_path_str, sr=None, mono=False)
        if segment_fs_orig != target_sr: 
            segment_audio = librosa.resample(segment_audio, orig_sr=segment_fs_orig, target_sr=target_sr, res_type='kaiser_fast')
        if segment_audio.ndim == 1: 
            segment_audio_pb = segment_audio
        elif segment_audio.shape[0] == 2 and segment_audio.shape[1] > 2: 
            segment_audio_pb = segment_audio.T
        else: 
            segment_audio_pb = segment_audio
        rir_waveform_pb = rir_waveform.astype(np.float32)
        if progress_callback:
            progress_callback(50, "Convolving segment with its specific RIR...")
        convolution = Convolution(rir_waveform_pb, sample_rate=rir_sr)
        board = Pedalboard([convolution])
        output_audio_convolved = board(segment_audio_pb, sample_rate=target_sr)
        sf.write(output_audio_path_str, output_audio_convolved, target_sr) # type: ignore
        if progress_callback: 
            progress_callback(100, "Specific RIR application complete.")
        self.log_status(f"Segment processed with specific RIR: {output_audio_path_str}")

    def _on_final_mix_waveform_marker_clicked(self, clicked_start_s, clicked_end_s):
        for i, seg_info in enumerate(self.final_mix_rir_settings):
            # Allow for small float precision issues
            if abs(seg_info['start_s'] - clicked_start_s) < 0.01 and \
               abs(seg_info['end_s'] - clicked_end_s) < 0.01:
                if self.final_mix_segment_list.currentRow() != i:
                    self.final_mix_segment_list.setCurrentRow(i)
                # _on_final_mix_segment_selected will be called automatically by setCurrentRow
                # or if it was already selected, we ensure its params are loaded.
                self._on_final_mix_segment_selected() 
                self.log_status(f"Selected segment {i+1} via waveform click.")
                return
        self.log_status(f"Clicked waveform region ({clicked_start_s:.2f}-{clicked_end_s:.2f}) does not match a defined segment in Row 3.", is_error=False)
# --- Main Application Execution ---
if __name__ == "__main__":
    app = QApplication(sys.argv)
    main_window = MainWindow()
    main_window.show()
    sys.exit(app.exec())