// Copyright 2020 Intel Corporation
// SPDX-License-Identifier: Apache-2.0

#pragma once

#include "../../observer/ObserverRegistry.h"
#include "../../sampler/Sampler.h"
#include "../common/simd.h"
#include "VdbGrid.h"
#include "VdbSamplerShared.h"
#include "VdbVolume.h"
#include "openvkl/devices/common/StructShared.h"
#include "openvkl/openvkl.h"
#include "openvkl/vdb.h"

namespace openvkl {
  namespace cpu_device {

    template <int W>
    using VdbSamplerBase = SamplerBase<W,
                                       VdbVolume,
                                       VdbIntervalIteratorFactory,
                                       VdbHitIteratorFactory>;

    template <int W>
    struct VdbSampler
        : public AddStructShared<VdbSamplerBase<W>, ispc::VdbSamplerShared>
    {
      explicit VdbSampler(Device *, VdbVolume<W> &volume);
      ~VdbSampler() override;

      void commit() override;

      VKLFeatureFlagsInternal getFeatureFlags() const override;

      // single attribute /////////////////////////////////////////////////////

      void computeSample(const vvec3fn<1> &objectCoordinates,
                         vfloatn<1> &samples,
                         unsigned int attributeIndex,
                         const vfloatn<1> &time) const override final;

      void computeSampleV(const vintn<W> &valid,
                          const vvec3fn<W> &objectCoordinates,
                          vfloatn<W> &samples,
                          unsigned int attributeIndex,
                          const vfloatn<W> &time) const override final;

      void computeSampleN(unsigned int N,
                          const vvec3fn<1> *objectCoordinates,
                          float *samples,
                          unsigned int attributeIndex,
                          const float *times) const override final;

      void computeGradientV(const vintn<W> &valid,
                            const vvec3fn<W> &objectCoordinates,
                            vvec3fn<W> &gradients,
                            unsigned int attributeIndex,
                            const vfloatn<W> &time) const override final;

      void computeGradientN(unsigned int N,
                            const vvec3fn<1> *objectCoordinates,
                            vvec3fn<1> *gradients,
                            unsigned int attributeIndex,
                            const float *times) const override final;

      // multi-attribute //////////////////////////////////////////////////////

      void computeSampleM(const vvec3fn<1> &objectCoordinates,
                          float *samples,
                          unsigned int M,
                          const unsigned int *attributeIndices,
                          const vfloatn<1> &time) const override final;

      void computeSampleMV(const vintn<W> &valid,
                           const vvec3fn<W> &objectCoordinates,
                           float *samples,
                           unsigned int M,
                           const unsigned int *attributeIndices,
                           const vfloatn<W> &time) const override final;

      void computeSampleMN(unsigned int N,
                           const vvec3fn<1> *objectCoordinates,
                           float *samples,
                           unsigned int M,
                           const unsigned int *attributeIndices,
                           const float *times) const override final;

      Observer<W> *newObserver(const char *type) override;

      ObserverRegistry<W> &getLeafAccessObserverRegistry()
      {
        return leafAccessObservers;
      }

     private:
      using VdbSamplerBase<W>::volume;

      ObserverRegistry<W> leafAccessObservers;
    };

  }  // namespace cpu_device
}  // namespace openvkl
