# Persefone.jl wrapper for AquaCrop.jl

module AquaCropWrapper

const CROPFILE = "crop_data.csv"

import AquaCrop
import CSV
using Dates: Date
using DataFrames: DataFrame

using Persefone:
    AbstractCropState,
    AbstractCropType,
    AnnualDate,
    cm,
    daynumber,
    SoilType,
    SimulationModel

import Persefone:
    stepagent!,
    croptype,
    cropname,
    cropheight,
    cropcover,
    cropyield,
    makecropstate,
    sow!,
    harvest!,
    isharvestable,
    setsoiltype!

using Unitful: @u_str

# We can't use Length{Float64} as that is a Union type
const Tlength = typeof(1.0cm)

# TODO: read crop names directly from AquaCrop.jl
# names extracted by running this command in the AquaCrop.jl src dir:
#     grep -nir 'crop_type ==' src/ | cut -d '=' -f 3 | sed 's/$/,/' | sort | uniq
const AQUACROP_CROPNAMES = [
    "alfalfaGDD",
    "barley",
    "barleyGDD",
    "cotton",
    "cottonGDD",
    "drybean",
    "drybeanGDD",
    "maize",
    "maizeGDD",
    "oat",
    "paddyrice",
    "paddyriceGDD",
    "potato",
    "potatoGDD",
    "quinoa",
    "rapeseed",
    "sorghum",
    "sorghumGDD",
    "soybean",
    "soybeanGDD",
    "sugarbeet",
    "sugarbeetGDD",
    "sugarcane",
    "sunflower",
    "sunflowerGDD",
    "tef",
    "tomato",
    "tomatoGDD",
    "wheat",
    "wheatGDD",
]

@kwdef struct CropType <: AbstractCropType
    name::String
    aquacrop_name::String
    group::String
    minsowdate::Union{Missing,AnnualDate}
    maxsowdate::Union{Missing,AnnualDate}
end

cropname(ct::CropType) = ct.name

"""
    readcropparameters(cropdirectory)

Parse a CSV file containing some parameters required to map AquaCrop
crop names to Persefone crop names as well as additional crop data
needed for Persefone (cropgroup, minsowdate, maxsowdate).
"""
function readcropparameters(cropdirectory::String)
    @debug "Reading extra crop parameters for AquaCrop crop model"
    croptypes = Dict{String, AbstractCropType}()
    df = CSV.read(joinpath(cropdirectory, CROPFILE), DataFrame;
                  missingstring="NA", types=[String, String, String, AnnualDate, AnnualDate])
    for row in eachrow(df)
        croptypes[row.persefone_name] =
            CropType(name = row.persefone_name,
                     aquacrop_name = row.aquacrop_name,
                     group = row.crop_group,
                     minsowdate = row.min_sow_date,
                     maxsowdate = row.max_sow_date)
    end
    return croptypes
end

function make_aquacrop_cropfield(croptype::CropType,
                                 model::SimulationModel,
                                 soiltype::SoilType)
    # TODO: get this mapping for soil types from a CSV file?
    soiltype_to_aquacrop = Dict(soiltype => replace(string(soiltype), "_" => " ")
                                for soiltype in instances(SoilType))
    aquacrop_soiltype = soiltype_to_aquacrop[soiltype]

    start_date = model.date            # TODO: maybe `today(model)`
    end_date = model.weather.lastdate  # TODO: maybe `lastdate(model)`
    start_daynum = daynumber(model.weather, start_date)
    end_daynum = daynumber(model.weather, end_date)
    dayrange = start_daynum:end_daynum
    input_Tmin = @view model.weather.mintemp[dayrange]
    input_Tmax = @view model.weather.maxtemp[dayrange]
    input_ETo = @view model.weather.evapotranspiration[dayrange]
    input_Rain = @view model.weather.precipitation[dayrange]

    # Generate the keyword object for the AquaCrop simulation
    kwargs = (
        ## Necessary keywords
        # runtype
        runtype = AquaCrop.NoFileRun(),
        # project input
        Simulation_DayNr1 = start_date,
        Simulation_DayNrN = end_date,
        Crop_Day1 = start_date,    # TODO: originally start_date + Week(1)
        Crop_DayN = end_date,
        # soil
        soil_type = aquacrop_soiltype,
        # crop
        crop_type = croptype.aquacrop_name,
        # Climate
        InitialClimDate = start_date,
        ## Optional keyworkds
        # Climate
        Tmin = input_Tmin,
        Tmax = input_Tmax,
        ETo = input_ETo,
        Rain = input_Rain,
        # change soil properties
        soil_layers = Dict("Thickness" => 5.0)
    )

    aquacrop_cropfield, all_ok = AquaCrop.start_cropfield(; kwargs...)
    if ! all_ok.logi
        @error "Failed calling AquaCrop.start_cropfield()\nAquaCrop error: $(all_ok.msg)"
    end

    return aquacrop_cropfield
end

mutable struct CropState <: AbstractCropState
    croptype::CropType
    height::Tlength  # TODO: remove height field, supply from cropstate
    soiltype::SoilType
    cropstate::AquaCrop.AquaCropField

    function CropState(croptype::CropType, soiltype::SoilType,
                       model::SimulationModel, height::Tlength=0.0cm)
        aquacrop_cropfield = make_aquacrop_cropfield(croptype, model, soiltype)
        return new(croptype, height, soiltype, aquacrop_cropfield)
    end
end

croptype(cs::CropState) = cs.croptype
cropname(cs::CropState) = cropname(croptype(cs))
cropheight(cs::CropState) = cs.height  # TODO: calculate from AquaCrop state info
cropcover(cs::CropState) = AquaCrop.canopycover(cs.cropstate)
cropyield(cs::CropState) = AquaCrop.dryyield(cs.cropstate)  # TODO: there is also freshyield
function isharvestable(cs::CropState)
    stage = get_aquacrop_stage(cs)
    return !ismissing(stage) && stage == 0
end
makecropstate(croptype::CropType, model::SimulationModel, soiltype::SoilType) =
    CropState(croptype, soiltype, model)
function setsoiltype!(cs::CropState, soiltype::SoilType)
    # TODO: this does not affect the actual soil type of the state of
    # the aquacrop simulation
    cs.soiltype = soiltype
    return nothing
end

function get_aquacrop_stage(cs::CropState)
    stages = cs.cropstate.dayout.Stage
    length(stages) > 0 ? last(stages) : missing
end

"""
    stepagent!(cropstate, model)

Update a crop state by one day.
"""
function stepagent!(cs::CropState, model::SimulationModel)
    AquaCrop.dailyupdate!(cs.cropstate)
end

"""
    sow!(cropstate, model, cropname)

Change the cropstate to sow the specified crop.
"""
function sow!(cs::CropState, model::SimulationModel, cropname::String)
    if cropname ∉ keys(model.crops)
        @error "cropname \"$cropname\" not found"
    end
    new_croptype = model.crops[cropname]
    if typeof(new_croptype) !== CropType
        error("Cannot change crop model of AquaCrop crop state.")
    end

    cs.croptype = new_croptype
    cs.height = 0.0cm
    # cs.soiltype stays the way it is
    cs.cropstate = make_aquacrop_cropfield(cs.croptype, model, soiltype)
end

"""
    harvest!(cropstate, model)

Harvest the crop of this cropstate.
"""
function harvest!(cs::CropState, model::SimulationModel)
    # TODO: implement this
    return 0.0u"g/m^2"
end

end # module