#include "scene.h"
#include "emitter/area.h"
#include "emitter/area2.h"
#include "bsdf/null.h"
#include "bsdf/diffuse.h"
#include "bsdf/texturedDiffuse.h"
#include "bsdf/phong.h"
#include "bsdf/roughconductor.h"
#include "bsdf/roughdielectric.h"
#include "bsdf/twosided.h"
#include "medium/homogeneous.h"
#include "medium/heterogeneous.h"
#include "integrator/path.h"
#include "integrator/ptracer.h"
#include "integrator/volpath_simple.h"
#include "integrator/volpath.h"
#include "integrator/direct.h"
#include "integrator/bdpt.h"
#include "integratorADps.h"
#include "tofIntegratorADps.h"
#include "integrator/differential/directAD.h"
#include "integrator/differential/pathAD.h"
#include "integrator/differential/volpathAD.h"
#include "integrator/differential/directADps.h"
#include "integrator/differential/pathADps.h"
#include "integrator/differential/bdptAD.h"
#include "integrator/differential/invDirectADps.h"
#include "integrator/differential/ptracerADps.h"
#include "integrator/differential/binnedPathADps.h"
#include "integrator/differential/binnedTofPathADps.h"
#include "integrator/differential/binnedPtracerADps.h"
#include "unit_test/ad_test.h"
#include "unit_test/volume_test.h"
#include "unit_test/ps_edge_test.h"
#include "unit_test/path_test.h"
#include "unit_test/ellipsoid_test.h"
#include "unit_test/tof_test.h"
#include "unit_test/pif_test.h"
#include "config.h"
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

namespace py = pybind11;

PYBIND11_MODULE(vredner, m) {
    m.doc() = "vredner"; // optional module docstring

    py::class_<ptr<float>>(m, "float_ptr")
        .def(py::init<std::size_t>());
    py::class_<ptr<int>>(m, "int_ptr")
        .def(py::init<std::size_t>());

    py::class_<Spectrum3f>(m, "Spectrum3f")
        .def(py::init<float, float, float>());

    py::class_<Camera>(m, "Camera")
        .def(py::init<int, int, ptr<float>, ptr<float>, float, int, ptr<float>>())
        .def("set_rect", &Camera::setCropRect)
        .def("set_pif", &Camera::setPIF);

    py::class_<Shape>(m, "Shape")
        .def(py::init<ptr<float>,
                      ptr<int>,
                      ptr<float>,
                      ptr<float>,
                      int,
                      int,
                      int,
                      int,
                      int,
                      int,
                      ptr<float>>())
        .def_readonly("num_vertices", &Shape::num_vertices)
        .def("has_uvs", &Shape::hasUVs)
        .def("has_normals", &Shape::hasNormals);

    py::class_<BSDF>(m, "BSDF");
    py::class_<DiffuseBSDF, BSDF>(m, "BSDF_diffuse")
        .def(py::init<Spectrum3f>())
        .def(py::init<Spectrum3f, ptr<float>>());
    py::class_<TexturedDiffuseBSDF, BSDF>(m, "BSDF_texturedDiffuse")
        .def(py::init<int, int, ptr<float>>());
    py::class_<NullBSDF, BSDF>(m, "BSDF_null")
        .def(py::init<>());
    py::class_<PhongBSDF, BSDF>(m, "BSDF_Phong")
        .def(py::init<Spectrum3f, Spectrum3f, float>())
        .def(py::init<Spectrum3f, Spectrum3f, float, ptr<float>, ptr<float>, ptr<float>>());
    py::class_<RoughConductorBSDF, BSDF>(m, "BSDF_roughconductor")
        .def(py::init<float, Spectrum3f, Spectrum3f>())
        .def(py::init<float, Spectrum3f, Spectrum3f, ptr<float>>())
        .def(py::init<float, Spectrum3f, Spectrum3f, ptr<float>, ptr<float>, ptr<float>>());
    py::class_<RoughDielectricBSDF, BSDF>(m, "BSDF_roughdielectric")
        .def(py::init<float, float, float>())
        .def(py::init<float, float, float, ptr<float>>())
        .def(py::init<float, float, float, ptr<float>, ptr<float>>());
    py::class_<TwosidedBSDF, BSDF>(m, "BSDF_twosided")
        .def(py::init<const BSDF *>());

    py::class_<PhaseFunction>(m, "Phase");
    py::class_<HGPhaseFunction, PhaseFunction>(m, "HG")
        .def(py::init<float>())
        .def(py::init<float, ptr<float>>());
    py::class_<IsotropicPhaseFunction, PhaseFunction>(m, "Isotropic")
        .def(py::init<>());

    py::class_<Emitter>(m, "Emitter");
    py::class_<AreaLight, Emitter>(m, "AreaLight")
        .def(py::init<int, Spectrum3f>())
        .def(py::init<int, Spectrum3f, ptr<float>>());
    py::class_<AreaLightEx, Emitter>(m, "AreaLightEx")
        .def(py::init<int, Spectrum3f, float>())
        .def(py::init<int, Spectrum3f, float, ptr<float>, ptr<float>>());

    py::class_<Medium>(m, "Medium");
    py::class_<Homogeneous, Medium>(m, "Homogeneous")
        .def(py::init<float, Spectrum3f, int>())
        .def(py::init<float, Spectrum3f, int, ptr<float>, ptr<float>>());
    py::class_<Heterogeneous, Medium>(m, "Heterogeneous")
        .def(py::init<const std::string &, ptr<float>, float, Spectrum3f, int>())
        .def(py::init<const std::string &, ptr<float>, float, Spectrum3f, int,
                      ptr<float>,  ptr<float>,  ptr<float>,  ptr<float>>())
        .def(py::init<const std::string &, ptr<float>, float, const std::string &, int>())
        .def(py::init<const std::string &, ptr<float>, float, const std::string &, int,
                      ptr<float>,  ptr<float>,  ptr<float>>());

    py::class_<Scene>(m, "Scene")
        .def(py::init<const Camera &,
                      const std::vector<const Shape*> &,
                      const std::vector<const BSDF*> &,
                      const std::vector<const Emitter*> &,
                      const std::vector<const PhaseFunction*> &,
                      const std::vector<const Medium*> &>())
        .def(py::init<const Camera &,
                      const std::vector<const Shape*> &,
                      const std::vector<const BSDF*> &,
                      const std::vector<const Emitter*> &,
                      const std::vector<const PhaseFunction*> &,
                      const std::vector<const Medium*> &,
                      bool>())
        .def("initEdges", &Scene::initEdgesPy)
        .def("initPathSpaceEdges", &Scene::initPathSpaceEdgesPy1_1)
        .def("initPathSpaceEdges", &Scene::initPathSpaceEdgesPy1_2)
        .def("initPathSpaceEdges", &Scene::initPathSpaceEdgesPy2_1)
        .def("initPathSpaceEdges", &Scene::initPathSpaceEdgesPy2_2)
        .def("initPathSpacePrimaryEdges", &Scene::initPathSpacePrimaryEdgesPy1_1);

    py::class_<RenderOptions>(m, "RenderOptions")
        .def(py::init<uint64_t, int, int, int, int, bool>())
        .def(py::init<uint64_t, int, int, int, int, bool, int>())
        .def(py::init<uint64_t, int, int, int, int, bool, int, float>())
        .def_readwrite("seed", &RenderOptions::seed)
        .def_readwrite("spp", &RenderOptions::num_samples)
        .def_readwrite("sppe", &RenderOptions::num_samples_primary_edge)
        .def_readwrite("sppse", &RenderOptions::num_samples_secondary_edge)
        .def_readwrite("max_bounces", &RenderOptions::max_bounces)
        .def_readwrite("quiet", &RenderOptions::quiet)
        .def_readwrite("mode", &RenderOptions::mode)
        .def_readwrite("ddistCoeff", &RenderOptions::ddistCoeff)
        .def_readwrite("sppse0", &RenderOptions::num_samples_secondary_edge_direct)
        .def_readwrite("sppse1", &RenderOptions::num_samples_secondary_edge_indirect)
        .def_readwrite("grad_threshold", &RenderOptions::grad_threshold)
        .def_readwrite("num_ellipsoidal_connections", &RenderOptions::num_ellipsoidal_connections)
        .def_readwrite("sppte", &RenderOptions::num_samples_time_edge)
        .def_readwrite("use_antithetic_boundary", &RenderOptions::use_antithetic_boundary)
        .def_readwrite("use_antithetic_interior", &RenderOptions::use_antithetic_interior);

    py::class_<GuidingOptions>(m, "GuidingOptions")
        .def(py::init<int, const std::vector<int>&>())
        .def(py::init<int, const std::vector<int>&, size_t, size_t>())
        .def(py::init<int, const std::vector<int>&, size_t, size_t, float>())
        .def_readwrite("params", &GuidingOptions::params)
        .def_readwrite("num_cam_path", &GuidingOptions::num_cam_path)
        .def_readwrite("num_light_path", &GuidingOptions::num_light_path)
        .def_readwrite("search_radius", &GuidingOptions::search_radius)
        .def_readwrite("quiet", &GuidingOptions::quiet);

    py::class_<Integrator>(m, "Integrator");
    py::class_<IntegratorAD, Integrator>(m, "IntegratorAD");
    py::class_<IntegratorAD_PathSpace, IntegratorAD>(m, "IntegratorAD_PathSpace");
    py::class_<TofIntegratorAD_PathSpace, IntegratorAD_PathSpace>(m, "TofIntegratorAD_PathSpace");

    py::class_<DirectIntegrator, Integrator>(m, "DirectIntegrator")
        .def(py::init<>())
        .def("render", &DirectIntegrator::render);

    py::class_<PathTracer, Integrator>(m, "PathTracer")
        .def(py::init<>())
        .def("render", &PathTracer::render);

    py::class_<ParticleTracer, Integrator>(m, "ParticleTracer")
        .def(py::init<>())
        .def("render", &ParticleTracer::render);

    py::class_<VolPathTracer, Integrator>(m, "VolPathTracer")
        .def(py::init<>())
        .def("render", &VolPathTracer::render);

    py::class_<VolPathTracerSimple, Integrator>(m, "VolPathTracerSimple")
        .def(py::init<>())
        .def("render", &VolPathTracerSimple::render);

    py::class_<DirectIntegratorAD, IntegratorAD>(m, "DirectAD")
        .def(py::init<>())
        .def("render", &DirectIntegratorAD::render);

    py::class_<InvDirectAD_PathSpace, IntegratorAD_PathSpace>(m, "InvDirectADps")
        .def(py::init<>())
        .def("render", &InvDirectAD_PathSpace::render)
        .def("preprocess", &InvDirectAD_PathSpace::preprocess);

    py::class_<ParticleTracerAD_PathSpace, IntegratorAD_PathSpace>(m, "ParticleTracerADps")
        .def(py::init<>())
        .def("render", &ParticleTracerAD_PathSpace::render)
        .def("preprocess", &ParticleTracerAD_PathSpace::preprocess);

    py::class_<PathTracerAD, IntegratorAD>(m, "PathTracerAD")
        .def(py::init<>())
        .def("render", &PathTracerAD::render);

    py::class_<VolPathTracerAD, IntegratorAD>(m, "VolPathTracerAD")
        .def(py::init<>())
        .def("render", &VolPathTracerAD::render);

    py::class_<DirectIntegratorAD_PathSpace, IntegratorAD_PathSpace>(m, "DirectADps")
        .def(py::init<>())
        .def("render", &DirectIntegratorAD_PathSpace::render)
        .def("preprocess", &DirectIntegratorAD_PathSpace::preprocess);

    py::class_<PathTracerAD_PathSpace, IntegratorAD_PathSpace>(m, "PathTracerADps")
        .def(py::init<>())
        .def("render", &PathTracerAD_PathSpace::render)
        .def("preprocess", &PathTracerAD_PathSpace::preprocess);

    py::class_<BidirectionalPathTracer, IntegratorAD_PathSpace>(m, "BidirectionalPathTracer")
        .def(py::init<>())
        .def("render", &BidirectionalPathTracer::render);

    py::class_<BidirectionalPathTracerAD, IntegratorAD_PathSpace>(m, "BidirectionalPathTracerAD")
        .def(py::init<>())
        .def("render", &BidirectionalPathTracerAD::render)
        .def("preprocess", &BidirectionalPathTracerAD::preprocess);

    py::class_<Binned_PathTracerAD_PathSpace, TofIntegratorAD_PathSpace>(m, "BinnedPathTracerADps")
        .def(py::init<>())
        .def("render", &Binned_PathTracerAD_PathSpace::render)
        .def("preprocess", &Binned_PathTracerAD_PathSpace::preprocess);

    py::class_<Binned_TofPathTracerAD_PathSpace, Binned_PathTracerAD_PathSpace>(m, "BinnedTofPathTracerADps")
        .def(py::init<>())
        .def("render", &Binned_TofPathTracerAD_PathSpace::render)
        .def("preprocess", &Binned_TofPathTracerAD_PathSpace::preprocess);

    py::class_<Binned_ParticleTracerAD_PathSpace, Binned_PathTracerAD_PathSpace>(m, "BinnedParticleTracerADps")
        .def(py::init<>())
        .def("render", &Binned_ParticleTracerAD_PathSpace::render)
        .def("preprocess", &Binned_ParticleTracerAD_PathSpace::preprocess);

    m.attr("nder") = nder;
    m.attr("angleEps") = AngleEpsilon;
    m.attr("edgeEps") = EdgeEpsilon;

    // Unit Tests
    m.def("ad_test", &ad_test, "");
    m.def("volume_test", &volume_test, "");
    m.def("ps_edge_test", &ps_edge_test, "");
    m.def("path_test", &path_test, "");
    m.def("compute_grid", &compute_grid, "");
    m.def("ellipsoid_test", &ellipsoid_test, "");
    m.def("tof_test", &tof_test, "");
    m.def("pif_test", &pif_test, "");
}
