Lighting issues when billboarding mesh with shader

Question

I'm trying to create a grass shader that, when provided with a mesh (a simple quad from multimesh) and texture, will display that texture always facing the camera. That part isn't too hard to do.
However, I also want to transfer lighting from the terrain onto the grass. To do this, I've been giving the grass the same normal as the terrain and then applying the same lighting function as is on the terrain shader. This works great when the grass has not been billboarded. However, once it has, it looses all lighting except for shadows. I suspect this has something to do with the transformations applied to the normal when I modify the model view matrix, but I haven't been able to fix it.

Shader code: (ignore the weird stuff I'm doing with UVs)

shader_type spatial;
render_mode ambient_light_disabled;

uniform vec4 ambientColor:hint_color = vec4(0);
//uniform vec4 albedo:hint_color = vec4(0);
//uniform vec4 tipColor:hint_color = vec4(0);

uniform sampler2D grassTexture:hint_albedo;

const vec3 UP = vec3(0,1,0);

void vertex() {
    // billboard sprite
    //MODELVIEW_MATRIX = INV_CAMERA_MATRIX * mat4(CAMERA_MATRIX[0],CAMERA_MATRIX[1],CAMERA_MATRIX[2],WORLD_MATRIX[3]);
    // billboard y only
    //MODELVIEW_MATRIX = INV_CAMERA_MATRIX * mat4(vec4(normalize(cross(vec3(0.0, 1.0, 0.0), CAMERA_MATRIX[2].xyz)),0.0),vec4(0.0, 1.0, 0.0, 0.0),vec4(normalize(cross(CAMERA_MATRIX[0].xyz, vec3(0.0, 1.0, 0.0))),0.0),WORLD_MATRIX[3]);
   // all taken from https://github.com/godotengine/godot/blob/0d9aecd96760ad85930cced7a40c0e6c3c19f3dc/scene/resources/material.cpp

    NORMAL = UP;
    vec3 vertex = VERTEX;

    // negative to fix some weird issue with mesh being inverted
    vertex.xz *= -INSTANCE_CUSTOM.x;
    vertex.y *= INSTANCE_CUSTOM.y;

    UV = UV2;
    UV2 = INSTANCE_CUSTOM.w > 200.0 ? vec2(0.0, UV2.y) : vec2(1.0, UV2.y);

    VERTEX = vertex;
}

void fragment() {
    ALBEDO = texture(grassTexture, UV).rgb;
    //ALPHA_SCISSOR = texture(grassTexture, UV).a > 0.1 ? 5.0 : 0.0;
}

// overthought toon shader
void light() {
    float lDot = dot(NORMAL, LIGHT)  * ATTENUATION.x;
    float lightIntensity = UV2.x == 1.0 ? (lDot > 0.15 ? 1.0 : (lDot > 0.02 ? 0.5 : 0.0)) : 0.5;
    vec3 light = lightIntensity * LIGHT_COLOR;

    DIFFUSE_LIGHT +=  length(light) > 0.0 ? (ALBEDO * LIGHT_COLOR * (ambientColor.rgb + light)) : (length(DIFFUSE_LIGHT) == 0.0 ? ALBEDO * ambientColor.rgb : vec3(0.0));
}

Shaded terrain works well:

Shaded grass works well without billboarding:

It all comes crashing down when the billboarding is enabled:

Answer 1

I was able to figure it out! As I suspected, the MODELVIEW_MATRIX was interfering with my normals. Once the matrix had been modified, it changed the lighting normals to reflect what was actually being shown on the screen. This meant that since all the quads were now facing towards the camera, they all had the same normal and did not calculate lighting how I wanted. I wanted the lighting to be calculated before the matrix was applied.

To fix this, I first store the original MODELVIEWMATRIX to be used in my lighting method. Then, I apply the vertex transformations manually with the billboarded MODELVIEWMATRIX. In my lighting function I also have to apply normal transformations manually, but this time with the original MODELVIEW_MATRIX.

Here's the finished working code, for anyone in the future who has the same issues:

shader_type spatial;
// skip vertex transform so we can apply matrix transformations ourselves
render_mode ambient_light_disabled, skip_vertex_transform;

uniform vec4 ambientColor:hint_color = vec4(0);
uniform sampler2D grassTexture:hint_albedo;

const vec3 UP = vec3(0,1,0);

varying mat4 modelviewMatrix;

void vertex() {
    // save the original MODELVIEW_MATRIX to be applied to lighting
    modelviewMatrix = MODELVIEW_MATRIX;

    NORMAL = UP;
    vec3 vertex = VERTEX;

    vertex.xz *= INSTANCE_CUSTOM.x;
    vertex.y *= INSTANCE_CUSTOM.y;

    UV = UV2;
    UV2 = INSTANCE_CUSTOM.w > 200.0 ? vec2(0.0, UV2.y) : vec2(1.0, UV2.y);
    VERTEX = vertex;
    // apply matrix transformation on our vertices so quad faces the camera
    // from https://docs.godotengine.org/en/stable/tutorials/shaders/shader_reference/spatial_shader.html#vertex-built-ins and https://github.com/godotengine/godot/blob/0d9aecd96760ad85930cced7a40c0e6c3c19f3dc/scene/resources/material.cpp
    VERTEX = ((INV_CAMERA_MATRIX * mat4(CAMERA_MATRIX[0],CAMERA_MATRIX[1],CAMERA_MATRIX[2],WORLD_MATRIX[3])) * vec4(VERTEX, 1.0)).xyz;

}

void fragment() {
    ALBEDO = texture(grassTexture, UV).rgb;
    ALPHA_SCISSOR = texture(grassTexture, UV).a > 0.1 ? 5.0 : 0.0;
}

void light() {
    // apply matrix transformations for lighting normals, however this time with the original MODELVIEW_MATRIX instead of the billboard one
    vec3 normal = normalize((modelviewMatrix * vec4(NORMAL, 0.0)).xyz);
    float lDot = dot(normal, LIGHT)  * ATTENUATION.x;
    float lightIntensity = UV2.x == 1.0 ? (lDot > 0.15 ? 1.0 : (lDot > 0.02 ? 0.5 : 0.0)) : 0.5;
    vec3 light = lightIntensity * LIGHT_COLOR;

    DIFFUSE_LIGHT +=  length(light) > 0.0 ? (ALBEDO * LIGHT_COLOR * (ambientColor.rgb + light)) : (length(DIFFUSE_LIGHT) == 0.0 ? ALBEDO * ambientColor.rgb : vec3(0.0));
}

Still needs a lot of work but it's so much closer!