# Grid.jl
# API of Grid class, that holds a matrix with concenctrations and 
# a respective matrix/matrices of alpha coefficients.
# Translated from C++'s Grid.hpp.

using LinearAlgebra

struct Grid{T}
    cols::Int
    rows::Int
    dim::Int
    domainCol::T
    domainRow::T
    deltaCol::T
    deltaRow::T
    concentrations::Ref{Matrix{T}}
    alphaX::Ref{Matrix{T}}
    alphaY::Union{Ref{Matrix{T}},Nothing}
    alphaX_t::Union{Ref{Matrix{T}},Nothing}
    alphaY_t::Union{Ref{Matrix{T}},Nothing}

    # Constructor for 1D-Grid
    function Grid{T}(length::Int, alphaX::Matrix{T}) where {T}
        if length <= 3
            throw(ArgumentError("Given grid length too small. Must be greater than 3."))
        end
        if size(alphaX, 1) != 1 || size(alphaX, 2) != length
            throw(ArgumentError("Given matrix of alpha coefficients mismatch with Grid dimensions!"))
        end

        alphaX_t = alphaX'

        new{T}(length, 1, 1, T(length), 0, T(1), 0, Ref(fill(T(0), 1, length)), alphaX, nothing, alphaX_t, nothing)
    end

    # Constructor for 2D-Grid
    function Grid{T}(rows::Int, cols::Int, alphaX::Matrix{T}, alphaY::Matrix{T}) where {T}
        if rows <= 3 || cols <= 3
            throw(ArgumentError("Given grid dimensions too small. Must each be greater than 3."))
        end
        if size(alphaX) != (rows, cols) || size(alphaY) != (rows, cols)
            throw(ArgumentError("Given matrices of alpha coefficients mismatch with Grid dimensions!"))
        end

        # Precompute alphaX_t and alphaY_t
        alphaX_t = alphaX'
        alphaY_t = alphaY'

        new{T}(cols, rows, 2, T(cols), T(rows), T(1), T(1), Ref(fill(T(0), rows, cols)), alphaX, alphaY, alphaX_t, alphaY_t)
    end

    function Grid{T}(rows::Int, cols::Int, dim::Int, domainCol::T, domainRow::T, deltaCol::T, deltaRow::T, concentrations::Ref{Matrix{T}}, alphaX::Ref{Matrix{T}}, alphaY::Union{Ref{Matrix{T}},Nothing}, alphaX_t::Union{Ref{Matrix{T}},Nothing}, alphaY_t::Union{Ref{Matrix{T}},Nothing}) where {T}
        new{T}(cols, rows, dim, domainCol, domainRow, deltaCol, deltaRow, concentrations, alphaX, alphaY, alphaX_t, alphaY_t)
    end
end

function clone(grid::Grid{T})::Grid{T} where {T}
    if grid.dim == 1
        return Grid{T}(1, grid.cols, grid.dim, grid.domainCol, grid.domainRow, grid.deltaCol, grid.deltaRow, Ref(copy(grid.concentrations[])), Ref(copy(grid.alphaX[])), nothing, Ref(copy(grid.alphaX_t[])), nothing)
    end
    Grid{T}(grid.rows, grid.cols, grid.dim, grid.domainCol, grid.domainRow, grid.deltaCol, grid.deltaRow, Ref(copy(grid.concentrations[])), Ref(copy(grid.alphaX[])), Ref(copy(grid.alphaY[])), Ref(copy(grid.alphaX_t[])), Ref(copy(grid.alphaY_t[])))
end

function getAlphaX(grid::Grid{T})::Matrix{T} where {T}
    grid.alphaX[]
end

function getAlphaY(grid::Grid{T})::Matrix{T} where {T}
    if grid.dim == 1
        error("Grid is 1D, so there is no alphaY matrix!")
    end

    grid.alphaY[]
end

function getAlphaX_t(grid::Grid{T})::Matrix{T} where {T}
    grid.alphaX_t[]
end

function getAlphaY_t(grid::Grid{T})::Matrix{T} where {T}
    if grid.dim == 1
        error("Grid is 1D, so there is no alphaY_t matrix!")
    end

    grid.alphaY_t[]
end

function getCols(grid::Grid{T})::Int where {T}
    grid.cols
end

function getConcentrations(grid::Grid{T})::Matrix{T} where {T}
    grid.concentrations[]
end

function getDeltaCol(grid::Grid{T})::T where {T}
    grid.deltaCol
end

function getDeltaRow(grid::Grid{T})::T where {T}
    grid.deltaRow
end

function getDim(grid::Grid{T})::Int where {T}
    grid.dim
end

function getRows(grid::Grid{T})::Int where {T}
    grid.rows
end

function setAlphaX!(grid::Grid{T}, new_alphaX::Matrix{T}) where {T}
    if size(new_alphaX) != size(grid.alphaX[])
        throw(ArgumentError("Given matrix of alpha coefficients mismatch with Grid dimensions!"))
    end

    grid.alphaX[] = new_alphaX
    grid.alphaX_t[] = new_alphaX'
end

function setAlphaY!(grid::Grid{T}, new_alphaY::Matrix{T}) where {T}
    if grid.dim == 1
        error("Grid is 1D, so there is no alphaY matrix!")
    end
    if size(new_alphaY) != size(grid.alphaY[])
        throw(ArgumentError("Given matrix of alpha coefficients mismatch with Grid dimensions!"))
    end

    grid.alphaY[] = new_alphaY
    grid.alphaY_t[] = new_alphaY'
end

function setConcentrations!(grid::Grid{T}, new_concentrations::Matrix{T}) where {T}
    if size(new_concentrations) != size(grid.concentrations[])
        throw(ArgumentError("Given matrix of concentrations mismatch with Grid dimensions!"))
    end

    grid.concentrations[] = new_concentrations
end