uniform shader image;
uniform float2 imageSize;
uniform float amount;
uniform float time;

const int num_iter = 16;
const float reci_num_iter_f = 1.0 / float(num_iter);
const float gamma = 2.2;
const float MAX_DIST_PX = 200.0;

float2 barrelDistortion(float2 p, float2 amt) {
  p = 2.0 * p - 1.0;
  // float BarrelPower = 1.125;
  const float maxBarrelPower = 3.0;
  float theta = atan(p.y, p.x);
  float radius = length(p);
  radius = pow(radius, 1.0 + maxBarrelPower * amt.x);
  p.x = radius * cos(theta);
  p.y = radius * sin(theta);
  return 0.5 * (p + 1.0);
}

float sat(float t) {
  return clamp(t, 0.0, 1.0);
}

float linterp(float t) {
  return sat(1.0 - abs(2.0 * t - 1.0));
}

float remap(float t, float a, float b) {
  return sat((t - a) / (b - a));
}

float3 spectrum_offset(float t) {
  float3 ret;
  float lo = step(t, 0.5);
  float hi = 1.0 - lo;
  float w = linterp(remap(t, 1.0 / 6.0, 5.0 / 6.0));
  ret = float3(lo, 1.0, hi) * float3(1.0 - w, w, 1.0 - w);

  return pow(ret, float3(1.0 / 2.2));
}

float nrand(float2 n) {
  return fract(sin(dot(n.xy, float2(12.9898, 78.233))) * 43758.5453);
}

float4 lin2srgb(float4 c) {
  return pow(c, float4(gamma));
}

float4 srgb2lin(float4 c) {
  return pow(c, float4(1.0 / gamma));
}

half4 main(float2 coord) {
  float2 uv = coord / imageSize;
  uv.y = 1.0 - uv.y;
  // resolution independent
  float2 max_distort = float2(amount);

  float2 oversiz = barrelDistortion(float2(1, 1), max_distort);
  uv = 2.0 * uv - 1.0;
  uv = uv / (oversiz * oversiz);
  uv = 0.5 * uv + 0.5;

  float4 sumcol = float4(0.0);
  float4 sumw = float4(0.0);
  float rnd = nrand(uv + fract(time));
  for (int i = 0; i < num_iter; ++i) {
    float t = (float(i) + rnd) * reci_num_iter_f;
    float4 w = float4(spectrum_offset(t), 1.0);
    sumw += w;
    float2 sampleUV = barrelDistortion(uv, max_distort * t);
    float2 sampleCoord = imageSize * float2(sampleUV.x, 1.0 - sampleUV.y); 
    sumcol += w * srgb2lin(unpremul(image.eval(sampleCoord)));
  }

  sumcol /= sumw;
  float4 outcol = lin2srgb(sumcol);
  outcol += rnd / 255.0;
  outcol.rgb *= outcol.a;
  return outcol;
}