//VERSION=3

function setup() {
  return {
    input: [{
      bands: ["B02", "B03", "B04", "B08", "B11"],
      units: "REFLECTANCE"
    }],
    output: {
      bands: 3,
      sampleType: "AUTO"
    }
  };
}

const MAX_INPUT = 0.35;
const CLIP_INPUT = 0.72;
const MID_LEVEL = 0.84;

const ROCK_DARKEN = 0.10;
const SNOW_BOOST = 0.025;
const GLOBAL_SAT = 1.01;

function clamp(x, min, max) {
  return Math.max(min, Math.min(max, x));
}

function smoothstep(edge0, edge1, x) {
  let t = clamp((x - edge0) / (edge1 - edge0), 0, 1);
  return t * t * (3 - 2 * t);
}

function mix(a, b, t) {
  return a * (1 - t) + b * t;
}

function compress(x) {
  if (x <= MAX_INPUT) {
    return (x / MAX_INPUT) * MID_LEVEL;
  }
  if (x >= CLIP_INPUT) {
    return 1.0;
  }
  let t = (x - MAX_INPUT) / (CLIP_INPUT - MAX_INPUT);
  return MID_LEVEL + (1.0 - MID_LEVEL) * (1.0 - Math.pow(1.0 - t, 2.0));
}

function saturate(rgb, amount) {
  let luma = rgb[0] * 0.2126 + rgb[1] * 0.7152 + rgb[2] * 0.0722;
  return [
    clamp(luma + (rgb[0] - luma) * amount, 0, 1),
    clamp(luma + (rgb[1] - luma) * amount, 0, 1),
    clamp(luma + (rgb[2] - luma) * amount, 0, 1)
  ];
}

function evaluatePixel(s) {
  let r = compress(s.B04);
  let g = compress(s.B03);
  let b = compress(s.B02);

  let brightness = (r + g + b) / 3.0;
  let ndvi = (s.B08 - s.B04) / (s.B08 + s.B04 + 1e-6);
  let ndsi = (s.B03 - s.B11) / (s.B03 + s.B11 + 1e-6);

  let snowMask =
    smoothstep(0.20, 0.48, ndsi) *
    smoothstep(0.40, 0.75, brightness);

  let vegetationMask = smoothstep(0.22, 0.55, ndvi);

  let rockMask =
    smoothstep(0.18, 0.55, brightness) *
    (1.0 - smoothstep(0.14, 0.38, ndsi)) *
    (1.0 - vegetationMask);

  let rgb = [r, g, b];
  rgb = saturate(rgb, GLOBAL_SAT);

  let rockStrength = ROCK_DARKEN * rockMask;
  rgb = [
    clamp(rgb[0] * (1.0 - rockStrength * 1.03), 0, 1),
    clamp(rgb[1] * (1.0 - rockStrength * 1.00), 0, 1),
    clamp(rgb[2] * (1.0 - rockStrength * 0.96), 0, 1)
  ];

  let snowLift = SNOW_BOOST * snowMask;
  rgb = [
    clamp(mix(rgb[0], rgb[0] * 1.02 + 0.006, snowLift), 0, 1),
    clamp(mix(rgb[1], rgb[1] * 1.02 + 0.006, snowLift), 0, 1),
    clamp(mix(rgb[2], rgb[2] * 1.03 + 0.008, snowLift), 0, 1)
  ];

  return rgb;
}