"""Identity and relatedness (spec QC-033, QC-064; plan §9.4).

* KING-robust kinship between every pair of samples
  [LIT: Manichaikul A. et al. 2010, Bioinformatics 26:2867, eq. for the within-family (robust) estimator]:
      φ_ij = [N(Aa,Aa) − 2·N(AA,aa)] / [2·min(N_Aa(i), N_Aa(j))] + 1/2 − [N_Aa(i) + N_Aa(j)] / [4·min(N_Aa(i), N_Aa(j))]
  with every count taken over the sites called in both samples.
* Duplicates: genotype concordance over jointly called sites >= 0.99 (plan §9.4) [OURS threshold].
* Mendelian errors of child/parent trios from the sample sheet (Parent1/Parent2).

Counts are formed with matrix products of per-genotype indicator matrices, so a cohort is O(n²·m) arithmetic in a
few BLAS calls. Relationship degrees use KING's published cut-offs (> 0.354 duplicate/MZ twin, > 0.177 first
degree, > 0.0884 second, > 0.0442 third).
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
import numpy.typing as npt

DEGREE_CUTS = ((0.354, "duplicate"), (0.177, "first"), (0.0884, "second"), (0.0442, "third"))
DUPLICATE_CONCORDANCE = 0.99


@dataclass
class PairStats:
    kinship: npt.NDArray[np.float64]  # (n, n); NaN on the diagonal or where undefined
    concordance: npt.NDArray[np.float64]  # (n, n)
    n_sites: npt.NDArray[np.int64]  # jointly called sites


def degree(phi: float) -> str:
    for cut, name in DEGREE_CUTS:
        if phi > cut:
            return name
    return "unrelated"


def select_sites(
    counts: npt.NDArray[np.int64],
    usable: npt.NDArray[np.bool_],
    *,
    min_call_rate: float = 0.95,
    min_maf: float = 0.01,
    max_sites: int = 100_000,
) -> npt.NDArray[np.intp]:
    """High-quality sites from per-locus (no-call, AA, AB, BB) ``counts``: ``usable`` (autosomal, not
    intensity-only), cohort call rate and MAF filters, thinned evenly to at most ``max_sites``."""
    called = counts[:, 1:].sum(axis=1)
    total = counts.sum(axis=1)
    with np.errstate(invalid="ignore", divide="ignore"):
        cr = np.where(total > 0, called / total, 0.0)
        af = np.where(called > 0, (counts[:, 2] + 2 * counts[:, 3]) / (2.0 * called), np.nan)
    maf = np.minimum(af, 1 - af)
    idx = np.flatnonzero(usable & (cr >= min_call_rate) & (maf >= min_maf))
    if len(idx) > max_sites:
        idx = idx[np.linspace(0, len(idx) - 1, max_sites).round().astype(np.intp)]
    return idx


def pair_stats(genotypes: npt.NDArray[np.uint8]) -> PairStats:
    """KING-robust kinship and concordance for every pair; ``genotypes`` is (sites × samples), 0 = no call."""
    g = np.asarray(genotypes)
    aa = (g == 1).astype(np.float32)
    ab = (g == 2).astype(np.float32)
    bb = (g == 3).astype(np.float32)
    c = (g > 0).astype(np.float32)
    # float32 products are exact while counts stay below 2**24 sites
    t = np.float64
    n_het_het = (ab.T @ ab).astype(t)
    n_opp = (aa.T @ bb).astype(t)
    n_opp = n_opp + n_opp.T
    het_i = (ab.T @ c).astype(t)  # [i, j] = hets of i over sites called in j
    het_j = het_i.T
    joint = (c.T @ c).astype(t)
    same = (aa.T @ aa + ab.T @ ab + bb.T @ bb).astype(t)
    mn = np.minimum(het_i, het_j)
    with np.errstate(invalid="ignore", divide="ignore"):
        phi = (n_het_het - 2 * n_opp) / (2 * mn) + 0.5 - (het_i + het_j) / (4 * mn)
        conc = same / joint
    phi[mn == 0] = np.nan
    np.fill_diagonal(phi, np.nan)
    np.fill_diagonal(conc, np.nan)
    return PairStats(kinship=phi, concordance=conc, n_sites=joint.astype(np.int64))


def mendel_errors(
    child: npt.NDArray[np.uint8], father: npt.NDArray[np.uint8] | None, mother: npt.NDArray[np.uint8] | None
) -> tuple[int, int]:
    """(errors, informative sites) for one child with one or two genotyped parents; autosomal sites only.

    A site is informative when the child and every given parent are called. Errors: the child carries an allele
    that no parent can transmit (duo: opposite homozygotes; trio: child genotype outside the parental combinations).
    """
    parents = [p for p in (father, mother) if p is not None]
    if not parents:
        return 0, 0
    ok = child > 0
    for p in parents:
        ok &= p > 0
    ch = child[ok].astype(np.int8) - 1  # B-allele count 0/1/2
    ps = [p[ok].astype(np.int8) - 1 for p in parents]
    if len(ps) == 1:
        par = ps[0]
        err = ((ch == 0) & (par == 2)) | ((ch == 2) & (par == 0))
    else:
        f, m = ps
        # B alleles a parent can transmit: hom-A → {0}, het → {0, 1}, hom-B → {1}
        lo = f // 2 + m // 2
        hi = (f + 1) // 2 + (m + 1) // 2
        err = (ch < lo) | (ch > hi)
    return int(err.sum()), int(ok.sum())