"""Decision layer: calibration, conflict resolution, per-entity expected-F0.5 selection."""
import numpy as np
import pandas as pd
from sklearn.isotonic import IsotonicRegression


def fit_calibrator(p, y):
    iso = IsotonicRegression(out_of_bounds="clip", y_min=1e-4, y_max=1 - 1e-4)
    iso.fit(p, y)
    return iso


def resolve_conflicts(df, prob_col="p"):
    """Each S2/S3 record may belong to at most one S1 entity: keep only its best-scoring
    claim (applied before selection, so a record never counts for two entities).

    TIES MUST BE BROKEN. The previous version kept every row with `p >= max(p)` for a
    record, which silently kept ALL tied claimants. Ties are not rare here: isotonic
    calibration is piecewise-constant, so many distinct raw scores collapse onto the
    same calibrated probability. On the full test set that let S3-709463074 be claimed
    by two entities at once - caught by the one-owner assert in run_test.py.

    Sorting by (c, -p) with a stable kind and keeping the first row per record gives
    exactly one owner, deterministically, with the original row order breaking ties."""
    d = df.sort_values(["c", prob_col], ascending=[True, False], kind="mergesort")
    return d[~d.duplicated(subset="c", keep="first")]


def expected_f05_select(probs, miss_mass, floor=0.0, beta=0.5):
    """probs: calibrated probabilities for one entity's candidates (any order).
    Returns indices (into probs) of the chosen set that maximises expected F-beta."""
    b2 = beta * beta
    o = np.argsort(-probs)
    ps = probs[o]
    total = ps.sum() + miss_mass
    best_k = 0
    best = float(np.prod(1 - ps) * np.exp(-miss_mass))          # P(entity is a singleton)
    s = 0.0
    for k in range(1, len(ps) + 1):
        if ps[k - 1] < floor:
            break
        s += ps[k - 1]
        ef = (1 + b2) * s / ((1 + b2) * s + b2 * (total - s) + (k - s))
        if ef > best:
            best, best_k = ef, k
    return o[:best_k]


def predict_sets(df, id_col, cand_col, prob_col, method, param, miss_mass=0.0, conflicts=True):
    """method: 'threshold' (param = threshold) or 'ef05' (param = probability floor)."""
    d = resolve_conflicts(df, prob_col) if conflicts else df
    out = {}
    if method == "threshold":
        d = d[d[prob_col] >= param]
        for q, g in d.groupby(id_col):
            out[q] = list(g[cand_col].values)
        return out
    for q, g in d.groupby(id_col):
        idx = expected_f05_select(g[prob_col].values, miss_mass, floor=param)
        if len(idx):
            out[q] = list(g[cand_col].values[idx])
    return out


def tune(df, truths, ids, id_col, cand_col, prob_col, method, grid, miss_mass, conflicts, scorer):
    best = (-1, None)
    for v in grid:
        m, _ = scorer(predict_sets(df, id_col, cand_col, prob_col, method, v, miss_mass, conflicts),
                      truths, ids)
        if m["f05"] > best[0]:
            best = (m["f05"], v)
    return best[1], best[0]