// 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.

#include "ortools/sat/integer_base.h"

#include <algorithm>
#include <cstdint>
#include <numeric>
#include <tuple>
#include <utility>
#include <vector>

#include "absl/log/check.h"
#include "ortools/util/bitset.h"

namespace operations_research::sat {

void LinearExpression2::SimpleCanonicalization() {
  if (coeffs[0] == 0) vars[0] = kNoIntegerVariable;
  if (coeffs[1] == 0) vars[1] = kNoIntegerVariable;

  // Corner case when the underlying variable is the same.
  if (vars[0] != kNoIntegerVariable && vars[1] != kNoIntegerVariable &&
      PositiveVariable(vars[0]) == PositiveVariable(vars[1])) {
    // Make sure variable are positive before merging.
    for (int i = 0; i < 2; ++i) {
      if (!VariableIsPositive(vars[i])) {
        coeffs[i] = -coeffs[i];
        vars[i] = NegationOf(vars[i]);
      }
    }

    coeffs[0] += coeffs[1];
    coeffs[1] = 0;
    vars[1] = kNoIntegerVariable;
    if (coeffs[0] == 0) vars[0] = kNoIntegerVariable;
  }

  // Make sure coeff are positive.
  for (int i = 0; i < 2; ++i) {
    if (coeffs[i] < 0) {
      coeffs[i] = -coeffs[i];
      vars[i] = NegationOf(vars[i]);
    }
  }

  // Make sure variable are sorted.
  if (vars[0] > vars[1]) {
    std::swap(vars[0], vars[1]);
    std::swap(coeffs[0], coeffs[1]);
  }
}

IntegerValue LinearExpression2::DivideByGcd() {
  const uint64_t gcd = std::gcd(coeffs[0].value(), coeffs[1].value());
  if (gcd > 1) {
    coeffs[0] /= gcd;
    coeffs[1] /= gcd;
    return IntegerValue(gcd);
  }
  return IntegerValue(1);
}

bool LinearExpression2::NegateForCanonicalization() {
  bool negate = false;
  if (coeffs[0] == 0) {
    if (coeffs[1] != 0) {
      negate = !VariableIsPositive(vars[1]);
    }
  } else {
    negate = !VariableIsPositive(vars[0]);
  }
  if (negate) Negate();
  return negate;
}

bool LinearExpression2::CanonicalizeAndUpdateBounds(IntegerValue& lb,
                                                    IntegerValue& ub) {
  SimpleCanonicalization();
  if (coeffs[0] == 0 || coeffs[1] == 0) return false;  // abort.

  const bool negated = NegateForCanonicalization();
  if (negated) {
    // We need to be able to negate without overflow.
    CHECK_GE(lb, kMinIntegerValue);
    CHECK_LE(ub, kMaxIntegerValue);
    std::swap(lb, ub);
    lb = -lb;
    ub = -ub;
  }

  // Do gcd division.
  const uint64_t gcd = std::gcd(coeffs[0].value(), coeffs[1].value());
  if (gcd > 1) {
    coeffs[0] /= gcd;
    coeffs[1] /= gcd;
    ub = FloorRatio(ub, IntegerValue(gcd));
    lb = CeilRatio(lb, IntegerValue(gcd));
  }

  CHECK(coeffs[0] != 0 || vars[0] == kNoIntegerVariable);
  CHECK(coeffs[1] != 0 || vars[1] == kNoIntegerVariable);

  return negated;
}

bool LinearExpression2::IsCanonicalized() const {
  for (int i : {0, 1}) {
    if ((vars[i] == kNoIntegerVariable) != (coeffs[i] == 0)) {
      return false;
    }
  }
  if (vars[0] >= vars[1]) {
    if (vars[0] == kNoIntegerVariable && vars[1] == kNoIntegerVariable) {
      return true;
    }
    return false;
  }

  if (vars[0] == kNoIntegerVariable) return true;

  return coeffs[0] > 0 && coeffs[1] > 0;
}

AffineExpression LinearExpression2::GetAffineLowerBound(
    int var_index, IntegerValue expr_lb, IntegerValue other_var_lb) const {
  DCHECK_GE(var_index, 0);
  DCHECK_LT(var_index, 2);
  const IntegerValue coeff = coeffs[var_index];
  const IntegerVariable var = vars[var_index];
  DCHECK_NE(var, kNoIntegerVariable);
  const IntegerVariable other_var = vars[1 - var_index];
  const IntegerValue other_coeff = coeffs[1 - var_index];
  if (other_var == kNoIntegerVariable) {
    return AffineExpression(expr_lb);
  }
  DCHECK_GT(other_coeff, 0);

  if (coeff == 1) {
    return AffineExpression(other_var, -other_coeff, expr_lb);
  }
  DCHECK_GT(coeff, 1);
  DCHECK_EQ(std::gcd(other_coeff.value(), coeff.value()), 1);

  // We need to handle the case where the coefficient is not unit. We have
  // a * x + b * y >= expr_lb  (with a > 1 and b >= 1).
  // x >= expr_lb / a - (b / a) * y
  //
  // We write b = ceil(b / a) * a - residual, with residual > 0 (since
  // GCD(a, b)==1). Then we have:
  // x >= expr_lb / a - ceil(b / a) * y + residual * y / a
  // x >= expr_lb / a - ceil(b / a) * y + residual * y_lb / a
  // x >= -ceil(b / a) * y + (expr_lb + residual * y_lb) / a
  // But now we have
  // x + ceil(b / a) * y >= (expr_lb + residual * y_lb) / a
  // And since the lhs is an integer value, this is equivalent to
  // x + ceil(b / a) * y >= ceil((expr_lb + residual * y_lb) / a)
  const IntegerValue ceil_coeff_ratio = CeilRatio(other_coeff, coeff);
  const IntegerValue residual = coeff * ceil_coeff_ratio - other_coeff;
  return AffineExpression(other_var, -ceil_coeff_ratio,
                          CeilRatio(expr_lb + residual * other_var_lb, coeff));
}

void LinearExpression2::MakeVariablesPositive() {
  SimpleCanonicalization();
  for (int i = 0; i < 2; ++i) {
    if (vars[i] != kNoIntegerVariable && !VariableIsPositive(vars[i])) {
      coeffs[i] = -coeffs[i];
      vars[i] = NegationOf(vars[i]);
    }
  }
}

std::pair<BestBinaryRelationBounds::AddResult,
          BestBinaryRelationBounds::AddResult>
BestBinaryRelationBounds::Add(LinearExpression2 expr, IntegerValue lb,
                              IntegerValue ub) {
  const bool negated = expr.CanonicalizeAndUpdateBounds(lb, ub);

  // We only store proper linear2.
  if (expr.coeffs[0] == 0 || expr.coeffs[1] == 0) {
    return {AddResult::INVALID, AddResult::INVALID};
  }

  auto [it, inserted] = best_bounds_.insert({expr, {lb, ub}});
  if (inserted) {
    std::pair<AddResult, AddResult> result = {
        lb > kMinIntegerValue ? AddResult::ADDED : AddResult::INVALID,
        ub < kMaxIntegerValue ? AddResult::ADDED : AddResult::INVALID};
    if (negated) std::swap(result.first, result.second);
    return result;
  }

  const auto [known_lb, known_ub] = it->second;

  std::pair<AddResult, AddResult> result = {
      lb > kMinIntegerValue ? AddResult::NOT_BETTER : AddResult::INVALID,
      ub < kMaxIntegerValue ? AddResult::NOT_BETTER : AddResult::INVALID};
  if (lb > known_lb) {
    result.first = (it->second.first == kMinIntegerValue) ? AddResult::ADDED
                                                          : AddResult::UPDATED;
    it->second.first = lb;
  }
  if (ub < known_ub) {
    result.second = (it->second.second == kMaxIntegerValue)
                        ? AddResult::ADDED
                        : AddResult::UPDATED;
    it->second.second = ub;
  }
  if (negated) std::swap(result.first, result.second);
  return result;
}

RelationStatus BestBinaryRelationBounds::GetStatus(LinearExpression2 expr,
                                                   IntegerValue lb,
                                                   IntegerValue ub) const {
  expr.CanonicalizeAndUpdateBounds(lb, ub);
  if (expr.coeffs[0] == 0 || expr.coeffs[1] == 0) {
    return RelationStatus::IS_UNKNOWN;
  }

  const auto it = best_bounds_.find(expr);
  if (it != best_bounds_.end()) {
    const auto [known_lb, known_ub] = it->second;
    if (lb <= known_lb && ub >= known_ub) return RelationStatus::IS_TRUE;
    if (lb > known_ub || ub < known_lb) return RelationStatus::IS_FALSE;
  }
  return RelationStatus::IS_UNKNOWN;
}

std::vector<std::pair<LinearExpression2, IntegerValue>>
BestBinaryRelationBounds::GetSortedNonTrivialUpperBounds() const {
  std::vector<std::pair<LinearExpression2, IntegerValue>> root_relations_sorted;
  root_relations_sorted.reserve(2 * best_bounds_.size());
  for (const auto& [expr, bounds] : best_bounds_) {
    if (bounds.first != kMinIntegerValue) {
      LinearExpression2 negated_expr = expr;
      negated_expr.Negate();
      root_relations_sorted.push_back({negated_expr, -bounds.first});
    }
    if (bounds.second != kMaxIntegerValue) {
      root_relations_sorted.push_back({expr, bounds.second});
    }
  }
  std::sort(root_relations_sorted.begin(), root_relations_sorted.end());
  return root_relations_sorted;
}

std::vector<std::tuple<LinearExpression2, IntegerValue, IntegerValue>>
BestBinaryRelationBounds::GetSortedNonTrivialBounds() const {
  std::vector<std::tuple<LinearExpression2, IntegerValue, IntegerValue>>
      root_relations_sorted;
  root_relations_sorted.reserve(best_bounds_.size());
  for (const auto& [expr, bounds] : best_bounds_) {
    root_relations_sorted.push_back({expr, bounds.first, bounds.second});
  }
  std::sort(root_relations_sorted.begin(), root_relations_sorted.end());
  return root_relations_sorted;
}

}  // namespace operations_research::sat