/******************************************************************************
 * Top contributors (to current version):
 *   Alex Ozdemir
 *
 * This file is part of the cvc5 project.
 *
 * Copyright (c) 2009-2025 by the authors listed in the file AUTHORS
 * in the top-level source directory and their institutional affiliations.
 * All rights reserved.  See the file COPYING in the top-level source
 * directory for licensing information.
 * ****************************************************************************
 *
 * Finding common zeros using split Groebner bases.
 *
 * The following procedures are from the paper "Split Groebner Bases for
 * Satisfiability Modulo Finite Fields", by Ozdemir et al. Hopefully it will be
 * published soon:
 *
 * * split: Split
 * * splitFindZero: SplitFindZero
 * * splitZeroExtend: SplitZeroExtend
 * * splitGb: SplitGb
 * * applyRule: ApplyRule
 * * admit: admit
 * * BitProp::getBitEqualities: extraProp
 */

#ifdef CVC5_USE_COCOA

#include "cvc5_private.h"

#ifndef CVC5__THEORY__FF__SPLIT_GB_H
#define CVC5__THEORY__FF__SPLIT_GB_H

// external includes
#include <CoCoA/ideal.H>
#include <CoCoA/ring.H>
#include <CoCoA/symbol.H>

// std includes
#include <memory>
#include <optional>
#include <unordered_set>
#include <variant>
#include <vector>

// internal includes
#include "theory/ff/cocoa_encoder.h"
#include "theory/ff/cocoa_util.h"
#include "theory/ff/multi_roots.h"
#include "util/resource_manager.h"

namespace cvc5::internal {
namespace theory {
namespace ff {

/** defined below */
class Gb;
/** defined below */
class BitProp;

/** A split GB. */
using SplitGb = std::vector<Gb>;

/** Find a model for these facts
 *
 * @param facts Finite field equalities and disequalities.
 * @param size The size of the field they are in.
 * @param env used for resource management
 *
 * @return A model, or not if none exists.
 * */
std::optional<FfModel> split(const std::vector<Node>& facts,
                             const FfSize& size,
                             const Env& env);

/** Compute a split Gb.
 *
 * @param generatorSets A collection of collections of polynomials.
 * @param bitProp A data structure that stores information about which CoCoA
 *                variables are bitsums of which other CoCoA variables.
 * @param env used for resource management
 *
 * @return A split Groebner basis for the ideal generated by generatorSets.
 * */
std::vector<Gb> splitGb(const std::vector<std::vector<Poly>>& generatorSets,
                        BitProp& bitProp,
                        const Env& env);

/** Whether to admit p into ideal i. */
bool admit(size_t i, const Poly& p);

/** Find a zero for this split Gb.
 *
 * @param splitBasis A split Groebner basis.
 * @param polyRing The polynomial ring it's in.
 * @param bitProp Info about which vars are bitsums of which. See splitGb.
 * @param env used for resource management
 *
 * @return A common zero for all polynomials in the ideal generated by
 *         splitBasis, or not if none exists.
 *
 * */
std::optional<Point> splitFindZero(SplitGb&& splitBasis,
                                   CoCoA::ring polyRing,
                                   BitProp& bitProp,
                                   const Env& env);

/** Extend curR into a zero for this split Gb.
 *
 * @param origPolys Generators for an ideal.
 * @param curBases A split basis for the same ideal.
 * @param curR A partial zero for the ideal.
 * @param bitProp Info about which vars are bitsums of which. See splitGb.
 * @param env used for resource management
 *
 * @return A common (full) zero, a conflict polynomial, or false (indicating
 *         that no common zeros exist). The conflict polynomial is guaranteed to
 *         not be in the ideal generated by the first of the split bases.
 * */
std::variant<Point, Poly, bool> splitZeroExtend(
    const std::vector<Poly>& origPolys,
    const SplitGb&& curBases,
    const PartialPoint&& curR,
    const BitProp& bitProp,
    const Env& env);

/** Apply a branching rule.
 *
 * @param gb A groebner basis.
 * @param polyRing The ring its ideal is in.
 * @param r A partial zero for the ideal.
 *
 * @return Something that can enumerate single-variable assignments that cover
 *         all common zeros for gb that extend r.
 * */
std::unique_ptr<AssignmentEnumerator> applyRule(const Gb& gb,
                                                const CoCoA::ring& polyRing,
                                                const PartialPoint& r);

/** Check wether this point is a zero. */
void checkZero(const SplitGb& origBases, const Point& zero);

/** Wraps a CoCoA GBasis, but supports an empty basis. */
class Gb
{
 public:
  Gb();
  /** A GB for the ideal generated by generators */
  Gb(const std::vector<Poly>& generators, const ResourceManager* rm);
  /** Does this GB's ideal contain p? */
  bool contains(const Poly& p) const;
  /** Is this GB's ideal equal to the whole ring? */
  bool isWholeRing() const;
  /** Reduce p modulo this GB's ideal. */
  Poly reduce(const Poly& p) const;
  /** Is this GB's ideal zero-dimensional? */
  bool zeroDimensional() const;
  /** Compute a minimal q such that q(p) = 0 mod this GB's ideal. */
  Poly minimalPolynomial(const Poly& p) const;
  /** List the elements of the GB */
  const std::vector<Poly>& basis() const;

 private:
  /** CoCoA's representation of the ideal. Empty if the ideal is 0. */
  std::optional<CoCoA::ideal> d_ideal;
  /** The GB. Empty if the ideal is 0. */
  std::vector<Poly> d_basis;
};

/** Propagator for bit equalities from bitsum equalities. */
class BitProp
{
 public:
  BitProp(const std::vector<Node>& facts, CocoaEncoder& encoder);
  BitProp();
  /** get all known bit equalities from thes split basis */
  std::vector<Poly> getBitEqualities(const SplitGb& sgb);

 private:
  /** is this node is bit constrained? */
  bool isBit(const Node& possibleBit, const SplitGb& splitBasis);

  /** terms that are known to be 0 or 1 */
  std::unordered_set<Node> d_bits;
  /** known bitsums */
  std::vector<Node> d_bitsums;
  /** the ambiant encoding of terms into Cocoa */
  CocoaEncoder* d_enc;
};

}  // namespace ff
}  // namespace theory
}  // namespace cvc5::internal

#endif /* CVC5__THEORY__FF__SPLIT_GB_H */

#endif /* CVC5_USE_COCOA */
