# 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}

Represents a computational grid with concentration values and alpha coefficients.
It can be either 1D or 2D and includes dimensions, domain sizes, delta values, and matrices for concentrations and alpha coefficients.

# Fields
- `cols::Int`: Number of columns in the grid.
- `rows::Int`: Number of rows in the grid.
- `dim::Int`: Dimension of the grid (1 for 1D, 2 for 2D).
- `domainCol::T`, `domainRow::T`: Size of the grid's domain in column and row direction.
- `deltaCol::T`, `deltaRow::T`: Delta values for columns and rows.
- `concentrations::Ref{Matrix{T}}`: Reference to the matrix holding concentration values.
- `alphaX::Ref{Matrix{T}}`: Reference to the matrix of alpha coefficients in the X direction.
- `alphaY::Union{Ref{Matrix{T}},Nothing}`: Reference to the matrix of alpha coefficients in the Y direction (if applicable).
- `alphaX_t`, `alphaY_t`: Transposed matrices of alpha coefficients.

# Constructors
- `Grid(length, alphaX)` creates a 1D grid with the given length and alphaX matrix.
- `Grid(rows, cols, alphaX, alphaY)` creates a 2D grid with the given rows, columns, and alphaX and alphaY matrices.
- `Grid(rows, cols, dim, domainCol, domainRow, deltaCol, deltaRow, concentrations, alphaX, alphaY, alphaX_t, alphaY_t)` creates a grid with the given parameters.
"""
struct Grid{T}
    cols::Int
    rows::Int
    dim::Int
    domainCol::T
    domainRow::T
    deltaCol::T
    deltaRow::T
    concentrations::Matrix{T}
    alphaX::Matrix{T}
    alphaY::Union{Matrix{T},Nothing}
    alphaX_t::Matrix{T}
    alphaY_t::Union{Matrix{T},Nothing}

    function Grid{T}(length::Int, alphaX::Matrix{T})::Grid{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,
            fill(T(0), 1, length),
            alphaX,
            nothing,
            alphaX_t,
            nothing,
        )
    end

    function Grid{T}(
        rows::Int,
        cols::Int,
        alphaX::Matrix{T},
        alphaY::Matrix{T},
    )::Grid{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

        alphaX_t = alphaX'
        alphaY_t = alphaY'

        new{T}(
            cols,
            rows,
            2,
            T(cols),
            T(rows),
            T(1),
            T(1),
            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::Matrix{T},
        alphaX::Matrix{T},
        alphaY::Union{Matrix{T},Nothing},
        alphaX_t::Matrix{T},
        alphaY_t::Union{Matrix{T},Nothing},
    )::Grid{T} where {T}
        new{T}(
            cols,
            rows,
            dim,
            domainCol,
            domainRow,
            deltaCol,
            deltaRow,
            concentrations,
            alphaX,
            alphaY,
            alphaX_t,
            alphaY_t,
        )
    end
end

"""
    clone(grid::Grid{T})::Grid{T} where {T}

Creates a deep copy of the specified grid. The new grid will have identical properties 
but will be a distinct object in memory.

# Arguments
- `grid::Grid{T}`: The grid to be cloned.

# Returns
A new `Grid{T}` object that is a deep copy of the input grid.
"""
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,
            copy(grid.concentrations),
            copy(grid.alphaX),
            nothing,
            copy(grid.alphaX_t),
            nothing,
        )
    end
    Grid{T}(
        grid.rows,
        grid.cols,
        grid.dim,
        grid.domainCol,
        grid.domainRow,
        grid.deltaCol,
        grid.deltaRow,
        copy(grid.concentrations),
        copy(grid.alphaX),
        copy(grid.alphaY),
        copy(grid.alphaX_t),
        copy(grid.alphaY_t),
    )
end

"""
    getAlphaX(grid::Grid{T})::Matrix{T} where {T}

Retrieves the alpha coefficients matrix in the X direction from the specified grid.

# Arguments
- `grid::Grid{T}`: The grid from which to retrieve the alphaX matrix.

# Returns
The `alphaX` matrix of the grid.
"""
function getAlphaX(grid::Grid{T})::Matrix{T} where {T}
    grid.alphaX
end

"""
    getAlphaY(grid::Grid{T})::Matrix{T} where {T}

Retrieves the alpha coefficients matrix in the Y direction from the specified grid.
Raises an error if the grid is 1D (no alphaY matrix).

# Arguments
- `grid::Grid{T}`: The grid from which to retrieve the alphaY matrix.

# Returns
The `alphaY` matrix of the grid, if applicable.
"""
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

"""
    getAlphaX_t(grid::Grid{T})::Matrix{T} where {T}

Retrieves the transposed alpha coefficients matrix in the X direction from the specified grid.

# Arguments
- `grid::Grid{T}`: The grid from which to retrieve the alphaX_t matrix.

# Returns
The transposed `alphaX_t` matrix of the grid.
"""
function getAlphaX_t(grid::Grid{T})::Matrix{T} where {T}
    grid.alphaX_t
end

"""
    getAlphaY_t(grid::Grid{T})::Matrix{T} where {T}

Retrieves the transposed alpha coefficients matrix in the Y direction from the specified grid.
Raises an error if the grid is 1D (no alphaY_t matrix).

# Arguments
- `grid::Grid{T}`: The grid from which to retrieve the alphaY_t matrix.

# Returns
The transposed `alphaY_t` matrix of the grid, if applicable.
"""
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

"""
    getCols(grid::Grid{T})::Int where {T}

Retrieves the number of columns in the specified grid.

# Arguments
- `grid::Grid{T}`: The grid from which to retrieve the number of columns.

# Returns
The number of columns in the grid.
"""
function getCols(grid::Grid{T})::Int where {T}
    grid.cols
end

"""
    getConcentrations(grid::Grid{T})::Matrix{T} where {T}

Retrieves the concentration matrix from the specified grid.

# Arguments
- `grid::Grid{T}`: The grid from which to retrieve the concentration matrix.

# Returns
The concentration matrix of the grid.
"""
function getConcentrations(grid::Grid{T})::Matrix{T} where {T}
    grid.concentrations
end

"""
    getDeltaCol(grid::Grid{T})::T where {T}

Retrieves the delta value for columns in the specified grid.

# Arguments
- `grid::Grid{T}`: The grid from which to retrieve the delta column value.

# Returns
The delta value for columns in the grid.
"""
function getDeltaCol(grid::Grid{T})::T where {T}
    grid.deltaCol
end

"""
    getDeltaRow(grid::Grid{T})::T where {T}

Retrieves the delta value for rows in the specified grid.

# Arguments
- `grid::Grid{T}`: The grid from which to retrieve the delta row value.

# Returns
The delta value for rows in the grid.
"""
function getDeltaRow(grid::Grid{T})::T where {T}
    grid.deltaRow
end

"""
    getDim(grid::Grid{T})::Int where {T}

Retrieves the dimension (1D or 2D) of the specified grid.

# Arguments
- `grid::Grid{T}`: The grid from which to retrieve the dimension.

# Returns
The dimension of the grid (1 for 1D, 2 for 2D).
"""
function getDim(grid::Grid{T})::Int where {T}
    grid.dim
end

"""
    getRows(grid::Grid{T})::Int where {T}

Retrieves the number of rows in the specified grid.

# Arguments
- `grid::Grid{T}`: The grid from which to retrieve the number of rows.

# Returns
The number of rows in the grid.
"""
function getRows(grid::Grid{T})::Int where {T}
    grid.rows
end

"""
    setAlphaX!(grid::Grid{T}, new_alphaX::Matrix{T})::Nothing where {T}

Sets the alpha coefficients matrix in the X direction for the specified grid.
Throws an error if the dimensions of the new matrix don't match the grid's dimensions.

# Arguments
- `grid::Grid{T}`: The grid for which to set the new alphaX matrix.
- `new_alphaX::Matrix{T}`: The new alpha coefficients matrix in the X direction.

# Returns
Nothing, but modifies the alphaX matrix of the grid.
"""
function setAlphaX!(grid::Grid{T}, new_alphaX::Matrix{T})::Nothing 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'
    return
end

"""
    setAlphaY!(grid::Grid{T}, new_alphaY::Matrix{T})::Nothing where {T}

Sets the alpha coefficients matrix in the Y direction for the specified grid.
Throws an error if the grid is 1D or if the dimensions of the new matrix don't match the grid's dimensions.

# Arguments
- `grid::Grid{T}`: The grid for which to set the new alphaY matrix.
- `new_alphaY::Matrix{T}`: The new alpha coefficients matrix in the Y direction.

# Returns
Nothing, but modifies the alphaY matrix of the grid.
"""
function setAlphaY!(grid::Grid{T}, new_alphaY::Matrix{T})::Nothing 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'
    return
end

"""
    setConcentrations!(grid::Grid{T}, new_concentrations::Matrix{T})::Nothing where {T}

Sets the concentration matrix for the specified grid.
Throws an error if the dimensions of the new matrix don't match the grid's dimensions.

# Arguments
- `grid::Grid{T}`: The grid for which to set the new concentrations matrix.
- `new_concentrations::Matrix{T}`: The new concentrations matrix.

# Returns
Nothing, but modifies the concentration matrix of the grid.
"""
function setConcentrations!(grid::Grid{T}, new_concentrations::Matrix{T})::Nothing where {T}
    if size(new_concentrations) != size(grid.concentrations)
        throw(
            ArgumentError("Given matrix of concentrations mismatch with Grid dimensions!"),
        )
    end

    grid.concentrations .= new_concentrations
    return
end