"""Highly variable gene (HVG) selection utilities."""

import re
from typing import List, Optional, Set, Tuple

import numpy as np


def parse_perturbation_targets(name: str, gene_names: Optional[List[str]] = None) -> List[str]:
    """Parse perturbation name into target gene symbols.

    Uses longest-match greedy parsing against the known gene dictionary to
    correctly handle compound gene names like ``SET_KLF1`` (two genes
    ``SET`` and ``KLF1`` separated by ``_``, the standard delimiter in
    Norman2019 and similar screens).

    If *gene_names* is provided, the parser greedily matches the longest
    known gene name at each position.  Otherwise it falls back to simple
    splitting on ``_`` and ``+``.

    Parameters
    ----------
    name : str
        Raw perturbation label, e.g. ``KLF1``, ``KLF1+CEBPE``, ``SET_KLF1``.
    gene_names : list of str, optional
        Known gene symbols sorted by the dataset's gene dictionary.
        When provided, longest-match parsing is used for correctness.

    Returns
    -------
    list of str
        Non-empty target gene symbols.
    """
    name = str(name)
    if gene_names is not None:
        # Longest-match greedy parsing: at each position, find the longest
        # gene name that matches the remaining string prefix.
        sorted_genes = sorted(gene_names, key=len, reverse=True)
        result = []
        remaining = name
        while remaining:
            matched = False
            for g in sorted_genes:
                if remaining.startswith(g):
                    result.append(g)
                    remaining = remaining[len(g):]
                    matched = True
                    break
            if not matched:
                # Skip unrecognized character (separator or noise)
                remaining = remaining[1:]
        return result
    else:
        parts = re.split(r"[_+]", name)
        return [p.strip() for p in parts if p.strip()]


def _is_single_gene_perturbation(name: str) -> bool:
    """Return True if the perturbation name targets exactly one gene.

    Uses the same parsing logic as :func:`parse_perturbation_targets` so that
    training and evaluation agree on which names are multi-gene.  Names like
    ``SET_KLF1`` are parsed as ``["SET", "KLF1"]`` → multi-gene.
    """
    return len(parse_perturbation_targets(name)) == 1


def _all_target_genes(pert_names: List[str]) -> dict:
    """Map each gene to the set of condition indices that target it.

    Parameters
    ----------
    pert_names : list of perturbation name strings

    Returns
    -------
    dict : gene_str → set of condition indices
    """
    gene_to_conds: dict = {}
    for i, name in enumerate(pert_names):
        for g in parse_perturbation_targets(name):
            gene_to_conds.setdefault(g, set()).add(i)
    return gene_to_conds


def gene_disjoint_split(
    pert_names: List[str],
    val_fraction: float = 0.15,
    seed: int = 42,
) -> Tuple[List[int], List[int]]:
    """Split condition indices so that val perturbation genes never appear in train.

    Strategy
    --------
    1. Collect all genes that appear **exclusively** in single-gene conditions.
       (Genes that only appear in multi-gene conditions cannot be held out
       without also removing all combo conditions containing them.)
    2. Randomly select ``val_fraction`` of those exclusive single-gene symbols
       as the val-gene set.
    3. Every condition targeting *any* val-gene → validation set.
       All other conditions → training set.

    This guarantees zero gene overlap between train and val *for genes that
    appear in single-gene conditions*.  Genes that only appear in combos may
    appear in both splits (since removing all their combos would discard too
    much data), but this is a known and acceptable limitation.

    Parameters
    ----------
    pert_names   : list of perturbation names (one per condition), e.g.
                   ['KLF1', 'CEBPE', 'KLF1+CEBPE', 'SET_KLF1', ...]
    val_fraction : fraction of exclusive single-gene symbols to hold out
    seed         : random seed for gene selection

    Returns
    -------
    train_idx : list[int]  condition indices for training
    val_idx   : list[int]  condition indices for validation
    """
    # Build gene → condition mapping
    gene_to_conds = _all_target_genes(pert_names)

    # Genes that appear ONLY in single-gene conditions (eligible for hold-out)
    exclusive_single_genes = sorted({
        g for g, conds in gene_to_conds.items()
        if all(_is_single_gene_perturbation(pert_names[i]) for i in conds)
    })

    rng = np.random.default_rng(seed)
    n_val = max(1, int(len(exclusive_single_genes) * val_fraction))
    val_genes = set(rng.choice(exclusive_single_genes, n_val, replace=False).tolist())

    train_idx, val_idx = [], []
    for i, name in enumerate(pert_names):
        genes = set(parse_perturbation_targets(name))
        if any(g in val_genes for g in genes):
            val_idx.append(i)
        else:
            train_idx.append(i)

    return train_idx, val_idx


def select_hvg_by_variance(
    X: np.ndarray,
    n_genes: int = 2000,
) -> np.ndarray:
    """Select top-n genes by across-cell variance.

    Parameters
    ----------
    X       : [n_cells, n_genes]  float array (log-normalised)
    n_genes : number of HVG to keep

    Returns
    -------
    hvg_idx : [n_genes]  integer indices into gene axis, sorted by variance desc
    """
    var = X.var(axis=0)
    n_genes = min(n_genes, X.shape[1])
    idx = np.argsort(var)[::-1][:n_genes]
    return idx.copy()


def normalize_counts(
    X: np.ndarray,
    target_sum: float = 1e4,
    log1p: bool = True,
) -> np.ndarray:
    """Library-size normalise and optionally log1p-transform raw count matrix.

    Parameters
    ----------
    X          : [n_cells, n_genes]  raw count matrix (dense float32)
    target_sum : scale each cell to this total count
    log1p      : apply log1p after normalisation

    Returns
    -------
    X_norm : [n_cells, n_genes]  float32
    """
    row_sums = X.sum(axis=1, keepdims=True).clip(min=1.0)
    X_norm = (X / row_sums) * target_sum
    if log1p:
        X_norm = np.log1p(X_norm)
    return X_norm.astype(np.float32)