text:
struct TsptwGenerator <: AbstractModelGenerator
    n_city::Int
    grid_size::Int # Maximum of positions
    max_tw_gap::Int # Maximum time windows gap allowed between the cities constituing the feasible tour
    max_tw::Int # Maximum time windows upper bound
    pruning::Bool
end

TsptwGenerator(n_city::Int, grid_size::Int, max_tw_gap::Int, max_tw::Int) = TsptwGenerator(n_city, grid_size, max_tw_gap, max_tw, true)

"""
    fill_with_generator!(cpmodel::CPModel, gen::TsptwGenerator)::CPModel

Fill a CPModel with the variables and constraints generated. We fill it directly instead of
creating temporary files for efficiency purpose !

It is the same generator as used in
"Combining Reinforcement Learning and Constraint Programming for Combinatorial Optimization":
Quentin Cappart, Thierry Moisan, Louis-Martin Rousseau, Isabeau Prémont-Schwarz & Andre Cire
https://arxiv.org/abs/2006.01610

Basicaly finds positions with a uniform distributions, then sets the time windows by creating a feasible tour and adding
some randomness by using uniform distributions with gap and the length of the time windows.
"""
function fill_with_generator!(cpmodel::CPModel, gen::TsptwGenerator; seed=nothing, dist = nothing, timeWindows=nothing)
    if !isnothing(seed)
        Random.seed!(seed)
    end

    x_pos = zeros(gen.n_city)
    y_pos = zeros(gen.n_city)

    ### Creating the TSPTW instance
    if isnothing(dist) || isnothing(timeWindows)
        pos_distribution = Uniform(0, gen.grid_size)
        x_pos = rand(pos_distribution, gen.n_city)
        y_pos = rand(pos_distribution, gen.n_city)

        dist = zeros(Int64, gen.n_city, gen.n_city)
        for i in 1:gen.n_city
            for j in 1:gen.n_city
                dist[i, j] = round(sqrt((x_pos[i] - x_pos[j])^2 + (y_pos[i] - y_pos[j])^2))
            end
        end

        timeWindows = zeros(Int64, gen.n_city, 2)
        timeWindows[1, :] = [0 10]

        random_solution = [1, shuffle(Vector(2:gen.n_city))...]

        for i in 2:gen.n_city
            prev_city = random_solution[i-1]
            cur_city = random_solution[i]

            cur_dist = dist[prev_city, cur_city]

            tw_lb_min = timeWindows[prev_city, 1] + cur_dist

            rand_tw_lb = rand(DiscreteUniform(tw_lb_min, tw_lb_min + gen.max_tw_gap))
            rand_tw_ub = rand(DiscreteUniform(rand_tw_lb, rand_tw_lb + gen.max_tw))

            timeWindows[cur_city, :] = [rand_tw_lb rand_tw_ub]
        end
    end

    cpmodel.adhocInfo = dist, timeWindows, hcat(x_pos, y_pos), gen.grid_size

    max_upper_tw = Base.maximum(timeWindows) * 2

    ### Filling the CPModel
    ## Creating variables
    m = [IntSetVar(1, gen.n_city, "m_"*string(i), cpmodel.trailer) for i in 1:gen.n_city] # Remaining cities to visit
    v = [IntVar(1, gen.n_city, "v_"*string(i), cpmodel.trailer) for i in 1:gen.n_city] # Last customer
    t = [IntVar(0, max_upper_tw, "t_"*string(i), cpmodel.trailer) for i in 1:gen.n_city] # Current time
    a = [IntVar(1, gen.n_city, "a_"*string(i), cpmodel.trailer) for i in 1:(gen.n_city-1)] # Action: serving customer a_i at stage i
    c = [IntVar(0, max_upper_tw, "c_"*string(i), cpmodel.trailer) for i in 1:gen.n_city] # Current cost
    total_cost = IntVar(0, max_upper_tw, "total_cost", cpmodel.trailer)
    addVariable!(cpmodel, total_cost; branchable=false)
    for i in 1:gen.n_city
        addVariable!(cpmodel, m[i]; branchable=false)
        addVariable!(cpmodel, v[i]; branchable=false)
        addVariable!(cpmodel, t[i]; branchable=false)
        if i != gen.n_city
            addVariable!(cpmodel, a[i]; branchable=true)
        end
        addVariable!(cpmodel, c[i]; branchable=false)
    end

    ## Intermediaries
    d = [IntVar(0, gen.grid_size * 2, "d_"*string(i), cpmodel.trailer) for i in 1:gen.n_city] # d[v_i, a_i]
    lowers = [IntVar(0, max_upper_tw, "td_"*string(i), cpmodel.trailer) for i in 1:(gen.n_city-1)] # t + d[v_i, a_i]
    lower_ai = [IntVar(0, max_upper_tw, "lower_ai_"*string(i), cpmodel.trailer) for i in 1:(gen.n_city-1)] # timeWindows[i, 1]
    upper_tw_minus_1 = [IntVar(timeWindows[i, 2] - 1, timeWindows[i, 2] - 1, "upper_tw_"*string(i), cpmodel.trailer) for i in 1:gen.n_city] # timeWindows[i, 2] + 1
    one_var = IntVar(1, 1, "one", cpmodel.trailer)
    upper_ai = [IntVar(0, max_upper_tw, "upper_ai_"*string(i), cpmodel.trailer) for i in 1:(gen.n_city-1)] # timeWindows[a_i, 2]
    j_index = [IntVarViewMul(one_var, j, "index_"*string(j)) for j in 1:gen.n_city]

    still_time = Array{BoolVar, 2}(undef, (gen.n_city, gen.n_city))
    j_in_m_i = Array{BoolVar, 2}(undef, (gen.n_city, gen.n_city))
    for i in 1:gen.n_city
        for j in 1:gen.n_city
            still_time[i, j] = BoolVar("s_t_"*string(i)*"_"*string(j), cpmodel.trailer) # t_i < upper_bound[j]
            j_in_m_i[i, j] = BoolVar(string(j)*"_in_m_"*string(i), cpmodel.trailer) # t_i < upper_bound[j]
        end
    end

    addVariable!(cpmodel, one_var; branchable=false)
    for i in 1:gen.n_city
        addVariable!(cpmodel, d[i]; branchable=false)
        addVariable!(cpmodel, upper_tw_minus_1[i]; branchable=false)
        if gen.pruning
            addVariable!(cpmodel, j_index[i]; branchable=false)
        end
        if i != gen.n_city
            addVariable!(cpmodel, lower_ai[i]; branchable=false)
            addVariable!(cpmodel, upper_ai[i]; branchable=false)
            addVariable!(cpmodel, lowers[i]; branchable=false)
        end
        for j in 1:gen.n_city
            if gen.pruning
                addVariable!(cpmodel, j_in_m_i[i, j]; branchable=false)
                addVariable!(cpmodel, still_time[i, j]; branchable=false)
            end
        end
    end

    ## Constraints
    # Initialization
    SeaPearl.addConstraint!(cpmodel, EqualConstant(t[1], 0, cpmodel.trailer))
    SeaPearl.addConstraint!(cpmodel, EqualConstant(v[1], 1, cpmodel.trailer))
    SeaPearl.addConstraint!(cpmodel, EqualConstant(c[1], 0, cpmodel.trailer))
    SeaPearl.addConstraint!(cpmodel, SetEqualConstant(m[1], Set{Int}(collect(2:gen.n_city)), cpmodel.trailer))

    # Variable definition
    for i in 1:(gen.n_city - 1)
        # m[i+1] = m[i] \ a[i]
        SeaPearl.addConstraint!(cpmodel, SetDiffSingleton(m[i+1], m[i], a[i], cpmodel.trailer))

        # v[i+1] = a[i]
        SeaPearl.addConstraint!(cpmodel, Equal(v[i+1], a[i], cpmodel.trailer))

        # t[i+1] = max(lowers[i], lower_ai[i])
        SeaPearl.addConstraint!(cpmodel, BinaryMaximumBC(t[i+1], lowers[i], lower_ai[i], cpmodel.trailer))

        # c[i + 1] = c[i] + d[i]
        SeaPearl.addConstraint!(cpmodel, SumToZero(AbstractIntVar[c[i], d[i], IntVarViewOpposite(c[i+1], "-c_"*string(i+1))], cpmodel.trailer))

        # upper_ai = timeWindows[a_i, 2]
        SeaPearl.addConstraint!(cpmodel, Element1D(timeWindows[:, 2], a[i], upper_ai[i], cpmodel.trailer))

        # lower_ai = timeWindows[a_i, 1]
        SeaPearl.addConstraint!(cpmodel, Element1D(timeWindows[:, 1], a[i], lower_ai[i], cpmodel.trailer))

        # d[i] = dist[v[i], a[i]]
        SeaPearl.addConstraint!(cpmodel, Element2D(dist, v[i], a[i], d[i], cpmodel.trailer))
        # lowers[i] = t[i] + d[i]
        SeaPearl.addConstraint!(cpmodel, SumToZero(AbstractIntVar[t[i], d[i], IntVarViewOpposite(lowers[i], "-td_"*string(i))], cpmodel.trailer))
    end
    # d[n] = dist[a[n-1], 1]
    SeaPearl.addConstraint!(cpmodel, Element2D(dist, a[gen.n_city-1], one_var, d[gen.n_city], cpmodel.trailer))
    # total_cost = c[n] + d[n]
    SeaPearl.addConstraint!(cpmodel, SumToZero(AbstractIntVar[c[gen.n_city], d[gen.n_city], IntVarViewOpposite(total_cost, "-total_cost")], cpmodel.trailer))

    # Validity constraints
    for i in 1:(gen.n_city - 1)
        # a[i] ∈ m[i]
        SeaPearl.addConstraint!(cpmodel, InSet(a[i], m[i], cpmodel.trailer))

        # lowers[i] <= upper_ai
        SeaPearl.addConstraint!(cpmodel, LessOrEqual(lowers[i], upper_ai[i], cpmodel.trailer))
    end

    # Pruning constraints
    if gen.pruning
        for i in 1:gen.n_city
            for j in 1:gen.n_city
                # still_time[i, j] = t[i] < upper_tw[j]
                SeaPearl.addConstraint!(cpmodel, isLessOrEqual(still_time[i, j], t[i], upper_tw_minus_1[j], cpmodel.trailer))

                # j_in_m_i[i, j] = j_index[j] ∈ m[i]
                SeaPearl.addConstraint!(cpmodel, ReifiedInSet(j_index[j], m[i], j_in_m_i[i, j], cpmodel.trailer))

                # t[i] >= upper[j] ⟹ j ∉ m[i]
                # ≡ t[i] < upper[j] ⋁ j ∉ m[i]
                # ≡ still_time[i, j] ⋁ ¬j_in_m_i[i, j]
                SeaPearl.addConstraint!(cpmodel, BinaryOr(still_time[i, j], BoolVarViewNot(j_in_m_i[i, j], "¬"*string(j)*"_in_m_"*string(i)), cpmodel.trailer))
            end
        end
    end

    # Objective function: min total_cost
    SeaPearl.addObjective!(cpmodel,total_cost)
    return dist, timeWindows
end

"""
    find_tsptw_dist_matrix(cpmodel::CPModel)

Find the distance matrix of a TSPTW instance generated by TsptwGenerator.
"""
# FIXME this is a very convoluted way to retrieve the distance matrix
function find_tsptw_dist_matrix(cpmodel::CPModel)
    for constraint in cpmodel.constraints
        if isa(constraint, Element2D)
            x, y = size(constraint.matrix)
            if x == y
                dist = constraint.matrix
                return dist
            end
        end
    end
    throw(ErrorException("The model given to find_tsptw_dist_matrix does not seem to be a TSPTW instance."))
end


text:
function get_boundary_filt(basin)

    mx = maximum(basin)
    mn = minimum(basin)
    # Kernel
    w = centered([1 1 1; 1 -8 1 ; 1 1 1])
    #w = centered([0 1 0; 1 -4 1 ; 0 1 0])

    # replicate for boundary conditions
    res = imfilter(basin,w,"replicate")

    # segmentation: take a low threshold in order to not loose some structure.
    # This detection can be improved.
    #res = abs.(res) .> (mx-mn)
    res = abs.(res) .> 0.
    return res
end

function get_list(basin,xg ,yg)

    num_att = length(unique(basin))

    # v_list empty container that will receive the different coordinates corresponding to each basins,
    # Each element of this v_set is an array of coordinates of different sizes.
    v_list = Vector{SVector{2,Float64}}()
    v_set = Vector{typeof(v_list)}()

    # original boundary
    bnd = get_boundary_filt(basin)
    I1=findall(bnd .== 1);
    for ind in I1; push!(v_list , [xg[ind[1]]; yg[ind[2]]]); end
    #@show v_list
    push!(v_set,v_list)

    # Merging !!
    basin_i_j = zeros(Int8,size(basin))
    for k in unique(basin)
        v_list = Vector{Vector{Float64}}()
        fill!(basin_i_j, zero(0)) # erase matrix
        I = findall(basin .== k)
        [basin_i_j[x] = 1 for x in I];
        # compute boundary
        bnd = get_boundary_filt(basin_i_j)
        I1=findall(bnd .== 1);
        # Get coordinates lists from matrix
        #push!(v_list , hcat([[xg[ind[1]]; yg[ind[2]]] for ind in I1 ]...))
        for ind in I1; push!(v_list , [xg[ind[1]]; yg[ind[2]]]); end
        push!(v_set,v_list)
    end

    return v_set
end

"""
    detect_wada_merge_method(xg,yg, basin)
The algorithm gives the maximum and minimum Haussdorff distances between merged basins. These two distances can help to decide if the basin has the Wada property.

[A. Daza, A. Wagemakers and M. A. F. Sanjuán, Ascertaining when a basin is Wada: the merging method, Sci. Rep., 8, 9954 (2018)]

## Arguments
* `basin` : the matrix containing the information of the basin.
* `xg`, `yg` : 1-dim range vector that defines the grid of the initial conditions to test.

## Example
```
max_dist,min_dist = detect_wada_merge_method(basin,xg,yg)
# grid resolution
epsilon = xg[2]-xg[1]
# if dmax is large then this is not Wada
@show dmax = max_dist/epsilon
@show dmin = min_dist/epsilon
```

"""
function detect_wada_merge_method(xg,yg,bsn::BasinInfo)
    ind  = findall(iseven.(bsn.basin) .== true)
    basin_test = deepcopy(bsn.basin)
    for k in ind; basin_test[k] =basin_test[k]+1; end
    return detect_wada_merge_method(xg,yg,basin_test)
end

function detect_wada_merge_method(xg,yg,basin)

    num_att = length(unique(basin))

    v_list=get_list(basin, xg, yg)

    # compute distances using combbinations of 2 elements from a collection
    ind = combinations(1:num_att,2)

    min_dist = Inf;
    max_dist = 0
    for (k, patron) in enumerate(ind)

       v01 = v_list[patron[1]]
       v02 = v_list[patron[2]]

       hd = haussdorff_dist(v01,v02)

       if hd  > max_dist
           max_dist = hd
       end

       if hd  < min_dist
           min_dist = hd
       end

  end

   return max_dist, min_dist
end

function haussdorff_dist(v01,v02)
       kdtree1 = KDTree(v01)
       idxs12, dists12 = knn(kdtree1, v02, 1, true)
       kdtree2 = KDTree(v02)
       idxs21, dists21 = knn(kdtree2, v01, 1, true)
       max12 = maximum(dists12);
       max21 = maximum(dists21);
       hd =max(max12[1],max21[1])
       return hd
end

mutable struct ds_info{I}
    bsn_nfo::BasinInfo # basin info for BA routine
    integ::I               # integrator
end

function reset_point_data!(ds_nfo)
    reset_bsn_nfo!(ds_nfo.bsn_nfo)
end

"""
    detect_wada_grid_method(integ, bsn_nfo::BasinInfo; max_iter=10)
The algorithm test for Wada basin in a dynamical system. It uses the dynamical system to look if all the atractors are represented in the boundary.

[A. Daza, A. Wagemakers, M. A. F. Sanjuán and J. A. Yorke, Testing for Basins of Wada, Sci. Rep., 5, 16579 (2015)]

## Arguments
* `integ` : the matrix containing the information of the basin.
* `bsn_nfo` : structure that holds the information of the basin as well as the map function. This structure is set when the basin is first computed with `basin_stroboscopic_map` or `basin_poincare_map`.

## Keyword arguments
* `max_iter` : set the maximum depth of subdivisions to look for an atractor. The number of points doubles at each step.

"""
function detect_wada_grid_method(integ, bsn_nfo::BasinInfo; max_iter=10)

   ds_nfo = ds_info(bsn_nfo, integ)
   num_att = bsn_nfo.Na

   if findfirst(x->x==-1, bsn_nfo.basin) != nothing
       @error "The basin contains escapes or undefined attractors, cannot test for Wada"
       return nothing
   end

   # helper function to obtain coordinates
   index_to_coord(p) = [bsn_nfo.xg[p[1]], bsn_nfo.yg[p[2]]]

   # We remove the atractors for this computation
   tmp_bsn = deepcopy(bsn_nfo.basin)
   ind  = findall(iseven.(tmp_bsn) .== true)
   [tmp_bsn[k] = tmp_bsn[k]+1 for k in ind ]

   # obtain points in the boundary
   bnd = get_boundary_filt(tmp_bsn)
   p1_ind = findall(bnd .> 0)

   # initialize empty array of indices and collection of empty sets of colors
   clr_mat = [Set{Int16}() for i=1:length(p1_ind)]
   p2_ind = typeof(p1_ind)(undef, length(p1_ind))
   W = zeros(num_att, max_iter)

   # Initialize matrices (step 1)
   for (k,p1) in enumerate(p1_ind)
       p2, nbgs = get_neighbor_and_colors(tmp_bsn, [p1[1], p1[2]])
       if length(nbgs) > 1
           # keep track of different colors and neighbor point
           push!(clr_mat[k],nbgs...)
           p2_ind[k] = p2
       else
           println("Not in the boundary: weird...")
       end
       W[length(clr_mat[k]),1] += 1
   end

   # Do the iteration!
   for n = 2:max_iter
       for k in 1:length(p1_ind)
           pc1 = index_to_coord(p1_ind[k])
           pc2 = index_to_coord(p2_ind[k])
           # update number of colors
           clr_mat[k]=divide_and_test_W(ds_nfo, pc1, pc2, n, clr_mat[k], num_att)
           # update W matrix
           W[length(clr_mat[k]),n] += 1
       end
       @show W

       # Stopping criterion: if W[Na] in % increases less than ε  then stop or if we have more than 95% boxes in Na
       if  (abs(W[num_att,n] - W[num_att,n-1]) <0.01*W[num_att,n] || W[num_att,n] == 0.0 || W[num_att,n]/sum(W[:,n])>0.95) &&  n > num_att # make at least num_att iterations for stats
           return W[:,n]./sum(W[:,1])
       end

   end

   return W[:,end]./sum(W[:,1])
end

function get_neighbor_and_colors(basin, p)
    radius = 1
    n=p[1]
    m=p[2]
    v=Int16[]
    p2=CartesianIndex(-1,-1)
    # check neihbors and collect basin colors
    for k=n-radius:n+radius, l=m-radius:m+radius
        try
            push!(v,basin[k,l])
            if k != n || l != m
                if basin[n,m] != basin[k,l]
                    p2=CartesianIndex(k,l)
                end
            end
        catch
            #println("error in push!")
        end
    end
    return p2,unique(v)
end

function divide_and_test_W(ds_nfo, p1, p2, nstep, clrs, Na)

    if length(clrs)  == Na
        return clrs
    end

    lerp(t) = (1-t)*p1 + t*p2 # linear interpolation function

    # generates coordinates for initial conditions to test with linear interpolation
    npts = 2^(nstep)
    vecs = Set([lerp(ratio) for ratio in range(0,1, step=1/npts)])  # create a set with the points interspeded
    npts = 2^(nstep-1)
    vecs_before = Set([lerp(ratio) for ratio in range(0,1, step=1/npts)])  # remove points already computed in the previous step
    pnt_to_test = setdiff(vecs,vecs_before) # get the points we have to test

    # get colors and update color set for this box!
    for pnt in pnt_to_test
        clr = get_color_point!(ds_nfo.bsn_nfo, ds_nfo.integ, pnt)
        push!(clrs, clr)
        if length(clrs)  == Na
            break
        end
    end

    return clrs

end

ext: jl


text:
include("run_pd_hyak.jl")
datafile = "../../data/T1_Spitzer_data.jld2"
foutput = "T1_pd_MCMC_nw50_ns50_limbprior_001.jld2"
numwalkers = 50
burnin = 1
thinning = 1
astep = 2.0
nsteps = 50
run_pd(datafile,foutput,numwalkers,burnin,thinning,astep,nsteps;nout=10000,limbprior=true)


text:
################################################################################
# TODO: replace this by making kernels @functor from Functors.jl
# this section automates the construction of parameter-dependent composite kernels
# parameters function returns either scalar or recursive vcat of parameter vectors
# useful for optimization algorithms which require the parameters in vector form
parameters(::Any) = []
parameters(::AbstractKernel{T}) where {T} = T == Union{} ? zeros(0) : zeros(T, 0)
nparameters(::Any) = 0

# checks if θ has the correct number of parameters to initialize a kernel of typeof(k)
function checklength(k::AbstractKernel, θ::AbstractVector)
    nt = length(θ)
    np = nparameters(k)
    if nt ≠ np
        throw(DimensionMismatch("length(θ) = $nt ≠ $np = nparameters(k)"))
    end
    return nt
end

# thanks to ffevotte in https://discourse.julialang.org/t/how-to-call-constructor-of-parametric-family-of-types-efficiently/38503/5
stripped_type(x) = stripped_type(typeof(x))
stripped_type(typ::DataType) = Base.typename(typ).wrapper
# fallback for zero-parameter kernels
function Base.similar(k::AbstractKernel, θ::AbstractVector)
    n = checklength(k, θ)
    # kernel = eval(Meta.parse(string(typeof(k).name)))
    kernel = stripped_type(k)
    if n == 0
        kernel()
    elseif n == 1
        kernel(θ[1])
    else
        kernel(θ)
    end
end

Base.similar(k::AbstractKernel, θ::Number) = similar(k, [θ])

# constructs a similar kernel to k with hyperparameters θ
function _similar_helper(k, θ)
    checklength(k, θ)
    args = Vector{AbstractKernel}(undef, length(k.args))
    for (i, k) in enumerate(k.args)
        n = nparameters(k)
        args[i] = similar(k, @view(θ[1:n]))
        θ = @view(θ[n+1:end])
    end
    return args
end


text:
"""
    (traces, log_norm_weights, lml_est) = importance_sampling(model::GenerativeFunction,
        model_args::Tuple, observations::ChoiceMap, num_samples::Int, verbose=false)

    (traces, log_norm_weights, lml_est) = importance_sampling(model::GenerativeFunction,
        model_args::Tuple, observations::ChoiceMap,
        proposal::GenerativeFunction, proposal_args::Tuple,
        num_samples::Int, verbose=false)

Run importance sampling, returning a vector of traces with associated log weights.

The log-weights are normalized.
Also return the estimate of the marginal likelihood of the observations (`lml_est`).
The observations are addresses that must be sampled by the model in the given model arguments.
The first variant uses the internal proposal distribution of the model.
The second variant uses a custom proposal distribution defined by the given generative function.
All addresses of random choices sampled by the proposal should also be sampled by the model function.
Setting `verbose=true` prints a progress message every sample.
"""
function importance_sampling(model::GenerativeFunction{T,U}, model_args::Tuple,
                             observations::ChoiceMap,
                             num_samples::Int, verbose=false) where {T,U}
    traces = Vector{U}(undef, num_samples)
    log_weights = Vector{Float64}(undef, num_samples)
    for i=1:num_samples
        verbose && println("sample: $i of $num_samples")
        (traces[i], log_weights[i]) = generate(model, model_args, observations)
    end
    log_total_weight = logsumexp(log_weights)
    log_ml_estimate = log_total_weight - log(num_samples)
    log_normalized_weights = log_weights .- log_total_weight
    return (traces, log_normalized_weights, log_ml_estimate)
end

function importance_sampling(model::GenerativeFunction{T,U}, model_args::Tuple,
                             observations::ChoiceMap,
                             proposal::GenerativeFunction, proposal_args::Tuple,
                             num_samples::Int, verbose=false) where {T,U}
    traces = Vector{U}(undef, num_samples)
    log_weights = Vector{Float64}(undef, num_samples)
    for i=1:num_samples
        verbose && println("sample: $i of $num_samples")
        (proposed_choices, proposal_weight, _) = propose(proposal, proposal_args)
        constraints = merge(observations, proposed_choices)
        (traces[i], model_weight) = generate(model, model_args, constraints)
        log_weights[i] = model_weight - proposal_weight
    end
    log_total_weight = logsumexp(log_weights)
    log_ml_estimate = log_total_weight - log(num_samples)
    log_normalized_weights = log_weights .- log_total_weight
    return (traces, log_normalized_weights, log_ml_estimate)
end

"""
    (trace, lml_est) = importance_resampling(model::GenerativeFunction,
        model_args::Tuple, observations::ChoiceMap, num_samples::Int,
        verbose=false)

    (traces, lml_est) = importance_resampling(model::GenerativeFunction,
        model_args::Tuple, observations::ChoiceMap,
        proposal::GenerativeFunction, proposal_args::Tuple,
        num_samples::Int, verbose=false)

Run sampling importance resampling, returning a single trace.

Unlike `importance_sampling`, the memory used constant in the number of samples.

Setting `verbose=true` prints a progress message every sample.
"""
function importance_resampling(model::GenerativeFunction{T,U}, model_args::Tuple,
                               observations::ChoiceMap,
                               num_samples::Int; verbose=false)  where {T,U,V,W}
    (model_trace::U, log_weight) = generate(model, model_args, observations)
    log_total_weight = log_weight
    for i=2:num_samples
        verbose && println("sample: $i of $num_samples")
        (cand_model_trace, log_weight) = generate(model, model_args, observations)
        log_total_weight = logsumexp(log_total_weight, log_weight)
        if bernoulli(exp(log_weight - log_total_weight))
            model_trace = cand_model_trace
        end
    end
    log_ml_estimate = log_total_weight - log(num_samples)
    return (model_trace::U, log_ml_estimate::Float64)
end

function importance_resampling(model::GenerativeFunction{T,U}, model_args::Tuple,
                               observations::ChoiceMap,
                               proposal::GenerativeFunction{V,W}, proposal_args::Tuple,
                               num_samples::Int; verbose=false)  where {T,U,V,W}
    (proposal_choices, proposal_weight, _) = propose(proposal, proposal_args)
    constraints = merge(observations, proposal_choices)
    (model_trace::U, model_weight) = generate(model, model_args, constraints)
    log_total_weight = model_weight - proposal_weight
    for i=2:num_samples
        verbose && println("sample: $i of $num_samples")
        (proposal_choices, proposal_weight, _) = propose(proposal, proposal_args)
        constraints = merge(observations, proposal_choices)
        (cand_model_trace, model_weight) = generate(model, model_args, constraints)
        log_weight = model_weight - proposal_weight
        log_total_weight = logsumexp(log_total_weight, log_weight)
        if bernoulli(exp(log_weight - log_total_weight))
            model_trace = cand_model_trace
        end
    end
    log_ml_estimate = log_total_weight - log(num_samples)
    return (model_trace::U, log_ml_estimate::Float64)
end

export importance_sampling, importance_resampling


text:
@testset "54.spiral-matrix.jl" begin
    @test spiral_order([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) == [1, 2, 3, 6, 9, 8, 7, 4, 5]
    @test spiral_order([[1], [2], [3], [4], [5], [6], [7], [8], [9], [10]]) ==
          [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
    @test spiral_order([[7], [9], [6]]) == [7, 9, 6]
    @test spiral_order([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]) ==
          [1, 2, 3, 4, 8, 12, 11, 10, 9, 5, 6, 7]
end


text:
@testset "Extraction" begin
# Period subset
d = DateTime(1961,1,1):Day(1):DateTime(1990,12,31)
Random.seed!(42)
data = randn(2, 2, 10957)
axisdata = AxisArray(data, Axis{:lon}(1:2), Axis{:lat}(1:2), Axis{:time}(d))
dimension_dict = Dict(["lon" => "lon", "lat" => "lat"])
C = ClimateTools.ClimGrid(axisdata, variable = "tasmax", dimension_dict=dimension_dict)

# When startmonth < endmonth
D = resample(C, 3, 6)
@test D[1].data[:,:,1] == C[1].data[:,:,60]
@test length(D[1][Axis{:time}]) == 3660 && Dates.month(D[1][Axis{:time}][end]) == 6 && Dates.month(D[1][Axis{:time}][1]) == 3

# When startmonth > endmonth
D = resample(C, 10, 2)
@test D[1].data[:,:,1] == C[1].data[:,:,1] && D[1].data[:,:,60] == C[1].data[:,:,274]
@test length(D[1][Axis{:time}]) == 4537 && Dates.month.(D[1][Axis{:time}][end]) == 12 && Dates.month.(D[1][Axis{:time}][1]) == 1
@test unique(Dates.month.(D[1][Axis{:time}][:])) == [1,2,10,11,12]

D = resample(C, "DJF") # hardcoded seasons
@test D[1].data[:,:,60] == C[1].data[:,:,365-30]
@test length(D[1][Axis{:time}]) == 2707 && Dates.month(D[1][Axis{:time}][1]) == 1 && Dates.month(D[1][Axis{:time}][end]) == 12
@test unique(Dates.month.(D[1][Axis{:time}][:])) == [1,2,12]

D = resample(C, "MAM") # hardcoded seasons
@test D[1].data[:,:,1] == C[1].data[:,:,60]
@test length(D[1][Axis{:time}]) == 2760 && Dates.month(D[1][Axis{:time}][1]) == 3 && Dates.month(D[1][Axis{:time}][end]) == 5
@test unique(Dates.month.(D[1][Axis{:time}][:])) == [3,4,5]

D = resample(C, "JJA") # hardcoded seasons
@test D[1].data[:,:,1] == C[1].data[:,:,152]
@test length(D[1][Axis{:time}]) == 2760 && Dates.month(D[1][Axis{:time}][1]) == 6 && Dates.month(D[1][Axis{:time}][end]) == 8
@test unique(Dates.month.(D[1][Axis{:time}][:])) == [6,7,8]

D = resample(C, "SON") # hardcoded seasons
@test D[1].data[:,:,1] == C[1].data[:,:,244]
@test length(D[1][Axis{:time}]) == 2730 && Dates.month(D[1][Axis{:time}][1]) == 9 && Dates.month(D[1][Axis{:time}][end]) == 11
@test unique(Dates.month.(D[1][Axis{:time}][:])) == [9,10,11]

@test_throws ErrorException D = resample(C, "df")

maps = ["rotated_pole", "rotated_mercator", "crs", "lambert_conformal_conic","rotated_latitude_longitude", "None"]

for imap in maps
  K = ["time", imap]
  if imap != "None"
    @test ClimateTools.get_mapping(K) == imap
  else
    @test ClimateTools.get_mapping(K) == "Regular_longitude_latitude"
  end
end

end

@testset "Exportation" begin

    filenc = joinpath(dirname(@__FILE__), "data", "sresa1b_ncar_ccsm3-example.nc")
    C = load(filenc, "tas")

    status = write(C, "test.nc")
    @test status == true
    run(`rm test.nc`)

    status = write(C, "test")
    @test status == true
    run(`rm test.nc`)

    # Custom ClimGrid
    data = Float32.(rand(2,2, 365))
    timev = DateTimeNoLeap(2001, 01, 01):Day(1):DateTimeNoLeap(2001, 12, 31)
    rlon = collect(range(20.0, stop=21.0))
    rlat = collect(range(40.0, stop=41.0))
    longrid, latgrid = ndgrid(rlon, rlat)
    variable = "pr"
    A = AxisArray(data, Axis{:rlon}(rlon), Axis{:rlat}(rlat), Axis{:time}(timev))
    dimension_dict = Dict(["lat" => "rlat", "lon" => "rlon"])

    timeattrib = Dict(["units" => "days since 2001-1-1",
    "calendar" => "365_day",
    "axis" => "T",
    "standard_name" => "time"])

    varattrib = Dict(["standard_name" => "precipitation_flux",
    "units" => "kg m-2 s-1"])

    globalattribs = Dict(["title"                 => "CanESM2 model output prepared for CMIP5 historical",
  "branch_time"           => 135415.0,
  "physics_version"       => 1,
  "tracking_id"           => "97542c1e-3841-4a49-bc21-76d6edb9fa1d",
  "branch_time_YMDH"      => "2221:01:01:00",
  "modeling_realm"        => "atmos",
  "creation_date"         => "2011-04-09T02:32:12Z",
  "history"               => "",
  "table_id"              => "Table day (28 March 2011) f9d6cfec5981bb8be1801b35a81002f0",
  "CCCma_runid"           => "IDE",
  "project_id"            => "CMIP5",
  "experiment"            => "historical",
  "forcing"               => "GHG,Oz,SA,BC,OC,LU,Sl,Vl (GHG includes CO2,CH4,N2O,CFC11,effective CFC12)",
  "contact"               => "cccma_info@ec.gc.ca",
  "Conventions"           => "CF-1.4",
  "institute_id"          => "CCCma",
  "parent_experiment_rip" => "r1i1p1",
  "experiment_id"         => "historical",
  "model_id"              => "CanESM2",
  "frequency"             => "day"])

  grid_mapping = Dict(["grid_mapping" => "Rotated_pole"])

  C = ClimGrid(A, variable=variable, longrid=longrid, latgrid=latgrid, dimension_dict=dimension_dict, timeattrib=timeattrib, varattribs=varattrib, globalattribs=globalattribs, grid_mapping=grid_mapping)

  status = write(C, "test.nc")
  @test status == true
  run(`rm test.nc`)

  # grid_mapping = Dict(["grid_mapping_name" => "Rotated_pole"])
  #
  # C = ClimGrid(A, variable=variable, longrid=longrid, latgrid=latgrid, dimension_dict=dimension_dict, timeattrib=timeattrib, varattribs=varattrib, globalattribs=globalattribs, grid_mapping=grid_mapping)
  #
  # status = write(C, "test.nc")
  # @test status == true
  # run(`rm test.nc`)

end


text:
@testset "HurdleModel" begin
  Random.seed!(0)
  m1 = HurdleModel(-Inf, MixtureModel(SkewT.([0, 3], 2, 4, -10)), .3)
  m2 = HurdleModel(-Inf, MixtureModel(Normal.([0, 3])), .3)
  for m in (m1, m2)
    rand(m)
    @test mean([rand(m) for _ in 1:10000] .=== -Inf) ≈ m1.prob atol=0.01
    @test pdf(m, 10.0) ≈ exp(logpdf(m, 10)) atol=0.01
    @test pdf(m, -Inf) ≈ exp(logpdf(m, -Inf)) atol=0.01
  end
end


text:
module TestFirmnessInSequence

using PSCOPF

using Test
using Dates
using DataStructures

@testset verbose=true "test_firmness_in_sequence" begin

    #=
    "D" indicates the decision time (TO_DECIDE state)
    Before "D", we have a FREE state
    After "D", we have a DECIDED state
          10h                13h             13h30    13h45     14h             14h30
          ECH1               ECH2            ECH3     ECH4      TS1              TS2
            |                  |               |        |        |                |
fuel_1_0    |                  |               |        |        |                |
commitment  |                  |               |        |        |                |
TS1:        D<------------------------------DMO----------------->|                |
TS2:        D         <---------------------------------DMO---------------------->|
power                          |               |        |        |                |
TS1:                           D  <----------------DP----------->|                |
TS2:                                           D     <--------------DP----------->|
                                               |        |        |                |
wind_1_0                                       |        |        |                |
commitment                                     |        |        |                |
= power                                        |        |        |                |
TS1:                                           D     <--DP=DMO-->|                |
TS2:                                                    D        X    <--DP=DMO-->|
                                                                 ^we cannot decide at this point
=#

    @testset "test_firmness_in_sequence" begin
        network = PSCOPF.Networks.Network()
        PSCOPF.Networks.add_new_bus!(network, "bus_1")
        PSCOPF.Networks.add_new_generator_to_bus!(network, "bus_1", "wind_1_0", PSCOPF.Networks.LIMITABLE,
                                                0., 100., 0., 0.,
                                                Dates.Second(20*60), Dates.Second(20*60))
        PSCOPF.Networks.add_new_generator_to_bus!(network, "bus_1", "fuel_1_0", PSCOPF.Networks.PILOTABLE,
                                                10., 100., 0., 0.,
                                                Dates.Second(4*60*60), Dates.Second(50*60))

        TS = [DateTime("2015-01-01T14:00:00"),
            DateTime("2015-01-01T14:30:00")]
        ECH = [DateTime("2015-01-01T10:00:00"),
            DateTime("2015-01-01T13:00:00"),
            DateTime("2015-01-01T13:30:00"),
            DateTime("2015-01-01T13:45:00")]
        firmness_history = Dict{DateTime,PSCOPF.Firmness}()

        struct MockRunnable <: PSCOPF.AbstractRunnable
        end
        function PSCOPF.run(runnable::MockRunnable, ech, firmness, TS, context::PSCOPF.AbstractContext)
            firmness_history[ech] = firmness
            return nothing
        end

        sequence = PSCOPF.Sequence(SortedDict(
            ECH[1]     => [MockRunnable()],
            ECH[2]     => [MockRunnable()],
            ECH[3]     => [MockRunnable()],
            ECH[4]     => [MockRunnable()]
        ))

        mode = PSCOPF.ManagementMode("test_firmness_in_sequence", Dates.Minute(0))
        exec_context = PSCOPF.PSCOPFContext(network, TS, mode)

        PSCOPF.run!(exec_context, sequence, check_context=false)

        @test length(firmness_history) == 4 # one for each executed step

        #ECH1
        firmness = firmness_history[ECH[1]]
        @test PSCOPF.get_commitment_firmness(firmness, "fuel_1_0", TS[1]) == PSCOPF.TO_DECIDE
        @test PSCOPF.get_commitment_firmness(firmness, "fuel_1_0", TS[2]) == PSCOPF.TO_DECIDE
        @test PSCOPF.get_power_level_firmness(firmness, "fuel_1_0", TS[1]) == PSCOPF.FREE
        @test PSCOPF.get_power_level_firmness(firmness, "fuel_1_0", TS[2]) == PSCOPF.FREE
        @test ismissing(PSCOPF.get_commitment_firmness(firmness, "wind_1_0", TS[1]))
        @test ismissing(PSCOPF.get_commitment_firmness(firmness, "wind_1_0", TS[2]))
        @test PSCOPF.get_power_level_firmness(firmness, "wind_1_0", TS[1]) == PSCOPF.FREE
        @test PSCOPF.get_power_level_firmness(firmness, "wind_1_0", TS[2]) == PSCOPF.FREE

        #ECH2
        firmness = firmness_history[ECH[2]]
        @test PSCOPF.get_commitment_firmness(firmness, "fuel_1_0", TS[1]) == PSCOPF.DECIDED
        @test PSCOPF.get_commitment_firmness(firmness, "fuel_1_0", TS[2]) == PSCOPF.DECIDED
        @test PSCOPF.get_power_level_firmness(firmness, "fuel_1_0", TS[1]) == PSCOPF.TO_DECIDE
        @test PSCOPF.get_power_level_firmness(firmness, "fuel_1_0", TS[2]) == PSCOPF.FREE
        @test ismissing(PSCOPF.get_commitment_firmness(firmness, "wind_1_0", TS[1]))
        @test ismissing(PSCOPF.get_commitment_firmness(firmness, "wind_1_0", TS[2]))
        @test PSCOPF.get_power_level_firmness(firmness, "wind_1_0", TS[1]) == PSCOPF.FREE
        @test PSCOPF.get_power_level_firmness(firmness, "wind_1_0", TS[2]) == PSCOPF.FREE

        #ECH3
        firmness = firmness_history[ECH[3]]
        @test PSCOPF.get_commitment_firmness(firmness, "fuel_1_0", TS[1]) == PSCOPF.DECIDED
        @test PSCOPF.get_commitment_firmness(firmness, "fuel_1_0", TS[2]) == PSCOPF.DECIDED
        @test PSCOPF.get_power_level_firmness(firmness, "fuel_1_0", TS[1]) == PSCOPF.DECIDED
        @test PSCOPF.get_power_level_firmness(firmness, "fuel_1_0", TS[2]) == PSCOPF.TO_DECIDE
        @test ismissing(PSCOPF.get_commitment_firmness(firmness, "wind_1_0", TS[1]))
        @test ismissing(PSCOPF.get_commitment_firmness(firmness, "wind_1_0", TS[2]))
        @test PSCOPF.get_power_level_firmness(firmness, "wind_1_0", TS[1]) == PSCOPF.TO_DECIDE
        @test PSCOPF.get_power_level_firmness(firmness, "wind_1_0", TS[2]) == PSCOPF.FREE

        #ECH4
        firmness = firmness_history[ECH[4]]
        @test PSCOPF.get_commitment_firmness(firmness, "fuel_1_0", TS[1]) == PSCOPF.DECIDED
        @test PSCOPF.get_commitment_firmness(firmness, "fuel_1_0", TS[2]) == PSCOPF.DECIDED
        @test PSCOPF.get_power_level_firmness(firmness, "fuel_1_0", TS[1]) == PSCOPF.DECIDED
        @test PSCOPF.get_power_level_firmness(firmness, "fuel_1_0", TS[2]) == PSCOPF.DECIDED
        @test ismissing(PSCOPF.get_commitment_firmness(firmness, "wind_1_0", TS[1]))
        @test ismissing(PSCOPF.get_commitment_firmness(firmness, "wind_1_0", TS[2]))
        @test PSCOPF.get_power_level_firmness(firmness, "wind_1_0", TS[1]) == PSCOPF.DECIDED
        @test PSCOPF.get_power_level_firmness(firmness, "wind_1_0", TS[2]) == PSCOPF.TO_DECIDE
    end

end

end #module

max_stars_repo_name: klorel/pscopf

max_issues_repo_name: klorel/pscopf

max_forks_repo_name: klorel/pscopf


text:
export Dense,
       outer,
       ⊗

#
# Dense storage
#

struct Dense{T} <: TensorStorage
  data::Vector{T}
  Dense{T}(data) where {T} = new{T}(convert(Vector{T},data))
  Dense{T}() where {T} = new{T}(Vector{T}())
end

# Convenient functions for Dense storage type
Base.@propagate_inbounds Base.getindex(D::Dense,i::Integer) = data(D)[i]
Base.@propagate_inbounds Base.setindex!(D::Dense,v,i::Integer) = (data(D)[i] = v)

Base.similar(D::Dense{T}) where {T} = Dense{T}(similar(data(D)))

# TODO: make this just take Int, the length of the data
Base.similar(D::Dense{T},dims) where {T} = Dense{T}(similar(data(D),dim(dims)))

# TODO: make this just take Int, the length of the data
Base.similar(::Type{Dense{T}},dims) where {T} = Dense{T}(similar(Vector{T},dim(dims)))

Base.similar(D::Dense,::Type{T}) where {T} = Dense{T}(similar(data(D),T))
Base.copy(D::Dense{T}) where {T} = Dense{T}(copy(data(D)))
Base.copyto!(D1::Dense,D2::Dense) = copyto!(data(D1),data(D2))

Base.fill!(D::Dense,v) = fill!(data(D),v)

Base.zeros(::Type{Dense{T}},dim::Int) where {T} = Dense{T}(zeros(T,dim))

# convert to complex
# TODO: this could be a generic TensorStorage function
Base.complex(D::Dense{T}) where {T} = Dense{complex(T)}(complex(data(D)))

Base.eltype(::Dense{T}) where {T} = eltype(T)
# This is necessary since for some reason inference doesn't work
# with the more general definition (eltype(Nothing) === Any)
Base.eltype(::Dense{Nothing}) = Nothing
Base.eltype(::Type{Dense{T}}) where {T} = eltype(T)

Base.promote_rule(::Type{Dense{T1}},
                  ::Type{Dense{T2}}) where {T1,T2} = Dense{promote_type(T1,T2)}
Base.convert(::Type{Dense{R}},
             D::Dense) where {R} = Dense{R}(convert(Vector{R},data(D)))

function Base.:*(D::Dense{<:El},x::S) where {El<:Number,S<:Number}
  return Dense{promote_type(El,S)}(x*data(D))
end

Base.:*(x::Number,D::Dense) = D*x

#
# DenseTensor (Tensor using Dense storage)
#

const DenseTensor{ElT,N,StoreT,IndsT} = Tensor{ElT,N,StoreT,IndsT} where {StoreT<:Dense}

# Basic functionality for AbstractArray interface
Base.IndexStyle(::Type{<:DenseTensor}) = IndexLinear()
Base.@propagate_inbounds Base.getindex(T::DenseTensor,i::Int) = store(T)[i]
Base.@propagate_inbounds Base.setindex!(T::DenseTensor,v,i::Int) = (store(T)[i] = v)

Base.fill!(T::DenseTensor,v) = fill!(store(T),v)

# How does Julia map from IndexCartesian to IndexLinear?
#Base.getindex(T::DenseTensor{<:Number,N},
#              i::Vararg{Int,N}) where {N} =
#store(T)[sum(i.*strides(T))+1-sum(strides(T))]
#Base.setindex!(T::DenseTensor{<:Number,N},
#               v,i::Vararg{Int,N}) where {N} =
#(store(T)[sum(i.*strides(T))+1-sum(strides(T))] = v)

# Get the specified value on the diagonal
function getdiag(T::DenseTensor{<:Number,N},ind::Int) where {N}
  return T[CartesianIndex(ntuple(_->ind,Val(N)))]
end

# Set the specified value on the diagonal
function setdiag!(T::DenseTensor{<:Number,N},val,ind::Int) where {N}
  T[CartesianIndex(ntuple(_->ind,Val(N)))] = val
end

# This is for BLAS/LAPACK
Base.strides(T::DenseTensor) = strides(inds(T))

# Needed for passing Tensor{T,2} to BLAS/LAPACK
function Base.unsafe_convert(::Type{Ptr{ElT}},
                             T::DenseTensor{ElT}) where {ElT}
  return Base.unsafe_convert(Ptr{ElT},data(store(T)))
end

# Convert a Dense Tensor to a Tensor with the specified storage
function Base.convert(::Type{<:Tensor{<:Any,<:Any,StoreR}},
                      T::DenseTensor) where {StoreR}
  return Tensor(convert(StoreR,store(T)),inds(T))
end

# Reshape a DenseTensor using the specified dimensions
# This returns a view into the same Tensor data
function Base.reshape(T::DenseTensor,dims)
  dim(T)==dim(dims) || error("Total new dimension must be the same as the old dimension")
  return Tensor(store(T),dims)
end
# This version fixes method ambiguity with AbstractArray reshape
function Base.reshape(T::DenseTensor,dims::Dims)
  dim(T)==dim(dims) || error("Total new dimension must be the same as the old dimension")
  return Tensor(store(T),dims)
end
function Base.reshape(T::DenseTensor,dims::Int...)
  return Tensor(store(T),tuple(dims...))
end

# Create an Array that is a view of the Dense Tensor
# Useful for using Base Array functions
array(T::DenseTensor) = reshape(data(store(T)),dims(inds(T)))
matrix(T::DenseTensor{<:Number,2}) = array(T)
vector(T::DenseTensor{<:Number,1}) = array(T)

# TODO: call permutedims!(R,T,perm,(r,t)->t)?
function Base.permutedims!(R::DenseTensor{<:Number,N},
                           T::DenseTensor{<:Number,N},
                           perm::NTuple{N,Int}) where {N}
  permutedims!(array(R),array(T),perm)
  return R
end

function scale!(T::DenseTensor,
                α::Number)
  A = array(T)
  # This is faster than A .*= α
  rmul!(A,α)
  return T
end

# Version that may overwrite the result or promote
# and return the result
# TODO: move to tensor.jl?
function permutedims!!(R::Tensor,
                       T::Tensor,
                       perm::NTuple{N,Int},
                       f=(r,t)->t) where {N}
  if !is_trivial_permutation(perm)
    permutedims!(R,T,perm,f)
  else
    RA = array(R)
    TA = array(T)
    RA .= f.(RA,TA)
  end
  return R
end

# TODO: move to tensor.jl?
function Base.permutedims(T::Tensor{<:Number,N},
                          perm::NTuple{N,Int}) where {N}
  Tp = similar(T,permute(inds(T),perm))
  Tp = permutedims!!(Tp,T,perm)
  return Tp
end

# TODO: move to tensor.jl?
function Base.:*(x::Number,
                 T::Tensor)
  return Tensor(x*store(T),inds(T))
end
Base.:*(T::Tensor, x::Number) = x*T

# For use in custom permutedims!
using Base.Cartesian: @nexprs,
                      @ntuple,
                      @nloops

#
# A generalized permutedims!(P,B,perm) that also allows
# a function to be applied elementwise
# TODO: benchmark to make sure it is similar to Base.permutedims!
#
# Based off of the permutedims! implementation in Julia's base:
# https://github.com/JuliaLang/julia/blob/91151ab871c7e7d6689d1cfa793c12062d37d6b6/base/multidimensional.jl#L1355
#
@generated function Base.permutedims!(TTP::DenseTensor{<:Number,N},
                                      TT::DenseTensor{<:Number,N},
                                      perm,
                                      f::Function) where {N}
  quote
    TP = array(TTP)
    T = array(TT)
    Base.checkdims_perm(TP, T, perm)

    #calculates all the strides
    native_strides = Base.size_to_strides(1, size(T)...)
    strides_1 = 0
    @nexprs $N d->(strides_{d+1} = native_strides[perm[d]])

    #Creates offset, because indexing starts at 1
    offset = 1 - sum(@ntuple $N d->strides_{d+1})

    ind = 1
    @nexprs 1 d->(counts_{$N+1} = strides_{$N+1}) # a trick to set counts_($N+1)
    @nloops($N, i, TP,
            d->(counts_d = strides_d), # PRE
            d->(counts_{d+1} += strides_{d+1}), # POST
            begin # BODY
                sumc = sum(@ntuple $N d->counts_{d+1})
                @inbounds TP[ind] = f(TP[ind],T[sumc+offset])
                ind += 1
            end)

    return TTP
  end
end
function Base.permutedims!(TP::DenseTensor{<:Number,0},
                           T::DenseTensor{<:Number,0},
                           perm,
                           f::Function)
  TP[] = f(TP[],T[])
  return TP
end

function outer!(R::DenseTensor,
                T1::DenseTensor,
                T2::DenseTensor)
  v1 = vec(T1)
  v2 = vec(T2)
  RM = reshape(R,dim(v1),dim(v2))
  RM .= v1 .* transpose(v2)
  return R
end

function outer!!(R::Tensor,
                 T1::Tensor,
                 T2::Tensor)
  outer!(R,T1,T2)
  return R
end

# TODO: call outer!!, make this generic
function outer(T1::DenseTensor{ElT1},
               T2::DenseTensor{ElT2}) where {ElT1,ElT2}
  array_outer = vec(array(T1))*transpose(vec(array(T2)))
  inds_outer = unioninds(inds(T1),inds(T2))
  return Tensor(Dense{promote_type(ElT1,ElT2)}(vec(array_outer)),inds_outer)
end
const ⊗ = outer

function contraction_output_type(TensorT1::Type{<:DenseTensor},
                                 TensorT2::Type{<:DenseTensor},
                                 indsR)
  return similar_type(promote_type(TensorT1,TensorT2),indsR)
end

function contraction_output(TensorT1::Type{<:DenseTensor},
                            TensorT2::Type{<:DenseTensor},
                            indsR)
  return similar(contraction_output_type(TensorT1,TensorT2,indsR),indsR)
end

# TODO: move to tensor.jl?
function contract(T1::Tensor{<:Any,N1},
                  labelsT1,
                  T2::Tensor{<:Any,N2},
                  labelsT2) where {N1,N2}
  # TODO: put the contract_inds logic into contraction_output,
  # call like R = contraction_ouput(T1,labelsT1,T2,labelsT2)
  indsR,labelsR = contract_inds(inds(T1),labelsT1,
                                inds(T2),labelsT2)
  R = contraction_output(typeof(T1),typeof(T2),indsR)
  # contract!! version here since the output R may not
  # be mutable (like UniformDiag)
  R = contract!!(R,labelsR,T1,labelsT1,T2,labelsT2)
  return R
end

# Move to tensor.jl? Is this generic for all storage types?
function contract!!(R::Tensor{<:Number,NR},
                    labelsR::NTuple{NR},
                    T1::Tensor{<:Number,N1},
                    labelsT1::NTuple{N1},
                    T2::Tensor{<:Number,N2},
                    labelsT2::NTuple{N2}) where {NR,N1,N2}
  if N1==0
    # TODO: replace with an add! function?
    # What about doing `R .= T1[] .* PermutedDimsArray(T2,perm)`?
    perm = getperm(labelsR,labelsT2)
    R = permutedims!!(R,T2,perm,(r,t2)->T1[]*t2)
  elseif N2==0
    perm = getperm(labelsR,labelsT1)
    R = permutedims!!(R,T1,perm,(r,t1)->T2[]*t1)
  elseif N1+N2==NR
    # TODO: permute T1 and T2 appropriately first (can be more efficient
    # then permuting the result of T1⊗T2)
    # TODO: implement the in-place version directly
    R = outer!!(R,T1,T2)
    labelsRp = unioninds(labelsT1,labelsT2)
    perm = getperm(labelsR,labelsRp)
    if !is_trivial_permutation(perm)
      R = permutedims!!(R,copy(R),perm)
    end
  else
    R = _contract!!(R,labelsR,T1,labelsT1,T2,labelsT2)
  end
  return R
end

# TODO: move to tensor.jl?
Base.copyto!(R::Tensor,T::Tensor) = copyto!(store(R),store(T))

# Move to tensor.jl? Overload this function
# for immutable storage types
function _contract!!(R::Tensor,labelsR,
                     T1::Tensor,labelsT1,
                     T2::Tensor,labelsT2)
  _contract!(R,labelsR,T1,labelsT1,T2,labelsT2)
  return R
end

# TODO: make sure this is doing type promotion correctly
# since we are calling BLAS (need to promote T1 and T2 to
# the same types)
function _contract!(R::DenseTensor,
                    labelsR,
                    T1::Tensor{<:Number,<:Any,StoreT1},
                    labelsT1,
                    T2::Tensor{<:Number,<:Any,StoreT2},
                    labelsT2) where {StoreT1<:Dense,StoreT2<:Dense}
  props = ContractionProperties(labelsT1,labelsT2,labelsR)
  compute_contraction_properties!(props,T1,T2,R)

  # We do type promotion here for BLAS (to ensure
  # we contract DenseComplex*DenseComplex)
  if StoreT1 !== StoreT2
    T1,T2 = promote(T1,T2)
  end

  _contract!(R,T1,T2,props)
  return R
end

function _contract!(CT::DenseTensor{El,NC},
                    AT::DenseTensor{El,NA},
                    BT::DenseTensor{El,NB},
                    props::ContractionProperties) where {El,NC,NA,NB}
  C = array(CT)
  A = array(AT)
  B = array(BT)

  tA = 'N'
  if props.permuteA
    AM = reshape(permutedims(A,NTuple{NA,Int}(props.PA)),props.dmid,props.dleft)
    tA = 'T'
  else
    #A doesn't have to be permuted
    if Atrans(props)
      AM = reshape(A,props.dmid,props.dleft)
      tA = 'T'
    else
      AM = reshape(A,props.dleft,props.dmid)
    end
  end

  tB = 'N'
  if props.permuteB
    BM = reshape(permutedims(B,NTuple{NB,Int}(props.PB)),props.dmid,props.dright)
  else
    if Btrans(props)
      BM = reshape(B,props.dright,props.dmid)
      tB = 'T'
    else
      BM = reshape(B,props.dmid,props.dright)
    end
  end

  # TODO: this logic may be wrong
  if props.permuteC
    # Need to copy here since we will be permuting
    # into C later
    CM = reshape(copy(C),props.dleft,props.dright)
  else
    if Ctrans(props)
      CM = reshape(C,props.dleft,props.dright)
      if tA=='N' && tB=='N'
        (AM,BM) = (BM,AM)
        tA = tB = 'T'
      elseif tA=='T' && tB=='T'
        (AM,BM) = (BM,AM)
        tA = tB = 'N'
      end
    else
      CM = reshape(C,props.dleft,props.dright)
    end
  end

  #BLAS.gemm!(tA,tB,promote_type(T,Tα)(α),AM,BM,promote_type(T,Tβ)(β),CM)

  # TODO: make sure this is fast with Tensor{ElT,2}, or
  # convert AM and BM to Matrix
  BLAS.gemm!(tA,tB,one(El),
             AM,BM,
             zero(El),CM)

  if props.permuteC
    permutedims!(C,reshape(CM,props.newCrange...),
                 NTuple{NC,Int}(props.PC))
  end
  return C
end

# Combine a bunch of tuples
# TODO: move this functionality to IndexSet, combine with unioninds?
@inline tuplejoin(x) = x
@inline tuplejoin(x, y) = (x..., y...)
@inline tuplejoin(x, y, z...) = (x..., tuplejoin(y, z...)...)

"""
permute_reshape(T::Tensor,pos)

Takes a permutation that is split up into tuples. Index positions
within the tuples are combined.

For example:

permute_reshape(T,(3,2),1)

First T is permuted as `permutedims(3,2,1)`, then reshaped such
that the original indices 3 and 2 are combined.
"""
function permute_reshape(T::DenseTensor{ElT,NT,IndsT},pos::Vararg{<:Any,N}) where {ElT,NT,IndsT,N}
  perm = tuplejoin(pos...)

  length(perm)≠NT && error("Index positions must add up to order of Tensor ($N)")
  isperm(perm) || error("Index positions must be a permutation")

  dimsT = dims(T)
  indsT = inds(T)
  if !is_trivial_permutation(perm)
    T = permutedims(T,perm)
  end
  N==NT && return T
  newdims = MVector(ntuple(_->eltype(IndsT)(1),Val(N)))
  for i ∈ 1:N
    if length(pos[i])==1
      # No reshape needed, just use the
      # original index
      newdims[i] = indsT[pos[i][1]]
    else
      newdim_i = 1
      for p ∈ pos[i]
        newdim_i *= dimsT[p]
      end
      newdims[i] = eltype(IndsT)(newdim_i)
    end
  end
  newinds = similar_type(IndsT,Val(N))(Tuple(newdims))
  return reshape(T,newinds)
end

# svd of an order-n tensor according to positions Lpos
# and Rpos
function LinearAlgebra.svd(T::DenseTensor{<:Number,N,IndsT},
                           Lpos::NTuple{NL,Int},
                           Rpos::NTuple{NR,Int};
                           kwargs...) where {N,IndsT,NL,NR}
  M = permute_reshape(T,Lpos,Rpos)
  UM,S,VM = svd(M;kwargs...)
  u = ind(UM,2)
  v = ind(VM,2)

  Linds = similar_type(IndsT,Val(NL))(ntuple(i->inds(T)[Lpos[i]],Val(NL)))
  Uinds = push(Linds,u)

  # TODO: do these positions need to be reversed?
  Rinds = similar_type(IndsT,Val(NR))(ntuple(i->inds(T)[Rpos[i]],Val(NR)))
  Vinds = push(Rinds,v)

  U = reshape(UM,Uinds)
  V = reshape(VM,Vinds)

  return U,S,V
end

# eigendecomposition of an order-n tensor according to
# positions Lpos and Rpos
function eigenHermitian(T::DenseTensor{<:Number,N,IndsT},
                        Lpos::NTuple{NL,Int},
                        Rpos::NTuple{NR,Int};
                        kwargs...) where {N,IndsT,NL,NR}
  M = permute_reshape(T,Lpos,Rpos)
  UM,D = eigenHermitian(M;kwargs...)
  u = ind(UM,2)
  Linds = similar_type(IndsT,Val(NL))(ntuple(i->inds(T)[Lpos[i]],Val(NL)))
  Uinds = push(Linds,u)
  U = reshape(UM,Uinds)
  return U,D
end

# qr decomposition of an order-n tensor according to
# positions Lpos and Rpos
function LinearAlgebra.qr(T::DenseTensor{<:Number,N,IndsT},
                          Lpos::NTuple{NL,Int},
                          Rpos::NTuple{NR,Int}) where {N,IndsT,NL,NR}
  M = permute_reshape(T,Lpos,Rpos)
  QM,RM = qr(M)
  q = ind(QM,2)
  r = ind(RM,1)
  # TODO: simplify this by permuting inds(T) by (Lpos,Rpos)
  # then grab Linds,Rinds
  Linds = similar_type(IndsT,Val(NL))(ntuple(i->inds(T)[Lpos[i]],Val(NL)))
  Qinds = push(Linds,r)
  Q = reshape(QM,Qinds)
  Rinds = similar_type(IndsT,Val(NR))(ntuple(i->inds(T)[Rpos[i]],Val(NR)))
  Rinds = pushfirst(Rinds,r)
  R = reshape(RM,Rinds)
  return Q,R
end

# polar decomposition of an order-n tensor according to positions Lpos
# and Rpos
function polar(T::DenseTensor{<:Number,N,IndsT},
               Lpos::NTuple{NL,Int},
               Rpos::NTuple{NR,Int}) where {N,IndsT,NL,NR}
  M = permute_reshape(T,Lpos,Rpos)
  UM,PM = polar(M)

  # TODO: turn these into functions
  Linds = similar_type(IndsT,Val(NL))(ntuple(i->inds(T)[Lpos[i]],Val(NL)))
  Rinds = similar_type(IndsT,Val(NR))(ntuple(i->inds(T)[Rpos[i]],Val(NR)))

  # Use sim to create "similar" indices, in case
  # the indices have identifiers. If not this should
  # act as an identity operator
  Rinds_sim = sim(Rinds)

  Uinds = unioninds(Linds,Rinds_sim)
  Pinds = unioninds(Rinds_sim,Rinds)

  U = reshape(UM,Uinds)
  P = reshape(PM,Pinds)
  return U,P
end

function LinearAlgebra.exp(T::DenseTensor{ElT,N},
                           Lpos::NTuple{NL,Int},
                           Rpos::NTuple{NR,Int};
                           ishermitian::Bool=false) where {ElT,N,
                                                           NL,NR}
  M = permute_reshape(T,Lpos,Rpos)
  indsTp = permute(inds(T),(Lpos...,Rpos...))
  if ishermitian
    expM = exp(Hermitian(matrix(M)))
    return Tensor(Dense{ElT}(vec(expM)),indsTp)
  else
    expM = exp(M)
    return reshape(expM,indsTp)
  end
end


text:
export energy_density_from,
    gas_pressure_from, temperature_from, mass_density_from, specific_energy_from

abstract type EquationOfState end

"""
    energy_density_from(eos::EquationOfState, ρ::Number, u::Number) -> Float

Compute the internal energy density for the given mass density `ρ` and specific
energy `u`."
"""
function energy_density_from end

"""
    gas_pressure_from(eos::EquationOfState, ρ::Number, u::Number) -> Float

Compute the gas pressure for the given mass density `ρ` and internal specific
energy `u`."
"""
function gas_pressure_from end

"""
    temperature_from(eos::EquationOfState, u::Number) -> Float

Compute the temperature for the given internal specific energy `u`."
"""
function temperature_from end

"""
    mass_density_from(eos::EquationOfState, P::Number, u::Number) -> Float

Compute the mass density for the given gas pressure `P` and specific energy `u`."
"""
function mass_density_from end

"""
    specific_energy_from(eos::EquationOfState, T::Number) -> Float

Compute the internal specific energy for the given temperature `T`."
"""
function specific_energy_from end

"""
    specific_energy_from(eos::EquationOfState, ρ::Number, P::Number) -> Float

Compute the internal specific energy for the given mass density `ρ` and
gas pressure `P`."
"""
function specific_energy_from end

"""
    soundspeed²_from(eos::EquationOfState, u::Number) -> Float

Compute the squared sound speed from the given internal specific energy `u`."
"""
function soundspeed²_from end

include("eos/ideal_gas.jl")

max_stars_repo_name: lars-frogner/Whirl

max_issues_repo_name: lars-frogner/Whirl

max_forks_repo_name: lars-frogner/Whirl


text:
function _label_branch!(d, m, ci::CI{SQF{T},S}) where {T<:Real,S}
    sqf = MOI.get(m, MOI.ConstraintFunction(), ci)
    for term in sqf.quadratic_terms
        d[term.variable_index_1.value] = true
        d[term.variable_index_2.value] = true
    end
    return
end
function _label_branch!(d, m, ci::CI{VECVAR,SOC_CONE}, start::Int = 1)
    vv = MOI.get(m, MOI.ConstraintFunction(), ci)
    for i = start:length(vv)
        d[vv.variables[i].value] = true
    end
    return
end

# Iterative deepening depth-first search for variable references
_label_branch_nl!(d, expr::T) where T<:Number = nothing
function _label_branch_nl!(d, ex::Expr)
    if ex.head == :ref
        d[ex.args[2]] = true
        return
    end
    for i = 2:length(ex.args)
        _label_branch_nl!(d, ex_arg)
    end
    return
end

is_scaled_ref(d::T) where T <: Real = true
function is_scaled_ref(d::Expr)
    is_affine = true
    h = ex.head
    if h == :*
        is_affine = is_scaled_ref(d.args[1]) && is_scaled_ref(d.args[1])
    elseif h != :ref
        is_affine = false
    end
    return is_affine
end
is_affine_expr(d::T) where T <: Real = true
function is_affine_expr(d::Expr)
    is_affine = true
    h = ex.head
    if h == :+
        is_affine = all(is_scaled_ref, d.args)
    elseif h == :*
        is_affine = is_scaled_ref(d)
    elseif h != :ref
        is_affine = false
    end
    return is_affine
end

_label_nl(d::Dict{Int,Bool}, ::Nothing) = false
function _label_nl(d::Dict{Int,Bool}, nl_data)
    nl_eval = nl_data.evaluator
    nl_bnds = nl_data.constraint_bounds
    MOI.initialize(nl_eval, Symbol[:ExprGraph])

    if nl_data.has_objective
        obj_expr = MOI.objective_expr(nl_eval)
        _label_branch_nl!(d, obj_expr)
    end
    for i = 1:length(nl_bnds)
        _label_branch_nl!(d, MOI.constraint_expr(nl_eval, i))
    end
    return is_affine_expr(obj_expr)
end


text:
"""
Bridges `CP.Reification{MOI.EqualTo}` to indicator constraints, both with equality
and inequalities (CP.DifferentFrom).
"""
struct ReificationEqualTo2IndicatorBridge{T <: Real} <: MOIBC.AbstractBridge
    indic_true::MOI.ConstraintIndex{MOI.VectorAffineFunction{T}, MOI.Indicator{MOI.ACTIVATE_ON_ONE, MOI.EqualTo{T}}}
    indic_false::MOI.ConstraintIndex{MOI.VectorAffineFunction{T}, MOI.Indicator{MOI.ACTIVATE_ON_ZERO, CP.DifferentFrom{T}}}
end

function MOIBC.bridge_constraint(
    ::Type{ReificationEqualTo2IndicatorBridge{T}},
    model,
    f::MOI.VectorOfVariables,
    s::CP.Reification{MOI.EqualTo{T}},
) where {T}
    return MOIBC.bridge_constraint(
        ReificationEqualTo2IndicatorBridge{T},
        model,
        MOI.VectorAffineFunction{T}(f),
        s,
    )
end

function MOIBC.bridge_constraint(
    ::Type{ReificationEqualTo2IndicatorBridge{T}},
    model,
    f::MOI.VectorAffineFunction{T},
    s::CP.Reification{MOI.EqualTo{T}},
) where {T <: Real}
    indic_true = MOI.add_constraint(
        model,
        f,
        MOI.Indicator{MOI.ACTIVATE_ON_ONE}(MOI.EqualTo(s.set.value))
    )
    indic_false = MOI.add_constraint(
        model,
        f,
        MOI.Indicator{MOI.ACTIVATE_ON_ZERO}(CP.DifferentFrom(s.set.value))
        # TODO: helper to build CP.\neq from MOI.EqTo, CP.Strictly from inequalities, like `!()`?
    )

    return ReificationEqualTo2IndicatorBridge{T}(indic_true, indic_false)
end

function MOI.supports_constraint(
    ::Type{ReificationEqualTo2IndicatorBridge{T}},
    ::Union{Type{MOI.VectorOfVariables}, Type{MOI.VectorAffineFunction{T}}},
    ::Type{CP.Reification{MOI.EqualTo{T}}},
) where {T <: Real}
    return true
end

function MOIB.added_constrained_variable_types(::Type{ReificationEqualTo2IndicatorBridge{T}}) where {T <: Real}
    return Tuple{Type}[]
end

function MOIB.added_constraint_types(::Type{ReificationEqualTo2IndicatorBridge{T}}) where {T <: Real}
    return [
        (MOI.VectorAffineFunction{T}, MOI.Indicator{MOI.ACTIVATE_ON_ONE, MOI.EqualTo{T}}),
        (MOI.VectorAffineFunction{T}, MOI.Indicator{MOI.ACTIVATE_ON_ZERO, CP.DifferentFrom{T}}),
    ]
end

function MOI.get(
    ::ReificationEqualTo2IndicatorBridge{T},
    ::MOI.NumberOfConstraints{
        MOI.VectorAffineFunction{T}, MOI.Indicator{MOI.ACTIVATE_ON_ONE, MOI.EqualTo{T}},
    },
) where {T <: Real}
    return 1
end

function MOI.get(
    ::ReificationEqualTo2IndicatorBridge{T},
    ::MOI.NumberOfConstraints{
        MOI.VectorAffineFunction{T}, MOI.Indicator{MOI.ACTIVATE_ON_ZERO, CP.DifferentFrom{T}},
    },
) where {T <: Real}
    return 1
end

function MOI.get(
    b::ReificationEqualTo2IndicatorBridge{T},
    ::MOI.ListOfConstraintIndices{
        MOI.VectorAffineFunction{T}, MOI.Indicator{MOI.ACTIVATE_ON_ONE, MOI.EqualTo{T}},
    },
) where {T <: Real}
    return [b.indic_true]
end

function MOI.get(
    b::ReificationEqualTo2IndicatorBridge{T},
    ::MOI.ListOfConstraintIndices{
        MOI.VectorAffineFunction{T}, MOI.Indicator{MOI.ACTIVATE_ON_ZERO, CP.DifferentFrom{T}},
    },
) where {T <: Real}
    return [b.indic_false]
end


text:
###############################################################################
#
# Compile data to nc files
#
###############################################################################
"""
    compile_RAW!(param::Array)
    compile_RAW!(
                dt::MOD15A2Hv006LAI{FT},
                params::Array{Any,1},
                zooms::Int
    ) where {FT<:AbstractFloat}

Compile RAW data to nc files, given
- `param` Array of parameters that pass into threads
- `dt` [`AbstractUngriddedData`](@ref) type ungridded data type
- `params` Array of parameters from [`query_RAW`](@ref), not that this `param`
    is different from the `param` used in [`compile_RAW!`](@ref)
- `zooms` Geometrical resolution zoom factor, `zooms=x` means `1/x` degree
"""
function compile_RAW!(param::Array)
    compile_RAW!(param...);

    return nothing
end

function compile_RAW!(
            dt::MOD09A1v006NIRv{FT},
            nlayer::Int,
            files::Array{String,1},
            dims::Array{Int,1},
            cache::String
) where {FT<:AbstractFloat}
    # if file exist, skip it
    if !isfile(cache)
        # create matrices
        res_lon = 360 / dims[1];
        res_lat = 180 / dims[2];
        matv::Array{FT ,2} = zeros(FT , dims[1], dims[2]);
        matn::Array{Int,2} = zeros(Int, dims[1], dims[2]);

        # iterate through files
        for _file in files
            if isfile(_file)
                _data = read_csv(_file);
                for i in eachindex((_data).lat)
                    _lat = lat_ind(_data.lat[i]; res=res_lat);
                    _lon = lon_ind(_data.lon[i]; res=res_lon);
                    _red = _data.red[i];
                    _nir = _data.nir[i];
                    matv[_lon,_lat] += _nir * (_nir - _red) / (_nir + _red);
                    matn[_lon,_lat] += 1;
                end
            end
        end

        # save file to nc file
        matv ./= matn;
        save_LUT!(matv, cache);
    end

    return nothing
end

function compile_RAW!(
            dt::MOD15A2Hv006LAI{FT},
            nlayer::Int,
            files::Array{String,1},
            dims::Array{Int,1},
            cache::String
) where {FT<:AbstractFloat}
    # if file exist, skip it
    if !isfile(cache)
        # create matrices
        res_lon = 360 / dims[1];
        res_lat = 180 / dims[2];
        matv::Array{FT ,2} = zeros(FT , dims[1], dims[2]);
        matn::Array{Int,2} = zeros(Int, dims[1], dims[2]);

        # iterate through files
        for _file in files
            if isfile(_file)
                _data = read_csv(_file);
                for i in eachindex((_data).lat)
                    _lat = lat_ind(_data.lat[i]; res=res_lat);
                    _lon = lon_ind(_data.lon[i]; res=res_lon);
                    _val = _data.lai[i];
                    matv[_lon,_lat] += _val;
                    matn[_lon,_lat] += 1;
                end
            end
        end

        # save file to nc file
        matv ./= matn;
        save_LUT!(matv, cache);
    end

    return nothing
end

function compile_RAW!(
            dt::MOD09A1v006NIRv{FT},
            params::Array{Any,1},
            year::Int,
            days::Int,
            zooms::Int
) where {FT<:AbstractFloat}
    # set time information for the gridding process
    date_start = DateTime("$(year)-01-01");
    date_stop  = DateTime("$(year)-12-31");
    date_dday  = Day(days);
    date_list  = date_start:date_dday:date_stop;

    # compile data per layer
    new_params = [];
    for layer in eachindex(date_list)
        _files = String[];
        for param in params
            if param[2] == layer
                _path     = param[3];
                _prefix   = param[4];
                _cache    = param[6];
                pos_start = findfirst(_prefix, _path);
                hv_string = _path[pos_start[1]:end];
                outp_file = joinpath(_cache, hv_string * ".csv");
                push!(_files, outp_file);
            end
        end
        _save = params[1][6] * "/" * string(layer)*"_"*string(zooms)*"X.nc";
        push!(new_params, [dt, layer, _files, [360*zooms,180*zooms], _save]);
    end

    # compile the data into cache files multi-threading manner
    @info "Compile gridded local data to global data..."
    @showprogress pmap(compile_RAW!, new_params);

    # compile the cache nc files into 1 file
    _fold = params[1][6][1:end-10] * "reprocessed/";
    _year = params[1][6][end-3:end];
    _file = _fold * "NIRv_MODIS_v006_" * _year * "_" *
            string(zooms) * "X_" * string(days) * "D.nc";
    _attr = Dict("longname" => "NIRv", "units" => "-");
    @show _file;

    # if file exists, skip it
    if !isfile(_file)
        _data = zeros(FT, (360*zooms,180*zooms,length(date_list)));
        for i in eachindex(date_list)
            view(_data,:,:,i) .= read_nc(FT, new_params[i][5], "Var");
        end
        save_LUT!(_data, _file, "NIRv", _attr);
    end

    return nothing
end

function compile_RAW!(
            dt::MOD15A2Hv006LAI{FT},
            params::Array{Any,1},
            year::Int,
            days::Int,
            zooms::Int
) where {FT<:AbstractFloat}
    # set time information for the gridding process
    date_start = DateTime("$(year)-01-01");
    date_stop  = DateTime("$(year)-12-31");
    date_dday  = Day(days);
    date_list  = date_start:date_dday:date_stop;

    # compile data per layer
    new_params = [];
    for layer in eachindex(date_list)
        _files = String[];
        for param in params
            if param[2] == layer
                _path     = param[3];
                _prefix   = param[4];
                _cache    = param[6];
                pos_start = findfirst(_prefix, _path);
                hv_string = _path[pos_start[1]:end];
                outp_file = joinpath(_cache, hv_string * ".csv");
                push!(_files, outp_file);
            end
        end
        _save = params[1][6] * "/" * string(layer)*"_"*string(zooms)*"X.nc";
        push!(new_params, [dt, layer, _files, [360*zooms,180*zooms], _save]);
    end

    # compile the data into cache files multi-threading manner
    @info "Compile gridded local data to global data..."
    @showprogress pmap(compile_RAW!, new_params);

    # compile the cache nc files into 1 file
    _fold = params[1][6][1:end-10] * "reprocessed/";
    _year = params[1][6][end-3:end];
    _file = _fold * "LAI_MODIS_v006_" * _year * "_" *
            string(zooms) * "X_" * string(days) * "D.nc";
    _attr = Dict("longname" => "LAI", "units" => "-");
    @show _file;

    # if file exists, skip it
    if !isfile(_file)
        _data = zeros(FT, (360*zooms,180*zooms,length(date_list)));
        for i in eachindex(date_list)
            view(_data,:,:,i) .= read_nc(FT, new_params[i][5], "Var");
        end
        save_LUT!(_data, _file, "LAI", _attr);
    end

    return nothing
end


text:
#####
#To DO:
#Add support for the whole MMVMNFA model family
#Stop using static arrays for tau and anything which scales up with n
#Change the initialisation procedure??
#Catch complex eigenvalues during initialisation
#####

using LinearAlgebra, Distributions, Random, DataFrames, DelimitedFiles, Plots, Distances, Clustering, ProgressMeter, StaticArrays

mutable struct MMVMNModel
    name::String
    pivec::Matrix{Float64}
    mu::Matrix{Float64}
    beta::Matrix{Float64} #Stochastic mean component
    B::Array{Float64}
    D::Array{Float64}
    Ψ::Union{Matrix{Float64},Nothing} #Parameters of the scaling density function
    g::Int64
    q::Int64
    tau::Matrix{Float64}
    logL::Float64
    BIC::Float64
end

"
This is the main fitting function for 3 (nested) classes of model.
"
function MFA_ECM(Y::Array{Float64, 2}, gmin::Int64, gmax::Int64 , qmin::Int64 , eta::Float64  ,itmax::Int64 , nkmeans::Int64 , nrandom::Int64 ,tol = 1e-05, conv_measure = "diff")

    n = size(Y,1); p = size(Y,2);
    if p==2
        qmax = 1
    else
        qmax = Int(floor(p + (1-sqrt(1+8*p))/2))
    end

    if qmax < qmin
        error("Invalid choice of q_min.")
    end

    n_it = Int((gmax-gmin+1)*(qmax-qmin+1)*(nkmeans+nrandom))
    prog = Progress(n_it)

    BIC_best = Inf
    best_model = missing

    n_comb = (gmax-gmin+1)*(qmax-qmin+1)
    combs = reshape(collect(Iterators.product(gmin:gmax, qmin:qmax)),(n_comb,1));
    #Threads.@threads
    Threads.@threads for i in 1:n_comb
        g = combs[i][1]; q = combs[i][2]
            initial_partitions = start_clust(Y, g, nkmeans,nrandom)
            maxinit = size(initial_partitions,2)
            if ismissing(maxinit)
                maxinit = 1
            end

            init_model_para = "placeholder"
            estd_model = "placeholder"
            for ii in 1:maxinit

            init_model_para = init_est_para_mfa(Y, g, q, initial_partitions[:, ii])
            if init_model_para == "failedModel"
                # Failed
            else
                estd_model = est_mfa(init_model_para, Y, itmax, tol, conv_measure, eta)

                if estd_model == "failedModel"
                    # Failed
                else
                    estd_q = q
                    d_model = (g - 1) + 2*g * p + g * (p * estd_q - estd_q * (estd_q - 1)/2)
                    estd_model.BIC = (-2 * estd_model.logL + d_model * log(n))
                    if (estd_model.BIC < BIC_best)
                        best_model = estd_model; BIC_best = estd_model.BIC; #model["BIC"] = BIC_best
                    end
                end
            end

            next!(prog)
            end
    end

  return best_model
end

"
Creates a matrix of starting values of dimension n x (n_kmeans + n_random). These are initial clusterings from which we calculate initial parameter values.
"
function start_clust(Y::Matrix{Float64}, g::Int64, nkmeans::Int64, nrandom::Int64)
    n = size(Y,1)

    init_clust = missing
    k_starts = Array{Float64,2}(undef,n,nkmeans)
    if !ismissing(nkmeans) && nkmeans > 0
        @inbounds for i in 1:nkmeans
            k_starts[:, i] = assignments(kmeans(Array(Y'), g))
        end
    end

  if !ismissing(nrandom) && nrandom > 0
    r_starts = reshape(rand(1:g, n * nrandom), n, nrandom)
  end
  nc = ifelse(ismissing(init_clust), 0, 1 ) #dummy here, will need to fix later

  starts = Array{Float64,2}(undef,n,nkmeans + nrandom)
  ind = 1

  if nkmeans > 0
    starts[:, ind:(ind - 1 + nkmeans)] = k_starts
    ind = ind + nkmeans
  end

  if nrandom > 0
    starts[:, ind:(ind - 1 + nrandom)] = r_starts
    ind = ind + nrandom
  end

  return starts
end

"
Generates initial parameter values for a given start (i.e. a given initial clustering)
"
function init_est_para_mfa(Y::Matrix{Float64}, g::Int64, q::Int64 , start::Union{Vector{Float64},Matrix{Float64}})

    p = size(Y,2)
    n = size(Y,1)

    B = @MArray zeros(Float64,p,q,g)
    pivec = @MArray zeros(Float64,1,g)
    mu = @MArray zeros(Float64,p,g)
    D = @MArray zeros(Float64,p,p,g)
    error = "no"
    @inbounds for i in 1:g
        @views indices = findall(start .== i)[:,1]
        pivec[i] = length(indices)/n
        @views mu[:, i] = mean(Y[indices,:], dims = 1)
        @views Si = cov(Y[indices, :])
        D[:,:,i] = Diagonal(Si)
        @views Di_sqrt = sqrt.(D[:,:,i])
        inv_Di_sqrt = Diagonal(1 ./ diag(Di_sqrt))

        lambda = try
                    eigen(inv_Di_sqrt * Si * inv_Di_sqrt)
                 catch e
                    error = "yes"
                    break
                 end
        eig_order = sortperm(lambda.values, rev = true) #Can get complex eigenvalues here...!
        lambda_vals = lambda.values[eig_order]

        if q == p
            sigma2 = 0
        else
            @views sigma2 = mean(lambda_vals[(q + 1):p])
        end

        if q == 1
            @views B[:,:,i] = Di_sqrt * lambda.vectors[:, eig_order[1:q]] * (lambda_vals[1:q] .- sigma2)
        else
            @views B[:,:,i] = Di_sqrt * lambda.vectors[:, eig_order[1:q]] * Diagonal((lambda_vals[1:q] .- sigma2))
        end
    end

    if error == "yes"
        return "failedModel"
    end

    return MFAModel(pivec,mu,mu,B,D,Nothing,g,q,pivec,0,1)
end

# Add this struct throughout as function output to be reference quickly.
#Then include a temporary struct which can be used to store variable values until the end of each M-step.
#struct MFAModel
#    pivec::Matrix{Float64}
#    mu::Matrix{Float64}
#    B::Array{Float64}
 #   D::Array{Float64}
 #   g::Int64
 #   q::Int64
 #   tau::Matrix{Float64}
 #   logL::Float64
 #   BIC::Float64
#end

function chol_inv(x)
  C <- chol(x)
  inv_x <- chol2inv(C)
  return(inv_x)
end

function MStep_amfa(Y::Matrix{Float64},model::MFAModel,eta::Float64)
    update = MFAModelM(model.pivec,model.mu,model.B,model.D,model.g,model.q,model.tau,model.logL,model.BIC)
    p = size(Y,2); n = size(Y,1)
    if p==2
        q_max = 1
    else
        q_max = Int(floor(p + (1 - sqrt(1 + 8*p))/2))
    end
    n_i = sum(model.tau, dims = 1)

    #CM-Step 1: Updating pi and mu_i
    update.pivec = n_i ./ n
    #mu_new = Array{Float64,2}(undef,p,g)
    @inbounds for i in 1:model.g
        @views update.mu[:,i] = sum(Y .* model.tau[:,i], dims = 1) ./ sum(model.tau[:,i])
    end

    #CM-Step 2: Updating q
    Vtilde = @MArray zeros(Float64,p,p,model.g)
    Lambda = @MMatrix zeros(Float64,q_max,model.g)
    #Lambda_new = Array{Float64,2}(undef,q_max,model.g)
    eigenvecs =  Array{Any,1}(undef,model.g)
    @inbounds for i in 1:model.g
        @views Ymu = Y .- update.mu[:,i]'
        @views Vi = (Array(Ymu') * (Ymu .* model.tau[:,i]))/n_i[i]
        @views D_is = Diagonal(1 ./ sqrt.(model.D[:,:,i]))
        Vtilde[:,:,i] = D_is*Vi*D_is
        @views V_eig_temp = eigvals(Array(Vtilde[:,:,i]))
        #if any(imag.(V_eig_temp) .!= 0)
        #    print("uh oh")
        #    print(imag.(V_eig_temp))
        #    print("converting to real number >:(")
        #end
        V_eig = real.(V_eig_temp)
        eig_order = sortperm(V_eig, rev = true)
        @views Lambda[:,i] = V_eig[eig_order][1:q_max]
        @views eigenvecs[i] = real.(eigvecs(Array(Vtilde[:,:,i]))[:,eig_order])
        if any(Lambda[:,i] .< 0)
            Lambda[(Lambda[:,i] .< 0),i] .= 1e-20
        end
    end

    #CM-Step 3: Updating B_i, then D_i in same loop.
    qi_break = false
    @inbounds for i in 1:model.g
        #Update B
        @views qi = sum(Lambda[:,i] .> 1)
        if qi > model.q
            qi = model.q
        end

        if qi == 0
            qi_break = true
           break
        end
        @views Uqi = eigenvecs[i][:,1:qi]
        @views Ri = I(model.q)[1:qi,:]
        @views update.B[:,:,i] = sqrt.(model.D[:,:,i]) * Uqi * Diagonal(sqrt.(Lambda[1:qi,i] - ones(1,qi))) * Ri

        #Update D
        C = Array(1.0*I(p)); Us = Uqi*(Diagonal(1 ./ Lambda[1:qi,i])-I(qi)); V = Array(Uqi')
        bvec = zeros(p,1)
        psitilde = Array{Float64,2}(undef,p,1) #Debugging up to here ...
        @inbounds for r in 1:p
            @views c = C[1:r,r]
            @views v = V[:,1:r]*c
            @views bvec[r:p] = Us[r:p,:]*v
            b = bvec[r] + 1
            @views a = Array(v')*Diagonal(1 ./ Lambda[1:qi, i] .- Lambda[1:qi,i])*v + Array(c')*Vtilde[1:r,1:r,i] * c
            psitilde[r] = max((((a .- b)/b^2 .+ 1)*update.D[r,r,i])[1],eta)
            omega = psitilde[r]/update.D[r,r,i] - 1
            if r<p && omega!=0
                ratio = omega/(1 + omega*b)
                @views C[1:r, (r+1):p] = C[1:r, (r+1):p] .- ratio .* c*Array(bvec[(r+1):p]')
            end
        end
        @views update.D[:,:,i] = Diagonal(psitilde[:,1])
    end
    if qi_break
        return "failedModel"
    else
        return MFAModel(update.pivec,update.mu,update.B,update.D,model.g,model.q,model.tau,model.logL,model.BIC)
    end
end

function logL_tau_mfa(Y::Matrix{Float64},model::MFAModel)

    p = size(Y,2)
    n = size(Y,1)
    Fji = @MMatrix zeros(Float64,n,model.g)
    @inbounds for i in 1:model.g

        @views inv_D = Diagonal(1 ./ model.D[:,:,i])
        @views B_inv_D = model.B[:,:,i] .* diag(inv_D)
        @views inv_O = inv( I(model.q) + B_inv_D' * model.B[:,:,i] )
        inv_S = inv_D - B_inv_D * inv_O * B_inv_D'
        @views logdetS = sum(log.(diag(model.D[:,:,i]))) - log(det(inv_O))

        @views mahal_dist = Distances.colwise(SqMahalanobis(inv_S), Array(Y'), model.mu[:,i])

        Fji[:, i] = -0.5 .* mahal_dist .- (p/2) .* log(2 * pi) .- 0.5 .* logdetS

    end

    Fji = Fji .+ log.(model.pivec)
    Fjmax = maximum(Fji, dims = 2)
    Fji = Fji .- Fjmax
    logL = sum(Fjmax) + sum(log.(sum(exp.(Fji), dims = 2)))
    Fji = exp.(Fji)
    tau = Fji ./ sum(Fji, dims = 2)
    #model.logL = loglike
    #model.tau = tau
    #return Dict([("loglike", loglike), ("tau", tau)])
    return MFAModel(model.pivec,model.mu,model.B,model.D,model.g,model.q,tau,logL,model.BIC) #model
end

max_stars_repo_name: john-c-davey/FactorMixtures

max_issues_repo_name: john-c-davey/FactorMixtures

max_forks_repo_name: john-c-davey/FactorMixtures


text:
# Hash table acceleration index
struct HashIndex{D <: HashDictionary} <: AbstractIndex
    dict::D
end

function HashIndex(a::AbstractArray)
    dict = HashDictionary{eltype(a), Vector{eltype(keys(a))}}()

    @inbounds for i in keys(a)
        value = a[i]
        vector = get!(Vector{eltype(keys(a))}, dict, value)
        push!(vector, i)
    end

    return HashIndex(dict)
end

# TODO: accelerate! by put elements in order so we can use ranges as dict values and dict
#       iteration mirrors array iteration for cache friendliness?

Base.summary(i::HashIndex) = "HashIndex ($(length(i.dict)) unique element$(length(i.dict) == 1 ? "" : "s"))"

# Accelerations
Base.in(x, a::AcceleratedArray{<:Any, <:Any, <:Any, <:HashIndex}) = haskey(a.index.dict, x)

function Base.count(f::Fix2{typeof(isequal)}, a::AcceleratedArray{<:Any, <:Any, <:Any, <:HashIndex})
    (hasindex, token) = gettoken(a.index.dict, f.x)
    if hasindex
        return length(@inbounds gettokenvalue(a.index.dict, token))
    else
        return 0
    end
end

function Base.findall(f::Fix2{typeof(isequal)}, a::AcceleratedArray{<:Any, <:Any, <:Any, <:HashIndex})
    (hasindex, token) = gettoken(a.index.dict, f.x)
    if hasindex
        return @inbounds gettokenvalue(a.index.dict, token)
    else
        return Vector{eltype(keys(a))}()
    end
end

# TODO: findall for arbitrary predicates by just checking each unique key? (Sometimes faster, sometimes slower?)

function Base.findfirst(f::Fix2{typeof(isequal)}, a::AcceleratedArray{<:Any, <:Any, <:Any, <:HashIndex})
    (hasindex, token) = gettoken(a.index.dict, f.x)
    if hasindex
        return @inbounds first(gettokenvalue(a.index.dict, token))
    else
        return nothing
    end
end

function Base.findlast(f::Fix2{typeof(isequal)}, a::AcceleratedArray{<:Any, <:Any, <:Any, <:HashIndex})
    (hasindex, token) = gettoken(a.index.dict, f.x)
    if hasindex
        return @inbounds last(gettokenvalue(a.index.dict, token))
    else
        return nothing
    end
end

function Base.filter(f::Fix2{typeof(isequal)}, a::AcceleratedArray{<:Any, <:Any, <:Any, <:HashIndex})
    (hasindex, token) = gettoken(a.index.dict, f.x)
    if hasindex
        return @inbounds parent(a)[(gettokenvalue(a.index.dict, token))]
    else
        return empty(a)
    end
end

# TODO: filter for arbitrary predicates by just checking each unique key? (Sometimes faster, sometimes slower?)

function Base.unique(a::AcceleratedArray{T, <:Any, <:Any, <:HashIndex}) where {T}
    out = Vector{T}()
    @inbounds for value in keys(a.index.dict)
        push!(out, value)
    end
    return out
end

function SplitApplyCombine.group(a::AcceleratedArray{<:Any, <:Any, <:Any, <:HashIndex}, b::AbstractArray)
    return map(inds -> @inbounds(b[inds]), a.index.dict)
end

function SplitApplyCombine.groupreduce(op::Callable, a::AcceleratedArray{<:Any, <:Any, <:Any, <:HashIndex}; kw...)
    return map(inds -> @inbounds(reduce(op, view(a, inds); kw...)), a.index.dict)
end

function SplitApplyCombine.groupfind(a::AcceleratedArray{<:Any, <:Any, <:Any, <:HashIndex})
    return a.index.dict
end

function SplitApplyCombine._innerjoin!(out, left::AbstractArray, right::AcceleratedArray{<:Any, <:Any, <:Any, <:HashIndex}, v::AbstractArray, ::typeof(isequal))
    @boundscheck if (axes(left)..., axes(right)...) != axes(v)
        throw(DimensionMismatch("innerjoin arrays do not have matching dimensions"))
    end

    dict = right.index.dict

    @inbounds for i_l ∈ keys(left)
        (hasindex, token) = gettoken(right.index.dict, @inbounds left[i_l])
        if hasindex
            for i_r ∈ gettokenvalue(dict, token)
                push!(out, v[Tuple(i_l)..., Tuple(i_r)...])
            end
        end
    end

    return out
end

function SplitApplyCombine.leftgroupjoin(lkey, ::typeof(identity), f, ::typeof(isequal), left, right::AcceleratedArray{<:Any, <:Any, <:Any, <:HashIndex})
    T = promote_op(f, eltype(left), eltype(right))
    K = promote_op(lkey, eltype(left))

    dict = right.index.dict
    out = HashDictionary{K, Vector{T}}()
    for a ∈ left
        key = lkey(a)
        group = get!(() -> T[], out, key)
        (hasindex, token) = gettoken(dict, key)
        if hasindex
            for b ∈ @inbounds gettokenvalue(dict, token)
                push!(group, f(a, b))
            end
        end
    end
    return out
end

#=
function SplitApplyCombine.innerjoin(lkey, ::typeof(identity), f, ::typeof(isequal), left, right::AcceleratedArray{<:Any, <:Any, <:Any, <:HashIndex})
    T = promote_op(f, eltype(left), eltype(right))
    dict = right.index.dict
    out = T[]
    for a ∈ left
        key = lkey(a)
        dict_index = Base.ht_keyindex(dict, key)
        if dict_index > 0 # -1 if key not found
            for b ∈ dict.vals[dict_index]
                push!(out, f(a, b))
            end
        end
    end
    return out
end
=#


text:
__precompile__()

module FilePaths

using Glob, MacroTools, Reexport, URIParser
using Glob: GlobMatch
@reexport using FilePathsBase

include("compat.jl")
include("glob.jl")
include("uri.jl")

end # end of module


text:
optimize(task:: MSKtask,server:: AbstractString,port:: AbstractString) = optimizermt(task,server,port)
optimize(task:: MSKtask,server:: AbstractString,port:: Int) = optimizermt(task,server,"$port")

function optimize(task:: MSKtask, fallback :: String)
    licfile = (if     haskey(ENV,"MOSEKLM_LICENSE_FILE") && isfile(ENV["MOSEKLM_LICENSE_FILE"]) ENV["MOSEKLM_LICENSE_FILE"]
               elseif haskey(ENV,"HOME") && isfile(joinpath(ENV["HOME"],"mosek","mosek.lic"))       joinpath(ENV["HOME"],"mosek","mosek.lic")
               elseif haskey(ENV,"PROFILE") && isfile(joinpath(ENV["PROFILE"],"mosek","mosek.lic")) joinpath(ENV["PROFILE"],"mosek","mosek.lic")
               else
                 nothing
               end)

    if licfile == nothing
        m = match(r"mosek://([a-zA-Z-]+(\.[a-zA-Z-]+)*)(:([0-9]+))?$",fallback)
        if m == nothing
            error("Invalid server specification $fallback")
        else
            server = m.captures[1]
            port   = m.captures[4]
            optimizermt(task,server,port)
        end
    else
        return optimize(task)
    end
end


text:
function simWithUncer(policyA1,EV,startingA,e,logy1,randYpath)

# This function takes the policy functions and value functions, along with
# starting assets and returns simulated paths of income, consumption,
# assets and value

#% ------------------------------------------------------------------------
# Declare global we need this file have access to
#global mu sigma rho T r Tretire
#global Agrid  Ygrid numSims interpMethod uncertaintyMethod hcIncPDF;
#global normBnd

#% ------------------------------------------------------------------------
# Initialise arrays that will hold the paths of income consumption, value
# and assets

# EV=exVal
# startingA=startA
# e=drawsNor
# logy1=logy1

# Arguments for output
y = Array(Float64,T, numSims);            # income
c = Array(Float64,T, numSims);            # consumption
v = Array(Float64,T, numSims);            # value
a = Array(Float64,T + 1,numSims);         # this is the path at the start of each period, so we include the 'start' of death

# Other arrays that will be used below
ly = Array(Float64,T, numSims);           # log income
ypathIndex = Array(Float64,T, numSims);   # holds the index (location) in the vector

#% ------------------------------------------------------------------------
# Obtain paths using the initial condition and the policy and value
# functions
#-------------------------------------------------------------------------%#
# Scenario where we have hardcoded the discrete
# number of income draws and their pdf

if (uncertaintyMethod == 0)     # we have hardcoded the income process

    #First get random discrete draws using subroutine getDiscreteDraws.
    #More details given in that routine. Given a seed (which
    #we set equal to the time period), it draws randomly from vector Ygrid(t, :)
    #with pdf hcIncPDF
    for t = 1:1:T
        ypathIndex[t, :]=vec(randYpath[t,:])
        y[t, :] = Ygrid[t, vec(ypathIndex[t, :])];
    end

    for s = 1:1: numSims              # loop through individuals
        a[1, s] = startingA;
        for t = 1:1:T                 # loop through time periods for a particular individual

            if (t >= Tretire)
                ixY = 1;
            else
                ixY = ypathIndex[t, s];
            end

            tV = EV[t, :, ixY];          #the relevant part of the value function
            tA1 = policyA1[t, :, ixY];   #the relevant part of the policy function

            intfV  = InterpIrregular(vec(Agrid[t, :]),vec(tV), BCnil, InterpLinear); # Define the interpolation func for Value
            intfP  = InterpIrregular(vec(Agrid[t, :]),vec(tA1), BCnil, InterpLinear); # Define the interpolation func for Polic

            v[t  , s]   =interp1extra(vec(Agrid[t, :]),vec(tV),a[t, s],intfV);
            a[t+1, s]   =interp1extra(vec(Agrid[t, :]),vec(tA1),a[t, s],intfP);
            c[t, s] = a[t, s]  + y[t, s] - (a[t+1, s]/(1+r));
        end   #t
    end # s

#----------------------------------------%
# Scenario where income draws are normally distributed

 elseif (uncertaintyMethod == 1)

     sig_inc = sigma/ ((1-rho^2)^0.5);
     for s = 1:1: numSims                                             # loop through individuals

         a[1, s] = startingA;
         ly[1, s] = truncate(logy1[s], -normBnd*sig_inc , normBnd*sig_inc );
         y[1, s] = exp(ly[1, s]);
         for t = 1:1:T                              # loop through time periods for a particular individual

            if (t >= Tretire)                      # set income to zero if the individual has retired
                 y[t, s] = 0;
            end

            if (t < Tretire)                       # first for the before retirement periods
                #clear tA1 tV;                     # necessary as the dimensions of these change as we wor through this file

                tA1  = squeeze(policyA1[t, :, :],1);   # the relevant part of the policy function
                tV   = squeeze(EV[t, :, :],1);         # the relevant part of the value function

                a[t+1, s] = interp2D(Agrid[t,:]', Ygrid[t, :]', tA1, a[t, s], y[t, s]);
                v[t  , s] = interp2D(Agrid[t,:]', Ygrid[t, :]', tV , a[t, s], y[t, s]);

                if (t != T)  # Get next year's income
                    ly[t+1, s] = (1 -rho) * mu + rho * ly[t, s] + e[t + 1, s];
                    ly[t+1, s] = truncate(ly[t+1, s], -normBnd*sig_inc,normBnd*sig_inc );
                    y[t+1, s] = exp( ly[t+1, s] );
                end # if (t != T)
            else                          # next for the post retirement periods
                #clear tA1 tV;
                tV  = EV[t, :, 1];         # the relevant part of the value function
                tA1 = policyA1[t, :, 1];  # the relevant part of the policy function

                intfV  = InterpIrregular(vec(Agrid[t, :]),vec(tV), BCnil, InterpLinear); # Define the interpolation func for Value
                intfP  = InterpIrregular(vec(Agrid[t, :]),vec(tA1), BCnil, InterpLinear); # Define the interpolation func for Polic

                v[t  , s] =interp1extra(vec(Agrid[t, :]),vec(tV),a[t, s],intfV);
                a[t+1, s] =interp1extra(vec(Agrid[t, :]),vec(tA1),a[t, s],intfP);
                if (t != T)
                    y[t+1, s] = 0;
                end #% if (t ~= T)
            end #% if (t < Tretire)

            c[t, s] = a[t, s]  + y[t, s] - (a[t+1, s]/(1+r));
        end   #t
     end # s
 end # elseif (uncertaintyMethod == 1)
# %-------------------------------------------------------------------------%
# %-------------------------------------------------------------------------%

 return  y, c, a, v

end

max_stars_repo_name: floswald/ucl-econ-julia

max_issues_repo_name: floswald/ucl-econ-julia

max_forks_repo_name: floswald/ucl-econ-julia


text:
using Flux: onehotbatch, crossentropy
using ProgressMeter
using Random
using BSON

const datadir = "sR_bson_subset"

const PHONES = split("h#	q	eh	dx	iy	r	ey	ix	tcl	sh	ow	z	s	hh	aw	m	t	er	l	w	aa	hv	ae	dcl	y	axr	d	kcl	k	ux	ng	gcl	g	ao	epi	ih	p	ay	v	n	f	jh	ax	en	oy	dh	pcl	ah	bcl	el	zh	uw	pau	b	uh	th	ax-h	em	ch	nx	eng")
phn2num = Dict(phone=>i for (i, phone) in enumerate(PHONES))
phn2num["sil"] = 1

function loadData()
  Xs, Ys = Vector(), Vector()
  println("Loading data")
  @showprogress for fname in readdir(datadir)
    BSON.@load joinpath(datadir, fname) mfccs labs

    mfccs = [mfccs[i,:] for i=1:size(mfccs, 1)]

    push!(Xs, mfccs)

    labs = [phn2num[lab] for lab in vec(labs)]
    labs = onehotbatch(labs, collect(1:61))

    labs = [labs[:,i] for i=1:size(labs, 2)]

    push!(Ys, labs)
  end

  return [x for x in Xs], [y for y in Ys]
end

function makeBatches(d, batchSize)
  batches = []

  for i=1:floor(Int64, length(d) / batchSize)
    startI = (i - 1) * batchSize + 1
    lastI = min(startI + batchSize - 1, length(d))

    batch = d[startI:lastI]

    batch = Flux.batchseq(batch, zeros(length(batch[1][1])))
    push!(batches, batch)
  end
  return batches
end

function shuffleData(Xs, Ys)
  indices = collect(1:length(Xs))
  shuffle!(indices)
  return Xs[indices], Ys[indices]
end

function prepData(Xs, Ys)
  XsShuffled, YsShuffled = shuffleData(Xs, Ys)
  XsBatched = makeBatches(XsShuffled, BATCH_SIZE)
  YsBatched = makeBatches(YsShuffled, BATCH_SIZE)
  return collect(zip(XsBatched, YsBatched))
end

function calculateCrossEntropy(wordName, model)
  BSON.@load "$(joinpath(datadir, wordName)).bson" mfccs labs
  mfccs = [mfccs[i,:] for i=1:size(mfccs, 1)]

  labs = [phn2num[lab] for lab in vec(labs)]
  labs = onehotbatch(labs, collect(1:61))

  labs = [labs[:,i] for i=1:size(labs, 2)]
  Flux.reset!(model)
  return Float64.(data.(crossentropy.(model.(mfccs), labs)))
end

function findBoundaries(wordName)
  BSON.@load "$(joinpath(datadir, wordName)).bson" labs

  boundaries = [0]
  curr = labs[1]

  for (i, l) in enumerate(labs)
    if l != curr
      push!(boundaries, i)
      curr = l
    end
  end

  return boundaries
end

function transcribe(wordName)
  BSON.@load "$(joinpath(datadir, wordName)).bson" labs
  phonemes = [labs[1]]
  for lab in labs[2:end]
    if lab != phonemes[length(phonemes)]
      push!(phonemes, lab)
    end
  end

  return phonemes
end

function makeTicks(boundaries, labels)
  ticks = []
  for (b, l) in zip(boundaries, labels)
    if l == "sil"
      l = "#"
    elseif l == "ao"
      l = "aa"
    end
    push!(ticks, uppercase("$(b)\n$(l)"))
  end

  return (boundaries, ticks)
end

function entropy(wordName, model)
  BSON.@load "$(joinpath(datadir, wordName)).bson" mfccs labs
  mfccs = [mfccs[i,:] for i=1:size(mfccs, 1)]

  yhat = model.(mfccs)
  Flux.reset!(model)
  return Float64.([-1 * sum(y .* log.(y)) for y in yhat])
end

function entropyPlot(wordName, model; legendLocation=:bottomright)
  ent = entropy(wordName, model)
  t = Float64.(collect(1:length(ent)))
  b = findBoundaries(wordName)

  ent = predict(loess(t, ent), t)

  p = plot(t, ent, lab="Entropy", title=wordName, ylab="Entropy", xlab="Frame number", legend=legendLocation)
  ticks = makeTicks(b, transcribe(wordName))
  p = vline!(b, lab="Boundaries", xticks=ticks)

  return p
end


text:
# ---------------------------------------------------------------------------- #
#
#   material.jl
#
#   Type for Material properties
#
#   λυτέος
#   Fall 2017
#
#   Max Opgenoord
#
# ---------------------------------------------------------------------------- #

"""
    Material

Material type:
Holds material properties and computes stiffness matrices.
"""
struct Material

  name::String   # Name of material
  E::Float64     # Young's modulus
  ν::Float64     # Poisson's ratio
  ρ::Float64     # Material density
  λ::Float64     # Lamé parameter
  μ::Float64     # Shear modulus (or rigidity)

  σt::Float64    # Maximum tensile stress at yield
  σc::Float64    # Maximum compressive stress at yield

  Cstiff::Array{Float64,4} # Stiffness matrices

end

"""
    Material( name::String )

Constructor for one of the default materials: PLA, Titanium, Al6061.
"""
function Material( name::String, dim::Int64 )

  if name == "PLA"
    E  = 3.5e9  # Pa
    ν  = 0.36
    ρ  = 1.24e3  # kg/m^3
    σt = 50e6   # Pa
    σc = 60e6   # Pa
  elseif name == "Onyx"
    E  = 1.4e9  # Pa
    ν  = 0.36
    ρ  = 1.18e3  # kg/m^3
    σt = 36e6   # Pa
    σc = 45e6   # Pa
  elseif name == "Titanium"
    # Ti6al4v
    E  = 113.8e9 # Pa
    ν  = 0.33
    ρ  = 4.43e3 # kg/m^3
    σt = 880e6   # Pa
    σc = 970e6   # Pa
  elseif name == "Al6061"
    E  = 68.9e9 # Pa
    ν  = 0.33
    ρ  = 2.7e3  # kg/m^3
    σt = 276e6  # Pa
    σc = 386e6  # Pa
  end

  λ = E * ν / ( (1 + ν) * (1 - 2*ν) )
  μ = E / ( 2 * (1 + ν) )

  Cstiff = compCstiff( λ, μ, dim)

  Material( name, E, ν, ρ, λ, μ, σt, σc, Cstiff )

end

"""
    Material( )

Empty constructor for `Material`.
"""
function Material(; name="Unobtainium", E = 68.9e9, ν = 0.33, ρ = 2.7e3, σt = 276e6, σc = 386e6, dim = 2  )

  λ = E * ν / ( (1 + ν) * (1 - 2*ν) )
  μ = E / ( 2 * (1 + ν) )

  Cstiff = compCstiff( λ, μ, dim)

  Material( name, E, ν, ρ, λ, μ, σt, σc, Cstiff )

end

"""
    compCstiff( λ::Float64, μ::Float64, dim::Int64 )

Computes stiffness matrix for dimension `dim`.
"""
function compCstiff( λ::Float64, μ::Float64, dim::Int64 )

  Cs = Array{Float64}(dim, dim, dim, dim)

  for ii in 1:dim, jj in 1:dim, kk in 1:dim, ll in 1:dim
    Cs[ii,jj,kk,ll] = λ*kdelt(ii,jj)*kdelt(kk,ll) +
        μ * (kdelt(ii,kk)*kdelt(jj,ll) + kdelt(ii,ll)*kdelt(jj,kk))
  end

  return Cs

end

"""
    kdelt( ii::Int64, jj::Int64 )

Returns Kronecker δ.
"""
function kdelt( ii::Int64, jj::Int64 )
  # Return kronecker δ
  return Int64( ii == jj )
end


text:
### A Pluto.jl notebook ###
# v0.14.6

using Markdown
using InteractiveUtils

# ╔═╡ cabdb8f7-3d78-422f-bc78-b87fd22eb1d8
begin
    import Pkg
    Pkg.activate(mktempdir())
    Pkg.add([
        Pkg.PackageSpec(name="Plots", version="1"),
        Pkg.PackageSpec(name="PlutoUI", version="0.7"),
    ])
    using Plots, PlutoUI

	Plots.gr(fmt=:png, lw=2)

	PlutoUI.TableOfContents()
end

# ╔═╡ 73adcc40-bdcc-11eb-06eb-970be2047c7c
md"""
# Fig 2.04 Exponential decay
"""

# ╔═╡ ad29e65b-99e2-4fe9-af31-4c29c97aed9b
plot([t-> 3 * exp(-t) t->3 * exp(-2t) t-> 3 * exp(-3t)], 0.0, 5.0, xlim = (0, 5), ylim=(0, 3.2), xlabel="Time", ylabel="Concentration", label = ["exp(-t)" "exp(-2t)" "exp(-3t)"], title= "Figure 2.4")

# ╔═╡ 0f5b8cf2-9b58-40d2-8db1-f61baaf353fe
md"""
## Appendix
Running environment and some auxillary functions.

**You can use [this helper tool](https://fonsp.com/article-test-3/pkghelper.html) to generate these commands!**
"""

# ╔═╡ fcfca624-1ff7-4a60-bd17-d7d32e46de40
begin
	# Convenience functions
    hill(x, k) = x / (x + k)
    hill(x, k, n) = hill(x^n, k^n)
    exprel(x) = ifelse(x≈zero(x), one(x), x / expm1(x))
end

# ╔═╡ Cell order:
# ╠═73adcc40-bdcc-11eb-06eb-970be2047c7c
# ╠═ad29e65b-99e2-4fe9-af31-4c29c97aed9b
# ╠═0f5b8cf2-9b58-40d2-8db1-f61baaf353fe
# ╠═fcfca624-1ff7-4a60-bd17-d7d32e46de40
# ╠═cabdb8f7-3d78-422f-bc78-b87fd22eb1d8

max_stars_repo_name: stevengogogo/pluto-notebooks

max_issues_repo_name: stevengogogo/pluto-notebooks

max_forks_repo_name: stevengogogo/pluto-notebooks


text:
#  String starting with $ is taken from the  enviroment and errors if the ENV variable is not set

string2type = Dict("Bool"=> Bool, "String" => String, "Int" => Int64, "Float" => Float64 )

replace_val_with_env(key, val::Number) = val
replace_val_with_env(key, val::Bool) = val

function replace_val_with_env(key,val::String)
    if startswith(val, '$')
        str = lstrip(val, '$')
        typestring = haskey(TYPE_INFO,key) ? TYPE_INFO[key] : error("No type_info.yaml entry for key $key")
        Type = string2type[typestring]
        haskey(MYENV, str) ? parse(Type,MYENV[str])  : error("No Env[$str]")
    else
        val
    end
end

function replace_env!(d::AbstractDict)
    for (name, namelist) in d
        for (key,val) in namelist
            d[name][key] = replace_val_with_env(key,val)
        end
    end
end