"""Single-cell allele counting functions."""

from __future__ import annotations

import logging
import timeit
from collections import defaultdict
from collections.abc import Iterator

import anndata as ad
import numpy as np
import pandas as pd
import polars as pl
from pysam.libcalignmentfile import AlignmentFile
from scipy.sparse import csr_matrix

# Local imports
from .count_alleles import find_read_aln_pos

logger = logging.getLogger(__name__)


class CountStatsSC:
    """Container for mutable single-cell counting statistics.

    Tracks allele counts and metadata per chromosome during counting.
    """

    def __init__(self) -> None:
        self.ref_count: defaultdict[tuple[int, int], int] = defaultdict(int)
        self.alt_count: defaultdict[tuple[int, int], int] = defaultdict(int)
        self.other_count: defaultdict[tuple[int, int], int] = defaultdict(int)

        # Keep track of metadata

        # Number
        self.num_snps: defaultdict[str, int] = defaultdict(int)
        self.num_barcodes: defaultdict[str, int] = defaultdict(int)
        self.reads_counted: defaultdict[str, int] = defaultdict(int)

        # Reads that were not counted
        self.reads_skipped_no_barcode: defaultdict[str, int] = defaultdict(int)
        self.reads_skipped_barcode_no_index: defaultdict[str, int] = defaultdict(int)
        self.reads_skipped_prev_counted: defaultdict[str, int] = defaultdict(int)
        self.reads_skipped_no_sequence: defaultdict[str, int] = defaultdict(int)
        self.reads_skipped_no_aln_pos: defaultdict[str, int] = defaultdict(int)
        self.reads_skipped_seq_error: defaultdict[str, int] = defaultdict(int)

    def stats_to_df(self) -> pd.DataFrame:
        """Convert statistics to a pandas DataFrame."""
        stat_attributes = [
            "num_snps",
            "num_barcodes",
            "reads_counted",
            "reads_skipped_no_barcode",
            "reads_skipped_barcode_no_index",
            "reads_skipped_prev_counted",
            "reads_skipped_no_sequence",
            "reads_skipped_no_aln_pos",
            "reads_skipped_seq_error",
        ]

        stat_df = pd.DataFrame({key: getattr(self, key) for key in stat_attributes}).reset_index(
            names="chrom"
        )

        return stat_df


def make_count_matrix(
    bam_file: str,
    df: pl.DataFrame,
    bc_dict: dict[str, int],
    include_samples: list[str] | None = None,
    include_features: list[str] | None = None,
) -> ad.AnnData:
    """Create sparse count matrix from BAM and variant data.

    Parameters
    ----------
    bam_file : str
        Path to BAM file with cell barcodes.
    df : pl.DataFrame
        DataFrame with variant positions from intersection.
    bc_dict : dict[str, int]
        Mapping of cell barcodes to integer indices.
    include_samples : list[str] | None, optional
        Sample columns to include from variant data, by default None.
    include_features : list[str] | None, optional
        Feature columns to include, by default None.

    Returns
    -------
    ad.AnnData
        AnnData object with count matrices in layers (ref, alt, other).
    """
    chrom_list = df.get_column("chrom").unique(maintain_order=True)
    # chrom_list = chrom_list[:3] # Testing purposes

    # Add genotypes annotations
    # Maybe do this automatically and parse feature col instead?
    snp_df_cols = ["chrom", "pos", "ref", "alt"]
    if include_samples is not None:
        snp_df_cols.extend(include_samples)

    # Might be more memory efficient to use pandas index instead...
    snp_df = df.select(snp_df_cols).unique(maintain_order=True).with_row_index()

    sc_counts = CountStatsSC()  # Class that holds total count data

    with AlignmentFile(bam_file, "rb") as bam:
        for chrom in chrom_list:
            chrom_df = snp_df.filter(pl.col("chrom") == chrom)

            start = timeit.default_timer()

            try:
                count_bc_snp_alleles(
                    bam=bam,
                    bc_dict=bc_dict,
                    chrom=chrom,
                    snp_list=chrom_df.select(["index", "pos", "ref", "alt"]).iter_rows(),
                    sc_counts=sc_counts,
                )

            except ValueError:
                logger.warning("Skipping %s: Contig not found!", chrom)
            else:
                logger.info(
                    "%s: Counted %d SNPs in %.2f seconds",
                    chrom,
                    chrom_df.height,
                    timeit.default_timer() - start,
                )

    # Create sparse matrices
    # sparse array is recommended...but doesnt work with adata
    sparse_ref = csr_matrix(
        (list(sc_counts.ref_count.values()), list(zip(*sc_counts.ref_count.keys()))),
        shape=(snp_df.shape[0], len(bc_dict)),
        dtype=np.uint16,
    )

    sparse_alt = csr_matrix(
        (list(sc_counts.alt_count.values()), list(zip(*sc_counts.alt_count.keys()))),
        shape=(snp_df.shape[0], len(bc_dict)),
        dtype=np.uint16,
    )

    sparse_other = csr_matrix(
        (list(sc_counts.other_count.values()), list(zip(*sc_counts.other_count.keys()))),
        shape=(snp_df.shape[0], len(bc_dict)),
        dtype=np.uint16,
    )

    # Create anndata With total as X
    adata = ad.AnnData(
        X=sparse_ref + sparse_alt + sparse_other,
        layers={"ref": sparse_ref, "alt": sparse_alt, "other": sparse_other},
    )

    # Annotate adata: Figure out what to add to adata here vs later
    adata.obs = snp_df.to_pandas()  # Maybe just switch to pandas? Should i set no copy?
    adata.obs["ref_count"] = adata.layers["ref"].sum(axis=1, dtype=np.uint16).T.A1
    adata.obs["alt_count"] = adata.layers["alt"].sum(axis=1, dtype=np.uint16).T.A1

    # Add barcode names
    adata.var_names = bc_dict.keys()

    # Add genotypes to anndata
    if include_samples is not None:
        adata.uns["samples"] = include_samples

    # TODO: Allow for other features besides 'region' using include_features
    # Could be case of no features, or feature is gene
    if "region" in df.columns:
        # Get unique snps and associated regions

        # Create dict during analysis step instead
        adata.uns["feature"] = (
            df.join(snp_df, on=["chrom", "pos", "ref", "alt"], how="left")
            .select(["region", "index"])
            .to_pandas()
        )

        # region_snp_dict = dict(
        #     df.join(snp_df, on=["chrom", "pos", "ref", "alt"], how="left"
        #            ).group_by("region", maintain_order=True
        #                      ).agg("index").iter_rows()
        # )
        # adata.uns["region_snps"] = region_snp_dict

    # Write out count stats
    adata.uns["count_stats"] = sc_counts.stats_to_df()

    return adata


def count_bc_snp_alleles(
    bam: AlignmentFile,
    bc_dict: dict[str, int],
    chrom: str,
    snp_list: Iterator[tuple[int, int, str, str]],
    sc_counts: CountStatsSC,
) -> None:
    """Count alleles at SNP positions for each cell barcode.

    Parameters
    ----------
    bam : AlignmentFile
        Open BAM file handle.
    bc_dict : dict[str, int]
        Mapping of cell barcodes to indices.
    chrom : str
        Chromosome to process.
    snp_list : Iterator[tuple[int, int, str, str]]
        Iterator of (index, pos, ref, alt) tuples.
    sc_counts : CountStatsSC
        Statistics container to update with counts.
    """
    read_set = set()  # Keep track of reads seen
    bc_set = set()

    for idx, pos, ref, alt in snp_list:
        for read in bam.fetch(chrom, pos - 1, pos):
            # If already counted allele or pair in read
            if read.query_name in read_set:
                sc_counts.reads_skipped_prev_counted[chrom] += 1
                continue

            # Check if there is a read barcode
            try:
                read_bc = str(read.get_tag("CB"))
            except KeyError:
                sc_counts.reads_skipped_no_barcode[chrom] += 1
                continue

            # If barcode not in dict
            if read_bc not in bc_dict:
                sc_counts.reads_skipped_barcode_no_index[chrom] += 1
                continue

            seq = read.query_sequence
            if seq is None:
                sc_counts.reads_skipped_no_sequence[chrom] += 1
                continue

            # Binary search for alignment position
            qpos = find_read_aln_pos(read, pos - 1)
            if qpos is None:
                sc_counts.reads_skipped_no_aln_pos[chrom] += 1
                continue

            try:
                if seq[qpos] == ref:
                    sc_counts.ref_count[(idx, bc_dict[read_bc])] += 1
                elif seq[qpos] == alt:
                    sc_counts.alt_count[(idx, bc_dict[read_bc])] += 1
                else:
                    sc_counts.other_count[(idx, bc_dict[read_bc])] += 1

            except (TypeError, IndexError) as e:
                # Narrow exception handling: only catch sequence access errors
                # Log the actual exception for debugging unexpected errors
                sc_counts.reads_skipped_seq_error[chrom] += 1
                logger.debug(
                    "Skipping read %s: sequence access error at %s:%d (qpos=%s): %s",
                    read.query_name,
                    chrom,
                    pos,
                    qpos,
                    e,
                )
                continue
            else:
                read_set.add(read.query_name)
                bc_set.add(read_bc)
                sc_counts.reads_counted[chrom] += 1

        sc_counts.num_snps[chrom] += 1  # Put here in case of error
    sc_counts.num_barcodes[chrom] = len(bc_set)  # Add unique barcodes
