# 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(grid::Grid{T}, row::Int, col::Int) where {T}
    return calcAlphaIntercell(getAlphaX(grid)[row, col+1],
        getAlphaX(grid)[row, col]) *
           getConcentrations(grid)[row, col+1] -
           (calcAlphaIntercell(getAlphaX(grid)[row, col+1],
        getAlphaX(grid)[row, col]) +
            calcAlphaIntercell(getAlphaX(grid)[row, col-1],
        getAlphaX(grid)[row, col])) *
           getConcentrations(grid)[row, col] +
           calcAlphaIntercell(getAlphaX(grid)[row, col-1],
        getAlphaX(grid)[row, col]) *
           getConcentrations(grid)[row, col-1]
end

function calcHorizontalChangeLeftBoundary(grid::Grid{T}, bc::Boundary{T}, row::Int, col::Int) where {T}
    if getBoundaryElementType(bc, LEFT, row) == CONSTANT
        return calcHorizontalChangeLeftBoundaryConstant(grid, bc, row, col)
    elseif getBoundaryElementType(bc, LEFT, row) == CLOSED
        return calcHorizontalChangeLeftBoundaryClosed(grid, row, col)
    else
        error("Undefined Boundary Condition Type!")
    end
end

function calcHorizontalChangeLeftBoundaryConstant(grid::Grid{T}, bc::Boundary{T}, row::Int, col::Int) where {T}
    return calcAlphaIntercell(getAlphaX(grid)[row, col+1],
        getAlphaX(grid)[row, col]) *
           getConcentrations(grid)[row, col+1] -
           (calcAlphaIntercell(getAlphaX(grid)[row, col+1],
        getAlphaX(grid)[row, col]) +
            2 * getAlphaX(grid)[row, col]) *
           getConcentrations(grid)[row, col] +
           2 * getAlphaX(grid)[row, col] *
           getBoundaryElementValue(bc, LEFT, row)
end

function calcHorizontalChangeLeftBoundaryClosed(grid::Grid{T}, row::Int, col::Int) where {T}
    return calcAlphaIntercell(getAlphaX(grid)[row, col+1],
        getAlphaX(grid)[row, col]) *
           (getConcentrations(grid)[row, col+1] -
            getConcentrations(grid)[row, col])
end


function calcHorizontalChangeRightBoundary(grid::Grid{T}, bc::Boundary{T}, row::Int, col::Int) where {T}
    if getBoundaryElementType(bc, RIGHT, row) == CONSTANT
        return calcHorizontalChangeRightBoundaryConstant(grid, bc, row, col)
    elseif getBoundaryElementType(bc, RIGHT, row) == CLOSED
        return calcHorizontalChangeRightBoundaryClosed(grid, row, col)
    else
        error("Undefined Boundary Condition Type!")
    end
end

function calcHorizontalChangeRightBoundaryConstant(grid::Grid{T}, bc::Boundary{T}, row::Int, col::Int) where {T}
    return 2 * getAlphaX(grid)[row, col] *
           getBoundaryElementValue(bc, RIGHT, row) -
           (calcAlphaIntercell(getAlphaX(grid)[row, col-1],
        getAlphaX(grid)[row, col]) +
            2 * getAlphaX(grid)[row, col]) *
           getConcentrations(grid)[row, col] +
           calcAlphaIntercell(getAlphaX(grid)[row, col-1],
        getAlphaX(grid)[row, col]) *
           getConcentrations(grid)[row, col-1]
end

function calcHorizontalChangeRightBoundaryClosed(grid::Grid{T}, row::Int, col::Int) where {T}
    return -(calcAlphaIntercell(getAlphaX(grid)[row, col-1],
        getAlphaX(grid)[row, col]) *
             (getConcentrations(grid)[row, col] -
              getConcentrations(grid)[row, col-1]))
end


function calcVerticalChange(grid::Grid{T}, row::Int, col::Int) where {T}
    return calcAlphaIntercell(getAlphaY(grid)[row+1, col],
        getAlphaY(grid)[row, col]) *
           getConcentrations(grid)[row+1, col] -
           (calcAlphaIntercell(getAlphaY(grid)[row+1, col],
        getAlphaY(grid)[row, col]) +
            calcAlphaIntercell(getAlphaY(grid)[row-1, col],
        getAlphaY(grid)[row, col])) *
           getConcentrations(grid)[row, col] +
           calcAlphaIntercell(getAlphaY(grid)[row-1, col],
        getAlphaY(grid)[row, col]) *
           getConcentrations(grid)[row-1, col]
end

function calcVerticalChangeTopBoundary(grid::Grid{T}, bc::Boundary{T}, row::Int, col::Int) where {T}
    if getBoundaryElementType(bc, TOP, col) == CONSTANT
        return calcVerticalChangeTopBoundaryConstant(grid, bc, row, col)
    elseif getBoundaryElementType(bc, TOP, col) == CLOSED
        return calcVerticalChangeTopBoundaryClosed(grid, row, col)
    else
        error("Undefined Boundary Condition Type!")
    end
end

function calcVerticalChangeTopBoundaryConstant(grid::Grid{T}, bc::Boundary{T}, row::Int, col::Int) where {T}
    return calcAlphaIntercell(getAlphaY(grid)[row+1, col],
        getAlphaY(grid)[row, col]) *
           getConcentrations(grid)[row+1, col] -
           (calcAlphaIntercell(getAlphaY(grid)[row+1, col],
        getAlphaY(grid)[row, col]) +
            2 * getAlphaY(grid)[row, col]) *
           getConcentrations(grid)[row, col] +
           2 * getAlphaY(grid)[row, col] *
           getBoundaryElementValue(bc, TOP, col)
end

function calcVerticalChangeTopBoundaryClosed(grid::Grid{T}, row::Int, col::Int) where {T}
    return calcAlphaIntercell(getAlphaY(grid)[row+1, col],
        getAlphaY(grid)[row, col]) *
           (getConcentrations(grid)[row+1, col] -
            getConcentrations(grid)[row, col])
end

function calcVerticalChangeBottomBoundary(grid::Grid{T}, bc::Boundary{T}, row::Int, col::Int) where {T}
    if getBoundaryElementType(bc, BOTTOM, col) == CONSTANT
        return calcVerticalChangeBottomBoundaryConstant(grid, bc, row, col)
    elseif getBoundaryElementType(bc, BOTTOM, col) == CLOSED
        return calcVerticalChangeBottomBoundaryClosed(grid, row, col)
    else
        error("Undefined Boundary Condition Type!")
    end
end

function calcVerticalChangeBottomBoundaryConstant(grid::Grid{T}, bc::Boundary{T}, row::Int, col::Int) where {T}
    return 2 * getAlphaY(grid)[row, col] *
           getBoundaryElementValue(bc, BOTTOM, col) -
           (calcAlphaIntercell(getAlphaY(grid)[row, col],
        getAlphaY(grid)[row-1, col]) +
            2 * getAlphaY(grid)[row, col]) *
           getConcentrations(grid)[row, col] +
           calcAlphaIntercell(getAlphaY(grid)[row, col],
        getAlphaY(grid)[row-1, col]) *
           getConcentrations(grid)[row-1, col]
end

function calcVerticalChangeBottomBoundaryClosed(grid::Grid{T}, row::Int, col::Int) where {T}
    return -(calcAlphaIntercell(getAlphaY(grid)[row, col],
        getAlphaY(grid)[row-1, col]) *
             (getConcentrations(grid)[row, col] -
              getConcentrations(grid)[row-1, col]))
end



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

    # matrix for concentrations at time t+1
    concentrations_t1 = zeros(1, colMax)

    # only one row in 1D case -> row constant at index 1
    row = 1

    # inner cells
    # independent of boundary condition type
    for col = 2:colMax-1
        concentrations_t1[row, col] = getConcentrations(grid)[row, col] +
                                      sx * calcHorizontalChange(grid, row, col)
    end

    # left boundary; hold column constant at index 1
    concentrations_t1[row, 1] = getConcentrations(grid)[row, 1] +
                                sx * calcHorizontalChangeLeftBoundary(grid, bc, row, 1)

    # right boundary; hold column constant at max index
    concentrations_t1[row, colMax] = getConcentrations(grid)[row, colMax] +
                                     sx * calcHorizontalChangeRightBoundary(grid, bc, row, colMax)

    # overwrite obsolete concentrations
    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)

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

    # inner cells
    # these are independent of the boundary condition type
    for row = 2:rowMax-1
        # inner cells
        for col = 2:colMax-1
            concentrations_t1[row, col] += sy * calcVerticalChange(grid, row, col) +
                                           sx * calcHorizontalChange(grid, row, col)
        end

        # left boundary without corners
        concentrations_t1[row, 1] += sx * calcHorizontalChangeLeftBoundary(grid, bc, row, 1) +
                                     sy * calcVerticalChange(grid, row, 1)
        # right boundary without corners
        concentrations_t1[row, colMax] += sx * calcHorizontalChangeRightBoundary(grid, bc, row, colMax) +
                                          sy * calcVerticalChange(grid, row, colMax)
    end

    # top / bottom without corners / looping over columns
    for col = 2:colMax-1
        # top
        concentrations_t1[1, col] += sy * calcVerticalChangeTopBoundary(grid, bc, 1, col) +
                                     sx * calcHorizontalChange(grid, 1, col)
        # bottom
        concentrations_t1[rowMax, col] += sy * calcVerticalChangeBottomBoundary(grid, bc, rowMax, col) +
                                          sx * calcHorizontalChange(grid, rowMax, col)
    end

    # corner top left
    concentrations_t1[1, 1] += sx * calcHorizontalChangeLeftBoundary(grid, bc, 1, 1) +
                               sy * calcVerticalChangeTopBoundary(grid, bc, 1, 1)

    # corner top right
    concentrations_t1[1, colMax] += sx * calcHorizontalChangeRightBoundary(grid, bc, 1, colMax) +
                                    sy * calcVerticalChangeTopBoundary(grid, bc, 1, colMax)

    # corner bottom left
    concentrations_t1[rowMax, 1] += sx * calcHorizontalChangeLeftBoundary(grid, bc, rowMax, 1) +
                                    sy * calcVerticalChangeBottomBoundary(grid, bc, rowMax, 1)

    # corner bottom right
    concentrations_t1[rowMax, colMax] += sx * calcHorizontalChangeRightBoundary(grid, bc, rowMax, colMax) +
                                         sy * calcVerticalChangeBottomBoundary(grid, bc, rowMax, colMax)

    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