#define PNANOVDB_HLSL
#define PNANOVDB_ADDRESS_32
#include "PNanoVDB.hlsl"

cbuffer CB0 : register(b1)
{
    float4x4 _InverseViewProjection;  // 64 bytes
    float4 _CameraPos_Density;        // xyz = CameraPos, w = DensityScale
    float4 _BBoxMin_Step;             // xyz = BBoxMin, w = StepSize
    float4 _BBoxMax_MaxSteps;         // xyz = BBoxMax, w = MaxSteps
    float4 _Rotation_Pad_LightSamples; // x = RotationY, yzw = pad, but we'll use differently
    float4 _LightDir_Samples;         // xyz = LightDir, w = LightSamples
};

StructuredBuffer<uint> buf : register(t1);

struct VSInput
{
    float2 position : POSITION;
    float2 texcoord : TEXCOORD0;
};

struct PSInput
{
    float4 position : SV_POSITION;
    float2 texcoord : TEXCOORD0;
    float3 worldPos : TEXCOORD1;
};

struct NanoVolume
{
    pnanovdb_grid_handle_t grid;
    pnanovdb_grid_type_t grid_type;
    pnanovdb_readaccessor_t acc;
};

void initVolume(inout NanoVolume volume)
{
    pnanovdb_grid_handle_t grid;
    grid.address.byte_offset = 0;
    
    pnanovdb_grid_type_t grid_type = pnanovdb_buf_read_uint32(buf, PNANOVDB_GRID_OFF_GRID_TYPE);
    pnanovdb_tree_handle_t tree = pnanovdb_grid_get_tree(buf, grid);
    pnanovdb_root_handle_t root = pnanovdb_tree_get_root(buf, tree);
    pnanovdb_readaccessor_t acc;
    pnanovdb_readaccessor_init(acc, root);

    volume.grid = grid;
    volume.grid_type = grid_type;
    volume.acc = acc;
}

float get_value_coord(inout pnanovdb_readaccessor_t acc, float3 pos)
{
    pnanovdb_vec3_t p = pos;
    pnanovdb_coord_t ijk = pnanovdb_hdda_pos_to_ijk(p);
    pnanovdb_address_t address = pnanovdb_readaccessor_get_value_address(PNANOVDB_GRID_TYPE_FLOAT, buf, acc, ijk);
    return pnanovdb_read_float(buf, address);
}

uint get_dim_coord(inout pnanovdb_readaccessor_t acc, float3 pos)
{
    pnanovdb_vec3_t p = pos;
    pnanovdb_coord_t ijk = pnanovdb_hdda_pos_to_ijk(p);
    return pnanovdb_readaccessor_get_dim(PNANOVDB_GRID_TYPE_FLOAT, buf, acc, ijk);
}

bool get_hdda_hit(inout pnanovdb_readaccessor_t acc, inout float tmin, float3 origin, float3 direction, float tmax, out float valueAtHit)
{
    pnanovdb_vec3_t p_origin = origin;
    pnanovdb_vec3_t p_direction = direction;
    float thit;
    bool hit = pnanovdb_hdda_tree_marcher(
        PNANOVDB_GRID_TYPE_FLOAT,
        buf,
        acc,
        p_origin, tmin,
        p_direction, tmax,
        thit,
        valueAtHit
    );
    tmin = thit;
    return hit;
}

float phase_function()
{
    return 1.0;
}

uint rand_xorshift(uint seed)
{
    seed ^= (seed << 13);
    seed ^= (seed >> 17);
    seed ^= (seed << 5);
    return seed;
}

float random_float(float3 pos)
{
    uint seed = asuint(pos.x + pos.y + pos.z);
    float res = float(rand_xorshift(seed)) * (1.0 / 4294967296.0);
    res = float(rand_xorshift(asuint(res))) * (1.0 / 4294967296.0);
    return res;
}

float volumetric_shadow(float3 pos, float densityScale, inout pnanovdb_readaccessor_t acc)
{
    float lightSamples = _LightDir_Samples.w;
    if (lightSamples < 1) { return 0.0; }
    
    float3 light_dir = _LightDir_Samples.xyz;
    
    float shadow = 1.0;
    float sigmaS = 0.0;
    float sigmaE = 0.0;
    float step_size = 1.0;
    float jitter = 0;
    
    int steps = 10;
    for (int step = 0; step < steps; step++)
    {
        float3 sample_pos = pos + (jitter + step_size) * light_dir;
        
        sigmaS = get_value_coord(acc, sample_pos) * densityScale;
        sigmaE = max(0.000001, sigmaS);
        sigmaE *= 0.3;
        shadow *= exp(-sigmaE * step_size);
        
        step_size *= 2.0;
    }
    
    return shadow;
}

PSInput MainVS(VSInput input)
{
    PSInput output;
    output.position = float4(input.position, 0.0, 1.0);
    output.texcoord = input.texcoord;
    
    float4 worldPosH = mul(_InverseViewProjection, float4(input.position, 0.5, 1.0));
    output.worldPos = worldPosH.xyz / worldPosH.w;
    
    return output;
}

bool intersectBox(float3 origin, float3 dir, float3 boxMin, float3 boxMax, out float tmin, out float tmax)
{
    float3 invDir = 1.0 / dir;
    float3 t1 = (boxMin - origin) * invDir;
    float3 t2 = (boxMax - origin) * invDir;
    
    tmin = max(max(min(t1.x, t2.x), min(t1.y, t2.y)), min(t1.z, t2.z));
    tmax = min(min(max(t1.x, t2.x), max(t1.y, t2.y)), max(t1.z, t2.z));
    
    return tmax >= tmin && tmax > 0;
}

float4 MainPS(PSInput input) : SV_TARGET
{
    float3 rayDir = normalize(input.worldPos - _CameraPos_Density.xyz);
    
    float tmin = 0.01;
    float tmax = 5000.0;
    
    NanoVolume volume;
    initVolume(volume);
    
    float3 color = float3(0, 0, 0);
    float transmittance = 1.0;
    float acc_density = 0.0;
    float3 ambient_light = 0.005;
    
    float _DensityScale = _CameraPos_Density.w;
    float _StepSize = _BBoxMin_Step.w;
    float _MaxSteps = _BBoxMax_MaxSteps.w;
    float _RotationY = _Rotation_Pad_LightSamples.x;
    float _LightSamples = _LightDir_Samples.w;
    
    float cosR = cos(_RotationY);
    float sinR = sin(_RotationY);
    float3x3 invRotY = float3x3(
        cosR, 0, sinR,
        0, 1, 0,
        -sinR, 0, cosR
    );
    
    float3 localCameraPos = mul(invRotY, _CameraPos_Density.xyz);
    float3 localRayDir = mul(invRotY, rayDir);
    
    float not_used;
    bool hit = get_hdda_hit(volume.acc, tmin, localCameraPos, localRayDir, tmax, not_used);
    if (!hit) { return float4(0, 0, 0, 0); }
    
    float skip = 0;
    
    for (int i = 0; i < (int)_MaxSteps; i++) {
        if (tmin >= tmax) break;
        
        float3 localPos = localCameraPos + localRayDir * tmin;
        
        uint dim = get_dim_coord(volume.acc, localPos);
        if (dim > 1) {
            float skip_step = 15.0;
            tmin += skip_step;
            skip = skip_step;
            continue;
        }
        
        float density = get_value_coord(volume.acc, localPos) * _DensityScale;
        
        if (density < 0.01) {
            float skip_step = 5.0;
            tmin += skip_step;
            skip = skip_step;
            continue;
        }
        
        if (skip > 0) {
            tmin -= skip * 0.8;
            localPos = localCameraPos + localRayDir * tmin;
            skip = 0;
        }
        
        float sigmaS = density;
        float sigmaE = max(0.000001, sigmaS);
        acc_density += sigmaS;
        
        float shadow = volumetric_shadow(localPos, _DensityScale, volume.acc);
        float3 S = sigmaS * phase_function() * shadow * float3(1, 1, 1);
        float3 Sint = (S - S * exp(-sigmaE * _StepSize)) / sigmaE;
        color += transmittance * Sint;
        transmittance *= exp(-sigmaE * _StepSize);
        
        if (acc_density > 1.0) break;
        
        if (transmittance < 0.05)
        {
            transmittance = 0;
            break;
        }
        
        tmin += _StepSize;
    }
    
    float3 final_color = (color + ambient_light) * acc_density;
    final_color = pow(final_color, 1.0 / 2.2);
    return float4(final_color, acc_density);
}