#include "RayTypes.h"

Texture2D<unorm float4> TFITexture : register(t0, space0);
StructuredBuffer<SphereHit> Hits : register(t1, space0);
StructuredBuffer<SphereRay> Rays : register(t2, space0);
StructuredBuffer<Sphere> Spheres : register(t3, space0);
StructuredBuffer<PixelIdData> PrevSkyBox : register(t4, space0);
RWStructuredBuffer<PixelIdData> SkyBox : register(u0, space1);

cbuffer UBO : register(b0, space2)
{
    int4 ubo_size : packoffset(c0);
    int spheresCount : packoffset(c1.x);
    int4 skybox_size : packoffset(c2);
};


    
float CatFromRadius(float radius)
{
    return log(radius / 0.25) / log(1.26);
}

float3 SampleTFI(float3 heading, float3 normal, float radius)
{
    float d = 1.0 - abs(dot(heading, normal));
    //float f = clamp(1200 - (d * d * 1200) - radius*100.0 - 400, 0, 1200 - 1);

    float f = clamp(d * 100.0 + CatFromRadius(radius) * 50.0 + 200.0, 0, 1200 - 1);

    int l = min(1200 - 1, (int) floor(f));
    int h = min(1200 - 1, (int) ceil(f));

    float3 a = TFITexture[int2(l, 0)].xyz;
    float3 b = TFITexture[int2(h, 0)].xyz;
    return fakehdr(lerp(a, b, f - l));
}

float3 SampleSkyBox(float3 heading)
{
    float2 uv = OctahedronUV(heading);
    int2 pos = (int2) floor(uv.xy * (skybox_size.xy + int2(1, 1)));
    return PrevSkyBox[clamp(pos.x, 0, skybox_size.x - 1) + clamp(pos.y, 0, skybox_size.y - 1) * skybox_size.x].id.xyz;
}

float3 DirectColor(SphereRay ray, SphereHit hit)
{
    float3 result;

    float3 heading = ray.Heading.xyz;

    if (hit.SphereId < spheresCount)
    {
        float3 spot = ray.Position.xyz + heading * hit.RayDistance;
        Sphere sphere = Spheres[hit.SphereId];
        float3 normal = normalize(spot - sphere.PositionAndRadius.xyz);
        float3 reflectionDir = heading - (normal * 2.0 * dot(heading, normal));
        float3 reflection = SampleSkyBox(reflectionDir);
        result = reflection * SampleTFI(heading, normal, sphere.PositionAndRadius.w);
    }
    else
    {
        result = SampleSkyBox(heading);
    }

    return result;
}

[numthreads(64, 1, 1)]
void main(uint3 GlobalInvocationID : SV_DispatchThreadID)
{
    uint n = GlobalInvocationID.x;

    SphereRay ray = Rays[n];
    SphereHit hit = Hits[n];

    PixelIdData pixel;

    pixel.id.xyz = lerp(float3(1.0, 1.0, 1.0), DirectColor(ray, hit), 0.8);

    SkyBox[n] = pixel;
}
