import { Fn, int, ivec2 } from 'three/tsl'
import type { Node } from 'three/webgpu'

export const xyz2m = Fn(([xyz]: [Node<'ivec3'>]) => {
        const x = xyz.x.toVar()
        const y = xyz.y.toVar()
        const z = xyz.z.toVar()
        x.assign(x.bitOr(x.shiftLeft(int(16))))
        y.assign(y.bitOr(y.shiftLeft(int(16))))
        z.assign(z.bitOr(z.shiftLeft(int(16))))
        x.assign(x.bitAnd(int(0xff0000ff | 0)))
        y.assign(y.bitAnd(int(0xff0000ff | 0)))
        z.assign(z.bitAnd(int(0xff0000ff | 0)))
        x.assign(x.bitOr(x.shiftLeft(int(8))))
        y.assign(y.bitOr(y.shiftLeft(int(8))))
        z.assign(z.bitOr(z.shiftLeft(int(8))))
        x.assign(x.bitAnd(int(0x0300f00f)))
        y.assign(y.bitAnd(int(0x0300f00f)))
        z.assign(z.bitAnd(int(0x0300f00f)))
        x.assign(x.bitOr(x.shiftLeft(int(4))))
        y.assign(y.bitOr(y.shiftLeft(int(4))))
        z.assign(z.bitOr(z.shiftLeft(int(4))))
        x.assign(x.bitAnd(int(0x030c30c3)))
        y.assign(y.bitAnd(int(0x030c30c3)))
        z.assign(z.bitAnd(int(0x030c30c3)))
        x.assign(x.bitOr(x.shiftLeft(int(2))))
        y.assign(y.bitOr(y.shiftLeft(int(2))))
        z.assign(z.bitOr(z.shiftLeft(int(2))))
        x.assign(x.bitAnd(int(0x09249249)))
        y.assign(y.bitAnd(int(0x09249249)))
        z.assign(z.bitAnd(int(0x09249249)))
        return x.bitOr(y.shiftLeft(int(1))).bitOr(z.shiftLeft(int(2)))
})

export const m2uv = Fn(([morton]: [Node<'int'>]) => {
        const px = morton.toVar()
        const py = morton.shiftRight(int(1)).toVar()
        px.assign(px.bitAnd(int(0x55555555)))
        py.assign(py.bitAnd(int(0x55555555)))
        px.assign(px.bitOr(px.shiftRight(int(1))))
        py.assign(py.bitOr(py.shiftRight(int(1))))
        px.assign(px.bitAnd(int(0x33333333)))
        py.assign(py.bitAnd(int(0x33333333)))
        px.assign(px.bitOr(px.shiftRight(int(2))))
        py.assign(py.bitOr(py.shiftRight(int(2))))
        px.assign(px.bitAnd(int(0x0f0f0f0f)))
        py.assign(py.bitAnd(int(0x0f0f0f0f)))
        px.assign(px.bitOr(px.shiftRight(int(4))))
        py.assign(py.bitOr(py.shiftRight(int(4))))
        px.assign(px.bitAnd(int(0x00ff00ff)))
        py.assign(py.bitAnd(int(0x00ff00ff)))
        px.assign(px.bitOr(px.shiftRight(int(8))))
        py.assign(py.bitOr(py.shiftRight(int(8))))
        px.assign(px.bitAnd(int(0x0000ffff)))
        py.assign(py.bitAnd(int(0x0000ffff)))
        return ivec2(px, py)
})

export const atlas = Fn(([p]: [Node<'ivec3'>]) => {
        return m2uv(xyz2m(p))
})