// Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.

#include "tensor_update.h"
#include "valueupdate.h"
#include <optional>

namespace vespalib::eval { struct Value; }

namespace document {

class TensorDataType;
class TensorFieldValue;

/*
 * An update for a subset of the cells in a tensor.
 *
 * The operand is represented as a tensor field value containing a
 * mapped (aka sparse) tensor.
 */
class TensorModifyUpdate final : public ValueUpdate, public TensorUpdate {
public:
    /** Declare all types of tensor modify updates. */
    enum class Operation { // Operation to be applied to matching tensor cells
        REPLACE  = 0,
        ADD      = 1,
        MULTIPLY = 2,
        MAX_NUM_OPERATIONS = 3
    };
private:
    Operation _operation;
    std::unique_ptr<const TensorDataType> _tensorType;
    std::unique_ptr<TensorFieldValue> _tensor;
    // When this is set, non-existing cells are created in the input tensor before applying the update.
    std::optional<double> _default_cell_value;

    friend ValueUpdate;
    TensorModifyUpdate();
    ACCEPT_UPDATE_VISITOR;
public:
    TensorModifyUpdate(Operation operation, std::unique_ptr<TensorFieldValue> tensor);
    TensorModifyUpdate(Operation operation, std::unique_ptr<TensorFieldValue> tensor, double default_cell_value);
    TensorModifyUpdate(const TensorModifyUpdate &rhs) = delete;
    TensorModifyUpdate &operator=(const TensorModifyUpdate &rhs) = delete;
    ~TensorModifyUpdate() override;

    bool operator==(const ValueUpdate &other) const override;
    Operation getOperation() const { return _operation; }
    const TensorFieldValue &getTensor() const { return *_tensor; }
    const std::optional<double>& get_default_cell_value() const { return _default_cell_value; }
    void checkCompatibility(const Field &field) const override;
    std::unique_ptr<vespalib::eval::Value> applyTo(const vespalib::eval::Value &tensor) const;
    std::unique_ptr<Value> apply_to(const Value &tensor,
                                    const ValueBuilderFactory &factory) const override;
    bool applyTo(FieldValue &value) const override;
    void printXml(XmlOutputStream &xos) const override;
    void print(std::ostream &out, bool verbose, const std::string &indent) const override;
    void deserialize(const DocumentTypeRepo &repo, const DataType &type, nbostream &stream) override;
};

}