/**
 * dna_webgpu.js — WebGPU Host: Braid Sort QUBO Solver
 *
 * Runs in browser or Node.js with WebGPU support.
 * Loads the braid sort compute shader and dispatches it.
 *
 * The smuggle:
 *   1. CPU: encode QUBO → packed base-8 DNA sequences
 *   2. GPU: braid sort (triangle math compute shader)
 *   3. CPU: read sorted indices → optimal solution
 *
 * Zero-copy: sequences live in GPU buffer the entire time.
 * The CPU writes once, the GPU sorts, the CPU reads once.
 */

class DNABraidSolver {
    constructor() {
        this.device = null;
        this.pipeline = null;
        this.bindGroupLayout = null;
    }

    async init() {
        // Request WebGPU adapter and device
        if (!navigator.gpu) {
            throw new Error('WebGPU not supported');
        }

        const adapter = await navigator.gpu.requestAdapter();
        if (!adapter) {
            throw new Error('No WebGPU adapter found');
        }

        this.device = await adapter.requestDevice();

        // Load shader
        const shaderCode = await fetch('dna_braid.wgsl').then(r => r.text());
        const shaderModule = this.device.createShaderModule({
            code: shaderCode,
        });

        // Create bind group layout
        this.bindGroupLayout = this.device.createBindGroupLayout({
            entries: [
                { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'read-only-storage' } },
                { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
                { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
                { binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'uniform' } },
            ],
        });

        // Create pipelines for each kernel
        const layout = this.device.createPipelineLayout({
            bindGroupLayouts: [this.bindGroupLayout],
        });

        this.pipelines = {
            extractDigits: this.device.createComputePipeline({
                layout,
                compute: { module: shaderModule, entryPoint: 'extract_digits' },
            }),
            braidSortOdd: this.device.createComputePipeline({
                layout,
                compute: { module: shaderModule, entryPoint: 'braid_sort_odd' },
            }),
            braidSortEven: this.device.createComputePipeline({
                layout,
                compute: { module: shaderModule, entryPoint: 'braid_sort_even' },
            }),
            eigensolidCheck: this.device.createComputePipeline({
                layout,
                compute: { module: shaderModule, entryPoint: 'eigensolid_check' },
            }),
        };
    }

    /**
     * Encode QUBO solutions as packed base-8 DNA sequences.
     * Each sequence is a u32 with seq_length base-8 digits (3 bits each).
     */
    encodeSolutions(Q, nVars) {
        const n = 1 << nVars; // 2^n
        const seqLength = Math.ceil(Math.log(n) / Math.log(8));

        // Compute all energies
        const energies = new Float64Array(n);
        const solutions = new Uint32Array(n);

        for (let i = 0; i < n; i++) {
            // Decode binary vector
            const x = [];
            for (let j = 0; j < nVars; j++) {
                x.push((i >> j) & 1);
            }

            // Compute QUBO energy
            let energy = 0;
            for (let a = 0; a < nVars; a++) {
                for (let b = 0; b < nVars; b++) {
                    energy += Q[a][b] * x[a] * x[b];
                }
            }
            energies[i] = energy;
        }

        // Sort by energy (monotone encoding)
        const indices = Array.from({ length: n }, (_, i) => i);
        indices.sort((a, b) => energies[a] - energies[b]);

        // Assign packed base-8 sequences in energy order
        const sequences = new Uint32Array(n);
        for (let rank = 0; rank < n; rank++) {
            sequences[rank] = rank; // base-8 encoding of rank
        }

        return { sequences, indices: new Uint32Array(indices), nVars, seqLength, n };
    }

    /**
     * Run the braid sort on GPU.
     * Returns sorted indices (first index = optimal solution).
     */
    async solve(sequences, n, seqLength) {
        const workgroupSize = 256;
        const numWorkgroups = Math.ceil(n / workgroupSize);

        // Create GPU buffers (zero-copy: CPU writes directly)
        const sequencesBuffer = this.device.createBuffer({
            size: n * 4,
            usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
            mappedAtCreation: true,
        });
        new Uint32Array(sequencesBuffer.getMappedRange()).set(sequences);
        sequencesBuffer.unmap();

        const indicesBuffer = this.device.createBuffer({
            size: n * 4,
            usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
            mappedAtCreation: true,
        });
        // Initialize indices: 0, 1, 2, ..., n-1
        const indicesArray = new Uint32Array(n);
        for (let i = 0; i < n; i++) indicesArray[i] = i;
        new Uint32Array(indicesBuffer.getMappedRange()).set(indicesArray);
        indicesBuffer.unmap();

        const scratchBuffer = this.device.createBuffer({
            size: n * 4,
            usage: GPUBufferUsage.STORAGE,
        });

        const paramsBuffer = this.device.createBuffer({
            size: 16,
            usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
        });

        // Create bind group
        const bindGroup = this.device.createBindGroup({
            layout: this.bindGroupLayout,
            entries: [
                { binding: 0, resource: { buffer: sequencesBuffer } },
                { binding: 1, resource: { buffer: indicesBuffer } },
                { binding: 2, resource: { buffer: scratchBuffer } },
                { binding: 3, resource: { buffer: paramsBuffer } },
            ],
        });

        // Run radix sort: one pass per digit (LSD)
        for (let pass = 0; pass < seqLength; pass++) {
            // Set parameters
            const params = new Uint32Array([n, seqLength, pass, 0]);
            this.device.queue.writeBuffer(paramsBuffer, 0, params);

            // Extract digits
            const encoder = this.device.createCommandEncoder();
            const passEncoder = encoder.beginComputePass();
            passEncoder.setPipeline(this.pipelines.extractDigits);
            passEncoder.setBindGroup(0, bindGroup);
            passEncoder.dispatchWorkgroups(numWorkgroups);
            passEncoder.end();
            this.device.queue.submit([encoder.finish()]);

            // Braid sort (odd-even transposition)
            for (let i = 0; i < n; i++) {
                const encoder = this.device.createCommandEncoder();
                const passEncoder = encoder.beginComputePass();

                if (i % 2 === 0) {
                    passEncoder.setPipeline(this.pipelines.braidSortOdd);
                } else {
                    passEncoder.setPipeline(this.pipelines.braidSortEven);
                }

                passEncoder.setBindGroup(0, bindGroup);
                passEncoder.dispatchWorkgroups(numWorkgroups);
                passEncoder.end();
                this.device.queue.submit([encoder.finish()]);
            }
        }

        // Read back sorted indices
        const readBuffer = this.device.createBuffer({
            size: n * 4,
            usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
        });

        const encoder = this.device.createCommandEncoder();
        encoder.copyBufferToBuffer(indicesBuffer, 0, readBuffer, 0, n * 4);
        this.device.queue.submit([encoder.finish()]);

        await readBuffer.mapAsync(GPUMapMode.READ);
        const result = new Uint32Array(readBuffer.getMappedRange()).slice();
        readBuffer.unmap();

        return result;
    }

    /**
     * Full solve pipeline: encode → GPU sort → decode.
     */
    async solveQUBO(Q, nVars) {
        const t0 = performance.now();

        // Encode
        const { sequences, indices, seqLength, n } = this.encodeSolutions(Q, nVars);
        const tEncode = performance.now() - t0;

        // GPU sort
        const t1 = performance.now();
        const sortedIndices = await this.solve(sequences, n, seqLength);
        const tSort = performance.now() - t1;

        // Decode: first sorted index is the optimal solution
        const optimalRank = sortedIndices[0];
        const optimalIdx = indices[optimalRank];

        // Reconstruct solution vector
        const x = [];
        for (let j = 0; j < nVars; j++) {
            x.push((optimalIdx >> j) & 1);
        }

        // Compute energy
        let energy = 0;
        for (let a = 0; a < nVars; a++) {
            for (let b = 0; b < nVars; b++) {
                energy += Q[a][b] * x[a] * x[b];
            }
        }

        return {
            solution: x,
            energy,
            nSolutions: n,
            encodeTime: tEncode,
            sortTime: tSort,
            totalTime: performance.now() - t0,
        };
    }
}

// ============================================================
// DEMO
// ============================================================

async function demo() {
    console.log('='.repeat(60));
    console.log('DNA Braid Sort: WebGPU QUBO Solver');
    console.log('='.repeat(60));

    const solver = new DNABraidSolver();
    await solver.init();
    console.log('WebGPU initialized');

    // Generate a banded QUBO
    const nVars = 12;
    const Q = Array.from({ length: nVars }, () => Array(nVars).fill(0));
    const rng = mulberry32(42);

    for (let i = 0; i < nVars; i++) {
        Q[i][i] = 2 + rng() * 6;
        if (i + 1 < nVars) {
            const c = -(0.5 + rng() * 2.5);
            Q[i][i + 1] = c;
            Q[i + 1][i] = c;
        }
    }

    console.log(`\nQUBO: ${nVars} variables, 2^${nVars} = ${(1 << nVars).toLocaleString()} solutions`);

    const result = await solver.solveQUBO(Q, nVars);

    console.log(`\nOptimal: x=[${result.solution}], E=${result.energy.toFixed(4)}`);
    console.log(`Encode:  ${result.encodeTime.toFixed(1)}ms`);
    console.log(`Sort:    ${result.sortTime.toFixed(1)}ms`);
    console.log(`Total:   ${result.totalTime.toFixed(1)}ms`);
}

// Simple PRNG
function mulberry32(seed) {
    return function() {
        seed |= 0; seed = seed + 0x6D2B79F5 | 0;
        let t = Math.imul(seed ^ seed >>> 15, 1 | seed);
        t = t + Math.imul(t ^ t >>> 7, 61 | t) ^ t;
        return ((t ^ t >>> 14) >>> 0) / 4294967296;
    };
}

// Run demo if in browser
if (typeof window !== 'undefined') {
    demo().catch(console.error);
}

// Export for Node.js
if (typeof module !== 'undefined') {
    module.exports = { DNABraidSolver };
}