How would I add diarisation to this whisperx script?

#!/usr/bin/env python3
import os
import sys
import json
import glob
import re
import time
import threading
import subprocess
from datetime import datetime, timedelta

try:
    import whisperx
    import torch
    import gc
    WHISPERX_AVAILABLE = True
except ImportError:
    WHISPERX_AVAILABLE = False
    print("ERROR: whisperx or torch library not found. This script requires them.")
    sys.exit(1)

# =============================================================================
# WHISPERX EMPTY TOKENS FIX - ADD THIS RIGHT AFTER IMPORTS
# =============================================================================
if WHISPERX_AVAILABLE:
    import whisperx.alignment
    
    # Store the original functions
    original_get_trellis = whisperx.alignment.get_trellis
    original_get_wildcard_emission = whisperx.alignment.get_wildcard_emission
    
    def patched_get_trellis(emission, tokens, blank_id):
        """Patched version that handles empty tokens"""
        # Fix empty tokens issue
        if isinstance(tokens, list):
            tokens = torch.tensor([0] if len(tokens) == 0 else tokens)
        elif isinstance(tokens, torch.Tensor):
            if tokens.numel() == 0:  # Empty tensor
                tokens = torch.tensor([0])
        else:
            tokens = torch.tensor([0] if len(tokens) == 0 else tokens)
        
        # Ensure tokens are the right type (long integers)
        if isinstance(tokens, torch.Tensor):
            tokens = tokens.long()
        
        return original_get_trellis(emission, tokens, blank_id)
    
    def patched_get_wildcard_emission(frame_emission, tokens, blank_id):
        """Patched version that handles empty tokens and ensures correct tensor type"""
        # Fix empty tokens issue
        if isinstance(tokens, list):
            tokens = torch.tensor([0] if len(tokens) == 0 else tokens)
        elif isinstance(tokens, torch.Tensor):
            if tokens.numel() == 0:  # Empty tensor
                tokens = torch.tensor([0])
            # Ensure correct tensor type for indexing
            tokens = tokens.long()
        else:
            tokens = torch.tensor([0] if len(tokens) == 0 else tokens).long()
        
        return original_get_wildcard_emission(frame_emission, tokens, blank_id)
    
    # Apply the patches
    whisperx.alignment.get_trellis = patched_get_trellis
    whisperx.alignment.get_wildcard_emission = patched_get_wildcard_emission
    
    print("✅ Applied WhisperX empty tokens fix patch")

# =============================================================================
# PROGRESS MONITORING CONFIGURATION
# =============================================================================
CONFIG = {
    # WhisperX Model Settings
    'model_size': 'large-v3-turbo',
    'batch_size': 32,#seems to be ignored due to not using gpu
    
    # Progress Monitoring
    'show_progress': True,
    'progress_update_interval': 30,   # Update every 30 seconds
    'estimate_total_time': True,      # Try to estimate total processing time
    'show_memory_usage': True,        # Show memory usage in progress
    
    # CPU Optimization
    'cpu_threads': 4,
    'compute_type': 'int8',
    
    # Subtitle Settings
    'target_subtitle_duration': 7.0,
    'min_split_gap': 0.3,
    
    # Output Settings
    'output_txt_filename': 'transcriptX.txt',
    'audio_extensions': ['*.wav', '*.mp3', '*.m4a', '*.aac', '*.flac', '*.ogg'],
}

class ProgressMonitor:
    """Monitor and display progress during WhisperX processing."""
    
    def __init__(self, total_duration=None):
        self.total_duration = total_duration
        self.start_time = time.time()
        self.current_phase = "Starting"
        self.phase_start_time = time.time()
        self.is_running = True
        self.transcription_progress = 0
        self.alignment_progress = 0
        self.estimated_total_time = None
        
        # Memory monitoring
        try:
            import psutil
            self.process = psutil.Process()
            self.memory_available = True
        except ImportError:
            self.memory_available = False
    
    def start_monitoring(self):
        """Start the progress monitoring thread."""
        if CONFIG['show_progress']:
            self.monitor_thread = threading.Thread(target=self._monitor_loop, daemon=True)
            self.monitor_thread.start()
    
    def stop_monitoring(self):
        """Stop the progress monitoring."""
        self.is_running = False
    
    def set_phase(self, phase_name, estimated_duration=None):
        """Set the current processing phase."""
        self.current_phase = phase_name
        self.phase_start_time = time.time()
        if estimated_duration:
            self.estimated_total_time = estimated_duration
        print(f"\n{'='*50}")
        print(f"PHASE: {phase_name}")
        print(f"{'='*50}")
    
    def update_transcription_progress(self, progress_pct):
        """Update transcription progress (0-100)."""
        self.transcription_progress = progress_pct
    
    def update_alignment_progress(self, progress_pct):
        """Update alignment progress (0-100)."""
        self.alignment_progress = progress_pct
    
    def _get_memory_info(self):
        """Get current memory usage."""
        if not self.memory_available:
            return ""
        
        try:
            memory_info = self.process.memory_info()
            memory_mb = memory_info.rss / (1024 * 1024)
            
            # System memory
            import psutil
            system_memory = psutil.virtual_memory()
            system_mb = system_memory.available / (1024 * 1024)
            
            return f" | Memory: {memory_mb:.0f}MB used, {system_mb:.0f}MB free"
        except:
            return ""
    
    def _monitor_loop(self):
        """Main monitoring loop that runs in a separate thread."""
        while self.is_running:
            elapsed = time.time() - self.start_time
            phase_elapsed = time.time() - self.phase_start_time
            
            # Create progress message
            progress_msg = f"[{elapsed/60:.1f}m] {self.current_phase}"
            
            if self.current_phase == "Transcribing":
                if self.transcription_progress > 0:
                    progress_msg += f" ({self.transcription_progress:.1f}%)"
                
                # Estimate remaining time based on progress
                if self.transcription_progress > 5 and self.total_duration:
                    estimated_transcription_time = (phase_elapsed / self.transcription_progress) * 100
                    remaining = estimated_transcription_time - phase_elapsed
                    progress_msg += f" | Est. {remaining/60:.1f}m remaining"
            
            elif self.current_phase == "Aligning":
                if self.alignment_progress > 0:
                    progress_msg += f" ({self.alignment_progress:.1f}%)"
            
            # Add duration info if available
            if self.total_duration:
                progress_msg += f" | Audio: {self.total_duration/60:.1f}m"
            
            # Add memory info
            progress_msg += self._get_memory_info()
            
            print(f"\r{progress_msg}", end="", flush=True)
            
            time.sleep(CONFIG['progress_update_interval'])
        
        print()  # New line when monitoring stops

def print_config():
    print("=" * 60)
    print("WHISPERX WITH PROGRESS MONITORING")
    print("=" * 60)
    for key, value in CONFIG.items():
        print(f"{key}: {value}")
    print("=" * 60)

def find_audio_file():
    """Find the first audio file in the current directory."""
    audio_files = []
    for ext in CONFIG['audio_extensions']:
        audio_files.extend(glob.glob(ext))
    
    if not audio_files:
        print("No audio files found in the current directory.")
        sys.exit(1)
    
    if len(audio_files) > 1:
        print(f"Warning: Multiple audio files found. Using '{audio_files[0]}'.")
    
    return audio_files[0]

def get_audio_duration(audio_path):
    """Get audio duration for progress estimation."""
    try:
        cmd = ['ffprobe', '-v', 'quiet', '-print_format', 'json', '-show_entries', 
               'format=duration', audio_path]
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        info = json.loads(result.stdout)
        duration = float(info['format']['duration'])
        return duration
    except Exception:
        # Fallback: estimate from file
        try:
            audio = whisperx.load_audio(audio_path)
            duration = len(audio) / 16000
            return duration
        except Exception:
            return None

class ProgressCallback:
    """Callback class to track WhisperX progress."""
    
    def __init__(self, monitor, total_segments_estimate=None):
        self.monitor = monitor
        self.total_segments_estimate = total_segments_estimate
        self.processed_segments = 0
    
    def __call__(self, **kwargs):
        """Called by WhisperX during processing."""
        # This is a simplified progress callback
        # WhisperX doesn't provide detailed progress info, so we estimate
        if 'segment' in kwargs:
            self.processed_segments += 1
            if self.total_segments_estimate:
                progress = (self.processed_segments / self.total_segments_estimate) * 100
                self.monitor.update_transcription_progress(min(progress, 99))

def run_whisperx_with_progress(audio_filepath, monitor):
    """Run WhisperX transcription with progress monitoring."""
    
    # Setup CPU optimization
    torch.set_num_threads(CONFIG['cpu_threads'])
    device = "cpu"
    
    # Phase 1: Load model
    monitor.set_phase("Loading Model")
    try:
        model = whisperx.load_model(
            CONFIG['model_size'], 
            device, 
            compute_type=CONFIG['compute_type']
        )
        print(f"Model '{CONFIG['model_size']}' loaded successfully")
    except Exception as e:
        print(f"Error loading model: {e}")
        try:
            model = whisperx.load_model(CONFIG['model_size'], device)
            print("Model loaded without compute_type")
        except Exception as e2:
            raise e2
    
    # Phase 2: Load audio
    monitor.set_phase("Loading Audio")
    audio = whisperx.load_audio(audio_filepath)
    audio_duration = len(audio) / 16000
    print(f"Audio loaded: {audio_duration:.1f} seconds ({audio_duration/60:.1f} minutes)")
    
    # Update monitor with actual duration
    monitor.total_duration = audio_duration
    
    # Phase 3: Transcription
    monitor.set_phase("Transcribing")
    
    # Estimate transcription time based on model size and audio duration
    # These are rough estimates based on typical performance
    if CONFIG['model_size'] == 'large-v3-turbo':
        estimated_ratio = 0.25  # Usually around 0.25x on decent CPU
    elif 'large' in CONFIG['model_size']:
        estimated_ratio = 0.3
    elif 'medium' in CONFIG['model_size']:
        estimated_ratio = 0.15
    else:
        estimated_ratio = 0.2
    
    estimated_transcription_time = audio_duration * estimated_ratio
    print(f"Estimated transcription time: {estimated_transcription_time/60:.1f} minutes")
    
    # Start transcription with progress tracking
    transcription_start = time.time()
    
    # We'll simulate progress by checking elapsed time vs estimated time
    def update_transcription_progress():
        while monitor.current_phase == "Transcribing" and monitor.is_running:
            elapsed = time.time() - transcription_start
            estimated_progress = min((elapsed / estimated_transcription_time) * 100, 95)
            monitor.update_transcription_progress(estimated_progress)
            time.sleep(5)
    
    # Start progress updater
    progress_thread = threading.Thread(target=update_transcription_progress, daemon=True)
    progress_thread.start()
    
    # Run transcription
    result = model.transcribe(
        audio, 
        batch_size=CONFIG['batch_size'],
        verbose=False  # Reduce verbose output for cleaner progress display
    )
    
    transcription_time = time.time() - transcription_start
    monitor.update_transcription_progress(100)
    print(f"\nTranscription completed in {transcription_time/60:.1f} minutes")
    
    # Cleanup transcription model
    monitor.set_phase("Cleaning up transcription model")
    del model
    gc.collect()
    
    # Phase 4: Alignment
    monitor.set_phase("Loading Alignment Model")
    language_code = result.get("language", "en")
    print(f"Detected language: {language_code}")
    
    try:
        model_a, metadata = whisperx.load_align_model(
            language_code=language_code, 
            device=device
        )
        print("Alignment model loaded")
    except Exception as e:
        print(f"Error loading alignment model: {e}")
        print("Proceeding without alignment.")
        monitor.stop_monitoring()
        for seg in result.get("segments", []):
            seg['words'] = [{'word': w, 'start': None, 'end': None, 'score': None} 
                           for w in seg.get('text', '').split()]
        return result
    
    # Phase 5: Alignment processing
    monitor.set_phase("Aligning")
    
    # Estimate alignment time (usually faster than transcription)
    estimated_alignment_time = transcription_time * 0.6  # Usually ~60% of transcription time
    print(f"Estimated alignment time: {estimated_alignment_time/60:.1f} minutes")
    
    def update_alignment_progress():
        alignment_start = time.time()
        while monitor.current_phase == "Aligning" and monitor.is_running:
            elapsed = time.time() - alignment_start
            estimated_progress = min((elapsed / estimated_alignment_time) * 100, 95)
            monitor.update_alignment_progress(estimated_progress)
            time.sleep(5)
    
    # Start alignment progress updater
    alignment_progress_thread = threading.Thread(target=update_alignment_progress, daemon=True)
    alignment_progress_thread.start()
    
    # Run alignment
    alignment_start = time.time()
    aligned_result = whisperx.align(
        result["segments"], 
        model_a, 
        metadata, 
        audio, 
        device, 
        return_char_alignments=False
    )
    
    alignment_time = time.time() - alignment_start
    monitor.update_alignment_progress(100)
    print(f"\nAlignment completed in {alignment_time/60:.1f} minutes")
    
    # Final cleanup
    monitor.set_phase("Final cleanup")
    del model_a
    gc.collect()
    
    return aligned_result

# Include all the original processing functions
def parse_whisperx_output(aligned_result):
    """Parse WhisperX output - same as original."""
    parsed_segments = []
    if not aligned_result or "segments" not in aligned_result:
        print("Warning: No segments found in WhisperX aligned result.")
        return []

    for segment_data in aligned_result["segments"]:
        seg_start_s = segment_data.get("start")
        seg_end_s = segment_data.get("end")
        seg_text = segment_data.get("text", "").strip()

        if seg_start_s is None or seg_end_s is None:
            continue

        timed_tokens = []
        if "words" in segment_data and isinstance(segment_data["words"], list):
            for word_info in segment_data["words"]:
                word_text = word_info.get("word", "").strip() 
                word_start_s = word_info.get("start")
                word_end_s = word_info.get("end")

                if not word_text or word_start_s is None or word_end_s is None:
                    continue
                
                word_start_s = max(word_start_s, seg_start_s)
                word_end_s = min(word_end_s, seg_end_s)
                if word_end_s <= word_start_s: 
                    word_end_s = word_start_s + 0.01

                timed_tokens.append({
                    "text": word_text, 
                    "start": word_start_s,
                    "end": word_end_s
                })
        
        if timed_tokens:
            reconstructed_text_builder = []
            for i, token_info in enumerate(timed_tokens):
                reconstructed_text_builder.append(token_info['text'])
                if i < len(timed_tokens) - 1:
                    next_token_text = timed_tokens[i+1]['text']
                    if token_info['text'] and next_token_text and not re.match(r"^[.,!?;:']", next_token_text):
                         reconstructed_text_builder.append(" ")
            seg_text_from_tokens = "".join(reconstructed_text_builder).strip()
        else:
            seg_text_from_tokens = seg_text

        if not seg_text_from_tokens:
            continue

        parsed_segments.append({
            "start": seg_start_s,
            "end": seg_end_s,
            "text": seg_text_from_tokens, 
            "tokens": timed_tokens 
        })
    return parsed_segments

def format_srt_timestamp(seconds):
    if not isinstance(seconds, (int, float)) or seconds < 0:
        seconds = 0 
    hours = int(seconds // 3600)
    minutes = int((seconds % 3600) // 60)
    seconds_val_float = seconds % 60 
    seconds_val_int = int(seconds_val_float)
    milliseconds = int(round((seconds_val_float - seconds_val_int) * 1000))
    
    if milliseconds >= 1000:
        seconds_val_int += 1; milliseconds = 0
        if seconds_val_int >= 60:
            seconds_val_int = 0; minutes +=1
            if minutes >= 60: minutes = 0; hours +=1
    return f"{hours:02d}:{minutes:02d}:{seconds_val_int:02d},{milliseconds:03d}"

def find_potential_split_points(sentence_tokens):
    split_candidates = []
    for i in range(len(sentence_tokens) - 1):
        current_token = sentence_tokens[i]
        next_token = sentence_tokens[i + 1]
        gap_duration = next_token['start'] - current_token['end']
        has_comma = current_token['text'].rstrip().endswith(',')
        score = gap_duration
        if has_comma:
            score += 0.5
        split_candidates.append((i, gap_duration, has_comma, score))
    split_candidates.sort(key=lambda x: x[3], reverse=True)
    return split_candidates

def validate_split_fragments(sentence_tokens, split_points):
    all_splits = [0] + [sp + 1 for sp in split_points] + [len(sentence_tokens)]
    all_splits = sorted(set(all_splits))
    for i in range(len(all_splits) - 1):
        start_idx = all_splits[i]
        end_idx = all_splits[i + 1]
        fragment_text = ""
        for j in range(start_idx, end_idx):
            fragment_text += sentence_tokens[j]['text']
            if j < end_idx - 1:
                next_token_text = sentence_tokens[j + 1]['text']
                if sentence_tokens[j]['text'] and next_token_text and not re.match(r"^[.,!?;:']", next_token_text):
                    fragment_text += " "
        alphanum_content = "".join(filter(str.isalnum, fragment_text.strip()))
        if len(alphanum_content) < 2:
            return False
    return True

def split_long_sentence(sentence_tokens, sentence_duration):
    if sentence_duration <= CONFIG['target_subtitle_duration']:
        return [(0, len(sentence_tokens))]
    
    target_parts = max(2, int(sentence_duration / CONFIG['target_subtitle_duration']) + 1)
    splits_needed = target_parts - 1
    
    split_candidates = find_potential_split_points(sentence_tokens)
    valid_candidates = [c for c in split_candidates if c[1] >= CONFIG['min_split_gap']]
    
    if len(valid_candidates) < splits_needed:
        splits_needed = len(valid_candidates)
    
    if splits_needed == 0:
        return [(0, len(sentence_tokens))]
    
    best_splits = None
    for attempt in range(min(3, len(valid_candidates) - splits_needed + 1)):
        candidate_splits = [c[0] for c in valid_candidates[attempt:attempt + splits_needed]]
        candidate_splits.sort()
        if validate_split_fragments(sentence_tokens, candidate_splits):
            best_splits = candidate_splits
            break
    
    if best_splits is None:
        return [(0, len(sentence_tokens))]
    
    ranges = []
    start_idx = 0
    for split_point in best_splits:
        ranges.append((start_idx, split_point + 1))
        start_idx = split_point + 1
    ranges.append((start_idx, len(sentence_tokens)))
    
    return ranges

# [Include the complete create_srt_file function - same as original]
def create_srt_file(parsed_segments, output_srt_file):
    with open(output_srt_file, 'w', encoding='utf-8') as f:
        srt_entry_count = 1 
        
        ABBREVIATIONS = {"Mr", "Mrs", "Ms", "Dr", "St", "Jr", "Sr", 
                         "Prof", "Rev", "Gen", "Sen", "Rep", "Gov", 
                         "Lt", "Col", "Capt", "Sgt", "Pvt", "Esq", "Hon", 
                         "Pres", "Sec", "Treas", "Det", "Insp", "Fr",
                         "No", "Op", "Vol", "Fig", "Ch", "Sec", "Dept", "Univ",
                         "approx", "etc", "eg", "ie", "vs"}

        all_final_srt_entries = [] 

        for segment_data in parsed_segments:
            full_segment_text = segment_data['text'] 
            segment_tokens = segment_data['tokens']

            if not full_segment_text.strip() or not segment_tokens:
                continue

            sentences_text_list = []
            current_sentence_buffer = ""
            delimiter_pattern = r'((?:[.?!]+)["\']?\s*)' 
            text_parts_for_splitting = re.split(delimiter_pattern, full_segment_text)
            
            current_part_idx = 0
            while current_part_idx < len(text_parts_for_splitting):
                text_chunk = text_parts_for_splitting[current_part_idx]; current_sentence_buffer += text_chunk; current_part_idx += 1
                if current_part_idx < len(text_parts_for_splitting): 
                    delimiter_chunk = text_parts_for_splitting[current_part_idx]; is_abbreviation_period = False
                    stripped_delimiter_text = delimiter_chunk.strip()
                    if stripped_delimiter_text.startswith("."): 
                        words_in_buffer = current_sentence_buffer.strip().split()
                        if words_in_buffer:
                            last_word_before_delim = words_in_buffer[-1]
                            cleaned_last_word = re.sub(r'[.,?!]+$', '', last_word_before_delim)
                            if cleaned_last_word in ABBREVIATIONS: is_abbreviation_period = True
                    next_text_lookahead_chunk = text_parts_for_splitting[current_part_idx+1] if current_part_idx + 1 < len(text_parts_for_splitting) else ""
                    alphanum_in_next_text = "".join(filter(str.isalnum, next_text_lookahead_chunk.strip()))
                    if not is_abbreviation_period and (len(alphanum_in_next_text) >= 2 or (current_part_idx + 1 >= len(text_parts_for_splitting) or not next_text_lookahead_chunk.strip())):
                        current_sentence_buffer += delimiter_chunk 
                        if current_sentence_buffer.strip(): sentences_text_list.append(current_sentence_buffer.strip())
                        current_sentence_buffer = "" 
                    else: current_sentence_buffer += delimiter_chunk
                current_part_idx += 1
            if current_sentence_buffer.strip(): sentences_text_list.append(current_sentence_buffer.strip())
            if not sentences_text_list and full_segment_text.strip(): sentences_text_list = [full_segment_text.strip()]

            token_cursor = 0 
            for sentence_str in sentences_text_list:
                sentence_str_clean = sentence_str.strip()
                if not sentence_str_clean: continue

                alphanum_content_check = "".join(filter(str.isalnum, sentence_str_clean))
                if len(alphanum_content_check) < 2:
                    num_words_in_filtered = len(sentence_str_clean.split())
                    token_cursor = min(token_cursor + num_words_in_filtered, len(segment_tokens))
                    continue

                sentence_tokens_list = []
                start_token_idx_for_sentence = -1
                end_token_idx_for_sentence = -1
                
                current_reconstructed_builder = []
                for i in range(token_cursor, len(segment_tokens)):
                    token = segment_tokens[i]
                    if start_token_idx_for_sentence == -1: 
                        start_token_idx_for_sentence = i
                    
                    sentence_tokens_list.append(token)
                    current_reconstructed_builder.append(token['text'])

                    temp_reconstructed_text = "".join(current_reconstructed_builder)
                    if temp_reconstructed_text.strip().replace(" ", "") == sentence_str_clean.replace(" ", ""):
                        end_token_idx_for_sentence = i
                        token_cursor = i + 1 
                        break 
                    if len(temp_reconstructed_text.replace(" ","")) > len(sentence_str_clean.replace(" ","")) + 5:
                        sentence_tokens_list = [] 
                        start_token_idx_for_sentence = -1
                        break
                else: 
                    if sentence_tokens_list and start_token_idx_for_sentence != -1: 
                        end_token_idx_for_sentence = len(segment_tokens) - 1 
                    token_cursor = len(segment_tokens)

                if sentence_tokens_list and start_token_idx_for_sentence != -1 and end_token_idx_for_sentence != -1:
                    sentence_start_time = segment_tokens[start_token_idx_for_sentence]['start']
                    sentence_end_time = segment_tokens[end_token_idx_for_sentence]['end']
                    sentence_duration = sentence_end_time - sentence_start_time
                    
                    relevant_tokens = segment_tokens[start_token_idx_for_sentence:end_token_idx_for_sentence + 1]
                    token_ranges = split_long_sentence(relevant_tokens, sentence_duration)
                    
                    for range_start, range_end in token_ranges:
                        part_tokens = relevant_tokens[range_start:range_end]
                        if not part_tokens:
                            continue
                            
                        part_text_builder = []
                        for j, token in enumerate(part_tokens):
                            part_text_builder.append(token['text'])
                            if j < len(part_tokens) - 1:
                                next_token_text = part_tokens[j + 1]['text']
                                if token['text'] and next_token_text and not re.match(r"^[.,!?;:']", next_token_text):
                                    part_text_builder.append(" ")
                        
                        part_text = "".join(part_text_builder).strip()
                        
                        part_alphanum = "".join(filter(str.isalnum, part_text))
                        if len(part_alphanum) >= 2:
                            part_start_time = part_tokens[0]['start']
                            part_end_time = part_tokens[-1]['end']
                            if part_end_time <= part_start_time:
                                part_end_time = part_start_time + 0.1
                            
                            all_final_srt_entries.append({
                                "text": part_text, 
                                "start": part_start_time, 
                                "end": part_end_time
                            })
                            
                elif sentence_str_clean: 
                    fallback_start = segment_data['start'] 
                    fallback_end = segment_data['end']   
                    if all_final_srt_entries: 
                        fallback_start = max(all_final_srt_entries[-1]['end'], fallback_start)
                    
                    relative_len = len(sentence_str_clean) / len(full_segment_text) if len(full_segment_text) > 0 else 0.1
                    estimated_duration = (segment_data['end'] - segment_data['start']) * relative_len
                    estimated_duration = max(0.2, estimated_duration)
                    
                    fallback_end = fallback_start + estimated_duration
                    fallback_end = min(fallback_end, segment_data['end']) 
                    if fallback_end <= fallback_start: fallback_end = fallback_start + 0.2

                    all_final_srt_entries.append({"text": sentence_str_clean, "start": fallback_start, "end": fallback_end})
            
        if not all_final_srt_entries:
            print("No subtitle entries to write after processing all segments.")
            return

        deduplicated_srt_entries = []
        previous_entry_text_compare = None 
        for entry in all_final_srt_entries: 
            current_text_stripped = entry["text"].strip()
            alphanum_content_final_check = "".join(filter(str.isalnum, current_text_stripped))
            if len(alphanum_content_final_check) < 2:
                continue

            if current_text_stripped == previous_entry_text_compare: 
                if deduplicated_srt_entries: 
                    deduplicated_srt_entries[-1]["end"] = max(deduplicated_srt_entries[-1]["end"], entry["end"])
                continue
            deduplicated_srt_entries.append(entry)
            previous_entry_text_compare = current_text_stripped 

        for entry_to_write in deduplicated_srt_entries:
            start_time_val = entry_to_write["start"]; end_time_val = entry_to_write["end"]
            if end_time_val <= start_time_val: end_time_val = start_time_val + 0.1 
            start_time_str = format_srt_timestamp(start_time_val); end_time_str = format_srt_timestamp(end_time_val)
            f.write(f"{srt_entry_count}\n{start_time_str} --> {end_time_str}\n{entry_to_write['text']}\n\n"); srt_entry_count += 1

def create_text_transcript_summary(parsed_segments):
    if not parsed_segments: return "No segments processed for summary."
    lines = []
    for segment_data in parsed_segments:
        minutes = int(segment_data['start'] // 60); seconds_val = segment_data['start'] % 60
        timestamp = f"[{minutes:02d}:{seconds_val:06.3f}]"; cleaned_text = segment_data['text'].strip() 
        lines.append(f"{timestamp} {cleaned_text}")
    return '\n'.join(lines)

def process_transcription_data(whisperx_result, srt_filename):
    """Process WhisperX result to create SRT and text summary."""
    parsed_segments = parse_whisperx_output(whisperx_result)
    if not parsed_segments:
        print("No valid segments extracted from WhisperX result.")
        return "No segments were found in the transcript to process." 
        
    create_srt_file(parsed_segments, srt_filename)
    return create_text_transcript_summary(parsed_segments)

def main():
    if not WHISPERX_AVAILABLE:
        sys.exit("Exiting due to missing WhisperX/PyTorch libraries.")

    print("Starting WhisperX with Progress Monitoring...")
    print_config()
    
    audio_file = find_audio_file()
    print(f"Found audio file: {audio_file}")
    
    # Get audio duration for progress estimation
    audio_duration = get_audio_duration(audio_file)
    if audio_duration:
        print(f"Audio duration: {audio_duration:.1f} seconds ({audio_duration/60:.1f} minutes)")
    
    base_filename, _ = os.path.splitext(audio_file)
    srt_filename = base_filename + ".srt"
    
    # Initialize progress monitor
    monitor = ProgressMonitor(audio_duration)
    
    try:
        total_start_time = time.time()
        
        # Start progress monitoring
        monitor.start_monitoring()
        
        # Run WhisperX processing with progress monitoring
        whisperx_result = run_whisperx_with_progress(audio_file, monitor)
        
        if not whisperx_result or "segments" not in whisperx_result or not whisperx_result["segments"]:
            print("Transcription failed: No segments were produced by WhisperX.")
            sys.exit(1)
        
        # Stop progress monitoring for processing phase
        monitor.stop_monitoring()
        
        # Save raw JSON output
        monitor.set_phase("Saving JSON output")
        raw_json_output_path = base_filename + "_whisperx_output.json"
        try:
            with open(raw_json_output_path, "w", encoding='utf-8') as f_json:
                json.dump(whisperx_result, f_json, indent=4, ensure_ascii=False)
            print(f"Raw WhisperX JSON output saved to: {raw_json_output_path}")
        except Exception as e:
            print(f"Error saving raw WhisperX JSON output: {e}")

        # Process into SRT and text
        monitor.set_phase("Creating SRT and text output")
        processed_summary_text = process_transcription_data(whisperx_result, srt_filename)
        
        total_time = time.time() - total_start_time
        minutes = int(total_time // 60)
        seconds_proc = total_time % 60
        
        # Calculate performance metrics
        processing_ratio = total_time / audio_duration if audio_duration else 0
        
        # Create summary header
        header = f"""WHISPERX WITH PROGRESS MONITORING SUMMARY
=====================================
Source audio: {audio_file}
Audio duration: {audio_duration:.1f} seconds ({audio_duration/60:.1f} minutes)
WhisperX JSON output: {raw_json_output_path} 
SRT output file: {srt_filename}
Processing time (total): {minutes} minutes and {seconds_proc:.2f} seconds
Processing ratio: {processing_ratio:.2f}x
Processed on: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}

CONFIGURATION USED:
Model: {CONFIG['model_size']}
Batch size: {CONFIG['batch_size']}
CPU threads: {CONFIG['cpu_threads']}
Compute type: {CONFIG['compute_type']}
Target subtitle duration: {CONFIG['target_subtitle_duration']}s
Min split gap: {CONFIG['min_split_gap']}s

PROGRESS MONITORING:
Update interval: {CONFIG['progress_update_interval']} seconds
Memory monitoring: {CONFIG['show_memory_usage']}
Time estimation: {CONFIG['estimate_total_time']}

"""
        
        # Save text summary
        with open(CONFIG['output_txt_filename'], 'w', encoding='utf-8') as f:
            f.write(header + processed_summary_text)
        
        # Print final results
        print(f"\n{'='*70}")
        print("PROGRESS MONITORING TRANSCRIPTION COMPLETE!")
        print(f"{'='*70}")
        print(f"Text transcript summary saved to: {CONFIG['output_txt_filename']}")
        print(f"SRT subtitle file saved to: {srt_filename}")
        print(f"JSON output saved to: {raw_json_output_path}")
        print(f"Total processing time: {minutes} minutes and {seconds_proc:.2f} seconds")
        
        if audio_duration:
            print(f"Audio duration: {audio_duration/60:.1f} minutes")
            print(f"Processing ratio: {processing_ratio:.2f}x")
            
            # Performance assessment
            if processing_ratio < 1.0:
                print("✅ Real-time performance achieved!")
            elif processing_ratio < 1.5:
                print("🟡 Near real-time performance")
            else:
                print("🔴 Slower than real-time")
            
            # Time projections for longer files
            if audio_duration < 3600:  # Less than 1 hour
                projected_2h = processing_ratio * 120  # 2 hours
                print(f"📊 Projected time for 2-hour audio: {projected_2h:.1f} minutes")
        
        print("\nSample of transcript:")
        summary_lines = processed_summary_text.split('\n')
        sample = '\n'.join(summary_lines[:min(5, len(summary_lines))])
        print(sample)
        
        # Final tips
        if processing_ratio > 2.0:
            print(f"\n{'='*50}")
            print("PERFORMANCE TIPS:")
            print("• Consider using 'medium' model for faster processing")
            print("• Close other applications to free CPU/memory")
            print("• Ensure adequate cooling for sustained processing")
            print("• Monitor system resources during long transcriptions")
    
    except KeyboardInterrupt:
        print("\n\nProcessing interrupted by user.")
        monitor.stop_monitoring()
        sys.exit(0)
    
    except Exception as e:
        print(f"\nAn error occurred during processing: {e}")
        monitor.stop_monitoring()
        import traceback
        traceback.print_exc()
        sys.exit(1)
    
    finally:
        # Ensure monitoring stops
        if 'monitor' in locals():
            monitor.stop_monitoring()

if __name__ == "__main__":
    main()