# 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.
# Translated from C++'s BTCS.hpp.

using Base.Threads
using LinearAlgebra

function calcBoundaryCoeff(
    alpha_center::T,
    alpha_side::T,
    sx::T,
    bcType::TYPE,
)::Tuple{T,T} where {T}
    alpha = calcAlphaIntercell(alpha_center, alpha_side)
    sideCoeff = -sx * alpha

    if bcType == CONSTANT
        centerCoeff = 1 + sx * (alpha + 2 * alpha_center)
    elseif bcType == CLOSED
        centerCoeff = 1 + sx * alpha
    else
        error("Undefined Boundary Condition Type!")
    end

    return (centerCoeff, sideCoeff)
end

function createCoeffMatrix(
    alpha::SubArray{T},
    alpha_left::SubArray{T},
    alpha_right::SubArray{T},
    bcLeft::Vector{BoundaryElement{T}},
    bcRight::Vector{BoundaryElement{T}},
    numCols::Int,
    rowIndex::Int,
    sx::T,
)::Tridiagonal{T} where {T}
    centerCoeffTop, rightCoeffTop =
        calcBoundaryCoeff(alpha[1], alpha[2], sx, getType(bcLeft[rowIndex]))

    centerCoeffBottom, leftCoeffBottom =
        calcBoundaryCoeff(alpha[numCols], alpha[numCols-1], sx, getType(bcRight[rowIndex]))

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


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

function calcExplicitConcentrationsBoundaryConstant(
    conc_center::T,
    conc_bc::T,
    alpha_center::T,
    alpha_neighbor::T,
    sy::T,
)::T where {T}
    alpha_center_neighbor = calcAlphaIntercell(alpha_center, alpha_neighbor)

    return sy * alpha_center_neighbor * conc_center +
           (1 - sy * (alpha_center_neighbor + alpha_center)) * conc_center +
           sy * alpha_center * conc_bc
end

function calcExplicitConcentrationsBoundaryClosed(
    conc_center::Vector{T},
    alpha::Vector{T},
    sy::T,
)::T where {T}
    return (sy .* alpha .+ (1 .- sy .* alpha)) .* conc_center
end

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

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

    if rowIndex > 1 && rowIndex < numRows
        @inbounds for i in eachindex(sv)
            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])

            sv[i] =
                sy * alpha_here_below * concentrations_t[i, rowIndex+1] +
                (1 - sy * (alpha_here_below + alpha_here_above)) *
                concentrations_t[i, rowIndex] +
                sy * alpha_here_above * concentrations_t[i, rowIndex-1]
        end
    end

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

    if getType(bcLeft) == CONSTANT
        sv[1] += 2 * sx * alphaX[rowIndex, 1] * getValue(bcLeft)
    end

    if getType(bcRight) == CONSTANT
        sv[end] += 2 * sx * alphaX[rowIndex, end] * getValue(bcRight)
    end
end

function BTCS_1D(
    grid::Grid{T},
    bc::Boundary{T},
    alpha_left::Matrix{T},
    alpha_right::Matrix{T},
    timestep::T,
) where {T}
    sx = timestep / (getDeltaCol(grid) * getDeltaCol(grid))

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

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

    A::Tridiagonal{T} = createCoeffMatrix(
        view(alpha, 1, :),
        view(alpha_left, 1, :),
        view(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

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 * getDeltaCol(grid) * getDeltaCol(grid))
    sy = timestep / (2 * getDeltaRow(grid) * getDeltaRow(grid))

    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 = copy(transpose(concentrations))

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

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

        concentrations_intermediate[i, :] = A \ localB
    end

    # Swap alphas, boundary conditions and sx/sy for column-wise calculation
    Threads.@threads for i = 1:cols
        localB = zeros(T, cols)
        A::Tridiagonal{T} = createCoeffMatrix(
            view(alphaY_t, i, :),
            view(alphaY_t_left, i, :),
            view(alphaY_t_right, i, :),
            bcTop,
            bcBottom,
            rows,
            i,
            sy,
        )
        writeSolutionVector!(
            localB,
            concentrations_intermediate,
            alphaY_t,
            alphaX_t,
            bcTop[i],
            bcBottom[i],
            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 getDim(grid) == 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 _ = 1:(iterations)
            BTCS_1D(grid, bc, alpha_left, alpha_right, timestep)

            stepCallback()
        end
    elseif getDim(grid) == 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 _ = 1:(iterations)
            BTCS_2D(
                grid,
                bc,
                alphaX_left,
                alphaX_right,
                alphaY_t_left,
                alphaY_t_right,
                timestep,
            )

            stepCallback()
        end
    else
        error("BTCS only implemented for 1D and 2D grids")
    end
end