# Pure Mojo Float64 CPU D3Q19 BGK+Guo solver. Gmsh mask is external input.
# x: streamwise; y: spanwise; z: up. x/y periodic, halfway floor/ceiling.
from std.math import sin, exp, sqrt, isfinite
from std.sys import argv
from std.io import FileHandle
comptime PI = 3.1415926535897932384626433832795

struct Flow:
    var nx: Int
    var ny: Int
    var nz: Int
    var nu: Float64
    var tau: Float64
    var g: Float64
    var u_ref: Float64
    var solid: List[Int]
    var cx: List[Int]
    var cy: List[Int]
    var cz: List[Int]
    var opposite: List[Int]
    var weight: List[Float64]
    var f: List[Float64]
    var next_f: List[Float64]
    var u: List[Float64]
    var v: List[Float64]
    var w: List[Float64]
    var rho: List[Float64]
    var mass: Float64
    var px: Float64
    var py: Float64
    var pz: Float64
    var force_x: Float64
    var bx: Float64
    var by: Float64
    var bz: Float64
    var car_bx: Float64
    var car_by: Float64
    var car_bz: Float64
    var density_min: Float64
    var density_max: Float64
    var speed_max: Float64
    var divergence_rms: Float64
    var divergence_cells: Int
    var fluid_cells: Int

    def __init__(out self, nx: Int, ny: Int, nz: Int, nu: Float64, g: Float64, u_ref: Float64):
        self.nx = nx
        self.ny = ny
        self.nz = nz
        self.nu = nu
        self.tau = 0.5 + 3.0*nu
        self.g = g
        self.u_ref = u_ref
        var n = nx*ny*nz
        self.solid = List[Int](length=n, fill=0)
        self.cx = [0,1,-1,0,0,0,0,1,-1,1,-1,1,-1,1,-1,0,0,0,0]
        self.cy = [0,0,0,1,-1,0,0,1,-1,-1,1,0,0,0,0,1,-1,1,-1]
        self.cz = [0,0,0,0,0,1,-1,0,0,0,0,1,-1,-1,1,1,-1,-1,1]
        self.opposite = [0,2,1,4,3,6,5,8,7,10,9,12,11,14,13,16,15,18,17]
        self.weight = [1.0/3.0,1.0/18.0,1.0/18.0,1.0/18.0,1.0/18.0,1.0/18.0,1.0/18.0,1.0/36.0,1.0/36.0,1.0/36.0,1.0/36.0,1.0/36.0,1.0/36.0,1.0/36.0,1.0/36.0,1.0/36.0,1.0/36.0,1.0/36.0,1.0/36.0]
        self.f = List[Float64](length=n*19, fill=0.0)
        self.next_f = List[Float64](length=n*19, fill=0.0)
        self.u = List[Float64](length=n, fill=0.0)
        self.v = List[Float64](length=n, fill=0.0)
        self.w = List[Float64](length=n, fill=0.0)
        self.rho = List[Float64](length=n, fill=0.0)
        self.mass = 0.0
        self.px = 0.0
        self.py = 0.0
        self.pz = 0.0
        self.force_x = 0.0
        self.bx = 0.0
        self.by = 0.0
        self.bz = 0.0
        self.car_bx = 0.0
        self.car_by = 0.0
        self.car_bz = 0.0
        self.density_min = 1.0e30
        self.density_max = -1.0e30
        self.speed_max = 0.0
        self.divergence_rms = 0.0
        self.divergence_cells = 0
        self.fluid_cells = 0

    def channel_mask(mut self):
        for y in range(self.ny):
            for x in range(self.nx):
                self.solid[x+self.nx*y] = 1
                self.solid[x+self.nx*(y+self.ny*(self.nz-1))] = 1

    def load_mask(mut self, path: String) raises:
        var content: String
        with open(path, 'r') as inp:
            content = inp.read()
        var lines = content.split('\n')
        var dims = lines[0].split(' ')
        if Int(dims[0]) != self.nx or Int(dims[1]) != self.ny or Int(dims[2]) != self.nz:
            raise Error('Mask dimensions mismatch')
        if len(lines) < self.nx*self.ny*self.nz+1:
            raise Error('Truncated geometry mask')
        for cell in range(self.nx*self.ny*self.nz):
            var value = Int(lines[cell+1])
            if value < 0 or value > 2:
                raise Error('Invalid mask value')
            self.solid[cell] = 2 if value != 0 else 0
        self.channel_mask()

    def initialize(mut self, mode: Int) raises:
        # 0=physical rest, 1=3D diagonal transverse shear wave, 2=Poiseuille.
        var k = 2.0*PI/Float64(self.nx)
        var h = Float64(self.nz-2)
        for z in range(self.nz):
            for y in range(self.ny):
                for x in range(self.nx):
                    var cell = x+self.nx*(y+self.ny*z)
                    if self.solid[cell] != 0:
                        continue
                    self.fluid_cells += 1
                    var ux: Float64 = 0.0
                    var uy: Float64 = 0.0
                    var uz: Float64 = 0.0
                    if mode == 1:
                        # a=(1,1,-2)/sqrt(6), k=(1,1,1)*2pi/N; a.k=0.
                        # Exact incompressible solution: a*U*sin(k.x)*exp(-nu|k|^2*t).
                        var amplitude = self.u_ref*sin(k*(Float64(x+y+z)+1.5))/sqrt(6.0)
                        ux = amplitude
                        uy = amplitude
                        uz = -2.0*amplitude
                    elif mode == 2:
                        var height = Float64(z)-0.5
                        ux = 4.0*self.u_ref*height*(h-height)/(h*h)
                    ux -= 0.5*self.g
                    var speed2 = ux*ux+uy*uy+uz*uz
                    for q in range(19):
                        var cu = Float64(self.cx[q])*ux+Float64(self.cy[q])*uy+Float64(self.cz[q])*uz
                        self.f[cell*19+q] = self.weight[q]*(1.0+3.0*cu+4.5*cu*cu-1.5*speed2)
        self.sample()

    def observe(mut self, density: Float64, ux: Float64, uy: Float64, uz: Float64) raises:
        if not isfinite(density) or density <= 0.0 or not isfinite(ux) or not isfinite(uy) or not isfinite(uz):
            raise Error('Non-finite or non-positive fluid state')
        self.density_min = min(self.density_min, density)
        self.density_max = max(self.density_max, density)
        self.speed_max = max(self.speed_max, sqrt(ux*ux+uy*uy+uz*uz))

    def advance(mut self) raises:
        var relaxation = 1.0/self.tau
        var force_prefactor = 1.0-0.5*relaxation
        for z in range(self.nz):
            for y in range(self.ny):
                for x in range(self.nx):
                    var cell = x+self.nx*(y+self.ny*z)
                    if self.solid[cell] != 0:
                        continue
                    var base = cell*19
                    var density: Float64 = 0.0
                    var mx: Float64 = 0.0
                    var my: Float64 = 0.0
                    var mz: Float64 = 0.0
                    for q in range(19):
                        var value = self.f[base+q]
                        density += value
                        mx += value*Float64(self.cx[q])
                        my += value*Float64(self.cy[q])
                        mz += value*Float64(self.cz[q])
                    var ux = mx/density+0.5*self.g
                    var uy = my/density
                    var uz = mz/density
                    self.observe(density, ux, uy, uz)
                    var speed2 = ux*ux+uy*uy+uz*uz
                    var fx = density*self.g
                    self.force_x += fx
                    for q in range(19):
                        var ex = Float64(self.cx[q])
                        var ey = Float64(self.cy[q])
                        var ez = Float64(self.cz[q])
                        var cu = ex*ux+ey*uy+ez*uz
                        var eq = self.weight[q]*density*(1.0+3.0*cu+4.5*cu*cu-1.5*speed2)
                        var forcing = self.weight[q]*force_prefactor*(3.0*(ex-ux)+9.0*cu*ex)*fx
                        var post = self.f[base+q]+relaxation*(eq-self.f[base+q])+forcing
                        var dx = (x+self.cx[q]+self.nx)%self.nx
                        var dy = (y+self.cy[q]+self.ny)%self.ny
                        var dz = (z+self.cz[q]+self.nz)%self.nz
                        var destination = dx+self.nx*(dy+self.ny*dz)
                        if self.solid[destination] != 0:
                            self.next_f[base+self.opposite[q]] = post
                            self.bx -= 2.0*ex*post
                            self.by -= 2.0*ey*post
                            self.bz -= 2.0*ez*post
                            if self.solid[destination] == 2:
                                self.car_bx -= 2.0*ex*post
                                self.car_by -= 2.0*ey*post
                                self.car_bz -= 2.0*ez*post
                        else:
                            self.next_f[destination*19+q] = post
        var swap = self.f^
        self.f = self.next_f^
        self.next_f = swap^

    def sample(mut self) raises:
        self.mass = 0.0
        self.px = 0.0
        self.py = 0.0
        self.pz = 0.0
        for cell in range(self.nx*self.ny*self.nz):
            if self.solid[cell] != 0:
                continue
            var density: Float64 = 0.0
            var mx: Float64 = 0.0
            var my: Float64 = 0.0
            var mz: Float64 = 0.0
            for q in range(19):
                var value = self.f[cell*19+q]
                density += value
                mx += value*Float64(self.cx[q])
                my += value*Float64(self.cy[q])
                mz += value*Float64(self.cz[q])
            self.rho[cell] = density
            self.u[cell] = mx/density+0.5*self.g
            self.v[cell] = my/density
            self.w[cell] = mz/density
            self.observe(density, self.u[cell], self.v[cell], self.w[cell])
            self.mass += density
            self.px += density*self.u[cell]
            self.py += density*self.v[cell]
            self.pz += density*self.w[cell]
        var div2: Float64 = 0.0
        self.divergence_cells = 0
        for z in range(self.nz):
            for y in range(self.ny):
                for x in range(self.nx):
                    var cell = x+self.nx*(y+self.ny*z)
                    var left = (x+self.nx-1)%self.nx+self.nx*(y+self.ny*z)
                    var right = (x+1)%self.nx+self.nx*(y+self.ny*z)
                    var back = x+self.nx*((y+self.ny-1)%self.ny+self.ny*z)
                    var front = x+self.nx*((y+1)%self.ny+self.ny*z)
                    var down = x+self.nx*(y+self.ny*((z+self.nz-1)%self.nz))
                    var up = x+self.nx*(y+self.ny*((z+1)%self.nz))
                    if self.solid[cell]+self.solid[left]+self.solid[right]+self.solid[back]+self.solid[front]+self.solid[down]+self.solid[up] == 0:
                        var divergence = 0.5*(self.u[right]-self.u[left]+self.v[front]-self.v[back]+self.w[up]-self.w[down])
                        div2 += divergence*divergence
                        self.divergence_cells += 1
        if self.divergence_cells > 0:
            self.divergence_rms = sqrt(div2/Float64(self.divergence_cells))

    def write_array(self, mut out: FileHandle, values: List[Float64]) raises:
        out.write('[')
        for i in range(self.nx*self.ny*self.nz):
            if i > 0:
                out.write(',')
            out.write(values[i])
        out.write(']')

    def write_fields(self, mut out: FileHandle, step: Int) raises:
        out.write('{"step":', step, ',"u":')
        self.write_array(out, self.u)
        out.write(',"v":')
        self.write_array(out, self.v)
        out.write(',"w":')
        self.write_array(out, self.w)
        out.write(',"rho":')
        self.write_array(out, self.rho)
        out.write('}')

    def write_audit(self, mut out: FileHandle, step: Int) raises:
        out.write('{"step":', step, ',"mass":', self.mass, ',"px":', self.px, ',"py":', self.py, ',"pz":', self.pz, ',"forceX":', self.force_x, ',"boundaryX":', self.bx, ',"boundaryY":', self.by, ',"boundaryZ":', self.bz, ',"carBoundaryX":', self.car_bx, ',"carBoundaryY":', self.car_by, ',"carBoundaryZ":', self.car_bz, '}')

    def write_header(self, mut out: FileHandle, name: String, steps: Int) raises:
        out.write('{"case":"', name, '","nx":', self.nx, ',"ny":', self.ny, ',"nz":', self.nz, ',"nu":', self.nu, ',"tau":', self.tau, ',"acceleration":', self.g, ',"u_ref":', self.u_ref, ',"steps":', steps, ',"solid":[')
        for i in range(self.nx*self.ny*self.nz):
            if i > 0:
                out.write(',')
            out.write(self.solid[i])
        out.write(']')

    def write_diagnostics(self, mut out: FileHandle) raises:
        out.write(',"diagnostics":{"minDensity":', self.density_min, ',"maxDensity":', self.density_max, ',"maxSpeed":', self.speed_max, ',"maxMach":', sqrt(3.0)*self.speed_max, ',"divergenceRms":', self.divergence_rms, ',"divergenceCells":', self.divergence_cells, ',"fluidCells":', self.fluid_cells, ',"finite":true,"divergenceDefinition":"centered lattice derivatives; RMS only where all six cardinal neighbors are fluid; dx=1"}')


def run_case(mode: String, n: Int, steps: Int, u0: Float64, mask: String, path: String) raises:
    var nx = n
    var ny = n
    var nz = n
    var g: Float64 = 0.0
    var init = 1
    if mode == 'channel':
        nx = 8
        ny = 8
        g = 8.0*0.08*u0/Float64((nz-2)*(nz-2))
        init = 2
    elif mode == 'car':
        nx = 96
        ny = 48
        nz = 32
        g = 8.0*0.08*u0/Float64((nz-2)*(nz-2))
        init = 0
    var flow = Flow(nx, ny, nz, 0.08, g, u0)
    if mode == 'channel':
        flow.channel_mask()
    elif mode == 'car':
        flow.load_mask(mask)
    flow.initialize(init)
    with open(path, 'w') as out:
        flow.write_header(out, mode, steps)
        out.write(',"initial":')
        flow.write_fields(out, 0)
        out.write(',"auditInitial":')
        flow.write_audit(out, 0)
        out.write(',"frames":[')
        flow.write_fields(out, 0)
        for step in range(1, steps+1):
            flow.advance()
            if mode == 'car' and step%100 == 0:
                flow.sample()
                out.write(',')
                flow.write_fields(out, step)
                print('car step', step, 'maxMach', sqrt(3.0)*flow.speed_max)
        flow.sample()
        if mode != 'car' or steps%100 != 0:
            out.write(',')
            flow.write_fields(out, steps)
        out.write('],"auditFinal":')
        flow.write_audit(out, steps)
        flow.write_diagnostics(out)
        out.write('}\n')
    print('Completed', mode, nx, ny, nz, steps)


def main() raises:
    var args = argv()
    if len(args) == 7:
        run_case(args[1], Int(args[2]), Int(args[3]), Float64(args[4]), args[5], args[6])
    else:
        raise Error('Usage: solver shear|channel|car N steps Uref mask.txt out.json')