No lighting and shadows in water shader - Vertex Displacement

Hello everybody, I'm very new to shaders and I've run into an issue. I'm writing a water shader which combines vertex displacement and a toon surface shader (which uses a custom lighting function). Both the toon shader and vertex displacement work as expected, but when combined, I lose all shading and the model is one solid colour. If I turn shaded-wireframe it's clear that the model isn't just a flat plane.

Here's the shader I'm using (a simplified version of the improved toon shader from ronja-tutorials)

 Shader "Custom/Toon SHader" {
     Properties {
         [Header(Colors)]
         _MainColor ("Main Color", Color) = (1,1,1,1)
         _ShadowColor ("Shadow Color", Color) = (0, 0, 0, 1)
         _SpecularColor ("Specular Color", Color) = (1,1,1,1)
         [HDR] _EmissionColor ("Emission Color", color) = (0 ,0 ,0 , 1)
 
         [Header(Shadow Settings)]
         [IntRange] _StepAmount ("Shadow Steps", Range(1, 16)) = 2
         _StepWidth ("Step Size", Range(0, 1)) = 0.25
         //_SpecularSize ("Specular Size", Range(0, 1)) = 0.1
         //_SpecularFalloff ("Specular Falloff", Range(0, 2)) = 1
 
         [Header(Textures)]
         _MainTex ("Texture", 2D) = "white" {}
 
         [Header(Vertex Parameters)]
         _NoiseTex ("Noise Texture", 2D) = "white" {}
         _NoiseIntensity ("Noise Intensity", Range(0, 10)) = 2
     }
     SubShader {
         Tags { "RenderType"="Opaque" }
         LOD 200
 
         CGPROGRAM
         // Physically based Standard lighting model, and enable shadows on all light types
         #pragma surface surf Stepped fullforwardshadows vertex:vert
 
         // Use shader model 3.0 target, to get nicer looking lighting
         #pragma target 3.0
 
         fixed4 _MainColor;
         fixed4 _ShadowColor;
         half3  _EmissionColor;
         fixed4 _SpecularColor;
 
         sampler2D _MainTex;
 
         float _StepWidth;
         float _StepAmount;
 
         sampler2D _NoiseTex;
         float _NoiseIntensity;
 
         void vert (inout appdata_base v)
         {
             float r = tex2Dlod(_NoiseTex, float4(v.texcoord.xy, 0, 0)).r * _NoiseIntensity;
             float g = tex2Dlod(_NoiseTex, float4(v.texcoord.xy, 0, 0)).g * _NoiseIntensity;
             
             float time = _SinTime.w;
             v.vertex.y -= lerp(r, g, time);
         }
 
         // Custom lighting function, must be prefixed by "Lighting"
         float4 LightingStepped(SurfaceOutput s, float3 lightDir, half3 viewDir, float shadowAttenuation){
             // dot is the Dot Product.
             // Returns 1 if the vectors are facing the same way, and -1 if facing opposite ways. (assuming they are all normalized)
             float lightDirection = dot(s.Normal, lightDir);
 
             // Shadow Steps
             lightDirection = lightDirection / _StepWidth;
             float lightIntensity = ceil(lightDirection);
 
 
             lightIntensity = lightIntensity / _StepAmount;
             lightIntensity = saturate(lightIntensity);
 
             // Step returns 0 if the first argument is greater, and 1 otherwise. 
             // In this case, 0 represents the middle ground, and light intensity will either be 1 or 0
             //lightIntensity = step(0, lightDirection);
 
         #ifdef USING_DIRECTIONAL_LIGHT
             //for directional lights, get a hard cut in the middle of the shadow attenuation
             float attenuationChange = fwidth(shadowAttenuation) * 0.5;
             float shadow = smoothstep(0.5 - attenuationChange, 0.5 + attenuationChange, shadowAttenuation);
         #else
             //for other light types (point, spot), put the cutoff near black, so the falloff doesn't affect the range
             float attenuationChange = fwidth(shadowAttenuation);
             float shadow = smoothstep(0, attenuationChange, shadowAttenuation);
         #endif
             lightIntensity = lightIntensity * shadow;
 
             float4 color;
             color.rgb = s.Albedo * lightIntensity * _LightColor0.rgb;
             color.a = 1;
             return color;
         }
 
         struct Input {
             float2 uv_MainTex;
         };
 
         void surf (Input i, inout SurfaceOutput o) {
             //sample and tint albedo texture
             fixed4 col = tex2D(_MainTex, i.uv_MainTex);
             col *= _MainColor;
             o.Albedo = col.rgb;
 
             float3 shadowColor = col.rgb * _ShadowColor;
             o.Emission = _EmissionColor + shadowColor;
         }
         ENDCG
     }
     FallBack "Diffuse"
 }
 

I'm using this shader on a procedural mesh that is just a large, flat plane. It seems like the lighting and shadows are calculated for the flat plane, and then the displacement happens afterwards. Is there a way to force the vertex displacement to happen before the lighting? Or is this an entirely different problem? Thanks in advance!