"""
    function temporalsubset(C::ClimGrid, startdate::Date, enddate::Date)

Returns the temporal subset of ClimGrid C. The temporal subset is defined by a start and end date.

"""
function temporalsubset(C::ClimGrid, datebeg::Tuple, dateend::Tuple)

    T = typeof(get_timevec(C)[1])
    timeV = get_timevec(C)
    idxtimebeg, idxtimeend = ClimateTools.timeindex(timeV, datebeg, dateend, T)

    startdate = ClimateTools.buildtimetype(datebeg, T)
    enddate = ClimateTools.buildtimetype(dateend, T)

    # some checkups
    @argcheck startdate <= enddate
    # @argcheck startdate >= C[1][Axis{:time}][1]
    # @argcheck enddate <= C[1][Axis{:time}][end]

    dataOut = C[1][Axis{:time}(idxtimebeg:idxtimeend)]

    # The following control ensure that a 1-timestep temporal subset returns a 3D Array with time information on the timestep. i.e. startdate == enddate
    if ndims(dataOut) == 2
        timeV = startdate
        latsymbol = Symbol(C.dimension_dict["lat"])
        lonsymbol = Symbol(C.dimension_dict["lon"])
        data2 = fill(NaN, (size(dataOut,1), size(dataOut, 2), 1))
        data2[:,:,1] = dataOut
        dataOut = AxisArray(data2, Axis{lonsymbol}(C[1][Axis{lonsymbol}][:]), Axis{latsymbol}(C[1][Axis{latsymbol}][:]), Axis{:time}(C[1][Axis{:time}][:]))

    end

    return ClimGrid(dataOut, longrid=C.longrid, latgrid=C.latgrid, msk=C.msk, grid_mapping=C.grid_mapping, dimension_dict=C.dimension_dict, timeattrib=C.timeattrib, model=C.model, frequency=C.frequency, experiment=C.experiment, run=C.run, project=C.project, institute=C.institute, filename=C.filename, dataunits=C.dataunits, latunits=C.latunits, lonunits=C.lonunits, variable=C.variable, typeofvar=C.typeofvar, typeofcal=C.typeofcal, varattribs=C.varattribs, globalattribs=C.globalattribs)

end

"""
    resample(C::ClimGrid, startmonth::Int64, endmonth::Int64)

Return a resampled subset of ClimGrid C based on months.
"""
function resample(C::ClimGrid, startmonth::Int64, endmonth::Int64)

    @argcheck startmonth >= minimum(Dates.month.(C[1][Axis{:time}][:]))
    @argcheck startmonth <= maximum(Dates.month.(C[1][Axis{:time}][:]))

    if startmonth <= endmonth
        # Each matrix [:,:,i] represent data for a day
        datain = C.data.data
        # Date vector
        datevecin = C[1][Axis{:time}][:]
        # Where are the data between startmonth and endmonth
        index = (Dates.month.(datevecin) .<= endmonth) .&  (Dates.month.(datevecin) .>= startmonth)
        # Keep only data between startmonth and endmonth
        dataout = datain[:,:,index]
        datevecout = datevecin[index]
        # Create the ClimGrid output
        lonsymbol = Symbol(C.dimension_dict["lon"])
        latsymbol = Symbol(C.dimension_dict["lat"])
        axisout = AxisArray(dataout, Axis{lonsymbol}(C[1][Axis{lonsymbol}][:]), Axis{latsymbol}(C[1][Axis{latsymbol}][:]),    Axis{:time}(datevecout))

    elseif endmonth < startmonth
        # Each matrix [:,:,i] represent data for a day
        datain = C.data.data
        # Date vector
        datevecin = C[1][Axis{:time}][:]
        # Where are the data between startmonth and endmonth
        index = (Dates.month.(datevecin) .<= endmonth) .|  (Dates.month.(datevecin) .>= startmonth)
        # Keep only data between startmonth and endmonth
        dataout = datain[:,:,index]
        datevecout = datevecin[index]
        # Create the ClimGrid output
        lonsymbol = Symbol(C.dimension_dict["lon"])
        latsymbol = Symbol(C.dimension_dict["lat"])
        axisout = AxisArray(dataout, Axis{lonsymbol}(C[1][Axis{lonsymbol}][:]), Axis{latsymbol}(C[1][Axis{latsymbol}][:]), Axis{:time}(datevecout))
    end

    return ClimGrid(axisout, longrid=C.longrid, latgrid=C.latgrid, msk=C.msk, grid_mapping=C.grid_mapping,dimension_dict=C.dimension_dict, timeattrib=C.timeattrib, model=C.model, frequency=C.frequency, experiment=C.experiment, run=C.run, project=C.project,institute=C.institute, filename=C.filename, dataunits=C.dataunits, latunits=C.latunits, lonunits=C.lonunits, variable=C.variable,typeofvar=C.typeofvar, typeofcal=C.typeofcal, varattribs=C.varattribs, globalattribs=C.globalattribs)
end

"""
    resample(C::ClimGrid, season::String)

Return a resampled subset of ClimGrid C for a given season. Season options are: "DJF" (December-February), "MAM" (March-May), "JJA" (June-August), "SON" (September-November)
"""
function resample(C::ClimGrid, season::String)
    if lowercase(season) == "djf"
      D = resample(C, 12, 2)
    elseif lowercase(season) == "mam"
      D = resample(C, 3, 5)
    elseif lowercase(season) == "jja"
      D = resample(C, 6, 8)
    elseif lowercase(season) == "son"
      D = resample(C, 9, 11)
    else
      error("Wrong season name. Options are DJF (December-February; Winter), MAM (March-May; Spring), JJA (June-August; Summer) and SON (September-November; Fall)")
    end

    return D
end

"""
    periodmean(C::ClimGrid; startdate::Tuple, enddate::Tuple)

Mean of array data over a given period.
"""
function periodmean(C::ClimGrid; start_date::Tuple=(Inf, ), end_date::Tuple=(Inf,), level=1)

    if start_date != (Inf,) || end_date != (Inf,)
        C = temporalsubset(C, start_date, end_date)
    end

    datain = C.data.data

    # Mean and squeeze
    if ndims(datain) == 2
        dataout = datain
    elseif ndims(datain) == 3
        if size(datain, 3) == 1 # already an average on single value
            dataout = dropdims(datain, dims=3)
        else
            dataout = dropdims(Images.meanfinite(datain, 3), dims=3)
        end
    elseif ndims(datain) == 4
        if size(datain, 4) == 1
            dataout = dropdims(datain[:, :, level, :], dims = 3)
        else
            dataout = dropdims(Images.meanfinite(datain, 3), dims=3)
        end
    end

    # Build output AxisArray
    FD = buildarray_climato(C, dataout)

    # Return climGrid type containing the indice
    return ClimGrid(FD, longrid=C.longrid, latgrid=C.latgrid, msk=C.msk, grid_mapping=C.grid_mapping, dimension_dict=C.dimension_dict, timeattrib=C.timeattrib, model=C.model, frequency=C.frequency, experiment=C.experiment, run=C.run, project=C.project, institute=C.institute, filename=C.filename, dataunits=C.dataunits, latunits=C.latunits, lonunits=C.lonunits, variable="periodmean", typeofvar=C.typeofvar, typeofcal="climatology", varattribs=C.varattribs, globalattribs=C.globalattribs)
end

"""
    timeindex(timeVec, start_date, end_date, T)

Return the index of time vector specified by start_date and end_date. T is the DateTime type (see NCDatasets.jl documentation).
"""
function timeindex(timeV, datebeg::Tuple, dateend::Tuple, T)

    # Start Date
    if !isinf(datebeg[1])
        # Build DateTime type
        start_date = ClimateTools.buildtimetype(datebeg, T)
        # @argcheck start_date >= timeV[1]
        idxtimebeg = findfirst(timeV .>= start_date)[1]
    else
        idxtimebeg = 1
    end
    # End date
    if !isinf(dateend[1])
        # Build DateTime type
        end_date = ClimateTools.buildtimetype(dateend, T)
        # @argcheck end_date <= timeV[end]
        idxtimeend = findlast(timeV .<= end_date)[1]
    else
        idxtimeend = length(timeV)
    end

    if !isinf(datebeg[1]) && !isinf(dateend[1])
        @argcheck start_date <= end_date
    end
    return idxtimebeg, idxtimeend
end

# function timeindex(C::ClimGrid, datebeg::Tuple, dateend::Tuple)
#
#     T = typeof(get_timevec(C)[1])
#
#     # Start Date
#     if !isinf(datebeg[1])
#
#         # Build DateTime type
#         start_date = ClimateTools.buildtimetype(datebeg, T)
#
#         # Check
#         @argcheck start_date >= timeV[1]
#
#         idxtimebeg = findfirst(timeV .>= start_date)[1]
#     else
#         idxtimebeg = 1
#     end
#
#     # End date
#     if !isinf(dateend[1])
#
#         # Build DateTime type
#         end_date = ClimateTools.buildtimetype(dateend, T)
#
#         @argcheck end_date <= timeV[end]
#
#         idxtimeend = findlast(timeV .<= end_date)[1]
#     else
#         idxtimeend = length(timeV)
#     end
#
#     if !isinf(datebeg[1]) && !isinf(dateend[1])
#         @argcheck start_date <= end_date
#     end
#
#     return idxtimebeg, idxtimeend
# end

"""
    buildtimetype(datetuple, f)

Returns the adequate DateTime for temporal subsetting using DateType *f*
"""
function buildtimetype(date_tuple, f)

    if length(date_tuple) == 1
        dateout = f(date_tuple[1], 01, 01)
    elseif length(date_tuple) == 2
        dateout = f(date_tuple[1], date_tuple[2], 01)
    elseif length(date_tuple) == 3
        dateout = f(date_tuple[1], date_tuple[2], date_tuple[3])
    elseif length(date_tuple) == 4
        dateout = f(date_tuple[1], date_tuple[2], date_tuple[3], date_tuple[4], 00, 00)
    elseif length(date_tuple) == 5
        dateout = f(date_tuple[1], date_tuple[2], date_tuple[3], date_tuple[4], date_tuple[5], 00)
    elseif length(date_tuple) == 6
        dateout = f(date_tuple[1], date_tuple[2], date_tuple[3], date_tuple[4], date_tuple[5], date_tuple[6])
    end

    return dateout
end

"""
    daymean(C::ClimGrid)

Returns the daily average given a sub-daily ClimGrid.
"""
function daymean(C::ClimGrid)

    datain = C[1].data
    timevec   = get_timevec(C)
    T = typeof(timevec[1])

    nbdays = unique(yearmonthday.(timevec))
    nbdays_len = length(nbdays)
    dataout = Array{typeof(datain[1])}(undef, (size(C[1], 1), size(C[1], 2), nbdays_len))
    newtime = Array{T}(undef, nbdays_len)

    # loop over year-month-days
    Threads.@threads for iday in 1:nbdays_len

        daytmp = nbdays[iday]
        datefind = T(daytmp[1], daytmp[2], daytmp[3])
        idx = findall(x -> Dates.year(x) == Dates.year(datefind) && Dates.month(x) == Dates.month(datefind) && Dates.day(x) == Dates.day(datefind), timevec)

        dataout[:, :, iday] = Statistics.mean(datain[:, :, idx], dims=3)
        newtime[iday] = datefind

    end

    # Build output AxisArray
    FD = buildarray_resample(C, dataout, newtime)

    return ClimGrid(FD, longrid=C.longrid, latgrid=C.latgrid, msk=C.msk, grid_mapping=C.grid_mapping, dimension_dict=C.dimension_dict, timeattrib=C.timeattrib, model=C.model, frequency="day", experiment=C.experiment, run=C.run, project=C.project, institute=C.institute, filename=C.filename, dataunits=C.dataunits, latunits=C.latunits, lonunits=C.lonunits, variable=C.variable, typeofvar=C.typeofvar, typeofcal=C.typeofcal, varattribs=C.varattribs, globalattribs=C.globalattribs)

end

"""
    daysum(C::ClimGrid)

Returns the daily sum given a sub-daily ClimGrid C.
"""
function daysum(C::ClimGrid)

    datain = C[1].data
    timevec   = get_timevec(C)
    T = typeof(timevec[1])

    nbdays = unique(yearmonthday.(timevec))
    nbdays_len = length(nbdays)
    dataout = Array{typeof(datain[1])}(undef, (size(C[1], 1), size(C[1], 2), nbdays_len))
    newtime = Array{T}(undef, nbdays_len)

    # loop over year-month-days
    Threads.@threads for iday in 1:nbdays_len

        daytmp = nbdays[iday]
        datefind = T(daytmp[1], daytmp[2], daytmp[3])
        idx = findall(x -> Dates.year(x) == Dates.year(datefind) && Dates.month(x) == Dates.month(datefind) && Dates.day(x) == Dates.day(datefind), timevec)

        dataout[:, :, iday] = Statistics.sum(datain[:, :, idx], dims=3)
        newtime[iday] = datefind

    end

    # Build output AxisArray
    FD = buildarray_resample(C, dataout, newtime)

    return ClimGrid(FD, longrid=C.longrid, latgrid=C.latgrid, msk=C.msk, grid_mapping=C.grid_mapping, dimension_dict=C.dimension_dict, timeattrib=C.timeattrib, model=C.model, frequency="day", experiment=C.experiment, run=C.run, project=C.project, institute=C.institute, filename=C.filename, dataunits=C.dataunits, latunits=C.latunits, lonunits=C.lonunits, variable=C.variable, typeofvar=C.typeofvar, typeofcal=C.typeofcal, varattribs=C.varattribs, globalattribs=C.globalattribs)

end

"""
    monthmean(C::ClimGrid)

Returns monthly means of ClimGrid C.
"""
function monthmean(C::ClimGrid)

    datain = C[1].data
    timevec = get_timevec(C)
    T = typeof(timevec[1])

    nbmonth = unique(yearmonth.(timevec))
    nbmonth_len = length(nbmonth)
    dataout = Array{typeof(datain[1])}(undef, (size(C[1], 1), size(C[1], 2), nbmonth_len))
    newtime = Array{T}(undef, nbmonth_len)

    # loop over year-month-days
    Threads.@threads for imonth in 1:nbmonth_len

        monthtmp = nbmonth[imonth]
        datefind = T(monthtmp[1], monthtmp[2])
        idx = findall(x -> Dates.year(x) == Dates.year(datefind) && Dates.month(x) == Dates.month(datefind), timevec)

        dataout[:, :, imonth] = Statistics.mean(datain[:, :, idx], dims=3)
        newtime[imonth] = datefind

    end

    # Build output AxisArray
    FD = buildarray_resample(C, dataout, newtime)

    return ClimGrid(FD, longrid=C.longrid, latgrid=C.latgrid, msk=C.msk, grid_mapping=C.grid_mapping, dimension_dict=C.dimension_dict, timeattrib=C.timeattrib, model=C.model, frequency="day", experiment=C.experiment, run=C.run, project=C.project, institute=C.institute, filename=C.filename, dataunits=C.dataunits, latunits=C.latunits, lonunits=C.lonunits, variable=C.variable, typeofvar=C.typeofvar, typeofcal=C.typeofcal, varattribs=C.varattribs, globalattribs=C.globalattribs)

end

"""
    monthsum(C::ClimGrid)

Returns monthly sums of ClimGrid C.
"""
function monthsum(C::ClimGrid)

    datain = C[1].data
    timevec = get_timevec(C)
    T = typeof(timevec[1])

    nbmonth = unique(yearmonth.(timevec))
    nbmonth_len = length(nbmonth)
    dataout = Array{typeof(datain[1])}(undef, (size(C[1], 1), size(C[1], 2), nbmonth_len))
    newtime = Array{T}(undef, nbmonth_len)

    # loop over year-month-days
    Threads.@threads for imonth in 1:nbmonth_len

        monthtmp = nbmonth[imonth]
        datefind = T(monthtmp[1], monthtmp[2])
        idx = findall(x -> Dates.year(x) == Dates.year(datefind) && Dates.month(x) == Dates.month(datefind), timevec)

        dataout[:, :, imonth] = Statistics.sum(datain[:, :, idx], dims=3)
        newtime[imonth] = datefind

    end

    # Build output AxisArray
    FD = buildarray_resample(C, dataout, newtime)

    return ClimGrid(FD, longrid=C.longrid, latgrid=C.latgrid, msk=C.msk, grid_mapping=C.grid_mapping, dimension_dict=C.dimension_dict, timeattrib=C.timeattrib, model=C.model, frequency="day", experiment=C.experiment, run=C.run, project=C.project, institute=C.institute, filename=C.filename, dataunits=C.dataunits, latunits=C.latunits, lonunits=C.lonunits, variable=C.variable, typeofvar=C.typeofvar, typeofcal=C.typeofcal, varattribs=C.varattribs, globalattribs=C.globalattribs)

end

"""
    timeresolution(timevec::Array{N,1} where N)

Return the time resolution of the vector timevec.

"""
function timeresolution(timevec::Array{N,1} where N)

    # timevec = (NetCDF.ncread(str, "time"))
    if length(timevec) > 1
        timediff = diff(timevec)[2]
        if timediff == 1. || timediff == 1
            return "24h"
        elseif round(timediff, digits=5) == round(12/24, digits=5)
            return "12h"
        elseif round(timediff, digits=5) == round(6/24, digits=5)
            return "6h"
        elseif round(timediff, digits=5) == round(3/24, digits=5)
            return "3h"
        elseif round(timediff, digits=5) == round(1/24, digits=5)
            return "1h"
        elseif round(timediff, digits=5) == 365.0 || round(timediff, digits=5) == 366.0 || round(timediff, digits=5) == 360.0
            return "Yearly"
        elseif timediff == 31.0 || timediff == 30.0 || timediff == 28.0 || timediff == 29.0
            return "Monthly"
        end
    else
        return "N/A"
    end
end

"""
    function pr_timefactor(rez::String)

Return the time factor that should be applied to precipitation to get accumulation for resolution "rez"

"""
function pr_timefactor(rez::String)

    if rez == "24h"
        return 86400.0
    elseif rez == "12h"
        return 43200.0
    elseif rez == "6h"
        return 21600.0
    elseif rez == "3h"
        return 10800.0
    elseif rez == "1h"
        return 3600.0
    elseif rez == "N/A"
        return 1.0
    end

end

"""
    function daymean_factor(rez::String)

Return the time factor that should be applied to precipitation to get accumulation for resolution "rez"

"""
function daymean_factor(rez::String)

    if rez == "24h" || rez == "day" || rez == "daily"
        return 1
    elseif rez == "12h"
        return 2
    elseif rez == "6h"
        return 4
    elseif rez == "3h"
        return 8
    elseif rez == "1h"
        return 24
    elseif rez == "N/A"
        return 1
    end

end

"""
    dropfeb29(C::ClimGrid)

Removes february 29th. Needed for bias correction.
"""
function dropfeb29(C::ClimGrid)
    date_vec = get_timevec(C) # [1][Axis{:time}][:]
    f = typeof(date_vec[1])
    feb29th = (Dates.month.(date_vec) .== Dates.month(f(2000, 2, 2))) .& (Dates.day.(date_vec) .== Dates.day(29))
    dataout = C[1].data[:, :, .!feb29th] # drop for data
    date_vec = date_vec[.!feb29th] # drop for calendar

    datevec_noleap = Array{DateTimeNoLeap}(undef, length(date_vec))

    for id = 1:size(date_vec, 1)
        datevec_noleap[id] = reinterpret(DateTimeNoLeap, date_vec[id])
    end

    dataout = buildarray_resample(C, dataout, datevec_noleap)

    return ClimGrid(dataout; longrid=C.longrid, latgrid=C.latgrid, msk=C.msk, grid_mapping=C.grid_mapping, dimension_dict=C.dimension_dict, timeattrib=C.timeattrib, model=C.model, frequency=C.frequency, experiment=C.experiment, run=C.run, project=C.project, institute=C.institute, filename=C.filename, dataunits=C.dataunits, latunits=C.latunits, lonunits=C.lonunits, variable=C.variable, typeofvar=C.typeofvar, typeofcal=C.typeofcal, varattribs=C.varattribs, globalattribs=C.globalattribs)
end

"""
    leapyears(datevec)

Returns the leap years contained into datevec.
"""
function leapyears(datevec)

    years = unique(Dates.year.(datevec))
    lyrs = years[Dates.isleapyear.(years)]

    return lyrs

end

function find_julianday_idx(julnb, ijulian, window)
    if ijulian <= window
        idx = @. (julnb >= 1) & (julnb <= (ijulian + window)) | (julnb >= (365 - (window - ijulian))) & (julnb <= 365)
    elseif ijulian > 365 - window
        idx = @. (julnb >= 1) & (julnb <= ((window-(365-ijulian)))) | (julnb >= (ijulian-window)) & (julnb <= 365)
    else
        idx = @. (julnb <= (ijulian + window)) & (julnb >= (ijulian - window))
    end
    return idx
end

"""
    yearmonthdayhour(dt::AbstractCFDateTime) -> (Int64, Int64, Int64, Int64)
Simultaneously return the year, month, day and hour parts of `dt`.
Author: Alexander-Barth (Github)
"""
yearmonthdayhour(dt::DT) where DT <: Dates.TimeType = (Dates.year(dt),Dates.month(dt), Dates.day(dt), Dates.hour(dt))


# This files contains all the gradient definitions for the game.

using Zygote
using Zygote: @adjoint, @nograd, literal_getproperty, @showgrad, hook
import Zygote: accum
import Base: sign, floor

@nograd get_obs, reset_ball!#wall_collision!, #, paddle_collision!

@adjoint Vector_(x, y) = Vector_(x, y), ∆ -> (∆.x, ∆.y)

@adjoint literal_getproperty(v::Vector_, ::Val{:x}) =
    getproperty(v, :x), ∆ -> (Vector_(∆, 0f0), )

@adjoint literal_getproperty(v::Vector_, ::Val{:y}) =
    getproperty(v, :y), ∆ -> (Vector_(0f0, ∆), )

@adjoint Player(length, position, velocity, score) =
    Player(length, position, velocity, score), ∆ -> (∆.length, ∆.position, ∆.velocity, ∆.score)

@adjoint literal_getproperty(p::Player, ::Val{:length}) =
    getproperty(p, :length), ∆ -> (Player(∆, zero(p.position), 0, 0), )

@adjoint literal_getproperty(p::Player, ::Val{:position}) =
    getproperty(p, :position), ∆ -> (Player(0f0, ∆, 0, 0), )

@adjoint literal_getproperty(p::Player, ::Val{:velocity}) =
    getproperty(p, :velocity), ∆ -> (Player(0f0, zero(p.position), ∆, 0), )

@adjoint literal_getproperty(p::Player, ::Val{:score}) =
    getproperty(p, :score), ∆ -> (Player(0f0, zero(p.position), 0, ∆), )

@adjoint Ball(position, velocity, radius) =
    Ball(position, velocity, radius), ∆ -> (∆.position, ∆.velocity, ∆.radius)

@adjoint literal_getproperty(b::Ball, ::Val{:position}) =
    getproperty(b, :position), ∆ -> (Ball(∆, zero(b.velocity), 0f0), )

@adjoint literal_getproperty(b::Ball, ::Val{:velocity}) =
    getproperty(b, :velocity), ∆ -> (Ball(zero(b.position), ∆, 0f0), )

@adjoint literal_getproperty(b::Ball, ::Val{:radius}) =
    getproperty(b, :radius), ∆ -> (Ball(zero(b.position), zero(b.velocity), ∆), )

@adjoint Arena(dims, margin) = Arena(dims, margin), ∆ -> (∆.dims, ∆.margin)

@adjoint literal_getproperty(a::Arena, ::Val{:dims}) =
    getproperty(a, :dims), ∆ -> (Arena(∆, zero(a.margin)), )

@adjoint literal_getproperty(a::Arena, ::Val{:margin}) =
    getproperty(a, :margin), ∆ -> (Arena(zero(a.dims), ∆), )

@adjoint Env(arena, player_a, player_b, ball) =
    Env(arena, player_a, player_b, ball), ∆ -> (∆.arena, ∆.player_a, ∆.player_b, ∆.ball)

@adjoint literal_getproperty(e::Env, ::Val{:arena}) =
    getproperty(e, :arena), ∆ -> (Env(∆, zero(e.player_a), zero(e.player_b), zero(e.ball)), )

@adjoint literal_getproperty(e::Env, ::Val{:player_a}) =
    getproperty(e, :player_a), ∆ -> (Env(zero(e.arena), ∆, zero(e.player_b), zero(e.ball)), )

@adjoint literal_getproperty(e::Env, ::Val{:player_b}) =
    getproperty(e, :player_b), ∆ -> (Env(zero(e.arena), zero(e.player_a), ∆, zero(e.ball)), )

@adjoint literal_getproperty(e::Env, ::Val{:ball}) =
    getproperty(e, :ball), ∆ -> (Env(zero(e.arena), zero(e.player_a), zero(e.player_b), ∆), )

x::Player + y::Player = Player(x.length, x.position + y.position, x.velocity + y.velocity, 0)
x::Ball + y::Ball = Ball(x.position + y.position, x.velocity + y.velocity, x.radius)
x::Arena + y::Arena = Arena(x.dims + y.dims, x.margin + y.margin)
x::Env + y::Env = Env(x.arena + y.arena, x.player_a + y.player_a, x.player_b + y.player_b, x.ball + y.ball)

#=
function Zygote.accum(x::Player, y::Player)
    ans = Player(x.length, x.position + y.position, x.velocity, 0)
    println("$x   +    $y    =  $ans\n\n\n\n")
    return ans
end
function Zygote.accum(x::Ball, y::Ball)
    ans = Ball(x.position + y.position, x.velocity + y.velocity, x.radius)
    println("$x   +    $y    =  $ans\n\n\n\n")
    return ans
end

function Zygote.accum(x::Env, y::Env)
    ball = Ball(x.ball.position + y.ball.position, x.ball.velocity + y.ball.velocity,
                x.ball.radius)
    player_a = Player(x.player_a.length, x.player_a.position + y.player_a.position, x.player_a.velocity, 0)
    player_b = Player(x.player_b.length, x.player_b.position + y.player_b.position, x.player_b.velocity, 0)

    arena = Arena(x.arena.dims + y.arena.dims, x.arena.margin + y.arena.margin)

    ans = Env(arena, player_a, player_b, ball)
    println("$x   +    $y    =  $ans\n\n\n\n")
    return ans
end
Zygote.accum(x::Arena, y::Arena) = println("x = $x | y = $y\n\n\n\n")
=#

#=
@adjoint update_paddle_position(player::Player, arena::Arena, dir) =
    update_paddle_position(player, arena, dir), ∆ -> begin
        dummy = player.position.y + dir * player.velocity  # dummy holds the new position of the paddle
        # If the new position is within the acceptable range of motion, then velocity should be the gradient, wrt dir
        # Otherwise the gradient must be zero.
        0 ≤ dummy ≤ arena.dims.y - player.length ?
            (zero(player), zero(arena), ∆*Vector_(0f0, player.velocity)) :
            (zero(player), zero(arena), ∆*zero(dir))
        end

@adjoint function reset_ball!(arena::Arena, ball::Ball)
    reset_ball!(arena, ball), ∆ -> begin
        return (zero(arena), zero(ball))
    end
end

@adjoint function paddle_collision!(new_position::Vector_, player_a_pos::Vector_, player_b_pos::Vector_,
                                    paddle_length::AbstractFloat, ball::Ball, arena::Arena)
   paddle_collision(new_position, player_a_pos, player_b_pos, paddle_length, ball, arena), ∆ -> begin
    if new_position.x > arena.dims.x - ball.radius
        lower_bound_b  = max(player_b_pos.y - 3ball.radius/4f0, 0)
        higher_bound_b = min(player_b_pos.y + paddle_length + 3ball.radius/4f0, arena.dims.y)

        if lower_bound_b < new_position.y < higher_bound_b

        end
    else
        return (zero(new_position), zero(player_a_pos), zero(player_b_pos), zero(paddle_length),
                zero(ball), zero(arena))
    end
end
end

@nograd wall_collision!, get_obs#, paddle_collision!

Zygote.accum(x::Arena, y::Arena) = println("x = $x | y = $y")
Zygote.accum(x::Player, y::Player) = println("x = $x | y = $y")
=#


struct UniqueValues{S, T1<:AbstractArray, T2<:AbstractArray{S}}
    values::T1
    unique::T2
    value2index::Dict{S, Int}
end

function UniqueValues(col, s = unique(sort(col)))
    value2index = Dict(zip(s, 1:length(s)))
    UniqueValues(col, s, value2index)
end

function (cs::UniqueValues)(s, el)
    scale = to_scale(s)
    @assert scale !== nothing
    @assert typeof(scale) !== typeof(s)
    cs(scale, el)
end

(cs::UniqueValues)(scale::Function, el) = scale(el)

function (cs::UniqueValues)(scale::AbstractArray, el)
    scale[(cs.value2index[el] - 1) % length(scale) + 1]
end

struct Group
    columns::NamedTuple
    f::FunctionOrAnalysis
end

Group(c::NamedTuple) = Group(c, tuple)

Group(v, f::FunctionOrAnalysis = tuple) = Group((color = v,), f)
Group(f::FunctionOrAnalysis = tuple; kwargs...) = Group(values(kwargs), f)

struct ByColumn end
const bycolumn = ByColumn()

Base.isless(::ByColumn, ::ByColumn) = false

combine_gog(f1, f2) = (args...) -> f1(to_tuple(f2(args...))...)
combine_gog(f1, f2::typeof(tuple)) = f1
combine_gog(f1::typeof(tuple), f2) = f2
combine_gog(f1::typeof(tuple), f2::typeof(tuple)) = tuple

Base.merge(g1::Group, g2::Group) = Group(merge(g1.columns, g2.columns), combine_gog(g1.f, g2.f))
Base.merge(f::FunctionOrAnalysis, g::Group) = merge(Group(f), g)
Base.merge(g::Group, f::FunctionOrAnalysis) = merge(g, Group(f))

Base.:*(g1::Group, g2::Group) = merge(g1, g2)

width(v::Tuple) = length(v)
width(v::NamedTuple) = maximum(width, v)
width(v::AbstractVector) = 1
width(v::AbstractArray) = mapreduce(length, *, axes(v)[2:end])

column_length(v::Union{Tuple, NamedTuple}) = column_length(v[1])
column_length(v::AbstractVector) = length(v)
column_length(v::AbstractArray) = length(axes(v)[1])

extract_view(v::Union{Tuple, NamedTuple}, idxs) = map(x -> extract_view(x, idxs), v)
extract_view(v::AbstractVector, idxs) = view(v, idxs)
extract_view(v::AbstractArray, idxs) = view(v, idxs, axes(v)[2:end]...)

extract_view(v::Tuple, idxs, n) = extract_view(v[n], idxs)
extract_view(v::AbstractVector, idxs, n) = view(v, idxs)
function extract_view(v::AbstractArray, idxs, n)
    ax = axes(v)[2:end]
    c = CartesianIndices(ax)[n]
    view(v, idxs, Tuple(c)...)
end

extract_view(v::NamedTuple, idxs, n) = map(t -> extract_view(t, idxs, n), v)

concatenate(v::Union{Tuple, NamedTuple}...) = map(concatenate, v...)
concatenate(v::AbstractArray...) = vcat(v...)


"""
    mutable struct OuterControl{
        A <: ActivePowerControl,
        R <: ReactivePowerControl
    } <: DynamicInverterComponent
        active_power::A
        reactive_power::R
        ext::Dict{String, Any}
        states::Vector{Symbol}
        n_states::Int
    end
Parameters of a Outer-Loop controller using a active power controller and a reactive power droop controller.

# Arguments
- `A <: ActivePowerControl`: Active power controller (typically droop or virtual inertia).
- `R <: ReactivePowerControl`: Reactive power controller (typically droop).
- `ext::Dict{String, Any}`
- `states::Vector{Symbol}`: Vector of states (will depend on the components).
- `n_states::Int`: Number of states (will depend on the components).
"""
mutable struct OuterControl{A <: ActivePowerControl, R <: ReactivePowerControl} <:
               DynamicInverterComponent
    active_power::A
    reactive_power::R
    ext::Dict{String, Any}
    states::Vector{Symbol}
    n_states::Int
end

function OuterControl(
    active_power::A,
    reactive_power::R,
    ext = Dict{String, Any}(),
) where {A <: ActivePowerControl, R <: ReactivePowerControl}
    OuterControl(
        active_power,
        reactive_power,
        ext,
        vcat(active_power.states, reactive_power.states),
        active_power.n_states + reactive_power.n_states,
    )
end

function OuterControl(;
    active_power,
    reactive_power,
    ext = Dict{String, Any}(),
    states = nothing,
    n_states = nothing,
)
    if states === nothing
        @assert n_states === nothing
        return OuterControl(active_power, reactive_power, ext)
    end
    @assert n_states !== nothing
    return OuterControl(active_power, reactive_power, ext, states, n_states)
end

get_active_power(value::OuterControl) = value.active_power
get_reactive_power(value::OuterControl) = value.reactive_power
get_ext(value::OuterControl) = value.ext
get_states(value::OuterControl) = value.states
get_n_states(value::OuterControl) = value.n_states


# ---
# jupyter:
#   jupytext:
#     formats: jl:light,ipynb
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.5.1
#   kernelspec:
#     display_name: Julia 1.4.2
#     language: julia
#     name: julia-1.4
# ---

# + [markdown] run_control={"frozen": false, "read_only": false}
# # Molecule Objects in  Psi4
#
# This tutorial provides an overview on creating and manipulating molecule objects in  Psi4, illustrated with an example parameterization of the Lennard-Jones potential for Helium dimer.
#
# In order to use  Psi4  within a Julia environment, we may install PyCall
# executing `using Pkg; Pkg.add("PyCall")`. We point PyCall to the Python installation that has  Psi4  installed, that is `ENV["PYTHON"] = ~/miniconda3/envs/p4env/bin/python`, rebuild PyCall, `using Pkg; Pkg.build("PyCall")` and restart Julia. The next time we start Julia we may import  Psi4  with pyimport:
#
# (check [here](https://github.com/dgasmith/psi4numpy/blob/master/README.md) for more details)

# + run_control={"frozen": false, "read_only": false}
using PyCall: pyimport
psi4 = pyimport("psi4")

# + [markdown] run_control={"frozen": false, "read_only": false}
# Unlike in  Psi4  input files, defining a molecule in Julia is done by passing the molecular coordinates as a triple-quoted string to the [`psi4.geometry()`](http://psicode.org/psi4manual/master/api/psi4.driver.geometry.html#psi4.driver.geometry "API Details") function:

# + run_control={"frozen": false, "read_only": false}
he = psi4.geometry("""
He
""")

# + [markdown] run_control={"frozen": false, "read_only": false}
# Here, not only does the variable `he` refer to the helium molecule, but also an instance of the [`psi4.core.Molecule`](http://psicode.org/psi4manual/master/psi4api.html#psi4.core.Molecule "Go to API")
# class in  Psi4; this will be discussed in more detail later.  For a more
# complicated system than an isolated atom, the coordinates can be given in Cartesian or Z-Matrix formats:

# + run_control={"frozen": false, "read_only": false}
h2o = psi4.geometry("""
O
H 1 0.96
H 1 0.96 2 104.5
""")

# + [markdown] run_control={"frozen": false, "read_only": false}
# Information like the molecular charge, multiplicity, and units are assumed to be 0, 1, and Angstroms, respectively, if not specified within the molecule definition.  This is done by adding one or more [molecule keywords](http://psicode.org/psi4manual/master/psithonmol.html#molecule-keywords "Go to Documentation") to the geometry string used to define the molecule.  Additionally,  Psi4  can detect molecular symmetry, or it can be specified manually.  For example, to define a doublet water cation in $C_{2V}$ symmetry using lengths in Bohr,

# + run_control={"frozen": false, "read_only": false}
doublet_h2o_cation = psi4.geometry("""
1 2
O
H 1 1.814
H 1 1.814 2 104.5

units bohr
symmetry c2v
""")

# + [markdown] run_control={"frozen": false, "read_only": false}
# where the line `1 2` defines the charge and multiplicity, respectively.  For systems of non-bonded fragments, the coordinates of each fragment are separated by a double-hyphen `"--"`; this allows for one fragment to be defined with Cartesian and another to be defined with Z-Matrix. For example, the hydronium-benzene complex can be defined with:

# + run_control={"frozen": false, "read_only": false}
hydronium_benzene = psi4.geometry("""
0 1
C          0.710500000000    -0.794637665924    -1.230622098778
C          1.421000000000    -0.794637665924     0.000000000000
C          0.710500000000    -0.794637665924     1.230622098778
C         -0.710500000000    -0.794637665924     1.230622098778
H          1.254500000000    -0.794637665924    -2.172857738095
H         -1.254500000000    -0.794637665924     2.172857738095
C         -0.710500000000    -0.794637665924    -1.230622098778
C         -1.421000000000    -0.794637665924     0.000000000000
H          2.509000000000    -0.794637665924     0.000000000000
H          1.254500000000    -0.794637665924     2.172857738095
H         -1.254500000000    -0.794637665924    -2.172857738095
H         -2.509000000000    -0.794637665924     0.000000000000
--
1 1
X  1  CC  3  30   2  A2
O  13 R   1  90   2  90
H  14 OH  13 TDA  1  0
H  14 OH  15 TDA  13 A1
H  14 OH  15 TDA  13 -A1

CC    = 1.421
CH    = 1.088
A1    = 120.0
A2    = 180.0
OH    = 1.05
R     = 4.0
units angstrom
""")

# + [markdown] run_control={"frozen": false, "read_only": false}
# For non-bonded fragments, the charge and multiplicity should be given explicitly for each fragment.  If not, the
# charge and multiplicity given (or inferred) for the first fragment is assumed to be the same for all fragments.  In
# addition to defining the coordinates outright, we have used variables within the geometry specification strings to
# define bond lengths, angles, and dihedrals in the molecule.  Similarly, we could define the X, Y, or Z Cartesian
# coordinate for any atom in our molecule.
#
# In order to define these variables after the molecule is built, as opposed to within the geometry specification
# itself, there are several ways to do so; one of which will be illustrated in the Lennard-Jones potential example
# below.
#
# When a Psi4 molecule is first built using ``psi4.geometry()``, it is in an unfinished state, as a user may wish to
# tweak the molecule. This can be solved by calling [``psi4.Molecule.update_geometry()``](http://psicode.org/psi4manual/master/psi4api.html#psi4.core.Molecule.update_geometry "Go to API"). This will update the molecule and restore sanity
# to chemistry.

# +
h2cch2 = psi4.geometry("""
H
C 1 HC
H 2 HC 1 A1
C 2 CC 3 A1 1 D1
H 4 HC 2 A1 1 D1
H 4 HC 2 A1 1 D2

HC = 1.08
CC = 1.4
A1 = 120.0
D1 = 180.0
D2 = 0.0
""")

println("Ethene has $(h2cch2.natom()) atoms")

h2cch2.update_geometry()
println("Ethene has $(h2cch2.natom()) atoms")
# -

# Finally, one can obtain useful information from a molecule by invoking one of several [`psi4.core.Molecule`](http://psicode.org/psi4manual/master/psi4api.html#psi4.core.Molecule "Go to Documentation") class methods on the molecule of interest.  For example, if we were interested in verifying that our doublet water cation from above is, in fact, a doublet, we could invoke
# ~~~python
# doublet_h2o_cation.multiplicity()
# ~~~
# Below, some useful Molecule class methods are tabulated; please refer to the documentation for more details.
#
# | Method | Description |
# |--------|-------------|
# | [center_of_mass()](http://psicode.org/psi4manual/master/psi4api.html#psi4.core.Molecule.center_of_mass "Go to Documentation") | Computes center of mass of molecule |
# | [molecular_charge()](http://psicode.org/psi4manual/master/psi4api.html#psi4.core.Molecule.molecular_charge "Go to Documentation") | Gets the molecular charge |
# | [multiplicity()](http://psicode.org/psi4manual/master/psi4api.html#psi4.core.Molecule.multiplicity "Go to Documentation") | Gets the total multiplicity |
# | [nuclear_repulsion_energy()](http://psicode.org/psi4manual/master/psi4api.html#psi4.core.Molecule.nuclear_repulsion_energy "Go to Documentation") | Computes the nuclear repulsion energy of the molecule |

# + [markdown] run_control={"frozen": false, "read_only": false}
# ## Example: Fitting Lennard-Jones Parameters from Potential Energy Scan
#
# In this example, we will compute and fit a potential energy curve for the Helium dimer.  To begin with, let's create a string representation for our He dimer in Z-Matrix format, with the variable `**R**` representing the distance between the He atoms. The stars surrounding the ``R`` aren't any special syntax, just a convenient marker for future substitution.

# + run_control={"frozen": false, "read_only": false}
# Define He Dimer
he_dimer = """
He
--
He 1 **R**
"""

# + [markdown] run_control={"frozen": false, "read_only": false}
# Now we can build a series of dimers with the He atoms at different separations, and compute the energy at each point:

# + run_control={"frozen": false, "read_only": false}
distances = [2.875, 3.0, 3.125, 3.25, 3.375, 3.5, 3.75, 4.0, 4.5, 5.0, 6.0, 7.0]
energies = []
for d in distances
    # Build a new molecule at each separation
    mol = psi4.geometry(replace(he_dimer,"**R**" => string(d)))

    # Compute the Counterpoise-Corrected interaction energy
    en = psi4.energy("MP2/aug-cc-pVDZ", molecule=mol, bsse_type="cp")

    # Place in a reasonable unit, Wavenumbers in this case
    en *= 219474.6

    # Append the value to our list
    push!(energies,en)
end

println("Finished computing the potential!")

# + [markdown] run_control={"frozen": false, "read_only": false}
# Next, we can use the [NumPy](http://www.numpy.org/) library to fit a curve to these points along the potential scan.  In this case, we will fit a Lennard-Jones potential.

# + run_control={"frozen": false, "read_only": false}
np   = pyimport("numpy")

# + run_control={"frozen": false, "read_only": false}
# Fit data in least-squares way to a -12, -6 polynomial
powers = [-12, -6]
x = np.power(np.reshape(np.array(distances),(-1, 1)), powers)
coeffs = np.linalg.lstsq(x, energies)[1]

# Build list of points
fpoints = np.reshape(np.linspace(2, 7, 50),(-1, 1))
fdata = np.power(fpoints, powers)

fit_energies = np.dot(fdata, coeffs)

# + [markdown] run_control={"frozen": false, "read_only": false}
# To visualize our results, we can use the [Matplotlib](http://matplotlib.org/) library.  Since we're working in a Jupyter notebook, we can also use the `%matplotlib inline` "magic" command so that the graphs will show up in the notebook itself (check [here](https://ipython.org/ipython-doc/3/interactive/magics.html) for a comprehensive list of magic commands).

# + run_control={"frozen": false, "read_only": false}
# %matplotlib inline
import PyPlot
plt = PyPlot

# + run_control={"frozen": false, "read_only": false}
plt.xlim((2, 7))  # X limits
plt.ylim((-7, 2))  # Y limits
plt.scatter(distances, energies)  # Scatter plot of the distances/energies
plt.plot(fpoints, fit_energies)  # Fit data
plt.plot([0,10], [0,0], "k-")  # Make a line at 0
# -


@testset "877.stone-game.jl" begin
    @test stone_game([5, 3, 4, 5]) == true
end


module DataReaders
using FileIO
using CSVFiles
using FeatherFiles
using Feather
using DataFrames
using HDF5
using JLD
using Parquet
using Dates

export DataReader, fit!, transform!
export Transformer
export datareaderrun

using TSML
using TSML.TSMLTypes
import TSML.TSMLTypes.fit! # to overload
import TSML.TSMLTypes.transform! # to overload
using TSML.Utils

const FILEFMT = Dict("csv"=>Val(:csv),"feather"=>Val(:feather),"hdf5"=>Val(:hdf5),
          "h5"=>Val(:hdf5),"parquet"=>Val(:parquet),"jld"=>Val(:jld))

const DATEFMT = "d/m/yyyy HH:MM"

mutable struct DataReader <: Transformer
    model
    args

    function DataReader(args=Dict())
        default_args=Dict(
            :filename => "",
            :dateformat => DATEFMT,
            :impl_args => Dict()
        )
        new(nothing,mergedict(default_args,args))
    end
end

function fit!(dtr::DataReader,x::T=[],y::Vector=[]) where {T<:Union{DataFrame,Vector,Matrix}}
    fname = dtr.args[:filename]
    fmt = dtr.args[:dateformat]
    (fname != "" && isfile(fname) && fmt != "") || error("missing filename or date format: (",fname,"), (",fmt,")")
    dtr.model = dtr.args
end

function transform!(dtr::DataReader,x::T=[]) where {T<:Union{DataFrame,Vector,Matrix}}
    fullname = dtr.args[:filename]
    fmt = dtr.args[:dateformat]
    (fullname != "" && isfile(fullname) && fmt != "") || error("missing filename or date format: (",fullname,"), (",fmt,")")
    fname = basename(fullname)
    fname != "" || error("filename is empty")
    fn,ext=split(fname,".")
    ext != "" || error("no filename extension format")
    ext in keys(FILEFMT)  || error("extension not recognized "*ext)
    # dispatch based on extension
    readfmt(FILEFMT[ext],fullname,fmt)
end

function readfmt(::T,fname::String,datefmt::String) where {T<:Union{Val{:csv},Val{:feather},Val{:parquet}}}
    df = load(fname) |> DataFrame
    ncol(df) == 2 || error("dataframe should have only two columns: Date,Value")
    rename!(df,names(df)[1]=>:Date,names(df)[2]=>:Value)
    if eltype(df[:Date]) <: DateTime
        return df
    else
        df[:Date] = DateTime.(df[:Date],datefmt)
        return df
    end
end

function readfmt(atype::Union{Val{:hdf5},Val{:jld}},fname::String,datefmt::String)
    # R library(hdf5r); dat=H5File$new(fname);
    # dat[["dateval/date"]];dat[["dateval/value"]]
    if atype == Val{:hdf5}
        fileopen = h5open
    else
        fileopen = jldopen
    end
    ldate = fileopen(fname, "r") do file
        read(file, "dateval/date")
    end
    lvalue = fileopen(fname, "r") do file
        read(file, "dateval/value")
    end
    DataFrame(date=DateTime.(ldate,datefmt),Value=lvalue)
end

end


"""
Adapted from sandreza/Learning/sandbox/oceananigans_converter.jl
https://github.com/sandreza/Learning/blob/master/sandbox/oceananigans_converter.jl
"""

using JLD2

struct LESbraryData{𝒮, 𝒯, 𝒰, 𝒱}
    # Initial conditions
    T⁰ :: 𝒮
    U⁰ :: 𝒮
    V⁰ :: 𝒮

    # Fields at each moment in time
    T :: 𝒯
    U :: 𝒯
    V :: 𝒯

    # Some second order statistics at each moment in time
    wT :: 𝒯
    wu :: 𝒯
    wv :: 𝒯
    # uu :: 𝒯
    # vv :: 𝒯
    # ww :: 𝒯

    # Subfilter fluxes
    # νₑ_∂z_u :: 𝒯
    # νₑ_∂z_v :: 𝒯
    # νₑ_∂z_w :: 𝒯
    # κₑ_∂z_T :: 𝒯

    # Simulation constants
    α  :: 𝒰
    β  :: 𝒰
    f⁰ :: 𝒰
     g :: 𝒰
     L :: 𝒰

    # Time and grid
    t  :: 𝒮
    zC :: 𝒮
    zF :: 𝒮

    # Boundary conditions
    θ_top    :: 𝒰
    u_top    :: 𝒰
    θ_bottom :: 𝒰
    u_bottom :: 𝒰

    # Info about the simulation
    info :: 𝒱
end

function ReadJLD2_LESbraryData(filename)
    les_data = jldopen(filename, "r")
    les_keys = keys(les_data)
    timeseries_keys = keys(les_data["timeseries"]["t"])
    t₀ = timeseries_keys[1]

    # hold the entries for easy constructor creation
    container = []

    # size of arrays
    Nz = length(collect(les_data["grid"]["zC"])) - 6
    Nt = length(timeseries_keys)

    # Initial Conditions
    T⁰ = zeros(typeof(les_data["timeseries/T/$t₀"][1]), Nz)
    U⁰ = zeros(typeof(les_data["timeseries/u/$t₀"][1]), Nz)
    V⁰ = zeros(typeof(les_data["timeseries/v/$t₀"][1]), Nz)

    # Timeseries
    T = zeros(typeof(les_data["timeseries/T/$t₀"][1]), Nz, Nt)
    U = zeros(typeof(les_data["timeseries/u/$t₀"][1]), Nz, Nt)
    V = zeros(typeof(les_data["timeseries/v/$t₀"][1]), Nz, Nt)
    t = zeros(typeof(les_data["timeseries/t/$t₀"]), Nt)

    # Second Order Statistics
    wT  = zeros(typeof(les_data["timeseries/wT/$t₀"][1]), Nz+1, Nt)
    # uu  = zeros(typeof(les_data["timeseries/uu/$t₀"][1]), Nz,   Nt)
    # vv  = zeros(typeof(les_data["timeseries/vv/$t₀"][1]), Nz,   Nt)
    # ww  = zeros(typeof(les_data["timeseries/ww/$t₀"][1]), Nz+1, Nt)
    wu  = zeros(typeof(les_data["timeseries/wu/$t₀"][1]), Nz+1, Nt)
    wv  = zeros(typeof(les_data["timeseries/wv/$t₀"][1]), Nz+1, Nt)

    # Subfilter momentum fluxes
    # νₑ_∂z_u  = zeros(typeof(les_data["timeseries/νₑ_∂z_u/$t₀"][1]), Nz+1, Nt)
    # νₑ_∂z_v  = zeros(typeof(les_data["timeseries/νₑ_∂z_v/$t₀"][1]), Nz+1, Nt)
    # νₑ_∂z_w  = zeros(typeof(les_data["timeseries/νₑ_∂z_w/$t₀"][1]), Nz, Nt)
    # κₑ_∂z_T  = zeros(typeof(les_data["timeseries/κₑ_∂z_T/$t₀"][1]), Nz+1, Nt)

    # grab arrays
    for j in 1:Nt
        key = timeseries_keys[j]

        # Fields
        @. T[:, j] = les_data["timeseries/T/$key"][1, 1, :]
        @. U[:, j] = les_data["timeseries/u/$key"][1, 1, :]
        @. V[:, j] = les_data["timeseries/v/$key"][1, 1, :]

        # Second Order Statistics
        @. wT[:, j] = les_data["timeseries/wT/$key"][1, 1, :]
        @. wu[:, j] = les_data["timeseries/wu/$key"][1, 1, :]
        @. wv[:, j] = les_data["timeseries/wv/$key"][1, 1, :]
        # @. uu[:, j] = les_data["timeseries/uu/$key"][1, 1, :]
        # @. vv[:, j] = les_data["timeseries/vv/$key"][1, 1, :]
        # @. ww[:, j] = les_data["timeseries/ww/$key"][1, 1, :]

        # Subfilter fluxes
        # @. νₑ_∂z_u[:, j] = les_data["timeseries/νₑ_∂z_u/$key"][1, 1, :]
        # @. νₑ_∂z_v[:, j] = les_data["timeseries/νₑ_∂z_v/$key"][1, 1, :]
        # @. νₑ_∂z_w[:, j] = les_data["timeseries/νₑ_∂z_w/$key"][1, 1, :]
        # @. κₑ_∂z_T[:, j] = les_data["timeseries/κₑ_∂z_T/$key"][1, 1, :]

        t[j] = les_data["timeseries"]["t"][key]
    end

    # Set initial Conditions
    @. T⁰ = T[:,1]
    @. U⁰ = U[:,1]
    @. V⁰ = V[:,1]

    # Push initial conditions current stuff into container
    push!(container, T⁰, V⁰, U⁰)

    # Push fields into container
    push!(container, T, U, V)

    # Push second order statistics into container
    # push!(container, wT, wu, wv, uu, vv, ww)
    push!(container, wT, wu, wv)

    # Push subfilter fluxes into container
    # push!(container, νₑ_∂z_u, νₑ_∂z_v, νₑ_∂z_w, κₑ_∂z_T)

    # Now grab parameter
    if haskey(les_data["buoyancy"], "model")
        α = les_data["buoyancy/model/equation_of_state/α"]
        β = les_data["buoyancy/model/equation_of_state/β"]
        g = les_data["buoyancy/model/gravitational_acceleration"]
    else
        α = les_data["buoyancy/equation_of_state/α"]
        β = les_data["buoyancy/equation_of_state/β"]
        g = les_data["buoyancy/gravitational_acceleration"]
    end

    f⁰ = les_data["coriolis"]["f"]
    L = les_data["grid"]["Lz"]

    # Push parameters to container
    push!(container, α, β, f⁰, g, L)

    # grab domain data
    zC = collect(les_data["grid"]["zC"])[4:end-3] # padding of 3 on each side
    zF = collect(les_data["grid"]["zF"])[4:end-3] # padding of 3 on each side

    # push
    push!(container, t, zC, zF)

    # now grab boundary condition data
    # θ_top = les_data["parameters/boundary_condition_θ_top"]
    buoyancy_flux = les_data["parameters/buoyancy_flux"]
    θ_top = buoyancy_flux / α / g
    u_top = les_data["parameters/boundary_condition_u_top"]
    θ_bottom = les_data["parameters/boundary_condition_θ_bottom"]
    u_bottom = les_data["parameters/boundary_condition_u_bottom"]

    # push to container
    push!(container, θ_top, u_top, θ_bottom, u_bottom)

    # Now construct types
    𝒮 = typeof(U⁰)
    𝒯 = typeof(T)
    𝒰 = typeof(α)
    𝒱 = typeof("string")

    # now create data string
    info_string = "The top boundary conditions are flux boundary conditions \n"
    info_string *= "The  bottom boundary condition for temperature is a gradient boundary condition \n"
    info_string *= "The grid data is assumed to be evenly spaced and a power of two \n"

    # push to container
    push!(container, info_string)

    close(les_data)

    return LESbraryData{𝒮, 𝒯, 𝒰, 𝒱}(container...)
end

# TESTING
# file = "2DaySuite/three_layer_constant_fluxes_hr48_Qu0.0e+00_Qb1.2e-07_f1.0e-04_Nh256_Nz128_free_convection_statistics.jld2"
# les = jldopen(file, "r")
# keys(les["timeseries"])
# les["timeseries"]["κₑ_∂z_T"]
# les["timeseries"]["v"]["0"][1,1,:]
# les["grid"]["Δx"]

# using Plots
# file = "2DaySuite/three_layer_constant_fluxes_hr48_Qu0.0e+00_Qb1.2e-07_f1.0e-04_Nh256_Nz128_free_convection_statistics.jld2"
# myles = ReadJLD2_LESbraryData(file)
#
# vbl = myles.κₑ_∂z_T
# xlims=(-1.2e-5, maximum(myles.κₑ_∂z_T))
# vbl = myles.wT
# xlims=(minimum(myles.wT), maximum(myles.wT))
#
# anim = @animate for i=1:288
#     plot(vbl[:,i], myles.zF, legend=false, xlims=xlims, size=(400,600))
# end
# gif(anim, "wT.gif")
#
# Δt = diff(myles.t, dims=1)'
# Nz,Nt = size(myles.T)
# dudt = (myles.u[:,2:Nt] .- myles.u[:,1:Nt-1]) ./ Δt # Nz x (Nt-1) array of approximate dUdt values


using SplineFEEC
using Test

@testset "SplineFEEC.jl" begin
    # Write your tests here.
end


###############################
### Result reduction functions
###############################

"""
Reduce buffers

Join buffers returned by multiple runs of simulation

**Arguments**
* `buffer` : list of `OpenStreetMapXSim.Road` objects containing informations about routes selected during the simulation run.

"""
function reduce_results(buffer::Array{OpenStreetMapXSim.Road,1}...)::Array{OpenStreetMapXSim.Road,1}
    reduced_buffer = OpenStreetMapXSim.Road[]
    buffer = vcat(buffer...)
    buffer  = [(Tuple(getfield(buffer[i],field) for field in fieldnames(typeof(buffer[i]))[1:end-1]), buffer[i].count) for i = 1:length(buffer)]
    unique_roads = Set(map(road -> road[1], buffer))
    for road in unique_roads
        count = sum(buffer[i][2] for i in findall(x->(x[1] == road), buffer))
        push!(reduced_buffer,OpenStreetMapXSim.Road(road...,count))
    end
    return reduced_buffer
end

"""
Reduce nodes statistics

Join nodes statistics returned by multiple runs of simulation

**Arguments**
* `nodes` : dictionary of `OpenStreetMapXSim.NodeStat` objects containing informations about each intersection in simulation.

"""
function reduce_results(nodes::Dict{Int64,OpenStreetMapXSim.NodeStat}...)::Dict{Int64,OpenStreetMapXSim.NodeStat}
    reduced_nodes = Dict{Int64,OpenStreetMapXSim.NodeStat}()
    ids = keys(merge(nodes...))
    for id in ids
        count = sum(node[id].count for node in nodes)
        lat = unique(node[id].latitude for node in nodes)[1]
        lon = unique(node[id].longitude for node in nodes)[1]
        agents_data = vcat([node[id].agents_data for node in nodes]...)
        reduced_nodes[id] = OpenStreetMapXSim.NodeStat(count,lat,lon,agents_data)
    end
    return reduced_nodes
end


# This file was generated, do not modify it. # hide
using MLJ
X, y = @load_iris
@load DecisionTreeClassifier


"""
	filter1d(cmd0::String="", arg1=[], kwargs...)

Time domain filtering of 1-D data tables.

Full option list at [`filter1d`](http://gmt.soest.hawaii.edu/doc/latest/filter1d.html)

Parameters
----------

- **F** : **filter_type** : -- Str --   Flags = type width[modifiers]

    Sets the filter type. Choose among convolution and non-convolution filters. Append the
    filter code followed by the full filter width in same units as time column.
    [`-F`](http://gmt.soest.hawaii.edu/doc/latest/filter1d.html#f)
- **D** : **inc** : -- Number --        Flags = increment

    ``increment`` is used when series is NOT equidistantly sampled. Then increment will be the abscissae resolution.
    [`-D`](http://gmt.soest.hawaii.edu/doc/latest/filter1d.html#d)
- **E** : **ends** : -- Bool or [] --

    Include Ends of time series in output. Default loses half the filter-width of data at each end.
    [`-E`](http://gmt.soest.hawaii.edu/doc/latest/filter1d.html#e)
- **N** : **time_col** : -- Int --      Flags = t_col

    Indicates which column contains the independent variable (time). The left-most column
    is # 0, the right-most is # (n_cols - 1). [Default is 0].
    [`-N`](http://gmt.soest.hawaii.edu/doc/latest/filter1d.html#n)
- **Q** : **quality** : -- Number --    Flags = q_factor

    Assess Quality of output value by checking mean weight in convolution. Enter q_factor between 0 and 1.
    [`-Q`](http://gmt.soest.hawaii.edu/doc/latest/filter1d.html#q)
- **S** : **symetry** : -- Number --    Flags = symmetry_factor

    Checks symmetry of data about window center. Enter a factor between 0 and 1.
    [`-S`](http://gmt.soest.hawaii.edu/doc/latest/filter1d.html#s)
- **T** : **equi_space** : -- List or Str --     Flags = [min/max/]inc[+a|n]

    Make evenly spaced time-steps from min to max by inc [Default uses input times].
    [`-T`](http://gmt.soest.hawaii.edu/doc/latest/filter1d.html#t)
- $(GMT.opt_V)
- $(GMT.opt_write)
- $(GMT.opt_append)
- $(GMT.opt_b)
- $(GMT.opt_d)
- $(GMT.opt_e)
- $(GMT.opt_f)
- $(GMT.opt_g)
- $(GMT.opt_h)
- $(GMT.opt_i)
- $(GMT.opt_o)
- $(GMT.opt_swap_xy)
"""
function filter1d(cmd0::String="", arg1=[]; kwargs...)

	length(kwargs) == 0 && return monolitic("filter1d", cmd0, arg1)	# Speedy mode

	d = KW(kwargs)

	cmd = parse_V_params("", d)
	cmd, = parse_b(cmd, d)
	cmd, = parse_d(cmd, d)
	cmd, = parse_e(cmd, d)
	cmd, = parse_f(cmd, d)
	cmd, = parse_g(cmd, d)
	cmd, = parse_h(cmd, d)
	cmd, = parse_i(cmd, d)
	cmd, = parse_o(cmd, d)
	cmd, = parse_swap_xy(cmd, d)

	cmd = add_opt(cmd, 'F', d, [:F :filter_type])
	cmd = add_opt(cmd, 'D', d, [:D :inc])
	cmd = add_opt(cmd, 'E', d, [:E :ends])
	cmd = add_opt(cmd, 'N', d, [:N :time_col])
	cmd = add_opt(cmd, 'Q', d, [:Q :quality])
	cmd = add_opt(cmd, 'S', d, [:S :symetry])
	cmd = add_opt(cmd, 'T', d, [:T :equi_space])

	cmd, got_fname, arg1 = find_data(d, cmd0, cmd, 1, arg1)
	return common_grd(d, cmd, got_fname, 1, "filter1d", arg1)		# Finish build cmd and run it
end

# ---------------------------------------------------------------------------------------------------
filter1d(arg1=[], cmd0::String=""; kw...) = filter1d(cmd0, arg1; kw...)


module DecisionTrees

using Printf

export children,
        class,
        entropy,
        id3,
        information_gain,
        nchildren,
        Node,
        value

include("algorithms.jl")
include("measures.jl")
include("node.jl")

end # module DecisionTrees


import Random.shuffle
import Random.shuffle!

function swap_degree!(Y::BinaryNetwork)
   iy = interactions(Y)
   i1, i2 = sample(iy, 2, replace=false)
   n1, n2 = (from=i1.from, to=i2.to), (from=i2.from, to=i1.to)

   while (n2 ∈ iy)|(n1 ∈ iy)
      i1, i2 = sample(iy, 2, replace=false)
      n1, n2 = (from=i1.from, to=i2.to), (from=i2.from, to=i1.to)
   end

   for old_i in [i1, i2, n1, n2]
      i = something(findfirst(isequal(old_i.from), species(Y; dims=1)), 0)
      j = something(findfirst(isequal(old_i.to), species(Y; dims=2)), 0)
      Y.A[i,j] = !Y.A[i,j]
   end
end

function swap_fill!(Y::BinaryNetwork)
   iy = interactions(Y)
   i1 = sample(iy)
   n1 = (from=sample(species(Y; dims=1)), to=sample(species(Y; dims=2)))

   while (n1 ∈ iy)
      n1 = (from=sample(species(Y; dims=1)), to=sample(species(Y; dims=2)))
   end

   for old_i in [i1, n1]
      i = something(findfirst(isequal(old_i.from), species(Y; dims=1)), 0)
      j = something(findfirst(isequal(old_i.to), species(Y; dims=2)), 0)
      Y.A[i,j] = !Y.A[i,j]
   end
end

function swap_vulnerability!(Y::BinaryNetwork)
   iy = interactions(Y)
   i1 = sample(iy)
   n1 = (from=sample(species(Y; dims=1)), to=i1.to)

   while (n1 ∈ iy)
      n1 = (from=sample(species(Y; dims=1)), to=i1.to)
   end

   for old_i in [i1, n1]
      i = something(findfirst(isequal(old_i.from), species(Y; dims=1)), 0)
      j = something(findfirst(isequal(old_i.to), species(Y; dims=2)), 0)
      Y.A[i,j] = !Y.A[i,j]
   end
end

function swap_generality!(Y::BinaryNetwork)
   iy = interactions(Y)
   i1 = sample(iy)
   n1 = (from=i1.from, to=sample(species(Y; dims=2)))

   while (n1 ∈ iy)
      n1 = (from=i1.from, to=sample(species(Y; dims=2)))
   end

   for old_i in [i1, n1]
      i = something(findfirst(isequal(old_i.from), species(Y; dims=1)), 0)
      j = something(findfirst(isequal(old_i.to), species(Y; dims=2)), 0)
      Y.A[i,j] = !Y.A[i,j]
   end
end

"""
    shuffle(N::BinaryNetwork; constraint::Symbol=:degree)

Return a shuffled copy of the network (the original network is not modified).
See `shuffle!` for a documentation of the keyword arguments.

#### References

Fortuna, M.A., Stouffer, D.B., Olesen, J.M., Jordano, P., Mouillot, D., Krasnov,
B.R., Poulin, R., Bascompte, J., 2010. Nestedness versus modularity in
ecological networks: two sides of the same coin? Journal of Animal Ecology 78,
811–817. https://doi.org/10.1111/j.1365-2656.2010.01688.x
"""
function shuffle(N::BinaryNetwork; constraint::Symbol=:degree)
   Y = copy(N)
   shuffle!(Y; constraint=constraint)
   return Y
end

"""
    shuffle!(N::BinaryNetwork; constraint::Symbol=:degree)

Shuffles interactions inside a network (the network is *modified*), under the
following `constraint`:

- :degree, which keeps the degree distribution intact
- :generality, which keeps the out-degree distribution intact
- :vulnerability, which keeps the in-degree distribution intact
- :fill, which moves interactions around freely

The function will take two interactions, and swap the species establishing them.
By repeating the process a large enough number of times, the resulting network
should be relatively random. Note that this function will conserve the degree
(when appropriate under the selected constraint) of *every* species. Calling the
function will perform **a single** shuffle. If you want to repeat the shuffling
a large enough number of times, you can use something like:

    [shuffle!(n) for i in 1:10_000]

If the keyword arguments are invalid, the function will throw an
`ArgumentError`.

#### References

Fortuna, M.A., Stouffer, D.B., Olesen, J.M., Jordano, P., Mouillot, D., Krasnov,
B.R., Poulin, R., Bascompte, J., 2010. Nestedness versus modularity in
ecological networks: two sides of the same coin? Journal of Animal Ecology 78,
811–817. https://doi.org/10.1111/j.1365-2656.2010.01688.x
"""
function shuffle!(N::BinaryNetwork; constraint::Symbol=:degree)
   constraint ∈ [:degree, :generality, :vulnerability, :fill] || throw(ArgumentError("The constraint argument you specificied ($(constraint)) is invalid -- see ?shuffle! for a list."))

   f = EcologicalNetworks.swap_degree!
   constraint == :generality && (f = EcologicalNetworks.swap_generality!)
   constraint == :vulnerability && (f = EcologicalNetworks.swap_vulnerability!)
   constraint == :fill && (f = EcologicalNetworks.swap_fill!)

   f(N)
end


module NLCOP

using Kombinator
using JuMP

export ExactNonlinearSolver, ApproximateNonlinearSolver

import Kombinator: solve

include("exact.jl")
include("approx.jl") # The actual NLCOP algorithm.

end


## DateString Aliasing

yrmo2dir(date::TimeType) = Dates.format(date,dateformat"yyyy/mm")
yrmo2str(date::TimeType) = Dates.format(date,dateformat"yyyymm")
yr2str(date::TimeType)   = Dates.format(date,dateformat"yyyy")
ymd2str(date::TimeType)  = Dates.format(date,dateformat"yyyymmdd")
mo2str(date::TimeType)   = Dates.format(date,dateformat"mm")

yrmo2str(yr::Integer,mo::Integer) = @sprintf("%04d%02d",yr,mo)
mo2str(mo::Integer) = @sprintf("%02d",mo)
dy2str(dy::Integer) = @sprintf("%02d",dy)

function extractdate(startdate::TimeType,finish::TimeType);

    yrs = Dates.year(start);  mos = Dates.month(start);  dys = Dates.day(start);
    yrf = Dates.year(finish); mof = Dates.month(finish); dyf = Dates.day(finish);
    ndy = Dates.value((finish-start)/Dates.day(1));
    dvecs = Date(yrs,mos); dvecf = Date(yrf,mof);

    dvec = collect(dvecs:Month(1):dvecf);

    return dvec,dys,dyf,ndy

end

bold() = Crayon(bold=true)
reset() = Crayon(reset=true)

function checkegrid(step::Real)

    if rem(360,step) == 0
          return true
    else; return false
    end

end


module JuLearn

using JLD

function include_path(path,folder)
	files=readdir(joinpath(path,folder))
	map(file->include(joinpath(path,folder,file)),files)
end

export Layer

abstract Layer

if ENV["JuLearnBackend"]=="native"
	path=Pkg.dir("JuLearn","src","native")
elseif ENV["JuLearnBackend"]=="BLAS"
	path=Pkg.dir("JuLearn","src","BLAS")
elseif ENV["JuLearnBackend"]=="ArrayFire"
	using ArrayFire
	path=Pkg.dir("JuLearn","src","ArrayFire")
end

include_path(path,"layers")
include_path(path,"essentials")

include_path(path,"utils")
include_path(path,"activation")
include_path(path,"chain")
include_path(path,"loss")
include_path(path,"rnn")

end # module


using SeisIO, SeisIO.RandSeis

# equivalence:
# FDSN wildcard         RegEx
#     ?           =>    .
#     *           =>    .*

function id_to_regex(cid::Array{UInt8,1})
  replace!(cid, 0x2e=>0x5c, 0x3f=>0x2e)
  L = length(cid)
  i = L
  while i > 0
    # replace: '*' => '.*'
    if cid[i] == 0x2a
      splice!(cid, i:i-1, 0x2e)
      i -= 1
    # now replace: '\' => '\.'
    elseif cid[i] == 0x5c
      if i == L
        push!(cid, 0x2e)
      else
        splice!(cid, i:i, [0x5c, 0x2e])
      end
    end
    i -= 1
  end
  return Regex(String(cid))
end

id_to_regex(id::AbstractString) = id_to_regex(copy(codeunits(id)))

function netsta_to_regex(id::AbstractString)
  cid = copy(codeunits(id))
  i = 0
  j = 0
  L = length(cid)
  while i < L
    i += 1
    if cid[i] == 0x2e
      j += 1
    end
    if j == 2
      deleteat!(cid, i:L)
      break
    end
  end
  return id_to_regex(cid)
end

function id_match(id::AbstractString, S::GphysData)
  j = findid(id, S.id)
  j > 0 && return [j]

  idr = id_to_regex(id)
  chans = Int64[]
  for (j, cid) in enumerate(S.id)
    if occursin(idr, cid)
      push!(chans, j)
    end
  end
  return chans
end


using MusicModels
using NoteSequences
using Test

include("musicencoders.jl")


"""
    insert_location(ex::Expr, location)

Insert a symbolic representation of `location` into the arguments of an `expression`.

Used in the `@at` macro for specifying the location of an `AbstractOperation`.
"""
function insert_location!(ex::Expr, location)
    if ex.head === :call && ex.args[1] ∈ operators
        push!(ex.args, ex.args[end])
        ex.args[3:end-1] .= ex.args[2:end-2]
        ex.args[2] = location
    end

    for arg in ex.args
        insert_location!(arg, location)
    end

    return nothing
end

"Fallback for when `insert_location` is called on objects other than expressions."
insert_location!(anything, location) = nothing

# A very special UnaryOperation
@inline interpolate_identity(x) = x
@unary interpolate_identity

interpolate_operation(L, x) = x

function interpolate_operation(L, x::AbstractField)
    L == location(x) && return x # Don't interpolate unecessarily
    return interpolate_identity(L, x)
end

"""
    @at location abstract_operation

Modify the `abstract_operation` so that it returns values at
`location`, where `location` is a 3-tuple of `Face`s and `Center`s.
"""
macro at(location, abstract_operation)
    insert_location!(abstract_operation, location)

    # We wrap it all in an interpolator to help "stubborn" binary operations
    # arrive in the right place.
    wrapped_operation = quote
        interpolate_operation($(esc(location)), $(esc(abstract_operation)))
    end

    return wrapped_operation
end


using Mjolnir
import Mjolnir: Defaults, AType, @trace, abstract

struct ConditionError end

function observe(result::Bool)
  result || throw(ConditionError())
  return
end

abstract(::Basic, ::AType{typeof(observe)}, ::AType{Bool}) = Nothing
Mjolnir.effectful(::AType{typeof(observe)}, x) = true

struct Multi
  algs::Vector{Any}
  Multi(algs...) = new(Any[algs...])
end

function infer(f, m::Multi)
  tr = trace(typeof(f))
  for alg in m.algs
    r = infer(f, alg, tr)
    r == nothing || return r
  end
end

default() = Multi(Trivial(), Quad(), Rejection())

infer(f) = infer(f, default())


function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{T, D},
    ::Type{S},
) where {
    T <: PSY.HybridSystem,
    D <: AbstractHybridFormulation,
    S <: PM.AbstractActivePowerModel,
}
    devices = get_available_components(T, sys)

    # Variables
    add_variables!(container, ActivePowerVariable, devices, D())
    add_variables!(container, ReactivePowerVariable, devices, D())
    add_variables!(container, ActivePowerInVariable, devices, D())
    add_variables!(container, ActivePowerOutVariable, devices, D())
    add_variables!(container, ComponentActivePowerVariable, devices, D())
    add_variables!(container, EnergyVariable, devices, D())
    if get_attribute(model, "reservation")
        add_variables!(container, ReservationVariable, devices, D())
    end

    # Parameters
    add_parameters!(container, ActivePowerTimeSeriesParameter, devices, model)

    # Initial Conditions
    initial_conditions!(container, devices, D())

    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ComponentActivePowerRangeExpressionLB,
        ComponentActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ComponentActivePowerRangeExpressionUB,
        ComponentActivePowerVariable,
        devices,
        model,
        S,
    )
    add_feedforward_arguments!(container, model, devices)

    if has_service_model(model)
        add_variables!(container, ComponentActivePowerReserveUpVariable, devices, D())
        add_variables!(container, ComponentActivePowerReserveDownVariable, devices, D())

        add_to_expression!(
            container,
            ComponentActivePowerRangeExpressionLB,
            ComponentActivePowerReserveDownVariable,
            devices,
            model,
            S,
        )
        add_to_expression!(
            container,
            ComponentActivePowerRangeExpressionUB,
            ComponentActivePowerReserveUpVariable,
            devices,
            model,
            S,
        )
        add_expressions!(container, ComponentReserveDownBalanceExpression, devices, model)
        add_expressions!(container, ComponentReserveUpBalanceExpression, devices, model)
    end
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{T, D},
    ::Type{S},
) where {
    T <: PSY.HybridSystem,
    D <: AbstractHybridFormulation,
    S <: PM.AbstractActivePowerModel,
}
    devices = get_available_components(T, sys)

    # Constraints
    #TODO: Change to Expressions for Services support
    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        OutputActivePowerVariableLimitsConstraint,
        ActivePowerOutVariable,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        InputActivePowerVariableLimitsConstraint,
        ActivePowerInVariable,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ComponentActivePowerVariableLimitsConstraint,
        ComponentActivePowerRangeExpressionLB,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ComponentActivePowerVariableLimitsConstraint,
        ComponentActivePowerRangeExpressionUB,
        devices,
        model,
        S,
    )

    add_constraints!(container, EnergyBalanceConstraint, devices, model, S)
    add_constraints!(container, PowerOutputRangeConstraint, devices, model, S)

    if has_service_model(model)
        add_constraints!(container, ReserveEnergyConstraint, devices, model, S)
        add_constraints!(container, RangeLimitConstraint, devices, model, S)
        add_constraints!(container, ComponentReserveUpBalance, devices, model, S)
        add_constraints!(container, ComponentReserveDownBalance, devices, model, S)
    end

    add_feedforward_constraints!(container, model, devices)

    # Cost Function
    objective_function!(container, devices, model, S)

    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{T, D},
    ::Type{S},
) where {T <: PSY.HybridSystem, D <: AbstractHybridFormulation, S <: PM.AbstractPowerModel}
    devices = get_available_components(T, sys)

    # Variables
    add_variables!(container, ActivePowerVariable, devices, D())
    add_variables!(container, ActivePowerInVariable, devices, D())
    add_variables!(container, ActivePowerOutVariable, devices, D())
    add_variables!(container, ComponentActivePowerVariable, devices, D())
    add_variables!(container, EnergyVariable, devices, D())

    add_variables!(container, ReactivePowerVariable, devices, D())
    add_variables!(container, ComponentReactivePowerVariable, devices, D())
    if get_attribute(model, "reservation")
        add_variables!(container, ReservationVariable, devices, D())
    end

    # Parameters
    add_parameters!(container, ActivePowerTimeSeriesParameter, devices, model)

    # Initial Conditions
    initial_conditions!(container, devices, D())

    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ReactivePowerBalance,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ComponentActivePowerRangeExpressionLB,
        ComponentActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ComponentActivePowerRangeExpressionUB,
        ComponentActivePowerVariable,
        devices,
        model,
        S,
    )

    add_feedforward_arguments!(container, model, devices)

    if has_service_model(model)
        add_variables!(container, ComponentActivePowerReserveUpVariable, devices, D())
        add_variables!(container, ComponentActivePowerReserveDownVariable, devices, D())
        add_to_expression!(
            container,
            ComponentActivePowerRangeExpressionLB,
            ComponentActivePowerReserveDownVariable,
            devices,
            model,
            S,
        )
        add_to_expression!(
            container,
            ComponentActivePowerRangeExpressionUB,
            ComponentActivePowerReserveUpVariable,
            devices,
            model,
            S,
        )
        add_expressions!(container, ComponentReserveDownBalanceExpression, devices, model)
        add_expressions!(container, ComponentReserveUpBalanceExpression, devices, model)
    end
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{T, D},
    ::Type{S},
) where {T <: PSY.HybridSystem, D <: AbstractHybridFormulation, S <: PM.AbstractPowerModel}
    devices = get_available_components(T, sys)

    # Constraints
    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ComponentActivePowerVariableLimitsConstraint,
        ComponentActivePowerRangeExpressionLB,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ComponentActivePowerVariableLimitsConstraint,
        ComponentActivePowerRangeExpressionUB,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        InputActivePowerVariableLimitsConstraint,
        ActivePowerInVariable,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        OutputActivePowerVariableLimitsConstraint,
        ActivePowerOutVariable,
        devices,
        model,
        S,
    )

    # Reactive power Constraints
    add_constraints!(
        container,
        ReactivePowerVariableLimitsConstraint,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ComponentReactivePowerVariableLimitsConstraint,
        ComponentReactivePowerVariable,
        devices,
        model,
        S,
    )

    # Constraints
    add_constraints!(container, EnergyBalanceConstraint, devices, model, S)
    add_constraints!(container, PowerOutputRangeConstraint, devices, model, S)
    add_constraints!(container, ReactivePowerConstraint, devices, model, S)
    add_constraints!(container, InterConnectionLimitConstraint, devices, model, S)
    if has_service_model(model)
        add_constraints!(container, ReserveEnergyConstraint, devices, model, S)
        add_constraints!(container, RangeLimitConstraint, devices, model, S)
        add_constraints!(container, ComponentReserveUpBalance, devices, model, S)
        add_constraints!(container, ComponentReserveDownBalance, devices, model, S)
    end
    add_feedforward_constraints!(container, model, devices)

    # Cost Function
    objective_function!(container, devices, model, S)

    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{T, D},
    ::Type{S},
) where {
    T <: PSY.HybridSystem,
    D <: StandardHybridDispatch,
    S <: PM.AbstractActivePowerModel,
}
    devices = get_available_components(T, sys)

    # Variables
    add_variables!(container, ActivePowerVariable, devices, D())
    add_variables!(container, ReactivePowerVariable, devices, D())
    add_variables!(container, ActivePowerInVariable, devices, D())
    add_variables!(container, ActivePowerOutVariable, devices, D())
    add_variables!(container, ComponentActivePowerVariable, devices, D())
    add_variables!(container, EnergyVariable, devices, D())
    if get_attribute(model, "reservation")
        add_variables!(container, ReservationVariable, devices, D())
    end

    # Parameters
    add_parameters!(container, ActivePowerTimeSeriesParameter, devices, model)

    # Initial Conditions
    initial_conditions!(container, devices, D())

    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ComponentActivePowerRangeExpressionLB,
        ComponentActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ComponentActivePowerRangeExpressionUB,
        ComponentActivePowerVariable,
        devices,
        model,
        S,
    )

    add_feedforward_arguments!(container, model, devices)

    if has_service_model(model)
        add_variables!(container, ComponentActivePowerReserveUpVariable, devices, D())
        add_variables!(container, ComponentActivePowerReserveDownVariable, devices, D())
        add_to_expression!(
            container,
            ComponentActivePowerRangeExpressionLB,
            ComponentActivePowerReserveDownVariable,
            devices,
            model,
            S,
        )
        add_to_expression!(
            container,
            ComponentActivePowerRangeExpressionUB,
            ComponentActivePowerReserveUpVariable,
            devices,
            model,
            S,
        )
        add_expressions!(container, ComponentReserveDownBalanceExpression, devices, model)
        add_expressions!(container, ComponentReserveUpBalanceExpression, devices, model)
    end
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{T, D},
    ::Type{S},
) where {
    T <: PSY.HybridSystem,
    D <: StandardHybridDispatch,
    S <: PM.AbstractActivePowerModel,
}
    devices = get_available_components(T, sys)

    # Constraints
    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ComponentActivePowerVariableLimitsConstraint,
        ComponentActivePowerRangeExpressionLB,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ComponentActivePowerVariableLimitsConstraint,
        ComponentActivePowerRangeExpressionUB,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        InputActivePowerVariableLimitsConstraint,
        ActivePowerInVariable,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        OutputActivePowerVariableLimitsConstraint,
        ActivePowerOutVariable,
        devices,
        model,
        S,
    )

    add_constraints!(container, EnergyBalanceConstraint, devices, model, S)
    add_constraints!(container, PowerOutputRangeConstraint, devices, model, S)

    if has_service_model(model)
        add_constraints!(container, ReserveEnergyConstraint, devices, model, S)
        add_constraints!(container, RangeLimitConstraint, devices, model, S)
        add_constraints!(container, ComponentReserveUpBalance, devices, model, S)
        add_constraints!(container, ComponentReserveDownBalance, devices, model, S)
    end

    add_feedforward_constraints!(container, model, devices)

    # Cost Function
    objective_function!(container, devices, model, S)

    return
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ArgumentConstructStage,
    model::DeviceModel{T, D},
    ::Type{S},
) where {T <: PSY.HybridSystem, D <: StandardHybridDispatch, S <: PM.AbstractPowerModel}
    devices = get_available_components(T, sys)

    # Variables
    add_variables!(container, ActivePowerVariable, devices, D())
    add_variables!(container, ActivePowerInVariable, devices, D())
    add_variables!(container, ActivePowerOutVariable, devices, D())
    add_variables!(container, ComponentActivePowerVariable, devices, D())
    add_variables!(container, EnergyVariable, devices, D())

    add_variables!(container, ReactivePowerVariable, devices, D())
    add_variables!(container, ComponentReactivePowerVariable, devices, D())
    if get_attribute(model, "reservation")
        add_variables!(container, ReservationVariable, devices, D())
    end

    # Parameters
    add_parameters!(container, ActivePowerTimeSeriesParameter, devices, model)

    # Initial Conditions
    initial_conditions!(container, devices, D())

    add_to_expression!(
        container,
        ActivePowerBalance,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ReactivePowerBalance,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ComponentActivePowerRangeExpressionLB,
        ComponentActivePowerVariable,
        devices,
        model,
        S,
    )
    add_to_expression!(
        container,
        ComponentActivePowerRangeExpressionUB,
        ComponentActivePowerVariable,
        devices,
        model,
        S,
    )

    add_feedforward_arguments!(container, model, devices)

    if has_service_model(model)
        add_variables!(container, ComponentActivePowerReserveUpVariable, devices, D())
        add_variables!(container, ComponentActivePowerReserveDownVariable, devices, D())
        add_to_expression!(
            container,
            ComponentActivePowerRangeExpressionLB,
            ComponentActivePowerReserveDownVariable,
            devices,
            model,
            S,
        )
        add_to_expression!(
            container,
            ComponentActivePowerRangeExpressionUB,
            ComponentActivePowerReserveUpVariable,
            devices,
            model,
            S,
        )
        add_expressions!(container, ComponentReserveDownBalanceExpression, devices, model)
        add_expressions!(container, ComponentReserveUpBalanceExpression, devices, model)
    end
end

function construct_device!(
    container::OptimizationContainer,
    sys::PSY.System,
    ::ModelConstructStage,
    model::DeviceModel{T, D},
    ::Type{S},
) where {T <: PSY.HybridSystem, D <: StandardHybridDispatch, S <: PM.AbstractPowerModel}
    devices = get_available_components(T, sys)
    # Constraints
    add_constraints!(
        container,
        ActivePowerVariableLimitsConstraint,
        ActivePowerVariable,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ComponentActivePowerVariableLimitsConstraint,
        ComponentActivePowerRangeExpressionLB,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ComponentActivePowerVariableLimitsConstraint,
        ComponentActivePowerRangeExpressionUB,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        InputActivePowerVariableLimitsConstraint,
        ActivePowerInVariable,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        OutputActivePowerVariableLimitsConstraint,
        ActivePowerOutVariable,
        devices,
        model,
        S,
    )

    # Reactive power Constraints
    add_constraints!(
        container,
        ReactivePowerVariableLimitsConstraint,
        ReactivePowerVariable,
        devices,
        model,
        S,
    )
    add_constraints!(
        container,
        ComponentReactivePowerVariableLimitsConstraint,
        ComponentReactivePowerVariable,
        devices,
        model,
        S,
    )

    # Constraints
    add_constraints!(container, EnergyBalanceConstraint, devices, model, S)
    add_constraints!(container, PowerOutputRangeConstraint, devices, model, S)
    add_constraints!(container, ReactivePowerConstraint, devices, model, S)
    add_constraints!(container, InterConnectionLimitConstraint, devices, model, S)

    if has_service_model(model)
        add_constraints!(container, ReserveEnergyConstraint, devices, model, S)
        add_constraints!(container, RangeLimitConstraint, devices, model, S)
        add_constraints!(container, ComponentReserveUpBalance, devices, model, S)
        add_constraints!(container, ComponentReserveDownBalance, devices, model, S)
    end

    add_feedforward_constraints!(container, model, devices)

    # Cost Function
    objective_function!(container, devices, model, S)

    return
end


###############################################################################
#
# Calculate gas exchange from flow rate only for required parameters
#
###############################################################################
"""
    leaf_gas_exchange_nonopt!(
                node::SPACSimple{FT},
                flow::FT,
                par::FT,
                rad::FT,
                la::FT,
                container::SPACContainer1L{FT}
    ) where {FT<:AbstractFloat}
    leaf_gas_exchange_nonopt!(
                node::SPACSimple{FT},
                flow::FT
    ) where {FT<:AbstractFloat}
    leaf_gas_exchange_nonopt!(
                node::SPACSimple{FT},
                f_sl::FT,
                f_sh::FT
    ) where {FT<:AbstractFloat}

Simulate leaf level gas exchange and fill it into the `container` for 1-layer
    or 2-layer canopy, given
- `node` [`SPACSimple`] type struct
- `flow` Flow rate per basal area into the leaves (e.g., for sunlit leaves)
- `f_sl` Flow rate per basal area into the sunlit leaves
- `f_sh` Flow rate per basal area into the shaded leaves
- `par` Leaf-level photosynthetic active radiation
- `rad` Leaf-level absorbed radiative energy
- `la` Leaf area of the leaves (total or each layer)
- `container` [`SPACContainer1L`] type container
"""
function leaf_gas_exchange_nonopt!(
            node::SPACSimple{FT},
            flow::FT,
            par::FT,
            rad::FT,
            la::FT,
            container::SPACContainer1L{FT}
) where {FT<:AbstractFloat}
    # if flow < 0
    if flow < 0
        container.an = FT(-Inf);

    # if flow >= 0
    else
        # 0. unpack required variables
        @unpack envir, g_max, width = node;
        @unpack P_AIR, p_H₂O, wind = envir;

        # 1. calculate leaf temperature from the flow rate
        t_leaf = max(200, leaf_temperature(node, rad, flow));

        # 2. update leaf photosynthetic variables and leaf-to-air VPD
        node.ps.apar = par;
        photosystem_temperature_dependence!(node.ps.PSM, envir, t_leaf);
        node.ps.p_H₂O_sat = saturation_vapor_pressure(t_leaf);
        d_leaf = node.ps.p_H₂O_sat - p_H₂O;

        # 3. update f_vis and f_st in leaf and calculate water potentials
        # TODO do not account for temperature effects for flow now
        # More reasonable functions need to be added
        # vc_temperature_effects!(node.hs.leaf, t_leaf);

        # if flow == 0
        if flow == 0
            g_lw = FT(0);
            g_lc = FT(1e-6);
            leaf_photosynthesis!(node.ps, envir, GCO₂Mode(), g_lc);
            container.an = node.ps.PSM.a_net;

        # if flow > 0 and reasonable
        elseif (d_leaf > 0) && (t_leaf > T_0(FT))
            t_cor = relative_diffusive_coefficient( (t_leaf+node.envir.t)/2 );
            g_lw  = flow / la / d_leaf * P_AIR;
            g_bw  = boundary_layer_conductance(wind, width);
            g_bc  = g_bw / FT(1.35);
            g_sw  = 1 / (1/g_lw - 1/(g_bw*t_cor));
            g_sc  = g_sw / FT(1.6);
            g_lc  = max(FT(1e-6), 1 / (1/g_bc + 1/g_sc));
            g_lim = 1 / (1/g_bw + 1/g_max);
            if g_lw < g_lim * t_cor
                leaf_photosynthesis!(node.ps, envir, GCO₂Mode(), g_lc);
                container.an = node.ps.PSM.a_net;
            else
                container.an = FT(-Inf);
            end

        # if flow > 0 and unrealistic
        else
            container.an = FT(-Inf);
        end
    end

    return nothing
end

function leaf_gas_exchange_nonopt!(
            node::SPACSimple{FT},
            flow::FT
) where {FT<:AbstractFloat}
    # unpack the data
    @unpack frac_sh, frac_sl, par_sh, par_sl, rad_sh, rad_sl = node.container2L;

    # calculate mean par and rad per leaf area, then gas exchange rate
    par_mean = par_sl * frac_sl + par_sh * frac_sh;
    rad_mean = rad_sl * frac_sl + rad_sh * frac_sh;
    leaf_gas_exchange_nonopt!(node, flow, par_mean, rad_mean, node.laba, node.container1L);

    node.containerOP = (node.ec - flow) * (node.container1L).an;

    return nothing
end

function leaf_gas_exchange_nonopt!(
            node::SPACSimple{FT},
            f_sl::FT,
            f_sh::FT
) where {FT<:AbstractFloat}
    # unpack the data
    @unpack frac_sh, frac_sl, la_sh, la_sl, par_sh, par_sl, rad_sh,
            rad_sl = node.container2L;

    # calculate gas exchangr for sunlit and shaded layers
    leaf_gas_exchange_nonopt!(node, f_sl, par_sl, rad_sl, la_sl, node.container2L.cont_sl);
    leaf_gas_exchange_nonopt!(node, f_sh, par_sh, rad_sh, la_sh, node.container2L.cont_sh);

    a_sum = frac_sl * node.container2L.cont_sl.an + frac_sh * node.container2L.cont_sh.an;
    e_sum = f_sl + f_sh;
    node.containerOP = (node.ec - e_sum) * a_sum;

    return nothing
end

###############################################################################
#
# Calculate gas exchange from flow rate considering temperature dynamics
#
###############################################################################
"""
    leaf_gas_exchange!(
                node::SPACSimple{FT},
                flow::FT,
                par::FT,
                rad::FT,
                la::FT,
                container::SPACContainer1L{FT}
    ) where {FT<:AbstractFloat}
    leaf_gas_exchange!(
                node::SPACSimple{FT},
                flow::FT
    ) where {FT<:AbstractFloat}
    leaf_gas_exchange!(
                node::SPACSimple{FT},
                f_sl::FT,
                f_sh::FT
    ) where {FT<:AbstractFloat}

Simulate leaf level gas exchange and fill it into the `container` for 1-layer
    or 2-layer canopy, given
- `node` [`SPACSimple`] type struct
- `flow` Flow rate per basal area into the leaves (e.g., for sunlit leaves)
- `f_sl` Flow rate per basal area into the sunlit leaves
- `f_sh` Flow rate per basal area into the shaded leaves
- `par` Leaf-level photosynthetic active radiation
- `rad` Leaf-level absorbed radiative energy
- `la` Leaf area of the leaves (total or each layer)
- `container` [`SPACContainer1L`] type container
"""
function leaf_gas_exchange!(
            node::SPACSimple{FT},
            flow::FT,
            par::FT,
            rad::FT,
            la::FT,
            container::SPACContainer1L{FT}
) where {FT<:AbstractFloat}
    # 0. unpack required variables
    @unpack envir = node;
    @unpack P_AIR, p_H₂O = envir;

    # 1. calculate leaf temperature from the flow rate
    t_leaf = max(200, leaf_temperature(node, rad, flow));

    # 2. update leaf photosynthetic variables and leaf-to-air VPD
    node.ps.apar = par;
    photosystem_temperature_dependence!(node.ps.PSM, envir, t_leaf);
    node.ps.p_H₂O_sat = saturation_vapor_pressure(t_leaf);
    d_leaf = node.ps.p_H₂O_sat - p_H₂O;

    # 3. update f_vis and f_st in leaf and calculate water potentials
    # TODO do not account for temperature effects for flow now
    # More reasonable functions need to be added
    # vc_temperature_effects!(node.hs.leaf, t_leaf);
    p_leaf = end_pressure(node.hs, flow);

    # 4. calculate photosynthesis
    g_lw = flow / la / d_leaf * P_AIR;
    g_lc = max(FT(1e-6), g_lw / FT(1.6));
    leaf_photosynthesis!(node.ps, envir, GCO₂Mode(), g_lc);
    container.ag = node.ps.PSM.a_gross;
    container.an = node.ps.PSM.a_net;
    container.c  = node.ps.p_CO₂_i;
    container.e  = flow;
    container.gh = g_lw;
    container.p  = p_leaf;
    container.t  = t_leaf;

    return nothing
end

function leaf_gas_exchange!(
            node::SPACSimple{FT},
            flow::FT
) where {FT<:AbstractFloat}
    # unpack the data
    @unpack container2L, laba = node;
    @unpack frac_sh, frac_sl, par_sh, par_sl, rad_sh, rad_sl = container2L;

    # calculate mean par and rad per leaf area, then gas exchange rate
    par_mean = par_sl * frac_sl + par_sh * frac_sh;
    rad_mean = rad_sl * frac_sl + rad_sh * frac_sh;
    leaf_gas_exchange!(node, flow, par_mean, rad_mean, laba, node.container1L);

    return nothing
end

function leaf_gas_exchange!(
            node::SPACSimple{FT},
            f_sl::FT,
            f_sh::FT
) where {FT<:AbstractFloat}
    # unpack the data
    @unpack frac_sh, frac_sl, la_sh, la_sl, par_sh, par_sl, rad_sh,
            rad_sl = node.container2L;

    # calculate gas exchangr for sunlit and shaded layers
    leaf_gas_exchange!(node, f_sl, par_sl, rad_sl, la_sl, node.container2L.cont_sl);
    leaf_gas_exchange!(node, f_sh, par_sh, rad_sh, la_sh, node.container2L.cont_sh);

    return nothing
end


using POMDPs
using ContinuousPOMDPTreeSearchExperiments
using POMDPToolbox
using POMCP
using POMCPOW
using LightDarkPOMDPs
using ProgressMeter
using PmapProgressMeter
using ParticleFilters
using JLD

N = 100

@everywhere begin
    using LightDarkPOMDPs
    using POMDPs
    using ContinuousPOMDPTreeSearchExperiments
    using POMDPToolbox
    using POMCP
    using POMCPOW
    using ParticleFilters

    pomdp = LightDark2DTarget(term_radius=0.1,
                              init_dist=SymmetricNormal2([2.0, 2.0], 5.0),
                              discount=0.95)
    p_rng = MersenneTwister(4)
    solvers = Dict{String, Union{Solver,Policy}}(
        "heuristic" => begin
            homing_policy = SimpleFeedback(gain=1.0, max_radius=10.0)
            exploring_up = ObsAdaptiveParticleFilter(pomdp, LowVarianceResampler(1000), 0.05, p_rng)
            homing_up = ObsAdaptiveParticleFilter(pomdp, LowVarianceResampler(1000), 0.0, p_rng)
            InfoGatherHeur(5.0, 0.03, homing_policy, exploring_up, homing_up)
        end,

        "vanilla_pomcp" => begin
            rng3 = copy(p_rng)
            action_gen = AdaptiveRadiusRandom(max_radius=6.0, to_zero_first=true, to_light_second=true, rng=rng3)
            rollout_policy = RadiusRandom(radius=10.0, rng=rng3)
            tree_queries = 50_000
            solver = POMCPDPWSolver(next_action=action_gen,
                                    tree_queries=tree_queries,
                                    c=5.0,
                                    max_depth=40,
                                    k_action=10.0,
                                    alpha_action=1/8,
                                    k_observation=4.0,
                                    alpha_observation=1/8,
                                    estimate_value=RolloutEstimator(rollout_policy),
                                    default_action=(b, ex)->Vec2(0.1, 0.1),
                                    rng=rng3
                                   )
        end,

        "modified_pomcp" => begin
            rng3 = copy(p_rng)
            action_gen = AdaptiveRadiusRandom(max_radius=6.0, to_zero_first=true, to_light_second=true, rng=rng3)
            rollout_policy = SimpleFeedback(max_radius=10.0)
            tree_queries = 50_000
            node_updater = ObsAdaptiveParticleFilter(pomdp, LowVarianceResampler(100), 0.05, rng3)
            # node_updater = SimpleParticleFilter(pomdp, LowVarianceResampler(100), rng=rng3)
            solver = POMCPDPWSolver(next_action=action_gen,
                                    tree_queries=tree_queries,
                                    c=5.0,
                                    max_depth=40,
                                    k_action=10.0,
                                    alpha_action=1/8,
                                    k_observation=4.0,
                                    alpha_observation=1/8,
                                    estimate_value=PORolloutEstimator(rollout_policy, node_updater),
                                    node_belief_updater=node_updater,
                                    default_action=(b, ex)->Vec2(0.1, 0.1),
                                    rng=rng3
                                   )
        end,

        "pomcpow" => begin
            rng3 = copy(p_rng)
            action_gen = AdaptiveRadiusRandom(max_radius=6.0, to_zero_first=true, to_light_second=true, rng=rng3)
            tree_queries = 1_000_000
            node_updater = ObsAdaptiveParticleFilter(pomdp, LowVarianceResampler(100), 0.05, rng3)
            solver = POMCPOWSolver(next_action=action_gen,
                                    tree_queries=tree_queries,
                                    criterion=MaxUCB(5.0),
                                    final_criterion=MaxTries(),
                                    max_depth=40,
                                    k_action=10.0,
                                    alpha_action=1/8,
                                    k_observation=4.0,
                                    alpha_observation=1/8,
                                    estimate_value=OneStepValue(),
                                    node_belief_updater=node_updater,
                                    rng=rng3
                                   )
        end,

        "greedy" => SimpleFeedback(gain=1.0, max_radius=10.0)
    )

    up_rng = MersenneTwister(5)
    standard_up = ObsAdaptiveParticleFilter(pomdp, LowVarianceResampler(100_000), 0.05, up_rng)
    updaters = Dict{String, Updater}(
        "vanilla_pomcp" => standard_up,
        "modified_pomcp" => standard_up,
        "greedy" => standard_up,
        "pomcpow" => standard_up,
        "heuristic" => updater(solvers["heuristic"])
    )
end

solver_keys = keys(solvers)
rewards = Dict{String, AbstractVector{Float64}}()

for (j, sk) in enumerate(solver_keys)
    s_rewards = SharedArray(Float64, N)
    prog = Progress(N, desc="$sk ($j of $(length(solver_keys)))...")
    pmap(prog, 1:N) do i
        sim_rng = MersenneTwister(i)
        if isa(solvers[sk], Policy)
            policy = deepcopy(solvers[sk])
        else
            policy = solve(solvers[sk], pomdp)
        end
        up = deepcopy(updaters[sk])
        sim = RolloutSimulator(max_steps=40, rng=sim_rng)
        s_rewards[i] = simulate(sim, pomdp, policy, up)
    end
    rewards[sk] = sdata(s_rewards)
end

for k in solver_keys
    println("$k mean: $(mean(rewards[k])) sem: $(std(rewards[k])/sqrt(N))")
end

filename = Pkg.dir("ContinuousPOMDPTreeSearchExperiments", "data", "compare_$(Dates.format(now(), "E_d_u_HH_MM")).jld")
# @save(filename, pomdp, solver_keys, solvers, rewards, updaters)