// Copyright 2010-2025 Google LLC
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef ORTOOLS_SAT_WORK_ASSIGNMENT_H_
#define ORTOOLS_SAT_WORK_ASSIGNMENT_H_

#include <stdint.h>
#include <sys/stat.h>

#include <array>
#include <deque>
#include <functional>
#include <limits>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>

#include "absl/base/thread_annotations.h"
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/container/node_hash_map.h"
#include "absl/functional/function_ref.h"
#include "absl/log/check.h"
#include "absl/synchronization/mutex.h"
#include "absl/types/span.h"
#include "ortools/sat/clause.h"
#include "ortools/sat/cp_model_mapping.h"
#include "ortools/sat/cp_model_utils.h"
#include "ortools/sat/integer.h"
#include "ortools/sat/integer_base.h"
#include "ortools/sat/integer_search.h"
#include "ortools/sat/lrat_proof_handler.h"
#include "ortools/sat/model.h"
#include "ortools/sat/restart.h"
#include "ortools/sat/sat_base.h"
#include "ortools/sat/sat_decision.h"
#include "ortools/sat/sat_parameters.pb.h"
#include "ortools/sat/sat_solver.h"
#include "ortools/sat/synchronization.h"
#include "ortools/sat/util.h"
#include "ortools/util/running_stat.h"
#include "ortools/util/strong_integers.h"
#include "ortools/util/time_limit.h"

namespace operations_research::sat {

class ProtoLiteral {
 public:
  ProtoLiteral() = default;
  ProtoLiteral(int var, IntegerValue lb) : proto_var_(var), lb_(lb) {}
  ProtoLiteral Negated() const {
    return ProtoLiteral(NegatedRef(proto_var_), -lb_ + 1);
  }
  int proto_var() const { return proto_var_; }
  IntegerValue lb() const { return lb_; }
  bool operator==(const ProtoLiteral& other) const {
    return proto_var_ == other.proto_var_ && lb_ == other.lb_;
  }
  bool operator!=(const ProtoLiteral& other) const { return !(*this == other); }
  template <typename H>
  friend H AbslHashValue(H h, const ProtoLiteral& literal) {
    return H::combine(std::move(h), literal.proto_var_, literal.lb_);
  }

  // Note you should only decode integer literals at the root level.
  Literal Decode(CpModelMapping*, IntegerEncoder*) const;

  // Encodes a literal as a ProtoLiteral. This can encode literals that occur in
  // the proto model, and also integer bounds literals.
  static std::optional<ProtoLiteral> Encode(Literal, CpModelMapping*,
                                            IntegerEncoder*);

  // As above, but will only encode literals that are boolean variables or their
  // negations (i.e. not integer bounds literals).
  static std::optional<ProtoLiteral> EncodeLiteral(Literal, CpModelMapping*);

 private:
  IntegerLiteral DecodeInteger(CpModelMapping*) const;
  static std::optional<ProtoLiteral> EncodeInteger(IntegerLiteral,
                                                   CpModelMapping*);

  int proto_var_ = std::numeric_limits<int>::max();
  IntegerValue lb_ = kMaxIntegerValue;
};

// ProtoTrail acts as an intermediate datastructure that can be synced
// with the shared tree and the local Trail as appropriate.
// It's intended that you sync a ProtoTrail with the tree on restart or when
// a subtree is closed, and with the local trail after propagation.
// Specifically it stores objective lower bounds, and literals that have been
// branched on and later proven to be implied by the prior decisions (i.e. they
// can be enqueued at this level).
// TODO(user): It'd be good to store an earlier level at which
// implications may be propagated.
class ProtoTrail {
 public:
  ProtoTrail();

  // Adds a new assigned level to the trail.
  void PushLevel(const ProtoLiteral& decision, IntegerValue objective_lb,
                 int node_id);

  // Asserts that the decision at `level` is implied by earlier decisions.
  void SetLevelImplied(int level);

  // Clear the trail, removing all levels.
  void Clear();

  // Set a lower bound on the objective at level.
  void SetObjectiveLb(int level, IntegerValue objective_lb);

  // Returns the maximum decision level stored in the trail.
  int MaxLevel() const { return decision_indexes_.size(); }

  // Returns the decision assigned at `level`.
  ProtoLiteral Decision(int level) const {
    CHECK_GE(level, 1);
    CHECK_LE(level, decision_indexes_.size());
    return literals_[decision_indexes_[level - 1]];
  }

  // Returns the node ID for the decision at `level`.
  int DecisionNodeId(int level) const;

  // Returns the node ids for decisions and implications at `level`.
  absl::Span<const int> NodeIds(int level) const;

  // Returns literals which may be propagated at `level`, this does not include
  // the decision.
  absl::Span<const ProtoLiteral> Implications(int level) const;

  // Adds a literal which is implied by the decisions from level 1 to `level`.
  // The caller must add a corresponding LRAT inferred clause (if LRAT is
  // enabled). This implication can then be used by other workers as an LRAT
  // imported clause, without proof.
  bool AddImplication(int level, ProtoLiteral implication) {
    auto it = assigned_at_level_.find(implication);
    if (it != assigned_at_level_.end() && it->second <= level) return false;
    MutableImplications(level).push_back(implication);
    assigned_at_level_[implication] = level;
    return true;
  }

  // Removes implications that are already assigned at an earlier level, as well
  // as duplicate implications at the same level.
  void NormalizeImplications();

  IntegerValue ObjectiveLb(int level) const {
    CHECK_GE(level, 1);
    return level_to_objective_lbs_[level - 1];
  }

  absl::Span<const ProtoLiteral> Literals() const { return literals_; }

  const std::vector<ProtoLiteral>& TargetPhase() const { return target_phase_; }

  // Returns the target phase and clears it.
  std::vector<ProtoLiteral> TakeTargetPhase() {
    return std::move(target_phase_);
  }
  void ClearTargetPhase() { target_phase_.clear(); }
  // Appends a literal to the target phase, returns false if the phase is full.
  bool AddPhase(const ProtoLiteral& lit) {
    if (target_phase_.size() >= kMaxPhaseSize) return false;
    if (!IsAssigned(lit)) {
      target_phase_.push_back(lit);
    }
    return true;
  }
  void SetTargetPhase(std::vector<ProtoLiteral> phase) {
    target_phase_ = std::move(phase);
  }
  bool IsAssigned(const ProtoLiteral& lit) const {
    return assigned_at_level_.contains(lit) ||
           assigned_at_level_.contains(lit.Negated());
  }

 private:
  // Store up to 4 KiB of literals in the target phase.
  static constexpr int kMaxPhaseSize = 256;

  std::vector<ProtoLiteral>& MutableImplications(int level) {
    return implications_[level - 1];
  }
  // Parallel vectors encoding the literals and node ids on the trail.
  std::vector<ProtoLiteral> literals_;
  std::vector<int> node_ids_;

  // Extra implications that can be propagated at each level but were never
  // branches in the shared tree.
  std::vector<std::vector<ProtoLiteral>> implications_;
  absl::flat_hash_map<ProtoLiteral, int> assigned_at_level_;

  // The index in the literals_/node_ids_ vectors for the start of each level.
  std::vector<int> decision_indexes_;

  // The objective lower bound of each level.
  std::vector<IntegerValue> level_to_objective_lbs_;

  std::vector<ProtoLiteral> target_phase_;
};

// Experimental thread-safe class for managing work assignments between workers.
// This API is intended to investigate Graeme Gange & Peter Stuckey's proposal
// "Scalable Parallelization of Learning Solvers".
// The core idea of this implementation is that workers can maintain several
// ProtoTrails, and periodically replace the "worst" one.
// With 1 assignment per worker, this leads to something very similar to
// Embarassingly Parallel Search.
class SharedTreeManager {
 public:
  explicit SharedTreeManager(Model* model);
  SharedTreeManager(const SharedTreeManager&) = delete;

  int NumWorkers() const { return num_workers_; }
  int NumNodes() const ABSL_LOCKS_EXCLUDED(mu_);
  int MaxPathDepth() const { return max_path_depth_; }

  // Syncs the state of path with the shared search tree.
  // Clears `path` and returns false if the assigned subtree is closed or a
  // restart has invalidated the path.
  bool SyncTree(ProtoTrail& path) ABSL_LOCKS_EXCLUDED(mu_);

  // Assigns a path prefix that the worker should explore.
  void ReplaceTree(ProtoTrail& path);

  // Asserts that the subtree in path up to `level` contains no improving
  // solutions. Clears path. The caller must add a corresponding LRAT inferred
  // clause first (stating that the negation of the decision at `level` is
  // implied by the previous decisions).
  void CloseTree(ProtoTrail& path, int level);

  // Attempts to split the tree repeatedly with the given decisions.
  // `path` will be extended with the accepted splits, the opposite branches
  // will be added as unassigned leaves.
  // Returns the number of splits accepted.
  int TrySplitTree(absl::Span<const ProtoLiteral> decisions, ProtoTrail& path)
      ABSL_LOCKS_EXCLUDED(mu_);

  void Restart() {
    absl::MutexLock l(mu_);
    RestartLockHeld();
  }

  void CloseLratProof();

 private:
  // Because it is quite difficult to get a flat_hash_map to release memory,
  // we store info we need only for open nodes implications via a unique_ptr.
  // Note to simplify code, the root will always have a NodeTrailInfo after it
  // is closed.
  struct NodeTrailInfo {
    // A map from literal to the best lower bound proven at this node, and to
    // the LRAT clause "decisions of node and its parents => literal" (with
    // literals of implied nodes filtered out).
    absl::flat_hash_map<int, std::pair<IntegerValue, ClauseId>> implications;
    // This is only non-empty for nodes where all but one descendent is closed
    // (i.e. mostly leaves).
    std::vector<ProtoLiteral> phase;
  };
  struct Node {
    ProtoLiteral decision;
    IntegerValue objective_lb = kMinIntegerValue;
    Node* parent = nullptr;
    std::array<Node*, 2> children = {nullptr, nullptr};
    // A node's id is related to its index in `nodes_`, see `node_id_offset_`.
    int id;
    // For each closed node, closing_clause_id is an LRAT clause stating that
    // all the parent decisions imply the negation of the node's decision.
    bool closed = false;
    ClauseId closing_clause_id = kNoClauseId;
    // A node is implied if its sibling is closed. The closing clause of the
    // sibling is also the clause stating that all the parent decisions imply
    // the node's decision.
    bool implied = false;
    bool implied_and_processed = false;
    // Only set for open, non-implied nodes.
    std::unique_ptr<NodeTrailInfo> trail_info;
  };
  bool IsValid(const ProtoTrail& path) const ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
  Node* GetSibling(const Node* node) const ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
  // Returns the NodeTrailInfo for `node` or it's closest non-closed,
  // non-implied ancestor. `node` must be valid, never returns nullptr.
  NodeTrailInfo* GetTrailInfo(Node* node);
  void ClearTrailInfo(Node* node, bool implications_only = false);
  bool TrySplitTreeLockHeld(ProtoLiteral decision, ProtoTrail& path)
      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
  void Split(std::vector<std::pair<Node*, int>>& nodes, ProtoLiteral lit)
      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
  Node* MakeSubtree(Node* parent, ProtoLiteral decision)
      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
  void ProcessNodeChanges() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
  void ProcessImpliedNode(Node* node) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
  void UpdateLratClausesInSubtree(Node* node, Node* n,
                                  std::vector<ClauseId>& clauses)
      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
  Node* GetNode(int node_id) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
  std::vector<std::pair<Node*, int>> GetAssignedNodes(const ProtoTrail& path)
      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
  void AssignLeaf(ProtoTrail& path, Node* leaf)
      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
  void RestartLockHeld() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
  std::string ShortStatus() const ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);

  // Returns true if `literal` is a decision of `node` or one of its ancestors.
  bool IsDecisionOfNodeOrAncestor(ProtoLiteral literal, const Node* node) const
      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);

  // Returns "non-implied decisions of `node` and its ancestors => implied". A
  // node is considered implied if its `implied` field is true, and if its
  // `implied_and_processed` field or `skip_unprocessed_implied_nodes` is
  // true.
  std::vector<Literal> ImplicationClause(
      const Node* node, ProtoLiteral implied,
      bool skip_unprocessed_implied_nodes = false) const
      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);

  // Returns "decisions of `node`'s ancestors => negation of `node`'s decision"
  // (with decisions of implied nodes filtered out -- a node is considered
  // implied if its `implied` field is true, and if its `implied_and_processed`
  // field or `skip_unprocessed_implied_nodes` is true).
  std::vector<Literal> ClosingClause(
      const Node* node, bool skip_unprocessed_implied_nodes = false) const
      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);

  // Adds `imported_clause` to the LRAT proof handler, as well as
  // `inferred_clause`, inferred from `imported_clause` with `lrat_proof` (which
  // should contain clauses proving that literals removed from `imported_clause`
  // can actually be removed). Then deletes `imported_clause` and returns the ID
  // of the inferred clause.
  ClauseId AddImportedAndInferredClauses(
      absl::Span<const Literal> imported_clause,
      absl::Span<const Literal> inferred_clause,
      std::vector<ClauseId>& lrat_proof) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);

  bool CheckLratInvariants() const ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);

  mutable absl::Mutex mu_;
  const SatParameters& params_;
  const int num_workers_;
  const int max_path_depth_;
  SharedResponseManager* const shared_response_manager_;
  ClauseIdGenerator clause_id_generator_;
  std::unique_ptr<LratProofHandler> lrat_proof_handler_;

  // Stores the node id of the root, this is used to handle global restarts.
  int node_id_offset_ ABSL_GUARDED_BY(mu_) = 0;

  // Stores the nodes in the search tree.
  std::deque<Node> nodes_ ABSL_GUARDED_BY(mu_);
  std::deque<Node*> unassigned_leaves_ ABSL_GUARDED_BY(mu_);

  // How many splits we should generate now to keep the desired number of
  // leaves.
  int num_splits_wanted_ ABSL_GUARDED_BY(mu_);

  // We limit the total nodes generated per restart to cap the RAM usage and
  // communication overhead. If we exceed this, workers become portfolio
  // workers when no unassigned leaves are available.
  const int max_nodes_;
  int num_leaves_assigned_since_restart_ ABSL_GUARDED_BY(mu_) = 0;

  // Temporary vectors used to maintain the state of the tree when nodes are
  // closed and/or children are updated. Each node to close is associated with
  // the clause stating that the literals of its parents imply the negation of
  // its own literal (with literals of implied nodes filtered out).
  std::vector<std::pair<Node*, ClauseId>> to_close_ ABSL_GUARDED_BY(mu_);
  std::vector<Node*> to_update_ ABSL_GUARDED_BY(mu_);

  int64_t num_restarts_ ABSL_GUARDED_BY(mu_) = 0;
  int num_closed_nodes_ ABSL_GUARDED_BY(mu_) = 0;
};

class SharedTreeWorker {
 public:
  explicit SharedTreeWorker(Model* model);
  SharedTreeWorker(const SharedTreeWorker&) = delete;
  SharedTreeWorker& operator=(const SharedTreeWorker&) = delete;

  SatSolver::Status Search(
      const std::function<void()>& feasible_solution_observer);

 private:
  // Syncs the assigned tree with the local trail, ensuring that any new
  // implications are synced. This is a noop if the search is deeper than the
  // assigned tree. Returns false if the problem is unsat.
  bool SyncWithLocalTrail();
  bool SyncWithSharedTree();
  Literal DecodeDecision(ProtoLiteral literal);
  std::optional<ProtoLiteral> EncodeDecision(Literal decision);
  bool NextDecision(LiteralIndex* decision_index);
  void MaybeProposeSplits();
  bool ShouldReplaceSubtree();
  bool FinishedMinRestarts() const {
    return assigned_tree_.MaxLevel() > 0 &&
           restart_policy_->NumRestarts() >=
               parameters_->shared_tree_worker_min_restarts_per_subtree();
  }

  // Add any implications to the clause database for the current level.
  // Return true if any new information was added.
  bool AddImplications();
  bool AddDecisionImplication(Literal literal, int level, ClauseId clause_id);

  void ClearAssignedTreeDecisionsAndImplications();

  // Adds the LRAT inferred clause "assigned tree decisions up to `level` =>
  // `literal`" if `lrat_proof_handler_` is not null.
  ClauseId AddLratClauseAndProofForImplication(
      Literal literal, int level,
      std::optional<absl::FunctionRef<ClauseId(int, int)>> root_literals = {});

  // Adds the LRAT imported clause "assigned tree decisions up to `level` =>
  // `literal`" if `lrat_proof_handler_` is not null.
  ClauseId ImportLratClauseForImplication(Literal literal, int level);

  std::vector<Literal>& DecisionReason(int level);

  bool CheckLratInvariants();

  SatParameters* parameters_;
  SharedResponseManager* shared_response_;
  TimeLimit* time_limit_;
  SharedTreeManager* manager_;
  CpModelMapping* mapping_;
  SatSolver* sat_solver_;
  Trail* trail_;
  BinaryImplicationGraph* binary_propagator_;
  ClauseManager* clause_manager_;
  ClauseIdGenerator* clause_id_generator_;
  LratProofHandler* lrat_proof_handler_;
  IntegerTrail* integer_trail_;
  IntegerEncoder* encoder_;
  const ObjectiveDefinition* objective_;
  ModelRandomGenerator* random_;
  IntegerSearchHelper* helper_;
  SearchHeuristics* heuristics_;
  SatDecisionPolicy* decision_policy_;
  RestartPolicy* restart_policy_;
  LevelZeroCallbackHelper* level_zero_callbacks_;
  RevIntRepository* reversible_int_repository_;

  int64_t num_trees_ = 0;

  ProtoTrail assigned_tree_;
  std::vector<Literal> assigned_tree_decisions_;
  // The i-th element contains the literals implied by the first i elements of
  // assigned_tree_decisions_, together with the IDs of the corresponding LRAT
  // clauses (or kNoClauseId if lrat_proof_handler_ is null).
  std::vector<std::vector<std::pair<Literal, ClauseId>>>
      assigned_tree_implications_;
  double next_split_dtime_ = 0;

  // For each literal on the trail, the ID of the LRAT clause stating that this
  // literal is implied by the previous decisions on the trail, or kNoClauseId
  // if there is no such clause.
  std::vector<ClauseId> trail_implication_clauses_;

  std::vector<ProtoLiteral> tmp_splits_;
  std::vector<Literal> reason_;
  // Stores the average LBD of learned clauses for each tree assigned since it
  // was assigned.
  // If a tree has worse LBD than the average over the last few trees we replace
  // the tree.
  RunningAverage assigned_tree_lbds_;
  double earliest_replacement_dtime_ = 0;

  // Stores the trail index of the last implication added to assigned_tree_.
  int reversible_trail_index_ = 0;
  // Stores the number of implications processed for each level in
  // assigned_tree_.
  std::deque<int> rev_num_processed_implications_;
};

}  // namespace operations_research::sat

#endif  // ORTOOLS_SAT_WORK_ASSIGNMENT_H_