#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, uint nRays, uint 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:
    uint _nRays;
    uint _nSamplesPerRay;
    uint _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;
};