#version 300 es
precision highp float;

uniform float iTime;
uniform vec2  iResolution;
out vec4 fragColor;

//–––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––
// PARAMETERS & CONSTANTS
//–––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––
const int   MAX_MARCHING_STEPS = 100;
const float MIN_DIST = 0.001;
const float MAX_DIST = 100.0;

// Object IDs: 0–2 for spheres, 3 for floor.
struct SceneResult {
  float dist;
  int objectId;    // 0,1,2 for spheres; 3 for floor.
  int sphereIndex; // Which sphere (if any) was hit.
};

//–––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––
// SDF FUNCTIONS
//–––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––

// A deformed sphere SDF (to simulate squash/stretch when bouncing).
// 'deform' holds the horizontal (x,z) scale in .x and vertical in .y.
float sdSphereDeformed(vec3 p, float r, vec2 deform) {
  vec3 q = vec3(p.x / deform.x, p.y / deform.y, p.z / deform.x);
  return length(q) - r;
}

// SDF for a horizontal plane (the floor) at y = -1.0.
float sdPlane(vec3 p, float h) {
  return p.y - h;
}

//–––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––
// ANIMATED RIPPLE FUNCTION
//–––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––
// This function computes a ripple that emanates from 'center' and
// depends on 'lt', the time (in seconds) elapsed since impact.
float rippleFromImpact(vec2 pos, vec2 center, float lt) {
    float dist = length(pos - center);
    float rippleSpeed = 1.5;   // How fast the ripple front moves outward.
    float frequency   = 10.0;  // Frequency of the ripple oscillation.
    float damping     = 2.0;   // Controls how quickly the ripple decays.
    float wave = sin(frequency * (dist - rippleSpeed * lt));
    // Envelope: maximum at impact (lt=0) and decays over time.
    float envelope = exp(-damping * lt) * (1.0 - smoothstep(0.0, 0.2, lt));
    return 0.1 * wave * envelope / (1.0 + 10.0 * dist);
}

//–––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––
// SCENE MAP: Three bouncing spheres and a rippled floor with animated ripples
//–––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––
SceneResult map(in vec3 p) {
  SceneResult res;
  res.dist = MAX_DIST;
  res.objectId = -1;
  res.sphereIndex = -1;
  
  // Parameters for the spheres.
  float sphereRadius = 0.3;
  float orbitRadius  = 0.8;
  
  // Process each of the 3 spheres.
  for (int i = 0; i < 3; i++) {
    // Each sphere gets its own phase offset.
    float phase = 6.2831853 * float(i) / 3.0;
    float t_i   = iTime + phase;
    // Bounce using the absolute value makes the motion symmetric.
    float bounce = abs(sin(t_i));
    
    // Compute the sphere’s center. They orbit in the xz–plane and bounce in y.
    vec3 center;
    center.x = orbitRadius * cos(iTime + phase);
    center.z = orbitRadius * sin(iTime + phase);
    center.y = -1.0 + sphereRadius + bounce;
    
    // Simulate squash/stretch on impact.
    float impact   = 1.0 - smoothstep(0.0, 0.2, bounce);
    float deformX  = mix(1.0, 1.2, impact); // horizontal expansion
    float deformY  = mix(1.0, 0.8, impact); // vertical squash
    
    float d = sdSphereDeformed(p - center, sphereRadius, vec2(deformX, deformY));
    if (d < res.dist) {
      res.dist       = d;
      res.objectId   = i;  // sphere id (0,1,2)
      res.sphereIndex = i;
    }
  }
  
  // Floor: basic SDF for a plane at y = -1.0.
  float dPlane = sdPlane(p, -1.0);
  // Add animated ripples from each sphere’s impact.
  // We simulate a ripple for a short duration after each impact.
  float period = 3.14;   // Approximate bounce period (in seconds).
  float T_ripple = 0.5;  // Duration (in seconds) for which a ripple is visible.
  for (int i = 0; i < 3; i++) {
    float phase = 6.2831853 * float(i) / 3.0;
    float t_i   = iTime + phase;
    // Compute time since impact for this sphere.
    float lt = mod(t_i, period);
    if (lt < T_ripple) {
      vec2 sphereXZ;
      sphereXZ.x = orbitRadius * cos(iTime + phase);
      sphereXZ.y = orbitRadius * sin(iTime + phase);
      dPlane += rippleFromImpact(p.xz, sphereXZ, lt);
    }
  }
  
  // Return whichever is closer: a sphere or the floor.
  if (dPlane < res.dist) {
    res.dist       = dPlane;
    res.objectId   = 3; // floor id
    res.sphereIndex = -1;
  }
  
  return res;
}

//–––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––
// ESTIMATING NORMALS
//–––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––
vec3 getNormal(vec3 p) {
  float h = 0.0001;
  return normalize(vec3(
    map(p + vec3(h, 0, 0)).dist - map(p - vec3(h, 0, 0)).dist,
    map(p + vec3(0, h, 0)).dist - map(p - vec3(0, h, 0)).dist,
    map(p + vec3(0, 0, h)).dist - map(p - vec3(0, 0, h)).dist
  ));
}

//–––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––
// RAY MARCHING
//–––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––
vec3 rayMarch(vec3 ro, vec3 rd, out int hitObject, out int sphereIndex) {
  float t = 0.0;
  hitObject   = -1;
  sphereIndex = -1;
  for (int i = 0; i < MAX_MARCHING_STEPS; i++) {
    vec3 p = ro + rd * t;
    SceneResult res = map(p);
    if (res.dist < MIN_DIST) {
      hitObject   = res.objectId;
      sphereIndex = res.sphereIndex;
      return p;
    }
    if (t > MAX_DIST) break;
    t += res.dist;
  }
  return ro + rd * t;
}

//–––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––
// LIGHTING & FRESNEL
//–––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––
vec3 getLightDir(vec3 p) {
  // A light that circles the scene.
  vec3 lightPos = vec3(2.0 * cos(iTime), 2.0, 2.0 * sin(iTime));
  return normalize(lightPos - p);
}

vec3 phongLighting(vec3 p, vec3 normal, vec3 viewDir, vec3 baseColor) {
  vec3 lightDir = getLightDir(p);
  float diff = max(dot(normal, lightDir), 0.0);
  vec3 reflDir = reflect(-lightDir, normal);
  float spec = pow(max(dot(viewDir, reflDir), 0.0), 32.0);
  return baseColor * diff + vec3(1.0) * spec;
}

float fresnel(vec3 I, vec3 N, float ior) {
  float cosi = clamp(dot(I, N), -1.0, 1.0);
  float etai = 1.0, etat = ior;
  if (cosi > 0.0) { float tmp = etai; etai = etat; etat = tmp; }
  float sint = etai / etat * sqrt(max(0.0, 1.0 - cosi * cosi));
  if (sint >= 1.0) {
    return 1.0;
  } else {
    float cost = sqrt(max(0.0, 1.0 - sint * sint));
    cosi = abs(cosi);
    float Rs = ((etat * cosi) - (etai * cost)) / ((etat * cosi) + (etai * cost));
    float Rp = ((etai * cosi) - (etat * cost)) / ((etai * cosi) + (etat * cost));
    return (Rs * Rs + Rp * Rp) / 2.0;
  }
}

//–––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––
// RAY-TRACING: REFLECTIONS & REFRACTIONS
//–––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––
vec3 trace(vec3 ro, vec3 rd) {
  vec3 color = vec3(0.0);
  vec3 attenuation = vec3(1.0);
  
  // Allow up to 2 bounces.
  for (int bounce = 0; bounce < 2; bounce++) {
    int hitObject, sphereIndex;
    vec3 hitPos = rayMarch(ro, rd, hitObject, sphereIndex);
    
    // If nothing is hit, return a dark background.
    if (hitObject == -1) {
      color += attenuation * vec3(0.1);
      break;
    }
    
    vec3 normal  = getNormal(hitPos);
    vec3 viewDir = normalize(-rd);
    
    if (hitObject == 3) {
      // Hit the floor: apply a chequerboard pattern.
      vec2 uvFloor = hitPos.xz * 2.0;
      float checker = mod(floor(uvFloor.x) + floor(uvFloor.y), 2.0);
      vec3 colorA = vec3(0.1, 0.1, 0.1);
      vec3 colorB = vec3(0.8, 0.8, 0.8);
      vec3 base = mix(colorA, colorB, checker);
      
      vec3 lit  = phongLighting(hitPos, normal, viewDir, base);
      color += attenuation * lit;
      rd = reflect(rd, normal);
      ro = hitPos + normal * MIN_DIST * 2.0;
      attenuation *= 0.8;
    } else {
      // Hit one of the spheres: treat it as tinted glass.
      vec3 sphereColor = (sphereIndex == 0) ? vec3(1,0,0) :
                         (sphereIndex == 1) ? vec3(0,1,0) :
                                              vec3(0,0,1);
      float ior = 1.5;
      float kr  = fresnel(rd, normal, ior);
      vec3 reflDir   = reflect(rd, normal);
      vec3 refrDir   = refract(rd, normal, 1.0 / ior);
      
      // Reflection: sample nearby surface shading.
      vec3 reflColor = vec3(0.0);
      {
        int dummy1, dummy2;
        vec3 rHit = rayMarch(hitPos + normal * MIN_DIST * 2.0, reflDir, dummy1, dummy2);
        vec3 rNorm = getNormal(rHit);
        reflColor = phongLighting(rHit, rNorm, -reflDir, sphereColor);
      }
      // Refraction: approximate by sampling the background.
      vec3 refrColor = vec3(0.0);
      {
        int dummy1, dummy2;
        vec3 rHit = rayMarch(hitPos - normal * MIN_DIST * 2.0, refrDir, dummy1, dummy2);
        vec3 rNorm = getNormal(rHit);
        refrColor = phongLighting(rHit, rNorm, -refrDir, sphereColor);
      }
      
      // Mix reflection and refraction using the Fresnel term.
      color += attenuation * mix(refrColor, reflColor, kr);
      break;
    }
  }
  
  return color;
}

//–––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––
// MAIN – Set up camera ray and output final colour.
//–––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––––
void main() {
  // Normalized pixel coordinates (with y scaled by iResolution.y).
  vec2 uv = (gl_FragCoord.xy - 0.5 * iResolution.xy) / iResolution.y;
  
  // A simple camera: positioned at (0,0,3) looking toward -z.
  vec3 ro = vec3(0.0, 0.0, 3.0);
  vec3 rd = normalize(vec3(uv, -1.5));
  
  vec3 col = trace(ro, rd);
  
  fragColor = vec4(col, 1.0);
}