import { chunk_lpv_get_irradiance_at } from "./chunk_lpv_get_irradiance_at.js";

// language=GLSL
export default `
#ifndef LPV_SHADER_CHUNK_COMMON
#define LPV_SHADER_CHUNK_COMMON


uniform usampler2D lpv_t_mesh_vertices;
uniform usampler2D lpv_t_mesh_neighbours;
uniform usampler3D lpv_t_mesh_lookup;

uniform sampler2D lpv_t_probe_positions;
uniform sampler2D lpv_t_probe_data;

uniform sampler2D lpv_t_probe_depth;
uniform uint lpv_u_probe_depth_resolution;

uniform uint lpv_u_mesh_tet_count;
uniform vec3 lpv_v3_bounds_min;
uniform vec3 lpv_v3_bounds_max;


const float lpv_min_thickness = 0.03; // in Meters
const float lpv_max_thickness = 0.50; // in Meters


#define SEARCH_STEP_LIMIT 32u
#define INVALID_TET 1073741823u


/** Slightly bump the location of the shadow test point away from the shadow casting surface.
    The shadow casting surface is the boundary for shadow, so the nearer an imprecise value is
    to it the more the light leaks.
*/
#define LPV_NORMAL_BIAS 0.0001f

ivec2 lpv_index_to_256_coordinate(uint index) {

    uint pixel_x = index % 256u;
    uint pixel_y = index / 256u;

    return ivec2(int(pixel_x), int(pixel_y));
}

uvec4 lpv_mesh_getVertices(uint tet_index) {
    ivec2 p = lpv_index_to_256_coordinate(tet_index);

    return texelFetch(lpv_t_mesh_vertices, p, 0);

}

uvec4 lpv_mesh_getNeighbours(uint tet_index) {
    ivec2 p = lpv_index_to_256_coordinate(tet_index);

    return texelFetch(lpv_t_mesh_neighbours, p, 0);
}

vec3[9] lpv_probe_getData(uint probe_index) {


    int slot = int(probe_index % 256u);
    int column = int(probe_index / 256u);

    int offset_x = int(slot * 9);

    vec3[9] result;

    for (int i = 0; i < 9; i++) {
        result[i] = texelFetch(lpv_t_probe_data, ivec2(offset_x + i, column), 0).rgb;
    }

    return result;
}

float sign_not_zero(float x) {
    return x >= 0.0 ? 1.0 : -1.0;
}

vec2 sign_not_zero(vec2 x) {
    return vec2(
        sign_not_zero(x.x),
        sign_not_zero(x.y)
    );
}

vec3 sign_not_zero(vec3 x) {
    return vec3(
        sign_not_zero(x.x),
        sign_not_zero(x.y),
        sign_not_zero(x.z)
    );
}

vec4 quadBlendWieghts(vec2 coords)
{
    vec4 res;
/* 0 0 0
    0 0 0
    1 0 0 */
    res.x = min(1.0f - coords.x, 1.0f - coords.y);
/* 1 0 0
    0 0 0
    0 0 1 */
    res.y = abs(coords.x - coords.y);
/* 0 0 1
    0 0 0
    0 0 0 */
    res.z = min(coords.x, coords.y);
/* 0 0 0
    0 0 1
    0 1 1 */
    res.w = ceil(coords.x - coords.y);
    //res.xyz /= (res.x + res.y + res.z);
    return res;
}

vec2 VecToSphereOct(vec3 v)
{
    float l1norm = abs(v.x) + abs(v.y) + abs(v.z);

    vec2 result = v.xz / l1norm;

    if (v.y < 0.0) {
        result = (1.0 - abs(result.yx)) * sign_not_zero(result.xy);
    }

    return result;
}

float SampleBlended(sampler2D tex, vec2 uv0, vec2 uv1, vec2 uv2, vec4 weights) {

    float samp0 = textureLod(tex, uv0, 0.0).r;
    float samp1 = textureLod(tex, uv1, 0.0).r;
    float samp2 = textureLod(tex, uv2, 0.0).r;

    return samp0 * weights.x + samp1 * weights.y + samp2 * weights.z;
}

vec2 lpv_probe_getDepthTriangular(uint probe_index, vec3 direction) {
    // get offset
    uint depth_tile_resolution = lpv_u_probe_depth_resolution;
    uvec2 atlas_size = uvec2(4096u);

    uint tiles_per_row = atlas_size.x / depth_tile_resolution;

    uint tile_x = probe_index % tiles_per_row;
    uint tile_y = probe_index / tiles_per_row;

    vec2 tile_offset = vec2(
        tile_x * depth_tile_resolution,
        tile_y * depth_tile_resolution
    );

    // convert direction to UV
    vec2 octahedral_uv = clamp(VecToSphereOct(direction) * 0.5 + 0.5, 0.0, 1.0);
    vec2 grid = octahedral_uv * vec2(depth_tile_resolution - 1u);

    vec2 gridFrac = fract(grid);
    vec2 gridFloor = floor(grid);

    vec4 weights = quadBlendWieghts(gridFrac);

    //3 nearest frames
    vec2 frame0 = gridFloor;
    vec2 frame1 = gridFloor + mix(vec2(0, 1), vec2(1, 0), weights.w);
    vec2 frame2 = gridFloor + vec2(1.0, 1.0);

    // move frames to atlas space
    frame0 += tile_offset;
    frame1 += tile_offset;
    frame2 += tile_offset;

    vec2 samp0 = texelFetch(lpv_t_probe_depth, ivec2(frame0), 0).rg;
    vec2 samp1 = texelFetch(lpv_t_probe_depth, ivec2(frame1), 0).rg;
    vec2 samp2 = texelFetch(lpv_t_probe_depth, ivec2(frame2), 0).rg;

    vec2 d0 = samp0 * weights.x;
    vec2 d1 = samp1 * weights.y;
    vec2 d2 = samp2 * weights.z;

    return (d0 + d1 + d2);
}

float lpv_bilinear_lerp(float v00, float v01, float v10, float v11, vec2 fraction) {

    float x0 = mix(v00, v01, fraction.x);
    float x1 = mix(v10, v11, fraction.x);

    return mix(x0, x1, fraction.y);
}

vec2 lpv_bilinear_lerp(vec2 v00, vec2 v01, vec2 v10, vec2 v11, vec2 fraction) {

    vec2 x0 = mix(v00, v01, fraction.x);
    vec2 x1 = mix(v10, v11, fraction.x);

    return mix(x0, x1, fraction.y);
}

vec2 lpv_sample_bilinear(sampler2D tex, ivec2 texel_position, vec2 fraction) {

    float texel_00 = texelFetch(tex, texel_position, 0).r;
    float texel_01 = texelFetch(tex, texel_position + ivec2(1, 0), 0).r;
    float texel_10 = texelFetch(tex, texel_position + ivec2(0, 1), 0).r;
    float texel_11 = texelFetch(tex, texel_position + ivec2(1, 1), 0).r;

    return vec2(
        lpv_bilinear_lerp(
            texel_00, texel_01,
            texel_10, texel_11,
            fraction
        ),
        lpv_bilinear_lerp(
            texel_00 * texel_00, texel_01 * texel_01,
            texel_10 * texel_10, texel_11 * texel_11,
            fraction
        )
    );
}

ivec2 wrapOctahedralTexelCoordinates(const in ivec2 texel, const in int texture_size) {
    ivec2 wrapped = ((texel % texture_size) + texture_size) % texture_size;

    int fx = (abs(texel.x / texture_size) + int(texel.x < 0));
    int fy = (abs(texel.y / texture_size) + int(texel.y < 0));

    if (((fx ^ fy) & 1) != 0) {
        return (texture_size - (wrapped + ivec2(1)));
        return wrapped;
    } else {
        return wrapped;
    }
}

vec2 lpv_probe_getDepthBilinear(uint probe_index, vec3 direction) {
    // get offset
    int depth_tile_resolution = int(lpv_u_probe_depth_resolution);
    ivec2 tile_resolution = ivec2(depth_tile_resolution);
    const ivec2 atlas_size = ivec2(4096);

    int tiles_per_row = atlas_size.x / depth_tile_resolution;

    int tile_x = int(probe_index) % tiles_per_row;
    int tile_y = int(probe_index) / tiles_per_row;

    ivec2 tile_offset = ivec2(
        tile_x * depth_tile_resolution,
        tile_y * depth_tile_resolution
    );

    // convert direction to UV
    vec2 octahedral_uv = clamp(VecToSphereOct(direction) * 0.5 + 0.5, 0.0, 1.0);
    vec2 grid = octahedral_uv * vec2(depth_tile_resolution) - 0.5;

    vec2 gridFrac = fract(grid);

    ivec2 texel_position = ivec2(floor(grid));

    ivec2 tile_p_00;
    ivec2 tile_p_01;
    ivec2 tile_p_10;
    ivec2 tile_p_11;

    tile_p_00 = wrapOctahedralTexelCoordinates(texel_position, depth_tile_resolution);
    tile_p_01 = wrapOctahedralTexelCoordinates(texel_position + ivec2(1, 0), depth_tile_resolution);
    tile_p_10 = wrapOctahedralTexelCoordinates(texel_position + ivec2(0, 1), depth_tile_resolution);
    tile_p_11 = wrapOctahedralTexelCoordinates(texel_position + ivec2(1, 1), depth_tile_resolution);

    vec2 texel_00 = texelFetch(lpv_t_probe_depth, tile_offset + tile_p_00, 0).rg;
    vec2 texel_01 = texelFetch(lpv_t_probe_depth, tile_offset + tile_p_01, 0).rg;
    vec2 texel_10 = texelFetch(lpv_t_probe_depth, tile_offset + tile_p_10, 0).rg;
    vec2 texel_11 = texelFetch(lpv_t_probe_depth, tile_offset + tile_p_11, 0).rg;

    return lpv_bilinear_lerp(
        texel_00, texel_01,
        texel_10, texel_11,
        gridFrac
    );
}

vec3 lpv_probe_getPosition(uint probe_index) {
    return texelFetch(lpv_t_probe_positions, lpv_index_to_256_coordinate(probe_index), 0).rgb;
}

mat3 lpv_mesh_makeMatrix(vec3 p0, vec3 p1, vec3 p2, vec3 p3) {

    return inverse(
        mat3(
            p0 - p3,
            p1 - p3,
            p2 - p3
        )
    );

}

vec4 lpv_mesh_getBarycentricCoordinates(uint tet_index, vec3 position) {
    uvec4 vertices = lpv_mesh_getVertices(tet_index);

    vec3 p0 = lpv_probe_getPosition(vertices[0]);
    vec3 p1 = lpv_probe_getPosition(vertices[1]);
    vec3 p2 = lpv_probe_getPosition(vertices[2]);
    vec3 p3 = lpv_probe_getPosition(vertices[3]);

    mat3 matrix = lpv_mesh_makeMatrix(p0, p1, p2, p3);

    vec3 mult = matrix * (position - p3);

    return vec4(mult, 1.0 - mult.x - mult.y - mult.z);
}

void lpv_walk_to_tetrahedron(
in vec3 position,
in uint tet_guess,
out uint tet_index,
out vec4 weights
) {
    tet_index = tet_guess;

    for (uint i = 0u; i < SEARCH_STEP_LIMIT; i++) {

        weights = lpv_mesh_getBarycentricCoordinates(tet_index, position);

        // Check if we're in the current "best guess" tetrahedron
        if (weights.x >= 0.0 && weights.y >= 0.0 && weights.z >= 0.0 && weights.w >= 0.0) {
            // success
            return;
        }

        uvec4 neighbors = lpv_mesh_getNeighbours(tet_index);

        uint next_tet;

        // Otherwise find the smallest barycentric coord and move in that direction
        if (weights.x < weights.y && weights.x < weights.z && weights.x < weights.w) {
            next_tet = neighbors[0];
        } else if (weights.y < weights.z && weights.y < weights.w) {
            next_tet = neighbors[1];
        } else if (weights.z < weights.w) {
            next_tet = neighbors[2];
        } else {
            next_tet = neighbors[3];
        }

        tet_index = next_tet;
    }

    tet_index = INVALID_TET;
}

float lpv_probe_getVisibilityMask(vec3 position, uint probe_index) {

    vec3 probe_position = lpv_probe_getPosition(probe_index);

    vec3 local_probe_offset = position - probe_position;

    float distToProbe = length(local_probe_offset);

    vec3 direction = local_probe_offset / distToProbe;

    vec2 temp = lpv_probe_getDepthBilinear(probe_index, direction);

    float mean = temp.x;
    float mean2 = temp.y; // mean of squared distances

    float variance = abs(mean * mean - mean2);

    // http://www.punkuser.net/vsm/vsm_paper.pdf; equation 5
    // Need the max in the denominator because biasing can cause a negative displacement
    float distance_delta = max(distToProbe - mean, 0.0);

    float chebyshevWeight = variance / (variance + distance_delta * distance_delta);

    // Increase contrast in the weight
    chebyshevWeight = max(chebyshevWeight * chebyshevWeight * chebyshevWeight, 0.0);

    return (distToProbe <= mean) ? 1.0 : chebyshevWeight;
}

vec4 lvp_mask_weights_by_visibility_by_depth(vec3 position, uint tet_index, vec4 weights) {

    uvec4 vertices = lpv_mesh_getVertices(tet_index);

    vec4 visibility = vec4(
        lpv_probe_getVisibilityMask(position, vertices[0]),
        lpv_probe_getVisibilityMask(position, vertices[1]),
        lpv_probe_getVisibilityMask(position, vertices[2]),
        lpv_probe_getVisibilityMask(position, vertices[3])
    );

    return visibility * weights;
}


vec4 lvp_mask_weights_by_visibility_by_normal(vec3 position, vec3 normal, uint tet_index, vec4 weights) {

    uvec4 vertices = lpv_mesh_getVertices(tet_index);

    vec4 visibility;

    for (int i = 0; i < 4; i++) {
        vec3 probe_position = lpv_probe_getPosition(vertices[i]);

        vec3 direction_to_probe = position - probe_position;

        visibility[i] = step(dot(direction_to_probe, normal), 0.0);
    }

    return visibility * weights;
}


vec4 lvp_mask_weights_by_visibility(in vec3 position, in vec3 normal, in vec3 view_direction, in uint tet_index, in vec4 barycentric) {

    uvec4 vertices = lpv_mesh_getVertices(tet_index);

    vec4 visibility;

    // Bias the position at which visibility is computed; this
    // avoids performing a shadow test *at* a surface, which is a
    // dangerous location because that is exactly the line between
    // shadowed and unshadowed. If the normal bias is too small,
    // there will be light and dark leaks. If it is too large,
    // then samples can pass through thin occluders to the other
    // side (this can only happen if there are MULTIPLE occluders
    // near each other, a wall surface won't pass through itself.)
    vec3 lookup_position = position + (normal - view_direction * 3.0) * LPV_NORMAL_BIAS;

    float weight_sum = 0.0;

    for (uint i = 0u; i < 4u; i++) {

        float weight = 1.0;

        uint probe_index = vertices[i];

        vec3 probe_position = lpv_probe_getPosition(probe_index);


        // Smooth backface test
        {

            // Computed without the biasing applied to the "dir" variable.
            // This test can cause reflection-map looking errors in the image
            // (stuff looks shiny) if the transition is poor.
            vec3 direction_to_probe = normalize(probe_position - position);

            // The naive soft backface weight would ignore a probe when
            // it is behind the surface. That's good for walls. But for small details inside of a
            // room, the normals on the details might rule out all of the probes that have mutual
            // visibility to the point. So, we instead use a "wrap shading" test below inspired by
            // NPR work.
//            float backface_term = max(0.0001, dot(direction_to_probe, normal));

            // The small offset at the end reduces the "going to zero" impact
            // where this is really close to exactly opposite
            float backface_term = max(0.0001, (dot(direction_to_probe, normal) + 1.0) * 0.5);
            weight *= backface_term * backface_term + 0.05;
//            weight *= backface_term;

        }

        // Moment visibility test (depth)
        {

//            weight *= lpv_probe_getVisibilityMask(lookup_position, probe_index);

        }


        // A tiny bit of light is really visible due to log perception, so
        // crush tiny weights but keep the curve continuous. This must be done
        // before the trilinear weights, because those should be preserved.
        const float crushThreshold = 0.2;

        if (weight < crushThreshold) {
            weight *= weight * weight * (1.0 / (crushThreshold * crushThreshold));
        }

        // Avoid zero weight
        weight = max(0.000001, weight);

        weight *= barycentric[i];

        weight_sum += weight;

        visibility[i] = weight;
    }

    // normalize
    visibility /= weight_sum;

    return visibility;
}

vec3[9] lpv_interpolate_probes(vec4 weights, uint tet_index) {

    uvec4 vertices = lpv_mesh_getVertices(tet_index);

    vec3[9] probe0 = lpv_probe_getData(vertices[0]);
    vec3[9] probe1 = lpv_probe_getData(vertices[1]);
    vec3[9] probe2 = lpv_probe_getData(vertices[2]);
    vec3[9] probe3 = lpv_probe_getData(vertices[3]);

    vec3[9] result;

    for (int i = 0; i < 9; i++) {

        result[i] = probe0[i] * weights[0]
        + probe1[i] * weights[1]
        + probe2[i] * weights[2]
        + probe3[i] * weights[3];

    }

    return result;
}

uint lpv_guess_initial_tet(vec3 position) {

    vec3 lpv_mesh_bounds_min = lpv_v3_bounds_min;
    vec3 lpv_mesh_bounds_max = lpv_v3_bounds_max;

    vec3 lookup_coordinates = (position - lpv_mesh_bounds_min) / (lpv_mesh_bounds_max - lpv_mesh_bounds_min);

    return textureLod(lpv_t_mesh_lookup, lookup_coordinates, 0.0).r;

}

${ chunk_lpv_get_irradiance_at }

vec4 lpv_renormalize_weights(in vec4 source) {

    float sum = source.x + source.y + source.z + source.w;

    if (sum <= 0.0001) {
        return vec4(0.0);
    }

    return source / sum;

}


vec3 lpv_mesh_interpolate_probe_irradiance(vec3 direction , vec4 weights, uint tet){

    uvec4 vertices = lpv_mesh_getVertices(tet);

    vec3[9] probe0 = lpv_probe_getData(vertices[0]);
    vec3[9] probe1 = lpv_probe_getData(vertices[1]);
    vec3[9] probe2 = lpv_probe_getData(vertices[2]);
    vec3[9] probe3 = lpv_probe_getData(vertices[3]);
    
    vec3 irradiance_0 = lpv_get_irradiance_at(direction, probe0);
    vec3 irradiance_1 = lpv_get_irradiance_at(direction, probe1);
    vec3 irradiance_2 = lpv_get_irradiance_at(direction, probe2);
    vec3 irradiance_3 = lpv_get_irradiance_at(direction, probe3);
    
    // move irradiance into preceptional space for interpolation
    vec3 irradiance_p_0 = sqrt(irradiance_0);
    vec3 irradiance_p_1 = sqrt(irradiance_1);
    vec3 irradiance_p_2 = sqrt(irradiance_2);
    vec3 irradiance_p_3 = sqrt(irradiance_3);

    vec3 interpolated_sqrt = irradiance_p_0 * weights[0]
    +irradiance_p_1 * weights[1]
    +irradiance_p_2 * weights[2]
    +irradiance_p_3 * weights[3]
    ;

    // convert back
    return interpolated_sqrt*interpolated_sqrt;
}

vec3 lpv_sample_irradiance(vec3 position, vec3 normal, vec3 view_direction) {

    // lookup nearby tet
    vec3 lpv_mesh_bounds_min = lpv_v3_bounds_min;
    vec3 lpv_mesh_bounds_max = lpv_v3_bounds_max;
    vec3 lookup_coordinates = (position - lpv_mesh_bounds_min) / (lpv_mesh_bounds_max - lpv_mesh_bounds_min);


    uint nearest_tet = lpv_guess_initial_tet(position);

    uint tet;
    vec4 barycentric_coordinates;
    lpv_walk_to_tetrahedron(position, nearest_tet, tet, barycentric_coordinates);

    // apply visibility term
    vec4 weights = lvp_mask_weights_by_visibility(position, normal, view_direction, tet, barycentric_coordinates);

    if (tet == INVALID_TET) {
        // do nothing
        return vec3(0.0);

    } else {

        vec3[9] lpv_values = lpv_interpolate_probes(weights, tet);

        vec3 irradiance = lpv_get_irradiance_at(normal, lpv_values);

        return irradiance;
    }
}

#endif
`