"""Reference CPU implementation. Coordinates are (row, col), native pixels."""
import itertools
import math
import numpy as np
from scipy import ndimage
from scipy.stats import beta, binom


def components(mask):
    mask = np.asarray(mask, dtype=bool)
    if mask.ndim != 2:
        raise ValueError('Expected a 2D mask')
    labels, n = ndimage.label(mask, structure=np.ones((3, 3), dtype=int))
    return [labels == i for i in range(1, n + 1)]


def centers(masks):
    return np.asarray([np.argwhere(m).mean(0) for m in masks], float).reshape(-1, 2)


def eligibility(pred_masks, gt_masks, radius=3.0):
    if radius < 0:
        raise ValueError('radius must be nonnegative')
    a, b = centers(pred_masks), centers(gt_masks)
    return np.linalg.norm(a[:, None] - b[None, :], axis=-1) <= radius


def matching(edges):
    """Deterministic maximum cardinality bipartite matching, not greedy matching."""
    edges = np.asarray(edges, bool)
    if edges.ndim != 2:
        raise ValueError('Expected two-dimensional eligibility')
    owner = [-1] * edges.shape[1]
    def augment(i, seen):
        for j in np.flatnonzero(edges[i]):
            if j in seen:
                continue
            seen.add(j)
            if owner[j] < 0 or augment(owner[j], seen):
                owner[j] = i
                return True
        return False
    for i in range(edges.shape[0]):
        augment(i, set())
    return sorted((i, j) for j, i in enumerate(owner) if i >= 0)


def candidate_pool(prob, base_threshold=0.1, cap=64):
    """Threshold ONCE; subsequent gate thresholds delete whole fixed components."""
    prob = np.asarray(prob, float)
    if prob.ndim != 2 or not np.isfinite(prob).all() or np.any((prob < 0) | (prob > 1)):
        raise ValueError('Expected finite 2D probabilities in [0,1]')
    if not 0 <= base_threshold <= 1 or not isinstance(cap, int) or cap < 1:
        raise ValueError('Invalid threshold/cap')
    masks = components(prob >= base_threshold)
    masks.sort(key=lambda m: -float(prob[m].max()))  # stable ties: raster component order
    masks = masks[:cap]
    features = []
    for m in masks:
        p = prob[m]
        ring = ndimage.binary_dilation(m, iterations=2) & ~m
        contrast = float(p.mean() - prob[ring].mean()) if ring.any() else 0.0
        features.append([float(p.max()), float(p.mean()), float(p.std()),
                         float(np.log1p(m.sum())), contrast])
    return masks, np.asarray(features, np.float32).reshape(-1, 5)


def canonical_labels(pred_masks, gt_mask, radius=3.0):
    gt = components(gt_mask)
    pairs = matching(eligibility(pred_masks, gt, radius))
    labels = np.zeros(len(pred_masks), np.float32)
    areas = np.ones(len(pred_masks), np.float32)
    for i, j in pairs:
        labels[i] = 1
        areas[i] = gt[j].sum()
    return labels, areas


def image_stats(pred_masks, gt_mask, radius=3.0):
    gt_mask = np.asarray(gt_mask, bool)
    gt = components(gt_mask)
    pairs = matching(eligibility(pred_masks, gt, radius))
    tp = len(pairs)
    union_pred = np.zeros_like(gt_mask)
    for m in pred_masks:
        if m.shape != gt_mask.shape:
            raise ValueError('Shape mismatch')
        union_pred |= m
    intersection = int((union_pred & gt_mask).sum())
    union = int((union_pred | gt_mask).sum())
    fp_pixels = int((union_pred & ~gt_mask).sum())
    matched = {i for i, _ in pairs}
    # Separate legacy-style unmatched-component area from literal background-pixel FP.
    unmatched_area = sum(int(m.sum()) for i, m in enumerate(pred_masks) if i not in matched)
    return dict(tp=tp, fp=len(pred_masks)-tp, fn=len(gt)-tp,
                intersection=intersection, union=union,
                pred_pixels=int(union_pred.sum()), gt_pixels=int(gt_mask.sum()),
                fp_pixels=fp_pixels, unmatched_area=unmatched_area,
                pixels=gt_mask.size, niou=(intersection/union if union else 1.0))


def aggregate(rows):
    if not rows:
        raise ValueError('Cannot evaluate an empty split')
    sums = {k: sum(r[k] for r in rows) for k in rows[0] if k != 'niou'}
    def ratio(a, b):
        return a/b if b else None
    tp, fp, fn = (sums[k] for k in ('tp', 'fp', 'fn'))
    return dict(images=len(rows), **sums,
                IoU=ratio(sums['intersection'], sums['union']),
                nIoU=float(np.mean([r['niou'] for r in rows])),
                Pd=ratio(tp, tp+fn), object_F1=ratio(2*tp, 2*tp+fp+fn),
                pixel_F1=ratio(2*sums['intersection'], sums['pred_pixels']+sums['gt_pixels']),
                FP_per_image=fp/len(rows), P_any_FA=sum(r['fp']>0 for r in rows)/len(rows),
                Fa_pixel=sums['fp_pixels']/sums['pixels'],
                Fa_unmatched_area=sums['unmatched_area']/sums['pixels'])


def cp_upper(successes, n, error):
    if n < 1 or not 0 < error < 1 or not 0 <= successes <= n:
        raise ValueError('Invalid binomial parameters')
    return 1.0 if successes == n else float(beta.ppf(1-error, successes+1, n-successes))


def calibrate(fp_counts, thresholds, cap, alpha=0.1, delta=0.05, risk='count', method='bonferroni'):
    """Finite-grid CP calibration. Rows MUST be iid held-out deployment units.

    Does not infer iid sampling from filenames. Empty fallback is explicit.
    """
    counts = np.asarray(fp_counts)
    ts = np.asarray(thresholds, float)
    if (counts.ndim != 2 or counts.shape[0] < 1 or counts.shape[1] != len(ts)
        or len(ts) < 1 or not np.isfinite(counts).all()
        or not np.isfinite(ts).all() or np.any((ts<0)|(ts>1))
        or len(np.unique(ts)) != len(ts) or not isinstance(cap, int) or cap < 1
        or np.any(counts < 0) or np.any(counts > cap) or np.any(counts != np.floor(counts))
        or not 0 < delta < 1 or not 0 <= alpha <= (cap if risk=='count' else 1)
        or risk not in ('count', 'any') or method not in ('bonferroni', 'fixed_sequence')):
        raise ValueError('Invalid calibration input')
    n, m = counts.shape
    if method == 'fixed_sequence':
        # Order fixed before inspecting labels; stop at the first failed test.
        order = np.argsort(-ts).tolist()
        tested = []
        chosen = None
        limit = cap if risk == 'count' else 1
        for j in order:
            values = counts[:, j] if risk == 'count' else (counts[:, j] > 0)
            mean = float(np.mean(values))/limit
            target = alpha/limit
            if target == 1:
                pvalue = 0.0  # Risk <= known support is deterministic.
            elif target == 0 or mean >= target:
                pvalue = 1.0
            elif risk == 'any':
                pvalue = float(binom.cdf(int(np.sum(values)), n, target))
            else:
                kl = ((mean*math.log(mean/target) if mean else 0.0)
                      + (1-mean)*math.log((1-mean)/(1-target)))
                pvalue = math.exp(-n*kl)
            tested.append(dict(threshold=float(ts[j]), pvalue=pvalue))
            if pvalue > delta:
                break
            chosen = j
        return dict(method=method, risk=risk, alpha=alpha, delta=delta, n=n, cap=cap,
                    thresholds=ts.tolist(), tested=tested,
                    threshold=None if chosen is None else float(ts[chosen]),
                    reject_all=chosen is None,
                    certified_risk_budget=0.0 if chosen is None else alpha,
                    status='reject_all_no_nontrivial_certificate' if chosen is None else 'certified_under_iid_assumptions')
    tails = cap if risk == 'count' else 1
    bounds = [sum(cp_upper(int((counts[:,j]>=k).sum()), n, delta/(m*tails))
                  for k in range(1, tails+1)) for j in range(m)]
    feasible = [j for j in range(m) if bounds[j] <= alpha]
    # Lowest threshold keeps largest nested set; no label-based utility search needed.
    chosen = min(feasible, key=lambda j: ts[j]) if feasible else None
    return dict(method=method, risk=risk, alpha=alpha, delta=delta, n=n, cap=cap,
                thresholds=ts.tolist(), upper_bounds=bounds,
                threshold=None if chosen is None else float(ts[chosen]),
                reject_all=chosen is None,
                selected_upper_bound=0.0 if chosen is None else bounds[chosen],
                status='reject_all_no_nontrivial_certificate' if chosen is None else 'certified_under_iid_assumptions')


def softplus_bound(scores, negative, tau=1.0):
    if tau <= 0:
        raise ValueError('tau must be positive')
    return float(np.logaddexp(0, np.asarray(scores)[np.asarray(negative, bool)]/tau).sum()/math.log(2))


def exact_mwis(weights, edges):
    """Tiny-instance oracle only; exponential, capped at 20 nodes."""
    w = np.asarray(weights, float)
    if len(w) > 20:
        raise ValueError('Oracle limited to 20 nodes')
    best, selection = 0.0, []
    for bits in itertools.product((0,1), repeat=len(w)):
        if any(bits[i] and bits[j] for i,j in edges):
            continue
        score = float(w @ bits)
        if score > best:
            best, selection = score, [i for i,b in enumerate(bits) if b]
    return best, selection


def best_path(node_scores, transition_costs):
    """Complete path through layered DAG. inf cost denotes forbidden edge."""
    if not node_scores or len(transition_costs) != len(node_scores)-1:
        raise ValueError('Inconsistent layers')
    dp = np.asarray(node_scores[0], float)
    if not len(dp):
        return -math.inf, []
    parents = []
    for scores, costs in zip(node_scores[1:], transition_costs):
        scores, costs = np.asarray(scores,float), np.asarray(costs,float)
        if costs.shape != (len(dp), len(scores)) or not len(scores):
            raise ValueError('Invalid transitions')
        val = dp[:,None]-costs
        par = val.argmax(0)
        dp = scores + val[par, np.arange(len(scores))]
        parents.append(par)
    k = int(dp.argmax())
    best = float(dp[k])
    if not np.isfinite(best):
        return -math.inf, []
    path = [k]
    for par in reversed(parents):
        k = int(par[k]); path.append(k)
    return best, path[::-1]