# BTCS.jl
# Implementation of heterogenous BTCS (backward time-centered space)
# solution of diffusion equation in 1D and 2D space using the
# alternating-direction implicit (ADI) method.

using LinearAlgebra
using SparseArrays
using Base.Threads
using CUDA

include("../Boundary.jl")
include("../Grid.jl")

function calcAlphaIntercell(alpha1::T, alpha2::T) where {T}
    return 2 / ((1 / alpha1) + (1 / alpha2))
end

function calcAlphaIntercell(alpha1::Matrix{T}, alpha2::Matrix{T}) where {T}
    return 2 ./ ((1 ./ alpha1) .+ (1 ./ alpha2))
end

function calcBoundaryCoeffConstant(alpha_center::T, alpha_side::T, sx::T) where {T}
    alpha = calcAlphaIntercell(alpha_center, alpha_side)
    centerCoeff = 1 + sx * (alpha + 2 * alpha_center)
    sideCoeff = -sx * alpha
    return (centerCoeff, sideCoeff)
end

function calcBoundaryCoeffClosed(alpha_center::T, alpha_side::T, sx::T) where {T}
    alpha = calcAlphaIntercell(alpha_center, alpha_side)
    centerCoeff = 1 + sx * alpha
    sideCoeff = -sx * alpha
    return (centerCoeff, sideCoeff)
end

# creates coefficient matrix for next time step from alphas in x-direction
function createCoeffMatrix(alpha::Matrix{T}, alpha_left::Vector{T}, alpha_right::Vector{T}, bcLeft::Vector{BoundaryElement{T}}, bcRight::Vector{BoundaryElement{T}}, numCols::Int, rowIndex::Int, sx::T)::Tridiagonal{T} where {T}
    # Precompute boundary condition type check for efficiency
    bcTypeLeft = getType(bcLeft[rowIndex])

    # Determine left side boundary coefficients based on boundary condition
    centerCoeffTop, rightCoeffTop = if bcTypeLeft == CONSTANT
        calcBoundaryCoeffConstant(alpha[rowIndex, 1], alpha[rowIndex, 2], sx)
    elseif bcTypeLeft == CLOSED
        calcBoundaryCoeffClosed(alpha[rowIndex, 1], alpha[rowIndex, 2], sx)
    else
        error("Undefined Boundary Condition Type on Left!")
    end

    # Precompute boundary condition type check for efficiency
    bcTypeRight = getType(bcRight[rowIndex])

    # Determine right side boundary coefficients based on boundary condition
    centerCoeffBottom, leftCoeffBottom = if bcTypeRight == CONSTANT
        calcBoundaryCoeffConstant(alpha[rowIndex, numCols], alpha[rowIndex, numCols-1], sx)
    elseif bcTypeRight == CLOSED
        calcBoundaryCoeffClosed(alpha[rowIndex, numCols], alpha[rowIndex, numCols-1], sx)
    else
        error("Undefined Boundary Condition Type on Right!")
    end

    dl = [-sx .* alpha_left; leftCoeffBottom]
    du = [rightCoeffTop; -sx .* alpha_right]
    d = [centerCoeffTop; 1 .+ sx .* (alpha_right + alpha_left); centerCoeffBottom]
    alpha_diagonal = Tridiagonal(dl, d, du)

    return alpha_diagonal
end


function calcExplicitConcentrationsBoundaryClosed(conc_center::T, alpha_center::T, alpha_neighbor::T, sy::T) where {T}
    alpha = calcAlphaIntercell(alpha_center, alpha_neighbor)
    sy * alpha * conc_center + (1 - sy * alpha) * conc_center
end


function calcExplicitConcentrationsBoundaryConstant(conc_center::T, conc_bc::T, alpha_center::T, alpha_neighbor::T, sy::T) where {T}
    alpha_center_neighbor = calcAlphaIntercell(alpha_center, alpha_neighbor)
    alpha_center_center = alpha_center == alpha_neighbor ? alpha_center_neighbor : calcAlphaIntercell(alpha_center, alpha_center)
    sy * alpha_center_neighbor * conc_center +
    (1 - sy * (alpha_center_center + 2 * alpha_center)) * conc_center +
    sy * alpha_center * conc_bc
end

function writeSolutionVector!(sv::Vector{T}, concentrations::Matrix{T}, alphaX::Matrix{T}, alphaY::Matrix{T}, bcLeft::Vector{BoundaryElement{T}}, bcRight::Vector{BoundaryElement{T}}, bcTop::Vector{BoundaryElement{T}}, bcBottom::Vector{BoundaryElement{T}}, rowIndex::Int, sx::T, sy::T) where {T}
    numRows = size(concentrations, 1)
    length = size(sv, 1)

    # Inner rows
    if rowIndex > 1 && rowIndex < numRows
        @inbounds for i = 1:length
            alpha_here_below = calcAlphaIntercell(alphaY[rowIndex, i], alphaY[rowIndex+1, i])
            alpha_here_above = alphaY[rowIndex+1, i] == alphaY[rowIndex-1, i] ? alpha_here_below : calcAlphaIntercell(alphaY[rowIndex-1, i], alphaY[rowIndex, i]) # calcAlphaIntercell is symmetric, so we can use it for both directions
            sv[i] = sy * alpha_here_below * concentrations[rowIndex+1, i] +
                    (1 - sy * (alpha_here_below + alpha_here_above)) * concentrations[rowIndex, i] +
                    sy * alpha_here_above * concentrations[rowIndex-1, i]
        end
    end

    # First row
    if rowIndex == 1
        @inbounds for i = 1:length
            if getType(bcTop[i]) == CONSTANT
                sv[i] = calcExplicitConcentrationsBoundaryConstant(concentrations[rowIndex, i], getValue(bcTop[i]), alphaY[rowIndex, i], alphaY[rowIndex+1, i], sy)
            elseif getType(bcTop[i]) == CLOSED
                sv[i] = calcExplicitConcentrationsBoundaryClosed(concentrations[rowIndex, i], alphaY[rowIndex, i], alphaY[rowIndex+1, i], sy)
            else
                error("Undefined Boundary Condition Type somewhere on Left or Top!")
            end
        end
    end

    # Last row
    if rowIndex == numRows
        @inbounds for i = 1:length
            if getType(bcBottom[i]) == CONSTANT
                sv[i] = calcExplicitConcentrationsBoundaryConstant(concentrations[rowIndex, i], getValue(bcBottom[i]), alphaY[rowIndex, i], alphaY[rowIndex-1, i], sy)
            elseif getType(bcBottom[i]) == CLOSED
                sv[i] = calcExplicitConcentrationsBoundaryClosed(concentrations[rowIndex, i], alphaY[rowIndex, i], alphaY[rowIndex-1, i], sy)
            else
                error("Undefined Boundary Condition Type somewhere on Right or Bottom!")
            end
        end
    end

    # First column - additional fixed concentration change from perpendicular dimension in constant BC case
    if getType(bcLeft[rowIndex]) == CONSTANT
        sv[1] += 2 * sx * alphaX[rowIndex, 1] * getValue(bcLeft[rowIndex])
    end

    # Last column - additional fixed concentration change from perpendicular dimension in constant BC case
    if getType(bcRight[rowIndex]) == CONSTANT
        sv[end] += 2 * sx * alphaX[rowIndex, end] * getValue(bcRight[rowIndex])
    end
end

# BTCS solution for 1D grid
function BTCS_1D(grid::Grid{T}, bc::Boundary{T}, alpha_left::Matrix{T}, alpha_right::Matrix{T}, timestep::T) where {T}
    sx = timestep / (grid.deltaCol * grid.deltaCol)

    alpha = getAlphaX(grid)
    bcLeft = getBoundarySide(bc, LEFT)
    bcRight = getBoundarySide(bc, RIGHT)
    length = getCols(grid)

    concentrations::Matrix{T} = getConcentrations(grid)

    A::Tridiagonal{T} = createCoeffMatrix(alpha, alpha_left[1, :], alpha_right[1, :], bcLeft, bcRight, length, 1, sx)

    b = concentrations[1, :]

    if getType(bcLeft[1]) == CONSTANT
        b[1] += 2 * sx * alpha[1, 1] * getValue(bcLeft[1])
    end
    if getType(bcRight[1]) == CONSTANT
        b[end] += 2 * sx * alpha[1, end] * getValue(bcRight[1])
    end

    concentrations[1, :] = A \ b
end

# BTCS solution for 2D grid
function BTCS_2D(grid::Grid{T}, bc::Boundary{T}, alphaX_left::Matrix{T}, alphaX_right::Matrix{T}, alphaY_t_left::Matrix{T}, alphaY_t_right::Matrix{T}, timestep::T) where {T}
    rows = getRows(grid)
    cols = getCols(grid)
    sx = timestep / (2 * grid.deltaCol * grid.deltaCol)
    sy = timestep / (2 * grid.deltaRow * grid.deltaRow)

    alphaX = getAlphaX(grid)
    alphaY = getAlphaY(grid)

    alphaX_t = getAlphaX_t(grid)
    alphaY_t = getAlphaY_t(grid)

    concentrations = getConcentrations(grid)
    concentrations_intermediate = similar(concentrations)
    concentrations_t_task = Threads.@spawn copy(transpose(concentrations))

    bcLeft = getBoundarySide(bc, LEFT)
    bcRight = getBoundarySide(bc, RIGHT)
    bcTop = getBoundarySide(bc, TOP)
    bcBottom = getBoundarySide(bc, BOTTOM)

    localBs = [zeros(T, cols) for _ in 1:Threads.nthreads()]
    Threads.@threads for i = 1:rows
        localB = localBs[Threads.threadid()]
        A::Tridiagonal{T} = createCoeffMatrix(alphaX, alphaX_left[i, :], alphaX_right[i, :], bcLeft, bcRight, cols, i, sx)
        writeSolutionVector!(localB, concentrations, alphaX, alphaY, bcLeft, bcRight, bcTop, bcBottom, i, sx, sy)

        concentrations_intermediate[i, :] = A \ localB
    end

    concentrations_intermediate = copy(transpose(concentrations_intermediate))
    concentrations_t = fetch(concentrations_t_task)

    localBs = [zeros(T, rows) for _ in 1:Threads.nthreads()]
    Threads.@threads for i = 1:cols
        localB = localBs[Threads.threadid()]
        # Swap alphas, boundary conditions and sx/sy for column-wise calculation
        A::Tridiagonal{T} = createCoeffMatrix(alphaY_t, alphaY_t_left[i, :], alphaY_t_right[i, :], bcTop, bcBottom, rows, i, sy)
        writeSolutionVector!(localB, concentrations_intermediate, alphaY_t, alphaX_t, bcTop, bcBottom, bcLeft, bcRight, i, sy, sx)

        concentrations_t[i, :] = (A \ localB)
    end

    concentrations = copy(transpose(concentrations_t))

    setConcentrations!(grid, concentrations)
end

function runBTCS(grid::Grid{T}, bc::Boundary{T}, timestep::T, iterations::Int, stepCallback::Function) where {T}
    if grid.dim == 1
        length = getCols(grid)
        alpha = getAlphaX(grid)

        alpha_left_task = Threads.@spawn calcAlphaIntercell(alpha[:, 1:(length-2)], alpha[:, 2:(length-1)])
        alpha_right_task = Threads.@spawn calcAlphaIntercell(alpha[:, 2:(length-1)], alpha[:, 3:length])

        alpha_left = fetch(alpha_left_task)
        alpha_right = fetch(alpha_right_task)

        for _ in 1:(iterations)
            BTCS_1D(grid, bc, alpha_left, alpha_right, timestep)

            stepCallback()
        end
    elseif grid.dim == 2
        rows = getRows(grid)
        cols = getCols(grid)
        alphaX = getAlphaX(grid)
        alphaY_t = getAlphaY_t(grid)

        alphaX_left_task = Threads.@spawn calcAlphaIntercell(alphaX[:, 1:(cols-2)], alphaX[:, 2:(cols-1)])
        alphaX_right_task = Threads.@spawn calcAlphaIntercell(alphaX[:, 2:(cols-1)], alphaX[:, 3:cols])
        alphaY_t_left_task = Threads.@spawn calcAlphaIntercell(alphaY_t[:, 1:(rows-2)], alphaY_t[:, 2:(rows-1)])
        alphaY_t_right_task = Threads.@spawn calcAlphaIntercell(alphaY_t[:, 2:(rows-1)], alphaY_t[:, 3:rows])

        alphaX_left = fetch(alphaX_left_task)
        alphaX_right = fetch(alphaX_right_task)
        alphaY_t_left = fetch(alphaY_t_left_task)
        alphaY_t_right = fetch(alphaY_t_right_task)

        for _ in 1:(iterations)
            BTCS_2D(grid, bc, alphaX_left, alphaX_right, alphaY_t_left, alphaY_t_right, timestep)

            stepCallback()
        end
    else
        error("Error: Only 1- and 2-dimensional grids are defined!")
    end

end