ortools.sat.python.cp_model

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