struct VertexOutput {
    @location(0) world_space_position: vec4<f32>,
    @builtin(position) position: vec4<f32>,
}

struct Uniforms {
    camera_to_world_matrix: mat4x4<f32>,
    world_to_ndc_matrix: mat4x4<f32>,
    camera_position: vec4<f32>,
    dem_min_corner: vec2<f32>,
    dem_cell_size: vec2<f32>,
    dem_z_range: vec2<f32>
}

@group(0) @binding(0)
var<uniform> uniforms: Uniforms;

@group(0) @binding(1)
var dem_texture: texture_2d<u32>;

@group(0) @binding(2)
var dembvh_texture: texture_2d<u32>;

//#include ray_intersection.wgsl

@vertex
fn vs_main(
    @location(0) position: vec4<f32>,
) -> VertexOutput {
    var result: VertexOutput;
    result.position = uniforms.world_to_ndc_matrix * position;
    result.world_space_position = position;
    return result;
}

const RGSS_WIGGLE = array(vec2(0.125, 0.375),
                          vec2(-0.125, -0.375),
                          vec2(0.375, 0.125),
                          vec2(-0.375, -0.125));

@fragment
fn fs_solid(vertex: VertexOutput) -> @location(0) vec4<f32> {
    var ray: Ray;
    ray.origin = vertex.world_space_position.xyz;
    ray.direction = normalize(vertex.world_space_position.xyz - uniforms.camera_position.xyz);
    // Spread rays into RGSS antialiasing pattern.
    let ray_dx = dpdy(ray.direction);
    let ray_dy = dpdy(ray.direction);
    //var ray_directions: array<vec3<f32>,4>;
    //ray_directions[0] = ray.direction + 0.125*ray_dx + 0.375*ray_dy;
    //ray_directions[1] = ray.direction - 0.125*ray_dx - 0.375*ray_dy;
    //ray_directions[2] = ray.direction + 0.375*ray_dx + 0.125*ray_dy;
    //ray_directions[3] = ray.direction - 0.375*ray_dx - 0.125*ray_dy;

    // Possibly these ray directions could be stored in a matrix and we could
    // evaluate them all at once instead of looping.

    let sun_direction = vec3<f32>(0.761904762, 0.380952381, 0.19047619);
    
    var color_accumulator = vec4<f32>(0);
    let rgss_value = vec2(0.125, 0.375);
    var wiggled_ray = ray;
    for(var i=0; i<4; i++) {
      let rgss_wiggle = select(rgss_value, rgss_value.yx, vec2(bool(i&2))) * select(1.0, -1.0, bool(i&1));
      wiggled_ray.direction = ray.direction + rgss_wiggle.x * ray_dx + rgss_wiggle.y * ray_dy;

       var root_node: BoundingNode;
       root_node.index = vec2<u32>(0);
       root_node.level = textureNumLevels(dembvh_texture) - 1;

       var hit_location: vec3<f32>;
       var hit_normal: vec3<f32>;
       if intersect_ray_with_node(wiggled_ray, root_node, &hit_location, &hit_normal) {
           var shadow_value = 0.0f;
           var shadow_ray :Ray;
           shadow_ray.origin = hit_location + hit_normal * 0.1;
           for(var j=0; i<8; i++) {
             let shadow_sample = 0.02 * (vec3(1.0, 0.0, 0.0) * f32(bool((i*4+j)&1))
                                         + vec3(0.0, 1.0, 0.0) * f32(bool((i*4+j)&2))
                                         + vec3(-1.0, 0.0, 0.0) * f32(bool((i*4+j)&4))
                                         + vec3(0.0, -1.0, 0.0) * f32(bool((i*4+j)&8))
                                         + vec3(0.0, 0.0, 1.0) * f32(bool((i*4+j)&16)));
               // Calculate light
               let sun_direction =
                 (uniforms.camera_to_world_matrix
                  * vec4<f32>(sun_direction + shadow_sample, 0.0)).xyz;
               shadow_ray.direction = sun_direction;
               var dummy0: vec3<f32>;
               var dummy1: vec3<f32>;
               shadow_value +=
                 select(0.25, 0.0, intersect_ray_with_node(shadow_ray,
                                                              root_node,
                                                              &dummy0,
                                                              &dummy1));
           }
           let sun_direction =
             (uniforms.camera_to_world_matrix
              * vec4<f32>(sun_direction, 0.0)).xyz;
           let sun_lambertian_value = dot(hit_normal, sun_direction);
           let fill_lambertian_value = dot(hit_normal, sun_direction * vec3(1.0, -1.0, 1.0));
           let fill_strength = 0.25;
           let l = fill_strength * fill_lambertian_value + (1.0 - fill_strength) * sun_lambertian_value * shadow_value;
           color_accumulator += 0.25 * vec4<f32>(vec3<f32>(l), 1.0);
        }
    }
    if color_accumulator.a == 0.0 {
        discard;
    }
    return color_accumulator;
}