import Metal

/// Host-side motion analysis, shared by all effects. Every camera frame it produces:
///
///   flow       which way things moved (small 2-channel texture, uv units per frame)
///   particles  a picture of particles that spawn where something moves and drift along the flow
///
/// Effects read these as `in.flow` and `in.particles`.
final class Analysis {
    private static let flowWidth = 160, flowHeight = 90
    private static let stateWidth = 128, stateHeight = 64  // one texel per particle
    private static let particleSpeed: Float = 1.5
    private static let trailFade: Float = 0.88

    // each pair is [this frame, previous frame], swapped at the start of every update
    private var luma: [MTLTexture]
    private var flowTextures: [MTLTexture]
    private var state: [MTLTexture]
    private var trail: [MTLTexture]

    private let lumaPipeline: MTLRenderPipelineState
    private let flowPipeline: MTLRenderPipelineState
    private let movePipeline: MTLRenderPipelineState
    private let fadePipeline: MTLRenderPipelineState
    private let dotPipeline: MTLRenderPipelineState

    var flow: MTLTexture { flowTextures[0] }
    var particles: MTLTexture { trail[0] }

    init(device: MTLDevice, queue: MTLCommandQueue, width: Int, height: Int) throws {
        func pair(_ w: Int, _ h: Int, _ format: MTLPixelFormat) -> [MTLTexture] {
            [makeTarget(device, width: w, height: h, format: format), makeTarget(device, width: w, height: h, format: format)]
        }
        luma = pair(Analysis.flowWidth, Analysis.flowHeight, .r16Float)
        flowTextures = pair(Analysis.flowWidth, Analysis.flowHeight, .rg16Float)
        state = pair(Analysis.stateWidth, Analysis.stateHeight, .rgba32Float)
        trail = pair(width, height, .r8Unorm)

        let source = shaderPrelude + analysisShaders
        lumaPipeline = try makePipeline(device, source: source, fragment: "lumaFragment", format: .r16Float)
        flowPipeline = try makePipeline(device, source: source, fragment: "flowFragment", format: .rg16Float)
        movePipeline = try makePipeline(device, source: source, fragment: "moveFragment", format: .rgba32Float)
        fadePipeline = try makePipeline(device, source: source, fragment: "fadeFragment", format: .r8Unorm)
        dotPipeline = try makePipeline(device, source: source, vertex: "dotVertex", fragment: "dotFragment",
                                       format: .r8Unorm, additive: true)

        // new textures hold garbage; start everything from zero
        if let commands = queue.makeCommandBuffer() {
            for texture in luma + flowTextures + state + trail { commands.pass(into: texture, clear: true) }
            commands.commit()
        }
    }

    func update(_ commands: MTLCommandBuffer, camera: MTLTexture, particles wantsParticles: Bool) {
        luma.swapAt(0, 1)
        commands.pass(into: luma[0]) { $0.fullscreen(lumaPipeline, textures: [camera]) }

        flowTextures.swapAt(0, 1)
        commands.pass(into: flowTextures[0]) {
            $0.fullscreen(flowPipeline, textures: [luma[0], luma[1], flowTextures[1]])
        }

        guard wantsParticles else { return }
        state.swapAt(0, 1)
        commands.pass(into: state[0]) {
            $0.fullscreen(movePipeline, textures: [state[1], flowTextures[0], luma[0], luma[1]],
                          uniforms: [Float.random(in: 0..<1000), Analysis.particleSpeed, 0, 0])
        }
        trail.swapAt(0, 1)
        commands.pass(into: trail[0]) { encoder in
            encoder.fullscreen(fadePipeline, textures: [trail[1]], uniforms: [Analysis.trailFade, 0, 0, 0])
            encoder.setRenderPipelineState(dotPipeline)
            encoder.setVertexTexture(state[0], index: 0)
            encoder.drawPrimitives(type: .point, vertexStart: 0, vertexCount: Analysis.stateWidth * Analysis.stateHeight)
        }
    }
}

private let analysisShaders = """
static float brightness(float3 c) { return dot(c, float3(0.299, 0.587, 0.114)); }
static float hash(float2 p) { return fract(sin(dot(p, float2(127.1, 311.7))) * 43758.5453); }

// Camera -> small, blurred brightness image. Flow is measured on this.
fragment float4 lumaFragment(VertexOut v [[stage_in]], texture2d<float> camera [[texture(0)]]) {
    constexpr sampler s(filter::linear, address::clamp_to_edge);
    float2 texel = 1.0 / float2(camera.get_width(), camera.get_height());
    float sum = 0.0;
    for (int y = 0; y < 4; y++) {
        for (int x = 0; x < 4; x++) {
            sum += brightness(camera.sample(s, v.uv + (float2(x, y) - 1.5) * 2.0 * texel).rgb);
        }
    }
    return float4(sum / 16.0);
}

// Lucas-Kanade optical flow: for each pixel, the shift that best explains how its 7x7
// neighbourhood changed between the previous and the current brightness image.
fragment float4 flowFragment(VertexOut v [[stage_in]],
                             texture2d<float> now [[texture(0)]],
                             texture2d<float> before [[texture(1)]],
                             texture2d<float> flowBefore [[texture(2)]]) {
    constexpr sampler s(filter::linear, address::clamp_to_edge);
    float2 texel = 1.0 / float2(now.get_width(), now.get_height());
    float a = 1e-3, b = 0.0, c = 1e-3, p = 0.0, q = 0.0;  // the 1e-3 keeps flat areas at zero flow
    for (int y = -3; y <= 3; y++) {
        for (int x = -3; x <= 3; x++) {
            float2 uv = v.uv + float2(x, y) * texel;
            float ix = (now.sample(s, uv + float2(texel.x, 0.0)).r - now.sample(s, uv - float2(texel.x, 0.0)).r) * 0.5;
            float iy = (now.sample(s, uv + float2(0.0, texel.y)).r - now.sample(s, uv - float2(0.0, texel.y)).r) * 0.5;
            float it = now.sample(s, uv).r - before.sample(s, uv).r;
            a += ix * ix; b += ix * iy; c += iy * iy;
            p += ix * it; q += iy * it;
        }
    }
    float2 shift = -float2(c * p - b * q, a * q - b * p) / (a * c - b * b);  // in small-image pixels
    float len = length(shift);
    if (len > 4.0) { shift *= 4.0 / len; }
    float2 smoothed = mix(flowBefore.sample(s, v.uv).xy, shift * texel, 0.6);
    return float4(smoothed, 0.0, 1.0);
}

// One texel per particle: xy = position (uv), z = frames left to live, w = 1 if it was born from motion.
fragment float4 moveFragment(VertexOut v [[stage_in]],
                             texture2d<float> state [[texture(0)]],
                             texture2d<float> flow [[texture(1)]],
                             texture2d<float> now [[texture(2)]],
                             texture2d<float> before [[texture(3)]],
                             constant float4 &u [[buffer(0)]]) {  // x = random seed, y = speed
    constexpr sampler s(filter::linear, address::clamp_to_edge);
    float4 particle = state.read(uint2(v.position.xy));
    float2 seed = v.position.xy * 0.37 + u.x;
    float2 pos = particle.xy + flow.sample(s, particle.xy).xy * u.y;
    pos += (float2(hash(seed), hash(seed + 7.3)) - 0.5) * 0.0015;
    float life = particle.z - 1.0;
    float born = particle.w;

    if (life <= 0.0 || pos.x < 0.0 || pos.x > 1.0 || pos.y < 0.0 || pos.y > 1.0) {
        // respawn: try a few random spots and take the first one where something is moving
        pos = float2(hash(seed + 1.1), hash(seed + 2.2));
        born = 0.0;
        life = 2.0 + hash(seed + 5.5) * 4.0;  // nothing moving: stay hidden and try again soon
        for (int i = 0; i < 6; i++) {
            float2 spot = float2(hash(seed + 3.3 + float(i)), hash(seed + 9.9 + float(i) * 1.7));
            if (abs(now.sample(s, spot).r - before.sample(s, spot).r) > 0.04) {
                pos = spot;
                born = 1.0;
                life = 20.0 + hash(seed + 5.5) * 70.0;
                break;
            }
        }
    }
    return float4(pos, life, born);
}

// Particle trails: fade what was there, then add a dot per particle on top.
fragment float4 fadeFragment(VertexOut v [[stage_in]], texture2d<float> trail [[texture(0)]],
                             constant float4 &u [[buffer(0)]]) {  // x = fade
    constexpr sampler s(filter::linear, address::clamp_to_edge);
    return trail.sample(s, v.uv) * u.x;
}

struct DotOut {
    float4 position [[position]];
    float size [[point_size]];
};

vertex DotOut dotVertex(uint id [[vertex_id]], texture2d<float> state [[texture(0)]]) {
    uint width = state.get_width();
    float4 particle = state.read(uint2(id % width, id / width));
    DotOut out;
    out.position = float4(particle.x * 2.0 - 1.0, 1.0 - particle.y * 2.0, 0.0, 1.0);
    if (particle.w < 0.5) { out.position.x = 4.0; }  // hidden particles are moved off screen
    out.size = 3.0;
    return out;
}

fragment float4 dotFragment(DotOut in [[stage_in]], float2 coord [[point_coord]]) {
    return float4(smoothstep(0.5, 0.1, length(coord - 0.5)) * 0.9);
}
"""