"""Type stubs for bunbun — Fast, unified genotype parser."""

from __future__ import annotations

from typing import Iterator, Optional

import numpy as np
import numpy.typing as npt
import polars as pl


class GenotypeData:
    """Unified genotype container returned by all bunbun readers."""

    @property
    def samples(self) -> pl.DataFrame: ...
    @property
    def variants(self) -> pl.DataFrame: ...
    @property
    def dosages(self) -> npt.NDArray[np.float32]: ...
    @property
    def shape(self) -> tuple[int, int]: ...
    @property
    def n_samples(self) -> int: ...
    @property
    def n_variants(self) -> int: ...
    @property
    def source_format(self) -> str: ...

    def allele_frequencies(self) -> npt.NDArray[np.float32]: ...
    def missing_rates(self) -> npt.NDArray[np.float32]: ...
    def mean_impute(self) -> None: ...
    def subset_samples(self, ids: list[str]) -> GenotypeData: ...
    def region(self, chrom: str, start: int, end: int) -> GenotypeData: ...


class GenotypeChunk:
    """A single chunk yielded by a GenotypeStream."""

    @property
    def variants(self) -> pl.DataFrame: ...
    @property
    def dosages(self) -> npt.NDArray[np.float32]: ...
    @property
    def n_samples(self) -> int: ...
    @property
    def n_variants(self) -> int: ...
    @property
    def offset(self) -> int:
        """Global variant index of the first variant in this chunk."""
        ...
    @property
    def total_variants(self) -> int:
        """Total variants in the full dataset."""
        ...
    @property
    def progress(self) -> float:
        """Progress fraction in [0, 1]."""
        ...
    @property
    def is_last(self) -> bool: ...


class GenotypeStream(Iterator[GenotypeChunk]):
    """Lazy iterator over genotype chunks. Only one chunk in memory at a time."""

    @property
    def samples(self) -> pl.DataFrame:
        """Sample metadata (available before iteration starts)."""
        ...
    @property
    def source_format(self) -> str: ...

    def __iter__(self) -> GenotypeStream: ...
    def __next__(self) -> GenotypeChunk: ...


# ── Full read ─────────────────────────────────────────────

def read(path: str) -> GenotypeData:
    """Auto-detect format and read a genotype fileset into memory."""
    ...

def read_plink(path: str) -> GenotypeData: ...
def read_vcf(path: str) -> GenotypeData: ...
def read_plink2(path: str) -> GenotypeData: ...
def read_bgen(path: str, sample_path: Optional[str] = None) -> GenotypeData: ...


# ── Streaming ─────────────────────────────────────────────

def stream(path: str, chunk_size: int = 5000) -> GenotypeStream:
    """Auto-detect format and stream genotype data in chunks."""
    ...

def stream_plink(path: str, chunk_size: int = 5000) -> GenotypeStream: ...
def stream_vcf(path: str, chunk_size: int = 5000) -> GenotypeStream: ...
def stream_plink2(path: str, chunk_size: int = 5000) -> GenotypeStream: ...
def stream_bgen(
    path: str, chunk_size: int = 5000, sample_path: Optional[str] = None
) -> GenotypeStream: ...


def supported_formats() -> dict[str, list[str]]: ...