//Simple Pass-Through Shader
export const PENDULUM_VERTEX_SHADER_WGSL = `
    attribute position : vec3<f32>;

    @vertex
    fn main(input : VertexInputs) -> FragmentInputs {
      var output : FragmentInputs;
      output.position = vec4<f32>(input.position, 1.0);
      return output;
    }
`;

// --- SHARED DATA STRUCTURES ---
// These structs map exactly to the Float32Array in the TypeScript code.
const SHARED_STRUCTS = `
    struct Params {
        width: f32,
        height: f32,
        time: f32,
        dt: f32,
        g: f32,
        m1: f32,
        m2: f32,
        l1: f32,
        l2: f32,
        damping: f32,
        trailDecay: f32,
        impulseM1: f32,
        impulseM2: f32,
        pad: f32         // <-- Padding for safe 16-byte memory alignment
    };

    struct State {
        theta1: f32,
        theta2: f32,
        v1: f32,
        v2: f32
    };
`;

//Fragment Shader to display the pixel buffer
export const PENDULUM_FRAGMENT_SHADER_WGSL = SHARED_STRUCTS + `
    var<storage, read> pixelBuffer : array<f32>;
    var<storage, read> p : Params;

    @fragment
    fn main(input : FragmentInputs) -> FragmentOutputs {
      let width = u32(p.width);
      let height = u32(p.height);

      if (width == 0u || height == 0u) {
          fragmentOutputs.color = vec4<f32>(0.5, 0.0, 0.0, 1.0);
          return fragmentOutputs;
      }

      let x = u32(input.position.x);
      let y = u32(input.position.y);

      if (x >= width || y >= height) {
          fragmentOutputs.color = vec4<f32>(0.0, 0.0, 0.0, 1.0);
          return fragmentOutputs;
      }

      let index = y * width + x;

      // --- THE MAGIC DECODING ---
      var val = pixelBuffer[index];
      var isLine1 = false;
      var isLine2 = false;

      // 1. Check for overlays (Lines)
      if (val >= 20.0) {
          isLine2 = true;
          val = val - 20.0;
      } else if (val >= 10.0) {
          isLine1 = true;
          val = val - 10.0;
      }

      // 2. Check which trail it is
      var isTrail2 = false;
      if (val >= 2.0) {
          isTrail2 = true;
          val = val - 2.0;
      }

      // 3. What remains is purely the fading intensity (0.0 to 1.0)
      let trailIntensity = val;

      // --- COLORS ---
      let bgColor    = vec3<f32>(0.1, 0.1, 0.15);
      let mass1Color = vec3<f32>(1.0, 0.0, 0.0); // Red
      let mass2Color = vec3<f32>(0.0, 1.0, 0.0); // Green
      let line1Color = vec3<f32>(1.0, 1.0, 0.0); // Yellow
      let line2Color = vec3<f32>(1.0, 0.0, 1.0); // Magenta

      var massColor = mass1Color;
      if (isTrail2) {
          massColor = mass2Color;
      }

      // Calculate background blending with the trail
      var finalColor = mix(bgColor, massColor, clamp(trailIntensity, 0.0, 1.0));

      // Overwrite with the line colors if necessary
      if (isLine1) { finalColor = line1Color; }
      if (isLine2) { finalColor = line2Color; }

      fragmentOutputs.color = vec4<f32>(finalColor, 1.0);
      return fragmentOutputs;
    }
`;

//Math for the double pendulum
//https://en.wikipedia.org/wiki/Double_pendulum
export const PENDULUM_PHYSIC_COMPUTE_SHADER_WGSL = SHARED_STRUCTS + `
    @group(0) @binding(0) var<storage, read_write> state : State;
    @group(0) @binding(1) var<storage, read> p : Params;

    @compute @workgroup_size(1)
    fn main() {
        let t1 = state.theta1;
        let t2 = state.theta2;
        let v1 = state.v1;
        let v2 = state.v2;

        let delta_t = t1 - t2;

        let num1 = -p.g * (2.0 * p.m1 + p.m2) * sin(t1)
                   - p.m2 * p.g * sin(t1 - 2.0 * t2)
                   - 2.0 * sin(delta_t) * p.m2 * (v2 * v2 * p.l2 + v1 * v1 * p.l1 * cos(delta_t));
        let den1 = p.l1 * (2.0 * p.m1 + p.m2 - p.m2 * cos(2.0 * delta_t));
        let a1 = num1 / den1;

        let num2 = 2.0 * sin(delta_t) * (v1 * v1 * p.l1 * (p.m1 + p.m2) + p.g * (p.m1 + p.m2) * cos(t1) + v2 * v2 * p.l2 * p.m2 * cos(delta_t));
        let den2 = p.l2 * (2.0 * p.m1 + p.m2 - p.m2 * cos(2.0 * delta_t));
        let a2 = num2 / den2;

        let new_v1 = (v1 + a1 * p.dt) * p.damping + p.impulseM1;
        let new_v2 = (v2 + a2 * p.dt) * p.damping + p.impulseM2;

        state.v1 = new_v1;
        state.v2 = new_v2;
        state.theta1 = t1 + new_v1 * p.dt;
        state.theta2 = t2 + new_v2 * p.dt;
    }
`;

//Pixel data to visualize the pendulum
export const PENDULUM_RENDER_COMPUTE_SHADER_WGSL = SHARED_STRUCTS + `
    @group(0) @binding(0) var<storage, read_write> pixelBuffer : array<f32>;
    @group(0) @binding(1) var<storage, read> p : Params;
    @group(0) @binding(2) var<storage, read> state : State;

    fn sdSegment(point: vec2<f32>, a: vec2<f32>, b: vec2<f32>) -> f32 {
        let pa = point - a;
        let ba = b - a;
        let h = clamp(dot(pa, ba) / dot(ba, ba), 0.0, 1.0);
        return length(pa - ba * h);
    }

    @compute @workgroup_size(64)
    fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
        let index = global_id.x;
        let width = u32(p.width);
        let height = u32(p.height);

        if (index >= width * height) { return; }

        let x = f32(index % width);
        let y = f32(index / width);
        let uv = vec2<f32>(x / p.width, y / p.height);

        let aspect = p.width / p.height;
        let uv_corr = vec2<f32>(uv.x * aspect, uv.y);

        // --- 1. EXTRACT & DECAY OLD MEMORY ---
        var memory = pixelBuffer[index];

        // Strip line overlays from the previous frame
        if (memory >= 20.0) { memory = memory - 20.0; }
        else if (memory >= 10.0) { memory = memory - 10.0; }

        // Check if the memory belongs to Trail 2
        var isTrail2 = false;
        if (memory >= 2.0) {
            isTrail2 = true;
            memory = memory - 2.0;
        }

        // Apply decay to the pure intensity
        memory = memory * p.trailDecay;

        // --- 2. CALCULATE GEOMETRY ---
        let origin = vec2<f32>(0.5 * aspect, 0.3);
        let displayScale = 0.15;

        let p1 = origin + vec2<f32>(sin(state.theta1), cos(state.theta1)) * p.l1 * displayScale;
        let p2 = p1     + vec2<f32>(sin(state.theta2), cos(state.theta2)) * p.l2 * displayScale;

        let dLine1 = sdSegment(uv_corr, origin, p1);
        let dLine2 = sdSegment(uv_corr, p1, p2);
        let dMass1 = length(uv_corr - p1);
        let dMass2 = length(uv_corr - p2);

        // --- 3. SMART LAYERING ---
        var baseVal = 0.0;

        // Base Layer (Masses & Trails)
        if (dMass1 < 0.02) {
            baseVal = 1.0; // Mass 1 = 1.0 (Trail 1 Max)
        } else if (dMass2 < 0.02) {
            baseVal = 3.0; // Mass 2 = 2.0 (Flag) + 1.0 (Trail 2 Max)
        } else {
            // Write fading memory back
            if (isTrail2) {
                baseVal = memory + 2.0;
            } else {
                baseVal = memory;
            }
        }

        // Overlay Layer (Lines)
        var overlay = 0.0;
        // Don't draw lines over the masses (Clean Z-Index)
        if (dMass1 < 0.02 || dMass2 < 0.02) {
            overlay = 0.0;
        } else if (dLine1 < 0.003) {
            overlay = 10.0;
        } else if (dLine2 < 0.003) {
            overlay = 20.0;
        }

        pixelBuffer[index] = baseVal + overlay;
    }
`;
