# Copyright (c) 2021 California Institute of Technology (“Caltech”).
# U.S. Government sponsorship acknowledged.
#
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# - Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
#
# - Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# - Neither the name of Caltech nor its operating division, the Jet Propulsion
#   Laboratory, nor the names of its contributors may be used to endorse or
#   promote products derived from this software without specific prior written
#   permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
# IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
# ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

@doc raw"""
This module implements functions to create your own fabulous basis functions,
    primarily for the analysis of galaxy redshift surveys.
"""
module CryoFaBs

export CryoFaB, AngularCryoFaB, AngRadCryoFaB
export isinsurvey
export estimate_density_contrast
export angular_cryo_window
export cryobin

using LinearAlgebra
using BlockDiagonals
using Healpix
using Statistics
using SparseArrays
using SpecialFunctions
using ThreadsX

include("Splines.jl")
using .Splines

using TimerOutputs

#include("SciPy.jl")
#using .SciPy
using FastTransforms
#using SphericalHarmonics

using Scratch
cachedir_greensmatrix = ""
function __init__()
    global cachedir_greensmatrix = @get_scratch!("greensmatrix")
end

################### CryoFaB basis ######################

@doc raw"""
The abstract type `CryoFaB` defines a common interface for the bases defined in
        this module.

            δnlm = cfb \ δθϕr[:]
"""
abstract type CryoFaB end  # Create Your Own Fabulous Basis
# TODO: CryoFaB should probably be an AbstractMatrix.

struct AngularCryoFaB <: CryoFaB
    # define angular mask
    mask::HealpixMap{Float64}
    nside::Int
    npixsurvey::Int
    hpix_full2survey::Array{Int,1}  # surveypixel = hpix_full2survey[fullskypixel]

    # transform
    ell::Vector{Float64}
    Trans::Matrix{Float64}
    TransInv::Matrix{Float64}
end

struct AngRadCryoFaB <: CryoFaB  # CryoFaB that separates angular and radial basis
    # angular transform
    angfab::AngularCryoFaB
    Angl::Vector{Float64}
    AngTrans::BlockDiagonal{Float64}
    AngTransInv::BlockDiagonal{Float64}

    sortbyl::Array{Int,1}  # to transpose the input vector between angular and radial transforms

    # define radial boundary
    rmin::Float64
    rmax::Float64
    rbounds::AbstractArray{Float64,1}  # bin boundaries
    reff::AbstractArray{Float64,1}  # effective center of bin

    # radial transform
    RadK::Array{Float64,1}
    RadTrans::BlockDiagonal{Float64}
    RadTransInv::BlockDiagonal{Float64}
end

include("io.jl")
include("cryofunks_angular.jl")
include("cryofunks_radial.jl")
include("cryowindow_angular.jl")
include("cryowindow_radial.jl")
include("cryowindow_sfb.jl")

function AngularCryoFaB(mask; nside_hires=npix2nside(length(mask)), method=:cryofunk)
    if !(typeof(mask) <: HealpixMap)
        mask = HealpixMap{Float64,Healpix.RingOrder}(mask)
    end
    nside = npix2nside(length(mask))
    npix = nside2npix(nside)
    npixsurvey = sum(mask .> 0)
    hpix_full2survey = fill(0, npix)
    @show typeof(hpix_full2survey) size(hpix_full2survey) npix
    p = 1
    for i=1:npix
        if mask[i] > 0
            hpix_full2survey[i] = p
            p += 1
        end
    end
    if method == :cryofunk
        ℓ, Y, Yinv = get_angular_cryofunks(mask; nside_hires=nside_hires)
    elseif method == :integerLM
        ℓ, Y, Yinv = get_angular_continuous_cryofunks(mask; nside_hires=nside_hires)
    elseif method == :continuousfunks
        ℓ, = get_angular_cryofunks(mask; nside_hires=nside_hires)
        Y, Yinv = get_angular_continuous_cryofunks(ℓ, mask; nside_hires=nside_hires)
    end
    AngularCryoFaB(mask, nside, npixsurvey, hpix_full2survey, ℓ, Y, Yinv)
end

function block_vector(v, n)
    bv = eltype(v)[]
    for i=1:n
        append!(bv, v)
    end
    return bv
end

function AngRadCryoFaB(mask, rmin, rmax, nbins)
    # angular part
    angfab = AngularCryoFaB(mask)
    Angℓ = block_vector(angfab.ell, nbins)
    AngTrans = BlockDiagonal(fill(angfab.Trans, nbins))
    AngTransInv = BlockDiagonal(fill(angfab.TransInv, nbins))

    # radial part
    Δr = (rmax - rmin) / nbins
    reff = range(rmin + Δr/2, rmax - Δr/2, length=nbins)
    kn = []
    gn = Matrix{Float64}[]
    gninv = Matrix{Float64}[]
    @timeit "radial cryofunk" for l=1:length(angfab.ell)
        k, g, ginv = get_radial_cryofunks(angfab.ell[l], reff, Δr)
        append!(kn, k)
        push!(gn, g)
        push!(gninv, ginv)
    end
    RadK = collect(kn)
    RadTrans = BlockDiagonal(gn)
    RadTransInv = BlockDiagonal(gninv)

    #sortbyl = sortperm(Angℓ)  # No! Need to keep the order of shells for equal ℓ the same.
    sortbyl = Int[]
    for iLM=1:length(angfab.ell)
        for ik=1:nbins
            j = (ik - 1) * angfab.npixsurvey + iLM
            push!(sortbyl, j)
        end
    end

    rbounds = range(rmin, rmax, length=nbins+1)

    AngRadCryoFaB(angfab, Angℓ, AngTrans, AngTransInv, sortbyl,
        rmin, rmax, rbounds, reff, RadK, RadTrans, RadTransInv)
end

Base.size(cfb::AngularCryoFaB) = size(cfb.Trans)
Base.size(cfb::AngularCryoFaB, i) = size(cfb.Trans,i)

Base.size(cfb::AngRadCryoFaB) = size(cfb.RadTrans,1), size(cfb.AngTrans,2)
Base.size(cfb::AngRadCryoFaB, i) = (i==1) ? size(cfb.RadTrans,1) : size(cfb.AngTrans,2)

function isinsurvey(cfb::AngularCryoFaB, θ, ϕ)
    p = ang2pix(cfb.mask, θ, ϕ)
    p = cfb.hpix_full2survey[p]
    if p == 0
        return false
    end
    return true
end

function coord2cell(cfb::AngularCryoFaB, θ, ϕ)
    p = ang2pix(cfb.mask, θ, ϕ)
    p = cfb.hpix_full2survey[p]
    if p == 0
        error("angle outside of mask: (θ,ϕ)=($θ,$ϕ)")
    end
    return p
end

function cell2coord(cfb::AngularCryoFaB, p)
    error("unimplemented")
end

function mode2idx(cfb::AngularCryoFaB, l)
    error("unimplemented")
end

function idx2mode(cfb::AngularCryoFaB, i)
    return cfb.ell[i]
end

function isinsurvey(cfb::AngRadCryoFaB, r, θ, ϕ)
    if !isinsurvey(cfb.angfab, θ, ϕ)
        return false
    end

    rbounds = get_rbounds(cfb)
    i = (r == cfb.rmax) ? length(cfb.reff) : searchsortedlast(rbounds, r)
    if !(1 <= i <= length(cfb.reff))
        return false
    end
    return true
end

function coord2cell(cfb::AngRadCryoFaB, r, θ, ϕ)
    p = coord2cell(cfb.angfab, θ, ϕ)

    rbounds = get_rbounds(cfb)
    i = (r == cfb.rmax) ? length(cfb.reff) : searchsortedlast(rbounds, r)
    if !(1 <= i <= length(cfb.reff))
        error("radial dist outside of survey: (r,θ,ϕ)=($r,$θ,$ϕ), $i, $(length(cfb.reff))")
    end

    v = p + (i - 1) * cfb.angfab.npixsurvey
    if v > size(cfb,2)
        @error r,θ,ϕ p i v size(cfb) cfb.angfab.npixsurvey
    end
    return v
end

function cell2coord(cfb::AngRadCryoFaB, v)
    error("unimplemented")
end

function idx2mode(cfb::AngRadCryoFaB, i)
    return cfb.RadK[i], cfb.Angl[cfb.sortbyl][i]
end

function mode2idx(cfb::AngRadCryoFaB, k, l)
    error("unimplemented")
end

Base.:\(cfb::AngularCryoFaB, v) = cfb.TransInv * v
Base.:*(cfb::AngularCryoFaB, v) = cfb.Trans * v
Base.adjoint(cfb::AngularCryoFaB) = cfb.TransInv'
Base.inv(cfb::AngularCryoFaB) = cfb.TransInv

Base.:\(cfb::AngRadCryoFaB, v::AbstractVector) = begin
    cfb.RadTransInv * (cfb.AngTransInv * v)[cfb.sortbyl]
end

Base.:*(cfb::AngRadCryoFaB, v::AbstractVector) = begin
    error("unimplemented")
end

function get_pixs_at_r(cfb::AngRadCryoFaB, i)
    return (1:cfb.angfab.npixsurvey) .+ (i-1) * cfb.angfab.npixsurvey
end

function get_rbounds(cfb::AngRadCryoFaB)
    range(cfb.rmin, cfb.rmax, length=length(cfb.reff)+1)
end

################### estimate density contrast ##########

@doc raw"""
calc_numgals_per_cell(rθϕ, cfb::CryoFaB)

Calculate the number of galaxies in `rθϕ` in each pixel or voxel in the survey
that `cfb` can act on. The radial coordinates are in `rθϕ[1,:]`, etc.
"""
function calc_numgals_per_cell_v1(rθϕ, cfb::CryoFaB)
    T = Float64  # need to allow floats for easy conversion to a density contrast
    N = fill(T(0), size(cfb,1))
    for i=1:size(rθϕ,2)
        p = coord2cell(cfb, rθϕ[:,i]...)
        #p = coord2cell(cfb, rθϕ[1,i], rθϕ[2,i], rθϕ[3,i])
        N[p] += 1
    end
    return N
end
function calc_numgals_per_cell(rθϕ, cfb::AngRadCryoFaB)
    T = Float64  # need to allow floats for easy conversion to a density contrast
    N = fill(T(0), size(cfb,1))
    #@time pix1 = coord2cell.(Ref(cfb), rθϕ[1,:], rθϕ[2,:], rθϕ[3,:])
    #@time pix2 = map((x,y,z)->coord2cell(cfb, x, y, z), rθϕ[1,:], rθϕ[2,:], rθϕ[3,:])
    @time pix3 = ThreadsX.map((r,θ,ϕ)->coord2cell(cfb, r, θ, ϕ), rθϕ[1,:], rθϕ[2,:], rθϕ[3,:])
    #@time pix4 = @strided coord2cell.(Ref(cfb), rθϕ[1,:], rθϕ[2,:], rθϕ[3,:])
    #@assert pix1 == pix2
    #@assert pix1 == pix3
    #@assert pix1 == pix4
    @time for p in pix3
        N[p] += 1
    end
    return N
end

function numgals2densitycontrast!(δ, cfb::AngRadCryoFaB)
    rbounds = get_rbounds(cfb)
    for i=2:length(rbounds)
        p = get_pixs_at_r(cfb, i-1)
        ntotz = sum(δ[p])
        nofz = ntotz / length(p)  # average number of galaxies per pixel
        @. δ[p] = δ[p] / nofz - 1
    end
    return δ
end

function numgals2densitycontrast(δ, cfb::AngRadCryoFaB)
    return numgals2densitycontrast!(deepcopy(δ), cfb)
end

@doc raw"""
estimate_density_contrast(rθϕ, cfb::CryoFaB)

Estimate the density contrast from the galaxy positions in `rθϕ`.
"""
function estimate_density_contrast(rθϕ, cfb::AngRadCryoFaB)
    δ = calc_numgals_per_cell(rθϕ, cfb)
    numgals2densitycontrast!(δ, cfb)
    return δ
end

function calc_angular_average(nr, cfb::CryoFaB)
    rbounds = get_rbounds(cfb)
    nbar = fill(float(eltype(nr)(NaN)), length(rbounds) - 1)
    for i=1:length(rbounds)-1
        p = get_pixs_at_r(cfb, i)
        nbar[i] = mean(nr[p])
    end
    return nbar
end

function cryobin(P, cfb::AngRadCryoFaB; ltol=0.2, ktol=minimum(cfb.RadK), ignore_l=[])
    kk, ll = idx2mode(cfb, 1:length(P))

    # ignore ℓ=0 and ℓ>2*nside
    s = fill(true, length(ll))
    sl = @. cfb.angfab.ell > 1.5 * cfb.angfab.nside
    ignore_l = [ignore_l; cfb.angfab.ell[sl]]
    for l in ignore_l
        @. s &= (ll != l)
    end

    # ignore large k
    Δr = (cfb.rmax - cfb.rmin) / length(cfb.reff)
    @. s &= (kk < 3 / Δr)

    return cryobin(P[s], kk[s], ll[s]; ltol=ltol, ktol=ktol)
end

function cryobin(P, kk, ll; ltol=0.5, ktol=minimum(abs,kk))
    kbin = float(eltype(kk))[]
    lbin = float(eltype(ll))[]
    Pbin = float(eltype(P))[]
    Nbin = Int[]
    while length(kk) >= 1
        k = kk[1]
        l = ll[1]
        sel = @. (abs(kk - k) <= ktol) & (abs(ll - l) <= ltol)
        #@show length(sel),sum(sel)
        #@show k,l kk kk[sel] ll[sel] P[sel]
        sel[1] = true  # include at least one in case 'sel' is false
        push!(kbin, mean(kk[sel]))
        push!(lbin, mean(ll[sel]))
        push!(Pbin, mean(P[sel]))
        push!(Nbin, sum(sel))
        kk = kk[.!sel]
        ll = ll[.!sel]
        P = P[.!sel]
    end
    return kbin, lbin, Pbin, Nbin
end

end

# vim: set sw=4 et sts=4 :


import OnlineStats: OnlineStat, Extrema, fit!, value
import ...Cubes.Axes: CategoricalAxis, RangeAxis
import IterTools
using WeightedOnlineStats
import ProgressMeter: next!

import WeightedOnlineStats: WeightedOnlineStat
abstract type TableAggregator end
struct OnlineAggregator{O,S}<:TableAggregator
    o::O
end
function OnlineAggregator(O::OnlineStat,s::Symbol)
    OnlineAggregator{typeof(O),s}(copy(O))
end
cubeeltype(t::OnlineAggregator)=Float64
function fitrow!(o::OnlineAggregator{T,S},r) where {T<:OnlineStat,S}
    v = getproperty(r,S)
    !ismissing(v) && fit!(o.o,v)
end
value(o::OnlineAggregator)=value(o.o)
struct WeightOnlineAggregator{O,S,W}<:TableAggregator
    o::O
    w::W
end
function WeightOnlineAggregator(O::WeightedOnlineStat,s::Symbol,w)
    WeightOnlineAggregator{typeof(O),s,typeof(w)}(copy(O),w)
end
cubeeltype(t::WeightOnlineAggregator{T}) where T = cubeeltype(T)
value(o::WeightOnlineAggregator)=value(o.o)
function fitrow!(o::WeightOnlineAggregator{T,S},r) where {T<:OnlineStat,S}
    v = getproperty(r,S)
    w = o.w(r)
    if !checkmiss(v) && !ismissing(w)
      fit!(o.o,v,w)
    end
end
checkmiss(v) = ismissing(v)
checkmiss(v::AbstractVector) = any(ismissing,v)
struct GroupedOnlineAggregator{O,S,BY,W,C}<:TableAggregator
    d::O
    w::W
    by::BY
    cloneobj::C
end
value(o::GroupedOnlineAggregator)=Dict(zip(keys(o.d),map(value,(values(o.d)))))
struct SymType{S}
end
SymType(s::Symbol)=SymType{s}()
(f::SymType{S})(x) where S=getproperty(x,S)
getbytypes(et,by) = Tuple{map(i->Base.nonmissingtype(Base.return_types(i,Tuple{et})[1]),by)...}
cubeeltype(t::GroupedOnlineAggregator{T}) where T=cubeeltype(T)
cubeeltype(t::Type{<:Dict{<:Any,T}}) where T = cubeeltype(T)
cubeeltype(t::Type{<:WeightedOnlineStat{T}}) where T = T
cubeeltype(t::Type{<:OnlineStat{Number}}) = Float64
cubeeltype(t::Type{<:WeightedCovMatrix{T}}) where T = T
cubeeltype(t::Type{<:Extrema{T}}) where T = T

function GroupedOnlineAggregator(O::OnlineStat,s::Symbol,by,w,iter)
    ost = typeof(O)
    et = eltype(iter)
    bytypes = Tuple{map(i->Base.nonmissingtype(Base.return_types(i,Tuple{et})[1]),by)...}
    d = Dict{bytypes,ost}()
    GroupedOnlineAggregator{typeof(d),s,typeof(by),typeof(w),ost}(d,w,by,O)
end

dicteltype(::Type{<:Dict{K,V}}) where {K,V} = V
dictktype(::Type{<:Dict{K,V}}) where {K,V} = K
actval(v) = v
actval(v::SentinelMissings.SentinelMissing) = v[]
function fitrow!(o::GroupedOnlineAggregator{T,S,BY,W},r) where {T,S,BY,W}
    v = getproperty(r,S)
    if !ismissing(v)
        w = o.w(r)
        if w===nothing
            bykey = map(i->actval(i(r)),o.by)
            if !any(ismissing,bykey)
                if haskey(o.d,bykey)
                    fit!(o.d[bykey],v)
                else
                    o.d[bykey] = copy(o.cloneobj)
                    fit!(o.d[bykey],v)
                end
            end
        else
           if !ismissing(w)
                bykey = map(i->actval(i(r)),o.by)
                if !any(ismissing,bykey)
                    if haskey(o.d,bykey)
                        fit!(o.d[bykey],v,w)
                    else
                        o.d[bykey] = copy(o.cloneobj)
                        fit!(o.d[bykey],v,w)
                    end
                end
            end
        end
    end
end
export TableAggregator, fittable, cubefittable
function TableAggregator(iter,O,fitsym;by=(),weight=nothing)
  !isa(by,Tuple) && (by=(by,))
  if !isempty(by)
    weight===nothing && (weight=(i->nothing))
    by = map(i->isa(i,Symbol) ? (SymType(i)) : i,by)
    GroupedOnlineAggregator(O,fitsym,by,weight,iter)
  else
    if weight===nothing
      OnlineAggregator(O,fitsym)
    else
      WeightOnlineAggregator(O,fitsym,weight)
    end
  end
end

function tooutaxis(::SymType{s},iter::CubeIterator{<:Any,<:Any,<:Any,<:Any,<:Any,S},k,ibc) where {s,S}
    ichosen = findfirst(i->i===s,fieldnames(S))
    if ichosen<=length(iter.inars)
        bycube = iter.dc.incubes[ichosen].cube
        if haskey(bycube.properties,"labels")
            idict=bycube.properties["labels"]
            axname=get(bycube.properties,"name","Label")
            outAxis=CategoricalAxis(axname,collect(String,values(idict)))
            convertdict=Dict(k=>i for (i,k) in enumerate(keys(idict)))
        else
            sort!(k)
            outAxis=CategoricalAxis("Label$(ibc)",k)
            convertdict=Dict(k=>i for (i,k) in enumerate(k))
        end
    else
       iax = findAxis(string(s),iter.dc.LoopAxes)
        outAxis=iter.dc.LoopAxes[iax]
        convertdict=Dict(k=>i for (i,k) in enumerate(outAxis.values))
    end
    outAxis,convertdict
end
function tooutaxis(f,iter::CubeIterator{<:Any,<:Any,<:Any,<:Any,<:Any,S},k,ibc) where S
    sort!(k)
    outAxis=CategoricalAxis("Category$(ibc)",k)
    convertdict=Dict(k=>i for (i,k) in enumerate(k))
    outAxis,convertdict
end

varsym(::WeightOnlineAggregator{<:Any,S}) where S = S
varsym(::OnlineAggregator{<:Any,S}) where S = S
varsym(::GroupedOnlineAggregator{<:Any,S}) where S = S
axt(::CategoricalAxis)=CategoricalAxis
axt(::RangeAxis)=RangeAxis
getStatType(::WeightOnlineAggregator{T}) where T = T
getStatType(::OnlineAggregator{T}) where T = T
getStatType(t::GroupedOnlineAggregator{T}) where T=getStatType(T)
getStatType(t::Type{<:Dict{<:Any,T}}) where T = T

getStatOutAxes(tab,agg)=getStatOutAxes(tab,agg,getStatType(agg))
getStatOutAxes(tab,agg,::Type{<:OnlineStat})=()
function getStatOutAxes(tab, agg, ::Type{<:Extrema})
    ( CategoricalAxis(:Extrema, ["min", "max"]), )
end
function getStatOutAxes(tab,agg,::Type{<:WeightedCovMatrix})
    varn = fieldnames(eltype(tab))
    s = varsym(agg)
    icube = findfirst(isequal(s),varn)
    ax = tab.dc.incubes[icube].axesSmall[1]
    oldname = YAXArrays.Cubes.Axes.axname(ax)
    coname = string("Co",oldname)
    v = ax.values
    axtype = axt(ax)
    a1 = axtype(oldname,copy(v))
    a2 = axtype(coname,copy(v))
    (a1,a2)
end
function getStatOutAxes(tab,agg,::Type{<:WeightedAdaptiveHist})
    nbin = getnbins(agg)
    a1 = RangeAxis("Bin",1:nbin)
    a2 = CategoricalAxis("Hist",["MidPoints","Frequency"])
    (a1,a2)
end
function getByAxes(iter,agg::GroupedOnlineAggregator)
    by = agg.by
    ntuple(length(by)) do ibc
      bc=agg.by[ibc]
      tooutaxis(bc,iter,unique(map(i->i[ibc],collect(keys(agg.d)))),ibc)
    end
end
getByAxes(iter,agg)=()

function tooutcube(
    agg,
    iter,
    post
  )
  outaxby = getByAxes(iter,agg)
  axby = map(i->i[1],outaxby)
  convdictall = map(i->i[2],outaxby)

  outaxstat = getStatOutAxes(iter,agg)
  outax = (outaxstat...,axby...)
  snew = map(length,outax)
  aout = fill!(zeros(Union{cubeeltype(agg),Missing},snew),missing)
  filloutar(aout,convdictall,agg,map(i->1:length(i),outaxstat),post)

  YAXArray(collect(outax),aout)
end
function filloutar(aout,convdictall,agg::GroupedOnlineAggregator,s,post)
    for (k,v) in agg.d
        i = CartesianIndices((s...,map((i,d)->d[convert(keytype(d),i)]:d[convert(keytype(d),i)],k,convdictall)...))
        aout[i.indices...].=post(v)
    end
end
function filloutar(aout,convdictall,agg,g,post)
    copyto!(aout,post(agg.o))
end
"""
    fittable(tab,o,fitsym;by=(),weight=nothing)

Loops through an iterable table `tab` and thereby fitting an OnlineStat `o` with the values
specified through `fitsym`. Optionally one can specify a field (or tuple) to group by.
Any groupby specifier can either be a symbol denoting the entry to group by or an anynymous
function calculating the group from a table row.

For example the following would caluclate a weighted mean over a cube weighted by grid cell
area and grouped by country and month:

````julia
fittable(iter,WeightedMean,:tair,weight=(i->abs(cosd(i.lat))),by=(i->month(i.time),:country))
````
"""
function fittable(tab,o,fitsym;by=(),weight=nothing,showprog=false)
  agg = TableAggregator(tab,o,fitsym,by=by,weight=weight)
  if showprog
    runfitrows_progress(agg,tab)
  else
    foreach(i->fitrow!(agg,i),tab)
  end
  agg
end
fittable(tab,o::Type{<:OnlineStat},fitsym;kwargs...)=fittable(tab,o(),fitsym;kwargs...)

@noinline function runfitrows_progress(agg,tab)
  p = Progress(length(tab)÷100,1)
  every = 0
  for row in tab
    fitrow!(agg,row)
    every += 1
    if every == 100
      next!(p)
      every=0
    end
  end
end

struct collectedValue{V,S,SY}
    value::V
    laststruct::S
end

function Base.getproperty(s::collectedValue{<:Any,<:Any,SY},sy::Symbol) where SY
    if sy == SY
        getfield(s,:value)
    else
        getproperty(getfield(s,:laststruct),sy)
    end
end

function collectval(row::Union{Tuple, Vector},::Val{SY}) where SY
    nvars = length(row)
    v = ntuple(i->getfield(row[i],SY),nvars) |> collect
    val = collectedValue{typeof(v),typeof(row[end]),SY}(v, row[end])
end

getpostfunction(s::OnlineStat)=getpostfunction(typeof(s))
getpostfunction(::Type{<:OnlineStat})=value
function getpostfunction(w::WeightedAdaptiveHist)
  nb = w.alg.b
  i->begin
    r = hcat(value(i)...)
    if size(r,1)<nb
      r = vcat(r,zeros(eltype(r),nb-size(r,1),2))
    end
    r
end
end
getnbins(f::GroupedOnlineAggregator)=f.cloneobj.alg.b
getnbins(f::TableAggregator)=f.o.alg.b

fitfun(o) = fitfun(typeof(o))
fitfun(::Type{<:Any}) = fittable

"""
    cubefittable(tab,o,fitsym;post=getpostfunction(o),kwargs...)

Executes [`fittable`](@ref) on the [`@CubeTable`](@ref) `tab` with the
(Weighted-)OnlineStat `o`, looping through the values specified by `fitsym`.
Finally, writes the results from the `TableAggregator` to an output data cube.
"""
function cubefittable(tab,o,fitsym;post=getpostfunction(o),kwargs...)
  agg=fitfun(o)(tab,o,fitsym;showprog=true,kwargs...)
  tooutcube(agg,tab,post)
end

function tupleeltypebyname(::Type{NamedTuple{names,tt}},s::Symbol) where {names, tt}
    i=findfirst(isequal(s),names)
    fieldtype(tt,i)
end


abstract type ModelSelectionInput end

struct SimulatedModelSelectionInput{AF<:AbstractFloat, CUD<:ContinuousUnivariateDistribution} <: ModelSelectionInput
	M::Int
	n_particles::Int
	threshold_schedule::AbstractArray{AF,1}
	model_prior::DiscreteUnivariateDistribution
	parameter_priors::AbstractArray{Array{CUD, 1},1}
	distance_simulation_input::AbstractArray{DistanceSimulationInput,1}
	max_iter::Int
end

struct EmulatedModelSelectionInput{AF<:AbstractFloat, CUD<:ContinuousUnivariateDistribution,
	    ER<:AbstractEmulatorRetraining, EPS<:AbstractEmulatedParticleSelection, ET<:AbstractEmulatorTraining} <: ModelSelectionInput
	M::Int
	n_particles::Int
	threshold_schedule::AbstractArray{AF,1}
	model_prior::DiscreteUnivariateDistribution
	parameter_priors::AbstractArray{Array{CUD,1},1}
	emulator_training_input::AbstractArray{EmulatorTrainingInput{ET},1}
	emulator_retraining::ER
	emulated_particle_selection::EPS
	max_batch_size::Int
	max_iter::Int
end

mutable struct CandidateModelTracker{AF<:AbstractFloat}
	n_accepted::Int
	n_tries::Int
	population::AbstractArray{AF,2}
	distances::AbstractArray{AF,1}
	weight_values::AbstractArray{AF,1}
end

function CandidateModelTracker(n_params::Int)
	return CandidateModelTracker(0, 0, zeros(0, n_params), zeros(0), zeros(0))
end

abstract type ModelSelectionTracker end

mutable struct SimulatedModelSelectionTracker <: ModelSelectionTracker
	M::Int64
	n_particles::Int64
	threshold_schedule::AbstractArray{Float64,1}
	model_prior::DiscreteUnivariateDistribution
	smc_trackers::AbstractArray{SimulatedABCSMCTracker,1}
	max_iter::Int
end

mutable struct EmulatedModelSelectionTracker{CUD<:ContinuousUnivariateDistribution, ET,
        ER<:AbstractEmulatorRetraining, EPS<:AbstractEmulatedParticleSelection} <: ModelSelectionTracker
	M::Int64
	n_particles::Int64
	threshold_schedule::AbstractArray{Float64,1}
	model_prior::DiscreteUnivariateDistribution
	smc_trackers::AbstractArray{EmulatedABCSMCTracker{CUD, ET, ER, EPS},1}
	max_batch_size::Int64
	max_iter::Int
end

"""
	ModelSelectionOutput

Contains results of a model selection computation, including which models are best supported by the observed data and the parameter poseteriors at each population for each model.

# Fields
- `M::Int64`: The number of models.
- `n_accepted::AbstractArray{AbstractArray{Int64,1},1}`: The number of parameters accepted by each model at each population. `n_accepted[i][j]` contains the number of acceptances for model `j` at population `i`.
- `threshold_schedule::AbstractArray{Float64,1}`: A set of maximum distances from the summarised model output to summarised observed data for a parameter vector to be included in the posterior.
- `smc_outputs::AbstractArray{ABCSMCOutput,1}`: A ['SimulatedABCSMCOutput']@(ref) or ['EmulatedABCSMCOutput']@(ref) for each model. Use to find details of the ABC results at each population.
- `completed_all_populations::Bool`: Indicates whether the algorithm completed all the populations successfully. A successful population is one where at least one model accepts at least one particle.
"""
struct ModelSelectionOutput
	M::Int64
	n_accepted::AbstractArray{AbstractArray{Int64,1},1}
	threshold_schedule::AbstractArray{Float64,1}
	smc_outputs::AbstractArray{ABCSMCOutput,1}
	completed_all_populations::Bool
end


###########
# Imports #
###########

# Julia Parallel Packages
using Distributed

# Julia Packages
using Dates
using Random
using Profile
using Printf
using Logging
using Statistics
using Distributions

# Register Number of processes
# Julia must be started with `julia -p auto -e "using IJulia; jupyterlab();"`
# On Computer with 6 cores add 6-1 = 5 processes
# NOTE: This must be called _before_ @everywhere is used those functions get propery distributed
if length(procs()) == 1
    # Remove existing process IDs so reloading works
    addprocs(length(Sys.cpu_info()));
    println("Started $(nprocs()) total processes - $(nworkers()) worker processes.")
else
    println("Found $(nprocs()) total processes - $(nworkers()) worker processes.")
end

# Satellite Dynamics Packages
@everywhere using SatelliteDynamics

# Load SatelliteTasking
@everywhere using SatelliteTasking

# Add Map Plots
using PGFMapPlots

#########################
# Simulation Definition #
#########################

# Set simulation time window
t_start = Epoch(2019, 1, 1, 0, 0, 0, tsys="UTC"); # Start of time span
t_end = Epoch(2019, 1, 2, 0, 0, 0, tsys="UTC"); # End of simulation time span

# Initialize Planning Problem
problem = SatPlanningProblem(t_start=t_start, t_end=t_end);

# Add constraints to problem
push!(problem.constraints, constraint_agility_single_axis)

# Create Spacecraft
line1 = "1     1U          19  1.00000000  .00000000  00000-0  00000-0 0    05"
line2 = "2     1  90.0000   0.0000 0010000   0.0000   0.0000 15.24162312    00"
sc1 = Spacecraft(id=1, tle=TLE(line1, line2), slew_rate=1.0);
push!(problem.spacecraft, sc1);

# Compute Orital Period
T = orbit_period(state(problem.spacecraft[1].tle, problem.spacecraft[1].tle.epoch)[1]);

println("Loadded $(length(problem.spacecraft)) spacecraft")

# Set Solve Parameters
problem.solve_discount = 0.99 # Typical values 0.99 - 0.9999
problem.solve_depth = 3
problem.solve_breadth = 3
problem.solve_horizon = T

println("Default Solve Parameters:\nsolve_discount: $(problem.solve_discount)\nsolve_depth: $(problem.solve_depth)\nsolve_breadth: $(problem.solve_breadth)\nsolve_horizon: $(problem.solve_horizon)")

######################
# Load Location Data #
######################

# Load Groundstation Data
stations = load_stations("./data/groundstations.json");
num_stations = length(stations);

# Load Image Data - [100, 200, 500, 1000, 2000, 5000, 10000, 12000, 15000]
num_requests = 1000
requests = load_requests("./data/landsat_test_$num_requests.json", id_offset=num_stations);
@assert num_requests == length(requests);

# Set Request Parameters
collect_duration = 30.0
set_collect_duration(requests, collect_duration)

# Add Locations To Problem
# add_locations(problem, stations)
add_locations(problem, requests)

#########################
# Compute Opportunities #
#########################

# Compute problem accesses
compute_access(problem, orbit_fraction=0.75);

println("")
println("Found $(length(problem.contacts)) passes in planning horizon")
println("Found $(length(problem.collects)) collects in planning horizon")
num_accessible_requests = sum([1 for req in problem.requests if length(problem.lt_loc_opps[req.id]) > 0])
pct_accessible_requests = floor(num_accessible_requests/length(problem.requests) * 10000)/100
println("Found $num_accessible_requests out of $(length(problem.requests)) requests have accesses. $pct_accessible_requests%")

# MDP Planning Preparation - (Pre-compute Actions)
println("Precomputing action space")
precompute_action_space(problem, enable_resources=false);

########################
# Define Work Function #
########################

@everywhere function run_mcts_sim(problem::SatPlanningProblem, n::Int, params::Tuple{Real, Real, Real, Real})
    println("Reporting in!")
    # Create reward storage
    hyper_time = Dict{Tuple{Real, Real, Real, Real}, AbstractVector{Real}}()
    hyper_reward = Dict{Tuple{Real, Real, Real, Real}, AbstractVector{Real}}()

    # Extract search point from params
    ec, ni, sd, ds = params

    # Copy problem
    # problem = copy(problem)
    problem.solve_discount = ds # Typical values 0.99 - 0.9999
    problem.solve_depth = sd
    problem.mcts_n_iterations = ni
    problem.mcts_exploration_constant = ec

    hyper_time[(ec, ni, sd, ds)]   = Real[]
    hyper_reward[(ec, ni, sd, ds)] = Real[]

    # Run specified number of iterations
    for idx in 1:n
        start_time = time()

        # Solve problem
        println("Iteration $idx - Solving Problem with parameters:($(problem.solve_discount), $(problem.solve_depth), $(problem.mcts_n_iterations), $(problem.mcts_exploration_constant))")
        mdp_mcts_states, mdp_mcts_plan, mdp_mcts_reward = satellite_plan_mdp_mcts(problem, parallel=false)
        end_time = time()

        # Extract reward
        mdp_mcts_feasible, mdp_mcts_realized_reward, mdp_mcts_nur, mdp_mcts_ndr, mdp_mcts_nr, mdp_mcts_nc = analyze_plan(problem, mdp_mcts_plan);

        # Update serach values
        push!(hyper_time[(ec, ni, sd, ds)], end_time-start_time)
        push!(hyper_reward[(ec, ni, sd, ds)], mdp_mcts_realized_reward)
    end

    return hyper_time, hyper_reward
end

######################################
# Define Hyperparameter Search Space #
######################################

# Define parameter space
ec_space = [0.01, 0.1, 1.0, 1.5, 2.0, 3.0, 5.0, 10.0]
ni_space = [10, 20, 30, 50, 100, 200, 500, 1000]
sd_space = [2, 3, 5, 10, 20, 30, 50] # solve depth
ds_space = [0.99, 0.995, 0.999] # discount
bin_iterations = 10

ec_space = [1.5, 2.0, 3.0]
ni_space = [10, 20]
sd_space = [10, 20] # solve depth
ds_space = [0.99, 0.995] # discount
bin_iterations = 3

# ec_space = [0.01, 0.1, 0.2]
# ni_space = [5, 10]
# sd_space = [2, 3, 5]
# ds_space = [0.99]
# bin_iterations = 10

# Create reward storage
hyper_time = Dict{Tuple{Real, Real, Real, Real}, AbstractVector{Real}}()
hyper_reward = Dict{Tuple{Real, Real, Real, Real}, AbstractVector{Real}}()

println("Total number of required simulations: $(length(ec_space)*length(ni_space)*length(sd_space)*length(ds_space)*bin_iterations)")
num_jobs = length(ec_space)*length(ni_space)*length(sd_space)*length(ds_space)
println("Total number of jobs: $num_jobs")

#######################
# Run Parallel Search #
#######################

fn(p, n) = x -> run_mcts_sim(p, n, x)

# Create work assignments
nw = nworkers()        # Number of workers
lw = num_jobs # Length of work

wa = floor(Int, lw/nw) # Average work per worker
wr = lw - nw*wa        # Remaining worker

println("Have $lw items of work and $nw workers")
println("Average work: $wa, Remaining work: $wr")

# Create search points
work = []
for ec in ec_space
    for ni in ni_space
        for sd in sd_space
            for ds in ds_space
                push!(work, (ec, ni, sd, ds))
            end
        end
    end
end

println("Work: $work")
# Construct assignments (function inputs) for Workers
# assignments = []
# println("Assignments: $assignments")

# push!(assignments, work[1:(wa+wr)])
# for i in 2:min(nw, length(work))
#     # @debug "Worker $i convering $(1+wr+(i-1)*wa):$(i*wa+wr)")
#     push!(assignments, work[(1+wr+(i-1)*wa):(i*wa+wr)])
# end

# Execute Work in parallel

println("Starting hyper parameter search")
hp_start = time()
results = pmap(fn(problem, bin_iterations), work)
# results = pmap(fn(problem, bin_iterations), assignments)
hp_end = time()
println("Parameter search finished after $(hp_end-hp_start) seconds.")

# Merge results
hyper_time = Dict{Tuple{Real, Real, Real, Real}, AbstractVector{Real}}()
hyper_reward = Dict{Tuple{Real, Real, Real, Real}, AbstractVector{Real}}()

for r in results
    ht, hr = r
    for params, time_values in ht
        hyper_time[params] = time_values
        hyper_reward[params] = hr[params]
    end
end

#################################
# Extract Best Hyper Parameters #
#################################

# Extract Maximum reward
max_max = -Inf
max_mean = -Inf
max_min = - Inf
max_hp = (-Inf, -Inf, -Inf, -Inf)

for (hparams, rvec) in hyper_reward
    minval, maxval = extrema(rvec)
    if length(rvec) > 0 && maxval > max_max
        max_max = maxval
        max_mean = mean(rvec)
        max_min = minval
        max_hp = hparams
    end
end

# Print Best Point
println("Best Hyperparameter Cominbation -- Min: $(max_min) - Mean: $(max_mean) - Max: $(max_max)")
println("Exploration Constant: $(max_hp[1])")
println("Number of Iterations: $(max_hp[2])")
println("Solve Depth: $(max_hp[3])")
println("Discount Facgtor: $(max_hp[4])")

# Save Rest to File
open("hyperparam_search_rewards_$(Dates.format(now(), "yyyy-mm-ddTHH:MM:SS")).dat") do file
    column_string = join(["Reward $idx" for idx in 1:bin_iterations], ",")
    write(file, "Exploration Constant, Number Of Iterations, Search Depth, Discount Factor, $column_string\n")
    for hparams, rewards in hyper_reward
        ec, ni, sd, ds = hparams
        reward_string = join(rewards, ",")
        write(file, "$(ec),$ni,$sd,$ds,$reward_string\n")
    end
end

open("hyperparam_search_times_$(Dates.format(now(), "yyyy-mm-ddTHH:MM:SS")).dat") do file
    column_string = join(["Runtime $idx" for idx in 1:bin_iterations], ",")
    write(file, "Exploration Constant, Number Of Iterations, Search Depth, Discount Factor, $column_string\n")
    for hparams, times in hyper_time
        ec, ni, sd, ds = hparams
        time_string = join(times, ",")
        write(file, "$(ec),$ni,$sd,$ds,$time_string\n")
    end
end


function lazy_lb(
    psi_container::PSIContainer,
    ts_data::Vector{DeviceTimeSeries},
    cons_name::Symbol,
    var_name::Symbol,
)
    time_steps = model_time_steps(psi_container)
    names = (d.name for d in ts_data)
    variable = get_variable(psi_container, var_name)
    lb_name = middle_rename(cons_name, PSI_NAME_DELIMITER, "lb")
    con_lb = add_cons_container!(psi_container, lb_name, names, time_steps)

    for data in ts_data
        if data.range.limits.min > -Inf && !isempty(data.range.additional_terms_lb)
            con_lb[data.name, :].data .= JuMP.AffExpr(0.0)
            continue
        end
        for t in time_steps
            expression_lb = JuMP.AffExpr(0.0, variable[data.name, t] => 1.0)
            for val in data.range.additional_terms_lb
                JuMP.add_to_expression!(
                    expression_lb,
                    get_variable(psi_container, val)[data.name, t],
                    -1.0,
                )
            end
            lb_val = max(0.0, data.range.limits.min)
            con_lb[data.name, t] =
                JuMP.@constraint(psi_container.JuMPmodel, expression_lb >= lb_val)
        end
    end
    return
end

@doc raw"""
    device_timeseries_ub(psi_container::PSIContainer,
                     ts_data::Vector{DeviceTimeSeries},
                     cons_name::Symbol,
                     var_name::Symbol)

Constructs upper bound for given variable and time series data and a multiplier.

# Constraint

```variable[name, t] <= ts_data[name].multiplier * ts_data[name].timeseries[t] ```

# LaTeX

`` x_t \leq r^{val} r_t, \forall t ``

# Arguments
* psi_container::PSIContainer : the psi_container model built in PowerSimulations
* ts_data::Vector{DeviceTimeSeries} : container of device time series data and scaling factors
* cons_name::Symbol : name of the constraint
* var_name::Symbol : the name of the variable
"""
function device_timeseries_ub(
    psi_container::PSIContainer,
    ts_data::Vector{DeviceTimeSeries},
    cons_name::Symbol,
    var_name::Symbol,
)
    time_steps = model_time_steps(psi_container)
    names = (d.name for d in ts_data)
    variable = get_variable(psi_container, var_name)
    ub_name = middle_rename(cons_name, PSI_NAME_DELIMITER, "ub")
    con_ub = add_cons_container!(psi_container, ub_name, names, time_steps)
    lazy_add_lb = false

    for data in ts_data
        for t in time_steps
            expression_ub = JuMP.AffExpr(0.0, variable[data.name, t] => 1.0)
            for val in data.range.additional_terms_ub
                JuMP.add_to_expression!(
                    expression_ub,
                    get_variable(psi_container, val)[data.name, t],
                )
            end
            con_ub[data.name, t] = JuMP.@constraint(
                psi_container.JuMPmodel,
                expression_ub <= data.multiplier * data.timeseries[t]
            )
        end
        if data.range.limits.min > -Inf || !isempty(data.range.additional_terms_lb)
            lazy_add_lb = true
        end
    end
    @debug lazy_add_lb
    lazy_add_lb && lazy_lb(psi_container, ts_data, cons_name, var_name)
    return
end

@doc raw"""
    device_timeseries_lb(psi_container::PSIContainer,
                     ts_data::Vector{DeviceTimeSeries},
                     cons_name::Symbol,
                     var_name::Symbol)

Constructs lower bound for given variable subject to time series data and a multiplier.

# Constraint

``` ts_data[name].multiplier * ts_data[name].timeseries[t] <= variable[name, t] ```

# LaTeX

`` r^{val} r_t \leq x_t, \forall t ``

where (name, data) in range_data.

# Arguments
* psi_container::PSIContainer : the psi_container model built in PowerSimulations
* ts_data::Vector{DeviceTimeSeries} : container of device time series data and scaling factors
* cons_name::Symbol : name of the constraint
* var_name::Symbol : the name of the variable
"""
function device_timeseries_lb(
    psi_container::PSIContainer,
    ts_data::Vector{DeviceTimeSeries},
    cons_name::Symbol,
    var_name::Symbol,
)
    time_steps = model_time_steps(psi_container)
    variable = get_variable(psi_container, var_name)
    lb_name = middle_rename(cons_name, PSI_NAME_DELIMITER, "lb")
    names = (d.name for d in ts_data)
    constraint = add_cons_container!(psi_container, lb_name, names, time_steps)

    for data in ts_data
        for t in time_steps
            expression_lb = JuMP.AffExpr(0.0, variable[data.name, t] => 1.0)
            for val in data.range.additional_terms_lb
                JuMP.add_to_expression!(
                    expression_lb,
                    get_variable(psi_container, val)[data.name, t],
                    -1.0,
                )
            end
            constraint[data.name, t] = JuMP.@constraint(
                psi_container.JuMPmodel,
                expression_lb >= data.multiplier * data.timeseries[t]
            )
        end
    end
    return
end

@doc raw"""
    device_timeseries_param_ub(psi_container::PSIContainer,
                                    ts_data::Vector{DeviceTimeSeries},
                                    cons_name::Symbol,
                                    param_reference::UpdateRef,
                                    var_name::Symbol)

Constructs upper bound for given variable using a parameter. The constraint is
    built with a time series data vector and a multiplier

# Constraint

``` variable[name, t] <= ts_data[name].multiplier * param[name, t] ```

# LaTeX

`` x^{var}_t \leq r^{val} x^{param}_t, \forall t ``

# Arguments
* psi_container::PSIContainer : the psi_container model built in PowerSimulations
* ts_data::Vector{DeviceTimeSeries} : container of device time series data and scaling factors
* cons_name::Symbol : name of the constraint
* param_reference::UpdateRef : UpdateRef to access the parameter
* var_name::Symbol : the name of the variable
"""
function device_timeseries_param_ub(
    psi_container::PSIContainer,
    ts_data::Vector{DeviceTimeSeries},
    cons_name::Symbol,
    param_reference::UpdateRef,
    var_name::Symbol,
)
    time_steps = model_time_steps(psi_container)
    names = (d.name for d in ts_data)
    variable = get_variable(psi_container, var_name)
    ub_name = middle_rename(cons_name, PSI_NAME_DELIMITER, "ub")
    con_ub = add_cons_container!(psi_container, ub_name, names, time_steps)
    container = add_param_container!(psi_container, param_reference, names, time_steps)
    multiplier = get_multiplier_array(container)
    param = get_parameter_array(container)
    lazy_add_lb = false

    for data in ts_data
        for t in time_steps
            expression_ub = JuMP.AffExpr(0.0, variable[data.name, t] => 1.0)
            for val in data.range.additional_terms_ub
                JuMP.add_to_expression!(
                    expression_ub,
                    get_variable(psi_container, val)[data.name, t],
                )
            end
            param[data.name, t] =
                PJ.add_parameter(psi_container.JuMPmodel, data.timeseries[t])
            con_ub[data.name, t] = JuMP.@constraint(
                psi_container.JuMPmodel,
                expression_ub <= data.multiplier * param[data.name, t]
            )
            multiplier[data.name, t] = data.multiplier
        end
        if data.range.limits.min > -Inf || !isempty(data.range.additional_terms_lb)
            lazy_add_lb = true
        end
    end
    @debug lazy_add_lb
    lazy_add_lb && lazy_lb(psi_container, ts_data, cons_name, var_name)
    return
end

@doc raw"""
    device_timeseries_param_lb(psi_container::PSIContainer,
                                    ts_data::Vector{DeviceTimeSeries},
                                    cons_name::Symbol,
                                    param_reference::UpdateRef,
                                    var_name::Symbol)

Constructs lower bound for given variable using a parameter. The constraint is
    built with a time series data vector and a multiplier

# Constraint

``` ts_data[name].multiplier * param[name, t] <= variable[name, t] ```

# LaTeX

`` r^{val} x^{param}_t \leq x^{var}_t, \forall t ``

# Arguments
* psi_container::PSIContainer : the psi_container model built in PowerSimulations
* ts_data::Vector{DeviceTimeSeries} : container of device time series data and scaling factors
* cons_name::Symbol : name of the constraint
* param_reference::UpdateRef : UpdateRef to access the parameter
* var_name::Symbol : the name of the variable
"""
function device_timeseries_param_lb(
    psi_container::PSIContainer,
    ts_data::Vector{DeviceTimeSeries},
    cons_name::Symbol,
    param_reference::UpdateRef,
    var_name::Symbol,
)
    time_steps = model_time_steps(psi_container)
    variable = get_variable(psi_container, var_name)
    lb_name = middle_rename(cons_name, PSI_NAME_DELIMITER, "lb")
    names = (d.name for d in ts_data)
    constraint = add_cons_container!(psi_container, lb_name, names, time_steps)
    container = add_param_container!(psi_container, param_reference, names, time_steps)
    multiplier = get_multiplier_array(container)
    param = get_parameter_array(container)

    for data in ts_data
        for t in time_steps
            expression_lb = JuMP.AffExpr(0.0, variable[data.name, t] => 1.0)
            for val in data.range.additional_terms_lb
                JuMP.add_to_expression!(
                    expression_lb,
                    get_variable(psi_container, val)[data.name, t],
                    -1.0,
                )
            end
            param[data.name, t] =
                PJ.add_parameter(psi_container.JuMPmodel, data.timeseries[t])
            constraint[data.name, t] = JuMP.@constraint(
                psi_container.JuMPmodel,
                expression_lb >= data.multiplier * param[data.name, t]
            )
            multiplier[data.name, t] = data.multiplier
        end
    end

    return

end

@doc raw"""
    device_timeseries_ub_bin(psi_container::PSIContainer,
                                    ts_data::Vector{DeviceTimeSeries},
                                    var_name::Symbol,
                                    binvar_name::Symbol)

Constructs upper bound for variable and time series or confines to 0 depending on binary variable.
    The upper bound is defined by a time series and a multiplier.

# Constraints

``` varcts[name, t] <= varbin[name, t]* ts_data[name].multiplier * ts_data[name].timeseries[t] ```

where (name, data) in range_data.

# LaTeX

`` x^{cts}_t \leq r^{val} r_t x^{bin}_t, \forall t ``

# Arguments
* psi_container::PSIContainer : the psi_container model built in PowerSimulations
* ts_data::Vector{DeviceTimeSeries} : container of device time series data and scaling factors
* cons_name::Symbol : name of the constraint
* var_name::Symbol :  name of the variable
* binvar_name::Symbol : name of binary variable
"""
function device_timeseries_ub_bin(
    psi_container::PSIContainer,
    ts_data::Vector{DeviceTimeSeries},
    cons_name::Symbol,
    var_name::Symbol,
    binvar_name::Symbol,
)
    time_steps = model_time_steps(psi_container)
    ub_name = middle_rename(cons_name, PSI_NAME_DELIMITER, "ub")
    varcts = get_variable(psi_container, var_name)
    varbin = get_variable(psi_container, binvar_name)
    names = (d.name for d in ts_data)
    con_ub = add_cons_container!(psi_container, ub_name, names, time_steps)
    for data in ts_data
        for t in time_steps
            forecast = data.timeseries[t]
            multiplier = data.multiplier
            expression_ub = JuMP.AffExpr(0.0, varcts[data.name, t] => 1.0)
            for val in data.range.additional_terms_ub
                JuMP.add_to_expression!(
                    expression_ub,
                    get_variable(psi_container, val)[data.name, t],
                )
            end
            con_ub[data.name, t] = JuMP.@constraint(
                psi_container.JuMPmodel,
                expression_ub <= varbin[data.name, t] * multiplier * forecast
            )
        end
    end
    return
end

@doc raw"""
    device_timeseries_ub_bigM(psi_container::PSIContainer,
                                    ts_data::Vector{DeviceTimeSeries},
                                    cons_name::Symbol,
                                    var_name::Symbol,
                                    param_reference::UpdateRef,
                                    binvar_name::Symbol,
                                    M_value::Float64 = 1e6)

Constructs upper bound for variable and time series and a multiplier or confines to 0 depending on binary variable.
    Uses BigM constraint type to allow for parameter since ParameterJuMP doesn't support var*parameter

# Constraints

``` varcts[name, t] - ts_data[name].multipliers * param[name, t] <= (1 - varbin[name, t]) * M_value ```

``` varcts[name, t] <= varbin[name, t]*M_value ```

# LaTeX

`` x^{cts}_t - r^{val} x^{param}_t \leq M(1 - x^{bin}_t ), forall t ``

`` x^{cts}_t \leq M x^{bin}_t, \forall t ``

# Arguments
* psi_container::PSIContainer : the psi_container model built in PowerSimulations
* ts_data::Vector{DeviceTimeSeries} : container of device time series data and scaling factors
* cons_name::Symbol : name of the constraint
* var_name::Symbol :  name of the variable
param_reference::UpdateRef : UpdateRef of access the parameters
* binvar_name::Symbol : name of binary variable
* M_value::Float64 : bigM
"""
function device_timeseries_ub_bigM(
    psi_container::PSIContainer,
    ts_data::Vector{DeviceTimeSeries},
    cons_name::Symbol,
    var_name::Symbol,
    param_reference::UpdateRef,
    binvar_name::Symbol,
    M_value::Float64 = 1e6,
)
    time_steps = model_time_steps(psi_container)
    ub_name = middle_rename(cons_name, PSI_NAME_DELIMITER, "ub")
    key_status = middle_rename(cons_name, PSI_NAME_DELIMITER, "status")

    varcts = get_variable(psi_container, var_name)
    varbin = get_variable(psi_container, binvar_name)
    names = (d.name for d in ts_data)
    con_ub = add_cons_container!(psi_container, ub_name, names, time_steps)
    con_status = add_cons_container!(psi_container, key_status, names, time_steps)
    container = add_param_container!(psi_container, param_reference, names, time_steps)
    multiplier = get_multiplier_array(container)
    param = get_parameter_array(container)

    for data in ts_data
        for t in time_steps
            expression_ub = JuMP.AffExpr(0.0, varcts[data.name, t] => 1.0)
            for val in data.range.additional_terms_ub
                JuMP.add_to_expression!(
                    expression_ub,
                    get_variable(psi_container, val)[data.name, t],
                )
            end
            param[data.name, t] =
                PJ.add_parameter(psi_container.JuMPmodel, data.timeseries[t])
            con_ub[data.name, t] = JuMP.@constraint(
                psi_container.JuMPmodel,
                expression_ub - param[data.name, t] * data.multiplier <=
                (1 - varbin[data.name, t]) * M_value
            )
            con_status[data.name, t] = JuMP.@constraint(
                psi_container.JuMPmodel,
                expression_ub <= varbin[data.name, t] * M_value
            )
            multiplier[data.name, t] = data.multiplier
        end
    end
    return
end


module Condor

#using Falcons
using Healpix
using PyCall
using Base.Threads
using StaticArrays
using ReferenceFrameRotations
using LinearAlgebra
using ProgressMeter
using HDF5
using Interpolations
const hp = pyimport("healpy")
np = pyimport("numpy")
pd = pyimport("pandas")
using WignerD

function __init__()
    copy!(hp, pyimport_conda("healpy", "healpy"))
end

include("./function/beam_func.jl")
include("./function/convolution_func.jl")
include("./function/adding_YN_func.jl")

export get_2Ddata, symmetrizer, dBi, gen_beammap, beam_pointor, AlmPair, gen_Blm, gen_GaussBeammap
export EffectiveBeamConvolution, GaussBeamConvolution, convolvor

end


using ContinuousPOMDPTreeSearchExperiments
using ParticleFilters
using ARDESPOT
using BasicPOMCP
using POMCPOW
using POMDPs
using DiscreteValueIteration
using QMDP
using MCTS
using VDPTag2
using POMDPToolbox
using DataFrames

file_contents = readstring(@__FILE__())

pomdp = VDPTagPOMDP()
dpomdp = AODiscreteVDPTagPOMDP(pomdp, 30, 0.5)

@show max_time = 1.0
@show max_depth = 10

solvers = Dict{String, Union{Solver,Policy}}(
    "to_next" => ToNextML(mdp(pomdp)),
    "manage_uncertainty" => ManageUncertainty(pomdp, 0.01),

    "pomcpow" => begin
        rng = MersenneTwister(13)
        ro = ToNextMLSolver(rng)
        solver = POMCPOWSolver(tree_queries=10_000_000,
                               criterion=MaxUCB(65.0),
                               final_criterion=MaxTries(),
                               max_depth=max_depth,
                               max_time=max_time,
                               k_action=12.0,
                               alpha_action=1/8,
                               k_observation=1.0,
                               alpha_observation=1/30,
                               estimate_value=FORollout(ro),
                               next_action=RootToNextMLFirst(rng),
                               check_repeat_obs=false,
                               check_repeat_act=false,
                               default_action=ReportWhenUsed(TagAction(false, 0.0)),
                               rng=rng
                              )
    end,

    "pft" => begin
        rng = MersenneTwister(13)
        m = 10
        node_updater = ObsAdaptiveParticleFilter(deepcopy(pomdp),
                                           LowVarianceResampler(m),
                                           0.05, rng)
        ro = ToNextML(mdp(pomdp), rng)
        solver = DPWSolver(n_iterations=typemax(Int),
                           exploration_constant=65.0,
                           depth=max_depth,
                           max_time=max_time,
                           k_action = 12.0,
                           alpha_action = 1/8,
                           k_state = 4.0,
                           alpha_state = 1/20,
                           check_repeat_state=false,
                           check_repeat_action=false,
                           estimate_value=RolloutEstimator(ro),
                           next_action=RootToNextMLFirst(rng),
                           rng=rng
                          )
        belief_mdp = GenerativeBeliefMDP(deepcopy(pomdp), node_updater)
        solve(solver, belief_mdp)
    end,

    "d_despot" => begin
        rng = MersenneTwister(13)
        ro = ToNextMLSolver(rng)
        b = IndependentBounds(DefaultPolicyLB(ro), VDPUpper())
        sol = DESPOTSolver(lambda=0.01,
                     K=500,
                     D=10,
                     max_trials=1_000_000,
                     T_max=0.1,
                     bounds=b,
                     random_source=MemorizingSource(500, 10, rng, min_reserve=8),
                     rng=rng)
        planner = solve(sol, dpomdp)
        translate_policy(planner, dpomdp, pomdp, dpomdp)
    end

)

@show N=500

alldata = DataFrame()
# for (k, solver) in solvers
test = ["pft"]
for (k, solver) in [(s, solvers[s]) for s in test]
    @show k
    if isa(solver, Solver)
        planner = solve(solver, pomdp)
    else
        planner = solver
    end
    sims = []
    for i in 1:N
        srand(planner, i+50_000)
        filter = SIRParticleFilter(deepcopy(pomdp), 100_000, rng=MersenneTwister(i+90_000))

        md = Dict(:solver=>k, :i=>i)
        sim = Sim(deepcopy(pomdp),
                  planner,
                  filter,
                  rng=MersenneTwister(i+70_000),
                  max_steps=100,
                  metadata=md
                 )

        push!(sims, sim)
    end

    data = run_parallel(sims)
    # data = run(sims)

    rs = data[:reward]
    println(@sprintf("reward: %6.3f ± %6.3f", mean(rs), std(rs)/sqrt(length(rs))))
end

#=
datestring = Dates.format(now(), "E_d_u_HH_MM")
copyname = Pkg.dir("ContinuousPOMDPTreeSearchExperiments", "icaps_2018", "data", "subhunt_table_$(datestring).jl")
write(copyname, file_contents)
filename = Pkg.dir("ContinuousPOMDPTreeSearchExperiments", "icaps_2018", "data", "subhunt_$(datestring).csv")
println("saving to $filename...")
writetable(filename, alldata)
println("done.")
=#


using ArgParse
s = ArgParseSettings()
@add_arg_table s begin
    "group1"
        help = "One end point of the edge for which to analyze context, where "*
               "| is used to separate mulitple ids. (i.e. ENCFF000ARO|ENCFF000ARK)."
        arg_type = ASCIIString
        required = true
    "group2"
        help = "The other end point of the edge for which to analyze context, where "*
               "| is used to separate mulitple ids. (i.e. ENCFF000ARO|ENCFF000ARK)."
        arg_type = ASCIIString
        required = true
    "data_bundles"
        help = "Location of the ChromNet data bundles (i.e. ENCODE_buildX.ChromNet.jld)"
        arg_type = ASCIIString
        nargs = '+'
        required = true
    "--output", "-o"
        help = "Where to save the edge context BED file"
        default = "-"
    "--load-cor"
        help = "A copy of the raw data correlation matrix"
        metavar = "COR_FILE"
    "--quiet", "-q"
        help = "Don't print anything"
        action = :store_true
end
args = parse_args(s)

using ChromNet
using JSON
using JLD
using SamIO
using ProgressMeter

# open the data bundle and load the metadata
dataBundles = [jldopen(x) for x in args["data_bundles"]]
metadata = Dict()
for bundle in dataBundles
    metadata = merge(metadata, read(bundle, "metadata"))
end
ids = vcat([read(bundle, "ids") for bundle in dataBundles]...)
numBins = round(Int, ceil(sum(ReferenceContigs_hg38.sizes) / 1000))

# find the indexes of the groups this edge is between
group1Inds = map(x->findfirst(ids, x), split(args["group1"], "|"))
group2Inds = map(x->findfirst(ids, x), split(args["group2"], "|"))
@assert all(group1Inds .> 0) "Some ids from group 1 were not found: $(group1Inds[group1Inds .== 0])"
@assert all(group2Inds .> 0) "Some ids from group 2 were not found: $(group2Inds[group2Inds .== 0])"

# compute the joint correlation matrix
C = nothing
if typeof(args["load-cor"]) != Void
    C,header = readdlm(args["load-cor"], header=true)
    @assert all(header .== ids) "Loaded correlation matrix does not match data! Perhaps a different file name ordering?"
else
    C = streaming_cor(dataBundles, quiet=args["quiet"])
end
cov2cor!(C)

# compute the edge impact
edgeScores = streaming_edgeimpact(dataBundles, C, group1Inds, group2Inds, quiet=args["quiet"])
map(close, dataBundles)

# send the data in BED format to the output stream
outStream = STDOUT
if args["output"] != "-"
    outStream = open(args["output"], "w")
end
contigIndex = 1
pos = 1
ctgs = ReferenceContigs_hg38
for i in 1:numBins
    while pos > ctgs.offsets[contigIndex] + ctgs.sizes[contigIndex]
        contigIndex += 1
    end
    println(outStream, ctgs.names[contigIndex], "\t", pos - ctgs.offsets[contigIndex], "\t", pos - ctgs.offsets[contigIndex] + 999, "\t*\t", edgeScores[i])
    pos += 1000
end


struct ModelClosure{M,A} <: AbstractConditionalModel{M,A,NamedTuple{(), Tuple{}}}
    model::M
    argvals::A
end

function Base.show(io::IO, mc::ModelClosure)
    println(io, "ModelClosure given")
    println(io, "    arguments    ", keys(argvals(mc)))
    println(io, "    observations ", keys(observations(mc)))
    println(io, model(mc))
end

export argvals
argvals(c::ModelClosure) = c.argvals

export observations
observations(c::ModelClosure) = NamedTuple()

export observed
function observed(mc::ModelClosure{M,A}) where {M,A}
    NamedTuple()
end

model(c::ModelClosure) = c.model

ModelClosure(m::AbstractModel) = ModelClosure(m,NamedTuple())

(m::AbstractModel)(nt::NamedTuple) = ModelClosure(m, nt)

(mc::ModelClosure)(nt::NamedTuple) = ModelClosure(model(mc), merge(mc.argvals, nt))

argvalstype(mc::ModelClosure{M,A}) where {M,A} = A
argvalstype(::Type{MC}) where {M,A,MC<:ModelClosure{M,A}} = A

obstype(::ModelClosure) = NamedTuple{(), Tuple{}}
obstype(::Type{<:ModelClosure}) = NamedTuple{(), Tuple{}}

type2model(::Type{MC}) where {M,MC<:ModelClosure{M}} = type2model(M)

import Base

Base.:|(m::ModelClosure, nt::NamedTuple) = ModelPosterior(m, nt)


abstract UniverseEstimate

type MuSigmaUniverse{T<:MuSigmaModel} <: UniverseEstimate
    Universe::T
    Date::Date
    Data::Timematr
    # EstDetails::Estimator

    function MuSigmaUniverse(univ::T, dat::Date, data::Timematr)
        return new(univ, dat, data)
    end
end

function MuSigmaUniverse{T}(univ::T, dat::Date, data::Timematr)
    return MuSigmaUniverse{T}(univ, dat, data)
end

## check for undefined universe
##-----------------------------

function isDef(univ::UniverseEstimate)
    return isDef(univ.Universe)
end

## display
##--------

import Base.Multimedia.display
function display(mod::UniverseEstimate)
    ## display type information
    typ = typeof(mod)
    println("\ntype: $typ")

    if !isDef(mod)
        println("Model is still undefined!")
        println("   mu: NA")
        println("sigma: NA")
    else
        ## display model
        print("model: ")
        display(typeof(mod.Universe))

        ## display mus
        println("Date: $(mod.Date)")

    end

    ## display data information
    print("Data: ")
    fromDate = mod.Data.idx[1];
    toDate = mod.Data.idx[end];
    nObs, nAss = size(mod.Data)
    println("From: $fromDate, To: $toDate, Size: $nObs x $nAss")
end

##############
## estimate ##
##############

macro defineEstimateForMuSigmaModel(myType)
    esc(quote
        function estimate(t::Type{$myType}, data::Timematr, dat::Date)
            ## estimate model and create UniverseEstimate instance
            modFit = fitModel(t, data, dat)
            return MuSigmaUniverse(modFit, dat, data)
        end
    end)
end

for t = (:(SampleMoments), :(MovWinSampleMoments), :(ExpWeighted))
    eval(macroexpand(:(@defineEstimateForMuSigmaModel($t))))
end


if basename(pwd()) == "aoc"
    cd("2019/2")
end
include("../IntCodeMachine/src/IntCodeMachine.jl")
using .IntCodeMachine

function part1()
    c = IntCode(readchomp("input.txt"))
    c[1] = 12
    c[2] = 2
    print(execute!(c)[0])
end

function part2()
    c = IntCode(readchomp("input.txt"))
    initmem = copy(c.memory)
    for x in 0:99, y in 0:99
        reset!(c, initmem)
        c[1] = x
        c[2] = y
        if execute!(c)[0] == 19690720
            return c[1] * 100 + c[2]
        end
    end
end


function FMM_mainMVP_pre!(output, beta, scatteringMatrices, φs::Vector{Float64}, ids::Vector{Int}, P, mFMM, pre_agg, translated_sum)
    #@simd does not have a positive effect in 0.6.0
    #calculate matrix-vector product - devectorized with pre-preconditioning

    @inbounds begin
    mul!(output, mFMM.Znear, beta)
    G = length(mFMM.groups)
    fill!(pre_agg,0.0)
    #first preagg all
    for ig2 = 1:G
        for is = 1:mFMM.groups[ig2].size
            indices = (mFMM.groups[ig2].point_ids[is]-1)*(2*P+1)
            for ii = 1:2*P+1
                for iq = 1:mFMM.Q
                    pre_agg[iq,ig2] += mFMM.Agg[ig2][iq,(is-1)*(2*P+1) + ii]*beta[indices + ii]
                end
            end
        end
    end

    for ig1 = 1:G
        #translate plane waves from ig2 to ig1
        fill!(translated_sum,0.0)
        for ig2 = 1:G
            if isempty(mFMM.Trans[(ig1-1)*G + ig2])
                continue
            else
                for iQ = 1:mFMM.Q
                    translated_sum[iQ] -= mFMM.Trans[(ig1-1)*G + ig2][iQ]*pre_agg[iQ,ig2]
                end
                #minus sign because the real equation is (I-ST)x=b
            end
        end
        #disaggregate from ig1 center to ig1's scatterers, producing -Tx
        for is = 1:mFMM.groups[ig1].size
            for ip = 1:2*P+1
                disagged = 0.0im
                for iq = 1:mFMM.Q
                    disagged += conj(mFMM.Agg[ig1][iq,(is-1)*(2*P+1) + ip])*translated_sum[iq]
                end
                output[(mFMM.groups[ig1].point_ids[is]-1)*(2*P+1) + ip] += disagged
            end
        end
    end
    #multiply by S to produce -S
    #temp can be moved outward, but for now this preallocation prevents most
    #dynamic mem alloc by this MVP
    temp = Array{Complex{Float64}}(undef, 2*P+1)
    for ic = 1:length(ids)
        rng = (ic-1)*(2*P+1) .+ (1:2*P+1)
        #copy values to avoid dynamic allocation - copyto!(temp,1:2*P+1,output,rng) doesn't work because output is a subarray
        for ip = 1:2*P+1
            temp[ip] = output[(ic-1)*(2*P+1) + ip]
        end
        v = view(output,rng)
        if φs[ic] == 0.0
            mul!(v, scatteringMatrices[ids[ic]], temp)
        else
            #rotate without matrix
            rotateMultipole!(temp,-φs[ic],P)
            mul!(v, scatteringMatrices[ids[ic]], temp)
            rotateMultipole!(v,φs[ic],P)
        end
    end
    #add identity matrix (Ix)
    output .+= beta
    end #inbounds
    return output
end

function FMM_mainMVP_transpose!(output, beta, scatteringMatrices, φs::Vector{Float64}, ids::Vector{Int}, P, mFMM, pre_agg, translated_sum)
    #calculate matrix^T-vector product - devectorized with pre-preconditioning

    @inbounds begin
    #here we first multiply by Xᵀ

    #Xβ storage can be moved outward, but for now this preallocation prevents most
    #dynamic mem alloc by this MVP
    Xβ = Array{Complex{Float64}}(undef, length(beta))
    temp = Array{Complex{Float64}}(undef, 2*P+1)
    for ic = 1:length(ids)
        rng = (ic-1)*(2*P+1) .+ (1:2*P+1)
        v = view(Xβ, rng)
        if φs[ic] == 0.0
            mul!(v, transpose(scatteringMatrices[ids[ic]]), view(beta, rng))
        else
            #rotate without matrix - transposed
            rotateMultipole!(temp, view(beta, rng), φs[ic], P)
            mul!(v, transpose(scatteringMatrices[ids[ic]]), temp)
            rotateMultipole!(v, -φs[ic], P)
        end
    end

    mul!(output, transpose(mFMM.Znear), Xβ)
    G = length(mFMM.groups)
    fill!(pre_agg,0.0)
    #first preagg all
    for ig2 = 1:G
        for is = 1:mFMM.groups[ig2].size
            indices = (mFMM.groups[ig2].point_ids[is]-1)*(2*P+1)
            for ii = 1:2*P+1
                for iq = 1:mFMM.Q
                    pre_agg[iq,ig2] += conj(mFMM.Agg[ig2][iq,(is-1)*(2*P+1) + ii])*Xβ[indices + ii]
                end
            end
        end
    end

    for ig1 = 1:G
        #translate plane waves from ig2 to ig1
        fill!(translated_sum,0.0)
        for ig2 = 1:G
            if isempty(mFMM.Trans[(ig2-1)*G + ig1]) #transposed
                continue
            else
                for iQ = 1:mFMM.Q
                    translated_sum[iQ] -= mFMM.Trans[(ig2-1)*G + ig1][iQ]*pre_agg[iQ,ig2] #transposed
                end
                #minus sign because the real equation is (I-TᵀXᵀ)β=b
            end
        end
        #disaggregate from ig1 center to ig1's scatterers
        for is = 1:mFMM.groups[ig1].size
            for ip = 1:2*P+1
                disagged = 0.0im
                for iq = 1:mFMM.Q
                    disagged += mFMM.Agg[ig1][iq,(is-1)*(2*P+1) + ip]*translated_sum[iq]
                end
                output[(mFMM.groups[ig1].point_ids[is]-1)*(2*P+1) + ip] += disagged
            end
        end
    end

    #add identity matrix (Iβ)
    output .+= beta
    end #inbounds
    return output
end


## We have different types:
##
## Basic: holds a ptr to a symengine object. Faster, so is default type
##
## BasicType{Val{:XXX}}: types that can be use to control dispatch
##
## SymbolicType: is a type union of the two
##
## Basic(x::BasicType) gives a basic object; BasicType(x::Basic) gives a BasicType object. (This name needs change)
##
## To control dispatch, one might have `N(b::Basic) = N(BasicType(b))` and then define `N` for types of interest

## Hold a reference to a SymEngine object
mutable struct Basic  <: Number
    ptr::Ptr{Cvoid}
    function Basic()
        z = new(C_NULL)
        ccall((:basic_new_stack, libsymengine), Nothing, (Ref{Basic}, ), z)
        _finalizer(basic_free, z)
        return z
    end
    function Basic(v::Ptr{Cvoid})
        z = new(v)
        return z
    end
end

basic_free(b::Basic) = ccall((:basic_free_stack, libsymengine), Nothing, (Ref{Basic}, ), b)

function convert(::Type{Basic}, x::Clong)
    a = Basic()
    ccall((:integer_set_si, libsymengine), Nothing, (Ref{Basic}, Clong), a, x)
    return a
end

function convert(::Type{Basic}, x::Culong)
    a = Basic()
    ccall((:integer_set_ui, libsymengine), Nothing, (Ref{Basic}, Culong), a, x)
    return a
end

function convert(::Type{Basic}, x::BigInt)
    a = Basic()
    ccall((:integer_set_mpz, libsymengine), Nothing, (Ref{Basic}, Ref{BigInt}), a, x)
    return a
end

function convert(::Type{Basic}, s::String)
    a = Basic()
    b = ccall((:basic_parse, libsymengine), Cuint, (Ref{Basic}, Ptr{Int8}), a, s)
    throw_if_error(b, s)
    return a
end

function convert(::Type{Basic}, ex::Expr)
    expr = copy(ex)
    Basic(string(expr))
end

convert(::Type{Basic}, ex::Symbol) = Basic(string(ex))

function convert(::Type{Basic}, x::Cdouble)
    a = Basic()
    ccall((:real_double_set_d, libsymengine), Nothing, (Ref{Basic}, Cdouble), a, x)
    return a
end

function convert(::Type{Basic}, x::BigFloat)
    if (x.prec <= 53)
        return convert(Basic, Cdouble(x))
    elseif have_mpfr
        a = Basic()
        ccall((:real_mpfr_set, libsymengine), Nothing, (Ref{Basic}, Ref{BigFloat}), a, x)
        return a
    else
        warn("SymEngine is not compiled with MPFR support. Converting will lose precision.")
        return convert(Basic, Cdouble(x))
    end
end

if Clong == Int32
    convert(::Type{Basic}, x::Union{Int8, Int16}) = Basic(convert(Clong, x))
    convert(::Type{Basic}, x::Union{UInt8, UInt16}) = Basic(convert(Culong, x))
else
    convert(::Type{Basic}, x::Union{Int8, Int16, Int32}) = Basic(convert(Clong, x))
    convert(::Type{Basic}, x::Union{UInt8, UInt16, UInt32}) = Basic(convert(Culong, x))
end
convert(::Type{Basic}, x::Union{Float16, Float32}) = Basic(convert(Cdouble, x))
convert(::Type{Basic}, x::Integer) = Basic(BigInt(x))
convert(::Type{Basic}, x::Rational) = Basic(numerator(x)) / Basic(denominator(x))
convert(::Type{Basic}, x::Complex) = Basic(real(x)) + Basic(imag(x)) * IM

Basic(x::T) where {T} = convert(Basic, x)
Basic(x::Basic) = x

Base.promote_rule(::Type{Basic}, ::Type{S}) where {S<:Number} = Basic
Base.promote_rule(::Type{S}, ::Type{Basic}) where {S<:Number} = Basic
if VERSION > VersionNumber("0.7.0-DEV")
    Base.promote_rule(::Type{S}, ::Type{Basic}) where {S<:AbstractIrrational} = Basic
else
    Base.promote_rule(::Type{S}, ::Type{Basic}) where {S<:Irrational} = Basic
end

## Class ID
get_type(s::Basic) = ccall((:basic_get_type, libsymengine), UInt, (Ref{Basic},), s)
function get_class_from_id(id::UInt)
    a = ccall((:basic_get_class_from_id, libsymengine), Ptr{UInt8}, (Int,), id)
    str = unsafe_string(a)
    ccall((:basic_str_free, libsymengine), Nothing, (Ptr{UInt8}, ), a)
    str
end

"Get SymEngine class of an object (e.g. 1=>:Integer, 1//2 =:Rational, sin(x) => :Sin, ..."
get_symengine_class(s::Basic) = Symbol(get_class_from_id(get_type(s)))

## Construct symbolic objects
## renamed, as `Symbol` conflicts with Base.Symbol
function _symbol(s::String)
    a = Basic()
    ccall((:symbol_set, libsymengine), Nothing, (Ref{Basic}, Ptr{Int8}), a, s)
    return a
end
_symbol(s::Symbol) = _symbol(string(s))

## use SymPy name here, but no assumptions
"""

`symbols(::Symbol)` construct symbolic value

Examples:
```
a = symbols(:a)
x = symbols("x")
x,y = symbols("x y")
x,y,z = symbols("x,y,z")
```

"""
symbols(s::Symbol) = _symbol(s)
function symbols(s::String)
    ## handle space or comma sparation
    s = replace(s, ","=> " ")
    by_space = split(s, r"\s+")
    Base.length(by_space) == 1 && return symbols(Symbol(s))
    tuple([_symbol(Symbol(o)) for o in by_space]...)
end

## Follow, somewhat, the python names: symbols to construct symbols, @vars

"""
Macro to define 1 or more variables in the main workspace.

Symbolic values are defined with `_symbol`. This is a convenience

Example
```
@vars x y z
```
"""
macro vars(x...)
    q=Expr(:block)
    if length(x) == 1 && isa(x[1],Expr)
        @assert x[1].head === :tuple "@syms expected a list of symbols"
        x = x[1].args
    end
    for s in x
        @assert isa(s,Symbol) "@syms expected a list of symbols"
        push!(q.args, Expr(:(=), esc(s), Expr(:call, :(SymEngine._symbol), Expr(:quote, s))))
    end
    push!(q.args, Expr(:tuple, map(esc, x)...))
    q
end

## We also have a wrapper type that can be used to control dispatch
## pros: wrapping adds overhead, so if possible best to use Basic
## cons: have to write methods meth(x::Basic, ...) = meth(BasicType(x),...)

## Parameterized type allowing for dispatch on Julia side by type of objecton SymEngine side
## Use as BasicType{Val{:Integer}}(...)
## To take advantage of this, define
## meth(x::Basic) = meth(BasicType(x))
## and then
## meth(x::BasicType{Val{:Integer}}) = ... or
## meth(x::BasicNumber) = ...
mutable struct BasicType{T} <: Number
    x::Basic
end

convert(::Type{Basic}, x::Basic) = x

SymbolicType = Union{Basic, BasicType}
convert(::Type{Basic}, x::BasicType) = x.x
Basic(x::BasicType) = x.x

BasicType(val::Basic) =  BasicType{Val{get_symengine_class(val)}}(val)
convert(::Type{BasicType{T}}, val::Basic) where {T} = BasicType{Val{get_symengine_class(val)}}(val)
# Needed for julia v0.4.7
convert(::Type{T}, x::Basic) where {T<:BasicType} = BasicType(x)

## We have Basic and BasicType{...}. We go back and forth with:
## Basic(b::BasicType) and BasicType(b::Basic)

# for mathops
Base.promote_rule(::Type{T}, ::Type{S} ) where {T<:BasicType, S<:Number} = T
Base.promote_rule(::Type{S}, ::Type{T} ) where {T<:BasicType, S<:Number} = T

# to intersperse BasicType and Basic in math ops
Base.promote_rule(::Type{T}, ::Type{Basic} ) where {T<:BasicType} = T
Base.promote_rule( ::Type{Basic}, ::Type{T} ) where {T<:BasicType} = T
if VERSION > VersionNumber("0.7.0-DEV")
    Base.promote_rule(::Type{S}, ::Type{T}) where {S<:AbstractIrrational, T<:BasicType} = T
else
    Base.promote_rule(::Type{S}, ::Type{T}) where {S<:Irrational, T<:BasicType} = T
end

## needed for mathops
convert(::Type{T}, val::Number) where {T<:BasicType} = T(Basic(val))
## Julia v0.6 errors with ambiguous error if this method is not defined.
convert(::Type{T}, val::T) where {T<:BasicType} = val

## some type unions possibly useful for dispatch
## Names here match those returned by get_symengine_class()
real_number_types = [:Integer, :RealDouble, :Rational, :RealMPFR]
complex_number_types = [:Complex, :ComplexDouble, :ComplexMPC]
number_types = vcat(real_number_types, complex_number_types)
BasicNumber = Union{[SymEngine.BasicType{Val{i}} for i in number_types]...}
BasicRealNumber = Union{[SymEngine.BasicType{Val{i}} for i in real_number_types]...}
BasicComplexNumber = Union{[SymEngine.BasicType{Val{i}} for i in complex_number_types]...}

op_types = [:Mul, :Add, :Pow, :Symbol, :Const]
BasicOp = Union{[SymEngine.BasicType{Val{i}} for i in op_types]...}

trig_types = [:Sin, :Cos, :Tan, :Csc, :Sec, :Cot, :ASin, :ACos, :ATan, :ACsc, :ASec, :ACot]
BasicTrigFunction =  Union{[SymEngine.BasicType{Val{i}} for i in trig_types]...}

###

" Return free symbols in an expression as a `Set`"
function free_symbols(ex::Basic)
    syms = CSetBasic()
    ccall((:basic_free_symbols, libsymengine), Nothing, (Ref{Basic}, Ptr{Cvoid}), ex, syms.ptr)
    convert(Vector, syms)
end
free_symbols(ex::BasicType) = free_symbols(Basic(ex))
_flat(A) = mapreduce(x->isa(x,Array) ? _flat(x) : x, vcat, A, init=Basic[])  # from rosetta code example
free_symbols(exs::Array{T}) where {T<:SymbolicType}  = unique(_flat([free_symbols(ex) for ex in exs]))
free_symbols(exs::Tuple) =  unique(_flat([free_symbols(ex) for ex in exs]))

"Return function symbols in an expression as a `Set`"
function function_symbols(ex::Basic)
    syms = CSetBasic()
    ccall((:basic_function_symbols, libsymengine), Nothing, (Ptr{Cvoid}, Ref{Basic}), syms.ptr, ex)
    convert(Vector, syms)
end
function_symbols(ex::BasicType) = function_symbols(Basic(ex))
function_symbols(exs::Array{T}) where {T<:SymbolicType} = unique(_flat([function_symbols(ex) for ex in exs]))
function_symbols(exs::Tuple) = unique(_flat([function_symbols(ex) for ex in exs]))

"Return name of function symbol"
function get_name(ex::Basic)
    a = ccall((:function_symbol_get_name, libsymengine), Cstring, (Ref{Basic}, ), ex)
    string = unsafe_string(a)
    ccall((:basic_str_free, libsymengine), Nothing, (Cstring, ), a)
    return string
end

"Return arguments of a function call as a vector of `Basic` objects"
function get_args(ex::Basic)
    args = CVecBasic()
    ccall((:basic_get_args, libsymengine), Nothing, (Ref{Basic}, Ptr{Cvoid}), ex, args.ptr)
    convert(Vector, args)
end

## so that Dicts will work
Base.hash(ex::Basic) = ccall((:basic_hash, libsymengine), UInt, (Ref{Basic}, ), ex)

function coeff(b::Basic, x::Basic, n::Basic)
    c = Basic()
    ccall((:basic_coeff, libsymengine), Nothing, (Ref{Basic}, Ref{Basic}, Ref{Basic}, Ref{Basic}), c, b, x, n)
    return c
end

coeff(b::Basic, x::Basic) = coeff(b, x, one(Basic))


function test_model(m::Model)
    @test length(m.cells) <= m.cfg["cells_max"]
    nconns = 0
    for cell in m.cells
        nconns += length(cell.conns)
        @test cell.input >= -1.0
        @test cell.input <= 1.0
        @test cell.output >= -1.0
        @test cell.output <= 1.0
        @test all(cell.state .>= -1.0)
        @test all(cell.state .<= 1.0)
        @test length(cell.state) == m.cfg["n_cell_state"]
        @test all(cell.params .>= -1.0)
        @test all(cell.params .<= 1.0)
        @test length(cell.params) == m.cfg["n_cell_params"]
        for conn in cell.conns
            @test conn.input >= -1.0
            @test conn.input <= 1.0
            @test all(conn.state .>= -1.0)
            @test all(conn.state .<= 1.0)
            @test length(conn.state) == m.cfg["n_conn_state"]
            @test all(conn.params .>= -1.0)
            @test all(conn.params .<= 1.0)
            @test length(conn.params) == m.cfg["n_conn_params"]
            @test conn.source[] == cell
            if typeof(conn.dest[]) == Cell
                @test conn.dest[] in m.cells
            else
                in_conns = false
                for c2 in m.cells
                    if conn.dest[] in c2.conns
                        in_conns = true
                        break
                    end
                end
                @test in_conns
            end
        end
    end
    @test nconns <= m.cfg["conns_max"]
end

@testset "Initialization tests" begin
    cfg = Config()
    c = Controller(cfg)
    nin = rand(5:10)
    nout = rand(5:10)

    @testset "Layered" begin
        m = Model(cfg, c)
        layered_init!(m, nin, nout)
        test_model(m)
    end

    @testset "Layered single cell" begin
        m = Model(cfg, c)
        layered_init!(m, nin, nout; nhidden=1)
        test_model(m)
    end

    @testset "Random cell" begin
        m = Model(cfg, c)
        random_init!(m, nin, nout)
        test_model(m)
    end

    @testset "Random single cell" begin
        m = Model(cfg, c)
        random_init!(m, nin, nout; nhidden=1)
        test_model(m)
    end
end

@testset "Model tests" begin
    cfg = Config()
    c = Controller(cfg)
    m = Model(cfg, c)
    nin = rand(5:10)
    nout = rand(5:10)
    nhidden = rand(5:10)
    nreward = rand(5:10)

    @testset "Initialize" begin
        random_init!(m, nin, nout; nhidden=nhidden, nreward=nreward)
        test_model(m)
        @test length(m.cells) == nin + nout + nhidden + nreward
        @test length(m.inputs) == nin
        @test length(m.outputs) == nout
    end

    @testset "Input" begin
        set_input!(m, rand(nin))
        test_model(m)
    end

    @testset "Step" begin
        for i in 1:5
            step!(m)
            @test m.state[1] == i
            test_model(m)
        end
    end

    @testset "Output" begin
        outputs = get_output(m)
        test_model(m)
        @test all(outputs .>= -1.0)
        @test all(outputs .<= 1.0)
        @test length(outputs) == nout
    end

    @testset "Reward" begin
        reward!(m, rand(nreward))
        test_model(m)
    end

    @testset "Full step" begin
        for i in 1:5
            println("Base controller step $i")
            @time outputs = step!(m, rand(nin))
            reward!(m, rand(nreward))
            test_model(m)
            @test all(outputs .>= -1.0)
            @test all(outputs .<= 1.0)
            @test length(outputs) == nout
        end
    end
end

function controller_test(cont_constructor::Function; cfg::Dict=Config())
    c = cont_constructor(cfg)
    m = Model(cfg, c)
    nin = rand(5:10)
    nout = rand(5:10)
    random_init!(m, nin, nout; nreward=nout)

    @testset "Full step" begin
        for i in 1:5
            println(string(cont_constructor, " step ", i))
            @time outputs = step!(m, rand(nin))
            reward!(m, rand(nout))
            test_model(m)
            @test all(outputs .>= -1.0)
            @test all(outputs .<= 1.0)
            @test length(outputs) == nout
        end
    end
end

@testset "Default controller" begin
    controller_test(x->Controller(x))
end

@testset "Const controller" begin
    controller_test(const_controller)
end

@testset "Random controller" begin
    controller_test(rand_controller)
end

@testset "Static controller" begin
    controller_test(static_controller;
                    cfg = Config(Config()))
end

# @testset "SNN controller" begin
#     controller_test(snn_controller;
#                     cfg = Config(Config(), "cfg/snn.yaml"))
# end

@testset "CGP controller" begin
    controller_test(cgp_controller)
end


"""
The Gr4j type contains the state variables (st, st_uh1, st_uh2), the inputs
(epot, prec) for one time step, the parameters (param) and the time step
length (tstep) for the GR4J model.
"""
type Gr4j <: Hydro

  st::Array{Float64,1}
  st_uh1::Array{Float64,1}
  st_uh2::Array{Float64,1}
  ord_uh1::Array{Float64,1}
  ord_uh2::Array{Float64,1}
  epot::Float64
  prec::Float64
  q_sim::Float64
  param::Array{Float64,1}
  tstep::Float64
  time::DateTime

end

"""
    Gr4j(tstep)
Constructor for GR4J with predefined state variables, parameters and inputs.
The time step (tstep) is given as a fraction of one day. Thus, for hourly input
data tstep should be set to 1/24.
"""
function Gr4j(tstep::Float64, time::DateTime)

  n_ord = ceil(Int64, 20.0 * 24.0 / tstep)

  st      = zeros(Float64, 2)
  st_uh1  = zeros(Float64, n_ord)
  st_uh2  = zeros(Float64, 2 * n_ord)
  ord_uh1 = zeros(Float64, n_ord)
  ord_uh2 = zeros(Float64, 2 * n_ord)

  epot = 0.0
  prec = 0.0

  q_sim = 0.0

  param  = [257.238, 1.012, 88.235, 2.208]

  st[1] = 0.3 * param[1]
  st[2] = 0.5 * param[3]

  D = 1.30434782 * (tstep / 24.0) + 1.19565217

  UH1(ord_uh1, param[4] * 24.0 / tstep, D)
  UH2(ord_uh2, param[4] * 24.0 / tstep, D)

  Gr4j(st, st_uh1, st_uh2, ord_uh1, ord_uh2, epot, prec, q_sim, param, tstep, time)

end

"""
    Gr4j(tstep, time, param)
Constructor for GR4J with predefined state variables and inputs. The parameter
values are given as input. The time step (tstep) is given as a fraction of one
day. Thus, for hourly input data tstep should be set to 1/24.
"""
function Gr4j(tstep::Float64, time::DateTime, param::Array{Float64,1})

  n_ord = ceil(Int64, 20.0 * 24.0 / tstep)

  st      = zeros(Float64, 2)
  st_uh1  = zeros(Float64, n_ord)
  st_uh2  = zeros(Float64, 2 * n_ord)
  ord_uh1 = zeros(Float64, n_ord)
  ord_uh2 = zeros(Float64, 2 * n_ord)

  st[1] = 0.3 * param[1]
  st[2] = 0.5 * param[3]

  D = 1.30434782 * (tstep / 24.0) + 1.19565217

  UH1(ord_uh1, param[4] * 24.0 / tstep, D)
  UH2(ord_uh2, param[4] * 24.0 / tstep, D)

  epot = 0.0
  prec = 0.0

  q_sim = 0.0

  Gr4j(st, st_uh1, st_uh2, ord_uh1, ord_uh2, epot, prec, q_sim, param, tstep, time)

end

"""
    init_states(mdata::Gr4j)
Initilize the state variables of the model.
"""
function init_states(mdata::Gr4j)

  mdata.st[1] = 0.3 * mdata.param[1]
  mdata.st[2] = 0.5 * mdata.param[3]

  for i in eachindex(mdata.st_uh1)
    mdata.st_uh1[i] = 0.
  end

  for i in eachindex(mdata.st_uh2)
    mdata.st_uh2[i] = 0.
  end

end

"""
    get_param_range(mdata::Gr4j)
Get allowed parameter ranges for the calibration of the model.
"""
function get_param_range(mdata::Gr4j)

  param_range_hydro = [(1.0, 20000.0), (-100.0, 100.0), (1.0, 20000.0), (0.5, 10.0)]

end

"""
    get_states(mdata::Gr4j)

Get state variables for computing penelty during calibration.

"""
function get_states(mdata::Gr4j)

  return []

end

"""
    assign_param(mdata::Gr4j, param::Array{Float64,1})
Assign parameter values to the Gr4j type.
"""
function assign_param(mdata::Gr4j, param::Array{Float64,1})

  for i in eachindex(mdata.param)
    mdata.param[i] = param[i]
  end

  D = 1.30434782 * (mdata.tstep / 24.0) + 1.19565217

  UH1(mdata.ord_uh1, param[4] * 24.0 / mdata.tstep, D)
  UH2(mdata.ord_uh2, param[4] * 24.0 / mdata.tstep, D)

end

"""
    get_states(hydro_out::Array{Gr4j,1})
Get state variables
"""
function get_states(hydro_out::Array{Gr4j,1})

  states = zeros(length(hydro_out), 2)

  for i in 1:length(hydro_out)

    states[i, 1] = hydro_out[i].st[1]
    states[i, 2] = hydro_out[i].st[2]

  end

  return states

end

"""
    enkf_hydro(mdata::Array{Gr4j,1}, obs_ens, q_sim)
Implementation of the ensemble Kalman filter for the GR4J model.
"""
function enkf_hydro(mdata::Array{Gr4j,1}, obs_ens, q_sim)

  nens = length(mdata)

  # Allocate arrays

  st     = zeros(Float64, length(mdata[1].st), nens)
  st_uh1 = zeros(Float64, length(mdata[1].st_uh1), nens)
  st_uh2 = zeros(Float64, length(mdata[1].st_uh2), nens)

  # Add states to arrays

  for iens = 1:nens

    st[:, iens]     = mdata[iens].st
    st_uh1[:, iens] = mdata[iens].st_uh1
    st_uh2[:, iens] = mdata[iens].st_uh2

  end

  # Run ensemble kalman filter

  st     = enkf(st, obs_ens, q_sim)
  st_uh1 = enkf(st_uh1, obs_ens, q_sim)
  st_uh2 = enkf(st_uh2, obs_ens, q_sim)

  # Check limits of states

  st[st .< 0] = 0.
  st_uh1[st_uh1 .< 0] = 0.
  st_uh2[st_uh2 .< 0] = 0.

  st[1, st[1, :] .> mdata[1].param[1]] = mdata[1].param[1]
  st[2, st[2, :] .> mdata[1].param[3]] = mdata[1].param[3]

  # Add arrays to states

  for iens = 1:nens

    mdata[iens].st = st[:, iens]
    mdata[iens].st_uh1 = st_uh1[:, iens]
    mdata[iens].st_uh2 = st_uh2[:, iens]

  end

  nothing

end

"""
    run_timestep(mdata::Gr4j)
Propagate the model one time step and return simulated dischage.
"""
function run_timestep(mdata::Gr4j)

  St     = mdata.st
  StUH1  = mdata.st_uh1
  StUH2  = mdata.st_uh2
  OrdUH1 = mdata.ord_uh1
  OrdUH2 = mdata.ord_uh2
  Param  = mdata.param
  P1     = mdata.prec
  E      = mdata.epot
  tstep  = mdata.tstep

  A = Param[1]
  B = 0.9

  # Interception and production store

  if P1 <= E
    EN = E - P1
    PN = 0.0
    WS = EN / A
    if WS > 13.0
      WS = 13.0
    end
    TWS = tanh(WS)
    Sr = St[1] / A
    ER = St[1] * (2.0-Sr)*TWS/(1.0+(1.0-Sr)*TWS)
    AE = ER + P1
    St[1] = St[1] - ER
    PR = 0.0
  else
    EN = 0.0
    AE = E
    PN = P1 - E
    WS = PN / A
    if WS > 13.0
      WS = 13.0
    end
    TWS = tanh(WS)
    Sr = St[1] / A
    PS = A*(1.0-Sr*Sr)*TWS/(1.0+Sr*TWS)
    PR = PN - PS
    St[1] = St[1] + PS
  end

  # Percolation from production store

  if St[1] < 0.0
    St[1] = 0.0
  end

  scale_param = 25.62891 * (24.0 / tstep)

  Sr = St[1] / Param[1]
  Sr = Sr * Sr
  Sr = Sr * Sr
  PERC = St[1] * (1.0-1.0/sqrt(sqrt(1.0 + Sr/scale_param)))

  St[1] = St[1] - PERC

  PR = PR + PERC

  # Split of effective rainfall into the two routing components

  PRHU1 = PR * B
  PRHU2 = PR * (1.0-B)

  # Convolution of unit hydrograph UH1

  NH = length(OrdUH1)

  for K = 1:NH-1
    StUH1[K] = StUH1[K+1] + OrdUH1[K]*PRHU1
  end

  StUH1[NH] = OrdUH1[NH] * PRHU1

  # Convolution of unit hydrograph UH2

  NH = length(OrdUH2)

  for K = 1:NH-1
    StUH2[K] = StUH2[K+1] + OrdUH2[K]*PRHU2
  end

  StUH2[NH] = OrdUH2[NH] * PRHU2

  # Potential intercatchment semi-exchange

  scale_param = (tstep / 24.0)

  Rr = St[2]/Param[3]
  EXCH = scale_param * Param[2]*Rr*Rr*Rr*sqrt(Rr)

  # Routing store

  AEXCH1 = EXCH

  if (St[2]+StUH1[1]+EXCH) < 0.0
    AEXCH1 = -St[2] - StUH1[1]
  end

  St[2] = St[2] + StUH1[1] + EXCH

  if St[2] < 0.0
    St[2] = 0.0
  end

  scale_param = (24.0 / tstep)
  Rr = St[2]^4 / (Param[3]^4 * scale_param)

  QR = St[2] * (1.0-1.0/sqrt(sqrt(1.0+Rr)))

  St[2] = St[2] - QR

  # Runoff from direct branch QD

  AEXCH2 = EXCH

  if (StUH2[1]+EXCH) < 0.0
    AEXCH2 = -StUH2[1]
  end

  QD = max(0.0,StUH2[1]+EXCH)

  # Total runoff

  Q = QR + QD

  if Q < 0.0
    Q = 0.0
  end

  # Update struct

  mdata.time    += Dates.Hour(tstep)
  mdata.st      = St
  mdata.st_uh1  = StUH1
  mdata.st_uh2  = StUH2
  mdata.ord_uh1 = OrdUH1
  mdata.ord_uh2 = OrdUH2
  mdata.q_sim   = Q

  return nothing

end

function SS1(I,C,D)

  FI = I
  if FI <= 0.0
    SS1 = 0.0
  elseif FI < C
    SS1 = (FI/C)^D
  else
    SS1 = 1.0
  end

end

function SS2(I,C,D)

  FI = I
  if FI <= 0.0
    SS2 = 0.0
  elseif FI <= C
    SS2 = 0.5*(FI/C)^D
  elseif FI < 2.0*C
    SS2 = 1.0-0.5*(2.0-FI/C)^D
  else
    SS2=1.0
  end

end

function UH1(OrdUH1,C,D)

  NH = length(OrdUH1)

  for I in 1:NH

    OrdUH1[I] = SS1(I,C,D)-SS1(I-1,C,D)

  end

end

function UH2(OrdUH2,C,D)

  NH = length(OrdUH2)

  for I in 1:NH

    OrdUH2[I] = SS2(I,C,D)-SS2(I-1,C,D)

  end

end

#**********************************************************************
#     SUBROUTINE MOD_GR4J(St,StUH1,StUH2,OrdUH1,OrdUH2,Param,P1,E,Q)
# Run on a single time step with the GR4J model
# Inputs:
#       St     Vector of model states in stores at the beginning of the time step [mm]
#       StUH1  Vector of model states in Unit Hydrograph 1 at the beginning of the time step [mm]
#       StUH2  Vector of model states in Unit Hydrograph 2 at the beginning of the time step [mm]
#       OrdUH1 Vector of ordinates in UH1 [-]
#       OrdUH2 Vector of ordinates in UH2 [-]
#       Param  Vector of model parameters [various units]
#       P1     Value of rainfall during the time step [mm/day]
#       E      Value of potential evapotranspiration during the time step [mm/day]
# Outputs:
#       St     Vector of model states in stores at the end of the time step [mm]
#       StUH1  Vector of model states in Unit Hydrograph 1 at the end of the time step [mm]
#       StUH2  Vector of model states in Unit Hydrograph 2 at the end of the time step [mm]
#       Q      Value of simulated flow at the catchment outlet for the time step [mm]
#**********************************************************************

#**********************************************************************
#     FUNCTION SS1(I,C,D)
# Values of the S curve (cumulative HU curve) of GR unit hydrograph UH1
# Inputs:
#    C: time constant
#    D: exponent
#    I: time-step
# Outputs:
#    SS1: Values of the S curve for I
#**********************************************************************

#**********************************************************************
#     FUNCTION SS2(I,C,D)
# Values of the S curve (cumulative HU curve) of GR unit hydrograph UH2
# Inputs:
#    C: time constant
#    D: exponent
#    I: time-step
# Outputs:
#    SS2: Values of the S curve for I
#**********************************************************************

#**********************************************************************
#     SUBROUTINE UH1(OrdUH1,C,D)
# Computation of ordinates of GR unit hydrograph UH1 using successive differences on the S curve SS1
# Inputs:
#    C: time constant
#    D: exponent
# Outputs:
#    OrdUH1: NH ordinates of discrete hydrograph
#**********************************************************************

#**********************************************************************
#     SUBROUTINE UH2(OrdUH2,C,D)
# Computation of ordinates of GR unit hydrograph HU2 using successive differences on the S curve SS2
# Inputs:
#    C: time constant
#    D: exponent
# Outputs:
#    OrdUH2: 2*NH ordinates of discrete hydrograph
#**********************************************************************


module exportData
include("../../Alessio_Data/loadAlessioData.jl");
using .loadAlessioData
export getExpData

function getExpData(ind; Vind=1, tInds=1:125)

        # get reference data
        basepath = "../../Alessio_Data/2011-05-28_Rec78-103_Pace_Apex/"

        # print some diagnostic notices
        print("Loading data for $(ind)...\t")

        # load in some data
        expStates = loadFile(basepath, ind)

        print("Done.\n")

        # form the space and time masks
        spaceMasks = [spatialMask(state) for state in expStates]
        timeMasks  = [temporalMask(state) for state in expStates]
        mutualSpaceMask = mutualMask(spaceMasks)
        mutualTimeMask = mutualMask(timeMasks)

        # determine mutualSpaceMask centroid for square center
        center = spatialMaskCentroid(mutualSpaceMask)
        samplePoint = Int.(center./0.06)
        t, V = sampleData(expStates, samplePoint, mutualTimeMask, mutualSpaceMask)
        APD, DI, APA = loadAlessioData.analyzeVoltage(t, V[1]; threshold=0.25)
        ii = 1:min(length(APD),length(DI))
        BCL = sum(APD[ii] .+ DI[ii])/length(APD[ii])

        # select EPI for data to fit?
        data = collect(transpose(V[Vind][tInds]));

        return t[tInds], data, BCL
end

function getExpData(ind; Vind=1, xInds=64:64, tInds=1:125)

        # get reference data
        basepath = "../../Alessio_Data/2011-05-28_Rec78-103_Pace_Apex/"

        # print some diagnostic notices
        print("Loading data for $(ind)...\t")

        # load in some data
        expStates = loadFile(basepath, ind)

        print("Done.\n")

        # form the space and time masks
        spaceMasks = [spatialMask(state) for state in expStates]
        timeMasks  = [temporalMask(state) for state in expStates]
        mutualSpaceMask = mutualMask(spaceMasks)
        mutualTimeMask = mutualMask(timeMasks)

        # determine mutualSpaceMask centroid for square center
        samplePoints = [[x,x] for x in xInds]
        BCLs = []
        data = []
        ts   = []
        for samplePoint in samplePoints
        	print("Sampling at $(samplePoint) index...\t")
        	t, V = sampleData(expStates, samplePoint, mutualTimeMask, mutualSpaceMask)
        	print("Done\n")

        	push!(ts, t)

		APD, DI, APA = loadAlessioData.analyzeVoltage(t, V[1]; threshold=0.25)
		ii = 1:min(length(APD),length(DI))
		BCL = sum(APD[ii] .+ DI[ii])/length(APD[ii])
		push!(BCLs, BCL)

		# select EPI for data to fit?
		if isempty(data)
			data = collect(transpose(V[Vind][tInds]))
		else
			data = vcat(data, collect(transpose(V[Vind][tInds])))
		end
	end
	t = ts[1]
	BCL = sum(BCLs)/length(BCLs)

	return t[tInds], data, BCL
end

end