export piso!

"""
    cpiso!(model, config; 
        limit_gradient=false, pref=nothing, ncorrectors=0, inner_loops=0)

Incompressible and transient variant of the SIMPLE algorithm to solving coupled momentum and mass conservation equations. 

# Input arguments

- `model` reference to a `Physics`` model defined by the user.
- `config` Configuration structure defined by the user with solvers, schemes, runtime and hardware structures configuration details.
- `limit_gradient` flag use to activate gradient limiters in the solver (default = `false`)
- `pref` Reference pressure value for cases that do not have a pressure defining BC. Incompressible solvers only (default = `nothing`)
- `ncorrectors` number of non-orthogonality correction loops (default = `0`)
- `inner_loops` number to inner loops used in transient solver based on PISO algorithm (default = `0`)

# Output

- `Ux` Vector of x-velocity residuals for each iteration.
- `Uy` Vector of y-velocity residuals for each iteration.
- `Uz` Vector of y-velocity residuals for each iteration.
- `p` Vector of pressure residuals for each iteration.
"""
function piso!(
    model, config; 
    limit_gradient=false, pref=nothing, ncorrectors=0, inner_loops=2)

    residuals = setup_incompressible_solvers(
        PISO, model, config; 
        limit_gradient=limit_gradient, 
        pref=pref,
        ncorrectors=ncorrectors, 
        inner_loops=inner_loops
        )
        
    return residuals
end

function PISO(
    model, turbulenceModel, ∇p, U_eqn, p_eqn, config; 
    limit_gradient=false, pref=nothing, ncorrectors=0, inner_loops=2
    )
    
    # Extract model variables and configuration
    (; U, p) = model.momentum
    (; nu) = model.fluid
    mesh = model.domain
    (; solvers, schemes, runtime, hardware) = config
    (; iterations, write_interval, dt) = runtime
    (; backend) = hardware
    
    mdotf = get_flux(U_eqn, 2)
    nueff = get_flux(U_eqn, 3)
    rDf = get_flux(p_eqn, 1)
    divHv = get_source(p_eqn, 1)

    @info "Initialise VTKWriter (Store mesh in host memory)"

    VTKMeshData = initialise_writer(model.domain)
    
    @info "Allocating working memory..."

    # Define aux fields 
    gradU = Grad{schemes.U.gradient}(U)
    gradUT = T(gradU)
    S = StrainRate(gradU, gradUT)
    S2 = ScalarField(mesh)

    n_cells = length(mesh.cells)
    Uf = FaceVectorField(mesh)
    pf = FaceScalarField(mesh)
    Hv = VectorField(mesh)
    rD = ScalarField(mesh)

    # Pre-allocate auxiliary variables
    TF = _get_float(mesh)
    TI = _get_int(mesh)
    prev = zeros(TF, n_cells)
    prev = _convert_array!(prev, backend) 

    # Pre-allocate vectors to hold residuals 
    R_ux = ones(TF, iterations)
    R_uy = ones(TF, iterations)
    R_uz = ones(TF, iterations)
    R_p = ones(TF, iterations)
    cellsCourant =adapt(backend, zeros(TF, length(mesh.cells)))
    
    # Initial calculations
    time = zero(TF) # assuming time=0
    interpolate!(Uf, U, config)   
    correct_boundaries!(Uf, U, U.BCs, time, config)
    flux!(mdotf, Uf, config)
    grad!(∇p, pf, p, p.BCs, time, config)
    limit_gradient && limit_gradient!(∇p, p, config)

    update_nueff!(nueff, nu, model.turbulence, config)

    xdir, ydir, zdir = XDir(), YDir(), ZDir()

    @info "Starting PISO loops..."

    progress = Progress(iterations; dt=1.0, showspeed=true)

    @time for iteration ∈ 1:iterations
        time = (iteration - 1)*dt

        solve_equation!(U_eqn, U, solvers.U, xdir, ydir, zdir, config; time=time)
          
        # Pressure correction
        inverse_diagonal!(rD, U_eqn, config)
        interpolate!(rDf, rD, config)
        remove_pressure_source!(U_eqn, ∇p, config)
        
        for i ∈ 1:inner_loops
            H!(Hv, U, U_eqn, config)
            
            # Interpolate faces
            interpolate!(Uf, Hv, config) # Careful: reusing Uf for interpolation
            correct_boundaries!(Uf, Hv, U.BCs, time, config)
            # div!(divHv, Uf, config)

            # new approach
            flux!(mdotf, Uf, config)
            div!(divHv, mdotf, config)
            
            # Pressure calculations (previous implementation)
            @. prev = p.values
            solve_equation!(p_eqn, p, solvers.p, config; ref=pref, time=time)
            if i == inner_loops
                explicit_relaxation!(p, prev, 1.0, config)
            else
                explicit_relaxation!(p, prev, solvers.p.relax, config)
            end

            grad!(∇p, pf, p, p.BCs, time, config) 
            limit_gradient && limit_gradient!(∇p, p, config)

            # nonorthogonal correction (experimental)
            for i ∈ 1:ncorrectors
                discretise!(p_eqn, p, config)       
                apply_boundary_conditions!(p_eqn, p.BCs, nothing, time, config)
                setReference!(p_eqn, pref, 1, config)
                nonorthogonal_face_correction(p_eqn, ∇p, rDf, config)
                update_preconditioner!(p_eqn.preconditioner, p.mesh, config)
                solve_system!(p_eqn, solvers.p, p, nothing, config)

                if i == ncorrectors
                    explicit_relaxation!(p, prev, 1.0, config)
                else
                    explicit_relaxation!(p, prev, solvers.p.relax, config)
                end
                grad!(∇p, pf, p, p.BCs, time, config) 
                limit_gradient && limit_gradient!(∇p, p, config)
            end

            # old approach - keep for now!
            # correct_velocity!(U, Hv, ∇p, rD, config)
            # interpolate!(Uf, U, config)
            # correct_boundaries!(Uf, U, U.BCs, time, config)
            # flux!(mdotf, Uf, config) # old approach

            # new approach
            interpolate!(Uf, U, config) # velocity from momentum equation
            correct_boundaries!(Uf, U, U.BCs, time, config)
            flux!(mdotf, Uf, config)
            correct_mass_flux(mdotf, p, rDf, config)
            correct_velocity!(U, Hv, ∇p, rD, config)

        end # corrector loop end
        
        # correct_mass_flux(mdotf, p, rDf, config) # new approach


    grad!(gradU, Uf, U, U.BCs, time, config)
    limit_gradient && limit_gradient!(gradU, U, config)

    turbulence!(turbulenceModel, model, S, S2, prev, time, config) 
    update_nueff!(nueff, nu, model.turbulence, config)

    residual!(R_ux, U_eqn, U.x, iteration, xdir, config)
    residual!(R_uy, U_eqn, U.y, iteration, ydir, config)
    if typeof(mesh) <: Mesh3
        residual!(R_uz, U_eqn, U.z, iteration, zdir, config)
    end
    residual!(R_p, p_eqn, p, iteration, nothing, config)
    maxCourant = max_courant_number!(cellsCourant, model, config)

    ProgressMeter.next!(
        progress, showvalues = [
            (:time, iteration*runtime.dt),
            (:Courant, maxCourant),
            (:Ux, R_ux[iteration]),
            (:Uy, R_uy[iteration]),
            (:Uz, R_uz[iteration]),
            (:p, R_p[iteration]),
            ]
        )

    if iteration%write_interval + signbit(write_interval) == 0
        model2vtk(model, VTKMeshData, @sprintf "timestep_%.6d" iteration)
    end

    end # end for loop

    return (Ux=R_ux, Uy=R_uy, Uz=R_uz, p=R_p)
end

### GRADIENT LIMITER - EXPERIMENTAL

function limit_gradient!(∇F, F, config)
    (; hardware) = config
    (; backend, workgroup) = hardware

    mesh = F.mesh
    (; cells, cell_neighbours, cell_faces, cell_nsign, faces) = mesh

    minPhi0 = maximum(F.values) # use min value so all values compared are larger
    maxPhi0 = minimum(F.values)
    (; x, y, z) = ∇F.result

    kernel! = _limit_gradient!(backend, workgroup)
    kernel!(x, y, z, F, cells, cell_neighbours, cell_faces, cell_nsign, faces, minPhi0, maxPhi0, ndrange=length(cells))
    KernelAbstractions.synchronize(backend)
end

function limit_gradient!(∇F::Grad{S,FF,R,I,M}, F, config) where {S,FF,R<:TensorField,I,M}
    (; hardware) = config
    (; backend, workgroup) = hardware

    mesh = F.mesh
    (; cells, cell_neighbours, cell_faces, cell_nsign, faces) = mesh

    # minPhi0 = maximum(F.values) # use min value so all values compared are larger
    # maxPhi0 = minimum(F.values)
    (; xx, yx, zx) = ∇F.result
    (; xy, yy, zy) = ∇F.result
    (; xz, yz, zz) = ∇F.result

    minPhi0 = maximum(F.x.values) # use min value so all values compared are larger
    maxPhi0 = minimum(F.x.values)
    kernel! = _limit_gradient!(backend, workgroup)
    kernel!(xx, yx, zx, F.x, cells, cell_neighbours, cell_faces, cell_nsign, faces, minPhi0, maxPhi0, ndrange=length(cells))
    KernelAbstractions.synchronize(backend)

    minPhi0 = maximum(F.y.values) # use min value so all values compared are larger
    maxPhi0 = minimum(F.y.values)
    kernel! = _limit_gradient!(backend, workgroup)
    kernel!(xy, yy, zy, F.y, cells, cell_neighbours, cell_faces, cell_nsign, faces, minPhi0, maxPhi0, ndrange=length(cells))
    KernelAbstractions.synchronize(backend)

    minPhi0 = maximum(F.z.values) # use min value so all values compared are larger
    maxPhi0 = minimum(F.z.values)
    kernel! = _limit_gradient!(backend, workgroup)
    kernel!(xz, yz, zz, F.z, cells, cell_neighbours, cell_faces, cell_nsign, faces, minPhi0, maxPhi0, ndrange=length(cells))
    KernelAbstractions.synchronize(backend)
end

# @kernel function _limit_gradient!(x, y, z, F, cells, cell_neighbours, cell_faces, cell_nsign, faces, minPhi, maxPhi)
#     cID = @index(Global)

#     cell = cells[cID]
#     faces_range = cell.faces_range
#     phiP = F[cID]
#     maxPhi = minPhi = phiP

#     for fi ∈ faces_range
#         nID = cell_neighbours[fi]
#         phiN = F[nID]
#         maxPhi = max(phiN, maxPhi)
#         minPhi = min(phiN, minPhi)
#     end

#     # g0 = ∇F[cID]
#     g0 = SVector{3}(x[cID] , y[cID] , z[cID])

#     cc = cell.centre

#     for fi ∈ faces_range 
#         fID = cell_faces[fi]
#         face = faces[fID]
#         nID = face.ownerCells[2]
#         # phiN = F[nID]
#         normal = face.normal
#         nsign = cell_nsign[fi]
#         na = nsign*normal

#         fc = face.centre 
#         cc_fc = fc - cc
#         n0 = cc_fc/norm(cc_fc)
#         gn = g0⋅n0
#         δϕ = g0⋅cc_fc
#         gτ = g0 - gn*n0
#         if (maxPhi > phiP) && (δϕ > maxPhi - phiP)
#             g0 = gτ + na*(maxPhi - phiP)
#         elseif (minPhi < phiP) && (δϕ < minPhi - phiP)
#             g0 = gτ + na*(minPhi - phiP)
#         end            
#     end
#     x.values[cID] = g0[1]
#     y.values[cID] = g0[2]
#     z.values[cID] = g0[3]
# end

@kernel function _limit_gradient!(x, y, z, F, cells, cell_neighbours, cell_faces, cell_nsign, faces, minPhi, maxPhi)
    cID = @index(Global)

    cell = cells[cID]
    faces_range = cell.faces_range
    phiP = F[cID]
    phiMax = phiMin = phiP
 
    for fi ∈ faces_range
        nID = cell_neighbours[fi]
        phiN = F[nID]
        phiMax = max(phiN, phiMax)
        phiMin = min(phiN, phiMin)
    end

    # g0 = ∇F[cID]
    grad0 = SVector{3}(x[cID] , y[cID] , z[cID])

    cc = cell.centre
    limiter = 1
    limiterf = 1
    for fi ∈ faces_range 
        fID = cell_faces[fi]
        nID = cell_neighbours[fi]
        face = faces[fID]
        cellN = cells[nID]
        # nID = face.ownerCells[2]
        # phiN = F[nID]
        normal = face.normal
        nsign = cell_nsign[fi]
        na = nsign*normal

        # r = fc - cc
        # fc = face.centre

        nc = cellN.centre
        r = nc - cc
        δϕ = r⋅grad0

        # rn = (nc - cc) ⋅ na
        # gradn = grad0⋅na
        # δϕ = rn* gradn
        if δϕ > 0
            limiterf = min(1, (phiMax - phiP)/δϕ)
        elseif δϕ < 0
            limiterf = min(1, (phiMin - phiP)/δϕ)
        else
            limiterf = 1
        end
        limiter = min(limiterf, limiter)
    end
    grad0 *= limiter
    x.values[cID] = grad0[1]
    y.values[cID] = grad0[2]
    z.values[cID] = grad0[3]
end