#version 450

// Simple voxel advection + noise injection for volumetric clouds.
// This is not a full fluid sim, but provides "fluid-like" evolving density in a voxel grid.

layout(local_size_x = 8, local_size_y = 8, local_size_z = 8) in;

layout(std430, set = 0, binding = 0) readonly buffer DensityIn
{
    float density[];
} vox_in;

layout(std430, set = 0, binding = 1) buffer DensityOut
{
    float density[];
} vox_out;

layout(push_constant) uniform Push
{
    vec4 wind_dt;          // xyz windVelocityLocal (units/sec), w dt_sec
    vec4 volume_size_time; // xyz volume size (units), w time_sec
    vec4 sim_params;       // x dissipation (1/sec), y noiseStrength, z noiseScale, w noiseSpeed
    vec4 emitter_params;   // xyz emitterUVW, w emitterRadius
    ivec4 misc;            // x gridResolution, y volumeType (0=cloud,1=smoke,2=flame)
} pc;

uint hash_u32(uint x)
{
    x ^= x >> 16;
    x *= 0x7feb352du;
    x ^= x >> 15;
    x *= 0x846ca68bu;
    x ^= x >> 16;
    return x;
}

float hash3_to_unit_float(ivec3 p)
{
    uint h = 0u;
    h ^= hash_u32(uint(p.x) * 73856093u);
    h ^= hash_u32(uint(p.y) * 19349663u);
    h ^= hash_u32(uint(p.z) * 83492791u);
    return float(h & 0x00FFFFFFu) / float(0x01000000u);
}

float smoothstep01(float x)
{
    x = clamp(x, 0.0, 1.0);
    return x * x * (3.0 - 2.0 * x);
}

float value_noise3(vec3 p)
{
    ivec3 i0 = ivec3(floor(p));
    ivec3 i1 = i0 + ivec3(1);

    vec3 f = fract(p);
    f = vec3(smoothstep01(f.x), smoothstep01(f.y), smoothstep01(f.z));

    float c000 = hash3_to_unit_float(ivec3(i0.x, i0.y, i0.z));
    float c100 = hash3_to_unit_float(ivec3(i1.x, i0.y, i0.z));
    float c010 = hash3_to_unit_float(ivec3(i0.x, i1.y, i0.z));
    float c110 = hash3_to_unit_float(ivec3(i1.x, i1.y, i0.z));
    float c001 = hash3_to_unit_float(ivec3(i0.x, i0.y, i1.z));
    float c101 = hash3_to_unit_float(ivec3(i1.x, i0.y, i1.z));
    float c011 = hash3_to_unit_float(ivec3(i0.x, i1.y, i1.z));
    float c111 = hash3_to_unit_float(ivec3(i1.x, i1.y, i1.z));

    float x00 = mix(c000, c100, f.x);
    float x10 = mix(c010, c110, f.x);
    float x01 = mix(c001, c101, f.x);
    float x11 = mix(c011, c111, f.x);

    float y0 = mix(x00, x10, f.y);
    float y1 = mix(x01, x11, f.y);

    return mix(y0, y1, f.z);
}

float fbm3(vec3 p)
{
    float sum = 0.0;
    float amp = 0.55;
    float freq = 1.0;
    for (int i = 0; i < 4; ++i)
    {
        sum += amp * value_noise3(p * freq);
        freq *= 2.02;
        amp *= 0.5;
    }
    return clamp(sum, 0.0, 1.0);
}

int idx3(ivec3 c, int res)
{
    return c.x + c.y * res + c.z * res * res;
}

float sample_density_trilinear(vec3 uvw, int res)
{
    uvw = clamp(uvw, vec3(0.0), vec3(1.0));

    int slice = res * res;
    float fres = float(res);
    vec3 g = uvw * (fres - 1.0);

    ivec3 base = ivec3(floor(g));
    base = clamp(base, ivec3(0), ivec3(res - 1));
    vec3 f = fract(g);

    ivec3 b1 = min(base + ivec3(1), ivec3(res - 1));
    ivec3 step = b1 - base; // 0 or 1 per axis

    int baseIndex = base.x + base.y * res + base.z * slice;
    int dx = step.x;
    int dy = step.y * res;
    int dz = step.z * slice;

    float d000 = vox_in.density[baseIndex];
    float d100 = vox_in.density[baseIndex + dx];
    float d010 = vox_in.density[baseIndex + dy];
    float d110 = vox_in.density[baseIndex + dy + dx];

    int baseIndexZ = baseIndex + dz;
    float d001 = vox_in.density[baseIndexZ];
    float d101 = vox_in.density[baseIndexZ + dx];
    float d011 = vox_in.density[baseIndexZ + dy];
    float d111 = vox_in.density[baseIndexZ + dy + dx];

    float x00 = mix(d000, d100, f.x);
    float x10 = mix(d010, d110, f.x);
    float x01 = mix(d001, d101, f.x);
    float x11 = mix(d011, d111, f.x);

    float y0 = mix(x00, x10, f.y);
    float y1 = mix(x01, x11, f.y);

    return mix(y0, y1, f.z);
}

void main()
{
    int res = max(pc.misc.x, 1);
    int vol_type = pc.misc.y;

    ivec3 c = ivec3(gl_GlobalInvocationID.xyz);
    if (c.x >= res || c.y >= res || c.z >= res)
    {
        return;
    }

    // Voxel center in [0,1].
    vec3 uvw = (vec3(c) + vec3(0.5)) / float(res);

    float dt = max(pc.wind_dt.w, 0.0);

    vec3 volSize = max(pc.volume_size_time.xyz, vec3(0.001));
    vec3 wind_uv = pc.wind_dt.xyz / volSize; // normalized per-second

    // Semi-Lagrangian advection: backtrace.
    vec3 back = uvw - wind_uv * dt;
    if (vol_type == 0)
    {
        back.xz = fract(back.xz);         // wrap XZ for continuous motion (clouds)
        back.y = clamp(back.y, 0.0, 1.0); // clamp Y
    }
    else
    {
        back = clamp(back, vec3(0.0), vec3(1.0)); // clamp for localized effects (smoke/flame)
    }

    float advected = sample_density_trilinear(back, res);

    // Dissipation.
    float dissipation = max(pc.sim_params.x, 0.0);
    advected *= exp(-dissipation * dt);

    // Inject new density from procedural noise to keep volume evolving.
    float time = pc.volume_size_time.w;
    float noise_scale = max(pc.sim_params.z, 0.001);
    float noise_speed = pc.sim_params.w;
    float n = fbm3(uvw * noise_scale + vec3(0.0, time * noise_speed, 0.0));
    float injected = 0.0;

    if (vol_type == 0)
    {
        // Clouds: broad slab with continuous XZ wrapping.
        injected = smoothstep(0.55, 0.80, n);

        // Height shaping (keep density within a slab).
        float low = smoothstep(0.00, 0.18, uvw.y);
        float high = 1.0 - smoothstep(0.78, 1.00, uvw.y);
        injected *= clamp(low * high, 0.0, 1.0);
    }
    else
    {
        // Smoke/flame: inject near an emitter in UVW space.
        vec3 e = clamp(pc.emitter_params.xyz, vec3(0.0), vec3(1.0));
        float r = max(pc.emitter_params.w, 1e-4);

        vec3 dpos = uvw - e;
        dpos.y *= 1.5; // squash vertically for a more column-like source
        float dist = length(dpos);
        float shape = 1.0 - smoothstep(r, r * 1.25, dist);

        if (vol_type == 1)
        {
            // Smoke: softer noise threshold.
            injected = smoothstep(0.45, 0.75, n) * shape;
        }
        else
        {
            // Flame: spikier + stronger flicker.
            float f = smoothstep(0.35, 0.90, n);
            injected = (f * f) * shape;
        }
    }

    float rate = clamp(pc.sim_params.y * dt, 0.0, 1.0);
    float out_d = advected;
    if (vol_type == 2)
    {
        // Flames flicker: blend toward injected field (avoid accumulating a "fog").
        out_d = mix(advected, injected, rate);
    }
    else
    {
        out_d = mix(advected, max(advected, injected), rate);
    }
    out_d = clamp(out_d, 0.0, 1.0);

    vox_out.density[idx3(c, res)] = out_d;
}