import os
import sys
from threading import Event, Thread
from typing import Any, List, Literal, Tuple, Union, Optional
from joblib import Parallel, delayed
import numpy as np
import pandas as pd
from janusx._optional_deps import format_missing_dependency_message
try:
    from sklearn.ensemble import GradientBoostingRegressor, RandomForestRegressor
    _HAS_SKLEARN = True
    _SKLEARN_IMPORT_ERROR: Exception | None = None
except Exception as _sklearn_exc:
    GradientBoostingRegressor = None  # type: ignore[assignment]
    RandomForestRegressor = None  # type: ignore[assignment]
    _HAS_SKLEARN = False
    _SKLEARN_IMPORT_ERROR = _sklearn_exc
from janusx.gfreader import (
    load_genotype_chunks,
    inspect_genotype_file,
    load_bed_2bit_packed,
)
from janusx.garfield.logreg import logreg
try:
    from janusx.script._common.progress import ProgressAdapter as _CliProgressAdapter
except Exception:
    _CliProgressAdapter = None

try:
    from tqdm import tqdm as _tqdm
except Exception:
    _tqdm = None


_TQDM_SPINNER_FRAMES: tuple[str, ...] = ("/", "-", "\\", "|")
_TQDM_REFRESH_SECONDS = 0.3
_TQDM_GREEN = "\033[32m"
_TQDM_RESET = "\033[0m"


def _require_sklearn(context: str) -> None:
    if _HAS_SKLEARN:
        return
    raise ImportError(
        format_missing_dependency_message(
            f"scikit-learn is required for {context}.",
            packages=("scikit-learn",),
            extra="ml",
            original_error=_SKLEARN_IMPORT_ERROR,
        )
    ) from _SKLEARN_IMPORT_ERROR


class _TqdmSpinnerAdapter:
    """
    TQDM wrapper that adds a rotating left spinner and green running style.
    """

    def __init__(
        self,
        *,
        total: Union[int, None],
        desc: str,
        unit: str,
        refresh_interval: float = _TQDM_REFRESH_SECONDS,
    ) -> None:
        if _tqdm is None:
            raise RuntimeError("tqdm is unavailable")
        self._desc = str(desc).strip() or "Progress"
        self._idx = 0
        self._refresh_interval = float(max(0.05, float(refresh_interval)))
        self._use_color = bool(getattr(sys.stdout, "isatty", lambda: False)())
        self._stop_event = Event()
        self._thread: Optional[Thread] = None
        kwargs: dict[str, Any] = {
            "total": total,
            "desc": self._render_desc(),
            "unit": str(unit),
            "leave": False,
            "file": sys.stdout,
            "mininterval": self._refresh_interval,
        }
        # Newer tqdm versions support `colour`; ignore silently if unsupported.
        if self._use_color:
            kwargs["colour"] = "green"
        try:
            self._bar = _tqdm(**kwargs)
        except TypeError:
            kwargs.pop("colour", None)
            self._bar = _tqdm(**kwargs)
        self._thread = Thread(target=self._refresh_loop, daemon=True)
        self._thread.start()

    def _render_desc(self) -> str:
        frame = _TQDM_SPINNER_FRAMES[self._idx % len(_TQDM_SPINNER_FRAMES)]
        text = f"{frame} {self._desc}"
        if self._use_color:
            return f"{_TQDM_GREEN}{text}{_TQDM_RESET}"
        return text

    def _refresh_loop(self) -> None:
        while not self._stop_event.wait(self._refresh_interval):
            self._idx = (self._idx + 1) % len(_TQDM_SPINNER_FRAMES)
            try:
                self._bar.set_description_str(self._render_desc(), refresh=False)
                self._bar.refresh()
            except Exception:
                break

    def update(self, n: int = 1) -> None:
        self._bar.update(int(n))

    def close(self) -> None:
        self._stop_event.set()
        if self._thread is not None:
            self._thread.join(timeout=max(1.0, self._refresh_interval * 4.0))
            self._thread = None
        self._bar.close()

try:
    from janusx.janusx import (
        bed_packed_row_flip_mask as _bed_packed_row_flip_mask,
        bed_packed_decode_rows_f32 as _bed_packed_decode_rows_f32,
    )
except Exception:
    _bed_packed_row_flip_mask = None
    _bed_packed_decode_rows_f32 = None

def load_all_genotype_int8(
    genofile: str,
    sampleid: np.ndarray,
    chunk_size: int = int(1e6),
    maf: float = 0.02,
    missing_rate: float = 0.05,
    mmap_window_mb: Union[int, None] = None,
    total_snps: Union[int, None] = None,
) -> tuple[np.ndarray, List[Any]]:
    """
    Load all filtered genotypes once into memory as SNP-major int8 matrix.

    Supported input: PLINK / VCF(.gz) / TXT matrix.
    """
    sample_arr = np.asarray(sampleid, dtype=str)
    low = str(genofile).lower()
    if low.endswith(".vcf") or low.endswith(".vcf.gz"):
        all_ids, n_all_snps = inspect_genotype_file(genofile)
        all_ids = np.asarray(all_ids, dtype=str)
        sameidx = np.isin(all_ids, sample_arr)
        if not np.any(sameidx):
            raise ValueError(
                "No overlapping samples between genotype input and selected sample IDs."
            )
        sample_sub = None
        n_use = int(np.sum(sameidx))
    else:
        _all_ids, n_all_snps = inspect_genotype_file(genofile)
        sameidx = None
        sample_sub = sample_arr.tolist()
        n_use = len(sample_sub)

    chunk_size = max(1, int(chunk_size))
    if total_snps is None:
        total_snps = int(n_all_snps)
    total_chunks = None
    if total_snps is not None and int(total_snps) > 0:
        total_chunks = (int(total_snps) + chunk_size - 1) // chunk_size

    chunks = load_genotype_chunks(
        genofile,
        chunk_size=chunk_size,
        maf=maf,
        missing_rate=missing_rate,
        impute=True,
        sample_ids=sample_sub,
        mmap_window_mb=mmap_window_mb,
    )
    M_list: List[np.ndarray] = []
    sites_list: List[Any] = []
    use_cli_progress = (
        _CliProgressAdapter is not None
        and total_chunks is not None
        and int(total_chunks) > 0
    )
    pbar = (
        _CliProgressAdapter(total=int(total_chunks), desc="Loading genotype")
        if use_cli_progress
        else (
            _tqdm(
                total=total_chunks,
                desc="Loading genotype",
                unit="chunk",
            )
            if _tqdm is not None
            else None
        )
    )
    try:
        for chunk, sites in chunks:
            if chunk.size == 0 or len(sites) == 0:
                if pbar is not None:
                    pbar.update(1)
                continue
            if sameidx is not None:
                chunk = chunk[:, sameidx]
            M_list.append(np.asarray(np.rint(chunk), dtype=np.int8))
            sites_list.extend(sites)
            if pbar is not None:
                pbar.update(1)
    finally:
        if pbar is not None:
            if use_cli_progress:
                pbar.finish()
                pbar.close()
            else:
                pbar.close()

    if not M_list:
        return np.empty((0, n_use), dtype=np.int8), []
    M = np.vstack(M_list)
    return np.ascontiguousarray(M, dtype=np.int8), sites_list


def _normalize_bed_prefix(genofile: str) -> str:
    p = str(genofile)
    low = p.lower()
    if low.endswith(".bed") or low.endswith(".bim") or low.endswith(".fam"):
        return p[:-4]
    return p


def _load_bim_chrom_pos_filtered(prefix: str, keep_mask: np.ndarray) -> list[tuple[str, int]]:
    keep = np.ascontiguousarray(np.asarray(keep_mask, dtype=np.bool_).reshape(-1))
    n_total = int(keep.shape[0])
    out: list[tuple[str, int]] = []
    bim_path = f"{prefix}.bim"
    lines = 0
    with open(bim_path, "r", encoding="utf-8", errors="replace") as fh:
        for raw in fh:
            if lines >= n_total:
                raise ValueError(
                    f"BIM has more rows than packed genotype rows: bim>{n_total} ({bim_path})"
                )
            idx = lines
            lines += 1
            if not keep[idx]:
                continue
            parts = raw.strip().split()
            if len(parts) < 4:
                raise ValueError(
                    f"Invalid BIM row at line {idx + 1}: expected >=4 columns."
                )
            chrom = str(parts[0])
            try:
                pos = int(parts[3])
            except Exception:
                try:
                    pos = int(float(parts[3]))
                except Exception:
                    pos = 0
            out.append((chrom, int(pos)))
    if lines != n_total:
        raise ValueError(
            f"BIM row count mismatch: bim={lines}, packed_rows={n_total} ({bim_path})"
        )
    if len(out) != int(np.sum(keep)):
        raise ValueError("Internal error: filtered BIM rows do not match keep mask.")
    return out


def load_all_genotype_packed_bed(
    genofile: str,
    maf: float = 0.02,
    missing_rate: float = 0.05,
) -> tuple[dict[str, Any], list[tuple[str, int]]]:
    """
    Load all filtered genotypes once into memory as BED-packed 2-bit matrix.

    Returned packed context can be decoded region-wise for selected samples
    without materializing a full SNP x sample int8 matrix.
    """
    if _bed_packed_row_flip_mask is None or _bed_packed_decode_rows_f32 is None:
        raise RuntimeError(
            "Rust packed BED decode helpers are unavailable. Rebuild JanusX extension."
        )
    prefix = _normalize_bed_prefix(genofile)
    if not (
        os.path.isfile(f"{prefix}.bed")
        and os.path.isfile(f"{prefix}.bim")
        and os.path.isfile(f"{prefix}.fam")
    ):
        raise ValueError(f"PLINK BED/BIM/FAM prefix not found: {prefix}")

    packed_raw, miss_raw, maf_raw, _std_raw, n_samples = load_bed_2bit_packed(prefix)
    packed_arr = np.ascontiguousarray(np.asarray(packed_raw, dtype=np.uint8))
    miss_arr = np.ascontiguousarray(np.asarray(miss_raw, dtype=np.float32).reshape(-1))
    maf_arr = np.ascontiguousarray(np.asarray(maf_raw, dtype=np.float32).reshape(-1))
    if packed_arr.shape[0] != maf_arr.shape[0] or miss_arr.shape[0] != maf_arr.shape[0]:
        raise ValueError("Packed BED arrays have inconsistent SNP dimensions.")

    keep = np.ones(maf_arr.shape[0], dtype=np.bool_)
    maf_thr = float(maf)
    if maf_thr > 0.0:
        keep &= (maf_arr >= maf_thr) & (maf_arr <= (1.0 - maf_thr))
    miss_thr = float(missing_rate)
    if miss_thr < 1.0:
        keep &= miss_arr <= miss_thr
    if not np.any(keep):
        raise ValueError("No SNP remains after BED packed filtering.")

    all_sites = _load_bim_chrom_pos_filtered(prefix, keep)
    if np.all(keep):
        packed = packed_arr
        maf_keep = maf_arr
        miss_keep = miss_arr
    else:
        packed = np.ascontiguousarray(packed_arr[keep], dtype=np.uint8)
        maf_keep = np.ascontiguousarray(maf_arr[keep], dtype=np.float32)
        miss_keep = np.ascontiguousarray(miss_arr[keep], dtype=np.float32)

    row_flip = np.ascontiguousarray(
        np.asarray(_bed_packed_row_flip_mask(packed, int(n_samples)), dtype=np.bool_).reshape(-1)
    )
    if row_flip.shape[0] != packed.shape[0]:
        raise ValueError("Packed row_flip length mismatch.")

    packed_ctx = {
        "packed": packed,
        "maf": maf_keep,
        "missing_rate": miss_keep,
        "row_flip": row_flip,
        "n_samples": int(n_samples),
    }
    return packed_ctx, all_sites


def _decode_packed_rows_int8(
    packed_ctx: dict[str, Any],
    row_indices: np.ndarray,
    sample_indices: np.ndarray,
) -> np.ndarray:
    if _bed_packed_decode_rows_f32 is None:
        raise RuntimeError("Rust packed BED decode helper is unavailable.")

    rows = np.ascontiguousarray(np.asarray(row_indices, dtype=np.int64).reshape(-1))
    sidx = np.ascontiguousarray(np.asarray(sample_indices, dtype=np.int64).reshape(-1))
    if rows.size == 0:
        return np.empty((0, int(sidx.size)), dtype=np.int8)

    decoded = _bed_packed_decode_rows_f32(
        packed_ctx["packed"],
        int(packed_ctx["n_samples"]),
        rows,
        packed_ctx["row_flip"],
        packed_ctx["maf"],
        sidx,
    )
    decoded_arr = np.asarray(decoded, dtype=np.float32)
    # Keep GARFIELD behavior consistent with dense preload path:
    # imputed hard-calls are rounded to int8.
    return np.ascontiguousarray(np.rint(decoded_arr), dtype=np.int8)

def ldprune(
    M: np.ndarray,
    sites,
    thr: float = 0.8,
    max_snps: int = 600,
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Lightweight LD pruning:
      1) Sort SNPs by variance in descending order (prefer informative SNPs)
      2) Iteratively add candidates and compute correlation only vs kept SNPs
      3) If LD(r^2) >= thr with kept SNPs, keep the SNP with larger MAF
      4) Keep at most ``max_snps`` SNPs

    Parameters
    ----------
    M : (m_snps, n_samples) genotype matrix
    sites : site metadata aligned with SNP rows
    thr : LD threshold on r^2 (r^2 >= thr is considered high LD)
    max_snps : maximum number of SNPs to keep

    Returns
    -------
    M_kept, sites_kept
    """
    m, n = M.shape
    if m <= 1:
        return M, sites

    # 1) Remove zero-variance SNPs (monomorphic).
    var = M.var(axis=1)
    non_mono = var > 0
    M = M[non_mono]
    sites = np.array(sites)[non_mono]
    var = var[non_mono]
    m = M.shape[0]
    if m <= 1:
        return M, sites

    # 2) Compute MAF (larger means closer to 0.5).
    af = np.clip(M.mean(axis=1, dtype=np.float64) / 2.0, 0.0, 1.0)
    maf = np.minimum(af, 1.0 - af)

    # 3) Sort by variance (larger variance first).
    order = np.argsort(var)[::-1]
    M = M[order]
    sites = sites[order]
    maf = maf[order]

    # 4) Standardize each SNP row to zero mean and unit variance.
    mu = M.mean(axis=1, keepdims=True)
    sigma = M.std(axis=1, keepdims=True)
    sigma[sigma == 0] = 1.0
    Z = (M - mu) / sigma  # (m, n)

    kept_idx = []
    for i in range(m):
        if len(kept_idx) == 0:
            kept_idx.append(i)
        else:
            # Compute correlation only against already kept SNPs.
            z_i = Z[i]  # (n,)
            Z_kept = Z[kept_idx]  # (k, n)
            # corr = dot / n because rows are standardized.
            corr = (Z_kept @ z_i) / n
            r2 = corr * corr
            conflict_mask = r2 >= thr
            if not np.any(conflict_mask):
                kept_idx.append(i)
            else:
                # For high-LD redundancy, keep the SNP with larger MAF.
                conflict_pos = np.flatnonzero(conflict_mask)
                conflict_kept_idx = np.array([kept_idx[p] for p in conflict_pos], dtype=int)
                if conflict_kept_idx.size > 0 and float(maf[i]) > float(np.max(maf[conflict_kept_idx])):
                    for p in sorted(conflict_pos.tolist(), reverse=True):
                        kept_idx.pop(p)
                    kept_idx.append(i)
        if len(kept_idx) >= max_snps:
            break

    kept_idx = np.array(kept_idx, dtype=int)
    return M[kept_idx], sites[kept_idx]

def getLogicgate(
    y: np.ndarray,
    M: np.ndarray,
    nsnp: int = 5,
    n_estimators: int = 200,
    sites: Any = None,
    core:Literal['rf','gbdt']='gbdt',
    response:Literal['binary', 'continuous']="continuous",
    gsetmode:bool = True,
    gset_groups: Union[None, List[np.ndarray], List[list[int]]] = None,
):
    """
    Find logic combinations (xcombine) in a genotype block using
    random forest + permutation importance + logistic regression.
    """
    _require_sklearn("GARFIELD core models (rf/gbdt)")
    if sites is not None:
        if len(sites) == 0:
            sites = np.arange(M.shape[0])
        else:
            first = sites[0]
            if hasattr(first, "chrom") and hasattr(first, "pos"):
                sites = np.array([f"{i.chrom}_{i.pos}" for i in sites])
            elif isinstance(first, (tuple, list)) and len(first) >= 2:
                sites = np.array([f"{str(i[0])}_{int(i[1])}" for i in sites])
            else:
                sites = np.array([str(i) for i in sites])
    else:
        sites = np.arange(M.shape[0])

    # If gset_groups are provided, normalize and apply ldprune within each group.
    if gset_groups is not None:
        group_indices: List[np.ndarray] = []
        for g in gset_groups:
            arr = np.asarray(g, dtype=int)
            if arr.size > 0:
                group_indices.append(arr)
        if len(group_indices) == 0:
            return None

        kept_union: List[int] = []
        for g_idx in group_indices:
            g_idx = g_idx[(g_idx >= 0) & (g_idx < M.shape[0])]
            if g_idx.size == 0:
                continue
            M_sub = M[g_idx]
            sites_sub = sites[g_idx]
            M_kept, sites_kept = ldprune(M_sub, sites_sub)
            # Map kept indices back to global indices via site strings
            kept_set = set(sites_kept.tolist())
            kept_global = [int(i) for i, s in enumerate(sites) if s in kept_set]
            kept_union.extend(kept_global)

        kept_union = sorted(set(kept_union))
        if len(kept_union) == 0:
            return None
        M = M[kept_union]
        sites = sites[kept_union]
        # Remap group indices to new local indices
        index_map = {old: new for new, old in enumerate(kept_union)}
        gset_groups = [
            np.array([index_map[i] for i in g if i in index_map], dtype=int)
            for g in gset_groups
        ]
        gset_groups = [g for g in gset_groups if g.size > 0]
        if len(gset_groups) == 0:
            return None
    else:
        M,sites = ldprune(M,sites)
    if core == 'rf':
        model = RandomForestRegressor(
            n_estimators=n_estimators,
            max_depth=nsnp,
            min_samples_leaf=nsnp,
            bootstrap=True,
            n_jobs=1,
            random_state=0,
        )
    elif core == 'gbdt':
        model = GradientBoostingRegressor(
            n_estimators=n_estimators,
            max_depth=nsnp,
            learning_rate=0.05,
            subsample=0.7,
            random_state=0,
        )
    model.fit(M.T, y)
    Imp = model.feature_importances_

    topk = nsnp
    if gsetmode and gset_groups is not None:
        order = np.argsort(Imp)[::-1]
        idx = []
        for g in gset_groups:
            if g.size == 0:
                continue
            # pick highest-importance SNP within this group
            best = g[np.argmax(Imp[g])]
            idx.append(best)
        idx = np.array(sorted(set(idx)), dtype=int)[:topk]  # at most k interaction tests
        if idx.size == 0:
            return None
    else:
        idx = np.argsort(Imp)[::-1][:topk]
    # 0/1/2 ->0/1
    Mchoice = (M[idx] / 2).astype(int).T

    resdict = logreg(Mchoice, y, response=response, tags=sites[idx])
    if resdict is None:
        return None
    indices = resdict.get("indices", [])
    if indices is None or len(indices) <= 1:
        return None
    if resdict.get("expression", "") == "1":
        return None
    return resdict


# ---------- Whole-genome traversal ----------
def _window_task(
    task: Tuple[str, int, int, np.ndarray, List[str]],
    y: np.ndarray,
    nsnp: int,
    n_estimators: int,
    response: Literal['binary', 'continuous'],
    gsetmode: bool,
) -> tuple[Union[dict[str,Any], None], Union[str, None]]:
    chrom, start, end, M_win, tags = task
    try:
        resdict = getLogicgate(
            y,
            M_win,
            nsnp=nsnp,
            n_estimators=n_estimators,
            sites=tags,
            response=response,
            gsetmode=gsetmode,
        )
    except Exception as e:
        emsg = str(e)
        if (
            "Found array with 0 feature(s)" in emsg
            and "minimum of 1 is required" in emsg
        ):
            return None, None
        return None, f"window {chrom}:{start}-{end} failed: {e}"
    if resdict is None:
        return None, None
    return resdict, None


def _resolve_stream_sample_subset(
    genofile: str,
    sampleid: np.ndarray,
) -> tuple[Union[list[str], None], Union[np.ndarray, None]]:
    """
    Keep streaming sample handling consistent with gwas.py:
    - For VCF input: read all samples in chunk reader, then subset by boolean mask.
    - For non-VCF input: pass sample_ids directly to chunk reader.
    """
    sample_arr = np.asarray(sampleid, dtype=str)
    low = str(genofile).lower()
    if low.endswith(".vcf") or low.endswith(".vcf.gz"):
        all_ids, _ = inspect_genotype_file(genofile)
        all_ids = np.asarray(all_ids, dtype=str)
        sameidx = np.isin(all_ids, sample_arr)
        if not np.any(sameidx):
            raise ValueError(
                "No overlapping samples between genotype input and selected sample IDs."
            )
        return None, sameidx
    return sample_arr.tolist(), None


def window(
    genofile: str,
    sampleid: np.ndarray,
    y: np.ndarray,
    step: int,
    windowsize: int,
    chunk_size: Union[int,float] = 1e5,
    maf: float = 0.02,
    missing_rate: float = 0.05,
    nsnp: int = 5,
    n_estimators: int = 200,
    response: Literal['binary', 'continuous'] = "continuous",
    gsetmode: bool = False,
    threads: int = 4,
    batch_size: int = 128,
    mmap_window_mb: Union[int, None] = None,
    collect_warnings: bool = False,
):
    """
    Whole-genome sliding-window traversal (load by chunksize order to reduce repeated IO).

    Window definition:
      - Move window start by ``step`` each time
      - Each window interval is [start, start + windowsize)

    Streaming strategy:
      - Keep at most two windows of data in memory
      - After processing two windows, drop the first and load the next

    Returns:
      list of dict
    """
    y = np.asarray(y, dtype=float).ravel()
    tasks: List[Tuple[str, int, int, np.ndarray, List[str]]] = []
    results: List[dict[str,Any]] = []

    total_windows = None
    step_int = int(step) if step else 0
    if step_int > 0 and not (str(genofile).endswith(".vcf") or str(genofile).endswith(".vcf.gz")):
        bim_path = f"{genofile}.bim"
        if os.path.isfile(bim_path):
            try:
                bim = pd.read_csv(bim_path, sep=r"\s+", header=None, usecols=[0, 3])
                if not bim.empty:
                    chr_minmax = bim.groupby(0)[3].agg(["min", "max"])
                    counts = (chr_minmax["max"] - chr_minmax["min"]) // step_int + 1
                    total_windows = int(counts.sum())
            except Exception:
                total_windows = None
    use_cli_progress = (
        _CliProgressAdapter is not None
        and total_windows is not None
        and int(total_windows) > 0
    )
    pbar = (
        _CliProgressAdapter(total=int(total_windows), desc="Loading windows")
        if use_cli_progress
        else (
            _TqdmSpinnerAdapter(total=total_windows, desc="Windows", unit="win")
            if _tqdm is not None
            else None
        )
    )
    sample_sub, sameidx = _resolve_stream_sample_subset(genofile, sampleid)

    window_warnings: List[str] = []

    def _consume_task_outputs(task_outputs: List[tuple[Union[dict[str,Any], None], Union[str, None]]]) -> None:
        for out, warn in task_outputs:
            if warn is not None:
                if collect_warnings:
                    window_warnings.append(str(warn))
                else:
                    print(f"[WARN] {warn}")
            if out is not None:
                results.append(out)

    buffer_M = None
    buffer_sites: List[Any] = []
    buffer_pos = None
    current_chrom = None
    window_start = None
    next_start = None

    def _append_segment(seg_M, seg_sites):
        nonlocal buffer_M, buffer_sites, buffer_pos
        if buffer_M is None:
            buffer_M = seg_M
            buffer_sites = list(seg_sites)
        else:
            buffer_M = np.vstack([buffer_M, seg_M])
            buffer_sites.extend(seg_sites)
        pos_arr = np.array([int(s.pos) for s in buffer_sites], dtype=np.int64)
        buffer_pos = pos_arr

    def _drop_before(pos_thr):
        nonlocal buffer_M, buffer_sites, buffer_pos
        if buffer_M is None:
            return
        mask = buffer_pos >= pos_thr
        if mask.sum() == 0:
            buffer_M = None
            buffer_sites = []
            buffer_pos = None
            return
        buffer_M = buffer_M[mask]
        buffer_sites = [s for i, s in enumerate(buffer_sites) if mask[i]]
        buffer_pos = buffer_pos[mask]

    def _emit_window(start_pos):
        if buffer_M is None:
            return None
        end_pos = start_pos + windowsize
        mask = (buffer_pos >= start_pos) & (buffer_pos < end_pos)
        if mask.sum() < 2:
            return None
        M_win = buffer_M[mask]
        sites_win = [s for i, s in enumerate(buffer_sites) if mask[i]]
        tags = [f"{s.chrom}_{s.pos}" for s in sites_win]
        return current_chrom, start_pos, end_pos, M_win, tags

    chunks = load_genotype_chunks(
        genofile,
        chunk_size=int(chunk_size),
        maf=maf,
        missing_rate=missing_rate,
        impute=True,
        sample_ids=sample_sub,
        mmap_window_mb=mmap_window_mb,
    )

    try:
        for chunk, sites in chunks:
            if not sites:
                continue
            if sameidx is not None:
                chunk = chunk[:, sameidx]
            # split by chromosome boundaries within the chunk
            idx = 0
            while idx < len(sites):
                chrom = str(sites[idx].chrom)
                j = idx + 1
                while j < len(sites) and str(sites[j].chrom) == chrom:
                    j += 1
                seg_sites = sites[idx:j]
                seg_M = chunk[idx:j]

                if current_chrom is None or chrom != current_chrom:
                    # flush remaining windows for previous chromosome
                    if current_chrom is not None and buffer_pos is not None:
                        max_pos = int(buffer_pos.max())
                        while window_start is not None and window_start <= max_pos:
                            res = _emit_window(window_start)
                            if pbar is not None:
                                pbar.update(1)
                            if res is not None:
                                tasks.append(res)
                                if len(tasks) >= batch_size:
                                    if threads <= 1:
                                        batch_outputs = [
                                            _window_task(t, y, nsnp, n_estimators, response, gsetmode)
                                            for t in tasks
                                        ]
                                    else:
                                        batch_outputs = Parallel(n_jobs=threads, backend="loky")(
                                            delayed(_window_task)(t, y, nsnp, n_estimators, response, gsetmode) for t in tasks
                                        )
                                    _consume_task_outputs(batch_outputs)
                                    tasks.clear()
                            window_start += step
                    # reset buffers for new chromosome
                    buffer_M = None
                    buffer_sites = []
                    buffer_pos = None
                    current_chrom = chrom
                    window_start = int(seg_sites[0].pos)
                    next_start = window_start + step

                _append_segment(seg_M, seg_sites)

                # process windows when we have two windows worth of data
                while buffer_pos is not None and buffer_pos.max() >= window_start + windowsize + step:
                    res = _emit_window(window_start)
                    if pbar is not None:
                        pbar.update(1)
                    if res is not None:
                        tasks.append(res)
                        if len(tasks) >= batch_size:
                            if threads <= 1:
                                batch_outputs = [
                                    _window_task(t, y, nsnp, n_estimators, response, gsetmode)
                                    for t in tasks
                                ]
                            else:
                                batch_outputs = Parallel(n_jobs=threads, backend="loky")(
                                    delayed(_window_task)(t, y, nsnp, n_estimators, response, gsetmode) for t in tasks
                                )
                            _consume_task_outputs(batch_outputs)
                            tasks.clear()
                    # drop the first window, advance
                    _drop_before(next_start)
                    window_start = next_start
                    next_start = window_start + step

                idx = j

        # flush remaining windows at end
        if current_chrom is not None and buffer_pos is not None:
            max_pos = int(buffer_pos.max())
            while window_start is not None and window_start <= max_pos:
                res = _emit_window(window_start)
                if pbar is not None:
                    pbar.update(1)
                if res is not None:
                    tasks.append(res)
                    if len(tasks) >= batch_size:
                        if threads <= 1:
                            batch_outputs = [
                                _window_task(t, y, nsnp, n_estimators, response, gsetmode)
                                for t in tasks
                            ]
                        else:
                            batch_outputs = Parallel(n_jobs=threads, backend="loky")(
                                delayed(_window_task)(t, y, nsnp, n_estimators, response, gsetmode) for t in tasks
                            )
                        _consume_task_outputs(batch_outputs)
                        tasks.clear()
                window_start += step

        if tasks:
            if threads <= 1:
                batch_outputs = [
                    _window_task(t, y, nsnp, n_estimators, response, gsetmode)
                    for t in tasks
                ]
            else:
                batch_outputs = Parallel(n_jobs=threads, backend="loky")(
                    delayed(_window_task)(t, y, nsnp, n_estimators, response, gsetmode) for t in tasks
                )
            _consume_task_outputs(batch_outputs)
    finally:
        if pbar is not None:
            if use_cli_progress:
                pbar.finish()
                pbar.close()
            else:
                pbar.close()

    if collect_warnings:
        return results, window_warnings
    return results


# ---------- Custom traversal ----------
chrposTuple = Tuple[str, int, int]
def _process(
    ChromPos: Union[chrposTuple,List[chrposTuple]],
    genofile: str,
    sample_sub: Union[list[str], None],
    sameidx: Union[np.ndarray, None],
    y: np.ndarray,
    maf: float = 0.02,
    missing_rate: float = 0.05,
    nsnp: int = 5,
    n_estimators: int = 200,
    response: Literal['binary', 'continuous'] = "continuous",
    gsetmode: bool = True,
    mmap_window_mb: Union[int, None] = None,
):
    """
    Run on one BED region:
      1) Extract all SNP genotypes from ``genofile`` in the region
      2) Call ``getLogicgate``
      3) Return ``(xcombine, site_tuple)`` or ``None``
    """
    # mmap_window_mb is not supported with bim_range; disable it here
    mmap_window_mb = None
    if isinstance(ChromPos,tuple):
        chrom, start, end = ChromPos
        chunks = load_genotype_chunks(
            genofile,
            chunk_size=int(1e6),
            maf=maf,
            missing_rate=missing_rate,
            impute=True,
            bim_range=(str(chrom), int(start), int(end)),
            sample_ids=sample_sub,
        )

        M_list = []
        sites_list = []

        for chunk, sites in chunks:
            # chunk: (m_chunk, n_samples)
            if sameidx is not None:
                chunk = chunk[:, sameidx]
            M_list.append(chunk)
            sites_list.extend(sites)
        if not M_list:
            return None
    elif isinstance(ChromPos,list):
        M_list = []
        sites_list = []
        gset_groups: List[np.ndarray] = []
        offset = 0
        # limit per-gene SNPs to bound memory
        max_gene_snps = max(200, int(nsnp) * 60)
        for chrom, start, end in ChromPos:
            chunks = load_genotype_chunks(
                genofile,
                chunk_size=int(1e6),
                maf=maf,
                missing_rate=missing_rate,
                impute=True,
                bim_range=(str(chrom), int(start), int(end)),
                sample_ids=sample_sub,
            )
            M_gene = None
            sites_gene: List[Any] = []
            for chunk, sites in chunks:
                if chunk.size == 0:
                    continue
                if sameidx is not None:
                    chunk = chunk[:, sameidx]
                if M_gene is None:
                    M_gene = chunk
                    sites_gene = list(sites)
                else:
                    M_gene = np.vstack([M_gene, chunk])
                    sites_gene.extend(sites)
                # prune periodically to bound memory per gene
                if M_gene.shape[0] > max_gene_snps:
                    M_gene, sites_kept = ldprune(
                        M_gene, np.array(sites_gene), max_snps=max_gene_snps
                    )
                    sites_gene = list(sites_kept)
            if M_gene is None or M_gene.shape[0] == 0:
                continue
            # final prune to limit per-gene SNPs
            if M_gene.shape[0] > max_gene_snps:
                M_gene, sites_kept = ldprune(
                    M_gene, np.array(sites_gene), max_snps=max_gene_snps
                )
                sites_gene = list(sites_kept)
            if M_gene.shape[0] == 0:
                continue

            M_list.append(M_gene)
            sites_list.extend(sites_gene)
            gset_groups.append(np.arange(offset, offset + M_gene.shape[0], dtype=int))
            offset += M_gene.shape[0]
        if not M_list:
            return None

    M = np.vstack(M_list)
    if M.shape[0] < 2:
        return None
    try:
        resdict = getLogicgate(
            y,
            M,
            nsnp=nsnp,
            n_estimators=n_estimators,
            sites=sites_list,
            response=response,
            gsetmode=gsetmode,
            gset_groups=gset_groups if gsetmode and isinstance(ChromPos, list) else None,
        )
    except Exception as e:
        # If one region fails, do not crash the whole parallel job.
        print(f"[WARN] process {ChromPos} failed: {e}")
        return None
    if resdict is None:
        return None
    return resdict

def _build_site_index(
    sites: List[Any],
) -> dict[str, tuple[np.ndarray, np.ndarray, bool]]:
    """
    Build per-chromosome positional index for fast range query.
    Returns {chrom: (pos, global_indices)} in original SNP order.
    """
    def _site_chrom_pos(site: Any) -> tuple[str, int]:
        if hasattr(site, "chrom") and hasattr(site, "pos"):
            return str(site.chrom), int(site.pos)
        if isinstance(site, (tuple, list)) and len(site) >= 2:
            return str(site[0]), int(site[1])
        s = str(site)
        if "_" in s:
            chrom, pos = s.rsplit("_", 1)
            try:
                return str(chrom), int(pos)
            except Exception:
                pass
        raise ValueError(f"Unsupported site format: {site!r}")

    by_chrom: dict[str, list[tuple[int, int]]] = {}
    for i, site in enumerate(sites):
        chrom, pos = _site_chrom_pos(site)
        if chrom not in by_chrom:
            by_chrom[chrom] = []
        by_chrom[chrom].append((pos, i))

    out: dict[str, tuple[np.ndarray, np.ndarray, bool]] = {}
    for chrom, pairs in by_chrom.items():
        pos = np.fromiter((p for p, _ in pairs), dtype=np.int64, count=len(pairs))
        idx = np.fromiter((j for _, j in pairs), dtype=np.int64, count=len(pairs))
        is_sorted = bool(pos.size <= 1 or np.all(pos[1:] >= pos[:-1]))
        out[chrom] = (pos, idx, is_sorted)
    return out

def _select_range_indices(
    site_index: dict[str, tuple[np.ndarray, np.ndarray, bool]],
    chrom: str,
    start: int,
    end: int,
) -> np.ndarray:
    item = site_index.get(str(chrom))
    if item is None:
        return np.empty(0, dtype=np.int64)
    pos, idx, is_sorted = item
    if not is_sorted:
        mask = (pos >= int(start)) & (pos <= int(end))
        return idx[mask]
    lo = int(np.searchsorted(pos, int(start), side="left"))
    hi = int(np.searchsorted(pos, int(end), side="right"))
    if hi <= lo:
        return np.empty(0, dtype=np.int64)
    return idx[lo:hi]

def _process_inmemory(
    ChromPos: Union[chrposTuple, List[chrposTuple]],
    all_genotype_i8: np.ndarray,
    all_sites: List[Any],
    site_index: dict[str, tuple[np.ndarray, np.ndarray, bool]],
    y: np.ndarray,
    nsnp: int = 5,
    n_estimators: int = 200,
    response: Literal['binary', 'continuous'] = "continuous",
    gsetmode: bool = True,
):
    if isinstance(ChromPos, tuple):
        chrom, start, end = ChromPos
        idx = _select_range_indices(site_index, str(chrom), int(start), int(end))
        if idx.size == 0:
            return None
        M = all_genotype_i8[idx]
        sites_list = [all_sites[int(i)] for i in idx]
        gset_groups = None
    elif isinstance(ChromPos, list):
        M_list: List[np.ndarray] = []
        sites_list: List[Any] = []
        gset_groups: List[np.ndarray] = []
        offset = 0
        max_gene_snps = max(200, int(nsnp) * 60)

        for chrom, start, end in ChromPos:
            idx = _select_range_indices(site_index, str(chrom), int(start), int(end))
            if idx.size == 0:
                continue
            M_gene = all_genotype_i8[idx]
            sites_gene: List[Any] = [all_sites[int(i)] for i in idx]
            if M_gene.shape[0] > max_gene_snps:
                M_gene, sites_kept = ldprune(
                    M_gene, np.array(sites_gene), max_snps=max_gene_snps
                )
                sites_gene = list(sites_kept)
            if M_gene.shape[0] == 0:
                continue

            M_list.append(M_gene)
            sites_list.extend(sites_gene)
            gset_groups.append(np.arange(offset, offset + M_gene.shape[0], dtype=int))
            offset += M_gene.shape[0]

        if not M_list:
            return None
        M = np.vstack(M_list)
    else:
        return None

    if M.shape[0] < 2:
        return None
    try:
        resdict = getLogicgate(
            y,
            M,
            nsnp=nsnp,
            n_estimators=n_estimators,
            sites=sites_list,
            response=response,
            gsetmode=gsetmode,
            gset_groups=gset_groups if gsetmode and isinstance(ChromPos, list) else None,
        )
    except Exception as e:
        print(f"[WARN] process {ChromPos} failed: {e}")
        return None
    if resdict is None:
        return None
    return resdict


def _process_inmemory_packed(
    ChromPos: Union[chrposTuple, List[chrposTuple]],
    packed_ctx: dict[str, Any],
    all_sites: List[Any],
    site_index: dict[str, tuple[np.ndarray, np.ndarray, bool]],
    sample_indices: np.ndarray,
    y: np.ndarray,
    nsnp: int = 5,
    n_estimators: int = 200,
    response: Literal['binary', 'continuous'] = "continuous",
    gsetmode: bool = True,
):
    if isinstance(ChromPos, tuple):
        chrom, start, end = ChromPos
        idx = _select_range_indices(site_index, str(chrom), int(start), int(end))
        if idx.size == 0:
            return None
        M = _decode_packed_rows_int8(packed_ctx, idx, sample_indices)
        sites_list = [all_sites[int(i)] for i in idx]
        gset_groups = None
    elif isinstance(ChromPos, list):
        M_list: List[np.ndarray] = []
        sites_list: List[Any] = []
        gset_groups: List[np.ndarray] = []
        offset = 0
        max_gene_snps = max(200, int(nsnp) * 60)

        for chrom, start, end in ChromPos:
            idx = _select_range_indices(site_index, str(chrom), int(start), int(end))
            if idx.size == 0:
                continue
            M_gene = _decode_packed_rows_int8(packed_ctx, idx, sample_indices)
            sites_gene: List[Any] = [all_sites[int(i)] for i in idx]
            if M_gene.shape[0] > max_gene_snps:
                M_gene, sites_kept = ldprune(
                    M_gene, np.array(sites_gene), max_snps=max_gene_snps
                )
                sites_gene = list(sites_kept)
            if M_gene.shape[0] == 0:
                continue

            M_list.append(M_gene)
            sites_list.extend(sites_gene)
            gset_groups.append(np.arange(offset, offset + M_gene.shape[0], dtype=int))
            offset += M_gene.shape[0]

        if not M_list:
            return None
        M = np.vstack(M_list)
    else:
        return None

    if M.shape[0] < 2:
        return None
    try:
        resdict = getLogicgate(
            y,
            M,
            nsnp=nsnp,
            n_estimators=n_estimators,
            sites=sites_list,
            response=response,
            gsetmode=gsetmode,
            gset_groups=gset_groups if gsetmode and isinstance(ChromPos, list) else None,
        )
    except Exception as e:
        print(f"[WARN] process {ChromPos} failed: {e}")
        return None
    if resdict is None:
        return None
    return resdict

# ---------- Main (collect results in one pass) ----------

def main(
    genofile,
    sampleid,
    y,
    bedlist,
    maf: float = 0.02,
    missing_rate: float = 0.05,
    nsnp=5,
    n_estimators=200,
    threads=4,
    response: Literal['binary', 'continuous'] = "continuous",
    gsetmode: bool = True,
    gsetmodes: Optional[List[bool]] = None,
    mmap_window_mb: Union[int, None] = None,
):
    y = np.asarray(y, dtype=float).ravel()
    if gsetmodes is not None and len(gsetmodes) != len(bedlist):
        raise ValueError("gsetmodes length must match bedlist length")
    sample_sub, sameidx = _resolve_stream_sample_subset(genofile, sampleid)
    results = Parallel(n_jobs=threads)(
        delayed(_process)(
            ChromPos,
            genofile,
            sample_sub,
            sameidx,
            y,
            maf,
            missing_rate,
            nsnp,
            n_estimators,
            response,
            gsetmodes[i] if gsetmodes is not None else gsetmode,
            mmap_window_mb,
        )
        for i, ChromPos in enumerate(bedlist)
    )
    return results

def main_inmemory(
    all_genotype_i8: np.ndarray,
    all_sites: List[Any],
    y: np.ndarray,
    bedlist,
    nsnp: int = 5,
    n_estimators: int = 200,
    threads: int = 4,
    response: Literal['binary', 'continuous'] = "continuous",
    gsetmode: bool = True,
    gsetmodes: Optional[List[bool]] = None,
):
    """
    Run GARFIELD search from preloaded in-memory int8 genotype matrix.
    """
    y = np.asarray(y, dtype=float).ravel()
    all_genotype_i8 = np.asarray(all_genotype_i8, dtype=np.int8)
    if all_genotype_i8.ndim != 2:
        raise ValueError("all_genotype_i8 must be a 2D SNP-major matrix.")
    if len(all_sites) != all_genotype_i8.shape[0]:
        raise ValueError("all_sites length must match all_genotype_i8.shape[0].")
    if gsetmodes is not None and len(gsetmodes) != len(bedlist):
        raise ValueError("gsetmodes length must match bedlist length")

    site_index = _build_site_index(all_sites)
    results = Parallel(n_jobs=threads, backend="threading")(
        delayed(_process_inmemory)(
            ChromPos,
            all_genotype_i8,
            all_sites,
            site_index,
            y,
            nsnp,
            n_estimators,
            response,
            gsetmodes[i] if gsetmodes is not None else gsetmode,
        )
        for i, ChromPos in enumerate(bedlist)
    )
    return results


def main_inmemory_packed(
    packed_ctx: dict[str, Any],
    all_sites: List[Any],
    sample_indices: np.ndarray,
    y: np.ndarray,
    bedlist,
    nsnp: int = 5,
    n_estimators: int = 200,
    threads: int = 4,
    response: Literal['binary', 'continuous'] = "continuous",
    gsetmode: bool = True,
    gsetmodes: Optional[List[bool]] = None,
):
    """
    Run GARFIELD search from preloaded BED-packed genotype matrix.
    """
    y = np.asarray(y, dtype=float).ravel()
    sidx = np.ascontiguousarray(np.asarray(sample_indices, dtype=np.int64).reshape(-1))
    if sidx.size != y.shape[0]:
        raise ValueError(
            f"sample_indices length mismatch: {sidx.size} != len(y)={y.shape[0]}"
        )
    packed = np.asarray(packed_ctx.get("packed"))
    if packed.ndim != 2:
        raise ValueError("packed_ctx['packed'] must be a 2D SNP-major packed matrix.")
    if len(all_sites) != int(packed.shape[0]):
        raise ValueError("all_sites length must match packed_ctx['packed'].shape[0].")
    if gsetmodes is not None and len(gsetmodes) != len(bedlist):
        raise ValueError("gsetmodes length must match bedlist length")

    site_index = _build_site_index(all_sites)
    results = Parallel(n_jobs=threads, backend="threading")(
        delayed(_process_inmemory_packed)(
            ChromPos,
            packed_ctx,
            all_sites,
            site_index,
            sidx,
            y,
            nsnp,
            n_estimators,
            response,
            gsetmodes[i] if gsetmodes is not None else gsetmode,
        )
        for i, ChromPos in enumerate(bedlist)
    )
    return results
