ForwardTracer

#pragma once

#include "SamplerBase.h"
#include "Camera.h"
#include <array>

class BidirSampler: public SamplerBase
{
private:
    struct Event
    {
        FPoint origin;
        FDisp norm;
        FDisp indir;
        FDisp outdir;
        float passrate;
        int face;
    };

    struct Emitter
    {
        float fraction;
        float area;
        float power;
        float dens;
    };

private:
    Camera camera;
    std::vector<Emitter> emitters;
    std::vector<Event> lightpath;
    std::vector<Event> lenspath;

    void PrepareLightStartGenerator();
    void GenerateTriangleInterior(float& u, float& v);
    void GenerateLightSubpath();
    void GenerateLensSubpath(int ix, int iy);
    float Geometric(Event const& from, Event const& to);
    float VisibleGeometric(Event const& from, Event const& to);
    double PathDensity(std::vector<Event*> const& lightpart, std::vector<Event*> const& lenspart);
    double PathFlux(std::vector<Event*> const& path);
    void RecordContribution(FDisp const& dir, float value);
    void RecordToFrame(int x, int y, float value);

public:
    BidirSampler(int width, int height);
    ~BidirSampler();

    virtual void Iterate() override;
};

////////////////////////////////////////////////////////////////////////////////

#include "BidirSampler.h"
#include <cmath>

float const passrate = 1.0f;
constexpr float invpi = 0.3183099f;

void BidirSampler::PrepareLightStartGenerator()
{
    emitters.resize(scene.FaceCount());
    float totalpower = 0;
    for (int i = 0; i < (int)emitters.size(); ++i) {
        FDisp ab = scene[i].mwt.xunit();
        FDisp ac = scene[i].mwt.yunit();
        FDisp n = cross(ab, ac);
        emitters[i].area = 0.5f * sqrtf(dot(n, n));
        emitters[i].power = scene[i].emissive * emitters[i].area;
        totalpower += emitters[i].power;
    }
    for (int i = 0; i < (int)emitters.size(); ++i) {
        emitters[i].fraction = emitters[i].power / totalpower;
        emitters[i].dens = scene[i].emissive / totalpower;
    }
}

void BidirSampler::GenerateTriangleInterior(float& u, float& v)
{
    GenerateTriangle(u);
    u = fabsf(u);
    GenerateUniform(v);
    v *= 1 - u;
}

void BidirSampler::GenerateLightSubpath()
{
    constexpr int MaxPathLen = 100;
    TraceRequest tr;
    {
        int ei = 0;
        float q;
        GenerateUniform(q);
        while (ei < (int)emitters.size() - 1 && q >= emitters[ei].fraction) {
            q -= emitters[ei].fraction;
            ei += 1;
        }
        float u;
        float v;
        GenerateTriangleInterior(u, v);
        lightpath.resize(1);
        Event& e = lightpath.front();
        e.face = ei;
        e.origin = scene[ei].mwt * FPoint{ u, v, 0 };
        e.norm = scene[ei].mwt.zunit();
        e.indir = e.norm;
        GenerateLambertian(e.norm, e.outdir);
        e.passrate = 1.0f;
        tr.origin = e.origin;
        tr.dir = e.outdir;
    }
    while (lightpath.size() < MaxPathLen && Trace(tr)) {
        lightpath.emplace_back();
        Event& e = lightpath.back();
        e.origin = tr.hit;
        e.indir = -tr.dir;
        e.face = tr.face;
        tr.origin = tr.hit;
        e.norm = scene[tr.face].mwt.zunit();
        GenerateLambertian(e.norm, tr.dir);
        e.outdir = tr.dir;
        e.passrate = scene[tr.face].albedo;
        {
            float q;
            GenerateUniform(q);
            if (e.passrate <= q) {
                break;
            }
        }
    }
}

void BidirSampler::GenerateLensSubpath(int ix, int iy)
{
    constexpr int MaxPathLen = 1;
    float dx;
    float dy;
    GenerateTriangle(dx);
    GenerateTriangle(dy);
    float cx = 2.0f * (ix + dx + 0.5f) / width - 1.0f;
    float cy = 1.0f - 2.0f * (iy + dx + 0.5f) / height;
    float cu = cx * camera.utan;
    float cv = cy * camera.vtan;
    TraceRequest tr;
    tr.origin = camera.mwc.origin();
    tr.dir = norm(camera.mwc * FDisp{ cu, cv, 1.0f });
    {
        lenspath.resize(1);
        Event& e = lenspath.front();
        e.origin = tr.origin;
        e.norm = camera.mwc.zunit();
        e.outdir = e.norm;
        e.indir = tr.dir;
        e.passrate = 1.0f;
        e.face = -1;
    }
    while (lenspath.size() < MaxPathLen && Trace(tr)) {
        lenspath.emplace_back();
        Event& e = lenspath.back();
        e.origin = tr.hit;
        e.outdir = -tr.dir;
        e.face = tr.face;
        tr.origin = tr.hit;
        e.norm = scene[tr.face].mwt.zunit();
        GenerateLambertian(e.norm, tr.dir);
        e.indir = tr.dir;
        e.passrate = scene[tr.face].albedo;
        {
            float q;
            GenerateUniform(q);
            if (e.passrate <= q) {
                break;
            }
        }
    }
}

float BidirSampler::Geometric(Event const& from, Event const& to)
{
    FDisp delta = to.origin - from.origin;
    FDisp dir = norm(delta);
    float cosa = dot(from.norm, dir);
    float cosb = -dot(to.norm, dir);
    float distsqr = dot(delta, delta);
    return cosa * cosb / distsqr;
}

float BidirSampler::VisibleGeometric(Event const& from, Event const& to)
{
    FDisp delta = to.origin - from.origin;
    FDisp dir = norm(delta);
    float cosa = dot(from.norm, dir);
    if (cosa < 0) {
        return 0;
    }
    float cosb = -dot(to.norm, dir);
    if (cosb < 0) {
        return 0;
    }
    if (!scene.Test(from.origin, delta)) {
        return 0;
    }
    float distsqr = dot(delta, delta);
    return cosa * cosb / distsqr;
}

double BidirSampler::PathDensity(std::vector<Event*> const& light, std::vector<Event*> const& lens)
{
    double dens = 1.0f;
    if (light.size() >= 1) {
        dens *= emitters[light[0]->face].dens;
    }
    if (light.size() >= 2) {
        dens *= invpi;
        for (int i = 1; i < (int)light.size() - 1; ++i) {
            float geom = Geometric(*light[i], *light[i - 1]);
            float edens = invpi * light[i]->passrate;
            dens *= geom * edens;
        }
        {
            int last = (int)light.size() - 1;
            float geom = Geometric(*light[last], *light[last - 1]);
            dens *= geom;
        }
    }
    if (lens.size() >= 2) {
        {
            float cosa = dot(lens[0]->norm, lens[0]->indir);
            float uvsqr = cosa * cosa;
            dens *= lens[0]->passrate * uvsqr * uvsqr / (4.0f * camera.utan * camera.vtan);
        }
        for (int i = 1; i < (int)lens.size() - 1; ++i) {
            float geom = Geometric(*lens[i], *lens[i - 1]);
            float edens = invpi * lens[i]->passrate;
            dens *= geom * edens;
        }
        {
            int last = (int)lens.size() - 1;
            float geom = Geometric(*lens[last], *lens[last - 1]);
            dens *= geom;
        }
    }
    return dens;
}

double BidirSampler::PathFlux(std::vector<Event*> const& path)
{
    if (path.size() < 2) {
        return 0.0f;
    }
    double flux;
    if (path[0]->face >= 0) {
        flux = scene[path[0]->face].emissive;
    } else {
        return 0;
    }
    for (int i = 1; i < (int)path.size() - 1; ++i) {
        float geom = Geometric(*path[i], *path[i - 1]);
        float fs = invpi;
        if (path[i]->face >= 0) {
            fs *= scene[path[i]->face].albedo;
        }
        flux *= geom * fs;
    }
    {
        int last = (int)path.size() - 1;
        FDisp delta = path[last]->origin - path[last - 1]->origin;
        FDisp dir = norm(delta);
        float cosa = dot(path[last]->norm, dir);
        float cosb = -dot(path[last - 1]->norm, dir);
        float distsqr = dot(delta, delta);
        float geom = cosa * cosb / distsqr;
        float uvsqr = cosa * cosa;
        float fs = uvsqr * uvsqr / (4.0f * camera.utan * camera.vtan);
        flux *= geom * fs;
    }
    return flux;
}

void BidirSampler::RecordContribution(FDisp const& dir, float value)
{
    FDisp cdir = camera.mcw * dir;
    float cdx, cdy, cdz;
    cdir.unpack(cdx, cdy, cdz);
    if (cdz >= 0) {
        return;
    }
    float cu = cdx / cdz;
    float cv = cdy / cdz;
    float cx = cu / camera.utan;
    float cy = cv / camera.vtan;
    float nx = 0.5f * (1.0f + cx) * width - 0.5f;
    float ny = 0.5f * (1.0f - cy) * height - 0.5f;
    float nfx = floorf(nx);
    float nfy = floorf(ny);
    float dx = nx - nfx;
    float dy = ny - nfy;
    int ix = (int)nfx;
    int iy = (int)nfy;
    RecordToFrame(ix, iy, value * (1 - dx) * (1 - dy));
    RecordToFrame(ix + 1, iy, value * dx * (1 - dy));
    RecordToFrame(ix, iy + 1, value * (1 - dx) * dy);
    RecordToFrame(ix + 1, iy + 1, value * dx * dy);
}

void BidirSampler::RecordToFrame(int x, int y, float value)
{
    if (x >= 0 && x < width && y >= 0 && y < height) {
        frame[y * width + x] += value;
    }
}

BidirSampler::BidirSampler(int width, int height)
    : SamplerBase(width, height)
{
    float aspect = (float)width / height;
    /*scene.LoadFromT3D("D:\\Program Files (x86)\\Unreal Tournament GOTY\\Maps\\CTF-Coret-FlagRoom.t3d");
    camera = Camera::Targeted(
    FPoint{ -160, 228, 410 },
    FPoint{ 288, -384, 240 },
    1.0f * sqrtf(aspect), 1.0f / sqrtf(aspect));*/
    scene.LoadFromT3D("D:\\Program Files (x86)\\Unreal Tournament GOTY\\Maps\\Box.t3d");
    camera = Camera::Targeted(
        FPoint{ -206, -219, 224 },
        FPoint{ -96, 0, 80 },
        1.0f * sqrtf(aspect), 1.0f / sqrtf(aspect));
    /*camera = Camera::Targeted(
        FPoint{ -50, -80, 224 },
        FPoint{ 0, 0, 80 },
        1.0f * sqrtf(aspect), 1.0f / sqrtf(aspect));*/
    PrepareLightStartGenerator();
}

BidirSampler::~BidirSampler()
{
}

void BidirSampler::Iterate()
{
    denominator += passrate;
    for (int iy = 0; iy < height; ++iy) {
        for (int ix = 0; ix < width; ++ix) {
            float q;
            GenerateUniform(q);
            if (passrate > q) {
                GenerateLightSubpath();
                GenerateLensSubpath(ix, iy);
                std::vector<Event*> path{ &lenspath[0] };
                std::vector<Event*> lightsection;
                std::vector<Event*> lenssection{ &lenspath[0] };
                for (int i = 0; i < (int)lightpath.size(); ++i) {
                    lightsection.insert(lightsection.end(), &lightpath[i]);
                    path.insert(path.end() - 1, &lightpath[i]);
                    if (VisibleGeometric(*lightsection.back(), *lenssection.back())) {
                        double dens = PathDensity(lightsection, lenssection);
                        if (dens > 0) {
                            double flux = PathFlux(path);
                            if (path.size() >= 2) {
                                int last = (int)path.size() - 1;
                                FDisp dir = path[last]->origin - path[last - 1]->origin;
                                RecordContribution(dir, (float)(flux / dens));
                            }
                        }
                    }
                }
                /*GenerateLensSubpath(ix, iy);
                std::vector<Event*> path;
                std::vector<Event*> lightsection;
                std::vector<Event*> lenssection;
                for (int i = 0; i < (int)lenspath.size(); ++i) {
                    lenssection.insert(lenssection.end(), &lenspath[i]);
                    path.insert(path.begin(), &lenspath[i]);
                    double dens = PathDensity(lightsection, lenssection);
                    if (dens > 0) {
                        double flux = PathFlux(path);
                        if (path.size() >= 2) {
                            int last = (int)path.size() - 1;
                            FDisp dir = path[last]->origin - path[last - 1]->origin;
                            RecordContribution(dir, (float)(flux / dens));
                        }
                    }
                }*/
            }
        }
    }
}