"""
Cost functions for all three QAOA problem modules.

All functions accept binary integer arguments (0 or 1) and return a scalar cost.
"""


def cost_standard(x1, x2, x3):
    """Unconstrained problem: C(x) = x1 + x2 - x2*x3 (3 qubits)."""
    return x1 + x2 - (x2 * x3)


def cost_equality(x1, x2, x3):
    """
    Equality-constrained problem (3 qubits).
    Minimize C(x) = x1 + x2 - x2*x3  subject to  x1 + x2 + x3 = 1.
    Penalty: A * (x1 + x2 + x3 - 1)^2,  A = 3.0
    """
    A = 3.0
    obj = x1 + x2 - (x2 * x3)
    penalty = A * (x1 + x2 + x3 - 1) ** 2
    return obj + penalty


def cost_inequality(x1, x2, x3, s):
    """
    Inequality-constrained problem via slack variable (4 qubits).
    Minimize C(x) = x1 + x2 - x2*x3  subject to  x1 + x2 + x3 >= 2.
    Slack variable formulation:  x1 + x2 + x3 - s = 2,  s in {0, 1}.
    Penalty: A * (x1 + x2 + x3 - s - 2)^2,  A = 4.0
    """
    A = 4.0
    obj = x1 + x2 - (x2 * x3)
    penalty = A * (x1 + x2 + x3 - s - 2) ** 2
    return obj + penalty