text:
#Written By Johnathan Bizzano
using ...HSMath

"Compile Time Converting the linear indexes from the output_mat index to the input_mat index for mapping in kernel"
function convert_index_matrix(input_shape, output_shape)
    output_len = prod(output_shape)
    output_indexes = reshape(collect(1:output_len), output_shape...)
    input_indexes = reshape(collect(1:prod(input_shape)), input_shape...)

    for idx in CartesianIndices(output_indexes)
        output_indexes[idx] = input_indexes[idx]
    end

    return output_indexes
end

"Compile time experimental determination of the summation deltas between dimensions"
function calculate_input_steps(output_shape, input_index_mat, input_type, stride)
    steps = Array{input_type}(undef, length(output_shape))
    first_input_entry_index = fill(1, length(output_shape))

    for dim in 1:length(output_shape)
        init = input_index_mat[first_input_entry_index...]
        first_input_entry_index[dim] = stride + 1
        after = input_index_mat[first_input_entry_index...]
        first_input_entry_index[dim] = 1
        steps[dim] = after - init                              #Calculate the step change
    end

    return reverse(steps)
end

"Mutlidimensional kernel operation"
function generate_compute_kernelND(T, N, pre_kernel, kernel, post_kernel)
    return :(
        function(input_array, kernel_array, input_steps, kernel_shape, idx, activation)
            input_index = idx
            kernel_index = 1
            index_vector = MArray{Tuple{$N}, $T}(undef)
            dim = 1

            for i in 1:length(index_vector)
                index_vector[i] = 1
            end

            $pre_kernel

            "Multidimensional Iteration. Adapted from HyperSphere.Hyperdimensional.Iteration"
            @label iter
            #Check to see if the index is less then the kernel bound
            if index_vector[dim] <= kernel_shape[end - dim + 1]
                if dim != N
                    dim += 1
                    @goto iter
                else
                    for i in 1:kernel_shape[1]
                        $kernel
                        kernel_index += 1
                        index_vector[end] += 1
                        input_index += input_steps[end]
                    end
                end
            end

            if dim != 1
                #Reset the input_index back to original position
                input_index -= (index_vector[dim] - 1) * input_steps[dim]
                index_vector[dim] = 1
                dim -= 1
                index_vector[dim] += 1
                input_index += input_steps[dim]
                @goto iter
            end

            $post_kernel
        end)
end

#pre_kernel runs before every kernel op. Void input and no args
#kernel runs in every kernel op. Takes in kernel_index and input_index
#post_kernel runs at the completion of every kernel op. Must return input_type value
"A convolutional neural network (CNN, or ConvNet) is a class of artificial neural network, most commonly applied to analyze visual imagery - Wikipedia"
function Kernel(filter_count, sig, initializer, kernel_shape, stride, constant_shape, activation, pre_kernel, kernel, post_kernel)
    return LayerGenerator(
        function (pb, sett, input_shape)
            device = sett.device
            output_shape .= input_shape .- kernel_shape .+ 1

            input_type = get_min_int(prod(input_shape))
            kernel_type = get_min_int(prod(kernel_shape))

            #Determine the mathmatics to perform cartesian operations upon the linear indexes at compile time
            converted_indexes = convert_index_matrix(input_shape, output_shape)
            input_steps = calculate_input_steps(output_shape, converted_indexes, input_type, stride)

            compute_kernelND = eval(generate_compute_kernelND(input_type, length(kernal_shape), pre_kernel, kernel, post_kernel))
            ref_kernel_shape = prealloc_readonly!(pb, kernel_type.(kernel_shape))
            input_steps = prealloc_readonly!(pb, input_steps)
            converted_indexes = prealloc_readonly!(pb, reshape(input_type.(converted_indexes), length(flat_outputs)))     #Flatten the indexes

            LayerDesign(sig, input_shape, (filter_count > 1) ? (output_shape..., filter_count) : output_shape, constant_shape, initializer,
                    function (input_matrixes, output_matrixes, kernel_matrixes)
                        flat_inputs = Ref(reshape(input_matrix, :))
                        input_steps = Ref(input_steps())
                        converted_indexes = converted_indexes()

                        output_matrix_collection = []
                        kernel_matrix_collection = []

                        for i in 1:filter_count
                            #Flatten matrixes for far faster operation
                            push!(output_matrix_collection, Ref(reshape(view(output_matrixes, output_shape..., i), :)))
                            push!(kernel_matrix_collection, Ref(reshape(view(kernel_matrixes, kernel_shape..., i), :)))
                        end

                        return function ()
                                    for i in eachindex(output_matrix_collection)
                                        broadcast!(compute_kernelND, output_matrix_collection[i], flat_inputs, kernel_matrix_collection[i], input_steps, ref_kernel_shape, converted_indexes, activation)
                                    end
                               end
                    end)

        end)
end

function Convolutional(filter_count, initializer; kernel_shape, stride=1, activation=Activation.None())
    return Kernel(filter_count, Sig(:ConvolutionalLayer), initializer, kernel_shape, stride, kernel_shape, activation,
                :(sum = 0.0),
                :(sum += kernel_array[kernel_index] * input_array[input_index]),
                :(activation(sum)))
end

function MaxPooling(; kernel_shape, stride=1, activation=Activation.None())
    return Kernel(1, Sig(:MaxPoolingLayer), initializer, kernel_shape, stride, (0), activation,
                :(maxVal = -Inf),
                :(maxVal += max(maxVal, input_array[input_index])),
                :(activation(maxVal)))
end

export ConvolutionalLayer, MaxPoolingLayer
ConvolutionalLayer(filter_count, initializer; kernel_shape, stride=1, activation=Activation.None()) = Convolutional(filter_count, initializer; kernel_shape=kernel_shape, stride=stride, activation=activation)
MaxPoolingLayer(; kernel_shape, stride=1, activation=Activation.None()) = MaxPooling(; kernel_shape=kernel_shape, stride=stride, activation=activation)


text:
# MIT License
#
# Copyright (c) 2018 Martin Biel
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# Scenario-dependent attributes #
# ========================== #
struct ScenarioDependentOptimizerAttribute <: MOI.AbstractOptimizerAttribute
    stage::Int
    scenario_index::Int
    attr::MOI.AbstractOptimizerAttribute
end

struct ScenarioDependentModelAttribute <: MOI.AbstractModelAttribute
    stage::Int
    scenario_index::Int
    attr::MOI.AbstractModelAttribute
end

struct ScenarioDependentVariableAttribute <: MOI.AbstractVariableAttribute
    stage::Int
    scenario_index::Int
    attr::MOI.AbstractVariableAttribute
end

struct ScenarioDependentConstraintAttribute <: MOI.AbstractConstraintAttribute
    stage::Int
    scenario_index::Int
    attr::MOI.AbstractConstraintAttribute
end

const AnyScenarioDependentAttribute = Union{ScenarioDependentOptimizerAttribute,
                                            ScenarioDependentModelAttribute,
                                            ScenarioDependentVariableAttribute,
                                            ScenarioDependentConstraintAttribute}

MOI.is_set_by_optimize(attr::AnyScenarioDependentAttribute) = MOI.is_set_by_optimize(attr.attr)

is_structure_independent(attr::MOI.AnyAttribute) = true
is_structure_independent(::Union{MOI.ObjectiveFunction,
                                 MOI.ObjectiveFunctionType,
                                 MOI.ObjectiveSense}) = false
is_structure_independent(attr::AnyScenarioDependentAttribute) = is_structure_independent(attr.attr)

ext: jl


text:
using DifferentialEquations

mpa2mm = FT(10^6/9.8); #mPa of water pressure to mm of water height

#functions to update the pressure and conductivity of a plant, after its storage volumes have already been updated
function update_pk_leaf!(tissue::LeafHydraulics{FT})
	tissue.p_storage = (tissue.v_storage ./ tissue.v_maximum .- FT(1)) / tissue.pv.slope
	tissue.p_element[1] = tissue.p_storage
	tissue.p_leaf = tissue.p_storage
	tissue.p_ups = tissue.p_storage[1]
	tissue.k_element = exp.(-(tissue.p_element ./ -tissue.vc.b) .^ tissue.vc.c) * tissue.k_sla
end

function update_pk_nonleaf!(tissue::Land.PlantHydraulics.AbstractHydraulicOrgan{FT})
	tissue.p_storage = (tissue.v_storage ./ tissue.v_maximum .- FT(1)) / tissue.pv.slope
	tissue.p_element = tissue.p_storage
	tissue.k_element = exp.(-(tissue.p_element ./ -tissue.vc.b) .^ tissue.vc.c) * tissue.k_max
end

function update_pk_tree!(plant_hs)

	for can_i in 1:plant_hs.n_canopy
		iLeaf = plant_hs.leaves[can_i]
		update_pk_leaf!(iLeaf)

		iBranch = plant_hs.branch[can_i]
		update_pk_nonleaf!(iBranch)
	end

	update_pk_nonleaf!(plant_hs.trunk)

	for root_i in 1:plant_hs.n_root
		iRoot = plant_hs.roots[root_i]
		update_pk_nonleaf!(iRoot)
	end

end

#conserve the total leaf water volume when changing LAI
#alternatively, this water could be removed or added to branches
function update_pv_leaf!(plant_hs, lai_old, lai_new)
	for can_i in 1:plant_hs.n_canopy
		iLeaf = plant_hs.leaves[can_i]
		iLeaf.v_storage *= (lai_old / lai_new)
		update_pk_leaf!(iLeaf)
	end
end

#The capacitance model solves a set of ODEs at each timestep
#described by dv/dt = A*v + b
#this function fills in the matrix A and vector b
#using the state of the plant
#note that it is designed to work with a system where each component only has one vertical slice

function create_deriv_mat(plant_hs)

	n_canopy= plant_hs.n_canopy
	n_tissue = plant_hs.n_canopy*2 + 1 + plant_hs.n_root

	ans_lhs = zeros(n_tissue, n_tissue);
	ans_rhs = zeros(n_tissue)

	i = 1
	for can_i in 1:plant_hs.n_canopy
		iLeaf = plant_hs.leaves[can_i]
		iBranch = plant_hs.branch[can_i]

		leaf_k_total = iLeaf.k_element[1]*iLeaf.area
		leaf_vmax_total = iLeaf.v_maximum*iLeaf.area

		kbl = 2/(1/leaf_k_total + 1/iBranch.k_element[1])
		#kbl = leaf_k_total

		dzbl = iBranch.Δh/2 * 1000/mpa2mm;
		#dzbl = 0;

		leaf_coef_leaf_in = -kbl/iLeaf.pv.slope/leaf_vmax_total
		ans_lhs[i,i] += leaf_coef_leaf_in
		ans_lhs[i+n_canopy,i] += -leaf_coef_leaf_in

		branch_coef_leaf_in = kbl/iBranch.pv.slope/iBranch.v_maximum[1]

		ans_lhs[i,i+n_canopy] += branch_coef_leaf_in
		ans_lhs[i+n_canopy,i+n_canopy] += -branch_coef_leaf_in

		#constant terms in pressure difference flow into leaf
		rhs_term1 = kbl*(-1/iBranch.pv.slope + 1/iLeaf.pv.slope)
		ans_rhs[i] += rhs_term1
		ans_rhs[i+n_canopy] += -rhs_term1

		ans_rhs[i] += -dzbl*kbl - iLeaf.q_out*iLeaf.area
		ans_rhs[i+n_canopy] += dzbl*kbl

		i = i+1
	end

	m1= ans_lhs*1;
	#println(ans_lhs)

	trunk = plant_hs.trunk
	itrunk = plant_hs.n_canopy*2+1;

	for can_i in 1:plant_hs.n_canopy
		iBranch = plant_hs.branch[can_i]
		ktb = 2/(1/iBranch.k_element[1] + 1/trunk.k_element[1])
		#ktb = iBranch.k_element[1]
		dztb = (iBranch.Δh/2 + trunk.Δh/2) * 1000/mpa2mm
		#dztb = iBranch.Δh* 1000/mpa2mm

		branch_coef_branch_in = -ktb/iBranch.pv.slope/iBranch.v_maximum[1]
		ans_lhs[i,i] += branch_coef_branch_in
		ans_lhs[itrunk,i] += -branch_coef_branch_in

		trunk_coef_branch_in = ktb/trunk.pv.slope/trunk.v_maximum[1]
		ans_lhs[i,itrunk] += trunk_coef_branch_in
		ans_lhs[itrunk,itrunk] += -trunk_coef_branch_in

		rhs_term1 = ktb*(-1/trunk.pv.slope + 1/iBranch.pv.slope)
		ans_rhs[i] += rhs_term1
		ans_rhs[itrunk] += -rhs_term1

		ans_rhs[i] += -dztb*ktb
		ans_rhs[itrunk] += dztb*ktb

		i = i+1

	end

	m2 = ans_lhs*1;
	#println(ans_lhs)

	i = itrunk+1

	total_k = sum([x.area/x.Δh for x in plant_hs.roots]);

	for root_i in 1:plant_hs.n_root

		iRoot = plant_hs.roots[root_i]
		area_frac = iRoot.area/iRoot.Δh/total_k  #total root area is 1
		krt = 2/(1/(iRoot.k_element[1]*area_frac) + 1/trunk.k_element[1])
		#krt = trunk.k_element[1]

		dzrt = (iRoot.Δh/2 + trunk.Δh/2)*1000/mpa2mm
		#dzrt = trunk.Δh/2*1000/mpa2mm

		trunk_coef_trunk_in = -krt/trunk.pv.slope/trunk.v_maximum[1]
		ans_lhs[itrunk,itrunk] += trunk_coef_trunk_in
		ans_lhs[i,itrunk] += -trunk_coef_trunk_in

		root_coef_trunk_in = krt/iRoot.pv.slope/iRoot.v_maximum[1]
		ans_lhs[itrunk,i] += root_coef_trunk_in
		ans_lhs[i,i] += -root_coef_trunk_in

		rhs_term1 = krt*(-1/iRoot.pv.slope + 1/trunk.pv.slope)
		ans_rhs[itrunk] += rhs_term1
		ans_rhs[i] += -rhs_term1

		ans_rhs[itrunk] += -dzrt*krt
		ans_rhs[i] += dzrt*krt

		ksr = iRoot.k_element[1]*area_frac
		dzsr = iRoot.Δh/2*1000/mpa2mm
		#dzsr = iRoot.Δh*1000/mpa2mm

		root_coef_root_in = -ksr/iRoot.pv.slope/iRoot.v_maximum[1]
		ans_lhs[i,i] += root_coef_root_in
		ans_rhs[i] += -dzsr*ksr + iRoot.p_ups*ksr + ksr/iRoot.pv.slope
		i = i+1

	end

	#println(ans_lhs)

	return ans_lhs, ans_rhs#, m1, m2

end

using LinearAlgebra

#this function solves the system of linear ODEs
#dx/dt = deriv_mat*x + deriv_const
#given the initial condition of x
#it returns the value of x after the time interval deltaT
#it also returns the integrals of the components of x over that time interval
function solve_odes(deriv_mat, deriv_const, init_cond, deltaT)

	e_val = eigvals(deriv_mat);
	e_vec = eigvecs(deriv_mat);

	inv_evec = inv(e_vec);
	inv_deriv = inv(deriv_mat);
	cprime = -inv_deriv * deriv_const;
	mateI = e_vec * diagm(exp.(e_val*deltaT)) * inv_evec
	new_x = cprime + mateI * (init_cond-cprime)
	int_x = cprime*deltaT + (inv_deriv * mateI - inv_deriv) * (init_cond-cprime)

	return new_x, int_x
end

#=
#slower way of solving ODEs
function solve_odes_exp(deriv_mat, deriv_const, init_cond, deltaT)
	inv_deriv = inv(deriv_mat);
	cprime = -inv_deriv * deriv_const;
	mateI = exp(deltaT*deriv_mat)
	new_x = cprime + mateI * (init_cond-cprime)
	int_x = cprime*deltaT + (inv_deriv * mateI - inv_deriv) * (init_cond-cprime)
	return new_x, int_x
end
=#

#this function takes a plant hydraulic system and returns the pressure of each leaf,branch,trunk,and root
#note that it is designed to work with a system where each component only has one vertical slice
function get_p_prof2(plant_hs)
	root_p = [x.p_element[1] for x in plant_hs.roots]
	trunk_p = plant_hs.trunk.p_element[1];
	branch_p = [x.p_element[1] for x in plant_hs.branch]
	leaf_p = [x.p_element[1] for x in plant_hs.leaves]

	ncan = plant_hs.n_canopy
	nroot = plant_hs.n_root
	n_tissue = ncan*2 + 1 + nroot

	p_list = vcat(leaf_p,branch_p,trunk_p,root_p)

	#=
	p_list = zeros(FT, n_tissue)
	p_list[1:ncan] = leaf_p;
	p_list[(ncan+1):(ncan*2)] = branch_p
	p_list[ncan*2+1] = trunk_p
	p_list[(ncan*2+2):end] = root_p
	=#

	return p_list
end

#same as previous function but returning current storage volume instead of pressure
function get_v_prof2(plant_hs)
	root_p = [x.v_storage[1] for x in plant_hs.roots]
	trunk_p = plant_hs.trunk.v_storage[1];
	branch_p = [x.v_storage[1] for x in plant_hs.branch]
	leaf_p = [x.v_storage[1]*x.area for x in plant_hs.leaves]

	ncan = plant_hs.n_canopy
	nroot = plant_hs.n_root
	n_tissue = ncan*2 + 1 + nroot

	p_list = vcat(leaf_p,branch_p,trunk_p,root_p)

	#=
	p_list = zeros(FT, n_tissue)
	p_list[1:ncan] = leaf_p;
	p_list[(ncan+1):(ncan*2)] = branch_p
	p_list[ncan*2+1] = trunk_p
	p_list[(ncan*2+2):end] = root_p
	=#
	return p_list
end

#sets volumes of each component to the values listed in newvols
function set_vol!(plant_hs, newvals)
	for ican in 1:plant_hs.n_canopy
		plant_hs.leaves[ican].v_storage = newvals[ican]/plant_hs.leaves[ican].area;
		plant_hs.branch[ican].v_storage[1] = newvals[ican+plant_hs.n_canopy];
	end

	plant_hs.trunk.v_storage[1] = newvals[plant_hs.n_canopy*2+1];

	for iroot in 1:plant_hs.n_root
		plant_hs.roots[iroot].v_storage[1] = newvals[plant_hs.n_canopy*2+1+iroot];
	end
end

#combines the above functions to advance the capacitance model by a time step
function update_cap_mat!(plant_hs, deltaT)

	volumes = get_v_prof2(plant_hs);
	dmat, dvec = create_deriv_mat(plant_hs)

	newvals, newints = solve_odes(dmat, dvec, volumes, deltaT)
	#newvals, newints = solve_odes_exp(dmat, dvec, volumes, deltaT)

	set_vol!(plant_hs,newvals)

	#compute the average value of root water uptake over the time step
	#for each root
	for iroot in 1:plant_hs.n_root
		rootI = plant_hs.roots[iroot]
		m = 1/rootI.pv.slope;
		ksr = rootI.k_element[1]*rootI.area #total root area is 1
		ps = rootI.p_ups
		dzr = rootI.Δh*1000/mpa2mm
		vroot = rootI.v_maximum[1]
		integral_rootvol = newints[plant_hs.n_canopy*2+1+iroot];
		rootI.q_in = ksr*(ps*deltaT - m/vroot*integral_rootvol + m*deltaT - dzr*deltaT)/deltaT
	end

	update_pk_tree!(plant_hs)

	return newvals,newints

end

function find_new_cap(plant_hs, deltaT)

	volumes = get_v_prof2(plant_hs);
	dmat, dvec = create_deriv_mat(plant_hs)

	newvals, newints = solve_odes(dmat, dvec, volumes, deltaT)

	return newvals,newints
end

function push_cap!(plant_hs, newvals, newints)
	set_vol!(plant_hs,newvals)

	total_k = sum([x.area/x.Δh for x in plant_hs.roots]);

	for iroot in 1:plant_hs.n_root
		rootI = plant_hs.roots[iroot]
		m = 1/rootI.pv.slope;
		ksr = rootI.k_element[1]*rootI.area/rootI.Δh/total_k #total root area is 1
		ps = rootI.p_ups
		dzr = rootI.Δh/2*1000/mpa2mm
		vroot = rootI.v_maximum[1]
		integral_rootvol = newints[plant_hs.n_canopy*2+1+iroot];

		rootI.q_in = ksr*(ps*deltaT - m/vroot*integral_rootvol + m*deltaT - dzr*deltaT)/deltaT
	end

	update_pk_tree!(plant_hs)
end

function get_v_max(plant_hs)
	root_p = [x.v_maximum[1] for x in plant_hs.roots]
	trunk_p = plant_hs.trunk.v_maximum[1];
	branch_p = [x.v_maximum[1] for x in plant_hs.branch]
	leaf_p = [x.v_maximum[1]*x.area for x in plant_hs.leaves]

	p_list = vcat(leaf_p,branch_p,trunk_p,root_p)

	return p_list
end

function soil_ode!(node::SPACMono{FT}, rain_in::FT, deltaT::FT, k_soil_max::FT, slope_index::FT)
	layer_thick = -1*diff(node.soil_bounds)*1000

	#fluxes_inout = -[x.q_in for x in node.plant_hs.roots] /KG_H_2_MOL_S * FT(1/node.ga*deltaT/(60*60));
	fluxes_inout = -[x.q_in for x in node.plant_hs.roots] *FT(18.02/1000)* FT(1/node.ga) #mol/s over basal to kg/s over ground;
	#node.swc[soil_layerI] = max(node.plant_hs.roots[i_root].sh.Θr,min(node.swc[soil_layerI],node.plant_hs.roots[i_root].sh.Θs));
	#fluxes_inout = zeros(FT,8);
	max_rain = (node.plant_hs.roots[1].sh.Θs - node.swc[1])*layer_thick[1]/(60*60)
	rain_rate = rain_in/(60*60);
	#fluxes_inout[1] += rain_in/(60*60)*(node.swc[1] < node.plant_hs.roots[1].sh.Θs) #mm/hour to mm/s;
	fluxes_inout[1] += min(rain_rate, max_rain)
	#fluxes_inout *= 0;

	wrc = node.plant_hs.roots[1].sh;
	midpoints = (node.soil_bounds[2:end] + node.soil_bounds[1:(end-1)]) /2 * 1000;
	layer_len = -1*diff(midpoints);

	function soil_derivative(smc_vec,p,t)
		k_soil = k_soil_max/(60*60) * [soil_k_ratio_swc(wrc, xi) for xi in smc_vec];
		p_soil_tot = [soil_p_25_swc(wrc,xi) for xi in smc_vec]*mpa2mm .+ midpoints;
		p_diff = -1*diff(p_soil_tot) #in mm
		q_down = k_soil[1:(end-1)] .* p_diff ./ layer_len * deltaT #mm/s * mm / mm * s
		free_drainage = k_soil[end]*slope_index*deltaT
		ans = fluxes_inout ./ layer_thick;
		ans[1:(end-1)] -= q_down ./ layer_thick[1:(end-1)];
		ans[2:end] += q_down ./ layer_thick[2:end];
		ans[end] -= free_drainage / layer_thick[end];

		#shortfall_to_sat = node.plant_hs.roots[1].sh.Θs .- smc_vec;
		#only_negative = ans .* (ans .< 0);
		#ans[smc_vec .>= node.plant_hs.roots[1].sh.Θs] = only_negative[smc_vec .>= node.plant_hs.roots[1].sh.Θs]
		return ans
	end
	u0 = node.swc;
	tspan = (FT(0.0),deltaT)
	problem1 = ODEProblem(soil_derivative,u0,tspan);
	sol1 = solve(problem1);
	finalsol = sol1.u[end];
	runoff = 0;

	for i_soil in eachindex(node.swc)
		node.swc[i_soil] = max(min(finalsol[i_soil], node.plant_hs.roots[1].sh.Θs),node.plant_hs.roots[1].sh.Θr);
	end

	return finalsol #in units of mm
end


text:
"""
    StreamDecoderPtr

A Julia type for a libflac stream decoder object.

This type primarily exists so that a finalizer to delete the decoder can be assigned.
"""
mutable struct StreamDecoderPtr  # type not immutable so that finalizer can be applied
    v::Ptr{Cvoid}
end

function StreamDecoderPtr()
    dec = StreamDecoderPtr(ccall((:FLAC__stream_decoder_new,libflac),Ptr{Cvoid},()))
    finalizer(x->ccall((:FLAC__stream_decoder_delete,libflac),Cvoid,(Ptr{Cvoid},),x.v), dec)
    dec
end

Base.unsafe_convert(::Type{Ptr{Cvoid}}, en::StreamDecoderPtr) = en.v

for (nm,typ) in (("ogg_serial_number", :Clong),
                 ("md5_checking", :Bool),
                 ("metadata_respond", :MetaDataType),
                 ("metadata_ignore", :MetaDataType))
    @eval begin
        function $(Symbol(string("set_", nm)))(dd::StreamDecoderPtr, val)
            ccall(($(string("FLAC__stream_decoder_set_",nm)), libflac), Bool,
                (Ptr{Cvoid}, $typ), dd, val)
        end
    end
end

set_md5_checking(dd::StreamDecoderPtr) = set_md5_checking(dd, true)

"""
    set_metadata_respond_all(dd::StreamDecoderPtr)

Set the stream decoder to respond to all metadata blocks.

Must be called before the decoder is initialized.
"""
set_metadata_respond_all(dd::StreamDecoderPtr) =
    ccall((:FLAC__stream_decoder_set_metadata_respond_all,libflac),Bool,(Ptr{Cvoid},),dd)

"""
    set_metadata_ignore_all(dd::StreamDecoderPtr)

Set the stream decoder to ignore all metadata blocks.

Must be called before the decoder is initialized.
"""
set_metadata_ignore_all(dd::StreamDecoderPtr) =
    ccall((:FLAC__stream_decoder_set_metadata_respond_all,libflac),Bool,(Ptr{Cvoid},),dd)

"""
    get_state_string(dd::StreamDecoderPtr)

Returns a character string describing the current state of the decoder
"""
get_state_string(dd::StreamDecoderPtr) =
    unsafe_string(ccall((:FLAC__stream_decoder_get_resolved_state_string,libflac), Ptr{UInt8},
                     (Ptr{Cvoid},), dd))

"""
    get_state(dd:StreamDecoderPtr)

Returns the decoder state as a `StreamDecoderState` enum value.
"""
get_state(dd)

@enum(StreamDecoderState,
      DecoderMetaDataSearch,
      DecoderMetaDataRead,
      DecoderFrameSyncSearch,
      DecoderFrameRead,
      DecoderEndOfStream,
      DecoderOggError,
      DecoderSeekError,
      DecoderAborted,
      DecoderMemoryAllocationError,
      DecoderUninitialized)

for (nm,typ) in (("state",:StreamDecoderState),
                 ("md5_checking",:Bool),
                 ("total_samples",:Int64),
                 ("channels",:Cint),
                 ("channel_assignment",:ChannelAssignment),
                 ("bits_per_sample",Cint),
                 ("sample_rate",Cint),
                 ("blocksize",Cint))
    @eval begin
        function $(Symbol(string("get_",nm)))(dd::StreamDecoderPtr)
            ccall(($(string("FLAC__stream_decoder_get_",nm)),libflac),$typ,(Ptr{Cvoid},),dd)
        end
    end
end

@enum(DecoderInitStatus,
      DecoderInitOK,
      DecoderInitUnsupportedContainer,
      DecoderInitInvalidCallbacks,
      DecoderInitMemoryAllocationError2,
      DecoderInitErrorOpeningFile,
      DecoderInitAlreadyInitialized)

"""
    debug_mcallback(d::Ptr{Cvoid}, mp::Ptr{Cvoid}, client::Ptr{Cvoid})

Debugging metadata callback function.  Prints a brief description of any
`Info`, `Padding`, `VorbisComment`, or `SeekTable` metadata blocks in the
stream.
"""
function debug_mcallback(d::Ptr{Cvoid}, mp::Ptr{Cvoid}, client::Ptr{Cvoid})
    typ = unsafe_load(reinterpret(Ptr{MetaDataType}, mp))
    println("Metadata callback on typ = ", typ)
    if typ == Info
        show(unsafe_load(reinterpret(Ptr{InfoMetaData}, mp)))
    elseif typ == Padding
        show(unsafe_load(reinterpret(Ptr{PaddingMetaData}, mp)))
    elseif typ == VorbisComment
        show(unsafe_load(reinterpret(Ptr{VorbisComment}, mp)))
    elseif typ == SeekTable
        show(unsafe_load(reinterpret(Ptr{SeekTable}, mp)))
    end
    nothing
end

"""
    silent_mcallback(d::Ptr{Cvoid}, mp::Ptr{Cvoid}, client::Ptr{Cvoid})

Silent metadata callback function.  Ignores all metadata blocks.
"""
function silent_mcallback(d::Ptr{Cvoid}, mp::Ptr{Cvoid}, client::Ptr{Cvoid})
    nothing
end

"""
    debug_ecallback(d::Ptr{Cvoid}, status::Int32, client::Ptr{Cvoid})

Error callback function.
"""
function debug_ecallback(d::Ptr{Cvoid}, status::Int32, client::Ptr{Cvoid})
    error("Got error callback with status = $status")
end

"""
    debug_wcallback(dd::Ptr{Cvoid}, hdr::Ptr{FrameHeader},
        buffer::Ptr{Ptr{Int32}}, client::Ptr{Cvoid})

Debugging write callback.  Prints information about every frame written.
Very verbose.  Use with caution.
"""
function debug_wcallback(dd::Ptr{Cvoid}, hdr::Ptr{FrameHeader},
             buffer::Ptr{Ptr{Int32}}, client::Ptr{Cvoid})
    fr = unsafe_load(hdr)
    println("Frame")
    println(" blocksize: ", fr.blocksize)
    println(" samplerate: ", fr.sample_rate)
    println(" channels: ", fr.channels)
    println(" channel assignment: ", fr.channel_assignment)
    println(" bits per sample: ", fr.bits_per_sample)
    println(" number type: ", fr.typ)
    println(" frame or sample number: ", fr.num)
    println(" crc: ", fr.crc)

    data = unsafe_wrap(unsafe_load(buffer), fr.blocksize)/(2^(fr.bits_per_sample - 1))
    zero(Int32)
end

"""
    initfile!(dd::StreamDecoderPtr, fnm::String; wcallback=debug_wcallback_c,
        mcallback=debug_mcallback_c, ecallback=debug_ecallback_c, client_data=nothing)

Initialize the `StreamDecoderPtr`, `dd`, to read the FLAC file `fnm`.

This function allows the user to override any of the default callback functions.
"""
function initfile!(dd::StreamDecoderPtr, fnm::String; wcallback=debug_wcallback_c,
                   mcallback=debug_mcallback_c, ecallback=debug_ecallback_c, client_data=nothing)
    status = ccall((:FLAC__stream_decoder_init_file,libflac),DecoderInitStatus,
                   (Ptr{Cvoid}, Ptr{UInt8}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}),
                   dd, fnm, wcallback, mcallback, ecallback, client_data)
    if status != DecoderInitOK
        error("decoder_init_file returned status $status")
    end
    dd
end

function process_single(dd::StreamDecoderPtr)
    disable_sigint() do
        ccall((:FLAC__stream_decoder_process_single,libflac), Bool, (Ptr{Cvoid},), dd)
    end
end
function process_metadata(dd::StreamDecoderPtr)
    disable_sigint() do
        ccall((:FLAC__stream_decoder_process_until_end_of_metadata,libflac), Bool, (Ptr{Cvoid},), dd)
    end
end
function process_stream(dd::StreamDecoderPtr)
    disable_sigint() do
        ccall((:FLAC__stream_decoder_process_until_end_of_stream,libflac), Bool, (Ptr{Cvoid},), dd)
    end
end
function seek_absolute(dd::StreamDecoderPtr, offset::UInt64)
    disable_sigint() do
        ccall((:FLAC__stream_decoder_seek_absolute,libflac), Bool, (Ptr{Cvoid},UInt64), dd, offset)
    end
end
function flush(dd::StreamDecoderPtr)
    disable_sigint() do
        ccall((:FLAC__stream_decoder_flush,libflac), Bool, (Ptr{Cvoid},), dd)
    end
end

function saving_mcallback(d::Ptr{Cvoid}, mp::Ptr{Cvoid}, client::Ptr{Cvoid})
    typ = unsafe_load(reinterpret(Ptr{MetaDataType}, mp))
    if typ == Info
        f = unsafe_pointer_to_objref(client)
        f.metadata = unsafe_load(reinterpret(Ptr{InfoMetaData}, mp))
    end
    return nothing
end

function buffering_wcallback(dd::Ptr{Cvoid}, hdr::Ptr{FrameHeader},
                   buffer::Ptr{Ptr{Int32}}, client::Ptr{Cvoid})
    fr = unsafe_load(hdr)
    f = unsafe_pointer_to_objref(client)

    max_val = Float32(2^(fr.bits_per_sample - 1))
    for chidx = 1:fr.channels
        ch_data = unsafe_load(buffer, chidx)
        f.chunk[1:fr.blocksize, chidx] = unsafe_wrap(Array, ch_data, fr.blocksize)/max_val
    end
    # Set the current blocksize and how much we've consumed from it
    f.chunk_len = fr.blocksize
    f.chunk_consumed = 0
    return zero(Int32)
end

mutable struct FLACDecoder
    # Filename of the file this actually represents
    filename::String

    # Pointer to our decoder object
    dec::StreamDecoderPtr

    # Metadata for this FLAC file
    metadata::InfoMetaData

    # Current chunk of audio data, its length (we allocate maximum framesize,
    # but frames can vary in size) and how much we've consumed from it, so we
    # know where to pick up next when reading in
    chunk::Array{Float32,2}
    chunk_len::UInt
    chunk_consumed::UInt

    function FLACDecoder(path::String)
        dec = StreamDecoderPtr()

        # Create initial, empty FLACDecoder object
        f = new(path, dec, InfoMetaData(), Array{Float32,2}(undef, 0, 0), 0, 0)

        # Open file, process metadata
        initfile!(dec, path, wcallback=buffering_wcallback_c, mcallback=saving_mcallback_c, client_data=pointer_from_objref(f))
        process_metadata(dec)

        # Initialize chunk storage
        f.chunk = Array{Float32,2}(undef, f.metadata.maxblocksize, f.metadata.channels)
        return f
    end
end

"""
    read(f::FLACDecoder, num_samples::Integer)

Read up to the specified number of samples from the given FLACDecoder,
"""
function read(f::FLACDecoder, num_samples::T) where T<:Integer
    # Allocate memory to hold all the read data
    data = Array{Float32,2}(undef, num_samples, f.metadata.channels)
    data_read = 0

    # Save our current position, and read until we hit an error or exceed our
    # num_samples read:
    while data_read < num_samples
        data_needed = num_samples - data_read
        data_avail = f.chunk_len - f.chunk_consumed

        # If we've got data to load in from a previous chunk, then do so
        if data_avail > 0
            # Take only as much as we need, or as much as we can
            amnt = min(data_needed, data_avail) - 1

            # Copy from the chunk into data
            d_start = data_read + 1
            f_start = f.chunk_consumed + 1
            data[d_start:d_start + amnt, :] = f.chunk[f_start:f_start + amnt, :]

            # Increment data_read and chunk_consumed for proper bookkeeping
            data_read += amnt + 1
            f.chunk_consumed += amnt + 1
        else
            # If we have no data ready for us, read it in!
            if !process_single(f.dec)
                # If this fails, quit out with whatever we've managed to get
                return data[1:data_read]
            end
            # Did we not actually get any new data?
            if f.chunk_len - f.chunk_consumed == 0
                return data[1:data_read]
            end
        end
    end

    return data
end

"""
    seek(f::FLACDecoder, offset::Int64)

Perform an absolute seek within the given FLAC stream.  Throws an
`ArgumentError` if the requested seek is impossible.  Will automatically
`flush()` the decoder stream if a seek error is encountered.
"""
function seek(f::FLACDecoder, offset::T) where T<:Integer
    if !seek_absolute(f.dec, UInt64(offset))
        # If we got a seek error, we need to flush the decoder before we
        # can continue decoding, as we've crunked the decoder state. Whoops.
        if get_state(f.dec) == 0x06
            flush(f.dec)
        end
        throw(ArgumentError("Could not seek to offset $offset"))
    end
end

"""
    length(f::FLACDecoder)

Return the total length of the FLAC file in samples
"""
function length(f::FLACDecoder)
    return f.metadata.totalsamples
end

"""
    size(f::FLACDecoder)

Return the size tuple of the FLAC file (length in samples, number of channels)
"""
function size(f::FLACDecoder)
    return (length(f), f.metadata.channels)
end

"""
    load(filename::File{format"FLAC"})

Opens and reads the given filename, returning an array of samples and as well
as the samplerate of the samples stored within the file.  This method is part
of the FileIO integration, do not call directly but instead import the FileIO
bindings and call `load()` with a filename ending in `.flac` to invoke this
method through the FileIO machinery.
"""
function load(file::File{format"FLAC"})
    f = FLACDecoder(file.filename)

    # Return data, fs
    return read(f, length(f)), f.metadata.samplerate
end

# Calculate cfunction versions of all our callbacks once, at runtime, as is necessary with cfunction's
debug_mcallback_c = Ptr{Cvoid}(C_NULL)
silent_mcallback_c = Ptr{Cvoid}(C_NULL)
saving_mcallback_c = Ptr{Cvoid}(C_NULL)
debug_wcallback_c = Ptr{Cvoid}(C_NULL)
buffering_wcallback_c = Ptr{Cvoid}(C_NULL)
debug_ecallback_c = Ptr{Cvoid}(C_NULL)
function init_decoder_cfunctions()
    global debug_mcallback_c     = @cfunction(debug_mcallback, Cvoid, (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}))
    global silent_mcallback_c    = @cfunction(silent_mcallback, Cvoid, (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}))
    global saving_mcallback_c    = @cfunction(saving_mcallback, Cvoid, (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}))
    global debug_wcallback_c     = @cfunction(debug_wcallback, Int32, (Ptr{Cvoid}, Ptr{FrameHeader}, Ptr{Ptr{Int32}}, Ptr{Cvoid}))
    global buffering_wcallback_c = @cfunction(buffering_wcallback, Int32, (Ptr{Cvoid}, Ptr{FrameHeader}, Ptr{Ptr{Int32}}, Ptr{Cvoid}))
    global debug_ecallback_c     = @cfunction(debug_ecallback, Cvoid, (Ptr{Cvoid}, Int32, Ptr{Cvoid}))
end


text:
module WNNMDenoise

using LinearAlgebra
using ImageCore
using ImageCore: NumberLike, GenericGrayImage, GenericImage
using ImageDistances
using Statistics
using BlockMatching
using Test
# using Dates # for benchmark

using ImageQualityIndexes # TODO: remove this
using ProgressMeter # TODO: remove this

include("utilities.jl")

export WNNM

"""
    WNNM(noise_level; kwargs...)

Weighted nuclear norm minimization denoising algorithm.

!!! note
    To keep consistent with the original implementation[2], here we stick
    to the 0-255 value range world. If the image is 0-1 range, then the default
    parameters may not work at all.

# Arguments

- `noise_level`: the (estimated) gaussian noise level, used in the first WNNM iteration.

# Keywords

The default values of these keywords come from the original reference
implementation [2]. Except that `patch_size` in our implementation has
to be odd number.

- `K`: the number of WNNM iterations
- `δ`: step value for each WNNM iteration

keywords that control the block matching step, it can be scalar or a vector of length `K`. If it is
a vector then it specifies the values for each WNNM iteration

- `num_patches`: the number of matched patches.
- `patch_size`: patch size.
- `patch_stride`: the stride of each sliding window.
- `window_size`: non-local search size in block matching.

keywords that control the internal WNNM solver, which is a small threshold-svd loop
on the extracted block matching results.

- `λ`: weight constant used to estimate the remained noise level of each patch.
- `C`: weight constant in WNNM solver used by the original implementation.

# Examples

```julia
using TestImages, WNNMDenoise, ImageTransformations

# 1. load a clean image
#    0-255 range is required
clean_img = Float64.(imresize(testimage("cameraman") .* 255, ratio=0.5));

# 2. add guassian noise
σₙ = 40;
noisy_img = clean_img .+ σₙ .* randn(Float32, size(clean_img))

# 3. denoise using the WNNM denoiser
f = WNNM(40)
denoised_img = f(copy(noisy_img), noisy_img)
```

# References

[1] Gu, Shuhang, et al. "Weighted nuclear norm minimization with application to image denoising." _Proceedings of the IEEE conference on computer vision and pattern recognition_. 2014.
[2] The MATLAB reference implementation: http://www4.comp.polyu.edu.hk/~cslzhang/code/WNNM_code.zip

"""
struct WNNM
    noise_level::Float64
    K::Int
    δ::Float64
    num_patches::Vector{Int}
    patch_size::Vector{Int}
    patch_stride::Vector{Int}
    λ::Float64
    C::Float64
    window_size::Int
end

function WNNM(noise_level;
              δ=0.1,
              C=2sqrt(2),
              window_size=60,
              patch_size=nothing,
              num_patches=nothing,
              K=nothing,
              λ=nothing,
              patch_stride=nothing)
    if noise_level <= 20
        isnothing(patch_size) && (patch_size = 5)
        isnothing(num_patches) && (num_patches = 70)
        isnothing(K) && (K = 8)
        isnothing(λ) && (λ = 0.56)
    elseif noise_level <= 40
        isnothing(patch_size) && (patch_size = 7)
        isnothing(num_patches) && (num_patches = 90)
        isnothing(K) && (K = 12)
        isnothing(λ) && (λ = 0.56)
    elseif noise_level <= 60
        isnothing(patch_size) && (patch_size = 7)
        isnothing(num_patches) && (num_patches = 120)
        isnothing(K) && (K = 14)
        isnothing(λ) && (λ = 0.58)
    else
        isnothing(patch_size) && (patch_size = 9)
        isnothing(num_patches) && (num_patches = 140)
        isnothing(K) && (K = 14)
        isnothing(λ) && (λ = 0.58)
    end

    @assert isodd(patch_size) "patch size is expected to be odd number, instead it is $patch_size"
    patch_size isa Number && (patch_size = fill(patch_size, K))
    isnothing(patch_stride) && (patch_stride = @. max(1, patch_size ÷ 2 - 1))
    patch_stride isa Number && (patch_stride = fill(patch_stride, K))

    num_patches = fill(num_patches - 10, K)
    drop_freq = 2
    for k in 2:K
        # drop by 10 for every 2 iteration
        num_patches[k] = (k - 1) % drop_freq == 0 ? num_patches[k - 1] - 10 : num_patches[k - 1]
    end

    WNNM(noise_level, K, δ, num_patches, patch_size, patch_stride, λ, C, window_size)
end

## Implementation

function (f::WNNM)(imgₑₛₜ, imgₙ; clean_img=nothing)
    # FIXME: this if branch is not type stable
    # if imgₑₛₜ === imgₙ
    #     imgₑₛₜ = copy(imgₙ)
    # else
    #     copyto!(imgₑₛₜ, imgₙ)
    # end
    copyto!(imgₑₛₜ, imgₙ)

    T = eltype(eltype(imgₑₛₜ))
    # outfile = open(joinpath("benchmark", "julia_$(Int(f.noise_level)).csv"), "w") # for benchmark
    # println(outfile, "iter,psnr,runtime") # for benchmark
    for iter in 1:f.K
        @. imgₑₛₜ = imgₑₛₜ + f.δ * (imgₙ - imgₑₛₜ) # This iteration can be done more sophisticatedly

        # The noise level for the first iteration is known (whether it is estimated outside or a
        # white noise). The noise is removed in each iteration, so we have to estimate a noise level
        # at a patch level; the denoising performance on each patch can be different, which means a
        # global noise level can be misleading.
        σₚ = iter == 1 ? f.noise_level : zero(f.noise_level)

        # calculating svd using blas threads is not an optimal parallel strategy
        # start_time = now() # for benchmark
        imgₑₛₜ .= with_blas_threads(1) do
            _estimate_img(imgₑₛₜ, imgₙ,
                f.patch_size[iter],
                f.patch_stride[iter],
                f.num_patches[iter],
                f.window_size;
                noise_level=f.noise_level,
                λ=T(f.λ),
                C=T(f.C),
                σₚ=σₚ,
            )
        end
        # duration = now() - start_time # for benchmark

        # TODO: remove this logging part when it is ready
        if !isnothing(clean_img)
            cur_psnr = assess_psnr(clean_img, imgₑₛₜ, 255)
            @info "Result" iter psnr = cur_psnr num_patches = f.num_patches[iter]
            # println(outfile, "$(iter), $(cur_psnr), $(duration.value/1000.)") # for benchmark
            display(Gray.(imgₑₛₜ ./ 255))
            sleep(0.1)
        end
    end
    # close(outfile) # for benchmark
    return imgₑₛₜ
end

function _estimate_img(imgₑₛₜ::AbstractMatrix, imgₙ,
        patch_size::Int,
        patch_stride::Int,
        num_patches::Int,
        window_size::Int;
        kwargs...)
    patch_size = (patch_size, patch_size)

    # We set stride in both dimension instead only in column dimension.
    # This gives less computation but the similar output speaking of PSNR.
    Δ = CartesianIndex((patch_stride, 1)) # original version in MATLAB
    Δ = CartesianIndex((patch_stride, patch_stride))

    r = CartesianIndex(patch_size .÷ 2)
    R = CartesianIndices(imgₑₛₜ)
    R = first(R) +r:Δ:last(R) -r

    imgₑₛₜ⁺ = zeros(eltype(imgₑₛₜ), axes(imgₑₛₜ))
    W = zeros(Int, axes(imgₑₛₜ))

    progress = Progress(length(R[1:patch_stride:end]))
    out_buffer = [Matrix{eltype(imgₑₛₜ)}(undef, prod(patch_size), num_patches) for i in 1:Threads.nthreads()]
    patch_group_buffer = [Matrix{eltype(imgₑₛₜ)}(undef, prod(patch_size), num_patches) for i in 1:Threads.nthreads()]
    patch_q_indices_buffer = [Matrix{CartesianIndex{2}}(undef, prod(patch_size), num_patches) for i in 1:Threads.nthreads()]
    m_buffer = [Vector{eltype(imgₑₛₜ)}(undef, prod(patch_size)) for i in 1:Threads.nthreads()]
    dist_buffer = [Vector{Float64}(undef, (window_size+1)^2) for i in 1:Threads.nthreads()]

    Threads.@threads for p in R
        tid = Threads.threadid()
        out = out_buffer[tid]
        patch_q_indices = patch_q_indices_buffer[tid]
        patch_group = patch_group_buffer[tid]
        m = m_buffer[tid]
        dist = dist_buffer[tid]

        fill!(out, zero(eltype(out)))
        patch_q_indices = _estimate_patch!(
            patch_q_indices, out, patch_group, m, dist,
            imgₑₛₜ, imgₙ, p,
            patch_size,
            num_patches,
            window_size,
            ;kwargs...)

        # Technically, there will be data racing here if multiple threads are involved, but as we've
        # observed, this doesn't affect the overall performance.
        view(W, patch_q_indices) .+= 1
        view(imgₑₛₜ⁺, patch_q_indices) .+= out
        next!(progress)
    end

    return imgₑₛₜ⁺ ./ max.(W, 1)
end

function _estimate_patch!(patch_q_indices, out, patch_group, m, dist,
                          imgₑₛₜ, imgₙ, p,
                          patch_size::Tuple,
                          num_patches::Int,
                          window_size::Int;
                          noise_level,
                          λ,
                          C,
                          σₚ=nothing)
    rₚ = CartesianIndex(patch_size .÷ 2)
    p_indices = p - rₚ:p + rₚ

    alg = FullSearch{Cityblock, 2}(
        Cityblock(),
        rₚ,
        CartesianIndex(window_size÷2, window_size÷2),
        CartesianIndex(1, 1)
    )
    q_inds = multi_match(alg, imgₑₛₜ, imgₑₛₜ, p, dist; num_patches=num_patches)
    # the memory allocation will be a hospot if we directly generate `patch_q_indices` using `vcat`,
    # thus we pre-allocate it and then copyto the buffer.
    @inbounds for i = 1:length(q_inds)
        q = q_inds[i]
        R = q - rₚ:q + rₚ
        copyto!(view(patch_q_indices, :, i), R)
        copyto!(view(patch_group, :, i), view(imgₑₛₜ, R))
    end
    # patch_group .= @view imgₑₛₜ[patch_q_indices]
    mean2!(m, patch_group)

    if σₚ == 0
        # Try: use the mean estimated σₚ of each patch
        σₚ = _estimate_noise_level(view(imgₑₛₜ, p_indices), view(imgₙ, p_indices), noise_level; λ=λ)
    end
    @. out = patch_group - m
    WNNM_optimizer!(out, out, eltype(out)(σₚ); C=C)
    out .+= m

    return patch_q_indices
end

@doc raw"""
    WNNM_optimizer(Y, σₚ; C, rank, fixed_point_num_iters=3)

Optimizes the weighted nuclear norm minimization problem with a fixed point estimation

```math
    \min_X \lVert Y - X \rVert^2_{F} + \lVert X \rVert_{w, *}
```

The weight `w` is specially chosen so that it satisfies the condition of Corollary 1 in [1].

# References

[1] Gu, S., Zhang, L., Zuo, W., & Feng, X. (2014). Weighted nuclear norm minimization with application to image denoising. In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_ (pp. 2862-2869).

"""
function WNNM_optimizer!(out, Y, σₚ; C, fixed_point_num_iters=3)
    # Apply Corollary 1 in [1] for image denoise purpose
    # Note: this solver is reserved to the denoising method and is not supposed to be used in other
    #       applications; it simply isn't designed so.

    n = size(Y, 2)
    F = svd!(Y)

    # For image denoising problems, it is natural to shrink large singular value less, i.e., to set
    # smaller weight to large singular value. For this reason, it uses `w = (C * sqrt(n))/(ΣX + eps())`
    # as the weights; inversely propotional to ΣX. With singular values ΣX sorted ascendingly, the
    # condition for Corollary 1 holds, and thus we could directly get the desired solution in a single
    # step.

    # Here we iterate more than once because we don't know what ΣX is; we have to iterate it a while
    # from ΣY to get an relatively good estimation of it.
    # TODO: could we set default σₚ as 0?
    # TODO: is this the best initialization we can get?
    σₚ² = σₚ*σₚ
    nσₚ² = n*σₚ²
    Csnσₚ² = C * sqrt(n) * σₚ²

    ΣX = @. sqrt(max(F.S*F.S - nσₚ², zero(σₚ)))
    for _ in 1:fixed_point_num_iters
        # the iterative algorithm proposed in section 2.2.2 in [1]
        # Step 1 in the iterative algorithm becomes trivial and a no-op

        # Step 2 degenerates to a soft thresholding; both P and Q are identity matrix.
        # all in one line to avoid unnecessary allocation for temporarily variable w
        @. ΣX = soft_threshold(F.S, Csnσₚ² / (ΣX + eps()))
    end

    rmul!(F.U, Diagonal(ΣX))
    mul!(out, F.U, F.Vt)
end

function _estimate_noise_level(patchₑₛₜ, patchₙ, σₙ; λ=0.56)
    # Estimate the noise level of given patch during the WNNM iteration
    # we still need to know the input noisy level σₙ to give an estimation
    λ * sqrt(abs(σₙ^2 - mse(patchₑₛₜ, patchₙ)))
end

end


text:
function lens_split(x :: Array, lens :: Array{Int64, 1})
    ind = cumsum(lens)
    n = length(ind)
    ind = [0; ind]
    [view(x, (ind[i]+1):ind[i+1]) for i in 1:n]
end

function init_node!(op :: Object, value :: Dict)
    value[:value] = Array{Float64, 1}(op.class.type_len)
    value[:d] = Array{Float64, 1}(op.class.type_len)
end

function init_node!(op :: DFunction, value :: Dict)
    value[:pre_matrix] = Array{Float64, 1}(size(op.f_matrix, 2))
    value[:pre_transform] = Array{Float64, 1}(op.out_class.type_len)
    value[:value] = Array{Float64, 1}(op.out_class.type_len)
    value[:inputs] = lens_split(value[:pre_matrix], [c.type_len for c in op.in_classes])
    value[:d] = Array{Float64, 1}(op.out_class.type_len)
    value[:pre_transform_d] = Array{Float64, 1}(op.out_class.type_len)
    value[:pre_matrix_d] = Array{Float64, 1}(size(op.f_matrix, 2))
    value[:input_ds] = lens_split(value[:pre_matrix_d], [c.type_len for c in op.in_classes])
end

function _init_tree!(tree :: Tree, _)
    init_node!(tree.op, tree.value)
    tree
end

function init_tree!(tree :: Tree)
    bottom_up(_init_tree!, tree)
end

function eval_node!(op :: Object, value :: Dict, value_list)
    value[:value] .= op.value
end

function eval_node!(op :: DFunction, value :: Dict, value_list)
    i = 1
    for v in value_list
        value[:inputs][i] .= v[:value]
        i += 1
    end
    ## value[:pre_matrix] .= vcat([v[:post_feature] for v in value_list]...)
    A_mul_B!(value[:pre_transform], op.f_matrix, value[:pre_matrix])
    value[:value] .= eval(op.f_transformation).(value[:pre_transform])
end

function eval_tree!(tree :: Tree)
    foreach(eval_tree!, tree.subtrees)
    eval_node!(tree.op, tree.value, (t.value for t in tree.subtrees))
    tree
end

max_stars_repo_name: Non-Contradiction/NextGeneration

max_issues_repo_name: Non-Contradiction/NextGeneration

max_forks_repo_name: Non-Contradiction/NextGeneration


text:
module SQLite

include("sql_exports.jl")
include("sql_imports.jl") # DBMS Structured SubQuery _to_sql _placeholder _placeholders

const DatabaseID = DBMS.SQLite

"""
    to_sql(query::Structured)::String
"""
to_sql(query::Structured)::String = _to_sql(DatabaseID(), query)

"""
    to_sql(subquery::SubQuery)::String
"""
to_sql(subquery::SubQuery)::String = _to_sql(DatabaseID(), subquery)

placeholder(nth::Int) = _placeholder(DatabaseID(), nth)
placeholders(dims::Int) = _placeholders(DatabaseID(), dims)

using .Octo.AdapterBase: FromItem, SqlPart, _sqlrepr
import .Octo.AdapterBase: sqlrepr
function sqlrepr(db::DatabaseID, clause::FromItem)::SqlPart
    _sqlrepr(db, clause; with_as=false)
end

export        AUTOINCREMENT
@sql_keywords AUTOINCREMENT

end # Octo.Adapters.SQLite


text:
# Parameters controlling how a plot appears

const title_font_desc = "'PT Sans','Helvetica Neue','Helvetica',sans-serif"
const label_font_desc = "'PT Sans Caption','Helvetica Neue','Helvetica',sans-serif"

# Choose highlight color by darkening the fill color
function default_discrete_highlight_color(fill_color::ColorValue)
    return RGB(1, 1, 1)
end

function default_continuous_highlight_color(fill_color::ColorValue)
    c = convert(LCHab, fill_color)
    return LCHab(max(0, c.l - 40), c.c, c.h)
end

function default_stroke_color(fill_color::ColorValue)
    fill_color = convert(LCHab, fill_color)
    c = LCHab(fill_color.l, fill_color.c, fill_color.h)
    LCHab(c.l - 15, c.c, c.h)
end

function default_lowlight_color(fill_color::ColorValue)
    fill_color = convert(LCHab, fill_color)
    c = LCHab(fill_color.l, fill_color.c, fill_color.h)
    LCHab(90, 20, c.h)
end

# Choose a middle color by darkening the fill color
function default_middle_color(fill_color::ColorValue)
    fill_color = convert(LCHab, fill_color)
    LCHab(fill_color.l + 40, fill_color.c, fill_color.h)
end

@varset Theme begin
    # If the color aesthetic is not mapped to anything, this is the color that
    # is used.
    default_color,         ColorOrNothing,  LCHab(70, 60, 240)

    # Default size when the size aesthetic is not mapped.
    default_point_size,    Measure,         0.9mm

    # Width of lines in the line geometry.
    line_width,            Measure,         0.3mm

    # Background color of the plot.
    panel_fill,            ColorOrNothing,  nothing

    # Border color of the plot panel.
    panel_stroke,          ColorOrNothing,  nothing

    # Opacity of the plot background panel.
    panel_opacity,         Float64,         1.0

    # Grid line color.
    grid_color,            ColorOrNothing,  color("#D0D0E0")
    grid_strokedash,       Maybe(Vector),   [0.5mm, 0.5mm]

    # Grid lines for focused item.
    grid_color_focused,    ColorOrNothing,  color("#A0A0A0")

    # Width of grid lines
    grid_line_width,       Measure,         0.2mm

    # Font name, size, and color used for tick labels, entries in keys, etc.
    minor_label_font,      String,          label_font_desc
    minor_label_font_size, Measure,         8pt
    minor_label_color,     ColorOrNothing,  color("#6c606b")

    # Font name, size and color used for axis labels, key title, etc.
    major_label_font,      String,          title_font_desc
    major_label_font_size, Measure,         11pt
    major_label_color,     ColorOrNothing,  color("#564a55")

    # Font name, size and color used for labels on plot elements.
    point_label_font,      String,          label_font_desc
    point_label_font_size, Measure,         8pt
    point_label_color,     ColorOrNothing,  color("#4c404b")

    # Font name, size and color used for key titles
    key_title_font,      String,          title_font_desc
    key_title_font_size, Measure,         11pt
    key_title_color,     ColorOrNothing,  color("#362a35")

    # Font name, size and color used for key entries.
    key_label_font,      String,          title_font_desc
    key_label_font_size, Measure,         8pt
    key_label_color,     ColorOrNothing,  color("#4c404b")

    # How many gradations to show in a continuous color key.
    key_color_gradations, Int,            40

    # Spacing between bars for Geom.bar.
    bar_spacing,           Measure,         0.0mm

    # Spacing between boxplots in Geom.boxplot.
    boxplot_spacing,       Measure,         1mm

    # Length of caps on error bars
    errorbar_cap_length,   Measure,         3mm

    # Lines are drawn in a slightly different color than fills, e.g. to
    # differentiate histogram bars from error bars.
    stroke_color,          Function,       default_stroke_color

    # Points, etc, are highlighted by stroking in slightly different color. This
    # is the stroke width.
    highlight_width,       Measure,         0.3mm

    # A function mapping fill color to stoke color for highlights.
    discrete_highlight_color,       Function,        default_discrete_highlight_color
    continuous_highlight_color,     Function,        default_continuous_highlight_color

    # A function mapping fill color to a duller background fill color. Used for
    # Geom.ribbon in particular so lines stand out against it.
    lowlight_color,        Function,        default_lowlight_color

    # Opacity of geometry filled with lowlight_color
    lowlight_opacity,      Float64,         0.6

    # A function mapping base fill color to the color of the median marker in a
    # boxplot.
    middle_color,          Function,        default_middle_color

    # Width of the middle line in a boxplot.
    middle_width,          Measure,         0.6mm

    # Horizontal position of the title of color key guides. One of :left,
    # :right, :center.
    guide_title_position,  Symbol,          :left

    # Shape used in color keys for color swatches. Either :square or :circle.
    colorkey_swatch_shape, Symbol,          :square

    # One of :left, :right, :top, :bottom, :none determining where color keys
    # and the like should be placed.
    key_position,          Symbol,          :right

    # TODO: This stuff is too incomprehensible to be in theme, I think. Put it
    # somewhere else.

    # Number of annealing iterations.
    label_placement_iterations,  Int,       1000

    # Penalty for a label not being contained within the plot frame.
    label_out_of_bounds_penalty, Float64,   10.0

    # Penalty for making a label hidden to avoid overlaps.
    label_hidden_penalty,        Float64,   0.5

    # Probability of proposing a visibility flip during label layout.
    label_visibility_flip_pr,    Float64,   0.2

end

const default_theme = Theme()


text:
@testset "Order" begin
    @test DNAMer("AA") < DNAMer("AC") < DNAMer("AG") < DNAMer("AT") < DNAMer("CA")
    @test RNAMer("AA") < RNAMer("AC") < RNAMer("AG") < RNAMer("AU") < RNAMer("CA")

    @test BigDNAMer("AA") < BigDNAMer("AC") < BigDNAMer("AG") < BigDNAMer("AT") < BigDNAMer("CA")
    @test BigRNAMer("AA") < BigRNAMer("AC") < BigRNAMer("AG") < BigRNAMer("AU") < BigRNAMer("CA")
end


text:
#=
    many of the types are
    Used from https://github.com/lanl-ansi/Juniper.jl
=#

mutable struct RegisteredFunction
    s         :: Symbol
    dimension :: Integer
    f         :: Function
    gradf     :: Union{Nothing, Function}
    grad2f    :: Union{Nothing, Function}
    autodiff  :: Bool
end

# Options for the solver (more details like defaults in solver.jl)
mutable struct SolverOptions
    nl_solver                           :: Any# needs to be set
    log_levels                          :: Vector{Symbol}
    silent                              :: Bool
    atol                                :: Float64
    incumbent_constr                    :: Bool
    obj_epsilon                         :: Float64

    time_limit                          :: Float64
    mip_gap                             :: Float64

    best_obj_stop                       :: Float64
    solution_limit                      :: Int64
    all_solutions                       :: Bool

    list_of_solutions                   :: Bool
    mip_solver                          :: Any
    allow_almost_solved                 :: Bool
    allow_almost_solved_integral        :: Bool
    registered_functions                :: Union{Nothing, Vector{RegisteredFunction}}
    processors                          :: Int64
end

mutable struct SolutionObj
    solution    :: Vector{Float64}
    objval      :: Float64
end

#OriginalProblem
mutable struct OriginalProblem
    nl_solver           :: Any
    nl_solver_options   :: Vector{Pair}
    model               :: JuMP.Model
    relaxation_status   :: MOI.TerminationStatusCode
    relaxation_objval   :: Float64
    relaxation_solution :: Vector{Float64}
    status              :: MOI.TerminationStatusCode
    objval              :: Float64
    best_bound          :: Float64
    x                   :: Vector{JuMP.VariableRef}
    primal_start        :: Vector{Real}
    num_constr          :: Int64
    num_nl_constr       :: Int64
    num_q_constr        :: Int64
    num_l_constr        :: Int64
    num_var             :: Int64
    l_var               :: Vector{Float64}
    u_var               :: Vector{Float64}
    has_nl_objective    :: Bool
    nlp_evaluator       :: MOI.AbstractNLPEvaluator
    objective           :: Union{SVF, SAF, SQF, Nothing}
    disc2var_idx        :: Vector{Int64}
    var2disc_idx        :: Vector{Int64}
    var_type            :: Vector{Symbol}
    obj_sense           :: Symbol
    num_disc_var        :: Int64
    solution            :: Vector{Float64}
    soltime             :: Float64
    solutions           :: Vector{SolutionObj}
    nsolutions          :: Int64
    mip_solver          :: Any
    mip_solver_options  :: Vector{Pair}
    relaxation_time     :: Float64
    start_time          :: Float64
    # Info
    nintvars            :: Int64
    nbinvars            :: Int64
    options             :: SolverOptions
    nnodes              :: Int64

    nlp_constr_type     :: Vector{Symbol}

    OriginalProblem() = new()
end

#OriginalProblem
mutable struct OAdata
    mip_model                 :: JuMP.Model
    nlp_model                 :: JuMP.Model
    ref_nlp_model             :: JuMP.Model
    ref_mip_model             :: JuMP.Model
    ref_feasibility_model     :: JuMP.Model

    mip_x                     :: Vector{JuMP.VariableRef}
    nlp_x                     :: Vector{JuMP.VariableRef}
    ref_nlp_x                 :: Vector{JuMP.VariableRef}
    ref_mip_x                 :: Vector{JuMP.VariableRef}
    ref_feasibility_x         :: Vector{JuMP.VariableRef}

    num_l_constr              :: Int

    ref_l_var                 :: Vector{Float64}
    ref_u_var                 :: Vector{Float64}
    ref_var_type              :: Vector{Symbol}
    ref_num_var               :: Int
    ref_num_nl_constr         :: Int
    ref_nlp_constr_type       :: Vector{Symbol}

    obj_sense                 :: Symbol

    oa_status                 :: Symbol
    oa_started                :: Bool
    incumbent                 :: Vector{Float64}
    new_incumbent             :: Bool
    total_time                :: Float64
    obj_val                   :: Float64
    obj_bound                 :: Float64
    obj_gap                   :: Float64
    oa_iter                   :: Int64

    milp_sol_available        :: Bool
    ref_mip_solution          :: Vector{Float64}
    mip_infeasible            :: Bool
    ref_nlp_solution          :: Vector{Float64}
    nlp_infeasible            :: Bool
    ref_feasibility_solution  :: Vector{Float64}

    int_idx                   :: Vector{Int64}
    prev_ref_mip_solution     :: Vector{Float64}

    ref_linear_le_constraints ::Vector{Tuple{SAF, MOI.LessThan{Float64}}}
    ref_linear_ge_constraints ::Vector{Tuple{SAF, MOI.GreaterThan{Float64}}}
    ref_linear_eq_constraints ::Vector{Tuple{SAF, MOI.EqualTo{Float64}}}

    OAdata() = new()
end

"""
    is_type_correct(x, var_type, atol)
Check whether a variable x has the correct type
"""
function is_type_correct(x, var_type, atol)
    if var_type != :Cont
        if !isapprox(round(x)-x, 0; atol=atol)
            return false
        end
    end
    return true
end

"""
    are_type_correct(sol, types, disc2var_idx, atol)
Check whether all variables have the correct type
"""
function are_type_correct(sol, types, disc2var_idx, atol)
    for i in disc2var_idx
        if !isapprox(round(sol[i])-sol[i], 0; atol=atol)
            return false
        end
    end
    return true
end

"""
    all_reasonable_type_correct(sol, disc2var_idx, reasonable_idx, atol)
Check whether all reasonable variable have are discrete alreadyy
"""
function all_reasonable_type_correct(sol, disc2var_idx, reasonable_idx, atol)
    for i in reasonable_idx
        idx = disc2var_idx[i]
        if !isapprox(round(sol[idx])-sol[idx], 0; atol=atol)
            return false
        end
    end
    return true
end

max_stars_repo_name: MdMahbubarRahman/Duran

max_issues_repo_name: MdMahbubarRahman/Duran

max_forks_repo_name: MdMahbubarRahman/Duran


text:
# Dynamic QSP reporting with Julia

## Preparations

using HetaSimulator, StatsPlots

## Loading

p = load_platform(".", rm_out = true)
m = models(p)[:nameless]

## Default simulation

Scenario(m, (0.,30.)) |> sim |> plot
Scenario(m, (0.,30.), parameters = [:IR_total=>55., :GLUT4_total=>50, :diabetes=>0.15]) |> sim |> plot

## Simulation scenarios

scn_csv = read_scenarios("./julia/scenarios.csv")
add_scenarios!(p, scn_csv)
p

p |> sim |> plot

## Advanced visualization

results_df = sim(p) |> DataFrame # simulate all scenarios, then convert them
first(results_df, 7)

#reproducing data from the original study

### Fig5

@df results_df plot(:t, :measuredIRp, group = :scenario, title = "Fig5 1B")
@df results_df plot(:t, :measuredIRint, group = :scenario, title = "Fig5 1C")
# 1D
@df results_df plot(:t, :measuredIRS1p, group = :scenario, title = "Fig5 2B")
@df results_df plot(:t, :measuredIRS1307, group = :scenario, title = "Fig5 2C")
# 2D
@df results_df plot(:t, :measuredPKB308p, group = :scenario, title = "Fig5 3B")
@df results_df plot(:t, :measuredPKB473p, group = :scenario, title = "Fig5 3C")
# 3D
#@df results_df plot(:t, :measuredAS160p, group = :scenario, title = "Fig5 4B")
#@df results_df plot(:t, :measuredmTORC1a, group = :scenario, title = "Fig5 4C")
#@df results_df plot(:t, :measuredS6Kp, group = :scenario, title = "Fig5 4D")
# 5B
#@df results_df plot(:t, :measuredmTORC2a, group = :scenario, title = "Fig5 5C")
#@df results_df plot(:t, :measuredS6p, group = :scenario, title = "Fig5 5D")

### Titration-like simulations

scn_titr_csv = read_scenarios("./julia/titration.csv")
add_scenarios!(p, scn_titr_csv)

results_titr = sim(p, saveat=[10.]);
results_titr_df = DataFrame(results_titr, add_parameters=true);
results_titr_subset = subset(results_titr_df, :is_titr=>x->x.===1.);
first(results_titr_subset, 7)

@df results_titr_subset plot(
    :insulin, :measuredmTORC2a,
    title = "Fig5 1A",
    xlabel="insulin, nM", ylabel="mTORC2a measured",
    group=:diabetes,
    xscale = :log10,
    legend = :topleft
)

## Multiple simulations

results_mc = mc(
    p,
    [
        :k1a => LogNormal(0.6331, 0.5),
        :k1basal => LogNormal(0.03683, 0.5),
        :k1c => LogNormal(0.8768, 0.5),
        :k1d => LogNormal(31.01, 0.5)
    ],
    100,                      # number of
    scenarios = [:base, :ir], # scenarios names to simulate
    saveat = 0:(0.1):30,      # time points
    abstol = 1e-7
)

plot(results_mc, vars=[:measuredIRp, :measuredIRint])

summary_mc = EnsembleSummary(results_mc; quantiles=[0.05,0.95])
plot(summary_mc, idxs=[1,2]) # to plot observables 1 - IRp, 2 - IRint


text:
# ## Non-zero inflated, continuous. e.g. LogNormal
# function expert_tn(d::e, tl::Real, yl::Real, yu::Real, tu::Real) where {e<:NonZIContinuousExpert}
#     # There is no zero inflation: only from the component
#     # return expert_tn_pos(d, tl, yl, yu, tu)

#     # Possibly coming from the zero probability mass
#     expert_tn = (tl == 0.) ? log.(0.0 + (1-0.0)*exp.(expert_tn_pos(d, tl, yl, yu, tu))) : log.(0.0 + (1-0.0)*exp.(expert_tn_pos(d, tl, yl, yu, tu)))
#     # Must be from the zero probability mass
#     expert_tn = (tu == 0.) ? log.(0.0) : expert_tn

#     return expert_tn
# end

# ## Zero inflated, continuous. e.g. ZILogNormal
# function expert_tn(d::e, tl::Real, yl::Real, yu::Real, tu::Real) where {e<:ZIContinuousExpert}
#     # Possibly coming from the zero probability mass
#     p0 = params(d)[1]
#     expert_tn = (tl == 0.) ? log.(p0 + (1-p0)*exp.(expert_tn_pos(d, tl, yl, yu, tu))) : log.(0.0 + (1-p0)*exp.(expert_tn_pos(d, tl, yl, yu, tu)))
#     # Must be from the zero probability mass
#     expert_tn = (tu == 0.) ? log.(p0) : expert_tn

#     return expert_tn
# end

# ## Non-zero inflated, discrete. e.g. Poisson
# function expert_tn(d::e, tl::Real, yl::Real, yu::Real, tu::Real) where {e<:NonZIDiscreteExpert}
#     # There is no zero inflation: only from the component
#     return expert_tn_pos(d, tl, yl, yu, tu)
# end

# ## Zero inflated, discrete. e.g. ZIPoisson
# function expert_tn(d::e, tl::Real, yl::Real, yu::Real, tu::Real) where {e<:ZIDiscreteExpert}
#     # Possibly coming from the zero probability mass
#     p0 = params(d)[1]
#     expert_tn = (tl == 0.) ? log.(p0 + (1-p0)*exp.(expert_tn_pos(d, tl, yl, yu, tu))) : log.(0.0 + (1-p0)*exp.(expert_tn_pos(d, tl, yl, yu, tu)))

#     return expert_tn
# end

max_stars_repo_name: sparktseung/JuliaLRMoE

max_issues_repo_name: sparktseung/JuliaLRMoE

max_forks_repo_name: sparktseung/JuliaLRMoE


text:
import DataStructures: OrderedDict

# This set of tests directly uses the high-level `convert` functions
# And checks the behaviour is as expected.

F.def_GLOBAL_LXDEFS!()
cmd   = st -> F.convert_md(st)
chtml = t -> F.convert_html(t)
conv  = st -> st |> cmd |> chtml

set_curpath(path) =
    (F.set_var!(F.LOCAL_VARS, "fd_rpath", path); F.locvar("fd_rpath"))

# convenience function that squeezes out all whitespaces and line returns out of a string
# and checks if the resulting strings are equal. When expecting a specific string +- some
# spaces, this is very convenient. Use == if want to check exact strings.
isapproxstr(s1::AbstractString, s2::AbstractString) =
    isequal(map(s->replace(s, r"\s|\n"=>""), String.((s1, s2)))...)

# stricter than isapproxstr, just strips the outside.
(//)(s1::String, s2::String) = strip(s1) == strip(s2)

# this is a slightly ridiculous requirement but apparently the `eval` blocks
# don't play well with Travis nor windows while testing, so you just need to forcibly
# specify that LinearAlgebra and Random are used (even though the included block says
# the same thing).
if get(ENV, "CI", "false") == "true" || Sys.iswindows()
    import Pkg; Pkg.add("LinearAlgebra"); Pkg.add("Random");
    using LinearAlgebra, Random;
end

# takes md input and outputs the html (good for integration testing)
function seval(st)
    F.def_GLOBAL_VARS!()
    F.def_GLOBAL_LXDEFS!()
    m = F.convert_md(st)
    h = F.convert_html(m)
    return h
end

function set_globals()
    F.def_GLOBAL_VARS!()
    F.def_GLOBAL_LXDEFS!()
    empty!(F.LOCAL_VARS)
    F.def_LOCAL_VARS!()
    set_curpath("index.md")
    # F.FD_ENV[:CUR_PATH_WITH_EVAL] = ""
end

function explore_md_steps(mds)
    F.def_GLOBAL_VARS!()
    F.def_GLOBAL_LXDEFS!()
    empty!(F.RSS_DICT)

    F.def_LOCAL_VARS!()
    F.def_PAGE_HEADERS!()
    F.def_PAGE_EQREFS!()
    F.def_PAGE_BIBREFS!()
    F.def_PAGE_FNREFS!()
    F.def_PAGE_LINK_DEFS!()

    steps = OrderedDict{Symbol,NamedTuple}()

    # tokenize
    tokens = F.find_tokens(mds, F.MD_TOKENS, F.MD_1C_TOKENS)
    F.validate_footnotes!(tokens)
    F.validate_headers!(tokens)
    hrules = F.find_hrules!(tokens)
    F.find_indented_blocks!(tokens, mds)
    steps[:tokenization] = (tokens=tokens,)

    # ocblocks
    blocks, tokens = F.find_all_ocblocks(tokens, F.MD_OCB)
    toks_pre_ocb   = deepcopy(tokens)

    blocks2, tokens = F.find_all_ocblocks(tokens, vcat(F.MD_OCB2, F.MD_OCB_MATH))
    append!(blocks, blocks2)
    F.deactivate_inner_blocks!(blocks)
    F.merge_indented_blocks!(blocks, mds)
    F.filter_indented_blocks!(blocks)
    steps[:ocblocks] = (blocks=blocks, tokens=tokens)

    # filter
    filter!(τ -> τ.name ∉ F.L_RETURNS, tokens)
    steps[:filter] = (tokens=tokens, blocks=blocks)

    F.validate_and_store_link_defs!(blocks)

    # latex commands
    lxdefs, tokens, braces, blocks = F.find_lxdefs(tokens, blocks)

    # lxdefs = cat(F.pastdef.(collect(values(F.GLOBAL_LXDEFS))), lxdefs, dims=1)

    lxcoms, _ = F.find_lxcoms(tokens, lxdefs, braces)
    steps[:latex] = (lxdefs=lxdefs, tokens=tokens, braces=braces,
                     blocks=blocks, lxcoms=lxcoms)

    dbb = F.find_double_brace_blocks(toks_pre_ocb)

    F.process_mddefs(blocks, false)

    sp_chars = F.find_special_chars(tokens)
    steps[:spchars] = (spchars=sp_chars,)

    fnrefs = filter(τ -> τ.name == :FOOTNOTE_REF, tokens)
    steps[:fnrefs] = (fnrefs=fnrefs,)

    if !F.locvar("indented_code")
        filter!(b -> b.name != :CODE_BLOCK_IND, blocks)
    end

    b2insert = F.merge_blocks(lxcoms, F.deactivate_divs(blocks),
                              sp_chars, fnrefs, dbb, hrules)
    steps[:b2insert] = (b2insert=b2insert,)

    inter_md, mblocks = F.form_inter_md(mds, b2insert, lxdefs)
    steps[:inter_md] = (inter_md=inter_md, mblocks=mblocks)

    inter_html = F.md2html(inter_md; stripp=false)
    steps[:inter_html] = (inter_html=inter_html,)

    hstring   = F.convert_inter_html(inter_html, mblocks, lxdefs)
    steps[:hstring] = (hstring=hstring,)

    return steps
end

function explore_h_steps(hs, allvars=F.PageVars())
    steps = OrderedDict{Symbol,NamedTuple}()

    tokens = F.find_tokens(hs, F.HTML_TOKENS, F.HTML_1C_TOKENS)
    steps[:tokenization] = (tokens=tokens,)

    hblocks, tokens = F.find_all_ocblocks(tokens, F.HTML_OCB)
    filter!(hb -> hb.name != :COMMENT, hblocks)
    steps[:ocblocks] = (hblocks=hblocks, tokens=tokens)

    qblocks = F.qualify_html_hblocks(hblocks)
    steps[:qblocks] = (qblocks=qblocks,)

    fhs = F.process_html_qblocks(hs, allvars, qblocks)
    steps[:fhs] = (fhs=fhs,)

    return steps
end

gotd() = (flush_td(); cd(td); F.FOLDER_PATH[] = td)

function fs1()
    F.FD_ENV[:STRUCTURE] = v"0.1"
    gotd()
    F.clear_dicts()
    empty!(F.PATHS)
    F.set_paths!()
    mkdir(F.PATHS[:src])
    mkdir(F.PATHS[:src_pages])
    mkdir(F.PATHS[:libs])
    mkdir(F.PATHS[:src_css])
    mkdir(F.PATHS[:src_html])
    mkdir(F.PATHS[:assets])
    F.def_GLOBAL_VARS!()
    F.def_LOCAL_VARS!()
    F.set_var!(F.LOCAL_VARS, "fd_rpath", "index.md")
end

function fs2()
    F.FD_ENV[:STRUCTURE] = v"0.2"
    gotd()
    F.clear_dicts()
    empty!(F.PATHS)
    F.set_paths!()
    mkdir(F.PATHS[:site])
    mkdir(F.PATHS[:assets])
    mkdir(F.PATHS[:css])
    mkdir(F.PATHS[:layout])
    mkdir(F.PATHS[:libs])
    mkdir(F.PATHS[:literate])
    F.def_GLOBAL_VARS!()
    F.def_LOCAL_VARS!()
    F.set_var!(F.LOCAL_VARS, "fd_rpath", "index.md")
end

fdi(s) = fd2html(s; internal=true)


text:
# MIT License
#
# Copyright (c) 2018 Martin Biel
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""
    Optimizer(; <keyword arguments>)

Return a progressive-hedging optimizer.

...
# Arguments
- `subproblem_optimizer::AbstractOptimizer`: MathOptInterface solver capable of solving quadratic programs.
- `penalty::AbstractPenalizer = Fixed()`: Specify penalty update procedure (Fixed, Adaptive)
- `execution::AbstractExecuter = Serial`: Specify how algorithm should be executed (Serial, Synchronous, Asynchronous). Distributed variants requires worker cores.
- `penaltyterm::PenaltyTerm = Quadratic`: Specify penaltyterm variant ([`Quadratic`](@ref), [`Linearized`](@ref), [`InfNorm`](@ref), [`ManhattanNorm`][@ref])
- <keyword arguments>: Algorithm specific parameters, consult individual docstrings (see above list) for list of possible arguments and default values.
...
"""
mutable struct Optimizer <: AbstractStructuredOptimizer
    subproblem_optimizer
    sub_params::Dict{MOI.AbstractOptimizerAttribute, Any}
    execution::AbstractExecution
    penalizer::AbstractPenalizer
    penaltyterm::AbstractPenaltyTerm
    parameters::ProgressiveHedgingParameters{Float64}

    status::MOI.TerminationStatusCode
    primal_status::MOI.ResultStatusCode
    dual_status::MOI.ResultStatusCode
    raw_status::String
    solve_time::Float64
    progressivehedging::Union{AbstractProgressiveHedging, Nothing}

    function Optimizer(; subproblem_optimizer = nothing,
                       execution::AbstractExecution = nworkers() == 1 ? Serial() : Synchronous(),
                       penalty::AbstractPenalizer = Fixed(),
                       penaltyterm::AbstractPenaltyTerm = Quadratic(),
                       kw...)
        return new(subproblem_optimizer,
                   Dict{MOI.AbstractOptimizerAttribute, Any}(),
                   execution,
                   penalty,
                   penaltyterm,
                   ProgressiveHedgingParameters{Float64}(; kw...),
                   MOI.OPTIMIZE_NOT_CALLED,
                   MOI.NO_SOLUTION,
                   MOI.NO_SOLUTION,
                   "Progressive-hedging optimizer has not been run.",
                   NaN,
                   nothing)
    end
end

# Interface #
# ========================== #
function supports_structure(optimizer::Optimizer, ::ScenarioDecompositionStructure)
    return true
end

function default_structure(::UnspecifiedInstantiation, optimizer::Optimizer)
    if optimizer.execution isa Serial && nworkers() == 1
        return ScenarioDecomposition()
    else
        return DistributedScenarioDecomposition()
    end
end

function check_loadable(optimizer::Optimizer, ::ScenarioDecompositionStructure)
    if optimizer.subproblem_optimizer === nothing
        msg = "Subproblem optimizer not set. Consider setting `SubProblemOptimizer` attribute."
        throw(UnloadableStructure{Optimizer, ScenarioDecompositionStructure}(msg))
    end
    return nothing
end

function ensure_compatible_execution!(optimizer::Optimizer, ::ScenarioDecompositionStructure{2, 1, <:Tuple{ScenarioProblems}})
    if !(optimizer.execution isa Serial)
        @warn "Distributed execution policies are not compatible with a single-core scenario-decomposition structure. Switching to `Serial` execution by default."
        MOI.set(optimizer, Execution(), Serial())
    end
    return nothing
end

function ensure_compatible_execution!(optimizer::Optimizer, ::ScenarioDecompositionStructure{2, 1, <:Tuple{DistributedScenarioProblems}})
    if optimizer.execution isa Serial
        @warn "Serial execution not compatible with distributed scenario-decomposition structure. Switching to `Synchronous` execution by default."
        MOI.set(optimizer, Execution(), Synchronous())
    end
    return nothing
end

function load_structure!(optimizer::Optimizer, structure::ScenarioDecompositionStructure, x₀::AbstractVector)
    # Sanity check
    check_loadable(optimizer, structure)
    # Restore structure if optimization has been run before
    restore_structure!(optimizer)
    # Ensure that execution policy is compatible
    ensure_compatible_execution!(optimizer, structure)
    # Set subproblem optimizers
    set_subproblem_optimizer!(structure, optimizer.subproblem_optimizer)
    # Set subproblem optimizer attributes
    for (attr, value) in optimizer.sub_params
        MOI.set(scenarioproblems(structure), attr, value)
    end
    # Create new progressive-hedging algorithm
    optimizer.progressivehedging = ProgressiveHedgingAlgorithm(structure,
                                                               x₀,
                                                               optimizer.execution,
                                                               optimizer.penalizer,
                                                               optimizer.penaltyterm;
                                                               type2dict(optimizer.parameters)...)
    return nothing
end

function restore_structure!(optimizer::Optimizer)
    if optimizer.progressivehedging !== nothing
        restore_subproblems!(optimizer.progressivehedging)
    end
    return nothing
end

function MOI.optimize!(optimizer::Optimizer)
    if optimizer.progressivehedging === nothing
        throw(StochasticProgram.UnloadedStructure{Optimizer}())
    end
    # Run progressive-hedging procedure
    optimizer.status = optimizer.progressivehedging()
    # Check if optimal
    if optimizer.status == MOI.OPTIMAL
        optimizer.primal_status = MOI.FEASIBLE_POINT
        optimizer.dual_status = MOI.FEASIBLE_POINT
        optimizer.raw_status = "Progressive-hedging procedure converged to optimal solution."
    end
    # Extract solve time
    optimizer.solve_time = optimizer.progressivehedging.progress.tlast - optimizer.progressivehedging.progress.tinit
    return nothing
end

function num_iterations(optimizer::Optimizer)
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return num_iterations(optimizer.progressivehedging)
end

function optimizer_name(optimizer::Optimizer)
    return "$(str(optimizer.execution))Progressive-hedging with $(str(optimizer.penalizer))"
end

# MOI #
# ========================== #
function MOI.get(optimizer::Optimizer, ::MOI.Silent)
    return !MOI.get(optimizer, MOI.RawParameter("log"))
end

function MOI.set(optimizer::Optimizer, attr::MOI.Silent, flag::Bool)
    MOI.set(optimizer, MOI.RawParameter("log"), !flag)
    optimizer.sub_params[attr] = flag
    return nothing
end

function MOI.get(optimizer::Optimizer, param::MOI.RawParameter)
    name = Symbol(param.name)
    if !(name in fieldnames(ProgressiveHedgingParameters))
        error("Unrecognized parameter name: $(name).")
    end
    return getfield(optimizer.parameters, name)
end

function MOI.set(optimizer::Optimizer, param::MOI.RawParameter, value)
    name = Symbol(param.name)
    if !(name in fieldnames(ProgressiveHedgingParameters))
        error("Unrecognized parameter name: $(name).")
    end
    setfield!(optimizer.parameters, name, value)
    return nothing
end

function MOI.get(optimizer::Optimizer, ::MOI.TimeLimitSec)
    limit = MOI.get(optimizer, MOI.RawParameter("time_limit"))
    return isinf(limit) ? nothing : limit
end

function MOI.set(optimizer::Optimizer, ::MOI.TimeLimitSec, limit::Union{Real, Nothing})
    limit = limit === nothing ? Inf : limit
    MOI.set(optimizer, MOI.RawParameter("time_limit"), limit)
    return
end

function MOI.get(optimizer::Optimizer, ::PrimalTolerance)
    return MOI.get(optimizer, MOI.RawParameter("ϵ₁"))
end

function MOI.set(optimizer::Optimizer, ::PrimalTolerance, limit::Real)
    MOI.set(optimizer, MOI.RawParameter("ϵ₁"), limit)
    return nothing
end

function MOI.get(optimizer::Optimizer, ::DualTolerance)
    return MOI.get(optimizer, MOI.RawParameter("ϵ₂"))
end

function MOI.set(optimizer::Optimizer, ::DualTolerance, limit::Real)
    MOI.set(optimizer, MOI.RawParameter("ϵ₂"), limit)
    return nothing
end

function MOI.get(optimizer::Optimizer, ::MasterOptimizer)
    return MOI.get(optimizer, SubProblemOptimizer())
end

function MOI.get(optimizer::Optimizer, ::SubProblemOptimizer)
    if optimizer.subproblem_optimizer === nothing
        return MOI.get(optimizer, MasterOptimizer())
    end
    return MOI.OptimizerWithAttributes(optimizer.subproblem_optimizer, collect(optimizer.sub_params))
end

function MOI.set(optimizer::Optimizer, ::SubProblemOptimizer, optimizer_constructor)
    optimizer.subproblem_optimizer = optimizer_constructor
    # Clear any old parameters
    empty!(optimizer.sub_params)
    return nothing
end

function MOI.get(optimizer::Optimizer, ::SubProblemOptimizerAttribute, attr::MOI.AbstractOptimizerAttribute)
    if !haskey(optimizer.sub_params, attr)
        error("Subproblem optimizer attribute $(attr) has not been set.")
    end
    return optimizer.sub_params[attr]
end

function MOI.set(optimizer::Optimizer, ::SubProblemOptimizerAttribute, attr::MOI.AbstractOptimizerAttribute, value)
    optimizer.sub_params[attr] = value
    return nothing
end

function MOI.get(optimizer::Optimizer, param::RawSubProblemOptimizerParameter)
    moi_param = MOI.RawParameter(param.name)
    if !haskey(optimizer.sub_params, moi_param)
        error("Subproblem optimizer attribute $(param.name) has not been set.")
    end
    return optimizer.sub_params[moi_param]
end

function MOI.set(optimizer::Optimizer, param::RawSubProblemOptimizerParameter, value)
    moi_param = MOI.RawParameter(param.name)
    optimizer.sub_params[moi_param] = value
    return nothing
end

function MOI.get(optimizer::Optimizer, ::Execution)
    return optimizer.execution
end

function MOI.set(optimizer::Optimizer, ::Execution, execution::AbstractExecution)
    optimizer.execution = execution
    return nothing
end

function MOI.get(optimizer::Optimizer, ::Penalizer)
    return optimizer.penalizer
end

function MOI.set(optimizer::Optimizer, ::Penalizer, penalizer::AbstractPenalizer)
    optimizer.penalizer = penalizer
    return nothing
end

function MOI.get(optimizer::Optimizer, ::PenaltyTerm)
    return optimizer.penaltyterm
end

function MOI.set(optimizer::Optimizer, ::PenaltyTerm, penaltyterm::AbstractPenaltyTerm)
    optimizer.penaltyterm = penaltyterm
    return nothing
end

function MOI.get(optimizer::Optimizer, param::ExecutionParameter)
    return MOI.get(optimizer.execution, param)
end

function MOI.set(optimizer::Optimizer, param::ExecutionParameter, value)
    MOI.set(optimizer.execution, param, value)
    return nothing
end

function MOI.get(optimizer::Optimizer, param::PenalizationParameter)
    return MOI.get(optimizer.penalizer, param)
end

function MOI.set(optimizer::Optimizer, param::PenalizationParameter, value)
    MOI.set(optimizer.penalizer, param, value)
    return nothing
end

function MOI.get(optimizer::Optimizer, ::MOI.TerminationStatus)
    return optimizer.status
end

function MOI.get(optimizer::Optimizer, ::MOI.PrimalStatus)
    return optimizer.primal_status
end

function MOI.get(optimizer::Optimizer, ::MOI.DualStatus)
    return optimizer.dual_status
end

function MOI.get(optimizer::Optimizer, ::MOI.RawStatusString)
    return optimizer.raw_status
end

function MOI.get(optimizer::Optimizer, ::MOI.VariablePrimal, index::MOI.VariableIndex)
    if optimizer.progressivehedging === nothing
        throw(StochasticProgram.UnloadedStructure{Optimizer}())
    end
    return decision(optimizer.progressivehedging, index)
end

function MOI.get(optimizer::Optimizer, ::MOI.ConstraintPrimal, ci::MOI.ConstraintIndex)
    if optimizer.progressivehedging === nothing
        throw(StochasticProgram.UnloadedStructure{Optimizer}())
    end
    return scalar_subproblem_reduction(optimizer.progressivehedging) do subproblem
        return MOI.get(subproblem.optimizer, MOI.ConstraintPrimal(), ci)
    end
end

function MOI.get(optimizer::Optimizer, ::MOI.ConstraintDual, ci::MOI.ConstraintIndex)
    if optimizer.progressivehedging === nothing
        throw(StochasticProgram.UnloadedStructure{Optimizer}())
    end
    return scalar_subproblem_reduction(optimizer.progressivehedging) do subproblem
        return MOI.get(subproblem.optimizer, MOI.ConstraintDual(), ci)
    end
end

function MOI.get(optimizer::Optimizer, ::MOI.ObjectiveValue)
    if optimizer.progressivehedging === nothing
        throw(StochasticProgram.UnloadedStructure{Optimizer}())
    end
    return objective_value(optimizer.progressivehedging)
end

function MOI.get(optimizer::Optimizer, ::MOI.DualObjectiveValue)
    if optimizer.progressivehedging === nothing
        throw(StochasticProgram.UnloadedStructure{Optimizer}())
    end
    if optimizer.status == MOI.OPTIMAL
        return objective_value(optimizer.progressivehedging)
    elseif optimizer.status == MOI.INFEASIBLE
        sense = MOI.get(optimizer.progressivehedging.structure, MOI.ObjectiveSense())
        return sense == MOI.MAX_SENSE ? Inf : -Inf
    elseif optimizer.status == MOI.DUAL_INFEASIBLE
        sense = MOI.get(optimizer.progressivehedging.structure, MOI.ObjectiveSense())
        return sense == MOI.MAX_SENSE ? -Inf : Inf
    else
        return NaN
    end
end

function MOI.get(optimizer::Optimizer, ::MOI.SolveTime)
    return optimizer.solve_time
end

function MOI.get(optimizer::Optimizer, attr::Union{MOI.AbstractOptimizerAttribute, MOI.AbstractModelAttribute})
    # Fallback to through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.get(optimizer.progressivehedging.structure, attr)
end

function MOI.get(optimizer::Optimizer, attr::MOI.AbstractConstraintAttribute, ci::MOI.ConstraintIndex)
    # Fallback to first-stage optimizer through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.get(optimizer.progressivehedging.structure, attr, ci)
end

function MOI.set(optimizer::Optimizer, attr::Union{MOI.AbstractOptimizerAttribute, MOI.AbstractModelAttribute}, value)
    # Fallback to first-stage optimizer through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.set(optimizer.progressivehedging.structure, attr, value)
end

function MOI.set(optimizer::Optimizer, attr::MOI.AbstractVariableAttribute, index::MOI.VariableIndex, value)
    # Fallback to first-stage optimizer through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.set(optimizer.progressivehedging.structure, attr, index, value)
end

function MOI.set(optimizer::Optimizer, attr::MOI.AbstractConstraintAttribute, ci::MOI.ConstraintIndex, value)
    # Fallback to first-stage optimizer through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.set(optimizer.progressivehedging.structure, attr, ci, value)
end

function MOI.get(optimizer::Optimizer, attr::ScenarioDependentModelAttribute)
    # Fallback to subproblem optimizer through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.get(optimizer.progressivehedging.structure, attr)
end

function MOI.get(optimizer::Optimizer, attr::ScenarioDependentVariableAttribute, index::MOI.VariableIndex)
    # Fallback to subproblem optimizer through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.get(optimizer.progressivehedging.structure, attr, index)
end

function MOI.get(optimizer::Optimizer, attr::ScenarioDependentConstraintAttribute, ci::MOI.ConstraintIndex)
    # Fallback to subproblem optimizer through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.get(optimizer.progressivehedging.structure, attr, ci)
end

function MOI.set(optimizer::Optimizer, attr::ScenarioDependentModelAttribute, value)
    # Fallback to subproblem optimizer through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.set(optimizer.progressivehedging.structure, attr, value)
end

function MOI.set(optimizer::Optimizer, attr::ScenarioDependentVariableAttribute, index::MOI.VariableIndex, value)
    # Fallback to subproblem optimizer through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.set(optimizer.progressivehedging.structure, attr, index, value)
end

function MOI.set(optimizer::Optimizer, attr::ScenarioDependentConstraintAttribute, ci::MOI.ConstraintIndex, value)
    # Fallback to subproblem optimizer through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.set(optimizer.progressivehedging.structure, attr, ci, value)
end

function MOI.is_empty(optimizer::Optimizer)
    return optimizer.progressivehedging === nothing
end

MOI.supports(::Optimizer, ::MOI.Silent) = true
MOI.supports(::Optimizer, ::MOI.TimeLimitSec) = true
MOI.supports(::Optimizer, ::MOI.RawParameter) = true
MOI.supports(::Optimizer, ::AbstractStructuredOptimizerAttribute) = true
MOI.supports(::Optimizer, ::RelativeTolerance) = false
MOI.supports(::Optimizer, ::RawInstanceOptimizerParameter) = true
MOI.supports(::Optimizer, ::AbstractProgressiveHedgingAttribute) = true

# High-level attribute setting #
# ========================== #
"""
    get_penalization_attribute(stochasticprogram::StochasticProgram, name::String)

Return the value associated with the penalization-specific attribute named `name` in `stochasticprogram`.

See also: [`set_penalization_attribute`](@ref), [`set_penalization_attributes`](@ref).
"""
function get_penalization_attribute(stochasticprogram::StochasticProgram, name::String)
    return return MOI.get(optimizer(stochasticprogram), RawPenalizationParameter(name))
end
"""
    set_penalization_attribute(stochasticprogram::StochasticProgram, name::Union{Symbol, String}, value)

Sets the penalization-specific attribute identified by `name` to `value`.

"""
function set_penalization_attribute(stochasticprogram::StochasticProgram, name::Union{Symbol, String}, value)
    return set_optimizer_attribute(stochasticprogram, RawPenalizationParameter(String(name)), value)
end
"""
    set_penalization_attributes(stochasticprogram::StochasticProgram, pairs::Pair...)

Given a list of `attribute => value` pairs or a collection of keyword arguments, calls
`set_penalization_attribute(stochasticprogram, attribute, value)` for each pair.

"""
function set_penalization_attributes(stochasticprogram::StochasticProgram, pairs::Pair...)
    for (name, value) in pairs
        set_penalization_attribute(stochasticprogram, name, value)
    end
end
function set_penalization_attributes(stochasticprogram::StochasticProgram; kw...)
    for (name, value) in kw
        set_penalization_attribute(stochasticprogram, name, value)
    end
end


text:
###########################################################################
#
#  Robust adaptive Metropolis (RAM)
#
#  Reference: Vihola M. Robust Adaptive Metropolis Algorithm with Coerced Acceptance Rate. Statistics and Computing,
#  2012, 22 (5), pp 997-1008.
#
#  Parameters:
#    -scale: for scaling the jumps
#    -rate: target acceptance rate
#
###########################################################################

export RAM

println("Loading RAM(scale, rate) sampler")

###########################################################################
#                  RAM type
###########################################################################

immutable RAM <: MCMCSampler
  scale::Float64
  rate::Float64

  function RAM(x::Float64, r::Float64)
    @assert x > 0 "scale should be > 0"
    @assert r > 0. && r < 1. "target acceptance rate ($r) should be between 0 and 1"
    new(x, r)
  end
end
RAM() = RAM(1., 0.234)
RAM(x::Float64) = RAM(x, 0.234)
RAM(;scale::Float64=1.0, rate::Float64=0.234) = RAM(scale, rate)

###########################################################################
#                  RAM task
###########################################################################

# RAM sampling
function SamplerTask(model::MCMCModel, sampler::RAM, runner::MCMCRunner)
	local pars, proposedPars
	local logTarget, proposedLogTarget
  local scale

	# hook inside Task to allow remote resetting
	task_local_storage(:reset, (resetPars::Vector{Float64}) -> (pars = copy(resetPars); logTarget = model.eval(pars)))

	# initialization
	scale = model.scale .* sampler.scale  # rescale model scale by sampler scale
	pars = copy(model.init)
	logTarget = model.eval(pars)
  @assert isfinite(logTarget) "Initial values out of model support, try other values"

	S = Float64[ i==j ? scale[i] : 0. for i in 1:model.size, j in 1:model.size]

	# main loop
	for i in 1:Inf
		rvec = randn(model.size)
		proposedPars = pars + S * rvec
		proposedLogTarget = model.eval(proposedPars)

	  ratio = proposedLogTarget-logTarget
		if ratio > 0 || (ratio > log(rand()))
			ms = MCMCSample(proposedPars, proposedLogTarget, pars, logTarget, {"accept" => true, "scale" => trace(S)})
      produce(ms)
      pars, logTarget = copy(proposedPars), copy(proposedLogTarget)
    else
			ms = MCMCSample(pars, logTarget, pars, logTarget, {"accept" => false, "scale" => trace(S)})
      produce(ms)
    end

    # scale tuning
		eta = min(1, model.size * i^(-2/3))  # decreasing influence with i

		SS = (rvec * rvec') / dot(rvec, rvec) * eta * (min(1, exp(ratio)) - sampler.rate)
		SS = S * (eye(model.size) + SS) * S'
		S = chol(SS)'  #'
  end
end


text:
"""
   micrm_params

   Type containing parameters for micrm simulations

    # Arguments
    N::Int64 - number of consumers
    M::Int64 - number of resources
    u::Matrix{Float64} - uptake matrix
    m::Vector{Float64} - consumer respiration/loss vector
    ρ::Vector{Float64} - resource inflow vector
    ω::Vector{Float64} - resource outflow vector
    l::Matrix{Float64} - leakage matrix
"""
struct micrm_params
        N::Int64
        M::Int64
        u::Matrix{Float64}
        m::Vector{Float64}
        ρ::Vector{Float64}
        ω::Vector{Float64}
        l::Matrix{Float64}
end

"""
    micrm_params(p::Dict{Symbol,Any})

    constructor function to convert parameter dictionary into `micrm_params` object. The parameter dictionary must have the following fields: `[:l,:ρ,:N,:M,:R,:ω,:u]`
"""
function micrm_params(p::Dict{Symbol,Any})
    micrm_params(p[:N],p[:M],p[:u],p[:m],p[:ρ],p[:ω],p[:l])
end

"""
    dx!(dx,x,p,t)

    Derivative function to simulate the micrm directly from a set parameters. To be used to generate an `ODEProblem` object.
"""
@views function dx!(dx,x,p,t)
    for i = 1:p.N #consumer
        dx[i] = -x[i] * p.m[i] #loss term
        for α = 1:p.M
            dx[i] += x[i] * x[α + p.N] * p.u[i,α] #uptake
            for β = 1:p.M
                dx[i] += -x[i] * x[α + p.N] * p.u[i,α] * p.l[α,β] #leakage
            end
        end
    end

    for α = 1:p.M #resource
        dx[α + p.N] = p.ρ[α] - (x[α + p.N] * p.ω[α])
        for i = 1:p.N
            dx[α + p.N] += -x[α + p.N] * x[i] * p.u[i,α]
            for β = 1:p.M
                dx[α + p.N] += x[β + p.N] * x[i] * p.u[i,β] * p.l[β,α]
            end
        end
    end
end


text:
module TestExtras

export @constinferred
export @timedtestset
export ConstInferred

include("constinferred.jl")
include("timedtest.jl")

using .ConstInferred: @constinferred
using .TimedTests: @timedtestset

end


text:
@testset "ZeroProblems" begin

struct MyProb <: zp.AbstractContinuationProblem end

pb = MyProb()

coll = zp.Var(:coll, 20)
x0 = zp.Var(:x0, 1, coll, 0)
x1 = zp.Var(:x1, 1, coll, -1)

u = (coll, x0, x1)
ui = [1:size(coll), 1:1, size(coll):size(coll)]
D = ((1), (2, 3), ())
ϕ = (sum, prod, maximum)
ϕi = [1:1, 2:2, 3:3]

prob = zp.ZeroProblem{D, length(u), typeof(ϕ)}(u, ui, ϕ, ϕi)

res = zeros(3)
zp.residual!(res, prob, pb, rand(20))

end