/*
 *  ******************************************************************************
 *  *
 *  *
 *  * This program and the accompanying materials are made available under the
 *  * terms of the Apache License, Version 2.0 which is available at
 *  * https://www.apache.org/licenses/LICENSE-2.0.
 *  *
 *  *  See the NOTICE file distributed with this work for additional
 *  *  information regarding copyright ownership.
 *  * 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.
 *  *
 *  * SPDX-License-Identifier: Apache-2.0
 *  *****************************************************************************
 */

//
// @author raver119@gmail.com
//

#ifndef LIBND4J_CONTEXT_H
#define LIBND4J_CONTEXT_H

#include <vector>
#include <array/NDArray.h>
#include <graph/Variable.h>
#include <graph/VariableSpace.h>
#include <graph/ContextPrototype.h>
#include <memory/Workspace.h>
#include <execution/Engine.h>

// CUDA-specific includes
#ifdef __CUDACC__

#include <cuda.h>
#include <cuda_runtime_api.h>
#include <cuda_runtime.h>
#include <cuda_device_runtime_api.h>
#endif

namespace sd {
    namespace graph {
        /**
         * This class defines input desired for any given node/operation within graph
         */
        class ND4J_EXPORT Context : public sd::graph::ContextPrototype {
        protected:
            sd::memory::Workspace* _workspace = nullptr;
            sd::graph::VariableSpace* _variableSpace = nullptr;
            std::pair<Nd4jLong, Nd4jLong> _executionTime;
            sd::random::RandomBuffer* _rng = nullptr;

            sd::DataType _dataType = sd::DataType::FLOAT32;
            // branch for divergent_op
            int _branch = 0;

            // temporary context for standalone ops execution
            LaunchContext* _context = nullptr;

            std::vector<sd::DataType> _dataTypes;

            // fields for fast execution (out-of-graph ops use)
            std::vector<NDArray*> _fastpath_in;
            std::vector<NDArray*> _fastpath_out;
            std::vector<NDArray*> _handles;

            bool _helpersAllowed = true;

            // in some cases we might be able to skip shape function for validation purposes
            bool _shapeFunctionOverride = false;

            // special flag used during conversion from Graph exec to FastPath exec
            bool _forbidFastPath = false;
        public:
            Context(ContextPrototype* prototype, VariableSpace* variableSpace);

            explicit Context(int nodeId, VariableSpace *variableSpace = nullptr);
            Context(int nodeId, VariableSpace *variableSpace, bool isInplace);

            // default destructor
            ~Context();

            // these methods are for execution timing
            void setOuterTime(Nd4jLong time);
            void setInnerTime(Nd4jLong time);
            Nd4jLong getOuterTime();
            Nd4jLong getInnerTime();

            sd::DataType dataType() override;

            sd::DataType dataType(int index) override;
            void setDataType(int index, sd::DataType type) override;
            // these methods are related to Workspace abstraction
            bool hasWorkspaceProvided();
            void attachWorkspace(sd::memory::Workspace* workspace);
            void forgetWorkspace();

            // these methods return full-time workspace
            sd::memory::Workspace* getWorkspace();
            sd::memory::Workspace* workspace();
            sd::memory::Workspace* fWorkspace();

            // this method returns workspace for temporary allocations
            sd::memory::Workspace* tWorkspace();

            // this method returns workspace for object allocations
            sd::memory::Workspace* oWorkspace();

            void setVariableSpace(VariableSpace* variableSpace);

            sd::random::RandomBuffer* getRNG();
            void setRNG(sd::random::RandomBuffer* rng);

            void setTargetEngine(samediff::Engine engine);

            VariableSpace *getVariableSpace();

            LaunchContext* launchContext();

            // these fields define, if we can execute specific node in-place, without generating new array


            // these variables are only for Divergent Nodes
            int getBranch();
            void setBranch(int branch);

            /**
             *
             * @return
             */
            Stash* getStash();

            /**
             *
             */
            void trackList(NDArrayList* list);


            /**
             * This method returns variable for a given input index for this block
             * @param idx
             * @return
             */
            Variable* getVariable(int idx);
            Variable* variable(int idx);

            /**
             * This method is shortcut to getVariable(int idx);
             *
             * + it check fastpath for array availability (preferred)
             * @return
             */
            NDArray* getNDArray(int idx);
            NDArray* array(int idx);


            /**
             * This method fetches variable from VariableSpace DIRECTLY
             * @param p
             * @return
             */
            Variable* variable(int node, int index);
            Variable* variable(std::pair<int,int>& p);
            Variable* variable(std::initializer_list<int> p);


            void pushNDArrayToVariableSpace(int nodeId, int index, NDArray* array, bool removable = true);
            void pushNDArrayToVariableSpace(std::pair<int, int>& pair, NDArray* array, bool removable = true);

            void pushNDArrayListToVariableSpace(int nodeId, int index, NDArrayList* list, bool track = true);
            void pushNDArrayListToVariableSpace(std::pair<int, int>& pair, NDArrayList* list, bool track = true);

            bool isValueAvailable(int idx = 0);

            Variable* ensureVariable(int idx = 0);

            unsigned long width() override;

            // methods used in java interop
            /**
             * This method checks if Context uses fastpath variable access
             * @return
             */
            bool isFastPath();

            /**
             * Method allows to forbid FastPath execution
             * @param reallyForbid
             */
            void forbidFastPath(bool reallyForbid);

#ifndef __JAVACPP_HACK__
            std::vector<NDArray*>& fastpath_in();
            std::vector<NDArray*>& fastpath_out();
#endif

            void setInputArray(int index, NDArray *array, bool removable = false);
            void setInputArray(int index, void *buffer, void const* shapeInfo, void *specialBuffer, void const* specialShapeInfo);
            void setInputArray(int index, void *buffer, void * shapeInfo, void *specialBuffer, void * specialShapeInfo);
            void setInputArray(int index, void *databuffer, void const* shapeInfo, void const* specialShapeInfo);

            void setOutputArray(int index, NDArray *array, bool removable = false);
            void setOutputArray(int index, void *buffer, const void * shapeInfo, void *specialBuffer, const void * specialShapeInfo);
            void setOutputArray(int index, void *buffer, void * shapeInfo, void *specialBuffer, void * specialShapeInfo);
            void setOutputArray(int index, void *databuffer, void const* shapeInfo, void const* specialShapeInfo);

            void setTArguments(double *arguments, int numberOfArguments);
            void setIArguments(Nd4jLong *arguments, int numberOfArguments);
            void setBArguments(bool *arguments, int numberOfArguments);
            void setDArguments(sd::DataType *arguments, int numberOfArguments);

            void setTArguments(const std::vector<double> &tArgs);
            void setIArguments(const std::vector<Nd4jLong> &tArgs);
            void setBArguments(const std::vector<bool> &tArgs);
            void setDArguments(const std::vector<sd::DataType> &dArgs);

            /**
             * This method purges fastpath in/out contents and releases all the handles.
             *
             * PLEASE NOTE: I/T/B/D args will stay intact
             */
            void clearFastPath();

            void setCudaContext(Nd4jPointer cudaStream, Nd4jPointer reductionPointer, Nd4jPointer allocationPointer);

            void allowHelpers(bool reallyAllow);
            bool helpersAllowed();

            void setShapeFunctionOverride(bool reallyOverride);
            bool shapeFunctionOverride();

            samediff::ExecutionMode executionMode();
            void setExecutionMode(samediff::ExecutionMode executionMode);

            bool isTraining();
            bool isInference();
        };
    }
}


#endif //LIBND4J_BLOCK_H