// Scene1 effect shader - ShaderToy conversion (raymarching cube & sphere)
// Source: Saturday cubism experiment by skal

#include "sequence_uniforms"
#include "camera_common"
#include "math/color"
#include "math/utils"
#include "math/sdf_shapes"
#include "math/common_utils"
#include "render/raymarching"

@group(0) @binding(2) var<uniform> uniforms: UniformsSequenceParams;
@group(0) @binding(3) var<uniform> camera: CameraParams;

const skyCol = vec3<f32>(0.176, 0.235, 0.25);
const skylineCol = vec3<f32>(0.5, 0.125, 0.025);
const sunCol = vec3<f32>(0.5, 0.163, 0.025);
const diffCol1 = vec3<f32>(0.4, 1.0, 1.0);
const diffCol2 = vec3<f32>(0.325, 1.0, 0.975);

// Lighting (normalized manually)
const sunDir1 = vec3<f32>(0.0, 0.04997, -0.99875); // normalize(0, 0.05, -1)
const lightPos1 = vec3<f32>(10.0, 10.0, 10.0);
const lightPos2 = vec3<f32>(-10.0, 10.0, -10.0);

fn rayPlane(ray: Ray, plane: vec4<f32>) -> f32 {
  return -(dot(ray.origin, plane.xyz) + plane.w) / dot(ray.direction, plane.xyz);
}

fn render0(ray: Ray) -> vec3<f32> {
  var col = vec3<f32>(0.0);
  let y = abs(ray.direction.y);
  var sf = 1.0001 - max(dot(sunDir1, ray.direction), 0.0);
  col += skyCol * pow(y, 8.0);
  col += clamp(vec3<f32>(mix(0.0025, 0.125, tanh_approx(0.005 / sf)) / y) * skylineCol, vec3<f32>(0.0), vec3<f32>(10.0));
  sf *= sf;
  col += sunCol * 0.00005 / sf;

  let tp1 = rayPlane(ray, vec4<f32>(0.0, -1.0, 0.0, 6.0));
  if (tp1 > 0.0) {
    let pos = ray.origin + tp1 * ray.direction;
    let pp = pos.xz;
    let db = sdBox2D(pp, vec2<f32>(5.0, 9.0)) - 3.0;
    col += vec3<f32>(4.0) * skyCol * y * y * smoothstep(0.25, 0.0, db);
    col += vec3<f32>(0.8) * skyCol * exp(-0.5 * max(db, 0.0));
  }

  return clamp(col, vec3<f32>(0.0), vec3<f32>(10.0));
}

const OBJ_BACKGROUND: f32 = 0.0;
const OBJ_CUBE: f32 = 1.0;
const OBJ_SPHERE: f32 = 2.0;
const OBJ_PLANE: f32 = 3.0;

fn df(p: vec3<f32>) -> f32 {
  // Cube
  var pc = p - vec3<f32>(-1.9, 0.0, 0.0);
  let dCube = sdBox(pc, vec3<f32>(1.6));

  // Sphere
  var ps = p - vec3<f32>(1.3, 0.0, 0.0);
  let dSphere = sdSphere(ps, 1.2);

  // Ground plane
  let dPlane = p.y + 1.0;

  // Union
  var d = min(dCube, dSphere);
  d = min(d, dPlane);

  return d;
}

fn dfWithID(p: vec3<f32>) -> RayMarchResult {
  // Cube
  var pc = p - vec3<f32>(-1.9, 0.0, 0.0);
  let dCube = sdBox(pc, vec3<f32>(1.6));

  // Sphere
  var ps = p - vec3<f32>(1.3, 0.0, 0.0);
  let dSphere = sdSphere(ps, 1.2);

  // Ground plane
  let dPlane = p.y + 1.0;

  // Find closest object
  var result: RayMarchResult;
  result.distance = dCube;
  result.object_id = OBJ_CUBE;

  if (dSphere < result.distance) {
    result.distance = dSphere;
    result.object_id = OBJ_SPHERE;
  }

  if (dPlane < result.distance) {
    result.distance = dPlane;
    result.object_id = OBJ_PLANE;
  }

  result.distance_max = result.distance;
  return result;
}

fn boxCol(col: vec3<f32>, nsp: vec3<f32>, rd: vec3<f32>, nnor: vec3<f32>, nrcol: vec3<f32>, nshd1: f32, nshd2: f32) -> vec3<f32> {
  var nfre = 1.0 + dot(rd, nnor);
  nfre *= nfre;

  let nld1 = normalize(lightPos1 - nsp);
  let nld2 = normalize(lightPos2 - nsp);

  var ndif1 = max(dot(nld1, nnor), 0.0);
  ndif1 *= ndif1;

  var ndif2 = max(dot(nld2, nnor), 0.0);
  ndif2 *= ndif2;

  var scol = vec3<f32>(0.0);
  let rf = smoothstep(1.0, 0.9, nfre);
  scol += diffCol1 * ndif1 * nshd1;
  scol += diffCol2 * ndif2 * nshd2;
  scol += 0.1 * (skyCol + skylineCol);
  scol += nrcol * 0.75 * mix(vec3<f32>(0.25), vec3<f32>(0.5, 0.5, 1.0), nfre);

  return mix(col, scol, rf * smoothstep(90.0, 20.0, dot(nsp, nsp)));
}

fn render1(ray: Ray) -> vec3<f32> {
  let skyCol_local = render0(ray);
  var col = skyCol_local;

  var init: RayMarchResult;
  init.distance = 0.0;
  init.distance_max = 0.0;
  init.object_id = OBJ_BACKGROUND;

  let result = rayMarchWithID(ray.origin, ray.direction, init);
  if (result.distance < MAX_RAY_LENGTH) {
    let nsp = reconstructPosition(ray, result);
    let nnor = normalWithID(nsp);

    let nref = reflect(ray.direction, nnor);
    var refl_init: RayMarchResult;
    refl_init.distance = 0.2;
    refl_init.distance_max = 0.2;
    refl_init.object_id = OBJ_BACKGROUND;
    let nrt_result = rayMarchWithID(nsp, nref, refl_init);
    let rRay = Ray(nsp, nref);
    var nrcol = render0(rRay);

    if (nrt_result.distance < MAX_RAY_LENGTH) {
      let nrsp = reconstructPosition(Ray(nsp, nref), nrt_result);
      let nrnor = normalWithID(nrsp);
      let nrref = reflect(nref, nrnor);
      nrcol = boxCol(nrcol, nrsp, nref, nrnor, render0(Ray(nrsp, nrref)), 1.0, 1.0);
    }

    let light_dist1 = distance(lightPos1, nsp);
    let light_dist2 = distance(lightPos2, nsp);
    let nshd1 = mix(0.0, 1.0, shadowWithStoredDistance(nsp, normalize(lightPos1 - nsp), light_dist1));
    let nshd2 = mix(0.0, 1.0, shadowWithStoredDistance(nsp, normalize(lightPos2 - nsp), light_dist2));

    col = boxCol(col, nsp, ray.direction, nnor, nrcol, nshd1, nshd2);
  }

  return col;
}

#include "render/fullscreen_vs"

@fragment fn fs_main(@builtin(position) p: vec4<f32>) -> @location(0) vec4<f32> {
  let coord = getScreenCoord(p, uniforms.resolution);
  let ray = getCameraRay(camera, coord);
  var col = render1(ray);
  col = aces_approx(col);
  col = sRGB(col);
  return vec4<f32>(col, 1.0);
}