fluid.lua

local fluid = {}

-- GLOBALS

N = 128   -- resolution
ITER = 8 -- Gauss–Seidel Iteration Method
DIFF = 0.0001
VISC = 0.0001


-- NOTE FROM CREATOR : x and y will be represented as i and j to sync with the formulas

local function bound(b, s, n)
    N = n or N
    --[[

        Checks if it hits the bounds, if so, invert the velocity

        (1,1)      ...    (N,1)
        .                  .
        .                  .
        (1,N)      ...    (N,N)
    ]]
    
    for i = 2, (N - 1) do
        for j = 2, (N - 1) do
            if fluid.boundary[i][j] == 1 then
                s[i][j] = 0 
            end
        end
    end


    for i = 2, (N - 1) do

        if fluid.boundary[i][1] == 0 then
            s[i][1] = (b == 2) and -s[i][2] or s[i][2]       -- Bottom Bound
        else
            s[i][1] = 0
        end

        if fluid.boundary[i][N] == 0 then
            s[i][N] = (b == 2) and -s[i][N - 1] or s[i][N - 1] -- Top Bound
        else
            s[i][N] = 0
        end

        if fluid.boundary[1][i] == 0 then
            s[1][i] = (b == 1) and -s[2][i] or s[2][i]       -- Left Bound
        else
            s[1][i] = 0
        end

        if fluid.boundary[N][i] == 0 then
            s[N][i] = (b == 1) and -s[N - 1][i] or s[N - 1][i] -- Right Bound
        else
            s[N][i] = 0
        end
         
    end

    -- -- Corner Cells  = avg of neighbours
    s[1][1] = 0.5 * (s[2][1] + s[1][2])
    s[1][N] = 0.5 * (s[2][N] + s[1][N - 1])
    s[N][1] = 0.5 * (s[N - 1][1] + s[N][2])
    s[N][N] = 0.5 * (s[N - 1][N] + s[N][N - 1])
end


-- Global force
function Force(Vx, Vy, Fx, Fy, x, y, radius, dt)
    --[[

        Vx, Vy = Velocity Field Array [ Gridwise ]
        Fx, Fy = External Force Factor Parameter
        x, y = Grid Position ( force center )
        radius = Range of influence
        dt = deltatime

    ]]
    

    for i = (x - radius), (x + radius) do
        for j = (y - radius), (y + radius) do
            if i > 1 and i < N and j > 1 and j < N then
                local dx = i - x
                local dy = j - y

                local distanceSq = dx * dx + dy * dy
                local radiusSq = radius * radius
                if distanceSq < radiusSq then
                    local falloff = exp(-distanceSq / radiusSq)  -- 1 - ( distance / radius ) for linear interpolation
                    Vx[i][j] = Vx[i][j] + Fx * falloff * dt
                    Vy[i][j] = Vy[i][j] + Fy * falloff * dt

                end
            end
        end
    end
    return Vx, Vy
end

function Diffuse(b, s, Bs, diff, dt, iter, N)
    --[[
        Laplacian ( using Laplacian Operator : 2X)
        ∇2X≈Xi−1,j​+Xi+1,j​+Xi,j−1​+Xi,j+1​−4Xi,j​

        b = boundary type (
            0 - scalar,
            1 - Vx ( Velocity x-axis )
            2 - Vy ( Velocity y-axis )
        )

        s = latest diffused state
        Bs = previous diffused state ( source state )
        iter = number of Gauss-Seidel iterations
    ]]


    local a = dt * diff * (N - 2) * (N - 2)

    for k = 1, iter do
        for j = 2, N - 1 do
            for i = 2, N - 1 do
                s[i][j] = (Bs[i][j] + a * (s[i - 1][j] + s[i + 1][j] + s[i][j - 1] + s[i][j + 1])) / (1 + 4 * a)
            end
        end
        bound(b, s, N)
    end
end

function Advect(b, a, Ba, Vx, Vy, dt, N)
    --[[
        / Semi-Lagrangian Advection (Stam’s Method)

        Transporting quantities across velocity field

        Ba = pre advection
        a = post advection

        bilinear interpolation

        note : this is backward tracing instead of forward integration,
        hence, the grid cell content depends on where it came from rather than
        where it might or will go


        ALGORTIHM GENERATED : 

        function advect(b, d, d0, Vx, Vy, dt, N):
            dt0 = dt * (N-2)
            for j = 2 to N-1:
                for i = 2 to N-1:
                    x = i - dt0 * Vx[i,j]
                    y = j - dt0 * Vy[i,j]
                    x = clamp(x, 1.5, N-0.5)
                    y = clamp(y, 1.5, N-0.5)

                    i0 = floor(x)
                    i1 = i0 + 1
                    j0 = floor(y)
                    j1 = j0 + 1

                    s1 = x - i0
                    s0 = 1 - s1
                    t1 = y - j0
                    t0 = 1 - t1

                    d[i,j] = s0*(t0*d0[i0,j0] + t1*d0[i0,j1]) +
                            s1*(t0*d0[i1,j0] + t1*d0[i1,j1])
            set_bnd(b, d, N)
    ]]

    Bdt = dt * (N - 2)

    for j = 2, N - 1 do
        for i = 2, N - 1 do
            local x = clamp((i - Bdt * Vx[i][j]), 1.5, N - 0.5)
            local y = clamp((j - Bdt * Vy[i][j]), 1.5, N - 0.5)
            local i0 = floor(x)
            local i1 = i0 + 1
            local j0 = floor(y)
            local j1 = j0 + 1

            local s1 = x - i0
            local s0 = 1 - s1

            local t1 = y - j0
            local t0 = 1 - t1

            a[i][j] = s0 * ( t0 * Ba[i0][j0] + t1 * Ba[i0][j1] ) +
                      s1 * ( t0 * Ba[i1][j0] + t1 * Ba[i1][j1] )
            
        end
    end
    bound(b, a, N)

end

function Project(Vx, Vy, p, div, N)

    --[[


        Vx, Vy = Velocity field components
        p = pressure field
        div = divergence map

        Information Stated below are sourced externally [ Algorithms are compiled using Artifical Intelligence ]
    
        Making fluid incompressible

        First for loop ( j x i ) :
            Computes divergence
            Measure how much fluid enters/leaves.
                Enters > Leaves = +ve Divergence
                Leaves > Enters = -ve Divergence

        
        Second for loop :
            Solve for pressure

            Gauss-Seidel relaxtion -> iteratively Poisson Equation
            ∇2p = div
            
            Pressure spreads out evenly - >
                High Pressure areas push back on lower ones

        
        Third for loop :
            Actual projection
            Push the velocity field in the opposite direction of pressure gradient
            balancing out any divergence

            ∇⋅V = 0 [ V is vector, but I am unable to represent that here]

    ]]

    for j = 2, N - 1 do
        for i = 2, N - 1 do
            div[i][j] = -0.5 * (Vx[i + 1][j] - Vx[i - 1][j] + Vy[i][j + 1] - Vy[i][j - 1]) / N
            p[i][j] = 0
        end
    end
    bound(0, div)
    bound(0, p, N)

    for k = 1, ITER do
        for j = 2, N - 1 do
            for i = 2, N - 1 do
                p[i][j] = (div[i][j] + p[i - 1][j] + p[i + 1][j] + p[i][j - 1] + p[i][j + 1]) / 4
            end     
        end
        bound(0 , p, N)
    end

    for j = 2, N - 1 do
        for i = 2, N - 1 do
            Vx[i][j] = Vx[i][j] - 0.5 * ( p[i + 1][j] - p[i - 1][j]) * N
            Vy[i][j] = Vy[i][j] - 0.5 * ( p[i][j + 1] - p[i][j - 1]) * N
        end
    end
    bound(1, Vx, N)
    bound(1, Vy, N)

end


function fluid:init()

    --[[
    
        fluid = {
        
            density -> Scalar field representing the amount of substance
                       each cell 
            Vx -> horizontal (x-axis) velocity
            Vy -> vertical (y-axis) velocity

            Vx0, Vy0 -> Temporary velocity field

            s -> Scalar Field of current density
            s0 -> Scalar Field of previous field
        }

    ]]
    self.density = {}
    self.Vx, self.Vy = {}, {}
    self.Vx0, self.Vy0 = {}, {}
    self.s = {}
    self.s0 = {}
    self.p = {}
    self.div = {}
    self.boundary = {}
    
    

    -- 2 dimensional array setup, i for colums and j for rows, using table in table method
    
    --[[
        { 
            { 0, 0, 0}
        }
        {
            { 0, 0, 0}
        }

        For boundary, 0 = fluid, 1 = solid

    ]]
    for i = 1, N do
        self.density[i] = {}
        self.Vx[i], self.Vy[i] = {}, {}
        self.Vx0[i], self.Vy0[i] = {}, {}     
        self.s[i], self.s0[i] = {}, {}
        self.p[i], self.div[i] = {}, {}
        self.boundary[i] = {}
        for j = 1, N do
            self.density[i][j] = 0
            self.Vx[i][j], self.Vy[i][j] = 0, 0
            self.Vx0[i][j], self.Vy0[i][j] = 0, 0
            self.s[i][j], self.s0[i][j] = 0, 0
            self.p[i][j], self.div[i][j] = 0, 0
            self.boundary[i][j] = 0
        end

    end

end

function fluid:simulate(dt)

    --[[
        add viscosity :- 
        spread out velocity vals to neighbouring cells
    ]]
    
    Diffuse(1, self.Vx0, self.Vx, VISC, dt, ITER, N)
    Diffuse(2, self.Vy0, self.Vy, VISC, dt, ITER, N)
    Project(self.Vx0, self.Vy0, self.p, self.div, N)

    --[[
        advect velocity :-
        moves the velocity field itself according to own flow ( self-advect )
        project to keep it incompressible post movement
    
    ]]
    
    Advect(1, self.Vx, self.Vx0, self.Vx0, self.Vy0, dt, N)
    Advect(2, self.Vy, self.Vy0, self.Vx0, self.Vy0, dt, N)
    Project(self.Vx, self.Vy, self.p, self.div, N)

    --[[
        Smoke v, passively moved by velocity field
    ]]
    Diffuse(0, self.s0, self.s, DIFF, dt, ITER, N)
    Advect(0, self.s, self.s0, self.Vx, self.Vy, dt, N)
end


-- Helper functions forward

function fluid:addDensity(x, y, amt)
    if x >= 1 and x <= N and y >= 1 and y <= N then
        self.s[x][y] = self.s[x][y] + amt
    end
end

function fluid:addVelocity(x, y, amtX, amtY)
    if x >= 1 and x <= N and y >= 1 and y <= N then
        self.Vx[x][y] = self.Vx[x][y] + amtX
        self.Vy[x][y] = self.Vy[x][y] + amtY
    end
end

function fluid:create(x, y, dx, dy, density, Vx, Vy)
    local normX = floor(x / size) + 1
    local normY = floor(y / size) + 1

    normX = clamp(normX, 2, N - 1)
    normY = clamp(normY, 2, N - 1)
    self:addDensity(normX, normY, density)
    self:addVelocity(normX, normY, Vx * dx , Vy * dy)
end

function fluid:clear(x, y)
    local normX = floor(x / size) + 1
    local normY = floor(y / size) + 1

    normX = clamp(normX, 2, N - 1)
    normY = clamp(normY, 2, N - 1)
    self:addDensity(normX, normY, -100)
end

function fluid:solidCircle(cx, cy, radius)

    --[[
    
        cx, cy - center of hte circle in grid coords
        radius - radius of circle

        Marking the circle points as 1 on the boundary to denote solid
        To make the fluid not interact with the solid, set the velocity
        to zero
    ]]
    
    for i = 1, N do
        for j = 1, N do
            local dx = i - cx
            local dy = j - cy
            local distSq = dx * dx + dy * dy

            if distSq <= radius * radius then
                self.boundary[i][j] = 1 
                self.Vx[i][j] = 0
                self.Vy[i][j] = 0
            end
        end
    end


end

function fluid:clearSolids()
        for i = 1, N do
            for j = 1, N do
                self.boundary[i][j] = 0 
            end
        end
end

function fluid:getNormalizedPressure(i , j)
    local pressure = self.s[i][j] or 0
    local minPressure = 0.1
    local maxPressure = -0.1

    local normalized = (pressure - minPressure) / (maxPressure - minPressure)
    
    return clamp(normalized, 0, 1)

end

function fluid:getPressureColor(i , j)
    

    local r, g, b
    local d = math.min(self.s[i][j], 1)

    local normalized = self:getNormalizedPressure(i , j)

    if normalized < 0.25 then
        -- Red to Yellow
        local t = normalized * 4
        r = 1.0
        g = t
        b = 0
    elseif normalized < 0.5 then
        -- Yellow to Green
        local t = (normalized - 0.25) * 4
        r = 1.0 - t
        g = 1.0
        b = 0
    elseif normalized < 0.75 then
        -- Green to Cyan
        local t = (normalized - 0.5) * 4
        r = 0
        g = 1.0
        b = t
    else
        -- Cyan to Blue
        local t = (normalized - 0.75) * 4
        r = 0
        g = 1.0 - t
        b = 1.0
    end
    return r, g, b, d
end
function fluid:createLaminarFlow(direction, strength, width, position)
    --[[
        Creates laminar flow (smooth, parallel flow)
        
        Parameters:
        direction - "left", "right", "up", "down"
        strength - velocity magnitude
        width - thickness of the flow stream
        position - position along the perpendicular axis
    ]]
    
    local center = math.floor(position or (N / 2))
    local halfWidth = math.floor((width or 5) / 2)
    
    for j = 2, N - 1 do
        for i = 2, N - 1 do
            -- Check if current cell is within the flow stream
            local inFlow = false
            
            if direction == "left" or direction == "right" then
                -- Horizontal flow: check vertical position
                inFlow = (j >= center - halfWidth and j <= center + halfWidth)
            else
                -- Vertical flow: check horizontal position
                inFlow = (i >= center - halfWidth and i <= center + halfWidth)
            end
            
            if inFlow and fluid.boundary[i][j] == 0 then
                if direction == "left" then
                    self.Vx[i][j] = self.Vx[i][j] - strength
                elseif direction == "right" then
                    self.Vx[i][j] = self.Vx[i][j] + strength
                elseif direction == "up" then
                    self.Vy[i][j] = self.Vy[i][j] - strength
                elseif direction == "down" then
                    self.Vy[i][j] = self.Vy[i][j] + strength
                end
                
                -- Add some density to visualize the flow
                self.s[i][j] = math.min(self.s[i][j] + (strength * 0.5), 100)
            end
        end
    end
end

function fluid:createChannelFlow(horizontal, strength, densityAmount)
    --[[
        Creates flow between two parallel boundaries (like in a channel)
        
        Parameters:
        horizontal - true for horizontal channel, false for vertical
        strength - flow velocity
        densityAmount - density to add for visualization
    ]]
    
    if horizontal then
        -- Horizontal channel flow (left to right)
        for j = 2, N - 1 do
            for i = 2, N - 1 do
                if fluid.boundary[i][j] == 0 then
                    -- Parabolic velocity profile (faster in center, slower near walls)
                    local normalizedPos = (j - 2) / (N - 3)
                    local velocityProfile = 4 * normalizedPos * (1 - normalizedPos)
                    
                    self.Vx[i][j] = self.Vx[i][j] + strength * velocityProfile
                    self.s[i][j] = math.min(self.s[i][j] + densityAmount * velocityProfile, 100)
                end
            end
        end
    else
        -- Vertical channel flow (top to bottom)
        for j = 2, N - 1 do
            for i = 2, N - 1 do
                if fluid.boundary[i][j] == 0 then
                    -- Parabolic velocity profile
                    local normalizedPos = (i - 2) / (N - 3)
                    local velocityProfile = 4 * normalizedPos * (1 - normalizedPos)
                    
                    self.Vy[i][j] = self.Vy[i][j] + strength * velocityProfile
                    self.s[i][j] = math.min(self.s[i][j] + densityAmount * velocityProfile, 100)
                end
            end
        end
    end
end

function fluid:createShearFlow(strength, position)
    --[[
        Creates shear flow (fluid moving in opposite directions)
        
        Parameters:
        strength - velocity magnitude
        position - y-position of the shear line
    ]]
    
    local shearLine = position or math.floor(N / 2)
    
    for j = 2, N - 1 do
        for i = 2, N - 1 do
            if fluid.boundary[i][j] == 0 then
                if j < shearLine then
                    -- Flow to the right above shear line
                    self.Vx[i][j] = self.Vx[i][j] + strength
                else
                    -- Flow to the left below shear line
                    self.Vx[i][j] = self.Vx[i][j] - strength
                end
                
                self.s[i][j] = math.min(self.s[i][j] + strength * 0.3, 100)
            end
        end
    end
end
return fluid