"""Solomon answer semantics (design note, not distributed).

Yes/no, entity and multi-label candidates are Nouls: one probability P(yes). Gold yes only when the
document clearly establishes it; not stated and conflicting are no. Choice (single, ordered) is a
distribution over the listed options only; reserved-gold items have no Solomon target.

Works for both readouts: four_collapsed (four-state letter logits A/B/C/D, collapsed) and two_letter (A/B).
"""
import numpy as np

YES, NO, NOT_STATED, CONFLICTING = 0, 1, 2, 3


def softmax(x, t=1.0):
    z = np.asarray(x, np.float64) / t
    z = z - z.max()
    e = np.exp(z)
    return e / e.sum()


def noul_gold(gold4):
    """Four-state (or two-state) gold -> 1 for yes, 0 for no."""
    return int(int(gold4) == YES)


def noul_logit(letter_logits):
    """Binary log-odds z = log P(yes)/P(no) of a Noul branch: letter A against everything else."""
    logits = np.asarray(letter_logits, np.float64)
    if len(logits) not in (2, 4):
        raise ValueError(f'Noul branch must have 2 or 4 letters, got {len(logits)}')
    rest = logits[1:] - logits[1:].max()
    return float(logits[YES] - (logits[1:].max() + np.log(np.exp(rest).sum())))


def p_yes(letter_logits, t=1.0):
    """P(yes) from a Noul branch: letter A of a 2-letter (two_letter) or 4-state (four_collapsed) readout.

    Temperature applies to the COLLAPSED binary logit, not to the letters: a Noul is a binary unit whose
    'no' mass may be spread over several reserved letters, so p_yes(t) = sigmoid(z/t) with z = noul_logit.
    At t = 1 this is exactly softmax over the letters at A (the two forms only differ once t != 1, where the
    letterwise form would decay toward 1/len(letters) instead of toward 1/2). Solomon.qualification.p_yes and
    abstention_refit/readout.py fit and evaluate the collapsed form, so the serving path must match it.
    """
    logits = np.asarray(letter_logits, np.float64)
    if len(logits) not in (2, 4):
        raise ValueError(f'Noul branch must have 2 or 4 letters, got {len(logits)}')
    if t == 1.0:
        return float(softmax(logits)[YES])
    z = noul_logit(logits) / float(t)
    return float(1.0 / (1.0 + np.exp(-z))) if z > -700 else 0.0


def noul_confidence(p):
    return max(p, 1.0 - p)


def listed_gold(gold, n_options):
    """Listed option index, or None when the old gold was a reserved slot (not stated / none-of-listed / conflicting)."""
    return int(gold) if isinstance(gold, (int, np.integer)) and 0 <= int(gold) < n_options else None


def listed_probs(letter_logits, n_options, t=1.0):
    """Choice distribution over the listed options only (reserved slots, if present, are discarded)."""
    return softmax(np.asarray(letter_logits, np.float64)[:n_options], t)


def complement_deviation(p, p_negated):
    """G3a under Solomon: a statement and its negation should sum to 1."""
    return abs(p - (1.0 - p_negated))