from functools import lru_cache
from typing import Any, Dict, List, Optional, Tuple, Union

import allel  # type: ignore
import bokeh
import dask.array as da
import numpy as np
import pandas as pd
import xarray as xr
import zarr  # type: ignore
from numpydoc_decorator import doc  # type: ignore

from ..util import (
    DIM_ALLELE,
    DIM_PLOIDY,
    DIM_SAMPLE,
    DIM_VARIANT,
    CacheMiss,
    Region,
    _apply_allele_mapping,
    _check_types,
    _da_compress,
    _da_concat,
    _da_from_zarr,
    _dask_apply_allele_mapping,
    _dask_compress_dataset,
    _dask_genotype_array_map_alleles,
    _init_zarr_store,
    _locate_region,
    _parse_multi_region,
    _parse_single_region,
    _simple_xarray_concat,
    _trim_alleles,
    _true_runs,
)
from . import base_params
from .genome_features import AnophelesGenomeFeaturesData, gplt_params
from .genome_sequence import AnophelesGenomeSequenceData
from .sample_metadata import AnophelesSampleMetadata


class AnophelesSnpData(
    AnophelesSampleMetadata, AnophelesGenomeFeaturesData, AnophelesGenomeSequenceData
):
    def __init__(
        self,
        site_filters_analysis: Optional[str] = None,
        default_site_mask: Optional[str] = None,
        **kwargs,
    ):
        # N.B., this class is designed to work cooperatively, and
        # so it's important that any remaining parameters are passed
        # to the superclass constructor.
        super().__init__(**kwargs)

        # If provided, this analysis version will override the
        # default value provided in the release configuration.
        self._site_filters_analysis_override = site_filters_analysis

        # These will vary between data resources.
        self._default_site_mask = default_site_mask

        # Set up caches.
        # TODO review type annotations here, maybe can tighten
        self._cache_snp_sites = None
        self._cache_snp_genotypes: Dict = dict()
        self._cache_site_filters: Dict = dict()
        self._cache_site_annotations = None
        self._cache_locate_site_class: Dict = dict()

    @property
    def _site_filters_analysis(self) -> Optional[str]:
        if self._site_filters_analysis_override:
            return self._site_filters_analysis_override
        else:
            # N.B., this will return None if the key is not present in the
            # config.
            return self.config.get("DEFAULT_SITE_FILTERS_ANALYSIS")

    @property
    def site_mask_ids(self) -> Tuple[str, ...]:
        """Identifiers for the different site masks that are available.
        These are values than can be used for the `site_mask` parameter in any
        method making using of SNP data.

        """
        return tuple(self.config.get("SITE_MASK_IDS", ()))  # ensure tuple

    @property
    def _site_annotations_zarr_path(self) -> str:
        return self.config["SITE_ANNOTATIONS_ZARR_PATH"]

    def _prep_site_mask_param(
        self,
        *,
        site_mask: base_params.site_mask,
    ) -> base_params.site_mask:
        if site_mask == base_params.DEFAULT:
            # Use whatever is the default site mask for this data resource.
            assert self._default_site_mask is not None
            return self._default_site_mask
        elif site_mask in self.site_mask_ids:
            return site_mask
        else:
            raise ValueError(
                f"Invalid site mask, must be one of f{self.site_mask_ids}."
            )

    def _prep_optional_site_mask_param(
        self,
        *,
        site_mask: Optional[base_params.site_mask],
    ) -> Optional[base_params.site_mask]:
        if site_mask is None:
            # This is allowed, it means don't apply any site mask to the data.
            return None
        else:
            return self._prep_site_mask_param(site_mask=site_mask)

    @doc(
        summary="Open SNP sites zarr",
        returns="Zarr hierarchy.",
    )
    def open_snp_sites(self) -> zarr.hierarchy.Group:
        # Here we cache the opened zarr hierarchy, to avoid small delays
        # reading zarr metadata.
        if self._cache_snp_sites is None:
            path = (
                f"{self._base_path}/{self._major_version_path}/snp_genotypes/all/sites/"
            )
            store = _init_zarr_store(fs=self._fs, path=path)
            root = zarr.open_consolidated(store=store)
            self._cache_snp_sites = root
        return self._cache_snp_sites

    @_check_types
    @doc(
        summary="Open SNP genotypes zarr for a given sample set.",
        returns="Zarr hierarchy.",
    )
    def open_snp_genotypes(
        self, sample_set: base_params.sample_set
    ) -> zarr.hierarchy.Group:
        # Here we cache the opened zarr hierarchy, to avoid small delays
        # reading zarr metadata.
        try:
            return self._cache_snp_genotypes[sample_set]
        except KeyError:
            release = self.lookup_release(sample_set=sample_set)
            release_path = self._release_to_path(release)
            path = f"{self._base_path}/{release_path}/snp_genotypes/all/{sample_set}/"
            store = _init_zarr_store(fs=self._fs, path=path)
            root = zarr.open_consolidated(store=store)
            self._cache_snp_genotypes[sample_set] = root
            return root

    def _require_site_filters_analysis(self):
        if not self._site_filters_analysis:
            raise NotImplementedError(
                "Site filters not available for this data resource."
            )

    @_check_types
    @doc(
        summary="Open site filters zarr.",
        returns="Zarr hierarchy.",
    )
    def open_site_filters(
        self,
        mask: base_params.site_mask = base_params.DEFAULT,
    ) -> zarr.hierarchy.Group:
        self._require_site_filters_analysis()
        mask_prepped = self._prep_site_mask_param(site_mask=mask)

        # Here we cache the opened zarr hierarchy, to avoid small delays
        # reading zarr metadata.
        try:
            return self._cache_site_filters[mask_prepped]
        except KeyError:
            path = f"{self._base_path}/{self._major_version_path}/site_filters/{self._site_filters_analysis}/{mask_prepped}/"
            store = _init_zarr_store(fs=self._fs, path=path)
            root = zarr.open_consolidated(store=store)
            self._cache_site_filters[mask_prepped] = root
            return root

    @doc(
        summary="Open site annotations zarr.",
        returns="Zarr hierarchy.",
    )
    def open_site_annotations(self) -> zarr.hierarchy.Group:
        if self._cache_site_annotations is None:
            path = f"{self._base_path}/{self._site_annotations_zarr_path}"
            store = _init_zarr_store(fs=self._fs, path=path)
            self._cache_site_annotations = zarr.open_consolidated(store=store)
        return self._cache_site_annotations

    def _site_filters_for_contig(
        self,
        *,
        contig: str,
        mask: base_params.site_mask,
        field: base_params.field,
        inline_array: base_params.inline_array,
        chunks: base_params.chunks,
    ):
        if contig in self.virtual_contigs:
            contigs = self.virtual_contigs[contig]
            arrs = [
                self._site_filters_for_contig(
                    contig=c,
                    mask=mask,
                    field=field,
                    inline_array=inline_array,
                    chunks=chunks,
                )
                for c in contigs
            ]
            d = da.concatenate(arrs)
            return d

        else:
            assert contig in self.contigs
            root = self.open_site_filters(mask=mask)
            z = root[f"{contig}/variants/{field}"]
            d = _da_from_zarr(z, inline_array=inline_array, chunks=chunks)
            return d

    def _site_filters_for_region(
        self,
        *,
        region: Region,
        mask: base_params.site_mask,
        field: base_params.field,
        inline_array: base_params.inline_array,
        chunks: base_params.chunks,
    ):
        d = self._site_filters_for_contig(
            contig=region.contig,
            mask=mask,
            field=field,
            inline_array=inline_array,
            chunks=chunks,
        )
        if region.start or region.end:
            pos = self._snp_sites_for_contig(
                contig=region.contig,
                field="POS",
                inline_array=inline_array,
                chunks=chunks,
            )
            loc_region = _locate_region(region, np.asarray(pos))
            d = d[loc_region]
        return d

    @_check_types
    @doc(
        summary="Access SNP site filters.",
        returns="""
            An array of boolean values identifying sites that pass the filters.
        """,
    )
    def site_filters(
        self,
        region: base_params.regions,
        mask: base_params.site_mask = base_params.DEFAULT,
        field: base_params.field = "filter_pass",
        inline_array: base_params.inline_array = base_params.inline_array_default,
        chunks: base_params.chunks = base_params.native_chunks,
    ) -> da.Array:
        mask_prepped = self._prep_site_mask_param(site_mask=mask)
        del mask

        # Resolve the region parameter to a standard type.
        regions: List[Region] = _parse_multi_region(self, region)
        del region

        # Load arrays and concatenate if needed.
        d = _da_concat(
            [
                self._site_filters_for_region(
                    region=r,
                    mask=mask_prepped,
                    field=field,
                    inline_array=inline_array,
                    chunks=chunks,
                )
                for r in regions
            ]
        )

        return d

    def _snp_sites_for_contig(
        self,
        *,
        contig: base_params.contig,
        field: base_params.field,
        inline_array: base_params.inline_array,
        chunks: base_params.chunks,
    ) -> da.Array:
        """Access SNP sites data for a single contig."""

        # Handle virtual contig.
        if contig in self.virtual_contigs:
            contigs = self.virtual_contigs[contig]
            arrs = []
            offset = 0
            for c in contigs:
                arr = self._snp_sites_for_contig(
                    contig=c,
                    field=field,
                    inline_array=inline_array,
                    chunks=chunks,
                )
                if field == "POS":
                    if offset > 0:
                        arr = arr + offset
                    offset += self.genome_sequence(region=c).shape[0]
                arrs.append(arr)
            return da.concatenate(arrs)

        # Handle contig in the reference genome.
        else:
            assert contig in self.contigs
            root = self.open_snp_sites()
            z = root[f"{contig}/variants/{field}"]
            ret = _da_from_zarr(z, inline_array=inline_array, chunks=chunks)
            return ret

    def _snp_sites_for_region(
        self,
        *,
        region: Region,
        field: base_params.field,
        inline_array: base_params.inline_array,
        chunks: base_params.chunks,
    ) -> da.Array:
        # Access data for the requested contig.
        ret = self._snp_sites_for_contig(
            contig=region.contig, field=field, inline_array=inline_array, chunks=chunks
        )

        # Deal with a region.
        if region.start or region.end:
            if field == "POS":
                pos = ret
            else:
                pos = self._snp_sites_for_contig(
                    contig=region.contig,
                    field="POS",
                    inline_array=inline_array,
                    chunks=chunks,
                )
            loc_region = _locate_region(region, np.asarray(pos))
            ret = ret[loc_region]

        return ret

    @_check_types
    @doc(
        summary="Access SNP site data (positions or alleles).",
        returns="""
            An array of either SNP positions ("POS"), reference alleles ("REF") or
            alternate alleles ("ALT").
        """,
    )
    def snp_sites(
        self,
        region: base_params.regions,
        field: base_params.field,
        site_mask: Optional[base_params.site_mask] = None,
        inline_array: base_params.inline_array = base_params.inline_array_default,
        chunks: base_params.chunks = base_params.native_chunks,
    ) -> da.Array:
        # Resolve the region parameter to a standard type.
        regions: List[Region] = _parse_multi_region(self, region)
        del region
        site_mask_prepped = self._prep_optional_site_mask_param(site_mask=site_mask)
        del site_mask

        # Access SNP sites and concatenate over regions.
        ret = _da_concat(
            [
                self._snp_sites_for_region(
                    region=r,
                    field=field,
                    chunks=chunks,
                    inline_array=inline_array,
                )
                for r in regions
            ],
            axis=0,
        )

        # Apply site mask if requested.
        if site_mask_prepped is not None:
            loc_sites = self.site_filters(
                region=regions,
                mask=site_mask_prepped,
                chunks=chunks,
                inline_array=inline_array,
            )
            ret = _da_compress(loc_sites, ret, axis=0)

        return ret

    def _snp_genotypes_for_contig(
        self,
        *,
        contig: base_params.contig,
        sample_set: base_params.sample_set,
        field: base_params.field,
        inline_array: base_params.inline_array,
        chunks: base_params.chunks,
    ) -> da.Array:
        """Access SNP genotypes for a single contig and a single sample set."""

        if contig in self.virtual_contigs:
            contigs = self.virtual_contigs[contig]
            arrs = [
                self._snp_genotypes_for_contig(
                    contig=c,
                    sample_set=sample_set,
                    field=field,
                    inline_array=inline_array,
                    chunks=chunks,
                )
                for c in contigs
            ]
            return da.concatenate(arrs)

        else:
            assert contig in self.contigs
            root = self.open_snp_genotypes(sample_set=sample_set)
            z = root[f"{contig}/calldata/{field}"]
            d = _da_from_zarr(z, inline_array=inline_array, chunks=chunks)
            return d

    @_check_types
    @doc(
        summary="Access SNP genotypes and associated data.",
        returns="""
            An array of either genotypes (GT), genotype quality (GQ), allele
            depths (AD) or mapping quality (MQ) values.
        """,
    )
    def snp_genotypes(
        self,
        region: base_params.regions,
        sample_sets: Optional[base_params.sample_sets] = None,
        sample_query: Optional[base_params.sample_query] = None,
        sample_query_options: Optional[base_params.sample_query_options] = None,
        sample_indices: Optional[base_params.sample_indices] = None,
        field: base_params.field = "GT",
        site_mask: Optional[base_params.site_mask] = None,
        inline_array: base_params.inline_array = base_params.inline_array_default,
        chunks: base_params.chunks = base_params.native_chunks,
    ) -> da.Array:
        # Check that either sample_query xor sample_indices are provided.
        base_params._validate_sample_selection_params(
            sample_query=sample_query, sample_indices=sample_indices
        )

        # Prepare parameters.
        prepared_sample_sets = self._prep_sample_sets_param(sample_sets=sample_sets)
        prepared_sample_query = self._prep_sample_query_param(sample_query=sample_query)
        prepared_regions: List[Region] = _parse_multi_region(self, region)
        prepared_site_mask = self._prep_optional_site_mask_param(site_mask=site_mask)

        # Delete original parameters to prevent accidental use.
        del sample_sets
        del sample_query
        del region
        del site_mask

        with self._spinner("Access SNP genotypes"):
            # Concatenate multiple sample sets and/or contigs.
            lx = []
            for r in prepared_regions:
                contig = r.contig
                ly = []

                for s in prepared_sample_sets:
                    y = self._snp_genotypes_for_contig(
                        contig=contig,
                        sample_set=s,
                        field=field,
                        inline_array=inline_array,
                        chunks=chunks,
                    )
                    ly.append(y)

                # Concatenate data from multiple sample sets.
                x = _da_concat(ly, axis=1)

                # Locate region - do this only once, optimisation.
                if r.start or r.end:
                    pos = self._snp_sites_for_contig(
                        contig=contig,
                        field="POS",
                        inline_array=inline_array,
                        chunks=chunks,
                    )
                    loc_region = _locate_region(r, np.asarray(pos))
                    x = x[loc_region]

                lx.append(x)

            # Concatenate data from multiple regions.
            d = _da_concat(lx, axis=0)

        # Apply site filters if requested.
        if prepared_site_mask is not None:
            loc_sites = self.site_filters(
                region=prepared_regions,
                mask=prepared_site_mask,
            )
            d = _da_compress(loc_sites, d, axis=0)

        # Apply the sample_query, if there is one.
        # Note: this might have been internally modified, e.g. `is_surveillance == True`.
        if prepared_sample_query is not None:
            # Note: the unfiltered Dask array `d` is not aligned with the filtered `sample_metadata`,
            #       so we cannot use filtered `sample_metadata` to get the relevant boolean filter.

            # Note: the unfiltered Dask array `d` does not contain sample identifiers,
            #       so we cannot use a list of relevant sample ids to produce the boolean filter directly.

            # Note: we can first determine the list of relevant sample ids using filtered `sample_metadata`,
            #       then use the unfiltered `general_metadata` to determine the appropriate boolean filter.

            df_filtered_samples = self.sample_metadata(
                sample_sets=prepared_sample_sets,
                sample_query=prepared_sample_query,
                sample_query_options=sample_query_options,
            )

            # Raise an error if no samples match the sample query.
            if len(df_filtered_samples) == 0:
                raise ValueError(
                    f"No samples found for query {prepared_sample_query!r}"
                )

            # Get the list of unfiltered samples, in order to produce an aligned boolean filter.
            df_unfiltered_samples = self.general_metadata(
                sample_sets=prepared_sample_sets
            )

            # Get a boolean array for unfiltered data, indicating which samples match the query.
            loc_samples = df_unfiltered_samples["sample_id"].isin(
                df_filtered_samples["sample_id"]
            )

            # Filter the Dask array using the boolean array.
            d = da.compress(loc_samples, d, axis=1)

        # Apply the sample_indices, if there are any.
        # Note: this might need to apply to the result of an internal sample_query, e.g. `is_surveillance == True`.
        if sample_indices is not None:
            d = da.take(d, sample_indices, axis=1)

        return d

    def _snp_variants_for_contig(
        self,
        *,
        contig: base_params.contig,
        inline_array: base_params.inline_array,
        chunks: base_params.chunks,
    ):
        if contig in self.virtual_contigs:
            contigs = self.virtual_contigs[contig]
            datasets = []
            offset = 0
            for c in contigs:
                dsc = self._snp_variants_for_contig(
                    contig=c,
                    inline_array=inline_array,
                    chunks=chunks,
                )
                if offset > 0:
                    dsc["variant_position"] = dsc["variant_position"] + offset
                offset += self.genome_sequence(region=c).shape[0]
                datasets.append(dsc)
            ret = _simple_xarray_concat(datasets, dim=DIM_VARIANT)
            return ret

        else:
            assert contig in self.contigs
            coords = dict()
            data_vars = dict()
            sites_root = self.open_snp_sites()

            # Set up variant_position.
            pos_z = sites_root[f"{contig}/variants/POS"]
            variant_position = _da_from_zarr(
                pos_z, inline_array=inline_array, chunks=chunks
            )
            coords["variant_position"] = [DIM_VARIANT], variant_position

            # Set up variant_allele.
            ref_z = sites_root[f"{contig}/variants/REF"]
            alt_z = sites_root[f"{contig}/variants/ALT"]
            ref = _da_from_zarr(ref_z, inline_array=inline_array, chunks=chunks)
            alt = _da_from_zarr(alt_z, inline_array=inline_array, chunks=chunks)
            variant_allele = da.concatenate([ref[:, None], alt], axis=1)
            variant_allele = variant_allele.rechunk((variant_allele.chunks[0], -1))
            data_vars["variant_allele"] = [DIM_VARIANT, DIM_ALLELE], variant_allele

            # Set up variant_contig.
            contig_index = self.contigs.index(contig)
            variant_contig = da.full_like(
                variant_position, fill_value=contig_index, dtype="u1"
            )
            coords["variant_contig"] = [DIM_VARIANT], variant_contig

            # Set up site filters arrays.
            for mask in self.site_mask_ids:
                filters_root = self.open_site_filters(mask=mask)
                z = filters_root[f"{contig}/variants/filter_pass"]
                d = _da_from_zarr(z, inline_array=inline_array, chunks=chunks)
                data_vars[f"variant_filter_pass_{mask}"] = [DIM_VARIANT], d

            # Set up attributes.
            attrs = {"contigs": self.contigs}

            # Create a dataset.
            dsc = xr.Dataset(data_vars=data_vars, coords=coords, attrs=attrs)

            return dsc

    @_check_types
    @doc(
        summary="Access SNP sites and site filters.",
        returns="A dataset containing SNP sites and site filters.",
    )
    def snp_variants(
        self,
        region: base_params.regions,
        site_mask: Optional[base_params.site_mask] = None,
        inline_array: base_params.inline_array = base_params.inline_array_default,
        chunks: base_params.chunks = base_params.native_chunks,
    ):
        # Normalise parameters.
        regions: List[Region] = _parse_multi_region(self, region)
        del region
        site_mask_prepped = self._prep_optional_site_mask_param(site_mask=site_mask)
        del site_mask

        # Access SNP data and concatenate multiple regions.
        lx = []
        for r in regions:
            # Access variants.
            x = self._snp_variants_for_contig(
                contig=r.contig,
                inline_array=inline_array,
                chunks=chunks,
            )

            # Handle region.
            if r.start or r.end:
                pos = x["variant_position"].values
                loc_region = _locate_region(r, pos)
                x = x.isel(variants=loc_region)

            lx.append(x)

        # Concatenate data from multiple regions.
        ds = _simple_xarray_concat(lx, dim=DIM_VARIANT)

        # Apply site filters.
        if site_mask_prepped is not None:
            ds = _dask_compress_dataset(
                ds, indexer=f"variant_filter_pass_{site_mask_prepped}", dim=DIM_VARIANT
            )

        return ds

    def _site_annotations_raw(
        self,
        *,
        contig,
        inline_array: base_params.inline_array = base_params.inline_array_default,
        chunks: base_params.chunks = base_params.native_chunks,
    ) -> xr.Dataset:
        # Open site annotations zarr.
        root = self.open_site_annotations()

        # Build a dataset.
        ds = xr.Dataset()
        for field in (
            "codon_degeneracy",
            "codon_nonsyn",
            "codon_position",
            "seq_cls",
            "seq_flen",
            "seq_relpos_start",
            "seq_relpos_stop",
        ):
            data = _da_from_zarr(
                root[field][contig],
                inline_array=inline_array,
                chunks=chunks,
            )
            ds[field] = "variants", data

        return ds

    @_check_types
    @doc(
        summary="Load site annotations.",
        returns="""
        A dataset of site annotations with 1 dimension: `variants` the number of sites in the selected region. It has 7 variables:
        `seq_cls`: The feature class. There are 11 possible values:
                *1*: Upstream,
                *2*: Downstream,
                *3*: 5' UTR,
                *4*: 3' UTR,
                *5*: CDS (first),
                *6*: CDS (mid),
                *7*: CDS (last),
                *8*: Intron (first),
                *9*: Intron (mid),
                *10*: Intron (last),
                *0*: Unknown,
        `seq_flen`: The length of the feature,
        `seq_relpos_start`: Relative position to the start of the feature. 0 if not in a feature,
        `seq_relpos_stop`: Relative position to the end of the feature. 0 if not in a feature,
        `codon_position`: Position within a triplet codon. -1 if not in a CDS,
        `codon_nonsyn`: Number of possible nucleotide changes at this position that would result in an amino acid change from the reference, 0 if not in a CDS,
        `codon_degeneracy`: The redundancy of the codon. Can take 5 different values:
                *1*: 0-fold degenerate, i.e., all nucleotides encode different amino acids,
                *2*: simple 2-fold degenerate, i.e., 2 different amino acids can be encoded depending on which nucleotide is present; each amino acid is encoded by two different nucleotides,
                *3*: complex 2-fold degenerate, i.e., 2 different amino acids can be encoded depending on which nucleotide is present but they are not paired,
                *4*: 4-fold degenerate, i.e., all nucleotides encode the same amino acid,
                *-1*: not in a CDS
        """,
    )
    def site_annotations(
        self,
        region: base_params.region,
        site_mask: Optional[base_params.site_mask] = None,
        inline_array: base_params.inline_array = base_params.inline_array_default,
        chunks: base_params.chunks = base_params.native_chunks,
    ) -> xr.Dataset:
        # Resolve region.
        resolved_region: Region = _parse_single_region(self, region)
        del region
        contig = resolved_region.contig
        site_mask_prepped = self._prep_optional_site_mask_param(site_mask=site_mask)
        del site_mask

        # Access site annotations.
        ds = self._site_annotations_raw(
            contig=contig, inline_array=inline_array, chunks=chunks
        )

        # N.B., site annotations data are provided for every position in the genome. We need to
        # therefore subset to SNP positions.
        pos = self.snp_sites(
            region=resolved_region,
            field="POS",
            site_mask=site_mask_prepped,
            inline_array=inline_array,
            chunks=chunks,
        )
        idx = (pos - 1).compute()
        ds = ds.isel(variants=idx)

        return ds

    def _locate_site_class(
        self,
        *,
        region: Region,
        site_mask: Optional[base_params.site_mask],
        site_class: base_params.site_class,
        inline_array: base_params.inline_array = base_params.inline_array_default,
        chunks: base_params.chunks = base_params.native_chunks,
    ):
        # Cache these data in memory to avoid repeated computation.
        cache_key = (region, site_mask, site_class)

        try:
            loc_ann = self._cache_locate_site_class[cache_key]

        except KeyError as exc:
            # Access site annotations data.
            ds_ann = self._site_annotations_raw(
                contig=region.contig,
                inline_array=inline_array,
                chunks=chunks,
            )
            codon_pos = ds_ann["codon_position"].data
            codon_deg = ds_ann["codon_degeneracy"].data
            seq_cls = ds_ann["seq_cls"].data
            seq_flen = ds_ann["seq_flen"].data
            seq_relpos_start = ds_ann["seq_relpos_start"].data
            seq_relpos_stop = ds_ann["seq_relpos_stop"].data
            site_class = site_class.upper()

            # Define constants used in site annotations data.
            SEQ_CLS_UNKNOWN = 0  # noqa
            SEQ_CLS_UPSTREAM = 1
            SEQ_CLS_DOWNSTREAM = 2
            SEQ_CLS_5UTR = 3
            SEQ_CLS_3UTR = 4
            SEQ_CLS_CDS_FIRST = 5
            SEQ_CLS_CDS_MID = 6
            SEQ_CLS_CDS_LAST = 7
            SEQ_CLS_INTRON_FIRST = 8
            SEQ_CLS_INTRON_MID = 9
            SEQ_CLS_INTRON_LAST = 10
            CODON_DEG_UNKNOWN = 0  # noqa
            CODON_DEG_0 = 1
            CODON_DEG_2_SIMPLE = 2
            CODON_DEG_2_COMPLEX = 3  # noqa
            CODON_DEG_4 = 4

            # Set up site selection.

            if site_class == "CDS_DEG_4":
                # 4-fold degenerate coding sites
                loc_ann = (
                    (
                        (seq_cls == SEQ_CLS_CDS_FIRST)
                        | (seq_cls == SEQ_CLS_CDS_MID)
                        | (seq_cls == SEQ_CLS_CDS_LAST)
                    )
                    & (codon_pos == 2)
                    & (codon_deg == CODON_DEG_4)
                )

            elif site_class == "CDS_DEG_2_SIMPLE":
                # 2-fold degenerate coding sites
                loc_ann = (
                    (
                        (seq_cls == SEQ_CLS_CDS_FIRST)
                        | (seq_cls == SEQ_CLS_CDS_MID)
                        | (seq_cls == SEQ_CLS_CDS_LAST)
                    )
                    & (codon_pos == 2)
                    & (codon_deg == CODON_DEG_2_SIMPLE)
                )

            elif site_class == "CDS_DEG_0":
                # non-degenerate coding sites
                loc_ann = (
                    (seq_cls == SEQ_CLS_CDS_FIRST)
                    | (seq_cls == SEQ_CLS_CDS_MID)
                    | (seq_cls == SEQ_CLS_CDS_LAST)
                ) & (codon_deg == CODON_DEG_0)

            elif site_class == "INTRON_SHORT":
                # short introns, excluding splice regions
                loc_ann = (
                    (
                        (seq_cls == SEQ_CLS_INTRON_FIRST)
                        | (seq_cls == SEQ_CLS_INTRON_MID)
                        | (seq_cls == SEQ_CLS_INTRON_LAST)
                    )
                    & (seq_flen < 100)
                    & (seq_relpos_start > 10)
                    & (seq_relpos_stop > 10)
                )

            elif site_class == "INTRON_LONG":
                # long introns, excluding splice regions
                loc_ann = (
                    (
                        (seq_cls == SEQ_CLS_INTRON_FIRST)
                        | (seq_cls == SEQ_CLS_INTRON_MID)
                        | (seq_cls == SEQ_CLS_INTRON_LAST)
                    )
                    & (seq_flen > 200)
                    & (seq_relpos_start > 10)
                    & (seq_relpos_stop > 10)
                )

            elif site_class == "INTRON_SPLICE_5PRIME":
                # 5' intron splice regions
                loc_ann = (
                    (seq_cls == SEQ_CLS_INTRON_FIRST)
                    | (seq_cls == SEQ_CLS_INTRON_MID)
                    | (seq_cls == SEQ_CLS_INTRON_LAST)
                ) & (seq_relpos_start < 2)

            elif site_class == "INTRON_SPLICE_3PRIME":
                # 3' intron splice regions
                loc_ann = (
                    (seq_cls == SEQ_CLS_INTRON_FIRST)
                    | (seq_cls == SEQ_CLS_INTRON_MID)
                    | (seq_cls == SEQ_CLS_INTRON_LAST)
                ) & (seq_relpos_stop < 2)

            elif site_class == "UTR_5PRIME":
                # 5' UTR
                loc_ann = seq_cls == SEQ_CLS_5UTR

            elif site_class == "UTR_3PRIME":
                # 3' UTR
                loc_ann = seq_cls == SEQ_CLS_3UTR

            elif site_class == "INTERGENIC":
                # intergenic regions, distant from a gene
                loc_ann = (
                    (seq_cls == SEQ_CLS_UPSTREAM) & (seq_relpos_stop > 10_000)
                ) | ((seq_cls == SEQ_CLS_DOWNSTREAM) & (seq_relpos_start > 10_000))

            else:
                raise NotImplementedError(site_class) from exc

            # N.B., site annotations data are provided for every position in the genome. We need to
            # therefore subset to SNP positions.
            pos = self.snp_sites(
                region=region,
                field="POS",
                site_mask=site_mask,
                inline_array=inline_array,
                chunks=chunks,
            )
            idx = (pos - 1).compute()
            loc_ann = da.take(loc_ann, idx, axis=0)

            # Compute site selection.
            with self._dask_progress(desc=f"Locate {site_class} sites"):
                loc_ann = loc_ann.compute()

            self._cache_locate_site_class[cache_key] = loc_ann

        return loc_ann

    def _snp_calls_for_contig(
        self,
        *,
        contig: base_params.contig,
        sample_set: base_params.sample_set,
        inline_array: base_params.inline_array,
        chunks: base_params.chunks,
    ) -> xr.Dataset:
        # Handle virtual contig.
        if contig in self.virtual_contigs:
            contigs = self.virtual_contigs[contig]
            datasets = [
                self._snp_calls_for_contig(
                    contig=c,
                    sample_set=sample_set,
                    inline_array=inline_array,
                    chunks=chunks,
                )
                for c in contigs
            ]
            ds = _simple_xarray_concat(datasets, dim=DIM_VARIANT)
            return ds

        # Handle contig in the reference genome.
        else:
            assert contig in self.contigs

            coords = dict()
            data_vars = dict()

            # Set up call arrays.
            calls_root = self.open_snp_genotypes(sample_set=sample_set)
            gt_z = calls_root[f"{contig}/calldata/GT"]
            call_genotype = _da_from_zarr(
                gt_z, inline_array=inline_array, chunks=chunks
            )
            gq_z = calls_root[f"{contig}/calldata/GQ"]
            call_gq = _da_from_zarr(gq_z, inline_array=inline_array, chunks=chunks)
            ad_z = calls_root[f"{contig}/calldata/AD"]
            call_ad = _da_from_zarr(ad_z, inline_array=inline_array, chunks=chunks)
            mq_z = calls_root[f"{contig}/calldata/MQ"]
            call_mq = _da_from_zarr(mq_z, inline_array=inline_array, chunks=chunks)
            data_vars["call_genotype"] = (
                [DIM_VARIANT, DIM_SAMPLE, DIM_PLOIDY],
                call_genotype,
            )
            data_vars["call_GQ"] = ([DIM_VARIANT, DIM_SAMPLE], call_gq)
            data_vars["call_MQ"] = ([DIM_VARIANT, DIM_SAMPLE], call_mq)
            data_vars["call_AD"] = (
                [DIM_VARIANT, DIM_SAMPLE, DIM_ALLELE],
                call_ad,
            )

            # Set up sample arrays.
            z = calls_root["samples"]
            sample_id = _da_from_zarr(z, inline_array=inline_array, chunks=chunks)
            # Decode to unicode strings, as it is stored as bytes objects.
            sample_id = sample_id.astype("U")
            coords["sample_id"] = [DIM_SAMPLE], sample_id

            # Create a dataset.
            ds = xr.Dataset(data_vars=data_vars, coords=coords)

            return ds

    @_check_types
    @doc(
        summary="Access SNP sites, site filters and genotype calls.",
        returns="""A dataset with 4 dimensions:
        `variants` the number of sites in the selected region,
        `allele` the number of alleles (4),
        `samples` the number of samples,
        and `ploidy` the ploidy (2). There are 3 coordinates:
        `variant_position` has `variants` values and contains the position of each site,
        `variant_contig` has `variants` values and contains the contig of each site,
        `sample_id` has `samples` values and contains the identifier of each sample. The data variables are:
        `variant_allele`, it has (`variants`, `alleles`) values and contains the reference followed by the 3 alternate alleles for each site,
        `variant_filter_pass_gamb_colu_arab`, it has (`variants`) values and contains whether the site passes the site filters for all of gambiae, coluzzii and arabiensis (only available for *Ag3*),
        `variant_filter_pass_gamb_colu`, it has (`variants`) values and contains whether the site passes the site filters for gambiae and coluzzii (only available for *Ag3*),
        `variant_filter_pass_arab`, it has (`variants`) values and contains whether the site passes the site filters for arabiensis (only available for *Ag3*),
        `variant_filter_pass_funestus`, it has (`variants`) values and contains whether the site passes the site filters for funestus (only available for *Af1*),
        `call_genotype`, it has (`variants`, `samples`, `ploidy`) values and contains both calls for each site and each sample,
        `call_GQ`, it has (`variants`, `samples`) values and contains the genotype quality for each site and each sample,
        `call_MQ`, it has (`variants`, `samples`) values and contains the mapping quality for each site and each sample,
        `call_AD`, it has (`variants*, `samples`, *alleles`) values and contains the allele depth for each site, each sample and each allele,
        `call_genotypes_mask`, it has (`variants`,`samples`, `ploidy`) values and contains whether the allele is absent for each site, each sample and each ploidy.
        """,
    )
    def snp_calls(
        self,
        region: base_params.regions,
        sample_sets: Optional[base_params.sample_sets] = None,
        sample_query: Optional[base_params.sample_query] = None,
        sample_query_options: Optional[base_params.sample_query_options] = None,
        sample_indices: Optional[base_params.sample_indices] = None,
        site_mask: Optional[base_params.site_mask] = None,
        site_class: Optional[base_params.site_class] = None,
        inline_array: base_params.inline_array = base_params.inline_array_default,
        chunks: base_params.chunks = base_params.native_chunks,
        cohort_size: Optional[base_params.cohort_size] = None,
        min_cohort_size: Optional[base_params.min_cohort_size] = None,
        max_cohort_size: Optional[base_params.max_cohort_size] = None,
        random_seed: base_params.random_seed = 42,
    ) -> xr.Dataset:
        # Check that either sample_query xor sample_indices are provided.
        base_params._validate_sample_selection_params(
            sample_query=sample_query, sample_indices=sample_indices
        )

        # Normalise parameters.
        prepared_regions = _parse_multi_region(self, region)
        prepared_site_mask = self._prep_optional_site_mask_param(site_mask=site_mask)

        # Note: `_prep_sample_selection_cache_params` converts `sample_query` and `sample_query_options` into `sample_indices`.
        # So `sample_query` and `sample_query_options` should not be used beyond this point. (`sample_indices` should be used instead.)
        (
            prepared_sample_sets,
            prepared_sample_indices,
        ) = self._prep_sample_selection_cache_params(
            sample_sets=sample_sets,
            sample_query=sample_query,
            sample_query_options=sample_query_options,
            sample_indices=sample_indices,
        )

        # Delete original parameters to prevent accidental use.
        del sample_sets
        del sample_query
        del sample_indices
        del region
        del site_mask

        # Convert lists to tuples to avoid CacheMiss "TypeError: unhashable type: 'list'".
        prepared_regions_tuple: Tuple[Region, ...] = tuple(prepared_regions)
        prepared_sample_sets_tuple: Optional[Tuple[str, ...]] = (
            tuple(prepared_sample_sets) if prepared_sample_sets is not None else None
        )
        prepared_sample_indices_tuple: Optional[Tuple[int, ...]] = (
            tuple(prepared_sample_indices)
            if prepared_sample_indices is not None
            else None
        )

        # Note: `_snp_calls` should only take `sample_indices`, not `sample_query`, to facilitate caching.
        return self._snp_calls(
            regions=prepared_regions_tuple,
            sample_sets=prepared_sample_sets_tuple,
            sample_indices=prepared_sample_indices_tuple,
            site_mask=prepared_site_mask,
            site_class=site_class,
            cohort_size=cohort_size,
            min_cohort_size=min_cohort_size,
            max_cohort_size=max_cohort_size,
            random_seed=random_seed,
            inline_array=inline_array,
            chunks=chunks,
        )

    # Here we cache to improve performance for functions which
    # access SNP calls more than once. For example, this currently
    # happens during access of biallelic SNP calls, because a
    # first computation of allele counts is required, before
    # then using that to filter SNP calls.
    #
    # We only cache up to 2 items because otherwise we can see
    # high memory usage.
    @lru_cache(maxsize=2)
    def _cached_snp_calls(
        self,
        *,
        regions: Tuple[Region, ...],
        sample_sets,
        site_mask,
        site_class,
        inline_array,
        chunks,
    ):
        # Access SNP calls and concatenate multiple sample sets and/or regions.
        with self._spinner("Access SNP calls"):
            lx = []
            for r in regions:
                ly = []
                for s in sample_sets:
                    y = self._snp_calls_for_contig(
                        contig=r.contig,
                        sample_set=s,
                        inline_array=inline_array,
                        chunks=chunks,
                    )
                    ly.append(y)

                # Concatenate data from multiple sample sets.
                x = _simple_xarray_concat(ly, dim=DIM_SAMPLE)

                # Add variants variables.
                v = self._snp_variants_for_contig(
                    contig=r.contig, inline_array=inline_array, chunks=chunks
                )
                x = xr.merge([v, x], compat="override", join="override")

                # Handle region, do this only once - optimisation.
                if r.start or r.end:
                    pos = x["variant_position"].values
                    loc_region = _locate_region(r, pos)
                    x = x.isel(variants=loc_region)

                # Handle site class.
                if site_class is not None:
                    loc_ann = self._locate_site_class(
                        region=r,
                        site_class=site_class,
                        site_mask=None,
                        inline_array=inline_array,
                        chunks=chunks,
                    )
                    assert x.sizes["variants"] == loc_ann.shape[0]
                    x = x.isel(variants=loc_ann)

                lx.append(x)

            # Concatenate data from multiple regions.
            ds = _simple_xarray_concat(lx, dim=DIM_VARIANT)

        if site_mask is not None:
            with self._spinner(desc="Apply site filters"):
                ds = _dask_compress_dataset(
                    ds,
                    indexer=f"variant_filter_pass_{site_mask}",
                    dim=DIM_VARIANT,
                )

        # Add call_genotype_mask.
        ds["call_genotype_mask"] = ds["call_genotype"] < 0

        return ds

    def _snp_calls(
        self,
        *,
        regions: Tuple[Region, ...],
        sample_sets,
        sample_indices,
        site_mask,
        site_class,
        cohort_size,
        min_cohort_size,
        max_cohort_size,
        random_seed,
        inline_array,
        chunks,
    ):
        ## Get SNP calls and concatenate multiple sample sets and/or regions.

        # Note: sample_sets should be "prepared" before being passed to this private function.

        # Note: `_snp_calls` should only take `sample_indices`, not `sample_query`.
        #       Use `_prep_sample_selection_cache_params` to convert `sample_query` to `sample_indices`.

        # Note: we don't cache different sample_indices subsets, which are selected below.

        ds = self._cached_snp_calls(
            regions=regions,
            sample_sets=sample_sets,
            site_mask=site_mask,
            site_class=site_class,
            inline_array=inline_array,
            chunks=chunks,
        )

        # Handle sample selection.
        if sample_indices is not None:
            # Note: `sample_indices` could be any tuple of integers, while the `ds` DataSet will contain data for all samples in the `sample_sets`.
            # In other words, the internal `sample_query` is not being applied to `ds`.
            # We need to get the filtered set of samples from `sample_metadata` and then select samples based on that set.

            # Get the relevant sample metadata.
            relevant_samples_df = self.sample_metadata(sample_sets=sample_sets)

            # We need to select only the samples that are identified by the `sample_indices` tuple relative to the results of `sample_metadata`.
            # However, the `ds` DataSet contains data for all samples in the `sample_sets`, regardless of any internal `sample_query`.

            # Get the samples identified via `sample_indices`.
            # Note: this might raise `IndexingError` if the user provides bad indices, e.g. "positional indexers are out-of-bounds".
            # Note: `sample_indices` needs to be a list rather than tuple for `iloc`, otherwise `IndexingError`, e.g. "Too many indexers".
            sample_indices_as_list = list(sample_indices)
            selected_samples_df = relevant_samples_df.iloc[sample_indices_as_list]

            # Get the selected sample ids from the sample metadata DataFrame.
            relevant_sample_ids = selected_samples_df["sample_id"].values

            # Get all the sample ids from the unfiltered Dataset.
            ds_sample_ids = ds.coords["sample_id"].values

            # Get the indices of samples in the Dataset that match the relevant sample ids.
            # Note: we use `[0]` to get the first element of the tuple returned by `np.where`.
            relevant_sample_indices = np.where(
                np.isin(ds_sample_ids, relevant_sample_ids)
            )[0]

            # Preserve the behaviour of raising a `ValueError` instead of empty results.
            if relevant_sample_indices.size == 0:
                raise ValueError("No relevant samples found.")

            # Select only the relevant samples from the Dataset.
            ds = ds.isel(samples=relevant_sample_indices)

        # Handle cohort size, overrides min and max.
        if cohort_size is not None:
            min_cohort_size = cohort_size
            max_cohort_size = cohort_size

        # Handle min cohort size.
        if min_cohort_size is not None:
            n_samples = ds.sizes["samples"]
            if n_samples < min_cohort_size:
                raise ValueError(
                    f"not enough samples ({n_samples}) for minimum cohort size ({min_cohort_size})"
                )

        # Handle max cohort size.
        if max_cohort_size is not None:
            n_samples = ds.sizes["samples"]
            if n_samples > max_cohort_size:
                rng = np.random.default_rng(seed=random_seed)
                loc_downsample = rng.choice(
                    n_samples, size=max_cohort_size, replace=False
                )
                loc_downsample.sort()
                ds = ds.isel(samples=loc_downsample)

        return ds

    def snp_dataset(self, *args, **kwargs):  # pragma: no cover
        """Deprecated, this method has been renamed to snp_calls()."""
        return self.snp_calls(*args, **kwargs)

    def _prep_region_cache_param(
        self, *, region: base_params.regions
    ) -> Union[dict, List[dict]]:
        """Obtain a normalised representation of a region parameter which can
        be used with the results cache."""

        # N.B., we need to convert to a dict, because cache saves params as
        # JSON.

        region_prepped: List[Region] = _parse_multi_region(self, region)
        if len(region_prepped) > 1:
            ret = [r.to_dict() for r in region_prepped]
        else:
            ret = region_prepped[0].to_dict()
        return ret

    def _results_cache_add_analysis_params(self, params: dict):
        super()._results_cache_add_analysis_params(params)
        params["site_filters_analysis"] = self._site_filters_analysis

    def _snp_allele_counts(
        self,
        *,
        region,
        sample_sets,
        sample_indices,
        site_mask,
        site_class,
        cohort_size,
        min_cohort_size,
        max_cohort_size,
        random_seed,
        inline_array,
        chunks,
    ):
        # Access SNP calls.
        ds_snps = self.snp_calls(
            region=region,
            sample_sets=sample_sets,
            sample_indices=sample_indices,
            site_mask=site_mask,
            site_class=site_class,
            cohort_size=cohort_size,
            min_cohort_size=min_cohort_size,
            max_cohort_size=max_cohort_size,
            random_seed=random_seed,
            inline_array=inline_array,
            chunks=chunks,
        )
        gt = ds_snps["call_genotype"]

        # Set up and run allele counts computation.
        gt = allel.GenotypeDaskArray(gt.data)
        ac = gt.count_alleles(max_allele=3)
        with self._dask_progress(desc="Compute SNP allele counts"):
            ac = ac.compute()

        # Return plain numpy array.
        results = dict(ac=ac.values)

        return results

    @_check_types
    @doc(
        summary="""
            Compute SNP allele counts. This returns the number of times each
            SNP allele was observed in the selected samples.
        """,
        returns="""
            A numpy array of shape (n_variants, 4), where the first column has
            the reference allele (0) counts, the second column has the first
            alternate allele (1) counts, the third column has the second
            alternate allele (2) counts, and the fourth column has the third
            alternate allele (3) counts.
        """,
        notes="""
            This computation may take some time to run, depending on your
            computing environment. Results of this computation will be cached
            and re-used if the `results_cache` parameter was set when
            instantiating the class.
        """,
    )
    def snp_allele_counts(
        self,
        region: base_params.regions,
        sample_sets: Optional[base_params.sample_sets] = None,
        sample_query: Optional[base_params.sample_query] = None,
        sample_query_options: Optional[base_params.sample_query_options] = None,
        sample_indices: Optional[base_params.sample_indices] = None,
        site_mask: Optional[base_params.site_mask] = None,
        site_class: Optional[base_params.site_class] = None,
        cohort_size: Optional[base_params.cohort_size] = None,
        min_cohort_size: Optional[base_params.min_cohort_size] = None,
        max_cohort_size: Optional[base_params.max_cohort_size] = None,
        random_seed: base_params.random_seed = 42,
        inline_array: base_params.inline_array = base_params.inline_array_default,
        chunks: base_params.chunks = base_params.native_chunks,
    ) -> np.ndarray:
        # Change this name if you ever change the behaviour of this function,
        # to invalidate any previously cached data.
        name = "snp_allele_counts_v2"

        # Check that either sample_query xor sample_indices are provided.
        base_params._validate_sample_selection_params(
            sample_query=sample_query, sample_indices=sample_indices
        )

        ## Normalize params for consistent hash value.

        # Note: `_prep_sample_selection_cache_params` converts `sample_query` and `sample_query_options` into `sample_indices`.
        # So `sample_query` and `sample_query_options` should not be used beyond this point. (`sample_indices` should be used instead.)
        (
            sample_sets_prepped,
            sample_indices_prepped,
        ) = self._prep_sample_selection_cache_params(
            sample_sets=sample_sets,
            sample_query=sample_query,
            sample_query_options=sample_query_options,
            sample_indices=sample_indices,
        )
        del sample_sets
        del sample_query
        del sample_query_options
        del sample_indices
        region_prepped = self._prep_region_cache_param(region=region)
        del region
        site_mask_prepped = self._prep_optional_site_mask_param(site_mask=site_mask)
        del site_mask
        params = dict(
            region=region_prepped,
            sample_sets=sample_sets_prepped,
            sample_indices=sample_indices_prepped,
            site_mask=site_mask_prepped,
            site_class=site_class,
            cohort_size=cohort_size,
            min_cohort_size=min_cohort_size,
            max_cohort_size=max_cohort_size,
            random_seed=random_seed,
        )

        try:
            results = self.results_cache_get(name=name, params=params)

        except CacheMiss:
            results = self._snp_allele_counts(
                **params, inline_array=inline_array, chunks=chunks
            )
            self.results_cache_set(name=name, params=params, results=results)

        ac = results["ac"]
        return ac

    @_check_types
    @doc(
        summary="""
            Plot SNPs in a given genome region. SNPs are shown as rectangles,
            with segregating and non-segregating SNPs positioned on different levels,
            and coloured by site filter.
        """,
        parameters=dict(
            max_snps="Maximum number of SNPs to show.",
        ),
    )
    def plot_snps(
        self,
        region: base_params.regions,
        sample_sets: Optional[base_params.sample_sets] = None,
        sample_query: Optional[base_params.sample_query] = None,
        sample_query_options: Optional[base_params.sample_query_options] = None,
        site_mask: base_params.site_mask = base_params.DEFAULT,
        cohort_size: Optional[base_params.cohort_size] = None,
        sizing_mode: gplt_params.sizing_mode = gplt_params.sizing_mode_default,
        width: gplt_params.width = gplt_params.width_default,
        track_height: gplt_params.height = 80,
        genes_height: gplt_params.genes_height = gplt_params.genes_height_default,
        max_snps: int = 200_000,
        show: gplt_params.show = True,
        gene_labels: Optional[gplt_params.gene_labels] = None,
        gene_labelset: Optional[gplt_params.gene_labelset] = None,
    ) -> gplt_params.optional_figure:
        # Plot SNPs track.
        fig1 = self.plot_snps_track(
            region=region,
            sample_sets=sample_sets,
            sample_query=sample_query,
            sample_query_options=sample_query_options,
            site_mask=site_mask,
            cohort_size=cohort_size,
            sizing_mode=sizing_mode,
            width=width,
            height=track_height,
            max_snps=max_snps,
            show=False,
        )
        fig1.xaxis.visible = False

        # Plot genes track.
        fig2 = self.plot_genes(
            region=region,
            sizing_mode=sizing_mode,
            width=width,
            height=genes_height,
            x_range=fig1.x_range,
            show=False,
            gene_labels=gene_labels,
            gene_labelset=gene_labelset,
        )

        # Layout tracks in a grid.
        fig = bokeh.layouts.gridplot(
            [fig1, fig2],
            ncols=1,
            toolbar_location="above",
            merge_tools=True,
            sizing_mode=sizing_mode,
        )

        if show:  # pragma: no cover
            bokeh.plotting.show(fig)
            return None
        else:
            return fig

    @_check_types
    @doc(
        summary="""
            Plot SNPs in a given genome region. SNPs are shown as rectangles,
            with segregating and non-segregating SNPs positioned on different levels,
            and coloured by site filter.
        """,
        parameters=dict(
            max_snps="Maximum number of SNPs to show.",
        ),
    )
    def plot_snps_track(
        self,
        region: base_params.region,
        sample_sets: Optional[base_params.sample_sets] = None,
        sample_query: Optional[base_params.sample_query] = None,
        sample_query_options: Optional[base_params.sample_query_options] = None,
        site_mask: base_params.site_mask = base_params.DEFAULT,
        cohort_size: Optional[base_params.cohort_size] = None,
        min_cohort_size: Optional[base_params.min_cohort_size] = None,
        max_cohort_size: Optional[base_params.max_cohort_size] = None,
        sizing_mode: gplt_params.sizing_mode = gplt_params.sizing_mode_default,
        width: gplt_params.width = gplt_params.width_default,
        height: gplt_params.height = 120,
        max_snps: int = 200_000,
        x_range: Optional[gplt_params.x_range] = None,
        show: gplt_params.show = True,
        output_backend: gplt_params.output_backend = gplt_params.output_backend_default,
    ) -> gplt_params.optional_figure:
        # Normalise params.
        site_mask_prepped = self._prep_site_mask_param(site_mask=site_mask)
        del site_mask

        # Resolve and check region.
        resolved_region: Region = _parse_single_region(self, region)
        del region

        if (
            (resolved_region.start is None)
            or (resolved_region.end is None)
            or ((resolved_region.end - resolved_region.start) > max_snps)
        ):
            raise ValueError("Region is too large, please provide a smaller region.")

        # Compute allele counts.
        ac = allel.AlleleCountsArray(
            self.snp_allele_counts(
                region=resolved_region,
                sample_sets=sample_sets,
                sample_query=sample_query,
                sample_query_options=sample_query_options,
                site_mask=None,
                cohort_size=cohort_size,
                min_cohort_size=min_cohort_size,
                max_cohort_size=max_cohort_size,
            )
        )
        an = ac.sum(axis=1)
        is_seg = ac.is_segregating()
        is_var = ac.is_variant()
        allelism = ac.allelism()

        # Obtain SNP variants data.
        ds_sites = self.snp_variants(
            region=resolved_region,
        ).compute()

        # Build a dataframe.
        pos = ds_sites["variant_position"].values
        alleles = ds_sites["variant_allele"].values.astype("U")
        cols = {
            "pos": pos,
            "allele_0": alleles[:, 0],
            "allele_1": alleles[:, 1],
            "allele_2": alleles[:, 2],
            "allele_3": alleles[:, 3],
            "ac_0": ac[:, 0],
            "ac_1": ac[:, 1],
            "ac_2": ac[:, 2],
            "ac_3": ac[:, 3],
            "an": an,
            "is_seg": is_seg,
            "is_var": is_var,
            "allelism": allelism,
        }

        for site_mask_id in self.site_mask_ids:
            cols[f"pass_{site_mask_id}"] = ds_sites[
                f"variant_filter_pass_{site_mask_id}"
            ].values

        data = pd.DataFrame(cols)

        # Find gaps in the reference genome.
        seq = self.genome_sequence(region=resolved_region.contig).compute()
        is_n = (seq == b"N") | (seq == b"n")
        n_starts, n_stops = _true_runs(is_n)

        # Create figure.
        xwheel_zoom = bokeh.models.WheelZoomTool(
            dimensions="width", maintain_focus=False
        )
        pos = data["pos"].values
        x_min = resolved_region.start or 1
        x_max = resolved_region.end or len(seq)
        if x_range is None:
            x_range = bokeh.models.Range1d(x_min, x_max, bounds="auto")

        tooltips = [
            ("Position", "$x{0,0}"),
            (
                "Alleles",
                "@allele_0 (@ac_0), @allele_1 (@ac_1), @allele_2 (@ac_2), @allele_3 (@ac_3)",
            ),
            ("No. alleles", "@allelism"),
            ("Allele calls", "@an"),
        ]

        for site_mask_id in self.site_mask_ids:
            tooltips.append((f"Pass {site_mask_id}", f"@pass_{site_mask_id}"))

        fig = bokeh.plotting.figure(
            title="SNPs",
            tools=["xpan", "xzoom_in", "xzoom_out", xwheel_zoom, "reset"],
            active_scroll=xwheel_zoom,
            active_drag="xpan",
            sizing_mode=sizing_mode,
            width=width,
            height=height,
            toolbar_location="above",
            x_range=x_range,
            y_range=(0.5, 2.5),
            tooltips=tooltips,
            output_backend=output_backend,
        )
        hover_tool = fig.select(type=bokeh.models.HoverTool)
        hover_tool.name = "snps"

        # Plot gaps in the reference genome.
        df_n_runs = pd.DataFrame(
            {"left": n_starts + 0.6, "right": n_stops + 0.4, "top": 2.5, "bottom": 0.5}
        )
        fig.quad(
            top="top",
            bottom="bottom",
            left="left",
            right="right",
            color="#cccccc",
            source=df_n_runs,
            name="gaps",
            line_width=0,
        )

        # Plot SNPs.
        color_pass = bokeh.palettes.Colorblind6[3]
        color_fail = bokeh.palettes.Colorblind6[5]
        data["left"] = data["pos"] - 0.4
        data["right"] = data["pos"] + 0.4
        data["bottom"] = np.where(data["is_seg"], 1.6, 0.6)
        data["top"] = data["bottom"] + 0.8
        data["color"] = np.where(
            data[f"pass_{site_mask_prepped}"], color_pass, color_fail
        )
        fig.quad(
            top="top",
            bottom="bottom",
            left="left",
            right="right",
            color="color",
            source=data,
            name="snps",
        )

        # Tidy plot.
        fig.yaxis.ticker = bokeh.models.FixedTicker(
            ticks=[1, 2],
        )
        fig.yaxis.major_label_overrides = {
            1: "Non-segregating",
            2: "Segregating",
        }
        fig.xaxis.axis_label = f"Contig {resolved_region.contig} position (bp)"
        fig.xaxis.ticker = bokeh.models.AdaptiveTicker(min_interval=1)
        fig.xaxis.minor_tick_line_color = None
        fig.xaxis[0].formatter = bokeh.models.NumeralTickFormatter(format="0,0")

        if show:  # pragma: no cover
            bokeh.plotting.show(fig)
            return None
        else:
            return fig

    @_check_types
    @doc(
        summary="Compute genome accessibility array.",
        returns="An array of boolean values identifying accessible genome sites.",
    )
    def is_accessible(
        self,
        region: base_params.region,
        site_mask: base_params.site_mask = base_params.DEFAULT,
        inline_array: base_params.inline_array = base_params.inline_array_default,
        chunks: base_params.chunks = base_params.native_chunks,
    ) -> np.ndarray:
        # Normalise params.
        resolved_region: Region = _parse_single_region(self, region)
        del region
        site_mask_prepped = self._prep_site_mask_param(site_mask=site_mask)
        del site_mask

        # Determine contig sequence length.
        seq_length = self.genome_sequence(resolved_region).shape[0]

        # Set up output.
        is_accessible = np.zeros(seq_length, dtype=bool)

        # Access SNP site positions.
        pos = self.snp_sites(region=resolved_region, field="POS").compute()
        if resolved_region.start:
            offset = resolved_region.start
        else:
            offset = 1

        # Access site filters.
        filter_pass = self._site_filters_for_region(
            region=resolved_region,
            mask=site_mask_prepped,
            field="filter_pass",
            inline_array=inline_array,
            chunks=chunks,
        ).compute()

        # Assign values from site filters.
        is_accessible[pos - offset] = filter_pass

        return is_accessible

    @_check_types
    @doc(
        summary="Access SNP calls at sites which are biallelic within the selected samples.",
        returns="""
        A dataset with 4 dimensions: `variants`: the number of sites, `alleles` the number of alleles (2), `samples`: the number of samples, `ploidy`: the ploidy (2). There are 3 coordinates:
        `sample_id` has `samples` values and contains the identifier of each sample,
        `variant_position` has `variants` values and contains the position of each site,
        `variant_contig` has `variants` values and contains the contig of each site. There are 3 data variables:
        `variant_allele`, it (`variants`, `alleles`) values and contains the nucleotide represented as a character for each site, each allele,
        `variant_allele_count`, it (`variants`, `alleles`) values and contains the number of occurences of each allele at each site,
        `call_genotype`, it (`variants`, `samples`, `ploidy`) values and contains and contains both calls for each site and each sample.
          """,
    )
    def biallelic_snp_calls(
        self,
        region: base_params.regions,
        sample_sets: Optional[base_params.sample_sets] = None,
        sample_query: Optional[base_params.sample_query] = None,
        sample_query_options: Optional[base_params.sample_query_options] = None,
        sample_indices: Optional[base_params.sample_indices] = None,
        site_mask: Optional[base_params.site_mask] = None,
        site_class: Optional[base_params.site_class] = None,
        inline_array: base_params.inline_array = base_params.inline_array_default,
        chunks: base_params.chunks = base_params.native_chunks,
        cohort_size: Optional[base_params.cohort_size] = None,
        min_cohort_size: Optional[base_params.min_cohort_size] = None,
        max_cohort_size: Optional[base_params.max_cohort_size] = None,
        random_seed: base_params.random_seed = 42,
        min_minor_ac: Optional[base_params.min_minor_ac] = None,
        max_missing_an: Optional[base_params.max_missing_an] = None,
        n_snps: Optional[base_params.n_snps] = None,
        thin_offset: base_params.thin_offset = 0,
    ) -> xr.Dataset:
        # Check that either sample_query xor sample_indices are provided.
        base_params._validate_sample_selection_params(
            sample_query=sample_query, sample_indices=sample_indices
        )

        # Perform an allele count.
        ac = self.snp_allele_counts(
            region=region,
            sample_sets=sample_sets,
            sample_query=sample_query,
            sample_query_options=sample_query_options,
            sample_indices=sample_indices,
            site_mask=site_mask,
            site_class=site_class,
            cohort_size=cohort_size,
            min_cohort_size=min_cohort_size,
            max_cohort_size=max_cohort_size,
            random_seed=random_seed,
            inline_array=inline_array,
            chunks=chunks,
        )

        # Locate biallelic SNPs.
        loc_bi = allel.AlleleCountsArray(ac).is_biallelic()

        # Remap alleles to squeeze out unobserved alleles.
        ac_bi = ac[loc_bi]
        allele_mapping = _trim_alleles(ac_bi)

        # Set up SNP calls.
        ds = self.snp_calls(
            region=region,
            sample_sets=sample_sets,
            sample_query=sample_query,
            sample_query_options=sample_query_options,
            sample_indices=sample_indices,
            site_mask=site_mask,
            site_class=site_class,
            cohort_size=cohort_size,
            min_cohort_size=min_cohort_size,
            max_cohort_size=max_cohort_size,
            random_seed=random_seed,
            inline_array=inline_array,
            chunks=chunks,
        )

        with self._spinner("Prepare biallelic SNP calls"):
            # Subset to biallelic sites.
            ds_bi = _dask_compress_dataset(ds, indexer=loc_bi, dim="variants")

            # Start building a new dataset.
            coords: Dict[str, Any] = dict()
            data_vars: Dict[str, Any] = dict()

            # Store sample IDs.
            coords["sample_id"] = ("samples",), ds_bi["sample_id"].data

            # Store contig.
            coords["variant_contig"] = ("variants",), ds_bi["variant_contig"].data

            # Store position.
            variant_position = ds_bi["variant_position"].data
            coords["variant_position"] = ("variants",), variant_position

            # Store alleles, transformed.
            variant_allele_dask = ds_bi["variant_allele"].data
            variant_allele_out = _dask_apply_allele_mapping(
                variant_allele_dask, allele_mapping, max_allele=1
            )
            data_vars["variant_allele"] = ("variants", "alleles"), variant_allele_out

            # Store allele counts, transformed.
            ac_out = _apply_allele_mapping(ac_bi, allele_mapping, max_allele=1)
            data_vars["variant_allele_count"] = ("variants", "alleles"), ac_out

            # Store genotype calls, transformed.
            gt_dask = ds_bi["call_genotype"].data
            gt_out = _dask_genotype_array_map_alleles(gt_dask, allele_mapping)
            data_vars["call_genotype"] = (
                (
                    "variants",
                    "samples",
                    "ploidy",
                ),
                gt_out,
            )

            # Build dataset.
            ds_out = xr.Dataset(coords=coords, data_vars=data_vars, attrs=ds.attrs)

            # Apply conditions.
            if max_missing_an is not None or min_minor_ac is not None:
                loc_out = np.ones(ds_out.sizes["variants"], dtype=bool)
                an = ac_out.sum(axis=1)

                # Apply missingness condition.
                if max_missing_an is not None:
                    an_missing = (ds_out.sizes["samples"] * ds_out.sizes["ploidy"]) - an
                    if isinstance(max_missing_an, float):
                        an_missing_frac = an_missing / an
                        loc_missing = an_missing_frac <= max_missing_an
                    else:
                        loc_missing = an_missing <= max_missing_an
                    loc_out &= loc_missing

                # Apply minor allele count condition.
                if min_minor_ac is not None:
                    ac_minor = ac_out.min(axis=1)
                    if isinstance(min_minor_ac, float):
                        ac_minor_frac = ac_minor / an
                        loc_minor = ac_minor_frac >= min_minor_ac
                    else:
                        loc_minor = ac_minor >= min_minor_ac
                    loc_out &= loc_minor

                # Apply selection from conditions.
                ds_out = _dask_compress_dataset(ds_out, indexer=loc_out, dim="variants")

            # Try to meet target number of SNPs.
            if n_snps is not None:
                if ds_out.sizes["variants"] > (n_snps * 2):
                    # Apply thinning.
                    thin_step = ds_out.sizes["variants"] // n_snps
                    loc_thin = slice(thin_offset, None, thin_step)
                    ds_out = ds_out.isel(variants=loc_thin)

                elif ds_out.sizes["variants"] < n_snps:
                    raise ValueError("Not enough SNPs.")

        return ds_out

    @_check_types
    @doc(
        summary="Load biallelic SNP genotypes.",
        returns=dict(
            gn="""
                An array of shape (variants, samples) where each value counts the
                number of alternate alleles per genotype call.
            """,
            samples="Sample identifiers.",
        ),
    )
    def biallelic_diplotypes(
        self,
        region: base_params.regions,
        sample_sets: Optional[base_params.sample_sets] = None,
        sample_query: Optional[base_params.sample_query] = None,
        sample_query_options: Optional[base_params.sample_query_options] = None,
        sample_indices: Optional[base_params.sample_indices] = None,
        site_mask: Optional[base_params.site_mask] = None,
        site_class: Optional[base_params.site_class] = None,
        cohort_size: Optional[base_params.cohort_size] = None,
        min_cohort_size: Optional[base_params.min_cohort_size] = None,
        max_cohort_size: Optional[base_params.max_cohort_size] = None,
        random_seed: base_params.random_seed = 42,
        min_minor_ac: Optional[base_params.min_minor_ac] = None,
        max_missing_an: Optional[base_params.max_missing_an] = None,
        n_snps: Optional[base_params.n_snps] = None,
        thin_offset: base_params.thin_offset = 0,
        inline_array: base_params.inline_array = base_params.inline_array_default,
        chunks: base_params.chunks = base_params.native_chunks,
    ) -> Tuple[np.ndarray, np.ndarray]:
        # Change this name if you ever change the behaviour of this function, to
        # invalidate any previously cached data.
        name = "biallelic_diplotypes"

        # Check that either sample_query xor sample_indices are provided.
        base_params._validate_sample_selection_params(
            sample_query=sample_query, sample_indices=sample_indices
        )

        ## Normalize params for consistent hash value.

        # Note: `_prep_sample_selection_cache_params` converts `sample_query` and `sample_query_options` into `sample_indices`.
        # So `sample_query` and `sample_query_options` should not be used beyond this point. (`sample_indices` should be used instead.)
        (
            prepared_sample_sets,
            prepared_sample_indices,
        ) = self._prep_sample_selection_cache_params(
            sample_sets=sample_sets,
            sample_query=sample_query,
            sample_query_options=sample_query_options,
            sample_indices=sample_indices,
        )
        prepared_region = self._prep_region_cache_param(region=region)
        prepared_site_mask = self._prep_optional_site_mask_param(site_mask=site_mask)

        # Delete original parameters to prevent accidental use.
        del sample_sets
        del sample_query
        del sample_query_options
        del sample_indices
        del region
        del site_mask

        params = dict(
            region=prepared_region,
            n_snps=n_snps,
            thin_offset=thin_offset,
            sample_sets=prepared_sample_sets,
            sample_indices=prepared_sample_indices,
            site_mask=prepared_site_mask,
            site_class=site_class,
            cohort_size=cohort_size,
            min_cohort_size=min_cohort_size,
            max_cohort_size=max_cohort_size,
            random_seed=random_seed,
            min_minor_ac=min_minor_ac,
            max_missing_an=max_missing_an,
        )

        # Try to retrieve results from the cache.
        try:
            results = self.results_cache_get(name=name, params=params)

        except CacheMiss:
            results = self._biallelic_diplotypes(
                inline_array=inline_array, chunks=chunks, **params
            )
            self.results_cache_set(name=name, params=params, results=results)

        # Unpack results.
        gn = results["gn"]
        samples = results["samples"]

        return gn, samples

    def _biallelic_diplotypes(
        self,
        *,
        region,
        sample_sets,
        sample_indices,
        site_mask,
        site_class,
        cohort_size,
        min_cohort_size,
        max_cohort_size,
        random_seed,
        max_missing_an,
        min_minor_ac,
        n_snps,
        thin_offset,
        inline_array,
        chunks,
    ):
        # Note: this function uses sample_indices and should not expect a sample_query.

        # Access biallelic SNPs.
        ds = self.biallelic_snp_calls(
            region=region,
            sample_sets=sample_sets,
            sample_indices=sample_indices,
            site_mask=site_mask,
            site_class=site_class,
            cohort_size=cohort_size,
            min_cohort_size=min_cohort_size,
            max_cohort_size=max_cohort_size,
            random_seed=random_seed,
            max_missing_an=max_missing_an,
            min_minor_ac=min_minor_ac,
            n_snps=n_snps,
            thin_offset=thin_offset,
            inline_array=inline_array,
            chunks=chunks,
        )

        # Load sample IDs
        samples = ds["sample_id"].values.astype("U")

        # Compute diplotypes as the number of alt alleles per genotype call.
        gt = allel.GenotypeDaskArray(ds["call_genotype"].data)
        with self._dask_progress(desc="Compute biallelic diplotypes"):
            gn = gt.to_n_alt().compute()

        return dict(samples=samples, gn=gn)
