#pragma once
#ifndef PATHSPACE_EDGE_MANAGER_H__
#define PATHSPACE_EDGE_MANAGER_H__

#include "edge_manager.h"
#include "sampler.h"
#include "utils.h"
#include <algorithm>
#include <iostream>
#include <vector>

struct Scene;

struct PathSpaceEdgeManager
{
    PathSpaceEdgeManager(const Scene& scene, const Eigen::Array<Float, -1, 1> &shapeWeights,
                         const Vector3i &direct_dims, const Eigen::Array<Float, -1, 1> &direct_data,
                         const Vector3i &indirect_dims, const Eigen::Array<Float, -1, 1> &indirect_data,
                         bool verbose = false);
    void initPrimaryEdges(int primary_dims, const Eigen::Array<Float, -1, 1> &data);

    const Edge& sampleEdgeRay(const Vector &rnd3, int &shape_id, RayAD &ray, Float &pdf, bool direct = true) const;
    const Edge& samplePrimaryEdgePoint(const Float& _rnd, int &shape_id, VectorAD& p, Float &pdf) const;


    const Scene &m_scene;

    DiscreteDistribution edge_distrbs[3];   // 0: direct; 1: indirect; 2: primary.
    std::vector<Vector2i> edge_indices[3];

    // For guided edge-ray sampling
    Vector3i grid_dims[2];
    int array_dim;
    DiscreteDistribution grid_distrbs[3];
};


#endif