#include "ps_edge_test.h"
#include <fstream>
#include "math_func.h"
#include "shape.h"
#include "scene.h"
#include "sampler.h"
#include "camera.h"
#include "rayAD.h"
#include "../emitter/area.h"
#include "../bsdf/diffuse.h"


static void test0()
{
    constexpr int nsamples = 10000;
    RndSampler sampler(321, 0);

    Vector n0, n1;
    for ( ; ; ) {
        n0 = squareToUniformSphere(sampler.next2D());
        n1 = squareToUniformSphere(sampler.next2D());
        Float val = n0.dot(n1);
        if ( val > -0.9f && val < 0.9f ) break;
    }

    std::ofstream fout("output0.txt");
    fout.setf(std::ios::scientific);
    fout.precision(6);

    fout << n0[0] << ' ' << n0[1] << ' ' << n0[2] << '\n' << n1[0] << ' ' << n1[1] << ' ' << n1[2] << std::endl;
    for ( int i = 0; i < nsamples; ++i ) {
        Float pdf;
        Vector dir = squareToEdgeRayDirection(sampler.next2D(), n0, n1, pdf);
        fout << dir[0] << ' ' << dir[1] << ' ' << dir[2] << ' ' << pdf << std::endl;
    }
}


static void test1()
{
    constexpr int nsamples = 1000;

    AreaLight area(0, Spectrum3f(10.0f, 10.0f, 10.0f));
    DiffuseBSDF diffuseBSDF(Spectrum3f(0.5f, 0.5f, 0.5f));
    Camera camera;
    RndSampler sampler(123, 0);

    float cubeVtxPositions[] = {
        -0.5f, -0.5f, -0.5f,
         0.5f, -0.5f, -0.5f,
         0.5f,  0.5f, -0.5f,
        -0.5f,  0.5f, -0.5f,
        -0.5f, -0.5f,  0.5f,
         0.5f, -0.5f,  0.5f,
         0.5f,  0.5f,  0.5f,
        -0.5f,  0.5f,  0.5f
    };

    int cubeVtxIndices[] = {
        0, 2, 1,
        0, 3, 2,
        0, 1, 5,
        0, 5, 4,
        1, 2, 6,
        1, 6, 5,
        4, 5, 6,
        4, 6, 7,
        3, 6, 2,
        3, 7, 6,
        0, 7, 3,
        0, 4, 7
    };

    Shape cube(cubeVtxPositions, cubeVtxIndices, nullptr, nullptr, 8, 12, -1, 0, -1, -1);

    Scene scene(camera,
                std::vector<const Shape*>{&cube},
                std::vector<const BSDF*>{&diffuseBSDF},
                std::vector<const Emitter*>{&area},
                std::vector<const PhaseFunction*>(),
                std::vector<const Medium*>());

    std::ofstream fout("output1.txt");
    fout.setf(std::ios::scientific);
    fout.precision(6);

    for ( int i = 0; i < nsamples; ++i ) {
        int shape_id;
        RayAD ray;
        Float pdf;

        const Edge &edge = scene.sampleEdgeRay(sampler.next3D(), shape_id, ray, pdf);
        const VectorAD &n0 = scene.shape_list[shape_id]->faceNormals[edge.f0],
                       &n1 = scene.shape_list[shape_id]->faceNormals[edge.f1];
        assert(math::signum(n0.val.dot(ray.dir.val))*math::signum(n1.val.dot(ray.dir.val)) < -0.5f);

        fout << ray.org.val[0] << ' ' << ray.org.val[1] << ' ' << ray.org.val[2] << ' '
             << ray.dir.val[0] << ' ' << ray.dir.val[1] << ' ' << ray.dir.val[2] << std::endl;
    }
}


static void test2()
{
    constexpr int nsamples = 5000;

    AreaLight area(0, Spectrum3f(10.0f, 10.0f, 10.0f));
    DiffuseBSDF diffuseBSDF(Spectrum3f(0.5f, 0.5f, 0.5f));
    Camera camera;
    RndSampler sampler(123, 0);

    float cubeVtxPositions1[] = {
        -0.5f, -0.5f, -0.5f,
         0.5f, -0.5f, -0.5f,
         0.5f,  0.5f, -0.5f,
        -0.5f,  0.5f, -0.5f,
        -0.5f, -0.5f,  0.5f,
         0.5f, -0.5f,  0.5f,
         0.5f,  0.5f,  0.5f,
        -0.5f,  0.5f,  0.5f
    };

    float cubeVtxPositions2[] = {
        -0.5f, -0.5f,  1.0f,
         0.5f, -0.5f,  1.0f,
         0.5f,  0.5f,  1.0f,
        -0.5f,  0.5f,  1.0f,
        -0.5f, -0.5f,  2.0f,
         0.5f, -0.5f,  2.0f,
         0.5f,  0.5f,  2.0f,
        -0.5f,  0.5f,  2.0f
    };

    int cubeVtxIndices[] = {
        0, 2, 1,
        0, 3, 2,
        0, 1, 5,
        0, 5, 4,
        1, 2, 6,
        1, 6, 5,
        4, 5, 6,
        4, 6, 7,
        3, 6, 2,
        3, 7, 6,
        0, 7, 3,
        0, 4, 7
    };

    Shape cube1(cubeVtxPositions1, cubeVtxIndices, nullptr, nullptr, 8, 12, -1, 0, -1, -1),
          cube2(cubeVtxPositions2, cubeVtxIndices, nullptr, nullptr, 8, 12, -1, 0, -1, -1);

    Scene scene(camera,
                std::vector<const Shape*>{&cube1, &cube2},
                std::vector<const BSDF*>{&diffuseBSDF},
                std::vector<const Emitter*>{&area},
                std::vector<const PhaseFunction*>(),
                std::vector<const Medium*>());

    std::ofstream fout("output2.txt");
    fout.setf(std::ios::scientific);
    fout.precision(6);

    for ( int i = 0; i < nsamples; ++i ) {
        PositionSamplingRecord pRec;
        Vector2i ind = scene.samplePosition(sampler.next2D(), pRec);
        fout << pRec.p[0] << ' ' << pRec.p[1] << ' ' << pRec.p[2] << ' '
             << pRec.n[0] << ' ' << pRec.n[1] << ' ' << pRec.n[2] << ' '
             << ind[0] << ' ' << ind[1] << std::endl;
    }
}


void ps_edge_test()
{
    test0();
    test1();
    test2();
}
