# Simulation.jl
# API of Simulation class, that holds all information regarding a
# specific simulation run like its timestep, number of iterations and output
# options. Simulation object also holds a predefined Grid and Boundary object.
# Translated from C++'s Simulation.hpp.

using Printf

include("Grid.jl")
include("Boundary.jl")
include("Core/BTCS.jl")
include("Core/FTCS.jl")

@enum APPROACH BTCS FTCS
@enum CONSOLE_OUTPUT CONSOLE_OUTPUT_OFF CONSOLE_OUTPUT_ON CONSOLE_OUTPUT_VERBOSE
@enum CSV_OUTPUT CSV_OUTPUT_OFF CSV_OUTPUT_ON CSV_OUTPUT_VERBOSE CSV_OUTPUT_XTREME

# Simulation class
struct Simulation{T}
    grid::Grid{T}
    bc::Boundary{T}
    approach::APPROACH

    iterations::Int
    timestep::T

    consoleOutput::CONSOLE_OUTPUT
    csvOutput::CSV_OUTPUT

    # Constructor
    function Simulation(grid::Grid{T}, bc::Boundary{T}, approach::APPROACH=BTCS, iterations::Int=1, timestep::T=0.1,
        consoleOutput::CONSOLE_OUTPUT=CONSOLE_OUTPUT_OFF, csvOutput::CSV_OUTPUT=CSV_OUTPUT_OFF) where {T}
        new{T}(grid, bc, approach, iterations, timestep, consoleOutput, csvOutput)
    end
end

function _adjustTimestep(grid::Grid{T}, approach::APPROACH, timestep::T, iterations::Int)::Tuple{T,Int} where {T}
    if approach == FTCS
        if getDim(grid) == 1
            deltaSquare = getDeltaCol(grid)
            maxAlpha = maximum(getAlphaX(grid))

            # Courant-Friedrichs-Lewy condition
            cfl = deltaSquare / (4 * maxAlpha)
        elseif getDim(grid) == 2
            deltaColSquare = getDeltaCol(grid) * getDeltaCol(grid)
            deltaRowSquare = getDeltaRow(grid) * getDeltaRow(grid)
            minDeltaSquare = min(deltaColSquare, deltaRowSquare)

            maxAlpha = min(maximum(getAlphaX(grid)), maximum(getAlphaY(grid)))

            cfl = minDeltaSquare / (4 * maxAlpha)
        end

        if timestep > cfl
            innerIterations = ceil(Int, timestep / cfl)
            iterations = iterations * innerIterations
            timestep = timestep / innerIterations
            println("Warning: Timestep is too large for FTCS approach. Adjusting timestep to ", timestep, " and iterations to ", iterations, ".")
        else
            timestep = timestep
        end
    end
    return timestep, iterations
end

function _createCSVfile(simulation::Simulation{T})::IOStream where {T}
    approachString = string(simulation.approach)
    rows = getRows(simulation.grid)
    cols = getCols(simulation.grid)
    numIterations = simulation.iterations

    filename = string(approachString, "_", rows, "_", cols, "_", numIterations, ".csv")

    appendIdent = 0
    while isfile(filename)
        appendIdent += 1
        filename = string(approachString, "_", rows, "_", cols, "_", numIterations, "-", appendIdent, ".csv")
    end

    # Write boundary conditions if required
    if simulation.csvOutput >= CSV_OUTPUT_XTREME
        open(filename, "w") do file
            _writeBoundarySideValues(file, simulation.bc, LEFT)
            _writeBoundarySideValues(file, simulation.bc, RIGHT)

            if getDim(simulation.grid) == 2
                _writeBoundarySideValues(file, simulation.bc, TOP)
                _writeBoundarySideValues(file, simulation.bc, BOTTOM)
            end

            write(file, "\n\n")
        end
    end

    file = open(filename, "a")
    return file
end

function _writeBoundarySideValues(file::IOStream, bc::Boundary{T}, side) where {T}
    values::Vector{BoundaryElement} = getBoundarySide(bc, side)
    formatted_values = join(map(getValue, values), " ")
    write(file, formatted_values, "\n")
end

function _printConcentrationsCSV(file::IOStream, simulation::Simulation{T}) where {T}
    concentrations = getConcentrations(simulation.grid)

    for row in eachrow(concentrations)
        formatted_row = [Printf.@sprintf("%.6g", x) for x in row]  # Format each element like is done in the C++ version using Eigen3
        println(file, join(formatted_row, " "))
    end
    println(file)
    println(file)
end

function _printConcentrations(simulation::Simulation{T}) where {T}
    println(getConcentrations(simulation.grid))
end

function run(simulation::Simulation{T}) where {T}
    file = nothing
    try
        if simulation.csvOutput > CSV_OUTPUT_OFF
            file = _createCSVfile(simulation)
        end

        function simulationStepCallback()
            if simulation.consoleOutput >= CONSOLE_OUTPUT_VERBOSE
                _printConcentrations(simulation)
            end
            if simulation.csvOutput >= CSV_OUTPUT_VERBOSE
                _printConcentrationsCSV(file, simulation)
            end
        end

        if simulation.approach == BTCS
            runBTCS(simulation.grid, simulation.bc, simulation.timestep, simulation.iterations, simulationStepCallback)
        elseif simulation.approach == FTCS
            timestep, iterations = _adjustTimestep(simulation.grid, simulation.approach, simulation.timestep, simulation.iterations)
            runFTCS(simulation.grid, simulation.bc, timestep, iterations, simulationStepCallback)
        else
            error("Undefined approach!")
        end

        if simulation.consoleOutput >= CONSOLE_OUTPUT_ON
            _printConcentrations(simulation)
        end

        if simulation.csvOutput >= CSV_OUTPUT_ON
            _printConcentrationsCSV(file, simulation)
        end

    finally
        if file !== nothing
            close(file)
        end
    end
end

function setIterations(simulation::Simulation{T}, iterations::Int)::Simulation{T} where {T}
    return Simulation(simulation.grid, simulation.bc, simulation.approach, iterations, simulation.timestep, simulation.consoleOutput, simulation.csvOutput)
end

function setOutputConsole(simulation::Simulation{T}, consoleOutput::CONSOLE_OUTPUT)::Simulation{T} where {T}
    return Simulation(simulation.grid, simulation.bc, simulation.approach, simulation.iterations, simulation.timestep, consoleOutput, simulation.csvOutput)
end

function setOutputCSV(simulation::Simulation{T}, csvOutput::CSV_OUTPUT)::Simulation{T} where {T}
    return Simulation(simulation.grid, simulation.bc, simulation.approach, simulation.iterations, simulation.timestep, simulation.consoleOutput, csvOutput)
end

function setTimestep(simulation::Simulation{T}, timestep::T)::Simulation{T} where {T}
    return Simulation(simulation.grid, simulation.bc, simulation.approach, simulation.iterations, timestep, simulation.consoleOutput, simulation.csvOutput)
end