# FTCS.jl
# Implementation of heterogenous FTCS (forward time-centered space)
# solution of diffusion equation in 1D and 2D space.
# Translated from C++'s FTCS.hpp. 

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

function calcHorizontalChange(alphaX_next::T, alphaX_prev::T, alphaX_current::T,
    concentration_next::T, concentration_prev::T, concentration_current::T) where {T}

    intercellAlpha_next = calcAlphaIntercell(alphaX_next, alphaX_current)
    intercellAlpha_prev = calcAlphaIntercell(alphaX_prev, alphaX_current)

    return intercellAlpha_next * concentration_next -
           (intercellAlpha_next + intercellAlpha_prev) * concentration_current +
           intercellAlpha_prev * concentration_prev
end

function calcHorizontalChangeLeftBoundary(boundaryType::TYPE, alphaX_next::T, alphaX_current::T,
    concentration_next::T, concentration_current::T, boundaryValue::T) where {T}
    intercellAlpha = calcAlphaIntercell(alphaX_next, alphaX_current)

    if boundaryType == CONSTANT
        return intercellAlpha * concentration_next -
               (intercellAlpha + 2 * alphaX_current) * concentration_current +
               2 * alphaX_current * boundaryValue
    elseif boundaryType == CLOSED
        return intercellAlpha * (concentration_next - concentration_current)
    else
        error("Undefined Boundary Condition Type!")
    end
end

function calcHorizontalChangeRightBoundary(boundaryType::TYPE, alphaX_prev::T, alphaX_current::T,
    concentration_prev::T, concentration_current::T, boundaryValue::T) where {T}
    intercellAlpha = calcAlphaIntercell(alphaX_prev, alphaX_current)

    if boundaryType == CONSTANT
        return 2 * alphaX_current * boundaryValue -
               (intercellAlpha + 2 * alphaX_current) * concentration_current +
               intercellAlpha * concentration_prev
    elseif boundaryType == CLOSED
        return -intercellAlpha * (concentration_current - concentration_prev)
    else
        error("Undefined Boundary Condition Type!")
    end
end

function calcVerticalChange(alphaY_above::T, alphaY_below::T, alphaY_current::T,
    concentration_above::T, concentration_below::T, concentration_current::T) where {T}
    intercellAlpha_above = calcAlphaIntercell(alphaY_above, alphaY_current)
    intercellAlpha_below = calcAlphaIntercell(alphaY_below, alphaY_current)

    return intercellAlpha_above * concentration_above -
           (intercellAlpha_above + intercellAlpha_below) * concentration_current +
           intercellAlpha_below * concentration_below
end

function calcVerticalChangeTopBoundary(boundaryType::TYPE, alphaY_above::T, alphaY_current::T,
    concentration_above::T, concentration_current::T, boundaryValue::T) where {T}
    intercellAlpha = calcAlphaIntercell(alphaY_above, alphaY_current)

    if boundaryType == CONSTANT
        return intercellAlpha * concentration_above -
               (intercellAlpha + 2 * alphaY_current) * concentration_current +
               2 * alphaY_current * boundaryValue
    elseif boundaryType == CLOSED
        return intercellAlpha * (concentration_above - concentration_current)
    else
        error("Undefined Boundary Condition Type!")
    end
end


function calcVerticalChangeBottomBoundary(boundaryType::TYPE, alphaY_current::T, alphaY_below::T,
    concentration_current::T, concentration_below::T, boundaryValue::T) where {T}
    intercellAlpha = calcAlphaIntercell(alphaY_current, alphaY_below)

    if boundaryType == CONSTANT
        return 2 * alphaY_current * boundaryValue -
               (intercellAlpha + 2 * alphaY_current) * concentration_current +
               intercellAlpha * concentration_below
    elseif boundaryType == CLOSED
        return -intercellAlpha * (concentration_current - concentration_below)
    else
        error("Undefined Boundary Condition Type!")
    end
end



function FTCS_1D(grid::Grid{T}, bc::Boundary{T}, timestep::T) where {T}
    colMax = getCols(grid)
    sx = timestep / (getDeltaCol(grid)^2)

    concentrations = getConcentrations(grid)
    alphaX = getAlphaX(grid)
    concentrations_t1 = copy(concentrations)

    row = 1

    # inner cells
    for col = 2:colMax-1
        concentrations_t1[row, col] += sx * calcHorizontalChange(alphaX[row, col+1], alphaX[row, col-1], alphaX[row, col],
            concentrations[row, col+1], concentrations[row, col-1], concentrations[row, col])
    end

    # left boundary
    leftBoundaryType = getBoundaryElementType(bc, LEFT, row)
    leftBoundaryValue = getBoundaryElementValue(bc, LEFT, row)
    concentrations_t1[row, 1] += sx * calcHorizontalChangeLeftBoundary(leftBoundaryType, alphaX[row, 2], alphaX[row, 1],
        concentrations[row, 2], concentrations[row, 1], leftBoundaryValue)

    # right boundary
    rightBoundaryType = getBoundaryElementType(bc, RIGHT, row)
    rightBoundaryValue = getBoundaryElementValue(bc, RIGHT, row)
    concentrations_t1[row, colMax] += sx * calcHorizontalChangeRightBoundary(rightBoundaryType, alphaX[row, colMax-1], alphaX[row, colMax],
        concentrations[row, colMax-1], concentrations[row, colMax], rightBoundaryValue)

    setConcentrations!(grid, concentrations_t1)
end


function FTCS_2D(grid::Grid{T}, bc::Boundary{T}, timestep::T) where {T}
    rowMax = getRows(grid)
    colMax = getCols(grid)
    sx = timestep / (getDeltaCol(grid)^2)
    sy = timestep / (getDeltaRow(grid)^2)

    alphaX = getAlphaX(grid)
    alphaY = getAlphaY(grid)
    concentrations = getConcentrations(grid)
    concentrations_t1 = copy(concentrations)

    # Currently boundary sides can only be set for the whole side, not for each cell individually
    leftBoundaryType = getBoundaryElementType(bc, LEFT, 1)
    rightBoundaryType = getBoundaryElementType(bc, RIGHT, 1)
    topBoundaryType = getBoundaryElementType(bc, TOP, 1)
    bottomBoundaryType = getBoundaryElementType(bc, BOTTOM, 1)

    Threads.@threads for row = 2:rowMax-1
        for col = 2:colMax-1
            concentrations_t1[row, col] += sy * calcVerticalChange(alphaY[row+1, col], alphaY[row-1, col], alphaY[row, col],
                concentrations[row+1, col], concentrations[row-1, col], concentrations[row, col]) +
                                           sx * calcHorizontalChange(alphaX[row, col+1], alphaX[row, col-1], alphaX[row, col],
                concentrations[row, col+1], concentrations[row, col-1], concentrations[row, col])
        end

        # Boundary conditions for each row
        # Left boundary without corners
        leftBoundaryValue = getBoundaryElementValue(bc, LEFT, row)
        concentrations_t1[row, 1] += sx * calcHorizontalChangeLeftBoundary(leftBoundaryType, alphaX[row, 2], alphaX[row, 1],
            concentrations[row, 2], concentrations[row, 1], leftBoundaryValue) +
                                     sy * calcVerticalChange(alphaY[row+1, 1], alphaY[row-1, 1], alphaY[row, 1],
            concentrations[row+1, 1], concentrations[row-1, 1], concentrations[row, 1])

        # Right boundary without corners
        rightBoundaryValue = getBoundaryElementValue(bc, RIGHT, row)
        concentrations_t1[row, colMax] += sx * calcHorizontalChangeRightBoundary(rightBoundaryType, alphaX[row, colMax-1], alphaX[row, colMax],
            concentrations[row, colMax-1], concentrations[row, colMax], rightBoundaryValue) +
                                          sy * calcVerticalChange(alphaY[row+1, colMax], alphaY[row-1, colMax], alphaY[row, colMax],
            concentrations[row+1, colMax], concentrations[row-1, colMax], concentrations[row, colMax])
    end

    # Handle top/bottom boundaries
    Threads.@threads for col = 2:colMax-1
        # Top boundary
        topBoundaryValue = getBoundaryElementValue(bc, TOP, col)
        concentrations_t1[1, col] += sy * calcVerticalChangeTopBoundary(topBoundaryType, alphaY[2, col], alphaY[1, col],
            concentrations[2, col], concentrations[1, col], topBoundaryValue) +
                                     sx * calcHorizontalChange(alphaX[1, col+1], alphaX[1, col-1], alphaX[1, col],
            concentrations[1, col+1], concentrations[1, col-1], concentrations[1, col])

        # Bottom boundary
        bottomBoundaryValue = getBoundaryElementValue(bc, BOTTOM, col)
        concentrations_t1[rowMax, col] += sy * calcVerticalChangeBottomBoundary(bottomBoundaryType, alphaY[rowMax, col], alphaY[rowMax-1, col],
            concentrations[rowMax, col], concentrations[rowMax-1, col], bottomBoundaryValue) +
                                          sx * calcHorizontalChange(alphaX[rowMax, col+1], alphaX[rowMax, col-1], alphaX[rowMax, col],
            concentrations[rowMax, col+1], concentrations[rowMax, col-1], concentrations[rowMax, col])
    end

    # Handle corners
    # Top left corner
    topBoundaryValue = getBoundaryElementValue(bc, TOP, 1)
    concentrations_t1[1, 1] += sx * calcHorizontalChange(alphaX[1, 2], alphaX[1, 1], alphaX[1, 1],
        concentrations[1, 2], concentrations[1, 1], concentrations[1, 1]) +
                               sy * calcVerticalChangeTopBoundary(topBoundaryType, alphaY[2, 1], alphaY[1, 1],
        concentrations[2, 1], concentrations[1, 1], topBoundaryValue)

    # Top right corner
    topBoundaryValue = getBoundaryElementValue(bc, TOP, colMax)
    concentrations_t1[1, colMax] += sx * calcHorizontalChange(alphaX[1, colMax-1], alphaX[1, colMax], alphaX[1, colMax],
        concentrations[1, colMax-1], concentrations[1, colMax], concentrations[1, colMax]) +
                                    sy * calcVerticalChangeTopBoundary(topBoundaryType, alphaY[2, colMax], alphaY[1, colMax],
        concentrations[2, colMax], concentrations[1, colMax], topBoundaryValue)

    # Bottom left corner
    bottomBoundaryValue = getBoundaryElementValue(bc, BOTTOM, 1)
    concentrations_t1[rowMax, 1] += sx * calcHorizontalChange(alphaX[rowMax, 2], alphaX[rowMax, 1], alphaX[rowMax, 1],
        concentrations[rowMax, 2], concentrations[rowMax, 1], concentrations[rowMax, 1]) +
                                    sy * calcVerticalChangeBottomBoundary(bottomBoundaryType, alphaY[rowMax, 1], alphaY[rowMax-1, 1],
        concentrations[rowMax, 1], concentrations[rowMax-1, 1], bottomBoundaryValue)

    # Bottom right corner
    bottomBoundaryValue = getBoundaryElementValue(bc, BOTTOM, colMax)
    concentrations_t1[rowMax, colMax] += sx * calcHorizontalChange(alphaX[rowMax, colMax-1], alphaX[rowMax, colMax], alphaX[rowMax, colMax],
        concentrations[rowMax, colMax-1], concentrations[rowMax, colMax], concentrations[rowMax, colMax]) +
                                         sy * calcVerticalChangeBottomBoundary(bottomBoundaryType, alphaY[rowMax, colMax], alphaY[rowMax-1, colMax],
        concentrations[rowMax, colMax], concentrations[rowMax-1, colMax], bottomBoundaryValue)

    setConcentrations!(grid, concentrations_t1)
end

function runFTCS(grid::Grid{T}, bc::Boundary{T}, timestep::T, iterations::Int, stepCallback::Function) where {T}
    if getDim(grid) == 1
        for _ = 1:iterations
            FTCS_1D(grid, bc, timestep)

            stepCallback()
        end
    elseif getDim(grid) == 2
        for _ = 1:iterations
            FTCS_2D(grid, bc, timestep)

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