### Persefone.jl - a model of agricultural landscapes and ecosystems in Europe.
###
### This file visualises model output using Makie.
###

const MIN_AXISLIMIT::Float64 = 1e-5

"""
    visualisemap(model, date, landcover)

Draw the model's land cover map and plot all individuals as points on it at
the specified date. If no date is passed, use the last date for which data
are available. Optionally, you can pass a landcover map image (this is needed
to reduce the frequency of disk I/O for Persefone Desktop).
Returns a Makie figure object.
"""
function visualisemap(model::SimulationModel,date=nothing,landcover=nothing)
    # load and plot the map
    # Note: if the landcover map is supplied, it needs to be rotr90'ed
    if isnothing(landcover)
        lcm = joinpath(@param(world.mapdirectory), @param(world.landcovermap))
        landcover = rotr90(load(lcm))
    end
    f = Figure()
    ax = Axis(f[1,1])
    hidedecorations!(ax)
    image!(f[1,1], landcover)
    ax.aspect = DataAspect()
    # check if there are individuals and plot them
    inds = @subset(@data("individuals"), :X .!= -1) #remove migrants
    if iszero(size(inds)[1])
        @debug "No individual data to map"
        return f
    end
    isnothing(date) && (date = inds.Date[end])
    #XXX other colour schemes: :tab10, :Accent_8, :Dark2_8, :Paired_12, :Set1_9
    # https://juliagraphics.github.io/ColorSchemes.jl/stable/catalogue/
    ncolors = max(2, length(@param(nature.targetspecies)))
    update_theme!(palette=(color=cgrad(:seaborn_bright, ncolors),), cycle=[:color])
    for s in @param(nature.targetspecies)
        points = @select!(@subset(inds, :Species .== s, :Date .== date),
                          :X, :Y)
        iszero(size(points)[1]) && continue
        # The origin in Makie is in the bottom-left rather than in the top-left as
        # on the model map, so we have to invert the Y coordinates
        @transform!(points, :Y = size(landcover)[2] .- :Y)
        scatter!(f[1,1], Matrix{Float32}(points), markersize=8)
    end
    f
end

"""
    populationtrends(model)

Plot a line graph of population sizes of each species over time.
Returns a Makie figure object.
"""
function populationtrends(model::SimulationModel)
    pops = @data("populations")
    ncolors = max(2, length(@param(nature.targetspecies)))
    update_theme!(palette=(color=cgrad(:seaborn_bright, ncolors),), cycle=[:color])
    f = Figure()
    dates = @param(core.startdate):@param(core.enddate)
    axlimits = (1, length(dates), 0, max(MIN_AXISLIMIT, maximum(pops[!,:Abundance]; init=0)*1.1))
    ax = Axis(f[1,1], xlabel="Date", ylabel="Population size",
              limits=axlimits, xticks=datetickmarks(dates))
    for s in @param(nature.targetspecies)
        points = @select!(@subset(pops, :Species .== s), :Abundance)
        iszero(size(points)[1]) && continue
        lines!(f[1,1], Vector{Float32}(points.Abundance), linewidth=3, label=s)
    end
    size(pops)[1] > 0 && axislegend("Species"; position=:lt)
    f
end

"""
    skylarkpopulation(model)

Plot a line graph of total population size and individual demographics of skylarks over time.
Returns a Makie figure object.
"""
function skylarkpopulation(model::SimulationModel)
    !("Skylark" in @param(nature.targetspecies)) && return nothing
    pops = @data("skylark_abundance")
    f = Figure()
    dates = @param(core.startdate):@param(core.enddate)
    axlimits = (1, length(dates), 0, max(MIN_AXISLIMIT, maximum(pops[!,:TotalAbundance]; init=0)))
    ax = Axis(f[1,1], xlabel="Date", ylabel="Population size",
              limits=axlimits, xticks=datetickmarks(dates))
    lines!(f[1,1], Vector{Float32}(pops.TotalAbundance), linewidth=3, label="Total abundance")
    lines!(f[1,1], Vector{Float32}(pops.Mating), linewidth=3, label="Mating")
    lines!(f[1,1], Vector{Float32}(pops.Breeding), linewidth=3, label="Breeding")
    lines!(f[1,1], Vector{Float32}(pops.Nonbreeding), linewidth=3, label="Non-breeding")
    lines!(f[1,1], Vector{Float32}(pops.Juvenile), linewidth=3, label="Juvenile")
    lines!(f[1,1], Vector{Float32}(pops.Migrants), linewidth=3, label="Migrants")
    axislegend("Demographic"; position=:lt)
    f
end

"""
    skylarkstats(model)

Plot various statistics from the skylark model: nesting habitat, territory size, mortality.
"""
function skylarkstats(model::SimulationModel)
    !("Skylark" in @param(nature.targetspecies)) && return nothing
    f = Figure()
    nestingdata = @data("skylark_breeding")
    mortalitydata = @subset(@data("mortality"), :Species .== "Skylark")
    # nesting habitat
    habitattype = unique(nestingdata.Landcover)
    htfrequency = [count(x -> x==h, nestingdata.Landcover) for h in habitattype]
    if ! isempty(habitattype)
        barplot(f[1,1], 1:length(habitattype), htfrequency,
                axis = (xticks=(1:length(habitattype), habitattype), title="Nesting habitat"))
    end
    croptype = unique(nestingdata.Crop)
    ctfrequency = [count(x -> x==c, nestingdata.Crop) for c in croptype]
    if ! isempty(croptype)
        barplot(f[2,1], 1:length(croptype), ctfrequency,
                axis = (xticks=(1:length(croptype), croptype), title="Nesting crop type"))
    end
    # mortality
    causes = unique(mortalitydata.Cause)
    mtfrequency = [count(x -> x==c, mortalitydata.Cause) for c in causes]
    if ! isempty(causes)
        barplot(f[1,2], 1:length(causes), mtfrequency,
                axis = (xticks=(1:length(causes), causes), title="Causes of Mortality"))
    end
    # territory size
    if ! isempty(nestingdata.Territory)
        hist(f[2,2], nestingdata.Territory, bins=20,
             axis=(ylabel="Frequency", title="Territory size (ha)"))
    end
    f
end

"""
    croptrends(model)

Plot a dual line graph of cropped area and average plant height per crop over time.
Returns a Makie figure object.
"""
function croptrends(model::SimulationModel)
    cropdata = @data("fields")
    f = Figure(size=(1200,1000))
    dates = @param(core.startdate):@param(core.enddate)
    # plot cropped area over time
    axlimitsarea = (1, length(dates), 0, max(MIN_AXISLIMIT, maximum(cropdata[!,:Area]; init=0)*1.1))
    ax1 = Axis(f[1,1], xlabel="Date", ylabel="Cropped area (ha)",
               limits=axlimitsarea, xticks=datetickmarks(dates))
    for c in unique(cropdata.Crop)
        points = @select!(@subset(cropdata, :Crop .== c), :Area)
        (iszero(size(points)[1]) || iszero(sum(points.Area))) && continue
        lines!(f[1,1], Vector{Float32}(points.Area), linewidth=3, label=c)
    end
    if ! isempty(cropdata.Crop)
        axislegend("Crop"; position=:rt)
    end
    # plot average crop height over time
    axlimitsheight = (1, length(dates), 0, max(MIN_AXISLIMIT, maximum(cropdata[!,:Height]; init=0)*1.1))
    ax2 = Axis(f[2,1], xlabel="Date", ylabel="Average plant height (cm)",
              limits=axlimitsheight, xticks=datetickmarks(dates))
    for c in unique(cropdata.Crop)
        points = @select!(@subset(cropdata, :Crop .== c), :Height)
        (iszero(size(points)[1]) || iszero(sum(points.Height))) && c != "no growth" && continue
        lines!(f[2,1], Vector{Float32}(points.Height), linewidth=3, label=c)
    end
    if ! isempty(cropdata.Crop)
        axislegend("Crop"; position=:rt)
    end
    f
end

"""
    datetickmarks(dates)

Given a vector of dates, construct a selection to use as tick mark locations.
Helper function for `[populationtrends](@ref)`
"""
function datetickmarks(dates)
    ticks = (Int[], String[])
    for i in 1:length(dates)
        if Day(dates[i]) == Day(1)
            push!(ticks[1], i)
            push!(ticks[2], Dates.format(dates[i], "u yy"))
        end
    end
    while length(ticks[1]) > 12
        deleteat!(ticks[1], 2:2:length(ticks[1]))
        deleteat!(ticks[2], 2:2:length(ticks[2]))            
    end
    ticks
end

#XXX add animation with record()? https://docs.makie.org/stable/explanations/animation/