#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<SphereHit> SecondaryHits : register(t4, space0);
StructuredBuffer<SphereRay> SecondaryRays : register(t5, space0);
StructuredBuffer<PixelIdData> SkyBox : register(t6, space0);
StructuredBuffer<SphereHit> ForwardHits : register(t7, space0);
StructuredBuffer<SphereRay> ForwardRays : register(t8, space0);
RWStructuredBuffer<PixelIdData> Pixels : 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 SkyBox[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;
    result = SampleSkyBox(heading);

    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);
        float f = length(result + reflection);
        float seeDot = 0.3 + 0.4 * abs(dot(normal, heading));
        result = lerp(result, reflection, seeDot) * SampleTFI(heading, normal, sphere.PositionAndRadius.w );

    }

    return result;
}

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

    float3 heading = ray.Heading.xyz;
    result = float3(0.9, 0.9, 0.9);

    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);
        float f = length(result + reflection);
        float seeDot = 0.3 + 0.4 * abs(dot(normal, heading));
        result = lerp(result, reflection, seeDot) * SampleTFI(heading, normal, sphere.PositionAndRadius.w );
    }

    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;
    float3 heading = ray.Heading.xyz;

    if (hit.SphereId < spheresCount)
    {
        float3 reflection = DirectColor(SecondaryRays[n], SecondaryHits[n]);
        float3 seeThrough = ForwardDirectColor(ForwardRays[n], ForwardHits[n]);
        float3 spot = ray.Position.xyz + heading * hit.RayDistance;
        Sphere sphere = Spheres[hit.SphereId];
        float3 normal = normalize(spot - sphere.PositionAndRadius.xyz);
        float f = length(seeThrough + reflection);
        float seeDot = 0.3 + 0.4 * abs(dot(normal, heading));
        pixel.id.xyz = lerp(seeThrough , reflection ,seeDot) * SampleTFI(heading, normal, sphere.PositionAndRadius.w);
        pixel.id.w = 1;
    }
    else
    {
    //    pixel.id.xyz = DirectColor(ray, hit);
       // pixel.id.xyz = lerp(DirectColor(ray, hit), float3(0.9, 0.9, 0.9), 0.8);
        pixel.id.xyz = float3(0.9, 0.9, 0.9);
    }
    Pixels[n] = pixel;
}
