#pragma once
#include "../utils/common.h"
#include "../msl_infer/Sampler.h"
#include "../msl_infer/Encoder.h"
#include "../msl_infer/Renderer.h"
#include "../msl_infer/Msl.h"

class InferPipeline {
  public:
    InferPipeline(sptr<Msl> net, unsigned int nRays, unsigned int nSamplesPerRay,
                  glm::vec2 depthRange, int encodeDim, int coordChns);

    void run(sptr<CudaArray<glm::vec4>> o_colors, sptr<CudaArray<glm::vec3>> rays, glm::vec3 origin,
             bool showPerf = false);

  private:
    unsigned int _nRays;
    unsigned int _nSamplesPerRay;
    unsigned int _coordChns;
    sptr<Msl> _net;
    sptr<Sampler> _sampler;
    sptr<Encoder> _encoder;
    sptr<Renderer> _renderer;
    sptr<CudaArray<float>> _coords;
    sptr<CudaArray<float>> _depths;
    sptr<CudaArray<float>> _encoded;
    sptr<CudaArray<glm::vec4>> _layeredColors;
};