#include <metal_stdlib>
using namespace metal;

struct Uniforms {
  float curvature;
  float scanline_intensity;
  float vignette;
  float padding;
  float2 resolution;
  float2 padding2;
};

struct FragmentIn {
  float4 position [[position]];
  float2 uv;
};

float2 curve_uv(float2 uv, float curvature) {
  uv = uv * 2.0 - 1.0;
  float2 offset = abs(uv.yx) / float2(curvature, curvature);
  uv = uv + uv * offset * offset;
  uv = uv * 0.5 + 0.5;
  return uv;
}

fragment float4 fragment_main(
  FragmentIn in [[stage_in]],
  constant Uniforms &uniforms [[buffer(0)]],
  texture2d<float> tex [[texture(0)]],
  sampler samp [[sampler(0)]]
) {
  float2 uv = in.uv;
  
  // Apply curvature
  if (uniforms.curvature > 0.0) {
    uv = curve_uv(uv, 6.0 / uniforms.curvature);
  }
  
  // Check if outside screen after curvature
  if (uv.x < 0.0 || uv.x > 1.0 || uv.y < 0.0 || uv.y > 1.0) {
    return float4(0.0, 0.0, 0.0, 1.0);
  }
  
  float4 color = tex.sample(samp, uv);
  
  // Scanlines
  float scanline = sin(uv.y * uniforms.resolution.y * 3.14159) * 0.5 + 0.5;
  scanline = pow(scanline, 1.5);
  color.rgb *= 1.0 - uniforms.scanline_intensity * (1.0 - scanline);
  
  // Vignette
  float2 vig_uv = uv * (1.0 - uv.yx);
  float vig = vig_uv.x * vig_uv.y * 15.0;
  vig = pow(vig, uniforms.vignette);
  color.rgb *= vig;
  
  // Slight RGB shift for chromatic aberration
  float2 offset = (uv - 0.5) * 0.002;
  color.r = tex.sample(samp, uv + offset).r;
  color.b = tex.sample(samp, uv - offset).b;
  
  return color;
}
