"""
src/dataset.py

Dataset Loader and Dynamic Padding Collate Function for Pre-cached Audio Features.
Handles cached .pt representations, token ID targets, attention masks, and CTC sequence lengths.
"""

from __future__ import annotations

import glob
import json
import os
import re
import unicodedata
import zipfile
from pathlib import Path
from typing import List, Dict, Any, Tuple, Optional

import torch
from torch.utils.data import Dataset
from tqdm import tqdm


# ==============================================================================
# 1. Korean Morpheme Tokenizer & Normalization
# ==============================================================================

# SentencePiece word prefix marker (U+2581 ' ')
WORD_PREFIX = " "

# Korean compatibility consonant (U+3131~U+314E) to standard conjoining jongseong (U+11A8~U+11C2) mapping
COMPAT_TO_JONGSEONG: Dict[str, str] = {
    "ㄱ": "\u11a8", "ㄲ": "\u11a9", "ㄳ": "\u11aa", "ㄴ": "\u11ab", "ㄵ": "\u11ac",
    "ㄶ": "\u11ad", "ㄷ": "\u11ae", "ㄹ": "\u11af", "ㄺ": "\u11b0", "ㄻ": "\u11b1",
    "ㄼ": "\u11b2", "ㄽ": "\u11b3", "ㄾ": "\u11b4", "ㄿ": "\u11b5", "ㅀ": "\u11b6",
    "ㅁ": "\u11b7", "ㅂ": "\u11b8", "ㅄ": "\u11b9", "ㅅ": "\u11ba", "ㅆ": "\u11bb",
    "ㅇ": "\u11bc", "ㅈ": "\u11bd", "ㅊ": "\u11be", "ㅋ": "\u11bf", "ㅌ": "\u11c0",
    "ㅍ": "\u11c1", "ㅎ": "\u11c2"
}

_PUNCT_RE = re.compile(r"[?!.,\"~:;()\[\]{}<>@#$%^&*_+=\/\\|`·…“”‘’]")
_SPACE_RE = re.compile(r"\s+")

def normalize_text(text: str) -> str:
    """
    Normalizes spoken Korean text by stripping punctuation marks (?, ., ,, !, etc.),
    preserving Korean characters, alphanumeric characters, and whitespace.
    """
    text = _PUNCT_RE.sub(" ", text)
    return _SPACE_RE.sub(" ", text).strip()

import sqlite3
import threading


def canonicalize_aihub_key(key: str) -> str:
    """Normalizes any AI Hub path or key (e.g. 'art/desi/C00890/U00004.wav') to 'C00890_U00004'."""
    k = key.replace("\\", "/").strip("/")
    if k.endswith((".wav", ".pcm", ".flac", ".json")):
        k = k.rsplit(".", 1)[0]
    if "/" in k:
        parts = k.split("/")
        if len(parts) >= 2:
            return f"{parts[-2]}_{parts[-1]}"
        return parts[-1]
    return k


class AIHubLabelIndexDB:
    """
    Zero-RAM SQLite-backed dictionary interface for AI Hub 1.71M lecture transcriptions.
    Replaces the 681MB in-memory JSON dict (~2.5-4.5GB RAM) with an on-disk B-Tree index (<5MB RAM).
    Thread-safe across main thread and PrefetchedStream worker threads.
    """
    def __init__(self, db_path: Path):
        self.db_path = Path(db_path)
        self._local = threading.local()
        self._count: Optional[int] = None

    def _get_conn(self) -> sqlite3.Connection:
        conn = getattr(self._local, "conn", None)
        if conn is None:
            conn = sqlite3.connect(f"file:{self.db_path}?mode=ro", uri=True, check_same_thread=False)
            conn.execute("PRAGMA cache_size = -4000")  # Cap SQLite page cache to ~4MB
            conn.execute("PRAGMA temp_store = MEMORY")
            self._local.conn = conn
        return conn

    def get(self, key: str, default: Optional[str] = None) -> Optional[str]:
        if not key:
            return default
        ckey = canonicalize_aihub_key(key)
        cur = self._get_conn().execute("SELECT text FROM labels WHERE key = ?", (ckey,))
        row = cur.fetchone()
        return row[0] if row is not None else default

    def __getitem__(self, key: str) -> str:
        val = self.get(key)
        if val is None:
            raise KeyError(key)
        return val

    def __contains__(self, key: str) -> bool:
        return self.get(key) is not None

    def __len__(self) -> int:
        if self._count is None:
            cur = self._get_conn().execute("SELECT COUNT(*) FROM labels")
            row = cur.fetchone()
            self._count = int(row[0]) if row else 0
        return self._count

    def values(self):
        cur = self._get_conn().execute("SELECT text FROM labels")
        for (txt,) in cur:
            yield txt


def load_or_build_aihub_label_index(
    tl_zip_path: str = "/Users/yeongchan/Downloads/137.한국어 대학 강의 데이터/01-1.정식개방데이터/Training/02.라벨링데이터/TL.zip",
    index_cache_path: str = "data/vocab/aihub_label_index.db",
    max_scan_entries: Optional[int] = None
) -> Any:
    """
    Scans TL.zip directly in memory without disk extraction and stores canonical keys
    in a lightweight SQLite database (data/vocab/aihub_label_index.db) consuming <5MB RAM.
    """
    project_root = Path(__file__).resolve().parent.parent
    cache_path = Path(index_cache_path)
    if not cache_path.is_absolute():
        cache_path = project_root / cache_path
    db_path = cache_path.with_suffix(".db")
    legacy_json_path = cache_path.with_suffix(".json")

    if db_path.exists() and db_path.stat().st_size > 1024 * 1024:
        db_index = AIHubLabelIndexDB(db_path)
        print(f"[AIHub Index] Opened low-memory SQLite label index ({len(db_index):,} utterances, <5MB RAM) at {db_path.name}")
        return db_index

    db_path.parent.mkdir(parents=True, exist_ok=True)
    tmp_db_path = db_path.with_suffix(".db.tmp")
    if tmp_db_path.exists():
        tmp_db_path.unlink()

    conn = sqlite3.connect(str(tmp_db_path))
    conn.execute("PRAGMA journal_mode = OFF")
    conn.execute("PRAGMA synchronous = OFF")
    conn.execute("CREATE TABLE labels (key TEXT PRIMARY KEY, text TEXT) WITHOUT ROWID")

    tl_zip = Path(tl_zip_path)
    if not tl_zip.exists():
        conn.close()
        raise FileNotFoundError(f"AI Hub TL.zip not found at {tl_zip}")

    print(f"[AIHub Index] Building low-memory SQLite index from {tl_zip.name} without disk extraction...")
    batch_rows: List[Tuple[str, str]] = []
    inserted = 0

    with zipfile.ZipFile(tl_zip, "r") as zf:
        infolist = zf.infolist()
        valid_infos = [
            info for info in infolist
            if info.filename.endswith(".json")
            and not info.filename.startswith("__MACOSX")
            and ".irx" not in info.filename
            and ".part" not in info.filename
            and ".DS_Store" not in info.filename
        ]
        if max_scan_entries and len(valid_infos) > max_scan_entries:
            valid_infos = valid_infos[:max_scan_entries]

        for info in tqdm(valid_infos, desc="Indexing TL.zip -> SQLite"):
            try:
                raw_bytes = zf.read(info)
                jdata = json.loads(raw_bytes.decode("utf-8"))
                raw_txt = jdata.get("06_transcription", {}).get("1_text", "").strip()
                if not raw_txt:
                    continue

                clean_txt = re.sub(r"\([^)]*\)/\(([^)]*)\)", r"\1", raw_txt)
                clean_txt = normalize_text(clean_txt)
                if not clean_txt:
                    continue

                src_path = jdata.get("01_dataset", {}).get("3_src_path", "") or info.filename
                ckey = canonicalize_aihub_key(src_path)
                batch_rows.append((ckey, clean_txt))

                if len(batch_rows) >= 20000:
                    conn.executemany("INSERT OR REPLACE INTO labels (key, text) VALUES (?, ?)", batch_rows)
                    conn.commit()
                    inserted += len(batch_rows)
                    batch_rows.clear()
            except Exception:
                continue

    if batch_rows:
        conn.executemany("INSERT OR REPLACE INTO labels (key, text) VALUES (?, ?)", batch_rows)
        conn.commit()
        inserted += len(batch_rows)
        batch_rows.clear()

    conn.close()
    tmp_db_path.replace(db_path)

    if legacy_json_path.exists():
        try:
            legacy_json_path.unlink()
            print(f"[AIHub Index] Removed legacy 681MB JSON cache ({legacy_json_path.name}) to free disk & RAM.")
        except Exception:
            pass

    db_index = AIHubLabelIndexDB(db_path)
    print(f"[AIHub Index] SQLite index ready at {db_path.name} ({len(db_index):,} unique utterances).")
    return db_index


def is_consonant_leading(token: str) -> bool:
    """
    Checks if a token starts with a standalone consonant / jamo (e.g. ᆫ, ᆯ, ᆷ, ᆸ, ᆺ, or compatibility ㄱ-ㅎ).
    """
    if not token:
        return False
    c = token[0]
    # Compatibility consonants: U+3131 ~ U+314E
    if "\u3131" <= c <= "\u314e":
        return True
    # Conjoining jamo: choseong (U+1100~U+115F) or jongseong (U+11A8~U+11FF)
    if "\u1100" <= c <= "\u11ff":
        return True
    return False

def merge_jamo_to_syllable(prev: str, curr: str) -> str:
    """
    Composes a preceding morpheme with a leading standalone jamo into a complete syllable:
      e.g. '쌓이' + 'ᆫ' -> '쌓인'
           '이' + 'ᆸ니다' -> '입니다'
           '만들' + '어' + '지' + 'ᆫ' -> '만들', '어', '진'
           '참여' + '하' + 'ᆯ' -> '참여', '할'
           '만들' + 'ᆫ다' -> '만든다' (ㄹ-탈락 규칙)
           '살' + 'ᆸ니다' -> '삽니다' (ㄹ-탈락 규칙)
           '그렇' + 'ᆫ가' -> '그런가' (ㅎ-탈락 규칙)
    Uses unicodedata NFD decomposition and NFC re-composition.
    """
    if not prev:
        return curr
    first_char = curr[0]
    first_mapped = COMPAT_TO_JONGSEONG.get(first_char, first_char)
    curr_mapped = first_mapped + curr[1:]

    prev_nfd = list(unicodedata.normalize("NFD", prev))
    curr_nfd = list(unicodedata.normalize("NFD", curr_mapped))

    # If prev ends in a jongseong and curr starts with a jongseong:
    # in Korean phonological rules (e.g. ㄹ-탈락, ㅎ-탈락), the stem jongseong is dropped
    if prev_nfd and ("\u11a8" <= prev_nfd[-1] <= "\u11c2"):
        if curr_nfd and ("\u11a8" <= curr_nfd[0] <= "\u11c2"):
            test_merged = unicodedata.normalize("NFC", "".join(prev_nfd + curr_nfd))
            if any(("\u1100" <= c <= "\u11ff" or "\u3131" <= c <= "\u314e") for c in test_merged):
                prev_nfd = prev_nfd[:-1]

    merged = unicodedata.normalize("NFC", "".join(prev_nfd + curr_nfd))
    return merged

def strip_word_prefix(tok: str) -> Tuple[bool, str]:
    """
    Checks if a token has a word-boundary prefix (' ', ' ▁', or space) and strips it.
    Returns (is_word_start, clean_token).
    """
    if tok.startswith(WORD_PREFIX) or tok.startswith(" ") or tok.startswith(" ▁") or tok.startswith(" "):
        clean = tok.lstrip(" \u2581▁")
        return True, clean
    return False, tok

def tokenize_korean_syllables(text: str, word_prefix: str = WORD_PREFIX) -> List[str]:
    clean_text = normalize_text(text)
    if not clean_text:
        return []
    tokens: List[str] = []
    for word in clean_text.split(" "):
        if not word:
            continue
        for idx, ch in enumerate(word):
            tokens.append((word_prefix + ch) if idx == 0 else ch)
    return tokens


def tokenize_korean_hybrid(
    text: str,
    token_to_id: Dict[str, int],
    word_prefix: str = WORD_PREFIX,
    max_tok_len: int = 16
) -> List[str]:
    """
    Greedy Longest-Match (Maximal Matching) Surface Morpheme Tokenizer with Syllable Fallback.
    Matches the longest registered surface morpheme/word in vocabulary first.
    Falls back to single-character KS X 1001 syllables when a word/morpheme is unregistered,
    guaranteeing 0.00% OOV with maximum semantic chunking efficiency.
    """
    clean_text = normalize_text(text)
    if not clean_text:
        return []
    tokens: List[str] = []
    for word in clean_text.split(" "):
        if not word:
            continue
        w_len = len(word)
        pos = 0
        first_in_word = True
        while pos < w_len:
            matched = False
            max_end = w_len if w_len - pos <= max_tok_len else pos + max_tok_len
            for end in range(max_end, pos, -1):
                sub = word[pos:end]
                candidate = (word_prefix + sub) if first_in_word else sub
                if candidate in token_to_id:
                    tokens.append(candidate)
                    pos = end
                    first_in_word = False
                    matched = True
                    break
            if not matched:
                sub = word[pos:pos + 1]
                candidate = (word_prefix + sub) if first_in_word else sub
                tokens.append(candidate)
                pos += 1
                first_in_word = False
    return tokens


# Backward compatibility alias
tokenize_korean_morphemes = tokenize_korean_hybrid


class KoreanMorphemeTokenizer:
    """
    Surface Morpheme + Syllable Fallback Hybrid Tokenizer using generated vocab.json with word boundary prefix.
    Utilizes O(L) Greedy Longest-Match dictionary lookup for 0.00% OOV and maximum semantic representation.
    """
    def __init__(self, vocab_json_path: str):
        with open(vocab_json_path, "r", encoding="utf-8") as f:
            data = json.load(f)

        self.token_to_id: Dict[str, int] = data["token_to_id"]
        # Convert integer keys back if serialized as strings in json
        self.id_to_token: Dict[int, str] = {int(k): v for k, v in data["id_to_token"].items()}
        self.blank_id: int = data.get("blank_id", 0)
        self.blank_token: str = data.get("blank_token", "<blank>")
        self.unk_id: int = data.get("unk_id", 1)
        self.unk_token: str = data.get("unk_token", "<unk>")
        self.word_prefix: str = WORD_PREFIX
        self._max_tok_len: int = max((len(k) for k in self.token_to_id.keys()), default=16)
        self._word_id_cache: Dict[str, Tuple[int, ...]] = {}

    def _encode_word(self, word: str) -> Tuple[int, ...]:
        cached = self._word_id_cache.get(word)
        if cached is not None:
            return cached
        w_len = len(word)
        pos = 0
        first_in_word = True
        ids: List[int] = []
        t2i = self.token_to_id
        wp = self.word_prefix
        max_tl = self._max_tok_len
        unk = self.unk_id

        while pos < w_len:
            matched = False
            max_end = w_len if w_len - pos <= max_tl else pos + max_tl
            for end in range(max_end, pos, -1):
                sub = word[pos:end]
                cand = (wp + sub) if first_in_word else sub
                tid = t2i.get(cand)
                if tid is not None:
                    ids.append(tid)
                    pos = end
                    first_in_word = False
                    matched = True
                    break
            if not matched:
                sub = word[pos:pos + 1]
                cand = (wp + sub) if first_in_word else sub
                ids.append(t2i.get(cand, unk))
                pos += 1
                first_in_word = False

        res = tuple(ids)
        if len(self._word_id_cache) < 65536:
            self._word_id_cache[word] = res
        return res

    def encode(self, text: str) -> List[int]:
        """
        Tokenizes text into surface morpheme tokens (with syllable fallback) and maps them to token IDs.
        """
        clean_text = normalize_text(text)
        if not clean_text:
            return []
        out_ids: List[int] = []
        for word in clean_text.split(" "):
            if word:
                out_ids.extend(self._encode_word(word))
        return out_ids

    def decode(self, token_ids: List[int], remove_blank: bool = True, collapse_repeats: bool = True) -> str:
        """
        Decodes a sequence of token IDs into text.
        For CTC output:
        - collapse_repeats: collapses consecutive duplicate token IDs
        - remove_blank: removes blank_id (0)
        - word boundary reconstruction: reconstructs spaces between words using prefix ' ' or ' ▁'
        """
        if collapse_repeats:
            collapsed = []
            prev = None
            for tid in token_ids:
                if tid != prev:
                    collapsed.append(tid)
                    prev = tid
            token_ids = collapsed

        words = []
        for tid in token_ids:
            if remove_blank and tid == self.blank_id:
                continue
            token_str = self.id_to_token.get(tid, "")
            if token_str and token_str != self.blank_token and token_str != self.unk_token:
                words.append(token_str)

        res = ""
        for w in words:
            is_ws, clean_w = strip_word_prefix(w)
            if is_ws:
                if res:
                    res += " " + clean_w
                else:
                    res = clean_w
            else:
                res += clean_w

        return res.strip()


_CONTRACTIONS: Tuple[Tuple[str, str], ...] = (
    ("하었", "했"), ("하 었", "했"), ("되었", "됐"), ("되 었", "됐"),
    ("나오었", "나왔"), ("나오 었", "나왔"), ("보았", "봤"), ("보 었", "봤"),
    ("주었", "줬"), ("주 었", "줬"), ("오었", "왔"), ("오 었", "왔"),
    ("가었", "갔"), ("가 었", "갔"), ("지었", "졌"), ("지 었", "졌"),
    ("받아들이어", "받아들여")
)


# ==============================================================================
# 2. Cached Audio Feature Dataset & Length-Bucketed Batch Sampler
# ==============================================================================

class _SharedFilePathAccessor:
    """Zero-RAM sequence wrapper over shared-memory byte & offset tensors."""
    def __init__(self, base_dir: str, names_buf: torch.Tensor, offsets_buf: torch.Tensor):
        self.base_dir = base_dir
        self.names_buf = names_buf
        self.offsets_buf = offsets_buf

    def __len__(self) -> int:
        return max(0, int(self.offsets_buf.numel()) - 1)

    def get_name(self, idx: int) -> str:
        s = int(self.offsets_buf[idx].item())
        e = int(self.offsets_buf[idx + 1].item())
        return bytes(self.names_buf[s:e].numpy()).decode("utf-8")

    def __getitem__(self, idx: int) -> str:
        return os.path.join(self.base_dir, self.get_name(idx))


class CachedAudioFeatureDataset(Dataset):
    """
    Loads pre-extracted (Time, 256) audio encoder representations and target token IDs.
    Uses OS shared-memory tensors (.share_memory_()) for file names and sizes so DataLoader
    worker processes share a single 68MB index with zero Copy-on-Write or pickling bloat.
    """
    def __init__(
        self,
        cache_dir: str,
        tokenizer: Optional[KoreanMorphemeTokenizer] = None,
        max_files: Optional[int] = None
    ):
        super().__init__()
        self.cache_dir = Path(cache_dir)
        self.cache_dir_str = str(self.cache_dir)
        self.tokenizer = tokenizer

        index_cache_file = self.cache_dir.parent / "vocab" / "feature_cache_index.pt"
        loaded_from_cache = False

        if max_files is None and index_cache_file.exists():
            try:
                cached_idx = torch.load(index_cache_file, map_location="cpu", weights_only=False)
                if (
                    isinstance(cached_idx, dict)
                    and cached_idx.get("cache_dir") == self.cache_dir_str
                    and int(cached_idx.get("count", 0)) > 0
                ):
                    self._names_buf = cached_idx["names_buf"].share_memory_()
                    self._offsets_buf = cached_idx["offsets_buf"].share_memory_()
                    self.file_sizes = cached_idx["sizes_buf"].share_memory_()
                    loaded_from_cache = True
            except Exception:
                loaded_from_cache = False

        if not loaded_from_cache:
            names_list: List[bytes] = []
            sizes_list: List[int] = []
            if self.cache_dir.exists():
                entries = [e for e in os.scandir(self.cache_dir) if e.name.endswith(".pt")]
                entries.sort(key=lambda e: e.name)
                if max_files is not None and max_files > 0:
                    entries = entries[:max_files]
                names_list = [e.name.encode("utf-8") for e in entries]
                sizes_list = [e.stat().st_size for e in entries]

            if names_list:
                lengths = [len(b) for b in names_list]
                offsets = [0]
                acc = 0
                for ln in lengths:
                    acc += ln
                    offsets.append(acc)
                joined = b"".join(names_list)
                self._names_buf = torch.frombuffer(bytearray(joined), dtype=torch.uint8).share_memory_()
                self._offsets_buf = torch.tensor(offsets, dtype=torch.int64).share_memory_()
                self.file_sizes = torch.tensor(sizes_list, dtype=torch.int32).share_memory_()

                if max_files is None:
                    try:
                        index_cache_file.parent.mkdir(parents=True, exist_ok=True)
                        tmp_idx = index_cache_file.with_suffix(".pt.tmp")
                        torch.save(
                            {
                                "cache_dir": self.cache_dir_str,
                                "count": len(names_list),
                                "names_buf": self._names_buf,
                                "offsets_buf": self._offsets_buf,
                                "sizes_buf": self.file_sizes,
                            },
                            tmp_idx
                        )
                        tmp_idx.replace(index_cache_file)
                    except Exception:
                        pass
            else:
                self._names_buf = torch.empty((0,), dtype=torch.uint8).share_memory_()
                self._offsets_buf = torch.zeros((1,), dtype=torch.int64).share_memory_()
                self.file_sizes = torch.empty((0,), dtype=torch.int32).share_memory_()
                print(f"[Dataset Warning] No cached .pt files found in {self.cache_dir}.")

        self.file_paths = _SharedFilePathAccessor(self.cache_dir_str, self._names_buf, self._offsets_buf)

    def __len__(self) -> int:
        return int(self.file_sizes.numel())

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        fname = self.file_paths.get_name(idx)
        file_path = os.path.join(self.cache_dir_str, fname)
        data = torch.load(file_path, map_location="cpu", weights_only=False)
        feature = data["feature"]  # Shape: (T, 256)
        raw_text = data.get("raw_text", "")

        if self.tokenizer is not None and raw_text:
            tok_ids = self.tokenizer.encode(raw_text)
            targets = torch.tensor(tok_ids, dtype=torch.long)
        elif "targets" in data and isinstance(data["targets"], torch.Tensor):
            targets = data["targets"].to(torch.long)
        else:
            targets = torch.tensor([], dtype=torch.long)

        if not isinstance(feature, torch.Tensor):
            feature = torch.tensor(feature, dtype=torch.bfloat16)

        # Strictly enforce bfloat16 (prevent MPS float16 underflow and float32 NaN instability)
        if feature.dtype != torch.bfloat16:
            feature = feature.to(torch.bfloat16)

        return {
            "feature": feature,
            "targets": targets,
            "raw_text": raw_text,
            "file_name": fname
        }


class LengthBucketBatchSampler(torch.utils.data.Sampler[List[int]]):
    """
    Groups samples of similar sequence lengths into batches using cached .pt file sizes
    as a zero-I/O O(1) proxy for frame length T (~512 bytes per time frame).
    Eliminates ~40% random-batching padding overhead while preserving stochasticity
    via mega-batch pool sorting and batch-order shuffling each epoch.
    """
    def __init__(
        self,
        dataset: Dataset,
        batch_size: int,
        shuffle: bool = True,
        pool_batches: int = 128,
        seed: int = 42
    ):
        super().__init__()
        self.dataset = dataset
        self.batch_size = max(1, batch_size)
        self.shuffle = shuffle
        self.pool_size = self.batch_size * max(1, pool_batches)
        self.seed = seed
        self.epoch = 0
        self.start_batch = 0

        # Resolve per-item file size proxy as compact int32 tensor
        if hasattr(dataset, "file_sizes"):
            fs = dataset.file_sizes
            self.lengths_t = fs if isinstance(fs, torch.Tensor) else torch.tensor(fs, dtype=torch.int32)
        elif hasattr(dataset, "dataset") and hasattr(dataset, "indices") and hasattr(dataset.dataset, "file_sizes"):
            base_sizes = dataset.dataset.file_sizes
            base_t = base_sizes if isinstance(base_sizes, torch.Tensor) else torch.tensor(base_sizes, dtype=torch.int32)
            idx_t = dataset.indices if isinstance(dataset.indices, torch.Tensor) else torch.tensor(dataset.indices, dtype=torch.long)
            self.lengths_t = base_t[idx_t]
        else:
            self.lengths_t = torch.zeros((len(dataset),), dtype=torch.int32)

    def set_epoch(self, epoch: int, start_batch: int = 0) -> None:
        """Sets current epoch (0-indexed) and optional starting batch offset for zero-I/O resume."""
        self.epoch = max(0, int(epoch))
        self.start_batch = max(0, int(start_batch))

    def __len__(self) -> int:
        n = int(self.lengths_t.numel())
        total = (n + self.batch_size - 1) // self.batch_size
        return max(0, total - self.start_batch)

    def __iter__(self):
        n = int(self.lengths_t.numel())
        start_b = self.start_batch
        self.start_batch = 0  # Reset after consuming once

        if self.shuffle:
            g = torch.Generator()
            g.manual_seed(self.seed + self.epoch)
            self.epoch += 1
            perm = torch.randperm(n, generator=g)

            # Vectorized mega-pool sorting in C++ (torch.argsort) instead of Python list sorting
            num_full_pools = n // self.pool_size
            if num_full_pools > 0:
                full_len = num_full_pools * self.pool_size
                perm_main = perm[:full_len].view(num_full_pools, self.pool_size)
                pool_lengths = self.lengths_t[perm_main]
                sort_idx = torch.argsort(pool_lengths, dim=1, stable=True)
                sorted_main = torch.gather(perm_main, 1, sort_idx).reshape(-1)
                if full_len < n:
                    rem = perm[full_len:]
                    rem_sorted = rem[torch.argsort(self.lengths_t[rem], stable=True)]
                    sorted_indices = torch.cat([sorted_main, rem_sorted], dim=0)
                else:
                    sorted_indices = sorted_main
            else:
                sorted_indices = perm[torch.argsort(self.lengths_t[perm], stable=True)]

            all_indices = sorted_indices.tolist()
            batches = [
                all_indices[b_start : b_start + self.batch_size]
                for b_start in range(0, n, self.batch_size)
            ]

            b_perm = torch.randperm(len(batches), generator=g).tolist()
            for bi in b_perm[start_b:]:
                yield batches[bi]
        else:
            sorted_indices = torch.argsort(self.lengths_t, stable=True).tolist()
            all_batches = [
                sorted_indices[b_start : b_start + self.batch_size]
                for b_start in range(0, n, self.batch_size)
            ]
            for b in all_batches[start_b:]:
                yield b


# ==============================================================================
# 3. Dynamic Padding Collate Function for CTC
# ==============================================================================

class PadCollate:
    """
    Collates a list of variable-length audio features and targets:
    - Caps outlier audio lengths to max_frames (default 512 = ~20.5s) to prevent MPS OOM spikes
    - Pads features to maximum time length in the batch -> (Batch, T_max, 256)
    - Generates key_padding_mask -> (Batch, T_max) for Conformer attention
    - Concatenates flat_targets for CTC Loss
    - Pre-computes active_ids (int32) and compact_targets (int64) in background DataLoader workers
      for FusedMetalCosineCTCLoss
    """
    def __init__(
        self,
        pad_feature_value: float = 0.0,
        pad_target_value: int = 0,
        pad_multiple: int = 16,
        max_frames: int = 512
    ):
        self.pad_feature_val = pad_feature_value
        self.pad_target_val = pad_target_value
        self.pad_multiple = pad_multiple
        self.max_frames = max_frames

    def __call__(self, batch: List[Dict[str, Any]]) -> Dict[str, Any]:
        # Filter out corrupted, empty, or CTC-invalid (T < L) entries
        batch = [
            item for item in batch
            if item["feature"].size(0) > 0 and 0 < item["targets"].size(0) <= item["feature"].size(0)
        ]
        if not batch:
            return {}

        # Cap outlier sequence lengths while guaranteeing CTC constraint (T >= target_len)
        if self.max_frames and self.max_frames > 0:
            for item in batch:
                t_len = item["feature"].size(0)
                if t_len > self.max_frames:
                    safe_cap = max(self.max_frames, int(item["targets"].size(0)) + 16)
                    if t_len > safe_cap:
                        item["feature"] = item["feature"][:safe_cap]

        batch_size = len(batch)
        feat_dim = batch[0]["feature"].size(-1)
        dtype = batch[0]["feature"].dtype

        # Extract sequence lengths
        input_lengths = torch.tensor([item["feature"].size(0) for item in batch], dtype=torch.long)
        target_lengths = torch.tensor([item["targets"].size(0) for item in batch], dtype=torch.long)

        actual_max_t = int(input_lengths.max().item())
        max_t = actual_max_t
        if self.pad_multiple > 1:
            max_t = ((max_t + self.pad_multiple - 1) // self.pad_multiple) * self.pad_multiple

        # Pad features: (B, max_t, 256)
        padded_features = torch.full((batch_size, max_t, feat_dim), self.pad_feature_val, dtype=dtype)
        # Key padding mask: True for padded positions (B, max_t)
        key_padding_mask = torch.zeros((batch_size, max_t), dtype=torch.bool)

        for i, item in enumerate(batch):
            t_len = item["feature"].size(0)
            padded_features[i, :t_len] = item["feature"]
            if t_len < max_t:
                key_padding_mask[i, t_len:] = True

        # Flattened targets for standard PyTorch CTCLoss compatibility: (sum(target_lengths),)
        flat_targets = torch.cat([item["targets"] for item in batch], dim=0)

        # Pre-compute unique active target IDs and compact targets in DataLoader worker thread
        uniq_t, inv_idx = torch.unique(flat_targets, sorted=True, return_inverse=True)
        active_ids = torch.cat([torch.zeros(1, dtype=torch.int32), uniq_t.to(torch.int32)])
        compact_targets = inv_idx + 1

        raw_texts = [item.get("raw_text", "") for item in batch]

        return {
            "features": padded_features,          # (B, T_max, 256)
            "input_lengths": input_lengths,        # (B,)
            "max_valid_len": actual_max_t,         # int (unbucketed max T in batch)
            "key_padding_mask": key_padding_mask,  # (B, T_max)
            "flat_targets": flat_targets,          # (sum(L),)
            "active_ids": active_ids,              # (U,) int32
            "compact_targets": compact_targets,    # (sum(L),) int64
            "target_lengths": target_lengths,      # (B,)
            "raw_texts": raw_texts
        }