load_smap.jl

using Unitful, Base.Filesystem, HDF5, Dates
using Unitful: °C, K, hr, d
using Base: tail


# Fix unitful handling of missing
import Base.*
*(::Missing, ::T) where {T<:Unitful.Units} = missing
*(::T, ::Missing) where {T<:Unitful.Units} = missing

@inline cleanread(x, key) = map(x -> x == -9999.0 ? missing : x,  read(x[string(key)]))
@inline ignoremissing(x) = ismissing(x) ? false : x


# Models

abstract type AbstractStressModel end

scale(model::AbstractStressModel) = model.scale

needsname(x) = needsdata(x)[1]
needsunit(x) = needsdata(x)[2]


struct DryDays{L,S} <: AbstractStressModel
    lowersm::L
    scale::S
end

@inline needsdata(::DryDays) = :sm_surface, 1
@inline subset(m::DryDays, sm) = sm < m.lowersm
@inline accumulate(m::DryDays, x) = abs(x - m.lowersm)


struct ColdDays{L,S} <: AbstractStressModel
    lowerct::L
    scale::S
end

@inline needsdata(::ColdDays) = :surface_temp, K
@inline subset(m::ColdDays, temp) = temp < m.lowerct
@inline accumulate(m::ColdDays, x) = abs(x - m.lowerct)


struct HotDays{U,S} <: AbstractStressModel
    upperct::U
    scale::S
end

@inline needsdata(::HotDays) = :surface_temp, K
@inline subset(m::HotDays, temp) = temp > m.upperct
@inline accumulate(m::HotDays, x) = x - m.upperct


struct Wilting{W,S} <: AbstractStressModel
    lowerwilt::W
    scale::S
end

@inline needsdata(::Wilting) = :land_fraction_wilting, 1
@inline subset(m::Wilting, wilting) = wilting > m.lowerwilt
@inline accumulate(m::Wilting, x) = x - m.lowerwilt


struct IntrinsicPopGrowth{P,K,S} <: AbstractStressModel
    p25::P
    H_A::K
    H_L::K
    T_0_5L::K
    H_H::K
    T_0_5H::K
    scale::S
end

const R = 1.987

@inline needsdata(::IntrinsicPopGrowth) = :surface_temp, K
@inline subset(m::IntrinsicPopGrowth, temp) = true
@inline accumulate(m::IntrinsicPopGrowth, x) = m.p25 * (x * exp(m.H_A/R * (1/K(25°C) - 1/x))) /
    (1 + exp(m.H_L/R * (1/m.T_0_5L - 1/x)) + exp(m.H_H/R * (1/m.T_0_5H - 1/x)))



# Intermediate storage
struct StressModelStorage{A}
    accumulator::A
end

StressModelStorage(model::AbstractStressModel, init) = begin
    typ = typeof(1/scale(model))
    a = Union{typ, Missing}[zero(typ) for i in CartesianIndices(init)]
    StressModelStorage(a)
end


# Framework

" run all the stress models for the timeseries "
function growth_stress(models, data_path, start_date, end_date, timespan)

    filepaths, dates = daterange_filenames(data_path, start_date, end_date)

    # Allocate memory
    smap_storage = setup_smap_storage(models, filepaths[1])
    init = smap_storage[1]
    storage = allocate_model_storage(models, init)
    sub = similar(init, Bool)
    string.(keys(smap_storage))

    day_counter = 0.0d
    timestep = 1.0d
    dayfrac = timestep / 1.0d
    

    for i in 1:length(dates)
        println(dates[i])
        i > 1 && update_smap_storage!(smap_storage, models, filepaths[i])
        run_stressmodels!(storage, models, smap_storage, sub)
        day_counter += timestep
    end

    stresses = collate_stresses(models, storage) ./ day_counter ./ 1d
    for (i, s) in enumerate(storage)
        s.accumulator .= zero(1/scale(models[i])) 
    end

    stresses
end

load_geophysical(filepath) = h5open(filepath)["Geophysical_Data"]

function setup_smap_storage(models, filepath)
    geophysical = load_geophysical(filepath)
    required = tuple(union(needsname.(models))...)
    rasters = cleanread.(Ref(geophysical), required)
    NamedTuple{required}(rasters)
end
function update_smap_storage!(smap, models, filepath)
    geophysical = load_geophysical(filepath)
    for i = 1:length(smap) 
        smap[i] .= cleanread(geophysical, keys(smap)[i])
    end
end

daterange_filenames(data_path, start_date, end_date) = begin
    filenames = readdir(data_path)

    pattern = Regex("_Vv401(1|0)_001.h5.iso.xml")
    datanames = [replace(f, "_Vv4011_001.h5.iso.xml" => "") for f in filenames if occursin(pattern, f)]

    smapformat = dateformat"\S\MAP_L4_\S\M_gph_yyyymmddTHHMMSS"
    dates = DateTime.(datanames, smapformat)

    # Find subset of all available dates between start and end dates
    subind = (dates .>= start_date) .& (dates .<= end_date)

    # Subset dates and filenames
    dates = dates[subind]
    filepaths = data_path .* datanames[subind] .* "_Vv4011_001.h5"
    filepaths, dates
end

collate_stresses(models::Tuple, storage::Tuple) = 
    storage[1].accumulator .* models[1].scale .+ collate_stresses(tail(models), tail(storage))
collate_stresses(::Tuple{}, ::Tuple{}) = 0

allocate_model_storage(models::Tuple, init) =
    (StressModelStorage(models[1], init), allocate_model_storage(tail(models), init)...)
allocate_model_storage(::Tuple{}, init) = ()


run_stressmodels!(storage::Tuple, models::Tuple, smap, sub) = begin
    m = models[1]
    modeldata = smap[needsname(m)] * needsunit(m)
    sub .= ignoremissing.(subset.(Ref(m), modeldata))
    storage[1].accumulator .+= accumulate.(Ref(m), modeldata)
    run_stressmodels!(tail(storage), tail(models), smap, sub)
end
run_stressmodels!(storage::Tuple{}, models::Tuple{}, smap, sub) = ()





#################################################################
# User Script

p25 = 3.377850e-01K^-1
H_A = 3.574560e+04K
H_L = -1.108990e+05K
T_0_5L = 2.359187e+02K
H_H = 3.276604e+05K
T_0_5H = 2.991132e+02K
R = 1.987

ipgrowth = IntrinsicPopGrowth(p25, H_A, H_L, T_0_5L, H_H, T_0_5H, 1/K)

lowersm = 0 # no data
uppersm = 1 # no data
lowerct = 7°C |> K  # Enriquez2017
lowerctm = -log(1.00) * K^-1
upperct = 30°C |> K # Kimura2004
upperctm = -log(1.15) * K^-1
lowerwilt = 0.5 # default?
wiltingm = -log(1.1)

# general options
start_date = Date("2016-01-01")
end_date = Date("2016-12-31")
data_path = "/home/raf/Data/drosophila/climate/SMAP_raw/" # raw climatic data
output_folder = "SMAP_based_layers" # computed suitability layers folder
cold = ColdDays(lowerct, lowerctm)
hot = HotDays(upperct, upperctm)
wilt = Wilting(lowerwilt, wiltingm)

models = (cold, hot, wilt)

stresses = @time growth_stress(models, data_path, start_date, end_date, 365.25d/12)