### Persefone.jl - a model of agricultural landscapes and ecosystems in Europe.
###
### This file contains functions for saving the model output. This includes logging
### facilities, but also functions for collecting and outputting data from the
### different submodels.
###

"Log output is saved to `simulation.log` in the output directory"
const LOGFILE = "simulation.log"

"All input data are copied to the `inputs` folder within the output directory"
const RECORDDIR = "inputs"

## Much of this code was adapted from the GeMM model by Leidinger et al.
## (https://github.com/CCTB-Ecomods/gemm/blob/master/src/output.jl)

"""
    createdatadir(outdir, overwrite)

Creates the output directory, dealing with possible conflicts.
"""
function createdatadir(outdir::String, overwrite::Union{Bool,String})
    if isdir(outdir)
        if overwrite == "ask"
            println("The chosen output directory $(outdir) already exists.")
            println("Type 'yes' to overwrite this directory. Otherwise, the simulation will abort.")
            print("Overwrite? ")
            answer = readline()
            overwrite = (answer == "yes" || answer == "y")
        end
        !overwrite ? Base.error("Output directory exists, will not overwrite. Aborting.") :
            @warn "Overwriting existing output directory $(outdir)."
        #TODO replace with exception
    end
    @debug "Setting up output directory $outdir."
    mkpath(outdir)
end

"""
    modellogger(loglevel, outdir, output="both")

Create a logger object that writes output to screen and/or a logfile.
This object is stored as `model.logger` and can then be used with `with_logger()`.
Note: requires [`createdatadir`](@ref) to be run first.
"""
function modellogger(loglevel::String, outdir::String, output::String="both")
    !isdir(outdir) && #TODO replace with exception
        Base.error("$(outdir) does not exist. Call `createdatadir()` before `modellogger()`.")
    loglevel == "debug" ? loglevel = Logging.Debug :
        loglevel == "warn" ? loglevel = Logging.Warn :
        loglevel == "info" ? loglevel = Logging.Info :
        Base.error("Invalid loglevel $loglevel, should be debug/info/warn.") #TODO make exception
    (output in ["file", "both"]) && (logfile = open(joinpath(outdir, LOGFILE), "w+"))
    if output == "both"
        return TeeLogger(ConsoleLogger(logfile, loglevel),
                         ConsoleLogger(stdout, loglevel))
    elseif output == "file"
        return ConsoleLogger(logfile, loglevel)
    elseif output == "screen"
        return ConsoleLogger(stdout, loglevel)
    elseif output == "none"
        return NullLogger()
    else
        Base.error("Invalid log output target $output, should be file/screen/both.")
    end
end

"""
    withtestlogger(model)

Replace the model logger with the currently active logger. This is intended to be used
in the testsuite to circumvent a [Julia issue](https://github.com/JuliaLang/julia/issues/48456),
where `@test_logs` doesn't work with local loggers.
"""
function withtestlogger(model::SimulationModel)
    # copied together from https://github.com/JuliaLang/julia/blob/master/base/logging.jl
    logstate = current_task().logstate
    logstate == nothing ? model.logger = global_logger() : model.logger = logstate.logger
    model
end

"""
    saveinputfiles(model)

Copy all input files into the output directory, including the actual parameter
settings used. This allows replicating a run in future.
"""
function saveinputfiles(model::SimulationModel)
    #XXX If this is a parallel run, we should save the global config to the top-level
    # output directory
    currentcommit = "********"
    try # if this is loaded as a package, the directory may not be a git repo
        currentcommit = read(`git rev-parse HEAD`, String)[1:8]
    catch
        @debug "saveinputfiles(): not a git repo."
    end
    mkpath(joinpath(@param(core.outdir), RECORDDIR))
    open(joinpath(@param(core.outdir), RECORDDIR, basename(@param(core.configfile))), "w") do f
        println(f, "#\n# --- Persefone configuration parameters ---")
        println(f, "# This file was generated automatically.")
        println(f, "# Simulation run on $(string(Dates.format(Dates.now(), "d u Y HH:MM:SS"))),")
        # Record the current git commit and versions of dependencies for reproducibility
        println(f, "# with Persefone $(pkgversion(Persefone)), git commit $(currentcommit),")
        println(f, "# running on Julia $(VERSION).\n#\n")
        if !isempty(strip(read(`git status -s`, String)))
            println(f, "# WARNING: Your repository contains uncommitted changes. This may")
            println(f, "#          compromise the reproducibility of this simulation run.\n")
        end
        TOML.print(f, prepareTOML(deepcopy(model.settings)))
    end
    # Copy the map and weather files to the output folder
    lcmap = joinpath(@param(world.mapdirectory), @param(world.landcovermap))
    ffmap = joinpath(@param(world.mapdirectory), @param(world.farmfieldsmap))
    wfile = joinpath(@param(world.mapdirectory), @param(world.weatherfile))
    #TODO replace errors with exceptions
    !(isfile(lcmap)) && Base.error("The map file $(lcmap) doesn't exist.")
    !(isfile(ffmap)) && Base.error("The map file $(ffmap) doesn't exist.")
    !(isfile(wfile)) && Base.error("The map file $(wfile) doesn't exist.")
    cp(lcmap, joinpath(@param(core.outdir), RECORDDIR, basename(lcmap)), force = true)
    cp(ffmap, joinpath(@param(core.outdir), RECORDDIR, basename(ffmap)), force = true)
    cp(wfile, joinpath(@param(core.outdir), RECORDDIR, basename(wfile)), force = true)
end

"""
    prepareTOML(dict)

An internal utility function to re-convert the one-dimensional dict created
by [`flattenTOML`](@ref) into the two-dimensional dict needed by `TOML.print`,
and convert any data types into TOML-compatible types where necessary.
"""
function prepareTOML(settings)
    # convert data types
    settings["core.loglevel"] == Logging.Debug ? settings["core.loglevel"] = "debug" :
        settings["core.loglevel"] == Logging.Warn ? settings["core.loglevel"] = "warn" :
        settings["core.loglevel"] = "info"
    settings["world.mapresolution"] = settings["world.mapresolution"] / m
    # convert dict structure
    fulldict = Dict{String, Dict{String, Any}}()
    for parameter in keys(settings)
        domain, param = split(parameter, ".")
        !(domain in keys(fulldict)) && (fulldict[domain] = Dict{String,Any}())
        fulldict[domain][param] = settings[parameter]
    end
    fulldict
end

"""
    DataOutput

A struct for organising model output. This is used to collect model data
in an in-memory dataframe or for CSV output. Submodels can register their
own output functions using [`newdataoutput!`](@ref).

Struct fields:
    - frequency: how often to call the output function (daily/monthly/yearly/end/never)
    - databuffer: a vector of vectors that temporarily saves data before it is stored permanently or written to file
    - datastore: a data frame that stores data until the end of the run
    - outputfunction: a function that takes a model object and returns data values to record (formatted as a vector of vectors)
    - plotfunction: a function that takes a model object and returns a Makie figure object (optional)
"""
mutable struct DataOutput
    frequency::String
    databuffer::Vector{Vector}
    datastore::DataFrame
    outputfunction::Union{Function,Nothing}
    plotfunction::Union{Function,Nothing} #XXX remove this? (#81)
end

"Retrieve the data stored in a DataOutput (assumes `core.storedata` is true)."
data(d::DataOutput) = d.datastore

##TODO what about output types that don't fit neatly into the standard CSV table format?
## (e.g. end-of-run stats, map data)

"""
    newdataoutput!(model, name, header, frequency, outputfunction, plotfunction)

Create and register a new data output. This function must be called by all submodels
that want to have their output functions called regularly.
"""
function newdataoutput!(model::SimulationModel, name::String,
                        header::Vector{String}, frequency::String,
                        outputfunction::Union{Function,Nothing}=nothing,
                        plotfunction::Union{Function,Nothing}=nothing) #XXX remove this? (#81)
    if !(frequency in ("daily", "monthly", "yearly", "end", "never"))
        Base.error("Invalid frequency '$frequency' for $name.") #TODO replace with exception
    end
    if frequency != "never"
        if @param(core.csvoutput)
            open(joinpath(@param(core.outdir), name*".csv"), "w") do f
                println(f, join(header, ","))
            end
        end
    end
    df = DataFrame()
    for h in header
        df[!,h] = Any[] #XXX allow specifying types?
    end
    ndo = DataOutput(frequency, [], df, outputfunction, plotfunction)
    model.dataoutputs[name] = ndo
end

"""
    outputdata(model, force=false)
 
Cycle through all registered data outputs and activate them according to their
configured frequency. If `force` is `true`, activate all outputs regardless
of their configuration.
"""
function outputdata(model::SimulationModel, force=false)
    #XXX enable output weekly, or on set dates?
    #XXX all output functions except for "end" are run on the first update
    # -> should they all be run on the last update, too?
    startdate = @param(core.startdate)
    isnextmonth = d -> (day(d) == day(startdate))
    isnextyear = d -> (month(d) == month(startdate) && day(d) == day(startdate))
    for o in keys(model.dataoutputs)
        output = model.dataoutputs[o]
        (!force && output.frequency == "never") && continue
        # check if this output should be activated today
        if force || (output.frequency == "daily") ||
            (output.frequency == "monthly" && isnextmonth(model.date)) ||
            (output.frequency == "yearly" && isnextyear(model.date)) ||
            (output.frequency == "end" && model.date == @param(core.enddate))
            !isnothing(output.outputfunction) && (output.databuffer = output.outputfunction(model))
            (isempty(output.databuffer) || output.databuffer == [[]]) && continue
            if @param(core.csvoutput)
                open(joinpath(@param(core.outdir), o*".csv"), "a") do f
                    for row in output.databuffer
                        println(f, join(row, ","))
                    end
                end                
            end
            if @param(core.storedata)
                for row in output.databuffer
                    push!(output.datastore, row)
                end
            end
            output.databuffer = []
        end
    end
end

"""
    record!(model, outputname, data)

Append an observation vector to the given output.
"""
function record!(model::SimulationModel, outputname::String, data::Vector)
    !(outputname in keys(model.dataoutputs)) && return #XXX should this be a warning?
    push!(model.dataoutputs[outputname].databuffer, data)
end

"""
    @record(outputname, data)

Record an observation / data point. Only use in scopes where `model` is available.
"""
macro record(args...)
    :(record!($(esc(:model)), $(map(esc, args)...)))
end

"""
    @data(outputname)

Return the data stored in the given output (assumes `core.storedata` is true).
Only use in scopes where `model` is available.
"""
macro data(outputname)
    :(data($(esc(:model)).dataoutputs[$(esc(outputname))]))
end
    
"""
    visualiseoutput(model)

Cycle through all data outputs and call their respective plot functions,
saving each figure to file.
"""
function visualiseoutput(model::SimulationModel) #XXX remove this? (#81)
    @debug "Visualising output."
    CairoMakie.activate!() # make sure we're using Cairo
    for o in keys(model.dataoutputs)
        output = model.dataoutputs[o]
        (output.frequency == "never" || isnothing(output.plotfunction)) && continue
        figure = output.plotfunction(model)
        isnothing(figure) ? continue :
            save(joinpath(@param(core.outdir), o*"."*@param(core.figureformat)), figure)
    end
end

"""
    savemodelobject(model, filename)

Serialise a model object and save it to file for later reference.
Includes the current model and Julia versions for compatibility checking.

WARNING: produces large files (>100 MB) and takes a while to execute.
"""
function savemodelobject(model::SimulationModel, filename::String)
    object = Dict("model"=>model,
                  "modelversion"=>pkgversion(Persefone),
                  "juliaversion"=>VERSION)
    !endswith(filename, ".dat") && (filename *= ".dat")
    filename = joinpath(@param(core.outdir), filename)
    serialize(filename, object)
    @debug "Saved model object to $(filename)."
end