"""IntuitiveAI finite-mechanism labs v1. Python 3.10+, standard library only."""
import math


def dpo_loss(policy_gap, reference_gap, beta):
    if not all(math.isfinite(v) for v in (policy_gap, reference_gap, beta)) or beta <= 0:
        raise ValueError("Finite gaps and beta > 0 required")
    margin = beta * (policy_gap - reference_gap)
    return max(0, -margin) + math.log1p(math.exp(-abs(margin)))


def dpo_gradient(theta, beta):
    """Uniform-reference, two-response model only: p(A)=sigmoid(theta)."""
    if not math.isfinite(theta) or not math.isfinite(beta) or beta <= 0:
        raise ValueError("Finite theta and beta > 0 required")
    z = beta * theta
    probability_negative = math.exp(-z) / (1 + math.exp(-z)) if z >= 0 else 1 / (1 + math.exp(z))
    return -beta * probability_negative


def lora_forward(W, A, B, x, alpha):
    """Contract: compatible nonempty finite matrices; W out×in, A r×in, B out×r."""
    ax = [sum(a * v for a, v in zip(row, x)) for row in A]
    return [sum(w * v for w, v in zip(row, x)) + alpha / len(A) * sum(b * v for b, v in zip(B[i], ax)) for i, row in enumerate(W)]


def merge(W, A, B, alpha):
    """Return a fresh matrix; do not alter W or apply the adapter twice."""
    return [[w + alpha / len(A) * sum(B[i][k] * A[k][j] for k in range(len(A))) for j, w in enumerate(row)] for i, row in enumerate(W)]


def quantize(values, bits):
    """Signed symmetric narrow range, ties-to-even, all-zero scale=1. Not NF4."""
    if bits not in (4, 8) or not values or not all(math.isfinite(v) for v in values):
        raise ValueError("Finite nonempty values; bits must be 4 or 8")
    limit = 2 ** (bits - 1) - 1
    maximum = max(abs(v) for v in values)
    scale = maximum / limit if maximum else 1.0
    codes = [max(-limit, min(limit, round(v / scale))) for v in values]
    return codes, scale


def advantages(rewards):
    if len(rewards) < 2 or not all(math.isfinite(v) for v in rewards):
        raise ValueError("At least two finite rewards required")
    mean = sum(rewards) / len(rewards)
    std = math.sqrt(sum((v - mean) ** 2 for v in rewards) / len(rewards))
    return [(v - mean) / std if std else 0.0 for v in rewards]


def surrogate(ratios, advantage, mask, epsilon=0.2):
    """Mean over valid completion tokens for ONE response; maximize this value.

    Ratios compare active to frozen OLD policy. No KL term (beta=0).
    Masked ratios are ignored; an empty completion is an input error.
    """
    if len(ratios) != len(mask) or not all(v in (0, 1) for v in mask) or not any(mask):
        raise ValueError("Matching binary mask with at least one valid token required")
    if not math.isfinite(advantage) or not math.isfinite(epsilon) or not 0 <= epsilon < 1:
        raise ValueError("Finite advantage and 0 <= epsilon < 1 required")
    valid = [ratio for ratio, keep in zip(ratios, mask) if keep]
    if not all(math.isfinite(v) and v > 0 for v in valid):
        raise ValueError("Valid token ratios must be positive and finite")
    terms = [min(r * advantage, max(1 - epsilon, min(1 + epsilon, r)) * advantage) for r in valid]
    return sum(terms) / len(terms)