vtk export

###########################FUNCTIONS##################################

##### exportVTK
function exportVTK(DataMat, NameMat ,filename, dh, Nxyz)

nx = dh[1]*(Nxyz[1]-1);
ny = dh[2]*(Nxyz[2]-1);
nz = dh[3]*(Nxyz[3]-1);

x, y, z = ndgrid(0:dh[1]:nx, 0:dh[2]:ny, 0:dh[3]:nz );

nr_of_elements = length(x);
fid = open(filename, "w"); 

#ASCII file header
println(fid, "# vtk DataFile Version 3.0");
println(fid, "VTK from Julia");
println(fid, "BINARY\n");
println(fid, "DATASET STRUCTURED_GRID");
println(fid, "DIMENSIONS $(size(x,1)) $(size(x,2)) $(size(x,3))");
println(fid, "POINTS $(nr_of_elements) double"); 
for  i = 1:nr_of_elements
    write(fid, bswap( x[i] ) )
    write(fid, bswap( y[i] ) )
    write(fid, bswap( z[i] ) )
end 
 
#append binary data
print(fid, "\nPOINT_DATA $(nr_of_elements)\n");

for i = 1:length(DataMat)

    if size(DataMat[i], 2) == 3
    #append a vector
        println(fid, "VECTORS $(NameMat[i]) double");
        for j = 1:nr_of_elements
            write(fid, bswap(DataMat[i][j, 1]))
            write(fid, bswap(DataMat[i][j, 2]))
            write(fid, bswap(DataMat[i][j, 3]))
        end
    elseif size(DataMat[i], 2) == 1
    #append a scalar
        println(fid, "SCALARS $(NameMat[i]) double");
        println(fid, "LOOKUP_TABLE default");
        for j = 1:nr_of_elements
            write(fid, bswap(DataMat[i][j]))
        end 
    else
    #data is no vector or scalar
        error("DataMat $i is not a vector or scalar, or not shaped the way it should be")
    end

end

close(fid);

end

##### ndgrid (based on https://github.com/JuliaLang/julia/blob/master/examples/ndgrid.jl)
function ndgrid{T}(vx::AbstractVector{T}, vy::AbstractVector{T}, vz::AbstractVector{T})
    m, n, o = length(vx), length(vy), length(vz)
    vx = reshape(vx, m, 1, 1)
    vy = reshape(vy, 1, n, 1)
    vz = reshape(vz, 1, 1, o)
    om = ones(Int, m)
    on = ones(Int, n)
    oo = ones(Int, o)
    return (vx[:, on, oo], vy[om, :, oo], vz[om, on, :])
end

##### Couette Flow
function Couette(L, Hy, Hz, r, R, dh, u = 1, addnoise = 0, rho = 1000)

w = u/r

eta = r/R;
A = -w*eta^2/(1-eta^2);
B = w*r^2/(1-eta^2);

Nxyz = int( [ceil(L/dh[1])+1, ceil(Hy/dh[2])+1, ceil(Hz/dh[3])+1, 2] );

V_x = zeros(Nxyz[1], Nxyz[2], Nxyz[3], Nxyz[4]);
V_y = zeros(Nxyz[1], Nxyz[2], Nxyz[3], Nxyz[4]);
V_z = zeros(Nxyz[1], Nxyz[2], Nxyz[3], Nxyz[4]);
P_real = zeros(Nxyz[1], Nxyz[2], Nxyz[3], Nxyz[4]);

C = -rho*(A^2*r^2/2 + 2*A*B*log(r) - B^2/(2*r^2));

for i = 1:Nxyz[2]
    for j= 1:Nxyz[3]
        dy = (i-1)*dh[2]-Hy/2;
        dz = (j-1)*dh[3]-Hz/2;
        ang = angle(complex(dy, dz));
        r2 = sqrt(dy^2 + dz^2);
        if r2>r && r2<R
            V_y[:, i, j, :] = -(A*r2 + B/r2)*sin(ang);
            V_z[:, i, j, :] = (A*r2 + B/r2)*cos(ang);
            P_real[:, i, j, :] = rho*(A^2*r2^2/2 + 2*A*B*log(r2) - B^2/(2*r2^2)) + C;
        end
    end
end

temp = (V_x.^2 + V_y.^2 + V_z.^2)
Mask = int(temp .!= 0)

if addnoise == 1
    V_x = V_x + 0.3*randn(Nxyz[1], Nxyz[2], Nxyz[3], Nxyz[4]).*Mask;
    V_y = V_y + 0.3*randn(Nxyz[1], Nxyz[2], Nxyz[3], Nxyz[4]).*Mask;
    V_z = V_z + 0.3*randn(Nxyz[1], Nxyz[2], Nxyz[3], Nxyz[4]).*Mask;
end

return V_x, V_y, V_z, P_real

end

#####################Script##############################

# Parameters
VTKfilename = "Couette.vtk"
L = 0.4
Hy = 2
Hz = 2
r = 0.2
R = 0.8
dh = [0.1, 0.05, 0.05] #spacing between points in each dimension
u_max = 2
addnoise = 0
rho = 1000

#
ParaviewData = Array(Any, 0)
ParaviewNames = Array(Any, 0)

# returns a 4d velocity and pressure field (4th dimensions = time) of a couette flow
vx, vy, vz, p = Couette(L, Hy, Hz, r, R, dh, u_max, addnoise, rho )

push!(ParaviewData, [vx[:, :, :, 1][:] vy[:, :, :, 1][:] vz[:, :, :, 1][:]])
push!(ParaviewData, p[:, :, :, 1][:])
push!(ParaviewNames, "Velocity")
push!(ParaviewNames, "Pressure")

Nxyz = size(vx)[1:3] #number of points in each direction
exportVTK(ParaviewData, ParaviewNames ,VTKfilename, dh, Nxyz)