"""Scientific metrics, conservative screening, and sequence-diverse selection."""
from __future__ import annotations
import math
from collections import defaultdict
import numpy as np
from .schema import Measurement, Requirement


def p_affinity(kd_nm: float) -> float:
    """Return minus log10 of a dissociation constant in molar units."""
    if kd_nm <= 0:
        raise ValueError("Kd must be positive")
    return 9.0 - math.log10(kd_nm)


def selectivity(kd_target_nm: float, kd_counter_nm: float) -> float:
    """Log10 selectivity, positive for preferential binding to the target."""
    return p_affinity(kd_target_nm) - p_affinity(kd_counter_nm)


def normalized_margin(m: Measurement, r: Requirement, conservative: bool = True) -> float:
    if m.endpoint != r.endpoint or m.unit != r.unit or m.assay != r.assay:
        raise ValueError("endpoint, unit, or assay mismatch")
    value = m.value
    if conservative:
        bound = m.lower if r.direction == "ge" else m.upper
        if bound is None:
            raise ValueError("conservative ranking requires an interval")
        value = bound
    return (value - r.threshold) / r.scale if r.direction == "ge" else (r.threshold - value) / r.scale


def candidate_score(measurements: list[Measurement], requirements: list[Requirement]) -> float:
    """Worst margin over the supplied computational ranking requirements."""
    values = []
    for r in requirements:
        if not r.required:
            continue
        matched = [m for m in measurements if m.endpoint == r.endpoint]
        if len(matched) != 1:
            raise ValueError(f"expected one declared estimate for {r.endpoint}")
        values.append(normalized_margin(matched[0], r))
    if not values:
        raise ValueError("no required endpoints")
    return min(values)


def identity(a: str, b: str) -> float:
    """One minus normalized Levenshtein distance, including unequal lengths."""
    row = list(range(len(b) + 1))
    for i, ca in enumerate(a, 1):
        nxt = [i]
        for j, cb in enumerate(b, 1):
            nxt.append(min(nxt[-1] + 1, row[j] + 1, row[j-1] + (ca != cb)))
        row = nxt
    return 1.0 - row[-1] / max(len(a), len(b), 1)


def diverse_select(sequences: dict[str, str], scores: dict[str, float], k: int,
                   penalty: float = 0.25, max_identity: float = 0.8) -> list[str]:
    """Greedy quality/diversity selection with an explicit redundancy ceiling."""
    if k < 0 or penalty < 0 or not 0 <= max_identity <= 1:
        raise ValueError("invalid selection settings")
    chosen = []
    remaining = set(scores) & set(sequences)
    while remaining and len(chosen) < k:
        utilities = {}
        for cid in remaining:
            sim = max((identity(sequences[cid], sequences[c]) for c in chosen), default=0.)
            if sim <= max_identity and math.isfinite(scores[cid]):
                utilities[cid] = scores[cid] - penalty * sim
        if not utilities:
            break
        best = min(utilities, key=lambda x: (-utilities[x], x))
        chosen.append(best)
        remaining.remove(best)
    return chosen


def conformal_radius(y_true, y_pred, scales, alpha=0.1) -> float:
    """Finite-sample split-conformal radius for normalized absolute residuals.

    Calibration rows must be independent of fitting and final selection. Returns
    infinity when the requested quantile exceeds the finite calibration sample.
    """
    y, p, s = map(lambda x: np.asarray(x, dtype=float), (y_true, y_pred, scales))
    if y.shape != p.shape or y.shape != s.shape or y.ndim != 1 or len(y) == 0:
        raise ValueError("expected equal nonempty one-dimensional arrays")
    if not 0 < alpha < 1 or np.any(s <= 0) or not np.all(np.isfinite([y,p,s])):
        raise ValueError("invalid calibration values")
    errors = np.abs(y-p)/s
    rank = math.ceil((len(errors)+1)*(1-alpha))
    return float(np.sort(errors)[rank-1]) if rank <= len(errors) else math.inf


def grouped_bootstrap(rows: list[dict], value: str, group: str, seed=2027, n=2000) -> dict:
    """Equal-weight group mean and percentile interval across independent tasks."""
    groups = defaultdict(list)
    for r in rows:
        groups[r[group]].append(float(r[value]))
    x = np.array([np.mean(v) for v in groups.values()])
    if len(x) < 2:
        raise ValueError("at least two independent groups are required")
    rng = np.random.default_rng(seed)
    draws = x[rng.integers(len(x), size=(n,len(x)))].mean(axis=1)
    return {"mean": float(x.mean()), "lower": float(np.quantile(draws,.025)),
            "upper": float(np.quantile(draws,.975)), "groups": len(x)}


def joint_success(measurements: list[Measurement], requirements: list[Requirement]) -> bool:
    """Require every registered endpoint to pass on experimental observations."""
    for r in requirements:
        if not r.required:
            continue
        rows = [m for m in measurements if m.endpoint == r.endpoint and m.kind == "experiment"]
        if len(rows) != 1 or rows[0].censor != "none":
            return False
        if normalized_margin(rows[0],r,conservative=False) < 0:
            return False
    return True