from __future__ import annotations

from typing import Literal
from pathlib import Path
from anndata import AnnData

import snapatac2
import snapatac2._snapatac2 as internal
from snapatac2.genome import Genome

__all__ = ['make_fragment_file', 'import_fragments', 'import_contacts', 'import_values']

def make_fragment_file(
    bam_file: Path,
    output_file: Path,
    is_paired: bool = True,
    barcode_tag: str | None = None,
    barcode_regex: str | None = None,
    umi_tag: str | None = None,
    umi_regex: str | None = None,
    shift_left: int = 4,
    shift_right: int = -5,
    min_mapq: int | None = 30,
    chunk_size: int = 50000000,
    chrM: list[str] | None = ["chrM", "M"],
    source: Literal["10x"] | None = None,
    compression: Literal["gzip", "zstandard"] | None = None,
    compression_level: int | None = None,
    tempdir: Path | None = None,
) -> dict[str, float]:
    """
    Convert a BAM file to a fragment file.

    Convert a BAM file to a fragment file by performing the following steps:

    1. Filtering: remove reads that are unmapped, not primary alignment, mapq < 30,
       fails platform/vendor quality checks, or optical duplicate.
       For paired-end sequencing, it also removes reads that are not properly aligned.
    2. Deduplicate: Sort the reads by cell barcodes and remove duplicated reads
       for each unique cell barcode.
    3. Output: Convert BAM records to fragments (if paired-end) or single-end reads.

    The bam file needn't be sorted or filtered.

    Note
    ----
    - When using `barcode_regex` or `umi_regex`, the regex must contain exactly one capturing group
      (Parentheses group the regex between them) that matches the barcodes or UMIs.
      Writting the correct regex is tricky. You can test your regex online at https://regex101.com/.
      BAM files produced by the 10X Genomics Cell Ranger pipeline are not supported,
      as they contain invalid BAM headers. Specifically, Cell Ranger ATAC <= 2.0 produces BAM
      files with no @VN tag in the header, and Cell Ranger ATAC >= 2.1 produces BAM files
      with invalid @VN tag in the header.
      It is recommended to use the fragment files produced by Cell Ranger ATAC instead.
    - This function generates large temporary files in `tempdir` during sorting.
      For large files, it is recommended to set `tempdir` to a location with
      sufficient space in order to avoid running out of disk space.

    Parameters
    ----------
    bam_file
        File name of the BAM file.
    output_file
        File name of the output fragment file.
    is_paired
        Indicate whether the BAM file contain paired-end reads
    barcode_tag
        Extract barcodes from TAG fields of BAM records, e.g., `barcode_tag="CB"`.
    barcode_regex
        Extract barcodes from read names of BAM records using regular expressions.
        Reguler expressions should contain exactly one capturing group 
        (Parentheses group the regex between them) that matches
        the barcodes. For example, `barcode_regex="(..:..:..:..):\\\\w+$"`
        extracts `bd:69:Y6:10` from
        `A01535:24:HW2MMDSX2:2:1359:8513:3458:bd:69:Y6:10:TGATAGGTTG`.
        You can test your regex on this website: https://regex101.com/.
    umi_tag
        Extract UMI from TAG fields of BAM records.
    umi_regex
        Extract UMI from read names of BAM records using regular expressions.
        See `barcode_regex` for more details.
    shift_left
        Insertion site correction for the left end. Note this has no effect on single-end reads.
    shift_right
        Insertion site correction for the right end. Note this has no effect on single-end reads.
    min_mapq
        Filter the reads based on MAPQ.
    chunk_size
        The size of data retained in memory when performing sorting. Larger chunk sizes
        result in faster sorting and greater memory usage.
    source
        The source of the BAM file. This is used for vendor-specific processing.
        For example, the BAM files generated by 10X Genomics needs special processing
        to fix the malformed BAM headers.
        Currently the only supported source is "10x".
    chrM
        A list of mitochondrial chromosome names, used to calculate QC metrics.
    compression
        Compression type. If `None`, it is inferred from the suffix.
    compression_level
        Compression level. 1-9 for gzip, 1-22 for zstandard.
        If `None`, it is set to 6 for gzip and 3 for zstandard.
    tempdir
        Location to store temporary files. If `None`, system temporary directory
        will be used.

    Returns
    -------
    dict[str, float]
        A dictionary containing the following metrics:

        - "sequenced_reads": number of reads in the input BAM file.
        - "sequenced_read_pairs": number of read pairs in the input BAM file.
        - "frac_duplicates": Fraction of high-quality read pairs that are deemed
          to be PCR duplicates. This metric is a measure of sequencing
          saturation and is a function of library complexity and sequencing
          depth. More specifically, this is the fraction of high-quality
          fragments with a valid barcode that align to the same genomic
          position as another read pair in the library.
        - "frac_confidently_mapped": Fraction of sequenced reads or read pairs with mapping quality > 30.
        - "frac_unmapped": Fraction of sequenced reads or read pairs that have
          a valid barcode but could not be mapped to the genome.
        - "frac_valid_barcode": Fraction of reads or read pairs with barcodes that match the whitelist after error correction.
        - "frac_nonnuclear": Fraction of sequenced read pairs that have a valid
          barcode and map to non-nuclear genome contigs, including mitochondria,
          with mapping quality > 30.
        - "frac_fragment_in_nucleosome_free_region": Fraction of high-quality
          fragments smaller than 147 basepairs.
        - "frac_fragment_flanking_single_nucleosome": Fraction of high-quality
          fragments between 147 and 294 basepairs.

    See Also
    --------
    import_fragments
    import_contacts
    import_values
    """
    if barcode_tag is None and barcode_regex is None:
        raise ValueError("Either barcode_tag or barcode_regex must be set.")
    if barcode_tag is not None and barcode_regex is not None:
        raise ValueError("Only one of barcode_tag or barcode_regex can be set.")

    if compression is None:
        _, compression = snapatac2._utils.get_file_format(output_file)

    return internal.make_fragment_file(
        bam_file, output_file, is_paired, shift_left, shift_right, chunk_size,
        barcode_tag, barcode_regex, umi_tag, umi_regex, min_mapq, chrM, source,
        compression, compression_level, tempdir,
    )

def import_fragments(
    fragment_file: Path | list[Path],
    chrom_sizes: Genome | dict[str, int],
    *,
    is_paired: bool = True,
    file: Path | list[Path] | None = None,
    min_num_fragments: int = 200,
    sorted_by_barcode: bool = True,
    whitelist: Path | list[str] | None = None,
    chrM: list[str] = ["chrM", "M"],
    chunk_size: int = 2000,
    tempdir: Path | None = None,
    backend: Literal['hdf5'] = 'hdf5',
    n_jobs: int = 8,
) -> internal.AnnData:
    """Import data from fragment files and compute basic QC metrics.

    A fragment refers to the sequence data originating from a distinct location
    in the genome. In single-ended sequencing, one read equates to a fragment.
    However, in paired-ended sequencing, a fragment is defined by a pair of reads.
    This function is designed to handle, store, and process input files with
    fragment data, further yielding a range of basic Quality Control (QC) metrics.
    These metrics include the total number of unique fragments, duplication rates,
    and the percentage of mitochondrial DNA detected.

    How fragments are stored is dependent on the sequencing approach utilized.
    For single-ended sequencing, fragments are found in `.obsm['fragment_single']`.
    In contrast, for paired-ended sequencing, they are located in
    `.obsm['fragment_paired']`.

    Diving deeper technically, the fragments are internally structured within a
    Compressed Sparse Row (CSR) matrix. In this configuration, each row signifies
    a specific cell, while each column represents a unique genomic position.
    Fragment starting positions dictate the column indices. Matrix values
    capture the lengths of the fragments for paired-end reads or the lengths of
    the reads for single-ended scenarios. It's important to note that for
    single-ended reads, the values are signed, with the sign providing information
    on the fragment's strand orientation. Additionally, it is worth noting that
    cells may harbor duplicate fragments, leading to the presence of duplicate
    column indices within the matrix. As a result, the matrix deviates from
    the standard CSR format, and it is not advisable to use the matrix for linear
    algebra operations.
    
    .. image:: /_static/images/func+import_data.svg
        :align: center

    Note
    ----
    - This function accepts both single-end and paired-end reads. The argument `is_paired`
      specifies the type of reads in the fragment file.
    - When `file` is not `None`, this function uses constant memory regardless of
      the size of the input file.
    - When `sorted_by_barcode` is `False`, this function will sort the fragment file
      first, during which temporary files will be created in `tempdir`. The size of
      temporary files is proportional to the number of records in the fragment file.
      For large fragment files, it is recommended to set `tempdir` to a location with
      sufficient space in order to avoid running out of disk space.
    - The QC metrics are computed only for reads that are included by the `whitelist`
      or `chrom_sizes`.

    Warning
    -------
    When the input to the function is a list of files, it employs multiprocessing
    to process these files concurrently. In this case, however, it is crucial to
    safeguard the entry point of the program by encapsulating the function call
    within `if __name__ == '__main__':`. This condition ensures that the module
    is being run as the main program and not being loaded as a module from
    another script. Without this protection, each subprocess might attempt to
    spawn its own subprocesses, leading to a cascade of process spawns—a situation
    that can cause the program to hang or crash due to infinite recursion.
    You don't need to do this in Jupyter notebook as it automatically does that.

    Parameters
    ----------
    fragment_file
        File name of the fragment file, optionally compressed with gzip or zstd.
        This can be a single file or a list of files.
        If it is a list of files, a separate AnnData object will be created for each file.
        A fragment file must contain at least 5 columns:
        chromosome, start, end, barcode, count.
        Optionally it can contain one more column indicating the strand of the fragment.
    chrom_sizes
        A Genome object or a dictionary containing chromosome sizes, for example,
        `{"chr1": 2393, "chr2": 2344, ...}`.
    is_paired
        Indicate whether the fragment file contains paired-end reads.
    file
        File name of the output h5ad file used to store the result. If provided,
        result will be saved to a backed AnnData, otherwise an in-memory AnnData
        is used.
        If `fragment_file` is a list of files, `file` must also be a list of files if provided.
    min_num_fragments
        Number of unique fragments threshold used to filter cells
    sorted_by_barcode
        Whether the fragment file has been sorted by cell barcodes.
        This function will be faster if `sorted_by_barcode==True`.
        Note the :func:`~snapatac2.pp.make_fragment_file` will always sort the
        fragment file by barcode.
    whitelist
        File name or a list of barcodes. If it is a file name, each line
        must contain a valid barcode. When provided, only barcodes in the whitelist
        will be retained.
    chrM
        A list of chromosome names that are considered mitochondrial DNA. This is
        used to compute the fraction of mitochondrial DNA.
    chunk_size
        Increasing the chunk_size may speed up I/O but will use more memory.
        The speed gain is usually not significant.
    tempdir
        Location to store temporary files. If `None`, system temporary directory
        will be used.
    backend
        The backend.
    n_jobs
        Number of jobs to run in parallel when `fragment_file` is a list.
        If `n_jobs=-1`, all CPUs will be used.

    Returns
    -------
    AnnData | ad.AnnData
        An annotated data matrix of shape `n_obs` x `n_vars`. Rows correspond to
        cells and columns to regions. If `file=None`, an in-memory AnnData will be
        returned, otherwise a backed AnnData is returned.

    See Also
    --------
    make_fragment_file
    import_contacts
    import_values

    Examples
    --------
    >>> import snapatac2 as snap
    >>> data = snap.pp.import_fragments(snap.datasets.pbmc500(downsample=True), chrom_sizes=snap.genome.hg38, sorted_by_barcode=False)
    >>> print(data)
    AnnData object with n_obs × n_vars = 585 × 0
        obs: 'n_fragment', 'frac_dup', 'frac_mito'
        uns: 'reference_sequences'
        obsm: 'fragment_paired'
    """
    chrom_sizes = chrom_sizes.chrom_sizes if isinstance(chrom_sizes, Genome) else chrom_sizes
    if len(chrom_sizes) == 0:
        raise ValueError("chrom_size cannot be empty")

    if whitelist is not None:
        if isinstance(whitelist, str) or isinstance(whitelist, Path):
            with open(whitelist, "r") as fl:
                whitelist = set([line.strip() for line in fl])
        else:
            whitelist = set(whitelist)

    if isinstance(fragment_file, list):
        n = len(fragment_file)
        if file is None:
            adatas = [AnnData() for _ in range(n)]
        else:
            if len(file) != n:
                raise ValueError("The length of 'file' must be the same as the length of 'fragment_file'")
            adatas = [internal.AnnData(filename=f, backend=backend) for f in file]

        snapatac2._utils.anndata_ipar(
            list(enumerate(adatas)),
            lambda x: internal.import_fragments(
                x[1], fragment_file[x[0]], is_paired, chrom_sizes, chrM, min_num_fragments,
                sorted_by_barcode, chunk_size, whitelist, tempdir,
            ),
            n_jobs=n_jobs,
        )
        return adatas
    else:
        adata = AnnData() if file is None else internal.AnnData(filename=file, backend=backend)
        internal.import_fragments(
            adata, fragment_file, is_paired, chrom_sizes, chrM, min_num_fragments,
            sorted_by_barcode, chunk_size, whitelist, tempdir,
        )
        return adata

def import_contacts(
    contact_file: Path,
    chrom_sizes: Genome | dict[str, int],
    *,
    file: Path | None = None,
    sorted_by_barcode: bool = True,
    bin_size: int = 500000,
    chunk_size: int = 200,
    tempdir: Path | None = None,
    backend: Literal['hdf5'] = 'hdf5',
) -> internal.AnnData:
    """Import chromatin contacts.

    Parameters
    ----------
    contact_file
        File name of the fragment file.
    file
        File name of the output h5ad file used to store the result. If provided,
        result will be saved to a backed AnnData, otherwise an in-memory AnnData
        is used.
    chrom_sizes
        A Genome object or a dictionary containing chromosome sizes, for example,
        `{"chr1": 2393, "chr2": 2344, ...}`.
    sorted_by_barcode
        Whether the fragment file has been sorted by cell barcodes.
        If `sorted_by_barcode == True`, this function makes use of small fixed amout of 
        memory. If `sorted_by_barcode == False` and `low_memory == False`,
        all data will be kept in memory. See `low_memory` for more details.
    bin_size
        The size of consecutive genomic regions used to record the counts.
    chunk_size
        Increasing the chunk_size speeds up I/O but uses more memory.
    tempdir
        Location to store temporary files. If `None`, system temporary directory
        will be used.
    backend
        The backend.

    Returns
    -------
    AnnData | ad.AnnData
        An annotated data matrix of shape `n_obs` x `n_vars`. Rows correspond to
        cells and columns to regions. If `file=None`, an in-memory AnnData will be
        returned, otherwise a backed AnnData is returned.
    """
    chrom_sizes = chrom_sizes.chrom_sizes if isinstance(chrom_sizes, Genome) else chrom_sizes
    if len(chrom_sizes) == 0:
        raise ValueError("chrom_size cannot be empty")

    adata = AnnData() if file is None else internal.AnnData(filename=file, backend=backend)
    internal.import_contacts(
        adata, contact_file, chrom_sizes, sorted_by_barcode, bin_size, chunk_size, tempdir
    )
    return adata

def import_values(
    input_dir: Path,
    chrom_sizes: Genome | dict[str, int],
    *,
    file: Path | None = None,
    whitelist: Path | list[str] | None = None,
    chunk_size: int = 200,
    backend: Literal['hdf5'] = 'hdf5',
) -> internal.AnnData:
    """Import values associated with base pairs, typically from experiments like
    whole-genome bisulfite sequencing (WGBS).

    Parameters
    ----------
    input_dir
        Directory containing the input files. Each file corresponds to a single cell.
    chrom_sizes
        A Genome object or a dictionary containing chromosome sizes, for example,
        `{"chr1": 2393, "chr2": 2344, ...}`.
    file
        File name of the output h5ad file used to store the result. If provided,
        result will be saved to a backed AnnData, otherwise an in-memory AnnData
        is used.
    whitelist
        File name or a list of barcodes. If it is a file name, each line
        must contain a valid barcode. When provided, only barcodes in the whitelist
        will be retained.
    chunk_size
        Increasing the chunk_size speeds up I/O but uses more memory.
    backend
        The backend.

    Returns
    -------
    AnnData | ad.AnnData
        An annotated data matrix of shape `n_obs` x `n_vars`. Rows correspond to
        cells and columns to regions. If `file=None`, an in-memory AnnData will be
        returned, otherwise a backed AnnData is returned.
    """
    chrom_sizes = chrom_sizes.chrom_sizes if isinstance(chrom_sizes, Genome) else chrom_sizes
    if len(chrom_sizes) == 0:
        raise ValueError("chrom_size cannot be empty")

    if whitelist is not None:
        if isinstance(whitelist, str) or isinstance(whitelist, Path):
            with open(whitelist, "r") as fl:
                whitelist = set([line.strip() for line in fl])
        else:
            whitelist = set(whitelist)

    adata = AnnData() if file is None else internal.AnnData(filename=file, backend=backend)
    internal.import_values(
        adata, input_dir, chrom_sizes, chunk_size, whitelist
    )
    return adata