ortools.sat.python.cp_model

   1# Copyright 2010-2022 Google LLC
   2# Licensed under the Apache License, Version 2.0 (the "License");
   3# you may not use this file except in compliance with the License.
   4# You may obtain a copy of the License at
   5#
   6#     http://www.apache.org/licenses/LICENSE-2.0
   7#
   8# Unless required by applicable law or agreed to in writing, software
   9# distributed under the License is distributed on an "AS IS" BASIS,
  10# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  11# See the License for the specific language governing permissions and
  12# limitations under the License.
  13
  14"""Methods for building and solving CP-SAT models.
  15
  16The following two sections describe the main
  17methods for building and solving CP-SAT models.
  18
  19* [`CpModel`](#cp_model.CpModel): Methods for creating
  20models, including variables and constraints.
  21* [`CPSolver`](#cp_model.CpSolver): Methods for solving
  22a model and evaluating solutions.
  23
  24The following methods implement callbacks that the
  25solver calls each time it finds a new solution.
  26
  27* [`CpSolverSolutionCallback`](#cp_model.CpSolverSolutionCallback):
  28  A general method for implementing callbacks.
  29* [`ObjectiveSolutionPrinter`](#cp_model.ObjectiveSolutionPrinter):
  30  Print objective values and elapsed time for intermediate solutions.
  31* [`VarArraySolutionPrinter`](#cp_model.VarArraySolutionPrinter):
  32  Print intermediate solutions (variable values, time).
  33* [`VarArrayAndObjectiveSolutionPrinter`]
  34      (#cp_model.VarArrayAndObjectiveSolutionPrinter):
  35  Print both intermediate solutions and objective values.
  36
  37Additional methods for solving CP-SAT models:
  38
  39* [`Constraint`](#cp_model.Constraint): A few utility methods for modifying
  40  constraints created by `CpModel`.
  41* [`LinearExpr`](#lineacp_model.LinearExpr): Methods for creating constraints
  42  and the objective from large arrays of coefficients.
  43
  44Other methods and functions listed are primarily used for developing OR-Tools,
  45rather than for solving specific optimization problems.
  46"""
  47
  48import collections
  49import itertools
  50import numbers
  51import threading
  52import time
  53from typing import (
  54    Callable,
  55    Dict,
  56    Iterable,
  57    List,
  58    Optional,
  59    Sequence,
  60    Tuple,
  61    Union,
  62    cast,
  63    overload,
  64)
  65import warnings
  66
  67import numpy as np
  68import pandas as pd
  69
  70from ortools.sat import cp_model_pb2
  71from ortools.sat import sat_parameters_pb2
  72from ortools.sat.python import cp_model_helper as cmh
  73from ortools.sat.python import swig_helper
  74from ortools.util.python import sorted_interval_list
  75
  76Domain = sorted_interval_list.Domain
  77
  78# The classes below allow linear expressions to be expressed naturally with the
  79# usual arithmetic operators + - * / and with constant numbers, which makes the
  80# python API very intuitive. See../ samples/*.py for examples.
  81
  82INT_MIN = -9223372036854775808  # hardcoded to be platform independent.
  83INT_MAX = 9223372036854775807
  84INT32_MAX = 2147483647
  85INT32_MIN = -2147483648
  86
  87# CpSolver status (exported to avoid importing cp_model_cp2).
  88UNKNOWN = cp_model_pb2.UNKNOWN
  89MODEL_INVALID = cp_model_pb2.MODEL_INVALID
  90FEASIBLE = cp_model_pb2.FEASIBLE
  91INFEASIBLE = cp_model_pb2.INFEASIBLE
  92OPTIMAL = cp_model_pb2.OPTIMAL
  93
  94# Variable selection strategy
  95CHOOSE_FIRST = cp_model_pb2.DecisionStrategyProto.CHOOSE_FIRST
  96CHOOSE_LOWEST_MIN = cp_model_pb2.DecisionStrategyProto.CHOOSE_LOWEST_MIN
  97CHOOSE_HIGHEST_MAX = cp_model_pb2.DecisionStrategyProto.CHOOSE_HIGHEST_MAX
  98CHOOSE_MIN_DOMAIN_SIZE = cp_model_pb2.DecisionStrategyProto.CHOOSE_MIN_DOMAIN_SIZE
  99CHOOSE_MAX_DOMAIN_SIZE = cp_model_pb2.DecisionStrategyProto.CHOOSE_MAX_DOMAIN_SIZE
 100
 101# Domain reduction strategy
 102SELECT_MIN_VALUE = cp_model_pb2.DecisionStrategyProto.SELECT_MIN_VALUE
 103SELECT_MAX_VALUE = cp_model_pb2.DecisionStrategyProto.SELECT_MAX_VALUE
 104SELECT_LOWER_HALF = cp_model_pb2.DecisionStrategyProto.SELECT_LOWER_HALF
 105SELECT_UPPER_HALF = cp_model_pb2.DecisionStrategyProto.SELECT_UPPER_HALF
 106
 107# Search branching
 108AUTOMATIC_SEARCH = sat_parameters_pb2.SatParameters.AUTOMATIC_SEARCH
 109FIXED_SEARCH = sat_parameters_pb2.SatParameters.FIXED_SEARCH
 110PORTFOLIO_SEARCH = sat_parameters_pb2.SatParameters.PORTFOLIO_SEARCH
 111LP_SEARCH = sat_parameters_pb2.SatParameters.LP_SEARCH
 112PSEUDO_COST_SEARCH = sat_parameters_pb2.SatParameters.PSEUDO_COST_SEARCH
 113PORTFOLIO_WITH_QUICK_RESTART_SEARCH = (
 114    sat_parameters_pb2.SatParameters.PORTFOLIO_WITH_QUICK_RESTART_SEARCH
 115)
 116HINT_SEARCH = sat_parameters_pb2.SatParameters.HINT_SEARCH
 117PARTIAL_FIXED_SEARCH = sat_parameters_pb2.SatParameters.PARTIAL_FIXED_SEARCH
 118RANDOMIZED_SEARCH = sat_parameters_pb2.SatParameters.RANDOMIZED_SEARCH
 119
 120# Type aliases
 121IntegralT = Union[numbers.Integral, np.integer, int]
 122NumberT = Union[numbers.Integral, np.integer, int, numbers.Number, np.double, float]
 123LiteralT = Union["IntVar", "_NotBooleanVariable", IntegralT, bool]
 124VariableT = Union["IntVar", IntegralT]
 125LinearExprT = Union["LinearExpr", "IntVar", IntegralT]
 126ObjLinearExprT = Union["LinearExpr", "IntVar", NumberT]
 127ArcT = Tuple[IntegralT, IntegralT, LiteralT]
 128_IndexOrSeries = Union[pd.Index, pd.Series]
 129
 130
 131def DisplayBounds(bounds: Sequence[int]) -> str:
 132    """Displays a flattened list of intervals."""
 133    out = ""
 134    for i in range(0, len(bounds), 2):
 135        if i != 0:
 136            out += ", "
 137        if bounds[i] == bounds[i + 1]:
 138            out += str(bounds[i])
 139        else:
 140            out += str(bounds[i]) + ".." + str(bounds[i + 1])
 141    return out
 142
 143
 144def ShortName(model: cp_model_pb2.CpModelProto, i: int) -> str:
 145    """Returns a short name of an integer variable, or its negation."""
 146    if i < 0:
 147        return "Not(%s)" % ShortName(model, -i - 1)
 148    v = model.variables[i]
 149    if v.name:
 150        return v.name
 151    elif len(v.domain) == 2 and v.domain[0] == v.domain[1]:
 152        return str(v.domain[0])
 153    else:
 154        return "[%s]" % DisplayBounds(v.domain)
 155
 156
 157def ShortExprName(
 158    model: cp_model_pb2.CpModelProto, e: cp_model_pb2.LinearExpressionProto
 159) -> str:
 160    """Pretty-print LinearExpressionProto instances."""
 161    if not e.vars:
 162        return str(e.offset)
 163    if len(e.vars) == 1:
 164        var_name = ShortName(model, e.vars[0])
 165        coeff = e.coeffs[0]
 166        result = ""
 167        if coeff == 1:
 168            result = var_name
 169        elif coeff == -1:
 170            result = f"-{var_name}"
 171        elif coeff != 0:
 172            result = f"{coeff} * {var_name}"
 173        if e.offset > 0:
 174            result = f"{result} + {e.offset}"
 175        elif e.offset < 0:
 176            result = f"{result} - {-e.offset}"
 177        return result
 178    # TODO(user): Support more than affine expressions.
 179    return str(e)
 180
 181
 182class LinearExpr:
 183    """Holds an integer linear expression.
 184
 185    A linear expression is built from integer constants and variables.
 186    For example, `x + 2 * (y - z + 1)`.
 187
 188    Linear expressions are used in CP-SAT models in constraints and in the
 189    objective:
 190
 191    * You can define linear constraints as in:
 192
 193    ```
 194    model.Add(x + 2 * y <= 5)
 195    model.Add(sum(array_of_vars) == 5)
 196    ```
 197
 198    * In CP-SAT, the objective is a linear expression:
 199
 200    ```
 201    model.Minimize(x + 2 * y + z)
 202    ```
 203
 204    * For large arrays, using the LinearExpr class is faster that using the python
 205    `sum()` function. You can create constraints and the objective from lists of
 206    linear expressions or coefficients as follows:
 207
 208    ```
 209    model.Minimize(cp_model.LinearExpr.Sum(expressions))
 210    model.Add(cp_model.LinearExpr.WeightedSum(expressions, coefficients) >= 0)
 211    ```
 212    """
 213
 214    @classmethod
 215    def Sum(cls, expressions: Sequence[LinearExprT]) -> LinearExprT:
 216        """Creates the expression sum(expressions)."""
 217        if len(expressions) == 1:
 218            return expressions[0]
 219        return _SumArray(expressions)
 220
 221    @overload
 222    @classmethod
 223    def WeightedSum(
 224        cls,
 225        expressions: Sequence[LinearExprT],
 226        coefficients: Sequence[IntegralT],
 227    ) -> LinearExprT:
 228        ...
 229
 230    @overload
 231    @classmethod
 232    def WeightedSum(
 233        cls,
 234        expressions: Sequence[ObjLinearExprT],
 235        coefficients: Sequence[NumberT],
 236    ) -> ObjLinearExprT:
 237        ...
 238
 239    @classmethod
 240    def WeightedSum(cls, expressions, coefficients):
 241        """Creates the expression sum(expressions[i] * coefficients[i])."""
 242        if LinearExpr.IsEmptyOrAllNull(coefficients):
 243            return 0
 244        elif len(expressions) == 1:
 245            return expressions[0] * coefficients[0]
 246        else:
 247            return _WeightedSum(expressions, coefficients)
 248
 249    @overload
 250    @classmethod
 251    def Term(
 252        cls,
 253        expressions: LinearExprT,
 254        coefficients: IntegralT,
 255    ) -> LinearExprT:
 256        ...
 257
 258    @overload
 259    @classmethod
 260    def Term(
 261        cls,
 262        expressions: ObjLinearExprT,
 263        coefficients: NumberT,
 264    ) -> ObjLinearExprT:
 265        ...
 266
 267    @classmethod
 268    def Term(cls, expression, coefficient):
 269        """Creates `expression * coefficient`."""
 270        if cmh.is_zero(coefficient):
 271            return 0
 272        else:
 273            return expression * coefficient
 274
 275    @classmethod
 276    def IsEmptyOrAllNull(cls, coefficients: Sequence[NumberT]) -> bool:
 277        for c in coefficients:
 278            if not cmh.is_zero(c):
 279                return False
 280        return True
 281
 282    @classmethod
 283    def RebuildFromLinearExpressionProto(
 284        cls,
 285        model: cp_model_pb2.CpModelProto,
 286        proto: cp_model_pb2.LinearExpressionProto,
 287    ) -> LinearExprT:
 288        """Recreate a LinearExpr from a LinearExpressionProto."""
 289        offset = proto.offset
 290        num_elements = len(proto.vars)
 291        if num_elements == 0:
 292            return offset
 293        elif num_elements == 1:
 294            return IntVar(model, proto.vars[0], None) * proto.coeffs[0] + offset
 295        else:
 296            variables = []
 297            coeffs = []
 298            all_ones = True
 299            for index, coeff in zip(proto.vars, proto.coeffs):
 300                variables.append(IntVar(model, index, None))
 301                coeffs.append(coeff)
 302                if not cmh.is_one(coeff):
 303                    all_ones = False
 304            if all_ones:
 305                return _SumArray(variables, offset)
 306            else:
 307                return _WeightedSum(variables, coeffs, offset)
 308
 309    def GetIntegerVarValueMap(self) -> Tuple[Dict[VariableT, IntegralT], int]:
 310        """Scans the expression, and returns (var_coef_map, constant)."""
 311        coeffs = collections.defaultdict(int)
 312        constant = 0
 313        to_process: List[Tuple[LinearExprT, IntegralT]] = [(self, 1)]
 314        while to_process:  # Flatten to avoid recursion.
 315            expr, coeff = to_process.pop()
 316            if cmh.is_integral(expr):
 317                constant += coeff * int(expr)
 318            elif isinstance(expr, _ProductCst):
 319                to_process.append((expr.Expression(), coeff * expr.Coefficient()))
 320            elif isinstance(expr, _Sum):
 321                to_process.append((expr.Left(), coeff))
 322                to_process.append((expr.Right(), coeff))
 323            elif isinstance(expr, _SumArray):
 324                for e in expr.Expressions():
 325                    to_process.append((e, coeff))
 326                constant += expr.Constant() * coeff
 327            elif isinstance(expr, _WeightedSum):
 328                for e, c in zip(expr.Expressions(), expr.Coefficients()):
 329                    to_process.append((e, coeff * c))
 330                constant += expr.Constant() * coeff
 331            elif isinstance(expr, IntVar):
 332                coeffs[expr] += coeff
 333            elif isinstance(expr, _NotBooleanVariable):
 334                constant += coeff
 335                coeffs[expr.Not()] -= coeff
 336            else:
 337                raise TypeError("Unrecognized linear expression: " + str(expr))
 338
 339        return coeffs, constant
 340
 341    def GetFloatVarValueMap(self) -> Tuple[Dict[VariableT, float], float, bool]:
 342        """Scans the expression. Returns (var_coef_map, constant, is_integer)."""
 343        coeffs = {}
 344        constant = 0
 345        to_process: List[Tuple[LinearExprT, Union[IntegralT, float]]] = [(self, 1)]
 346        while to_process:  # Flatten to avoid recursion.
 347            expr, coeff = to_process.pop()
 348            if cmh.is_integral(expr):  # Keep integrality.
 349                constant += coeff * int(expr)
 350            elif cmh.is_a_number(expr):
 351                constant += coeff * float(expr)
 352            elif isinstance(expr, _ProductCst):
 353                to_process.append((expr.Expression(), coeff * expr.Coefficient()))
 354            elif isinstance(expr, _Sum):
 355                to_process.append((expr.Left(), coeff))
 356                to_process.append((expr.Right(), coeff))
 357            elif isinstance(expr, _SumArray):
 358                for e in expr.Expressions():
 359                    to_process.append((e, coeff))
 360                constant += expr.Constant() * coeff
 361            elif isinstance(expr, _WeightedSum):
 362                for e, c in zip(expr.Expressions(), expr.Coefficients()):
 363                    to_process.append((e, coeff * c))
 364                constant += expr.Constant() * coeff
 365            elif isinstance(expr, IntVar):
 366                if expr in coeffs:
 367                    coeffs[expr] += coeff
 368                else:
 369                    coeffs[expr] = coeff
 370            elif isinstance(expr, _NotBooleanVariable):
 371                constant += coeff
 372                if expr.Not() in coeffs:
 373                    coeffs[expr.Not()] -= coeff
 374                else:
 375                    coeffs[expr.Not()] = -coeff
 376            else:
 377                raise TypeError("Unrecognized linear expression: " + str(expr))
 378        is_integer = cmh.is_integral(constant)
 379        if is_integer:
 380            for coeff in coeffs.values():
 381                if not cmh.is_integral(coeff):
 382                    is_integer = False
 383                    break
 384        return coeffs, constant, is_integer
 385
 386    def __hash__(self):
 387        return object.__hash__(self)
 388
 389    def __abs__(self):
 390        raise NotImplementedError(
 391            "calling abs() on a linear expression is not supported, "
 392            "please use CpModel.AddAbsEquality"
 393        )
 394
 395    def __add__(self, arg):
 396        if cmh.is_zero(arg):
 397            return self
 398        return _Sum(self, arg)
 399
 400    def __radd__(self, arg):
 401        if cmh.is_zero(arg):
 402            return self
 403        return _Sum(self, arg)
 404
 405    def __sub__(self, arg):
 406        if cmh.is_zero(arg):
 407            return self
 408        return _Sum(self, -arg)
 409
 410    def __rsub__(self, arg):
 411        return _Sum(-self, arg)
 412
 413    def __mul__(self, arg):
 414        arg = cmh.assert_is_a_number(arg)
 415        if cmh.is_one(arg):
 416            return self
 417        elif cmh.is_zero(arg):
 418            return 0
 419        return _ProductCst(self, arg)
 420
 421    def __rmul__(self, arg):
 422        arg = cmh.assert_is_a_number(arg)
 423        if cmh.is_one(arg):
 424            return self
 425        elif cmh.is_zero(arg):
 426            return 0
 427        return _ProductCst(self, arg)
 428
 429    def __div__(self, _):
 430        raise NotImplementedError(
 431            "calling / on a linear expression is not supported, "
 432            "please use CpModel.AddDivisionEquality"
 433        )
 434
 435    def __truediv__(self, _):
 436        raise NotImplementedError(
 437            "calling // on a linear expression is not supported, "
 438            "please use CpModel.AddDivisionEquality"
 439        )
 440
 441    def __mod__(self, _):
 442        raise NotImplementedError(
 443            "calling %% on a linear expression is not supported, "
 444            "please use CpModel.AddModuloEquality"
 445        )
 446
 447    def __pow__(self, _):
 448        raise NotImplementedError(
 449            "calling ** on a linear expression is not supported, "
 450            "please use CpModel.AddMultiplicationEquality"
 451        )
 452
 453    def __lshift__(self, _):
 454        raise NotImplementedError(
 455            "calling left shift on a linear expression is not supported"
 456        )
 457
 458    def __rshift__(self, _):
 459        raise NotImplementedError(
 460            "calling right shift on a linear expression is not supported"
 461        )
 462
 463    def __and__(self, _):
 464        raise NotImplementedError(
 465            "calling and on a linear expression is not supported, "
 466            "please use CpModel.AddBoolAnd"
 467        )
 468
 469    def __or__(self, _):
 470        raise NotImplementedError(
 471            "calling or on a linear expression is not supported, "
 472            "please use CpModel.AddBoolOr"
 473        )
 474
 475    def __xor__(self, _):
 476        raise NotImplementedError(
 477            "calling xor on a linear expression is not supported, "
 478            "please use CpModel.AddBoolXor"
 479        )
 480
 481    def __neg__(self):
 482        return _ProductCst(self, -1)
 483
 484    def __bool__(self):
 485        raise NotImplementedError(
 486            "Evaluating a LinearExpr instance as a Boolean is not implemented."
 487        )
 488
 489    def __eq__(self, arg):
 490        if arg is None:
 491            return False
 492        if cmh.is_integral(arg):
 493            arg = cmh.assert_is_int64(arg)
 494            return BoundedLinearExpression(self, [arg, arg])
 495        else:
 496            return BoundedLinearExpression(self - arg, [0, 0])
 497
 498    def __ge__(self, arg):
 499        if cmh.is_integral(arg):
 500            arg = cmh.assert_is_int64(arg)
 501            return BoundedLinearExpression(self, [arg, INT_MAX])
 502        else:
 503            return BoundedLinearExpression(self - arg, [0, INT_MAX])
 504
 505    def __le__(self, arg):
 506        if cmh.is_integral(arg):
 507            arg = cmh.assert_is_int64(arg)
 508            return BoundedLinearExpression(self, [INT_MIN, arg])
 509        else:
 510            return BoundedLinearExpression(self - arg, [INT_MIN, 0])
 511
 512    def __lt__(self, arg):
 513        if cmh.is_integral(arg):
 514            arg = cmh.assert_is_int64(arg)
 515            if arg == INT_MIN:
 516                raise ArithmeticError("< INT_MIN is not supported")
 517            return BoundedLinearExpression(self, [INT_MIN, arg - 1])
 518        else:
 519            return BoundedLinearExpression(self - arg, [INT_MIN, -1])
 520
 521    def __gt__(self, arg):
 522        if cmh.is_integral(arg):
 523            arg = cmh.assert_is_int64(arg)
 524            if arg == INT_MAX:
 525                raise ArithmeticError("> INT_MAX is not supported")
 526            return BoundedLinearExpression(self, [arg + 1, INT_MAX])
 527        else:
 528            return BoundedLinearExpression(self - arg, [1, INT_MAX])
 529
 530    def __ne__(self, arg):
 531        if arg is None:
 532            return True
 533        if cmh.is_integral(arg):
 534            arg = cmh.assert_is_int64(arg)
 535            if arg == INT_MAX:
 536                return BoundedLinearExpression(self, [INT_MIN, INT_MAX - 1])
 537            elif arg == INT_MIN:
 538                return BoundedLinearExpression(self, [INT_MIN + 1, INT_MAX])
 539            else:
 540                return BoundedLinearExpression(
 541                    self, [INT_MIN, arg - 1, arg + 1, INT_MAX]
 542                )
 543        else:
 544            return BoundedLinearExpression(self - arg, [INT_MIN, -1, 1, INT_MAX])
 545
 546
 547class _Sum(LinearExpr):
 548    """Represents the sum of two LinearExprs."""
 549
 550    def __init__(self, left, right):
 551        for x in [left, right]:
 552            if not cmh.is_a_number(x) and not isinstance(x, LinearExpr):
 553                raise TypeError("Not an linear expression: " + str(x))
 554        self.__left = left
 555        self.__right = right
 556
 557    def Left(self):
 558        return self.__left
 559
 560    def Right(self):
 561        return self.__right
 562
 563    def __str__(self):
 564        return f"({self.__left} + {self.__right})"
 565
 566    def __repr__(self):
 567        return f"Sum({repr(self.__left)}, {repr(self.__right)})"
 568
 569
 570class _ProductCst(LinearExpr):
 571    """Represents the product of a LinearExpr by a constant."""
 572
 573    def __init__(self, expr, coeff):
 574        coeff = cmh.assert_is_a_number(coeff)
 575        if isinstance(expr, _ProductCst):
 576            self.__expr = expr.Expression()
 577            self.__coef = expr.Coefficient() * coeff
 578        else:
 579            self.__expr = expr
 580            self.__coef = coeff
 581
 582    def __str__(self):
 583        if self.__coef == -1:
 584            return "-" + str(self.__expr)
 585        else:
 586            return "(" + str(self.__coef) + " * " + str(self.__expr) + ")"
 587
 588    def __repr__(self):
 589        return "ProductCst(" + repr(self.__expr) + ", " + repr(self.__coef) + ")"
 590
 591    def Coefficient(self):
 592        return self.__coef
 593
 594    def Expression(self):
 595        return self.__expr
 596
 597
 598class _SumArray(LinearExpr):
 599    """Represents the sum of a list of LinearExpr and a constant."""
 600
 601    def __init__(self, expressions, constant=0):
 602        self.__expressions = []
 603        self.__constant = constant
 604        for x in expressions:
 605            if cmh.is_a_number(x):
 606                if cmh.is_zero(x):
 607                    continue
 608                x = cmh.assert_is_a_number(x)
 609                self.__constant += x
 610            elif isinstance(x, LinearExpr):
 611                self.__expressions.append(x)
 612            else:
 613                raise TypeError("Not an linear expression: " + str(x))
 614
 615    def __str__(self):
 616        constant_terms = (self.__constant,) if self.__constant != 0 else ()
 617        exprs_str = " + ".join(
 618            map(repr, itertools.chain(self.__expressions, constant_terms))
 619        )
 620        if not exprs_str:
 621            return "0"
 622        return f"({exprs_str})"
 623
 624    def __repr__(self):
 625        exprs_str = ", ".join(map(repr, self.__expressions))
 626        return f"SumArray({exprs_str}, {self.__constant})"
 627
 628    def Expressions(self):
 629        return self.__expressions
 630
 631    def Constant(self):
 632        return self.__constant
 633
 634
 635class _WeightedSum(LinearExpr):
 636    """Represents sum(ai * xi) + b."""
 637
 638    def __init__(self, expressions, coefficients, constant=0):
 639        self.__expressions = []
 640        self.__coefficients = []
 641        self.__constant = constant
 642        if len(expressions) != len(coefficients):
 643            raise TypeError(
 644                "In the LinearExpr.WeightedSum method, the expression array and the "
 645                " coefficient array must have the same length."
 646            )
 647        for e, c in zip(expressions, coefficients):
 648            c = cmh.assert_is_a_number(c)
 649            if cmh.is_zero(c):
 650                continue
 651            if cmh.is_a_number(e):
 652                e = cmh.assert_is_a_number(e)
 653                self.__constant += e * c
 654            elif isinstance(e, LinearExpr):
 655                self.__expressions.append(e)
 656                self.__coefficients.append(c)
 657            else:
 658                raise TypeError("Not an linear expression: " + str(e))
 659
 660    def __str__(self):
 661        output = None
 662        for expr, coeff in zip(self.__expressions, self.__coefficients):
 663            if not output and cmh.is_one(coeff):
 664                output = str(expr)
 665            elif not output and cmh.is_minus_one(coeff):
 666                output = "-" + str(expr)
 667            elif not output:
 668                output = f"{coeff} * {expr}"
 669            elif cmh.is_one(coeff):
 670                output += f" + {expr}"
 671            elif cmh.is_minus_one(coeff):
 672                output += f" - {expr}"
 673            elif coeff > 1:
 674                output += f" + {coeff} * {expr}"
 675            elif coeff < -1:
 676                output += f" - {-coeff} * {expr}"
 677        if self.__constant > 0:
 678            output += f" + {self.__constant}"
 679        elif self.__constant < 0:
 680            output += f" - {-self.__constant}"
 681        if output is None:
 682            output = "0"
 683        return output
 684
 685    def __repr__(self):
 686        exprs_str = ", ".join(map(repr, self.__expressions))
 687        coeffs_str = ", ".join(map(repr, self.__coefficients))
 688        return f"WeightedSum([{exprs_str}], [{coeffs_str}], {self.__constant})"
 689
 690    def Expressions(self):
 691        return self.__expressions
 692
 693    def Coefficients(self):
 694        return self.__coefficients
 695
 696    def Constant(self):
 697        return self.__constant
 698
 699
 700class IntVar(LinearExpr):
 701    """An integer variable.
 702
 703    An IntVar is an object that can take on any integer value within defined
 704    ranges. Variables appear in constraint like:
 705
 706        x + y >= 5
 707        AllDifferent([x, y, z])
 708
 709    Solving a model is equivalent to finding, for each variable, a single value
 710    from the set of initial values (called the initial domain), such that the
 711    model is feasible, or optimal if you provided an objective function.
 712    """
 713
 714    def __init__(
 715        self,
 716        model: cp_model_pb2.CpModelProto,
 717        domain: Union[int, Domain],
 718        name: Optional[str],
 719    ):
 720        """See CpModel.NewIntVar below."""
 721        self.__negation: Optional[_NotBooleanVariable] = None
 722        # Python do not support multiple __init__ methods.
 723        # This method is only called from the CpModel class.
 724        # We hack the parameter to support the two cases:
 725        # case 1:
 726        #     model is a CpModelProto, domain is a Domain, and name is a string.
 727        # case 2:
 728        #     model is a CpModelProto, domain is an index (int), and name is None.
 729        if cmh.is_integral(domain) and name is None:
 730            self.__index: int = int(domain)
 731            self.__var: cp_model_pb2.IntegerVariableProto = model.variables[domain]
 732        else:
 733            self.__index: int = len(model.variables)
 734            self.__var: cp_model_pb2.IntegerVariableProto = model.variables.add()
 735            self.__var.domain.extend(cast(Domain, domain).FlattenedIntervals())
 736            self.__var.name = name
 737
 738    def Index(self) -> int:
 739        """Returns the index of the variable in the model."""
 740        return self.__index
 741
 742    def Proto(self) -> cp_model_pb2.IntegerVariableProto:
 743        """Returns the variable protobuf."""
 744        return self.__var
 745
 746    def IsEqualTo(self, other: ...) -> bool:
 747        """Returns true if self == other in the python sense."""
 748        if not isinstance(other, IntVar):
 749            return False
 750        return self.Index() == other.Index()
 751
 752    def __str__(self) -> str:
 753        if not self.__var.name:
 754            if (
 755                len(self.__var.domain) == 2
 756                and self.__var.domain[0] == self.__var.domain[1]
 757            ):
 758                # Special case for constants.
 759                return str(self.__var.domain[0])
 760            else:
 761                return "unnamed_var_%i" % self.__index
 762        return self.__var.name
 763
 764    def __repr__(self) -> str:
 765        return "%s(%s)" % (self.__var.name, DisplayBounds(self.__var.domain))
 766
 767    def Name(self) -> str:
 768        if not self.__var or not self.__var.name:
 769            return ""
 770        return self.__var.name
 771
 772    def Not(self) -> "_NotBooleanVariable":
 773        """Returns the negation of a Boolean variable.
 774
 775        This method implements the logical negation of a Boolean variable.
 776        It is only valid if the variable has a Boolean domain (0 or 1).
 777
 778        Note that this method is nilpotent: `x.Not().Not() == x`.
 779        """
 780
 781        for bound in self.__var.domain:
 782            if bound < 0 or bound > 1:
 783                raise TypeError("Cannot call Not on a non boolean variable: %s" % self)
 784        if self.__negation is None:
 785            self.__negation = _NotBooleanVariable(self)
 786        return self.__negation
 787
 788
 789class _NotBooleanVariable(LinearExpr):
 790    """Negation of a boolean variable."""
 791
 792    def __init__(self, boolvar: IntVar):
 793        self.__boolvar: IntVar = boolvar
 794
 795    def Index(self) -> int:
 796        return -self.__boolvar.Index() - 1
 797
 798    def Not(self) -> IntVar:
 799        return self.__boolvar
 800
 801    def __str__(self) -> str:
 802        return "not(%s)" % str(self.__boolvar)
 803
 804    def __bool__(self) -> bool:
 805        raise NotImplementedError(
 806            "Evaluating a literal as a Boolean value is not implemented."
 807        )
 808
 809
 810class BoundedLinearExpression:
 811    """Represents a linear constraint: `lb <= linear expression <= ub`.
 812
 813    The only use of this class is to be added to the CpModel through
 814    `CpModel.Add(expression)`, as in:
 815
 816        model.Add(x + 2 * y -1 >= z)
 817    """
 818
 819    def __init__(self, expr: LinearExprT, bounds: Sequence[int]):
 820        self.__expr: LinearExprT = expr
 821        self.__bounds: Sequence[int] = bounds
 822
 823    def __str__(self):
 824        if len(self.__bounds) == 2:
 825            lb, ub = self.__bounds
 826            if lb > INT_MIN and ub < INT_MAX:
 827                if lb == ub:
 828                    return str(self.__expr) + " == " + str(lb)
 829                else:
 830                    return str(lb) + " <= " + str(self.__expr) + " <= " + str(ub)
 831            elif lb > INT_MIN:
 832                return str(self.__expr) + " >= " + str(lb)
 833            elif ub < INT_MAX:
 834                return str(self.__expr) + " <= " + str(ub)
 835            else:
 836                return "True (unbounded expr " + str(self.__expr) + ")"
 837        elif (
 838            len(self.__bounds) == 4
 839            and self.__bounds[0] == INT_MIN
 840            and self.__bounds[1] + 2 == self.__bounds[2]
 841            and self.__bounds[3] == INT_MAX
 842        ):
 843            return str(self.__expr) + " != " + str(self.__bounds[1] + 1)
 844        else:
 845            return str(self.__expr) + " in [" + DisplayBounds(self.__bounds) + "]"
 846
 847    def Expression(self) -> LinearExprT:
 848        return self.__expr
 849
 850    def Bounds(self) -> Sequence[int]:
 851        return self.__bounds
 852
 853    def __bool__(self) -> bool:
 854        if isinstance(self.__expr, LinearExpr):
 855            coeffs_map, constant = self.__expr.GetIntegerVarValueMap()
 856            all_coeffs = set(coeffs_map.values())
 857            same_var = set([0])
 858            eq_bounds = [0, 0]
 859            different_vars = set([-1, 1])
 860            ne_bounds = [INT_MIN, -1, 1, INT_MAX]
 861            if (
 862                len(coeffs_map) == 1
 863                and all_coeffs == same_var
 864                and constant == 0
 865                and (self.__bounds == eq_bounds or self.__bounds == ne_bounds)
 866            ):
 867                return self.__bounds == eq_bounds
 868            if (
 869                len(coeffs_map) == 2
 870                and all_coeffs == different_vars
 871                and constant == 0
 872                and (self.__bounds == eq_bounds or self.__bounds == ne_bounds)
 873            ):
 874                return self.__bounds == ne_bounds
 875
 876        raise NotImplementedError(
 877            f'Evaluating a BoundedLinearExpression "{self}" as a Boolean value'
 878            + " is not supported."
 879        )
 880
 881
 882class Constraint:
 883    """Base class for constraints.
 884
 885    Constraints are built by the CpModel through the Add<XXX> methods.
 886    Once created by the CpModel class, they are automatically added to the model.
 887    The purpose of this class is to allow specification of enforcement literals
 888    for this constraint.
 889
 890        b = model.NewBoolVar('b')
 891        x = model.NewIntVar(0, 10, 'x')
 892        y = model.NewIntVar(0, 10, 'y')
 893
 894        model.Add(x + 2 * y == 5).OnlyEnforceIf(b.Not())
 895    """
 896
 897    def __init__(
 898        self,
 899        cp_model: "CpModel",
 900    ):
 901        self.__index: int = len(cp_model.Proto().constraints)
 902        self.__cp_model: "CpModel" = cp_model
 903        self.__constraint: cp_model_pb2.ConstraintProto = (
 904            cp_model.Proto().constraints.add()
 905        )
 906
 907    @overload
 908    def OnlyEnforceIf(self, boolvar: Iterable[LiteralT]) -> "Constraint":
 909        ...
 910
 911    @overload
 912    def OnlyEnforceIf(self, *boolvar: LiteralT) -> "Constraint":
 913        ...
 914
 915    def OnlyEnforceIf(self, *boolvar) -> "Constraint":
 916        """Adds an enforcement literal to the constraint.
 917
 918        This method adds one or more literals (that is, a boolean variable or its
 919        negation) as enforcement literals. The conjunction of all these literals
 920        determines whether the constraint is active or not. It acts as an
 921        implication, so if the conjunction is true, it implies that the constraint
 922        must be enforced. If it is false, then the constraint is ignored.
 923
 924        BoolOr, BoolAnd, and linear constraints all support enforcement literals.
 925
 926        Args:
 927          *boolvar: One or more Boolean literals.
 928
 929        Returns:
 930          self.
 931        """
 932        for lit in ExpandGeneratorOrTuple(boolvar):
 933            if (isinstance(lit, bool) and lit) or (cmh.is_integral(lit) and lit == 1):
 934                # Always true. Do nothing.
 935                pass
 936            elif (isinstance(lit, bool) and not lit) or (
 937                cmh.is_integral(lit) and lit == 0
 938            ):
 939                self.__constraint.enforcement_literal.append(
 940                    self.__cp_model.NewConstant(0).Index()
 941                )
 942            else:
 943                self.__constraint.enforcement_literal.append(
 944                    cast(Union[IntVar, _NotBooleanVariable], lit).Index()
 945                )
 946        return self
 947
 948    def WithName(self, name: str) -> "Constraint":
 949        """Sets the name of the constraint."""
 950        if name:
 951            self.__constraint.name = name
 952        else:
 953            self.__constraint.ClearField("name")
 954        return self
 955
 956    def Name(self) -> str:
 957        """Returns the name of the constraint."""
 958        if not self.__constraint or not self.__constraint.name:
 959            return ""
 960        return self.__constraint.name
 961
 962    def Index(self) -> int:
 963        """Returns the index of the constraint in the model."""
 964        return self.__index
 965
 966    def Proto(self) -> cp_model_pb2.ConstraintProto:
 967        """Returns the constraint protobuf."""
 968        return self.__constraint
 969
 970
 971class IntervalVar:
 972    """Represents an Interval variable.
 973
 974    An interval variable is both a constraint and a variable. It is defined by
 975    three integer variables: start, size, and end.
 976
 977    It is a constraint because, internally, it enforces that start + size == end.
 978
 979    It is also a variable as it can appear in specific scheduling constraints:
 980    NoOverlap, NoOverlap2D, Cumulative.
 981
 982    Optionally, an enforcement literal can be added to this constraint, in which
 983    case these scheduling constraints will ignore interval variables with
 984    enforcement literals assigned to false. Conversely, these constraints will
 985    also set these enforcement literals to false if they cannot fit these
 986    intervals into the schedule.
 987    """
 988
 989    def __init__(
 990        self,
 991        model: cp_model_pb2.CpModelProto,
 992        start: Union[cp_model_pb2.LinearExpressionProto, int],
 993        size: Optional[cp_model_pb2.LinearExpressionProto],
 994        end: Optional[cp_model_pb2.LinearExpressionProto],
 995        is_present_index: Optional[int],
 996        name: Optional[str],
 997    ):
 998        self.__model: cp_model_pb2.CpModelProto = model
 999        # As with the IntVar::__init__ method, we hack the __init__ method to
1000        # support two use cases:
1001        #   case 1: called when creating a new interval variable.
1002        #      {start|size|end} are linear expressions, is_present_index is either
1003        #      None or the index of a Boolean literal. name is a string
1004        #   case 2: called when querying an existing interval variable.
1005        #      start_index is an int, all parameters after are None.
1006        if size is None and end is None and is_present_index is None and name is None:
1007            self.__index: int = cast(int, start)
1008            self.__ct: cp_model_pb2.ConstraintProto = model.constraints[self.__index]
1009        else:
1010            self.__index: int = len(model.constraints)
1011            self.__ct: cp_model_pb2.ConstraintProto = self.__model.constraints.add()
1012            self.__ct.interval.start.CopyFrom(start)
1013            self.__ct.interval.size.CopyFrom(size)
1014            self.__ct.interval.end.CopyFrom(end)
1015            if is_present_index is not None:
1016                self.__ct.enforcement_literal.append(is_present_index)
1017            if name:
1018                self.__ct.name = name
1019
1020    def Index(self) -> int:
1021        """Returns the index of the interval constraint in the model."""
1022        return self.__index
1023
1024    def Proto(self) -> cp_model_pb2.IntervalConstraintProto:
1025        """Returns the interval protobuf."""
1026        return self.__ct.interval
1027
1028    def __str__(self):
1029        return self.__ct.name
1030
1031    def __repr__(self):
1032        interval = self.__ct.interval
1033        if self.__ct.enforcement_literal:
1034            return "%s(start = %s, size = %s, end = %s, is_present = %s)" % (
1035                self.__ct.name,
1036                ShortExprName(self.__model, interval.start),
1037                ShortExprName(self.__model, interval.size),
1038                ShortExprName(self.__model, interval.end),
1039                ShortName(self.__model, self.__ct.enforcement_literal[0]),
1040            )
1041        else:
1042            return "%s(start = %s, size = %s, end = %s)" % (
1043                self.__ct.name,
1044                ShortExprName(self.__model, interval.start),
1045                ShortExprName(self.__model, interval.size),
1046                ShortExprName(self.__model, interval.end),
1047            )
1048
1049    def Name(self) -> str:
1050        if not self.__ct or not self.__ct.name:
1051            return ""
1052        return self.__ct.name
1053
1054    def StartExpr(self) -> LinearExprT:
1055        return LinearExpr.RebuildFromLinearExpressionProto(
1056            self.__model, self.__ct.interval.start
1057        )
1058
1059    def SizeExpr(self) -> LinearExprT:
1060        return LinearExpr.RebuildFromLinearExpressionProto(
1061            self.__model, self.__ct.interval.size
1062        )
1063
1064    def EndExpr(self) -> LinearExprT:
1065        return LinearExpr.RebuildFromLinearExpressionProto(
1066            self.__model, self.__ct.interval.end
1067        )
1068
1069
1070def ObjectIsATrueLiteral(literal: LiteralT) -> bool:
1071    """Checks if literal is either True, or a Boolean literals fixed to True."""
1072    if isinstance(literal, IntVar):
1073        proto = literal.Proto()
1074        return len(proto.domain) == 2 and proto.domain[0] == 1 and proto.domain[1] == 1
1075    if isinstance(literal, _NotBooleanVariable):
1076        proto = literal.Not().Proto()
1077        return len(proto.domain) == 2 and proto.domain[0] == 0 and proto.domain[1] == 0
1078    if cmh.is_integral(literal):
1079        return int(literal) == 1
1080    return False
1081
1082
1083def ObjectIsAFalseLiteral(literal: LiteralT) -> bool:
1084    """Checks if literal is either False, or a Boolean literals fixed to False."""
1085    if isinstance(literal, IntVar):
1086        proto = literal.Proto()
1087        return len(proto.domain) == 2 and proto.domain[0] == 0 and proto.domain[1] == 0
1088    if isinstance(literal, _NotBooleanVariable):
1089        proto = literal.Not().Proto()
1090        return len(proto.domain) == 2 and proto.domain[0] == 1 and proto.domain[1] == 1
1091    if cmh.is_integral(literal):
1092        return int(literal) == 0
1093    return False
1094
1095
1096class CpModel:
1097    """Methods for building a CP model.
1098
1099    Methods beginning with:
1100
1101    * ```New``` create integer, boolean, or interval variables.
1102    * ```Add``` create new constraints and add them to the model.
1103    """
1104
1105    def __init__(self):
1106        self.__model: cp_model_pb2.CpModelProto = cp_model_pb2.CpModelProto()
1107        self.__constant_map = {}
1108
1109    # Naming.
1110    def Name(self) -> str:
1111        """Returns the name of the model."""
1112        if not self.__model or not self.__model.name:
1113            return ""
1114        return self.__model.name
1115
1116    def SetName(self, name: str):
1117        """Sets the name of the model."""
1118        self.__model.name = name
1119
1120    # Integer variable.
1121
1122    def NewIntVar(self, lb: IntegralT, ub: IntegralT, name: str) -> IntVar:
1123        """Create an integer variable with domain [lb, ub].
1124
1125        The CP-SAT solver is limited to integer variables. If you have fractional
1126        values, scale them up so that they become integers; if you have strings,
1127        encode them as integers.
1128
1129        Args:
1130          lb: Lower bound for the variable.
1131          ub: Upper bound for the variable.
1132          name: The name of the variable.
1133
1134        Returns:
1135          a variable whose domain is [lb, ub].
1136        """
1137
1138        return IntVar(self.__model, Domain(lb, ub), name)
1139
1140    def NewIntVarFromDomain(self, domain: Domain, name: str) -> IntVar:
1141        """Create an integer variable from a domain.
1142
1143        A domain is a set of integers specified by a collection of intervals.
1144        For example, `model.NewIntVarFromDomain(cp_model.
1145             Domain.FromIntervals([[1, 2], [4, 6]]), 'x')`
1146
1147        Args:
1148          domain: An instance of the Domain class.
1149          name: The name of the variable.
1150
1151        Returns:
1152            a variable whose domain is the given domain.
1153        """
1154        return IntVar(self.__model, domain, name)
1155
1156    def NewBoolVar(self, name: str) -> IntVar:
1157        """Creates a 0-1 variable with the given name."""
1158        return IntVar(self.__model, Domain(0, 1), name)
1159
1160    def NewConstant(self, value: IntegralT) -> IntVar:
1161        """Declares a constant integer."""
1162        return IntVar(self.__model, self.GetOrMakeIndexFromConstant(value), None)
1163
1164    def NewIntVarSeries(
1165        self,
1166        name: str,
1167        index: pd.Index,
1168        lower_bounds: Union[IntegralT, pd.Series],
1169        upper_bounds: Union[IntegralT, pd.Series],
1170    ) -> pd.Series:
1171        """Creates a series of (scalar-valued) variables with the given name.
1172
1173        Args:
1174          name (str): Required. The name of the variable set.
1175          index (pd.Index): Required. The index to use for the variable set.
1176          lower_bounds (Union[int, pd.Series]): A lower bound for variables in the
1177            set. If a `pd.Series` is passed in, it will be based on the
1178            corresponding values of the pd.Series.
1179          upper_bounds (Union[int, pd.Series]): An upper bound for variables in the
1180            set. If a `pd.Series` is passed in, it will be based on the
1181            corresponding values of the pd.Series.
1182
1183        Returns:
1184          pd.Series: The variable set indexed by its corresponding dimensions.
1185
1186        Raises:
1187          TypeError: if the `index` is invalid (e.g. a `DataFrame`).
1188          ValueError: if the `name` is not a valid identifier or already exists.
1189          ValueError: if the `lowerbound` is greater than the `upperbound`.
1190          ValueError: if the index of `lower_bound`, or `upper_bound` does not match
1191          the input index.
1192        """
1193        if not isinstance(index, pd.Index):
1194            raise TypeError("Non-index object is used as index")
1195        if not name.isidentifier():
1196            raise ValueError("name={} is not a valid identifier".format(name))
1197        if (
1198            cmh.is_integral(lower_bounds)
1199            and cmh.is_integral(upper_bounds)
1200            and lower_bounds > upper_bounds
1201        ):
1202            raise ValueError(
1203                f"lower_bound={lower_bounds} is greater than"
1204                f" upper_bound={upper_bounds} for variable set={name}"
1205            )
1206
1207        lower_bounds = _ConvertToIntegralSeriesAndValidateIndex(lower_bounds, index)
1208        upper_bounds = _ConvertToIntegralSeriesAndValidateIndex(upper_bounds, index)
1209        return pd.Series(
1210            index=index,
1211            data=[
1212                # pylint: disable=g-complex-comprehension
1213                IntVar(
1214                    model=self.__model,
1215                    name=f"{name}[{i}]",
1216                    domain=Domain(lower_bounds[i], upper_bounds[i]),
1217                )
1218                for i in index
1219            ],
1220        )
1221
1222    def NewBoolVarSeries(
1223        self,
1224        name: str,
1225        index: pd.Index,
1226    ) -> pd.Series:
1227        """Creates a series of (scalar-valued) variables with the given name.
1228
1229        Args:
1230          name (str): Required. The name of the variable set.
1231          index (pd.Index): Required. The index to use for the variable set.
1232
1233        Returns:
1234          pd.Series: The variable set indexed by its corresponding dimensions.
1235
1236        Raises:
1237          TypeError: if the `index` is invalid (e.g. a `DataFrame`).
1238          ValueError: if the `name` is not a valid identifier or already exists.
1239        """
1240        return self.NewIntVarSeries(
1241            name=name, index=index, lower_bounds=0, upper_bounds=1
1242        )
1243
1244    # Linear constraints.
1245
1246    def AddLinearConstraint(
1247        self, linear_expr: LinearExprT, lb: IntegralT, ub: IntegralT
1248    ) -> Constraint:
1249        """Adds the constraint: `lb <= linear_expr <= ub`."""
1250        return self.AddLinearExpressionInDomain(linear_expr, Domain(lb, ub))
1251
1252    def AddLinearExpressionInDomain(
1253        self, linear_expr: LinearExprT, domain: Domain
1254    ) -> Constraint:
1255        """Adds the constraint: `linear_expr` in `domain`."""
1256        if isinstance(linear_expr, LinearExpr):
1257            ct = Constraint(self)
1258            model_ct = self.__model.constraints[ct.Index()]
1259            coeffs_map, constant = linear_expr.GetIntegerVarValueMap()
1260            for t in coeffs_map.items():
1261                if not isinstance(t[0], IntVar):
1262                    raise TypeError("Wrong argument" + str(t))
1263                c = cmh.assert_is_int64(t[1])
1264                model_ct.linear.vars.append(t[0].Index())
1265                model_ct.linear.coeffs.append(c)
1266            model_ct.linear.domain.extend(
1267                [
1268                    cmh.capped_subtraction(x, constant)
1269                    for x in domain.FlattenedIntervals()
1270                ]
1271            )
1272            return ct
1273        elif cmh.is_integral(linear_expr):
1274            if not domain.Contains(int(linear_expr)):
1275                return self.AddBoolOr([])  # Evaluate to false.
1276            else:
1277                return self.AddBoolAnd([])  # Evaluate to true.
1278        raise TypeError(
1279            "Not supported: CpModel.AddLinearExpressionInDomain("
1280            + str(linear_expr)
1281            + " "
1282            + str(domain)
1283            + ")"
1284        )
1285
1286    def Add(self, ct: Union[BoundedLinearExpression, bool]) -> Constraint:
1287        """Adds a `BoundedLinearExpression` to the model.
1288
1289        Args:
1290          ct: A [`BoundedLinearExpression`](#boundedlinearexpression).
1291
1292        Returns:
1293          An instance of the `Constraint` class.
1294        """
1295        if isinstance(ct, BoundedLinearExpression):
1296            return self.AddLinearExpressionInDomain(
1297                ct.Expression(), Domain.FromFlatIntervals(ct.Bounds())
1298            )
1299        elif ct and isinstance(ct, bool):
1300            return self.AddBoolOr([True])
1301        elif not ct and isinstance(ct, bool):
1302            return self.AddBoolOr([])  # Evaluate to false.
1303        raise TypeError("Not supported: CpModel.Add(" + str(ct) + ")")
1304
1305    # General Integer Constraints.
1306
1307    @overload
1308    def AddAllDifferent(self, expressions: Iterable[LinearExprT]) -> Constraint:
1309        ...
1310
1311    @overload
1312    def AddAllDifferent(self, *expressions: LinearExprT) -> Constraint:
1313        ...
1314
1315    def AddAllDifferent(self, *expressions):
1316        """Adds AllDifferent(expressions).
1317
1318        This constraint forces all expressions to have different values.
1319
1320        Args:
1321          *expressions: simple expressions of the form a * var + constant.
1322
1323        Returns:
1324          An instance of the `Constraint` class.
1325        """
1326        ct = Constraint(self)
1327        model_ct = self.__model.constraints[ct.Index()]
1328        expanded = ExpandGeneratorOrTuple(expressions)
1329        model_ct.all_diff.exprs.extend(
1330            [self.ParseLinearExpression(x) for x in expanded]
1331        )
1332        return ct
1333
1334    def AddElement(
1335        self, index: VariableT, variables: Sequence[VariableT], target: VariableT
1336    ) -> Constraint:
1337        """Adds the element constraint: `variables[index] == target`.
1338
1339        Args:
1340          index: The index of the variable that's being constrained.
1341          variables: A list of variables.
1342          target: The value that the variable must be equal to.
1343
1344        Returns:
1345          An instance of the `Constraint` class.
1346        """
1347
1348        if not variables:
1349            raise ValueError("AddElement expects a non-empty variables array")
1350
1351        if cmh.is_integral(index):
1352            return self.Add(list(variables)[int(index)] == target)
1353
1354        ct = Constraint(self)
1355        model_ct = self.__model.constraints[ct.Index()]
1356        model_ct.element.index = self.GetOrMakeIndex(index)
1357        model_ct.element.vars.extend([self.GetOrMakeIndex(x) for x in variables])
1358        model_ct.element.target = self.GetOrMakeIndex(target)
1359        return ct
1360
1361    def AddCircuit(self, arcs: Sequence[ArcT]) -> Constraint:
1362        """Adds Circuit(arcs).
1363
1364        Adds a circuit constraint from a sparse list of arcs that encode the graph.
1365
1366        A circuit is a unique Hamiltonian path in a subgraph of the total
1367        graph. In case a node 'i' is not in the path, then there must be a
1368        loop arc 'i -> i' associated with a true literal. Otherwise
1369        this constraint will fail.
1370
1371        Args:
1372          arcs: a list of arcs. An arc is a tuple (source_node, destination_node,
1373            literal). The arc is selected in the circuit if the literal is true.
1374            Both source_node and destination_node must be integers between 0 and the
1375            number of nodes - 1.
1376
1377        Returns:
1378          An instance of the `Constraint` class.
1379
1380        Raises:
1381          ValueError: If the list of arcs is empty.
1382        """
1383        if not arcs:
1384            raise ValueError("AddCircuit expects a non-empty array of arcs")
1385        ct = Constraint(self)
1386        model_ct = self.__model.constraints[ct.Index()]
1387        for arc in arcs:
1388            tail = cmh.assert_is_int32(arc[0])
1389            head = cmh.assert_is_int32(arc[1])
1390            lit = self.GetOrMakeBooleanIndex(arc[2])
1391            model_ct.circuit.tails.append(tail)
1392            model_ct.circuit.heads.append(head)
1393            model_ct.circuit.literals.append(lit)
1394        return ct
1395
1396    def AddMultipleCircuit(self, arcs: Sequence[ArcT]) -> Constraint:
1397        """Adds a multiple circuit constraint, aka the 'VRP' constraint.
1398
1399        The direct graph where arc #i (from tails[i] to head[i]) is present iff
1400        literals[i] is true must satisfy this set of properties:
1401        - #incoming arcs == 1 except for node 0.
1402        - #outgoing arcs == 1 except for node 0.
1403        - for node zero, #incoming arcs == #outgoing arcs.
1404        - There are no duplicate arcs.
1405        - Self-arcs are allowed except for node 0.
1406        - There is no cycle in this graph, except through node 0.
1407
1408        Args:
1409          arcs: a list of arcs. An arc is a tuple (source_node, destination_node,
1410            literal). The arc is selected in the circuit if the literal is true.
1411            Both source_node and destination_node must be integers between 0 and the
1412            number of nodes - 1.
1413
1414        Returns:
1415          An instance of the `Constraint` class.
1416
1417        Raises:
1418          ValueError: If the list of arcs is empty.
1419        """
1420        if not arcs:
1421            raise ValueError("AddMultipleCircuit expects a non-empty array of arcs")
1422        ct = Constraint(self)
1423        model_ct = self.__model.constraints[ct.Index()]
1424        for arc in arcs:
1425            tail = cmh.assert_is_int32(arc[0])
1426            head = cmh.assert_is_int32(arc[1])
1427            lit = self.GetOrMakeBooleanIndex(arc[2])
1428            model_ct.routes.tails.append(tail)
1429            model_ct.routes.heads.append(head)
1430            model_ct.routes.literals.append(lit)
1431        return ct
1432
1433    def AddAllowedAssignments(
1434        self,
1435        variables: Sequence[VariableT],
1436        tuples_list: Iterable[Sequence[IntegralT]],
1437    ) -> Constraint:
1438        """Adds AllowedAssignments(variables, tuples_list).
1439
1440        An AllowedAssignments constraint is a constraint on an array of variables,
1441        which requires that when all variables are assigned values, the resulting
1442        array equals one of the  tuples in `tuple_list`.
1443
1444        Args:
1445          variables: A list of variables.
1446          tuples_list: A list of admissible tuples. Each tuple must have the same
1447            length as the variables, and the ith value of a tuple corresponds to the
1448            ith variable.
1449
1450        Returns:
1451          An instance of the `Constraint` class.
1452
1453        Raises:
1454          TypeError: If a tuple does not have the same size as the list of
1455              variables.
1456          ValueError: If the array of variables is empty.
1457        """
1458
1459        if not variables:
1460            raise ValueError(
1461                "AddAllowedAssignments expects a non-empty variables array"
1462            )
1463
1464        ct = Constraint(self)
1465        model_ct = self.__model.constraints[ct.Index()]
1466        model_ct.table.vars.extend([self.GetOrMakeIndex(x) for x in variables])
1467        arity = len(variables)
1468        for t in tuples_list:
1469            if len(t) != arity:
1470                raise TypeError("Tuple " + str(t) + " has the wrong arity")
1471            ar = []
1472            for v in t:
1473                ar.append(cmh.assert_is_int64(v))
1474            model_ct.table.values.extend(ar)
1475        return ct
1476
1477    def AddForbiddenAssignments(
1478        self,
1479        variables: Sequence[VariableT],
1480        tuples_list: Iterable[Sequence[IntegralT]],
1481    ) -> Constraint:
1482        """Adds AddForbiddenAssignments(variables, [tuples_list]).
1483
1484        A ForbiddenAssignments constraint is a constraint on an array of variables
1485        where the list of impossible combinations is provided in the tuples list.
1486
1487        Args:
1488          variables: A list of variables.
1489          tuples_list: A list of forbidden tuples. Each tuple must have the same
1490            length as the variables, and the *i*th value of a tuple corresponds to
1491            the *i*th variable.
1492
1493        Returns:
1494          An instance of the `Constraint` class.
1495
1496        Raises:
1497          TypeError: If a tuple does not have the same size as the list of
1498                     variables.
1499          ValueError: If the array of variables is empty.
1500        """
1501
1502        if not variables:
1503            raise ValueError(
1504                "AddForbiddenAssignments expects a non-empty variables array"
1505            )
1506
1507        index = len(self.__model.constraints)
1508        ct = self.AddAllowedAssignments(variables, tuples_list)
1509        self.__model.constraints[index].table.negated = True
1510        return ct
1511
1512    def AddAutomaton(
1513        self,
1514        transition_variables: Sequence[VariableT],
1515        starting_state: IntegralT,
1516        final_states: Sequence[IntegralT],
1517        transition_triples: Sequence[Tuple[IntegralT, IntegralT, IntegralT]],
1518    ) -> Constraint:
1519        """Adds an automaton constraint.
1520
1521        An automaton constraint takes a list of variables (of size *n*), an initial
1522        state, a set of final states, and a set of transitions. A transition is a
1523        triplet (*tail*, *transition*, *head*), where *tail* and *head* are states,
1524        and *transition* is the label of an arc from *head* to *tail*,
1525        corresponding to the value of one variable in the list of variables.
1526
1527        This automaton will be unrolled into a flow with *n* + 1 phases. Each phase
1528        contains the possible states of the automaton. The first state contains the
1529        initial state. The last phase contains the final states.
1530
1531        Between two consecutive phases *i* and *i* + 1, the automaton creates a set
1532        of arcs. For each transition (*tail*, *transition*, *head*), it will add
1533        an arc from the state *tail* of phase *i* and the state *head* of phase
1534        *i* + 1. This arc is labeled by the value *transition* of the variables
1535        `variables[i]`. That is, this arc can only be selected if `variables[i]`
1536        is assigned the value *transition*.
1537
1538        A feasible solution of this constraint is an assignment of variables such
1539        that, starting from the initial state in phase 0, there is a path labeled by
1540        the values of the variables that ends in one of the final states in the
1541        final phase.
1542
1543        Args:
1544          transition_variables: A non-empty list of variables whose values
1545            correspond to the labels of the arcs traversed by the automaton.
1546          starting_state: The initial state of the automaton.
1547          final_states: A non-empty list of admissible final states.
1548          transition_triples: A list of transitions for the automaton, in the
1549            following format (current_state, variable_value, next_state).
1550
1551        Returns:
1552          An instance of the `Constraint` class.
1553
1554        Raises:
1555          ValueError: if `transition_variables`, `final_states`, or
1556            `transition_triples` are empty.
1557        """
1558
1559        if not transition_variables:
1560            raise ValueError(
1561                "AddAutomaton expects a non-empty transition_variables array"
1562            )
1563        if not final_states:
1564            raise ValueError("AddAutomaton expects some final states")
1565
1566        if not transition_triples:
1567            raise ValueError("AddAutomaton expects some transition triples")
1568
1569        ct = Constraint(self)
1570        model_ct = self.__model.constraints[ct.Index()]
1571        model_ct.automaton.vars.extend(
1572            [self.GetOrMakeIndex(x) for x in transition_variables]
1573        )
1574        starting_state = cmh.assert_is_int64(starting_state)
1575        model_ct.automaton.starting_state = starting_state
1576        for v in final_states:
1577            v = cmh.assert_is_int64(v)
1578            model_ct.automaton.final_states.append(v)
1579        for t in transition_triples:
1580            if len(t) != 3:
1581                raise TypeError("Tuple " + str(t) + " has the wrong arity (!= 3)")
1582            tail = cmh.assert_is_int64(t[0])
1583            label = cmh.assert_is_int64(t[1])
1584            head = cmh.assert_is_int64(t[2])
1585            model_ct.automaton.transition_tail.append(tail)
1586            model_ct.automaton.transition_label.append(label)
1587            model_ct.automaton.transition_head.append(head)
1588        return ct
1589
1590    def AddInverse(
1591        self,
1592        variables: Sequence[VariableT],
1593        inverse_variables: Sequence[VariableT],
1594    ) -> Constraint:
1595        """Adds Inverse(variables, inverse_variables).
1596
1597        An inverse constraint enforces that if `variables[i]` is assigned a value
1598        `j`, then `inverse_variables[j]` is assigned a value `i`. And vice versa.
1599
1600        Args:
1601          variables: An array of integer variables.
1602          inverse_variables: An array of integer variables.
1603
1604        Returns:
1605          An instance of the `Constraint` class.
1606
1607        Raises:
1608          TypeError: if variables and inverse_variables have different lengths, or
1609              if they are empty.
1610        """
1611
1612        if not variables or not inverse_variables:
1613            raise TypeError("The Inverse constraint does not accept empty arrays")
1614        if len(variables) != len(inverse_variables):
1615            raise TypeError(
1616                "In the inverse constraint, the two array variables and"
1617                " inverse_variables must have the same length."
1618            )
1619        ct = Constraint(self)
1620        model_ct = self.__model.constraints[ct.Index()]
1621        model_ct.inverse.f_direct.extend([self.GetOrMakeIndex(x) for x in variables])
1622        model_ct.inverse.f_inverse.extend(
1623            [self.GetOrMakeIndex(x) for x in inverse_variables]
1624        )
1625        return ct
1626
1627    def AddReservoirConstraint(
1628        self,
1629        times: Iterable[LinearExprT],
1630        level_changes: Iterable[LinearExprT],
1631        min_level: int,
1632        max_level: int,
1633    ) -> Constraint:
1634        """Adds Reservoir(times, level_changes, min_level, max_level).
1635
1636        Maintains a reservoir level within bounds. The water level starts at 0, and
1637        at any time, it must be between min_level and max_level.
1638
1639        If the affine expression `times[i]` is assigned a value t, then the current
1640        level changes by `level_changes[i]`, which is constant, at time t.
1641
1642         Note that min level must be <= 0, and the max level must be >= 0. Please
1643         use fixed level_changes to simulate initial state.
1644
1645         Therefore, at any time:
1646             sum(level_changes[i] if times[i] <= t) in [min_level, max_level]
1647
1648        Args:
1649          times: A list of affine expressions which specify the time of the filling
1650            or emptying the reservoir.
1651          level_changes: A list of integer values that specifies the amount of the
1652            emptying or filling.
1653          min_level: At any time, the level of the reservoir must be greater or
1654            equal than the min level.
1655          max_level: At any time, the level of the reservoir must be less or equal
1656            than the max level.
1657
1658        Returns:
1659          An instance of the `Constraint` class.
1660
1661        Raises:
1662          ValueError: if max_level < min_level.
1663
1664          ValueError: if max_level < 0.
1665
1666          ValueError: if min_level > 0
1667        """
1668
1669        if max_level < min_level:
1670            raise ValueError("Reservoir constraint must have a max_level >= min_level")
1671
1672        if max_level < 0:
1673            raise ValueError("Reservoir constraint must have a max_level >= 0")
1674
1675        if min_level > 0:
1676            raise ValueError("Reservoir constraint must have a min_level <= 0")
1677
1678        ct = Constraint(self)
1679        model_ct = self.__model.constraints[ct.Index()]
1680        model_ct.reservoir.time_exprs.extend(
1681            [self.ParseLinearExpression(x) for x in times]
1682        )
1683        model_ct.reservoir.level_changes.extend(
1684            [self.ParseLinearExpression(x) for x in level_changes]
1685        )
1686        model_ct.reservoir.min_level = min_level
1687        model_ct.reservoir.max_level = max_level
1688        return ct
1689
1690    def AddReservoirConstraintWithActive(
1691        self,
1692        times: Iterable[LinearExprT],
1693        level_changes: Iterable[LinearExprT],
1694        actives: Iterable[LiteralT],
1695        min_level: int,
1696        max_level: int,
1697    ) -> Constraint:
1698        """Adds Reservoir(times, level_changes, actives, min_level, max_level).
1699
1700        Maintains a reservoir level within bounds. The water level starts at 0, and
1701        at any time, it must be between min_level and max_level.
1702
1703        If the variable `times[i]` is assigned a value t, and `actives[i]` is
1704        `True`, then the current level changes by `level_changes[i]`, which is
1705        constant,
1706        at time t.
1707
1708         Note that min level must be <= 0, and the max level must be >= 0. Please
1709         use fixed level_changes to simulate initial state.
1710
1711         Therefore, at any time:
1712             sum(level_changes[i] * actives[i] if times[i] <= t) in [min_level,
1713             max_level]
1714
1715
1716        The array of boolean variables 'actives', if defined, indicates which
1717        actions are actually performed.
1718
1719        Args:
1720          times: A list of affine expressions which specify the time of the filling
1721            or emptying the reservoir.
1722          level_changes: A list of integer values that specifies the amount of the
1723            emptying or filling.
1724          actives: a list of boolean variables. They indicates if the
1725            emptying/refilling events actually take place.
1726          min_level: At any time, the level of the reservoir must be greater or
1727            equal than the min level.
1728          max_level: At any time, the level of the reservoir must be less or equal
1729            than the max level.
1730
1731        Returns:
1732          An instance of the `Constraint` class.
1733
1734        Raises:
1735          ValueError: if max_level < min_level.
1736
1737          ValueError: if max_level < 0.
1738
1739          ValueError: if min_level > 0
1740        """
1741
1742        if max_level < min_level:
1743            raise ValueError("Reservoir constraint must have a max_level >= min_level")
1744
1745        if max_level < 0:
1746            raise ValueError("Reservoir constraint must have a max_level >= 0")
1747
1748        if min_level > 0:
1749            raise ValueError("Reservoir constraint must have a min_level <= 0")
1750
1751        ct = Constraint(self)
1752        model_ct = self.__model.constraints[ct.Index()]
1753        model_ct.reservoir.time_exprs.extend(
1754            [self.ParseLinearExpression(x) for x in times]
1755        )
1756        model_ct.reservoir.level_changes.extend(
1757            [self.ParseLinearExpression(x) for x in level_changes]
1758        )
1759        model_ct.reservoir.active_literals.extend(
1760            [self.GetOrMakeBooleanIndex(x) for x in actives]
1761        )
1762        model_ct.reservoir.min_level = min_level
1763        model_ct.reservoir.max_level = max_level
1764        return ct
1765
1766    def AddMapDomain(
1767        self, var: IntVar, bool_var_array: Iterable[IntVar], offset: IntegralT = 0
1768    ):
1769        """Adds `var == i + offset <=> bool_var_array[i] == true for all i`."""
1770
1771        for i, bool_var in enumerate(bool_var_array):
1772            b_index = bool_var.Index()
1773            var_index = var.Index()
1774            model_ct = self.__model.constraints.add()
1775            model_ct.linear.vars.append(var_index)
1776            model_ct.linear.coeffs.append(1)
1777            model_ct.linear.domain.extend([offset + i, offset + i])
1778            model_ct.enforcement_literal.append(b_index)
1779
1780            model_ct = self.__model.constraints.add()
1781            model_ct.linear.vars.append(var_index)
1782            model_ct.linear.coeffs.append(1)
1783            model_ct.enforcement_literal.append(-b_index - 1)
1784            if offset + i - 1 >= INT_MIN:
1785                model_ct.linear.domain.extend([INT_MIN, offset + i - 1])
1786            if offset + i + 1 <= INT_MAX:
1787                model_ct.linear.domain.extend([offset + i + 1, INT_MAX])
1788
1789    def AddImplication(self, a: LiteralT, b: LiteralT) -> Constraint:
1790        """Adds `a => b` (`a` implies `b`)."""
1791        ct = Constraint(self)
1792        model_ct = self.__model.constraints[ct.Index()]
1793        model_ct.bool_or.literals.append(self.GetOrMakeBooleanIndex(b))
1794        model_ct.enforcement_literal.append(self.GetOrMakeBooleanIndex(a))
1795        return ct
1796
1797    @overload
1798    def AddBoolOr(self, literals: Iterable[LiteralT]) -> Constraint:
1799        ...
1800
1801    @overload
1802    def AddBoolOr(self, *literals: LiteralT) -> Constraint:
1803        ...
1804
1805    def AddBoolOr(self, *literals):
1806        """Adds `Or(literals) == true`: Sum(literals) >= 1."""
1807        ct = Constraint(self)
1808        model_ct = self.__model.constraints[ct.Index()]
1809        model_ct.bool_or.literals.extend(
1810            [self.GetOrMakeBooleanIndex(x) for x in ExpandGeneratorOrTuple(literals)]
1811        )
1812        return ct
1813
1814    @overload
1815    def AddAtLeastOne(self, literals: Iterable[LiteralT]) -> Constraint:
1816        ...
1817
1818    @overload
1819    def AddAtLeastOne(self, *literals: LiteralT) -> Constraint:
1820        ...
1821
1822    def AddAtLeastOne(self, *literals):
1823        """Same as `AddBoolOr`: `Sum(literals) >= 1`."""
1824        return self.AddBoolOr(*literals)
1825
1826    @overload
1827    def AddAtMostOne(self, literals: Iterable[LiteralT]) -> Constraint:
1828        ...
1829
1830    @overload
1831    def AddAtMostOne(self, *literals: LiteralT) -> Constraint:
1832        ...
1833
1834    def AddAtMostOne(self, *literals):
1835        """Adds `AtMostOne(literals)`: `Sum(literals) <= 1`."""
1836        ct = Constraint(self)
1837        model_ct = self.__model.constraints[ct.Index()]
1838        model_ct.at_most_one.literals.extend(
1839            [self.GetOrMakeBooleanIndex(x) for x in ExpandGeneratorOrTuple(literals)]
1840        )
1841        return ct
1842
1843    @overload
1844    def AddExactlyOne(self, literals: Iterable[LiteralT]) -> Constraint:
1845        ...
1846
1847    @overload
1848    def AddExactlyOne(self, *literals: LiteralT) -> Constraint:
1849        ...
1850
1851    def AddExactlyOne(self, *literals):
1852        """Adds `ExactlyOne(literals)`: `Sum(literals) == 1`."""
1853        ct = Constraint(self)
1854        model_ct = self.__model.constraints[ct.Index()]
1855        model_ct.exactly_one.literals.extend(
1856            [self.GetOrMakeBooleanIndex(x) for x in ExpandGeneratorOrTuple(literals)]
1857        )
1858        return ct
1859
1860    @overload
1861    def AddBoolAnd(self, literals: Iterable[LiteralT]) -> Constraint:
1862        ...
1863
1864    @overload
1865    def AddBoolAnd(self, *literals: LiteralT) -> Constraint:
1866        ...
1867
1868    def AddBoolAnd(self, *literals):
1869        """Adds `And(literals) == true`."""
1870        ct = Constraint(self)
1871        model_ct = self.__model.constraints[ct.Index()]
1872        model_ct.bool_and.literals.extend(
1873            [self.GetOrMakeBooleanIndex(x) for x in ExpandGeneratorOrTuple(literals)]
1874        )
1875        return ct
1876
1877    @overload
1878    def AddBoolXOr(self, literals: Iterable[LiteralT]) -> Constraint:
1879        ...
1880
1881    @overload
1882    def AddBoolXOr(self, *literals: LiteralT) -> Constraint:
1883        ...
1884
1885    def AddBoolXOr(self, *literals):
1886        """Adds `XOr(literals) == true`.
1887
1888        In contrast to AddBoolOr and AddBoolAnd, it does not support
1889            .OnlyEnforceIf().
1890
1891        Args:
1892          *literals: the list of literals in the constraint.
1893
1894        Returns:
1895          An `Constraint` object.
1896        """
1897        ct = Constraint(self)
1898        model_ct = self.__model.constraints[ct.Index()]
1899        model_ct.bool_xor.literals.extend(
1900            [self.GetOrMakeBooleanIndex(x) for x in ExpandGeneratorOrTuple(literals)]
1901        )
1902        return ct
1903
1904    def AddMinEquality(
1905        self, target: LinearExprT, exprs: Iterable[LinearExprT]
1906    ) -> Constraint:
1907        """Adds `target == Min(exprs)`."""
1908        ct = Constraint(self)
1909        model_ct = self.__model.constraints[ct.Index()]
1910        model_ct.lin_max.exprs.extend(
1911            [self.ParseLinearExpression(x, True) for x in exprs]
1912        )
1913        model_ct.lin_max.target.CopyFrom(self.ParseLinearExpression(target, True))
1914        return ct
1915
1916    def AddMaxEquality(
1917        self, target: LinearExprT, exprs: Iterable[LinearExprT]
1918    ) -> Constraint:
1919        """Adds `target == Max(exprs)`."""
1920        ct = Constraint(self)
1921        model_ct = self.__model.constraints[ct.Index()]
1922        model_ct.lin_max.exprs.extend([self.ParseLinearExpression(x) for x in exprs])
1923        model_ct.lin_max.target.CopyFrom(self.ParseLinearExpression(target))
1924        return ct
1925
1926    def AddDivisionEquality(
1927        self, target: LinearExprT, num: LinearExprT, denom: LinearExprT
1928    ) -> Constraint:
1929        """Adds `target == num // denom` (integer division rounded towards 0)."""
1930        ct = Constraint(self)
1931        model_ct = self.__model.constraints[ct.Index()]
1932        model_ct.int_div.exprs.append(self.ParseLinearExpression(num))
1933        model_ct.int_div.exprs.append(self.ParseLinearExpression(denom))
1934        model_ct.int_div.target.CopyFrom(self.ParseLinearExpression(target))
1935        return ct
1936
1937    def AddAbsEquality(self, target: LinearExprT, expr: LinearExprT) -> Constraint:
1938        """Adds `target == Abs(var)`."""
1939        ct = Constraint(self)
1940        model_ct = self.__model.constraints[ct.Index()]
1941        model_ct.lin_max.exprs.append(self.ParseLinearExpression(expr))
1942        model_ct.lin_max.exprs.append(self.ParseLinearExpression(expr, True))
1943        model_ct.lin_max.target.CopyFrom(self.ParseLinearExpression(target))
1944        return ct
1945
1946    def AddModuloEquality(
1947        self, target: LinearExprT, var: LinearExprT, mod: LinearExprT
1948    ) -> Constraint:
1949        """Adds `target = var % mod`."""
1950        ct = Constraint(self)
1951        model_ct = self.__model.constraints[ct.Index()]
1952        model_ct.int_mod.exprs.append(self.ParseLinearExpression(var))
1953        model_ct.int_mod.exprs.append(self.ParseLinearExpression(mod))
1954        model_ct.int_mod.target.CopyFrom(self.ParseLinearExpression(target))
1955        return ct
1956
1957    def AddMultiplicationEquality(
1958        self,
1959        target: LinearExprT,
1960        *expressions: Union[Iterable[LinearExprT], LinearExprT],
1961    ) -> Constraint:
1962        """Adds `target == expressions[0] * .. * expressions[n]`."""
1963        ct = Constraint(self)
1964        model_ct = self.__model.constraints[ct.Index()]
1965        model_ct.int_prod.exprs.extend(
1966            [
1967                self.ParseLinearExpression(expr)
1968                for expr in ExpandGeneratorOrTuple(expressions)
1969            ]
1970        )
1971        model_ct.int_prod.target.CopyFrom(self.ParseLinearExpression(target))
1972        return ct
1973
1974    # Scheduling support
1975
1976    def NewIntervalVar(
1977        self, start: LinearExprT, size: LinearExprT, end: LinearExprT, name: str
1978    ) -> IntervalVar:
1979        """Creates an interval variable from start, size, and end.
1980
1981        An interval variable is a constraint, that is itself used in other
1982        constraints like NoOverlap.
1983
1984        Internally, it ensures that `start + size == end`.
1985
1986        Args:
1987          start: The start of the interval. It can be an affine or constant
1988            expression.
1989          size: The size of the interval. It can be an affine or constant
1990            expression.
1991          end: The end of the interval. It can be an affine or constant expression.
1992          name: The name of the interval variable.
1993
1994        Returns:
1995          An `IntervalVar` object.
1996        """
1997
1998        self.Add(start + size == end)
1999
2000        start_expr = self.ParseLinearExpression(start)
2001        size_expr = self.ParseLinearExpression(size)
2002        end_expr = self.ParseLinearExpression(end)
2003        if len(start_expr.vars) > 1:
2004            raise TypeError(
2005                "cp_model.NewIntervalVar: start must be affine or constant."
2006            )
2007        if len(size_expr.vars) > 1:
2008            raise TypeError("cp_model.NewIntervalVar: size must be affine or constant.")
2009        if len(end_expr.vars) > 1:
2010            raise TypeError("cp_model.NewIntervalVar: end must be affine or constant.")
2011        return IntervalVar(self.__model, start_expr, size_expr, end_expr, None, name)
2012
2013    def NewIntervalVarSeries(
2014        self,
2015        name: str,
2016        index: pd.Index,
2017        starts: Union[LinearExprT, pd.Series],
2018        sizes: Union[LinearExprT, pd.Series],
2019        ends: Union[LinearExprT, pd.Series],
2020    ) -> pd.Series:
2021        """Creates a series of interval variables with the given name.
2022
2023        Args:
2024          name (str): Required. The name of the variable set.
2025          index (pd.Index): Required. The index to use for the variable set.
2026          starts (Union[LinearExprT, pd.Series]): The start of each interval in the
2027            set. If a `pd.Series` is passed in, it will be based on the
2028            corresponding values of the pd.Series.
2029          sizes (Union[LinearExprT, pd.Series]): The size of each interval in the
2030            set. If a `pd.Series` is passed in, it will be based on the
2031            corresponding values of the pd.Series.
2032          ends (Union[LinearExprT, pd.Series]): The ends of each interval in the
2033            set. If a `pd.Series` is passed in, it will be based on the
2034            corresponding values of the pd.Series.
2035
2036        Returns:
2037          pd.Series: The interval variable set indexed by its corresponding
2038          dimensions.
2039
2040        Raises:
2041          TypeError: if the `index` is invalid (e.g. a `DataFrame`).
2042          ValueError: if the `name` is not a valid identifier or already exists.
2043          ValueError: if the all the indexes do not match.
2044        """
2045        if not isinstance(index, pd.Index):
2046            raise TypeError("Non-index object is used as index")
2047        if not name.isidentifier():
2048            raise ValueError("name={} is not a valid identifier".format(name))
2049
2050        starts = _ConvertToLinearExprSeriesAndValidateIndex(starts, index)
2051        sizes = _ConvertToLinearExprSeriesAndValidateIndex(sizes, index)
2052        ends = _ConvertToLinearExprSeriesAndValidateIndex(ends, index)
2053        interval_array = []
2054        for i in index:
2055            interval_array.append(
2056                self.NewIntervalVar(
2057                    start=starts[i],
2058                    size=sizes[i],
2059                    end=ends[i],
2060                    name=f"{name}[{i}]",
2061                )
2062            )
2063        return pd.Series(index=index, data=interval_array)
2064
2065    def NewFixedSizeIntervalVar(
2066        self, start: LinearExprT, size: IntegralT, name: str
2067    ) -> IntervalVar:
2068        """Creates an interval variable from start, and a fixed size.
2069
2070        An interval variable is a constraint, that is itself used in other
2071        constraints like NoOverlap.
2072
2073        Args:
2074          start: The start of the interval. It can be an affine or constant
2075            expression.
2076          size: The size of the interval. It must be an integer value.
2077          name: The name of the interval variable.
2078
2079        Returns:
2080          An `IntervalVar` object.
2081        """
2082        size = cmh.assert_is_int64(size)
2083        start_expr = self.ParseLinearExpression(start)
2084        size_expr = self.ParseLinearExpression(size)
2085        end_expr = self.ParseLinearExpression(start + size)
2086        if len(start_expr.vars) > 1:
2087            raise TypeError(
2088                "cp_model.NewIntervalVar: start must be affine or constant."
2089            )
2090        return IntervalVar(self.__model, start_expr, size_expr, end_expr, None, name)
2091
2092    def NewFixedSizeIntervalVarSeries(
2093        self,
2094        name: str,
2095        index: pd.Index,
2096        starts: Union[LinearExprT, pd.Series],
2097        sizes: Union[IntegralT, pd.Series],
2098    ) -> pd.Series:
2099        """Creates a series of interval variables with the given name.
2100
2101        Args:
2102          name (str): Required. The name of the variable set.
2103          index (pd.Index): Required. The index to use for the variable set.
2104          starts (Union[LinearExprT, pd.Series]): The start of each interval in the
2105            set. If a `pd.Series` is passed in, it will be based on the
2106            corresponding values of the pd.Series.
2107          sizes (Union[IntegralT, pd.Series]): The fixed size of each interval in
2108            the set. If a `pd.Series` is passed in, it will be based on the
2109            corresponding values of the pd.Series.
2110
2111        Returns:
2112          pd.Series: The interval variable set indexed by its corresponding
2113          dimensions.
2114
2115        Raises:
2116          TypeError: if the `index` is invalid (e.g. a `DataFrame`).
2117          ValueError: if the `name` is not a valid identifier or already exists.
2118          ValueError: if the all the indexes do not match.
2119        """
2120        if not isinstance(index, pd.Index):
2121            raise TypeError("Non-index object is used as index")
2122        if not name.isidentifier():
2123            raise ValueError("name={} is not a valid identifier".format(name))
2124
2125        starts = _ConvertToLinearExprSeriesAndValidateIndex(starts, index)
2126        sizes = _ConvertToIntegralSeriesAndValidateIndex(sizes, index)
2127        interval_array = []
2128        for i in index:
2129            interval_array.append(
2130                self.NewFixedSizeIntervalVar(
2131                    start=starts[i],
2132                    size=sizes[i],
2133                    name=f"{name}[{i}]",
2134                )
2135            )
2136        return pd.Series(index=index, data=interval_array)
2137
2138    def NewOptionalIntervalVar(
2139        self,
2140        start: LinearExprT,
2141        size: LinearExprT,
2142        end: LinearExprT,
2143        is_present: LiteralT,
2144        name: str,
2145    ) -> IntervalVar:
2146        """Creates an optional interval var from start, size, end, and is_present.
2147
2148        An optional interval variable is a constraint, that is itself used in other
2149        constraints like NoOverlap. This constraint is protected by a presence
2150        literal that indicates if it is active or not.
2151
2152        Internally, it ensures that `is_present` implies `start + size ==
2153        end`.
2154
2155        Args:
2156          start: The start of the interval. It can be an integer value, or an
2157            integer variable.
2158          size: The size of the interval. It can be an integer value, or an integer
2159            variable.
2160          end: The end of the interval. It can be an integer value, or an integer
2161            variable.
2162          is_present: A literal that indicates if the interval is active or not. A
2163            inactive interval is simply ignored by all constraints.
2164          name: The name of the interval variable.
2165
2166        Returns:
2167          An `IntervalVar` object.
2168        """
2169
2170        # Add the linear constraint.
2171        self.Add(start + size == end).OnlyEnforceIf(is_present)
2172
2173        # Creates the IntervalConstraintProto object.
2174        is_present_index = self.GetOrMakeBooleanIndex(is_present)
2175        start_expr = self.ParseLinearExpression(start)
2176        size_expr = self.ParseLinearExpression(size)
2177        end_expr = self.ParseLinearExpression(end)
2178        if len(start_expr.vars) > 1:
2179            raise TypeError(
2180                "cp_model.NewIntervalVar: start must be affine or constant."
2181            )
2182        if len(size_expr.vars) > 1:
2183            raise TypeError("cp_model.NewIntervalVar: size must be affine or constant.")
2184        if len(end_expr.vars) > 1:
2185            raise TypeError("cp_model.NewIntervalVar: end must be affine or constant.")
2186        return IntervalVar(
2187            self.__model, start_expr, size_expr, end_expr, is_present_index, name
2188        )
2189
2190    def NewOptionalIntervalVarSeries(
2191        self,
2192        name: str,
2193        index: pd.Index,
2194        starts: Union[LinearExprT, pd.Series],
2195        sizes: Union[LinearExprT, pd.Series],
2196        ends: Union[LinearExprT, pd.Series],
2197        are_present: Union[LiteralT, pd.Series],
2198    ) -> pd.Series:
2199        """Creates a series of interval variables with the given name.
2200
2201        Args:
2202          name (str): Required. The name of the variable set.
2203          index (pd.Index): Required. The index to use for the variable set.
2204          starts (Union[LinearExprT, pd.Series]): The start of each interval in the
2205            set. If a `pd.Series` is passed in, it will be based on the
2206            corresponding values of the pd.Series.
2207          sizes (Union[LinearExprT, pd.Series]): The size of each interval in the
2208            set. If a `pd.Series` is passed in, it will be based on the
2209            corresponding values of the pd.Series.
2210          ends (Union[LinearExprT, pd.Series]): The ends of each interval in the
2211            set. If a `pd.Series` is passed in, it will be based on the
2212            corresponding values of the pd.Series.
2213          are_present (Union[LiteralT, pd.Series]): The performed literal of each
2214            interval in the set. If a `pd.Series` is passed in, it will be based on
2215            the corresponding values of the pd.Series.
2216
2217        Returns:
2218          pd.Series: The interval variable set indexed by its corresponding
2219          dimensions.
2220
2221        Raises:
2222          TypeError: if the `index` is invalid (e.g. a `DataFrame`).
2223          ValueError: if the `name` is not a valid identifier or already exists.
2224          ValueError: if the all the indexes do not match.
2225        """
2226        if not isinstance(index, pd.Index):
2227            raise TypeError("Non-index object is used as index")
2228        if not name.isidentifier():
2229            raise ValueError("name={} is not a valid identifier".format(name))
2230
2231        starts = _ConvertToLinearExprSeriesAndValidateIndex(starts, index)
2232        sizes = _ConvertToLinearExprSeriesAndValidateIndex(sizes, index)
2233        ends = _ConvertToLinearExprSeriesAndValidateIndex(ends, index)
2234        are_present = _ConvertToLiteralSeriesAndValidateIndex(are_present, index)
2235
2236        interval_array = []
2237        for i in index:
2238            interval_array.append(
2239                self.NewOptionalIntervalVar(
2240                    start=starts[i],
2241                    size=sizes[i],
2242                    end=ends[i],
2243                    is_present=are_present[i],
2244                    name=f"{name}[{i}]",
2245                )
2246            )
2247        return pd.Series(index=index, data=interval_array)
2248
2249    def NewOptionalFixedSizeIntervalVar(
2250        self,
2251        start: LinearExprT,
2252        size: IntegralT,
2253        is_present: LiteralT,
2254        name: str,
2255    ) -> IntervalVar:
2256        """Creates an interval variable from start, and a fixed size.
2257
2258        An interval variable is a constraint, that is itself used in other
2259        constraints like NoOverlap.
2260
2261        Args:
2262          start: The start of the interval. It can be an affine or constant
2263            expression.
2264          size: The size of the interval. It must be an integer value.
2265          is_present: A literal that indicates if the interval is active or not. A
2266            inactive interval is simply ignored by all constraints.
2267          name: The name of the interval variable.
2268
2269        Returns:
2270          An `IntervalVar` object.
2271        """
2272        size = cmh.assert_is_int64(size)
2273        start_expr = self.ParseLinearExpression(start)
2274        size_expr = self.ParseLinearExpression(size)
2275        end_expr = self.ParseLinearExpression(start + size)
2276        if len(start_expr.vars) > 1:
2277            raise TypeError(
2278                "cp_model.NewIntervalVar: start must be affine or constant."
2279            )
2280        is_present_index = self.GetOrMakeBooleanIndex(is_present)
2281        return IntervalVar(
2282            self.__model,
2283            start_expr,
2284            size_expr,
2285            end_expr,
2286            is_present_index,
2287            name,
2288        )
2289
2290    def NewOptionalFixedSizeIntervalVarSeries(
2291        self,
2292        name: str,
2293        index: pd.Index,
2294        starts: Union[LinearExprT, pd.Series],
2295        sizes: Union[IntegralT, pd.Series],
2296        are_present: Union[LiteralT, pd.Series],
2297    ) -> pd.Series:
2298        """Creates a series of interval variables with the given name.
2299
2300        Args:
2301          name (str): Required. The name of the variable set.
2302          index (pd.Index): Required. The index to use for the variable set.
2303          starts (Union[LinearExprT, pd.Series]): The start of each interval in the
2304            set. If a `pd.Series` is passed in, it will be based on the
2305            corresponding values of the pd.Series.
2306          sizes (Union[IntegralT, pd.Series]): The fixed size of each interval in
2307            the set. If a `pd.Series` is passed in, it will be based on the
2308            corresponding values of the pd.Series.
2309          are_present (Union[LiteralT, pd.Series]): The performed literal of each
2310            interval in the set. If a `pd.Series` is passed in, it will be based on
2311            the corresponding values of the pd.Series.
2312
2313        Returns:
2314          pd.Series: The interval variable set indexed by its corresponding
2315          dimensions.
2316
2317        Raises:
2318          TypeError: if the `index` is invalid (e.g. a `DataFrame`).
2319          ValueError: if the `name` is not a valid identifier or already exists.
2320          ValueError: if the all the indexes do not match.
2321        """
2322        if not isinstance(index, pd.Index):
2323            raise TypeError("Non-index object is used as index")
2324        if not name.isidentifier():
2325            raise ValueError("name={} is not a valid identifier".format(name))
2326
2327        starts = _ConvertToLinearExprSeriesAndValidateIndex(starts, index)
2328        sizes = _ConvertToIntegralSeriesAndValidateIndex(sizes, index)
2329        are_present = _ConvertToLiteralSeriesAndValidateIndex(are_present, index)
2330        interval_array = []
2331        for i in index:
2332            interval_array.append(
2333                self.NewOptionalFixedSizeIntervalVar(
2334                    start=starts[i],
2335                    size=sizes[i],
2336                    is_present=are_present[i],
2337                    name=f"{name}[{i}]",
2338                )
2339            )
2340        return pd.Series(index=index, data=interval_array)
2341
2342    def AddNoOverlap(self, interval_vars: Iterable[IntervalVar]) -> Constraint:
2343        """Adds NoOverlap(interval_vars).
2344
2345        A NoOverlap constraint ensures that all present intervals do not overlap
2346        in time.
2347
2348        Args:
2349          interval_vars: The list of interval variables to constrain.
2350
2351        Returns:
2352          An instance of the `Constraint` class.
2353        """
2354        ct = Constraint(self)
2355        model_ct = self.__model.constraints[ct.Index()]
2356        model_ct.no_overlap.intervals.extend(
2357            [self.GetIntervalIndex(x) for x in interval_vars]
2358        )
2359        return ct
2360
2361    def AddNoOverlap2D(
2362        self,
2363        x_intervals: Iterable[IntervalVar],
2364        y_intervals: Iterable[IntervalVar],
2365    ) -> Constraint:
2366        """Adds NoOverlap2D(x_intervals, y_intervals).
2367
2368        A NoOverlap2D constraint ensures that all present rectangles do not overlap
2369        on a plane. Each rectangle is aligned with the X and Y axis, and is defined
2370        by two intervals which represent its projection onto the X and Y axis.
2371
2372        Furthermore, one box is optional if at least one of the x or y interval is
2373        optional.
2374
2375        Args:
2376          x_intervals: The X coordinates of the rectangles.
2377          y_intervals: The Y coordinates of the rectangles.
2378
2379        Returns:
2380          An instance of the `Constraint` class.
2381        """
2382        ct = Constraint(self)
2383        model_ct = self.__model.constraints[ct.Index()]
2384        model_ct.no_overlap_2d.x_intervals.extend(
2385            [self.GetIntervalIndex(x) for x in x_intervals]
2386        )
2387        model_ct.no_overlap_2d.y_intervals.extend(
2388            [self.GetIntervalIndex(x) for x in y_intervals]
2389        )
2390        return ct
2391
2392    def AddCumulative(
2393        self,
2394        intervals: Iterable[IntervalVar],
2395        demands: Iterable[LinearExprT],
2396        capacity: LinearExprT,
2397    ) -> Constraint:
2398        """Adds Cumulative(intervals, demands, capacity).
2399
2400        This constraint enforces that:
2401
2402            for all t:
2403              sum(demands[i]
2404                if (start(intervals[i]) <= t < end(intervals[i])) and
2405                (intervals[i] is present)) <= capacity
2406
2407        Args:
2408          intervals: The list of intervals.
2409          demands: The list of demands for each interval. Each demand must be >= 0.
2410            Each demand can be an integer value, or an integer variable.
2411          capacity: The maximum capacity of the cumulative constraint. It must be a
2412            positive integer value or variable.
2413
2414        Returns:
2415          An instance of the `Constraint` class.
2416        """
2417        ct = Constraint(self)
2418        model_ct = self.__model.constraints[ct.Index()]
2419        model_ct.cumulative.intervals.extend(
2420            [self.GetIntervalIndex(x) for x in intervals]
2421        )
2422        for d in demands:
2423            model_ct.cumulative.demands.append(self.ParseLinearExpression(d))
2424        model_ct.cumulative.capacity.CopyFrom(self.ParseLinearExpression(capacity))
2425        return ct
2426
2427    # Support for model cloning.
2428    def Clone(self) -> "CpModel":
2429        """Reset the model, and creates a new one from a CpModelProto instance."""
2430        clone = CpModel()
2431        clone.Proto().CopyFrom(self.Proto())
2432        clone.RebuildConstantMap()
2433        return clone
2434
2435    def RebuildConstantMap(self):
2436        """Internal method used during model cloning."""
2437        for i, var in enumerate(self.__model.variables):
2438            if len(var.domain) == 2 and var.domain[0] == var.domain[1]:
2439                self.__constant_map[var.domain[0]] = i
2440
2441    def GetBoolVarFromProtoIndex(self, index: int) -> IntVar:
2442        """Returns an already created Boolean variable from its index."""
2443        if index < 0 or index >= len(self.__model.variables):
2444            raise ValueError(f"GetBoolVarFromProtoIndex: out of bound index {index}")
2445        var = self.__model.variables[index]
2446        if len(var.domain) != 2 or var.domain[0] < 0 or var.domain[1] > 1:
2447            raise ValueError(
2448                f"GetBoolVarFromProtoIndex: index {index} does not reference"
2449                + " a Boolean variable"
2450            )
2451
2452        return IntVar(self.__model, index, None)
2453
2454    def GetIntVarFromProtoIndex(self, index: int) -> IntVar:
2455        """Returns an already created integer variable from its index."""
2456        if index < 0 or index >= len(self.__model.variables):
2457            raise ValueError(f"GetIntVarFromProtoIndex: out of bound index {index}")
2458        return IntVar(self.__model, index, None)
2459
2460    def GetIntervalVarFromProtoIndex(self, index: int) -> IntervalVar:
2461        """Returns an already created interval variable from its index."""
2462        if index < 0 or index >= len(self.__model.constraints):
2463            raise ValueError(
2464                f"GetIntervalVarFromProtoIndex: out of bound index {index}"
2465            )
2466        ct = self.__model.constraints[index]
2467        if not ct.HasField("interval"):
2468            raise ValueError(
2469                f"GetIntervalVarFromProtoIndex: index {index} does not reference an"
2470                + " interval variable"
2471            )
2472
2473        return IntervalVar(self.__model, index, None, None, None, None)
2474
2475    # Helpers.
2476
2477    def __str__(self):
2478        return str(self.__model)
2479
2480    def Proto(self) -> cp_model_pb2.CpModelProto:
2481        """Returns the underlying CpModelProto."""
2482        return self.__model
2483
2484    def Negated(self, index: int) -> int:
2485        return -index - 1
2486
2487    def GetOrMakeIndex(self, arg: VariableT) -> int:
2488        """Returns the index of a variable, its negation, or a number."""
2489        if isinstance(arg, IntVar):
2490            return arg.Index()
2491        elif (
2492            isinstance(arg, _ProductCst)
2493            and isinstance(arg.Expression(), IntVar)
2494            and arg.Coefficient() == -1
2495        ):
2496            return -arg.Expression().Index() - 1
2497        elif cmh.is_integral(arg):
2498            arg = cmh.assert_is_int64(arg)
2499            return self.GetOrMakeIndexFromConstant(arg)
2500        else:
2501            raise TypeError("NotSupported: model.GetOrMakeIndex(" + str(arg) + ")")
2502
2503    def GetOrMakeBooleanIndex(self, arg: LiteralT) -> int:
2504        """Returns an index from a boolean expression."""
2505        if isinstance(arg, IntVar):
2506            self.AssertIsBooleanVariable(arg)
2507            return arg.Index()
2508        elif isinstance(arg, _NotBooleanVariable):
2509            self.AssertIsBooleanVariable(arg.Not())
2510            return arg.Index()
2511        elif cmh.is_integral(arg):
2512            cmh.assert_is_boolean(arg)
2513            return self.GetOrMakeIndexFromConstant(int(arg))
2514        else:
2515            raise TypeError(
2516                "NotSupported: model.GetOrMakeBooleanIndex(" + str(arg) + ")"
2517            )
2518
2519    def GetIntervalIndex(self, arg: IntervalVar) -> int:
2520        if not isinstance(arg, IntervalVar):
2521            raise TypeError("NotSupported: model.GetIntervalIndex(%s)" % arg)
2522        return arg.Index()
2523
2524    def GetOrMakeIndexFromConstant(self, value: IntegralT) -> int:
2525        if value in self.__constant_map:
2526            return self.__constant_map[value]
2527        index = len(self.__model.variables)
2528        self.__model.variables.add(domain=[value, value])
2529        self.__constant_map[value] = index
2530        return index
2531
2532    def VarIndexToVarProto(self, var_index: int) -> cp_model_pb2.IntegerVariableProto:
2533        if var_index >= 0:
2534            return self.__model.variables[var_index]
2535        else:
2536            return self.__model.variables[-var_index - 1]
2537
2538    def ParseLinearExpression(
2539        self, linear_expr: LinearExprT, negate: bool = False
2540    ) -> cp_model_pb2.LinearExpressionProto:
2541        """Returns a LinearExpressionProto built from a LinearExpr instance."""
2542        result: cp_model_pb2.LinearExpressionProto = (
2543            cp_model_pb2.LinearExpressionProto()
2544        )
2545        mult = -1 if negate else 1
2546        if cmh.is_integral(linear_expr):
2547            result.offset = int(linear_expr) * mult
2548            return result
2549
2550        if isinstance(linear_expr, IntVar):
2551            result.vars.append(self.GetOrMakeIndex(linear_expr))
2552            result.coeffs.append(mult)
2553            return result
2554
2555        coeffs_map, constant = cast(LinearExpr, linear_expr).GetIntegerVarValueMap()
2556        result.offset = constant * mult
2557        for t in coeffs_map.items():
2558            if not isinstance(t[0], IntVar):
2559                raise TypeError("Wrong argument" + str(t))
2560            c = cmh.assert_is_int64(t[1])
2561            result.vars.append(t[0].Index())
2562            result.coeffs.append(c * mult)
2563        return result
2564
2565    def _SetObjective(self, obj: ObjLinearExprT, minimize: bool):
2566        """Sets the objective of the model."""
2567        self.ClearObjective()
2568        if isinstance(obj, IntVar):
2569            self.__model.objective.coeffs.append(1)
2570            self.__model.objective.offset = 0
2571            if minimize:
2572                self.__model.objective.vars.append(obj.Index())
2573                self.__model.objective.scaling_factor = 1
2574            else:
2575                self.__model.objective.vars.append(self.Negated(obj.Index()))
2576                self.__model.objective.scaling_factor = -1
2577        elif isinstance(obj, LinearExpr):
2578            coeffs_map, constant, is_integer = obj.GetFloatVarValueMap()
2579            if is_integer:
2580                if minimize:
2581                    self.__model.objective.scaling_factor = 1
2582                    self.__model.objective.offset = constant
2583                else:
2584                    self.__model.objective.scaling_factor = -1
2585                    self.__model.objective.offset = -constant
2586                for v, c in coeffs_map.items():
2587                    self.__model.objective.coeffs.append(c)
2588                    if minimize:
2589                        self.__model.objective.vars.append(v.Index())
2590                    else:
2591                        self.__model.objective.vars.append(self.Negated(v.Index()))
2592            else:
2593                self.__model.floating_point_objective.maximize = not minimize
2594                self.__model.floating_point_objective.offset = constant
2595                for v, c in coeffs_map.items():
2596                    self.__model.floating_point_objective.coeffs.append(c)
2597                    self.__model.floating_point_objective.vars.append(v.Index())
2598        elif cmh.is_integral(obj):
2599            self.__model.objective.offset = int(obj)
2600            self.__model.objective.scaling_factor = 1
2601        else:
2602            raise TypeError("TypeError: " + str(obj) + " is not a valid objective")
2603
2604    def Minimize(self, obj: ObjLinearExprT):
2605        """Sets the objective of the model to minimize(obj)."""
2606        self._SetObjective(obj, minimize=True)
2607
2608    def Maximize(self, obj: ObjLinearExprT):
2609        """Sets the objective of the model to maximize(obj)."""
2610        self._SetObjective(obj, minimize=False)
2611
2612    def HasObjective(self) -> bool:
2613        return self.__model.HasField("objective") or self.__model.HasField(
2614            "floating_point_objective"
2615        )
2616
2617    def ClearObjective(self):
2618        self.__model.ClearField("objective")
2619        self.__model.ClearField("floating_point_objective")
2620
2621    def AddDecisionStrategy(
2622        self,
2623        variables: Sequence[IntVar],
2624        var_strategy: cp_model_pb2.DecisionStrategyProto.VariableSelectionStrategy,
2625        domain_strategy: cp_model_pb2.DecisionStrategyProto.DomainReductionStrategy,
2626    ) -> None:
2627        """Adds a search strategy to the model.
2628
2629        Args:
2630          variables: a list of variables this strategy will assign.
2631          var_strategy: heuristic to choose the next variable to assign.
2632          domain_strategy: heuristic to reduce the domain of the selected variable.
2633            Currently, this is advanced code: the union of all strategies added to
2634            the model must be complete, i.e. instantiates all variables. Otherwise,
2635            Solve() will fail.
2636        """
2637
2638        strategy = self.__model.search_strategy.add()
2639        for v in variables:
2640            strategy.variables.append(v.Index())
2641        strategy.variable_selection_strategy = var_strategy
2642        strategy.domain_reduction_strategy = domain_strategy
2643
2644    def ModelStats(self) -> str:
2645        """Returns a string containing some model statistics."""
2646        return swig_helper.CpSatHelper.ModelStats(self.__model)
2647
2648    def Validate(self) -> str:
2649        """Returns a string indicating that the model is invalid."""
2650        return swig_helper.CpSatHelper.ValidateModel(self.__model)
2651
2652    def ExportToFile(self, file: str) -> bool:
2653        """Write the model as a protocol buffer to 'file'.
2654
2655        Args:
2656          file: file to write the model to. If the filename ends with 'txt', the
2657            model will be written as a text file, otherwise, the binary format will
2658            be used.
2659
2660        Returns:
2661          True if the model was correctly written.
2662        """
2663        return swig_helper.CpSatHelper.WriteModelToFile(self.__model, file)
2664
2665    def AssertIsBooleanVariable(self, x: LiteralT) -> None:
2666        if isinstance(x, IntVar):
2667            var = self.__model.variables[x.Index()]
2668            if len(var.domain) != 2 or var.domain[0] < 0 or var.domain[1] > 1:
2669                raise TypeError("TypeError: " + str(x) + " is not a boolean variable")
2670        elif not isinstance(x, _NotBooleanVariable):
2671            raise TypeError("TypeError: " + str(x) + " is not a boolean variable")
2672
2673    def AddHint(self, var: IntVar, value: int) -> None:
2674        """Adds 'var == value' as a hint to the solver."""
2675        self.__model.solution_hint.vars.append(self.GetOrMakeIndex(var))
2676        self.__model.solution_hint.values.append(value)
2677
2678    def ClearHints(self):
2679        """Remove any solution hint from the model."""
2680        self.__model.ClearField("solution_hint")
2681
2682    def AddAssumption(self, lit: LiteralT) -> None:
2683        """Add the literal 'lit' to the model as assumptions."""
2684        self.__model.assumptions.append(self.GetOrMakeBooleanIndex(lit))
2685
2686    def AddAssumptions(self, literals: Iterable[LiteralT]) -> None:
2687        """Add the literals to the model as assumptions."""
2688        for lit in literals:
2689            self.AddAssumption(lit)
2690
2691    def ClearAssumptions(self) -> None:
2692        """Remove all assumptions from the model."""
2693        self.__model.ClearField("assumptions")
2694
2695
2696@overload
2697def ExpandGeneratorOrTuple(
2698    args: Union[Tuple[LiteralT, ...], Iterable[LiteralT]]
2699) -> Union[Iterable[LiteralT], LiteralT]:
2700    ...
2701
2702
2703@overload
2704def ExpandGeneratorOrTuple(
2705    args: Union[Tuple[LinearExprT, ...], Iterable[LinearExprT]]
2706) -> Union[Iterable[LinearExprT], LinearExprT]:
2707    ...
2708
2709
2710def ExpandGeneratorOrTuple(args):
2711    if hasattr(args, "__len__"):  # Tuple
2712        if len(args) != 1:
2713            return args
2714        if cmh.is_a_number(args[0]) or isinstance(args[0], LinearExpr):
2715            return args
2716    # Generator
2717    return args[0]
2718
2719
2720def EvaluateLinearExpr(
2721    expression: LinearExprT, solution: cp_model_pb2.CpSolverResponse
2722) -> int:
2723    """Evaluate a linear expression against a solution."""
2724    if cmh.is_integral(expression):
2725        return int(expression)
2726    if not isinstance(expression, LinearExpr):
2727        raise TypeError("Cannot interpret %s as a linear expression." % expression)
2728
2729    value = 0
2730    to_process = [(expression, 1)]
2731    while to_process:
2732        expr, coeff = to_process.pop()
2733        if cmh.is_integral(expr):
2734            value += int(expr) * coeff
2735        elif isinstance(expr, _ProductCst):
2736            to_process.append((expr.Expression(), coeff * expr.Coefficient()))
2737        elif isinstance(expr, _Sum):
2738            to_process.append((expr.Left(), coeff))
2739            to_process.append((expr.Right(), coeff))
2740        elif isinstance(expr, _SumArray):
2741            for e in expr.Expressions():
2742                to_process.append((e, coeff))
2743            value += expr.Constant() * coeff
2744        elif isinstance(expr, _WeightedSum):
2745            for e, c in zip(expr.Expressions(), expr.Coefficients()):
2746                to_process.append((e, coeff * c))
2747            value += expr.Constant() * coeff
2748        elif isinstance(expr, IntVar):
2749            value += coeff * solution.solution[expr.Index()]
2750        elif isinstance(expr, _NotBooleanVariable):
2751            value += coeff * (1 - solution.solution[expr.Not().Index()])
2752        else:
2753            raise TypeError(f"Cannot interpret {expr} as a linear expression.")
2754
2755    return value
2756
2757
2758def EvaluateBooleanExpression(
2759    literal: LiteralT, solution: cp_model_pb2.CpSolverResponse
2760) -> bool:
2761    """Evaluate a boolean expression against a solution."""
2762    if cmh.is_integral(literal):
2763        return bool(literal)
2764    elif isinstance(literal, IntVar) or isinstance(literal, _NotBooleanVariable):
2765        index: int = cast(Union[IntVar, _NotBooleanVariable], literal).Index()
2766        if index >= 0:
2767            return bool(solution.solution[index])
2768        else:
2769            return not solution.solution[-index - 1]
2770    else:
2771        raise TypeError(f"Cannot interpret {literal} as a boolean expression.")
2772
2773
2774class CpSolver:
2775    """Main solver class.
2776
2777    The purpose of this class is to search for a solution to the model provided
2778    to the Solve() method.
2779
2780    Once Solve() is called, this class allows inspecting the solution found
2781    with the Value() and BooleanValue() methods, as well as general statistics
2782    about the solve procedure.
2783    """
2784
2785    def __init__(self):
2786        self.__solution: Optional[cp_model_pb2.CpSolverResponse] = None
2787        self.parameters: sat_parameters_pb2.SatParameters = (
2788            sat_parameters_pb2.SatParameters()
2789        )
2790        self.log_callback: Optional[swig_helper.LogCallback] = None
2791        self.__solve_wrapper: Optional[swig_helper.SolveWrapper] = None
2792        self.__lock: threading.Lock = threading.Lock()
2793
2794    def Solve(
2795        self,
2796        model: CpModel,
2797        solution_callback: Optional["CpSolverSolutionCallback"] = None,
2798    ) -> cp_model_pb2.CpSolverStatus:
2799        """Solves a problem and passes each solution to the callback if not null."""
2800        with self.__lock:
2801            self.__solve_wrapper = swig_helper.SolveWrapper()
2802
2803        self.__solve_wrapper.SetParameters(self.parameters)
2804        if solution_callback is not None:
2805            self.__solve_wrapper.AddSolutionCallback(solution_callback)
2806
2807        if self.log_callback is not None:
2808            self.__solve_wrapper.AddLogCallback(self.log_callback)
2809
2810        self.__solution = self.__solve_wrapper.Solve(model.Proto())
2811
2812        if solution_callback is not None:
2813            self.__solve_wrapper.ClearSolutionCallback(solution_callback)
2814
2815        with self.__lock:
2816            self.__solve_wrapper = None
2817
2818        return self.__solution.status
2819
2820    def SolveWithSolutionCallback(
2821        self, model: CpModel, callback: "CpSolverSolutionCallback"
2822    ) -> cp_model_pb2.CpSolverStatus:
2823        """DEPRECATED Use Solve() with the callback argument."""
2824        warnings.warn(
2825            "SolveWithSolutionCallback is deprecated; use Solve() with"
2826            + "the callback argument.",
2827            DeprecationWarning,
2828        )
2829        return self.Solve(model, callback)
2830
2831    def SearchForAllSolutions(
2832        self, model: CpModel, callback: "CpSolverSolutionCallback"
2833    ) -> cp_model_pb2.CpSolverStatus:
2834        """DEPRECATED Use Solve() with the right parameter.
2835
2836        Search for all solutions of a satisfiability problem.
2837
2838        This method searches for all feasible solutions of a given model.
2839        Then it feeds the solution to the callback.
2840
2841        Note that the model cannot contain an objective.
2842
2843        Args:
2844          model: The model to solve.
2845          callback: The callback that will be called at each solution.
2846
2847        Returns:
2848          The status of the solve:
2849
2850          * *FEASIBLE* if some solutions have been found
2851          * *INFEASIBLE* if the solver has proved there are no solution
2852          * *OPTIMAL* if all solutions have been found
2853        """
2854        warnings.warn(
2855            "SearchForAllSolutions is deprecated; use Solve() with"
2856            + "enumerate_all_solutions = True.",
2857            DeprecationWarning,
2858        )
2859        if model.HasObjective():
2860            raise TypeError(
2861                "Search for all solutions is only defined on satisfiability problems"
2862            )
2863        # Store old parameter.
2864        enumerate_all = self.parameters.enumerate_all_solutions
2865        self.parameters.enumerate_all_solutions = True
2866
2867        self.Solve(model, callback)
2868
2869        # Restore parameter.
2870        self.parameters.enumerate_all_solutions = enumerate_all
2871        return self.__solution.status
2872
2873    def StopSearch(self) -> None:
2874        """Stops the current search asynchronously."""
2875        with self.__lock:
2876            if self.__solve_wrapper:
2877                self.__solve_wrapper.StopSearch()
2878
2879    def _Solution(self) -> cp_model_pb2.CpSolverResponse:
2880        """Checks Solve() has been called, and returns the solution."""
2881        if self.__solution is None:
2882            raise RuntimeError("Solve() has not been called.")
2883        return self.__solution
2884
2885    def Value(self, expression: LinearExprT) -> int:
2886        """Returns the value of a linear expression after solve."""
2887        return EvaluateLinearExpr(expression, self._Solution())
2888
2889    def Values(self, variables: _IndexOrSeries) -> pd.Series:
2890        """Returns the values of the input variables.
2891
2892        If `variables` is a `pd.Index`, then the output will be indexed by the
2893        variables. If `variables` is a `pd.Series` indexed by the underlying
2894        dimensions, then the output will be indexed by the same underlying
2895        dimensions.
2896
2897        Args:
2898          variables (Union[pd.Index, pd.Series]): The set of variables from which to
2899            get the values.
2900
2901        Returns:
2902          pd.Series: The values of all variables in the set.
2903        """
2904        solution = self._Solution()
2905        return _AttributeSeries(
2906            func=lambda v: solution.solution[v.Index()],
2907            values=variables,
2908        )
2909
2910    def BooleanValue(self, literal: LiteralT) -> bool:
2911        """Returns the boolean value of a literal after solve."""
2912        return EvaluateBooleanExpression(literal, self._Solution())
2913
2914    def BooleanValues(self, variables: _IndexOrSeries) -> pd.Series:
2915        """Returns the values of the input variables.
2916
2917        If `variables` is a `pd.Index`, then the output will be indexed by the
2918        variables. If `variables` is a `pd.Series` indexed by the underlying
2919        dimensions, then the output will be indexed by the same underlying
2920        dimensions.
2921
2922        Args:
2923          variables (Union[pd.Index, pd.Series]): The set of variables from which to
2924            get the values.
2925
2926        Returns:
2927          pd.Series: The values of all variables in the set.
2928        """
2929        solution = self._Solution()
2930        return _AttributeSeries(
2931            func=lambda literal: EvaluateBooleanExpression(literal, solution),
2932            values=variables,
2933        )
2934
2935    def ObjectiveValue(self) -> float:
2936        """Returns the value of the objective after solve."""
2937        return self._Solution().objective_value
2938
2939    def BestObjectiveBound(self) -> float:
2940        """Returns the best lower (upper) bound found when min(max)imizing."""
2941        return self._Solution().best_objective_bound
2942
2943    def StatusName(self, status: ... = None) -> str:
2944        """Returns the name of the status returned by Solve()."""
2945        if status is None:
2946            status = self._Solution().status
2947        return cp_model_pb2.CpSolverStatus.Name(status)
2948
2949    def NumBooleans(self) -> int:
2950        """Returns the number of boolean variables managed by the SAT solver."""
2951        return self._Solution().num_booleans
2952
2953    def NumConflicts(self) -> int:
2954        """Returns the number of conflicts since the creation of the solver."""
2955        return self._Solution().num_conflicts
2956
2957    def NumBranches(self) -> int:
2958        """Returns the number of search branches explored by the solver."""
2959        return self._Solution().num_branches
2960
2961    def WallTime(self) -> float:
2962        """Returns the wall time in seconds since the creation of the solver."""
2963        return self._Solution().wall_time
2964
2965    def UserTime(self) -> float:
2966        """Returns the user time in seconds since the creation of the solver."""
2967        return self._Solution().user_time
2968
2969    def ResponseStats(self) -> str:
2970        """Returns some statistics on the solution found as a string."""
2971        return swig_helper.CpSatHelper.SolverResponseStats(self._Solution())
2972
2973    def ResponseProto(self) -> cp_model_pb2.CpSolverResponse:
2974        """Returns the response object."""
2975        return self._Solution()
2976
2977    def SufficientAssumptionsForInfeasibility(self) -> Sequence[int]:
2978        """Returns the indices of the infeasible assumptions."""
2979        return self._Solution().sufficient_assumptions_for_infeasibility
2980
2981    def SolutionInfo(self) -> str:
2982        """Returns some information on the solve process.
2983
2984        Returns some information on how the solution was found, or the reason
2985        why the model or the parameters are invalid.
2986
2987        Raises:
2988          RuntimeError: if Solve() has not been called.
2989        """
2990        return self._Solution().solution_info
2991
2992
2993class CpSolverSolutionCallback(swig_helper.SolutionCallback):
2994    """Solution callback.
2995
2996    This class implements a callback that will be called at each new solution
2997    found during search.
2998
2999    The method OnSolutionCallback() will be called by the solver, and must be
3000    implemented. The current solution can be queried using the BooleanValue()
3001    and Value() methods.
3002
3003    It inherits the following methods from its base class:
3004
3005    * `ObjectiveValue(self)`
3006    * `BestObjectiveBound(self)`
3007    * `NumBooleans(self)`
3008    * `NumConflicts(self)`
3009    * `NumBranches(self)`
3010    * `WallTime(self)`
3011    * `UserTime(self)`
3012
3013    These methods returns the same information as their counterpart in the
3014    `CpSolver` class.
3015    """
3016
3017    def __init__(self):
3018        swig_helper.SolutionCallback.__init__(self)
3019
3020    def OnSolutionCallback(self) -> None:
3021        """Proxy for the same method in snake case."""
3022        self.on_solution_callback()
3023
3024    def BooleanValue(self, lit: LiteralT) -> bool:
3025        """Returns the boolean value of a boolean literal.
3026
3027        Args:
3028            lit: A boolean variable or its negation.
3029
3030        Returns:
3031            The Boolean value of the literal in the solution.
3032
3033        Raises:
3034            RuntimeError: if `lit` is not a boolean variable or its negation.
3035        """
3036        if not self.HasResponse():
3037            raise RuntimeError("Solve() has not been called.")
3038        if cmh.is_integral(lit):
3039            return bool(lit)
3040        elif isinstance(lit, IntVar) or isinstance(lit, _NotBooleanVariable):
3041            return self.SolutionBooleanValue(
3042                cast(Union[IntVar, _NotBooleanVariable], lit).Index()
3043            )
3044        raise TypeError(f"Cannot interpret {lit} as a boolean expression.")
3045
3046    def Value(self, expression: LinearExprT) -> int:
3047        """Evaluates an linear expression in the current solution.
3048
3049        Args:
3050            expression: a linear expression of the model.
3051
3052        Returns:
3053            An integer value equal to the evaluation of the linear expression
3054            against the current solution.
3055
3056        Raises:
3057            RuntimeError: if 'expression' is not a LinearExpr.
3058        """
3059        if not self.HasResponse():
3060            raise RuntimeError("Solve() has not been called.")
3061
3062        value = 0
3063        to_process = [(expression, 1)]
3064        while to_process:
3065            expr, coeff = to_process.pop()
3066            if cmh.is_integral(expr):
3067                value += int(expr) * coeff
3068            elif isinstance(expr, _ProductCst):
3069                to_process.append((expr.Expression(), coeff * expr.Coefficient()))
3070            elif isinstance(expr, _Sum):
3071                to_process.append((expr.Left(), coeff))
3072                to_process.append((expr.Right(), coeff))
3073            elif isinstance(expr, _SumArray):
3074                for e in expr.Expressions():
3075                    to_process.append((e, coeff))
3076                    value += expr.Constant() * coeff
3077            elif isinstance(expr, _WeightedSum):
3078                for e, c in zip(expr.Expressions(), expr.Coefficients()):
3079                    to_process.append((e, coeff * c))
3080                value += expr.Constant() * coeff
3081            elif isinstance(expr, IntVar):
3082                value += coeff * self.SolutionIntegerValue(expr.Index())
3083            elif isinstance(expr, _NotBooleanVariable):
3084                value += coeff * (1 - self.SolutionIntegerValue(expr.Not().Index()))
3085            else:
3086                raise TypeError(
3087                    f"Cannot interpret {expression} as a linear expression."
3088                )
3089
3090        return value
3091
3092    def Response(self) -> cp_model_pb2.CpSolverResponse:
3093        """Returns the current solution response."""
3094        return swig_helper.SolutionCallback.Response(self)
3095
3096
3097class ObjectiveSolutionPrinter(CpSolverSolutionCallback):
3098    """Display the objective value and time of intermediate solutions."""
3099
3100    def __init__(self):
3101        CpSolverSolutionCallback.__init__(self)
3102        self.__solution_count = 0
3103        self.__start_time = time.time()
3104
3105    def on_solution_callback(self) -> None:
3106        """Called on each new solution."""
3107        current_time = time.time()
3108        obj = self.ObjectiveValue()
3109        print(
3110            "Solution %i, time = %0.2f s, objective = %i"
3111            % (self.__solution_count, current_time - self.__start_time, obj)
3112        )
3113        self.__solution_count += 1
3114
3115    def solution_count(self) -> int:
3116        """Returns the number of solutions found."""
3117        return self.__solution_count
3118
3119
3120class VarArrayAndObjectiveSolutionPrinter(CpSolverSolutionCallback):
3121    """Print intermediate solutions (objective, variable values, time)."""
3122
3123    def __init__(self, variables):
3124        CpSolverSolutionCallback.__init__(self)
3125        self.__variables: Sequence[IntVar] = variables
3126        self.__solution_count: int = 0
3127        self.__start_time: float = time.time()
3128
3129    def on_solution_callback(self) -> None:
3130        """Called on each new solution."""
3131        current_time = time.time()
3132        obj = self.ObjectiveValue()
3133        print(
3134            "Solution %i, time = %0.2f s, objective = %i"
3135            % (self.__solution_count, current_time - self.__start_time, obj)
3136        )
3137        for v in self.__variables:
3138            print("  %s = %i" % (v, self.Value(v)), end=" ")
3139        print()
3140        self.__solution_count += 1
3141
3142    def solution_count(self) -> int:
3143        """Returns the number of solutions found."""
3144        return self.__solution_count
3145
3146
3147class VarArraySolutionPrinter(CpSolverSolutionCallback):
3148    """Print intermediate solutions (variable values, time)."""
3149
3150    def __init__(self, variables: Sequence[IntVar]):
3151        CpSolverSolutionCallback.__init__(self)
3152        self.__variables: Sequence[IntVar] = variables
3153        self.__solution_count: int = 0
3154        self.__start_time: float = time.time()
3155
3156    def on_solution_callback(self) -> None:
3157        """Called on each new solution."""
3158        current_time = time.time()
3159        print(
3160            "Solution %i, time = %0.2f s"
3161            % (self.__solution_count, current_time - self.__start_time)
3162        )
3163        for v in self.__variables:
3164            print("  %s = %i" % (v, self.Value(v)), end=" ")
3165        print()
3166        self.__solution_count += 1
3167
3168    def solution_count(self) -> int:
3169        """Returns the number of solutions found."""
3170        return self.__solution_count
3171
3172
3173def _GetIndex(obj: _IndexOrSeries) -> pd.Index:
3174    """Returns the indices of `obj` as a `pd.Index`."""
3175    if isinstance(obj, pd.Series):
3176        return obj.index
3177    return obj
3178
3179
3180def _AttributeSeries(
3181    *,
3182    func: Callable[[IntVar], IntegralT],
3183    values: _IndexOrSeries,
3184) -> pd.Series:
3185    """Returns the attributes of `values`.
3186
3187    Args:
3188      func: The function to call for getting the attribute data.
3189      values: The values that the function will be applied (element-wise) to.
3190
3191    Returns:
3192      pd.Series: The attribute values.
3193    """
3194    return pd.Series(
3195        data=[func(v) for v in values],
3196        index=_GetIndex(values),
3197    )
3198
3199
3200def _ConvertToIntegralSeriesAndValidateIndex(
3201    value_or_series: Union[IntegralT, pd.Series], index: pd.Index
3202) -> pd.Series:
3203    """Returns a pd.Series of the given index with the corresponding values.
3204
3205    Args:
3206      value_or_series: the values to be converted (if applicable).
3207      index: the index of the resulting pd.Series.
3208
3209    Returns:
3210      pd.Series: The set of values with the given index.
3211
3212    Raises:
3213      TypeError: If the type of `value_or_series` is not recognized.
3214      ValueError: If the index does not match.
3215    """
3216    if cmh.is_integral(value_or_series):
3217        result = pd.Series(data=value_or_series, index=index)
3218    elif isinstance(value_or_series, pd.Series):
3219        if value_or_series.index.equals(index):
3220            result = value_or_series
3221        else:
3222            raise ValueError("index does not match")
3223    else:
3224        raise TypeError("invalid type={}".format(type(value_or_series)))
3225    return result
3226
3227
3228def _ConvertToLinearExprSeriesAndValidateIndex(
3229    value_or_series: Union[LinearExprT, pd.Series], index: pd.Index
3230) -> pd.Series:
3231    """Returns a pd.Series of the given index with the corresponding values.
3232
3233    Args:
3234      value_or_series: the values to be converted (if applicable).
3235      index: the index of the resulting pd.Series.
3236
3237    Returns:
3238      pd.Series: The set of values with the given index.
3239
3240    Raises:
3241      TypeError: If the type of `value_or_series` is not recognized.
3242      ValueError: If the index does not match.
3243    """
3244    if cmh.is_integral(value_or_series):
3245        result = pd.Series(data=value_or_series, index=index)
3246    elif isinstance(value_or_series, pd.Series):
3247        if value_or_series.index.equals(index):
3248            result = value_or_series
3249        else:
3250            raise ValueError("index does not match")
3251    else:
3252        raise TypeError("invalid type={}".format(type(value_or_series)))
3253    return result
3254
3255
3256def _ConvertToLiteralSeriesAndValidateIndex(
3257    value_or_series: Union[LiteralT, pd.Series], index: pd.Index
3258) -> pd.Series:
3259    """Returns a pd.Series of the given index with the corresponding values.
3260
3261    Args:
3262      value_or_series: the values to be converted (if applicable).
3263      index: the index of the resulting pd.Series.
3264
3265    Returns:
3266      pd.Series: The set of values with the given index.
3267
3268    Raises:
3269      TypeError: If the type of `value_or_series` is not recognized.
3270      ValueError: If the index does not match.
3271    """
3272    if cmh.is_integral(value_or_series):
3273        result = pd.Series(data=value_or_series, index=index)
3274    elif isinstance(value_or_series, pd.Series):
3275        if value_or_series.index.equals(index):
3276            result = value_or_series
3277        else:
3278            raise ValueError("index does not match")
3279    else:
3280        raise TypeError("invalid type={}".format(type(value_or_series)))
3281    return result