#include "constants.wgsl";
#include "colormap.wgsl";

@group(2) @binding(0) var<uniform> params: UniformParams;
@group(2) @binding(1) var data_buffer: texture_3d<f32>;

struct UniformParams {
    origin: vec3f
}

struct VertexOutput {
    @location(0) position: vec4f,
    @location(1) r_range: vec2f,
    @location(2) idx: vec3f
}

struct UniformParams {
    // Model-View-Projection matrix
    mvp: mat4x4f,
    origin: vec3f   
}

@vertex
fn vertex(
    @location(0) position: vec3f, 
    @location(1) r_range: vec2f,
    @location(2) idx: vec3f
) -> VertexOutput {

    var out: VertexOutput;

    // Transform position
    out.position = params.mvp * vec4f(position, 1.0);
    out.r_range = r_range;
    out.idx = idx;

    return out;
}

fn polar_forward(cartesian: vec3f) -> vec3f {
    let r = length(cartesian.xy - params.origin.xy);
    let theta = atan2(cartesian.y, cartesian.x);
    let z = cartesian.z;
    return vec3f(r, theta, z);
}

@fragment
fn fragment(input: VertexOutput) -> location(0) vec4f {
    // Sample data texture
    let value = textureSample(data_texture, data_sampler, input.idx).r;
    let ear = polar_forward(input.position);
    var color = linear_colormap(value);

    let r = ear.z;
    // Valid range
    let r_range = input.r_range;

    let outside_lower_bound = step(r, r_range.x);
    let outside_upper_bound = step(r_range.y, r);
    let is_outside = outside_lower_bound + outside_upper_bound;
    color.a *= 1.0 - is_outside;

    return color;
}