function Sim_Model(Cell,Input,Def,Tk,SList,SOC,A0,B0,C0,D0,t)
    """
    Function to simulate generated reduced-order models.

    Sim_Model(Cell,Input,Def,Tk,SOC,A,B,C,D)

    """
    #Determine time span and allocate arrays
    tlength = size(Input,1)

    A = Array{Float64}(undef,size(A0[1]))
    B = Array{Float64}(undef,size(B0[1]))
    C = Array{Float64}(undef,size(C0[1]))
    D = Array{Float64}(undef,size(D0[1]))

    #Selecting SS Models
    for i in 1:length(SList)
        if SOC == SList[i] # Fix this, round?
            A = A0[i]
            B = B0[i]
            C = C0[i]
            D = D0[i]
        end
    end

    #Capturing Indices
    tfstr = Array{String}(undef,0,1)
    for i in 1:length(Cell.Transfer.tfs)
        t1 = t2 = Array{String}(undef,0,1)
        for j in 1:length(Cell.Transfer.Locs[i])
            t1 ="$(Cell.Transfer.tfs[i])_$(Cell.Transfer.Elec[i])"
            t2 = [t2; t1]
        end
        tfstr = [tfstr; t2]
    end

    CeInd = findall(isequal("C_e_Na"), tfstr)
    ϕ_ẽInd = findall(isequal("Phi_e_Na"), tfstr)
    CsePosInd = findall(isequal("C_se_Pos"), tfstr)
    CseNegInd = findall(isequal("C_se_Neg"), tfstr)
    ϕ_sNegInd = findall(isequal("Phi_s_Neg"), tfstr)
    ϕ_sPosInd = findall(isequal("Phi_s_Pos"), tfstr)
    ϕ_sePosInd = findall(isequal("Phi_se_Pos"), tfstr)
    ϕ_seNegInd = findall(isequal("Phi_se_Neg"), tfstr)
    FluxNegInd = findall(isequal("Flux_Neg"), tfstr)
    FluxPosInd = findall(isequal("Flux_Pos"), tfstr)

    #@infiltrate cond = true

    CeNeg = Cell.Transfer.Locs[1].<=Cell.Neg.L
    CeNegOffset = Cell.Transfer.Locs[1].<Cell.Neg.L
    CeSep = Cell.Transfer.Locs[1].<=Cell.Neg.L+Cell.Sep.L
    CeSepOffset = Cell.Transfer.Locs[1].<Cell.Neg.L+Cell.Sep.L
    CePos = Cell.Transfer.Locs[1].<=Cell.Const.Ltot
    CeNegInd = findall(CeNeg .== 1)
    CeSepInd = findall(CeSep.-CeNegOffset .== 1)
    CePosInd = findall(CePos.-CeSepOffset .== 1)

    csegain_neg = C[CseNegInd[1][1],1] #First Column in C Array (zeros column)
    csegain_pos = C[CsePosInd[1][1],1] #First Column in C Array (zeros column)

    #Memory Allocation
    θ_neg = Array{Float64}(undef,tlength,1) .= 0.
    θ_pos = Array{Float64}(undef,tlength,1) .= 0.
    jeq_neg = Array{Float64}(undef,tlength,1) .= 0.
    jeq_pos = Array{Float64}(undef,tlength,1) .= 0.
    x = Array{Float64}(undef,tlength+1,size(A,1)) .= 0.
    y = Array{Float64}(undef,tlength,size(C,1)) .= 0.
    Cse_Neg = Array{Float64}(undef,tlength,size(CseNegInd,1)) .= 0.
    Cse_Pos = Array{Float64}(undef,tlength,size(CsePosInd,1)) .= 0.
    Ce = Array{Float64}(undef,tlength,size(CeInd,1)) .= Cell.Const.ce0
    η0 = Array{Float64}(undef,tlength,1) .= 0.
    η_neg = Array{Float64}(undef,tlength,size(FluxNegInd,1)) .= 0.
    ηL = Array{Float64}(undef,tlength,1) .= 0.
    η_pos = Array{Float64}(undef,tlength,size(FluxPosInd,1)) .= 0.
    ϕ_ẽ1 = Array{Float64}(undef,tlength,size(ϕ_ẽInd,1)) .= 0.
    ϕ_ẽ2 = Array{Float64}(undef,tlength,size(CeInd,1)) .= 0.
    ϕ_se_neg_0 = Array{Float64}(undef,tlength,1) .= 0. #Replace with length of ϕ_seNegInd @ zero
    jNeg = Array{Float64}(undef,tlength,size(FluxNegInd,1)) .= 0.
    j0 = Array{Float64}(undef,tlength,1) .= 0.
    jPos = Array{Float64}(undef,tlength,size(FluxPosInd,1)) .= 0.
    jL = Array{Float64}(undef,tlength,1) .= 0.
    j0_CC_neg = Array{Float64}(undef,tlength,1) .= 0.
    Rtot_neg = Array{Float64}(undef,tlength,1) .= 0.
    Rtot_pos = Array{Float64}(undef,tlength,1) .= 0.
    j0_CC_pos = Array{Float64}(undef,tlength,1) .= 0.
    Uocp_Neg = Array{Float64}(undef,tlength,1) .= 0.
    Uocp_Pos = Array{Float64}(undef,tlength,1) .= 0.
    Cell_V = Array{Float64}(undef,tlength,1) .= 0.
    ϕ_e = Array{Float64}(undef,tlength,size(CeInd,1)) .= 0.
    Cell_SOC = Array{Float64}(undef,tlength,1) .= 0
    Iapp = Array{Float64}(undef,tlength+1,1) .= 0

    #Defining SOC
    SOC_Neg = SOC * (Cell.Neg.θ_100-Cell.Neg.θ_0) + Cell.Neg.θ_0
    SOC_Pos = SOC * (Cell.Pos.θ_100-Cell.Pos.θ_0) + Cell.Pos.θ_0
    θ_neg[1] = SOC_Neg
    θ_pos[1] = SOC_Pos
    Cell_SOC[1] = (SOC_Neg-Cell.Neg.θ_0)/(Cell.Neg.θ_100-Cell.Neg.θ_0)

    #Loop through time
    #Compute dependent variables (voltage, flux, etc.)
    for i in 0:(tlength-1)
        cs_neg_avg = x[i+1,1] * csegain_neg + SOC_Neg * Cell.Neg.cs_max < 0. ? 0. : x[i+1,1] * csegain_neg + SOC_Neg * Cell.Neg.cs_max #Zero if < 0
        cs_pos_avg = x[i+1,1] * csegain_pos + SOC_Pos * Cell.Pos.cs_max < 0. ? 0. : x[i+1,1] * csegain_pos + SOC_Pos * Cell.Pos.cs_max #Zero if < 0

        if cs_neg_avg > Cell.Neg.cs_max
            cs_neg_avg = Cell.Neg.cs_max
        end

        if cs_pos_avg > Cell.Pos.cs_max
            cs_pos_avg = Cell.Pos.cs_max
        end

        #Reaction Rates
        if Cell.Const.CellTyp == "Doyle_94"
            k_neg = Cell.Neg.k_norm#/Cell.Neg.cs_max/Cell.Const.ce0^(1-Cell.Neg.α)
            k_pos = Cell.Pos.k_norm#/Cell.Pos.cs_max/Cell.Const.ce0^(1-Cell.Pos.α)
            jeq_neg[i+1] = k_neg*sqrt(cs_neg_avg*Cell.Const.ce0*(Cell.Neg.cs_max-cs_neg_avg))
            jeq_pos[i+1] = k_pos*sqrt(cs_pos_avg*Cell.Const.ce0*(Cell.Pos.cs_max-cs_pos_avg))
        else
            #jeq_neg = Cell.Neg.k_norm*Cell.Neg.cs_max*(Cell.Const.ce0*(cs_neg_avg/Cell.Neg.cs_max*(1-cs_neg_avg/Cell.Neg.cs_max)))^(1-Cell.Neg.α)
            #jeq_pos = Cell.Pos.k_norm*Cell.Pos.cs_max*(Cell.Const.ce0*(cs_pos_avg/Cell.Pos.cs_max*(1-cs_pos_avg/Cell.Pos.cs_max)))^(1-Cell.Pos.α)
            k_neg = Cell.Neg.k_norm
            k_pos = Cell.Pos.k_norm
            jeq_neg[i+1] = Cell.Neg.k_norm*(Cell.Const.ce0*cs_neg_avg*(Cell.Neg.cs_max-cs_neg_avg))^(1-Cell.Neg.α)
            jeq_pos[i+1] = Cell.Pos.k_norm*(Cell.Const.ce0*cs_pos_avg*(Cell.Pos.cs_max-cs_pos_avg))^(1-Cell.Pos.α)
        end

        θ_neg[i+1] = cs_neg_avg/Cell.Neg.cs_max
        θ_pos[i+1] = cs_pos_avg/Cell.Pos.cs_max
        Cell_SOC[i+1] = (θ_neg[i+1]-Cell.Neg.θ_0)/(Cell.Neg.θ_100-Cell.Neg.θ_0)

        javg_neg = Iapp[i+1]/(Cell.Neg.as*F*Cell.Neg.L*Cell.Const.CC_A)
        javg_pos = Iapp[i+1]/(Cell.Pos.as*F*Cell.Pos.L*Cell.Const.CC_A)

        Arr_Factor = ((1/Cell.Const.T_ref)-(1/Tk[i+1]))/R
        κneg = Cell.Const.κf(mean(Ce[i+1,CeNegInd]))*exp(Cell.Const.Ea_κ*Arr_Factor)
        κpos = Cell.Const.κf(mean(Ce[i+1,CePosInd]))*exp(Cell.Const.Ea_κ*Arr_Factor)
        κsep = Cell.Const.κf(mean(Ce[i+1,CeSepInd]))*exp(Cell.Const.Ea_κ*Arr_Factor)
        σ_neg = Cell.Neg.σ*exp(Cell.Const.Ea_κ*Arr_Factor)
        σ_pos = Cell.Pos.σ*exp(Cell.Const.Ea_κ*Arr_Factor)
        κ_eff_Neg = κneg*(Cell.Neg.ϵ_e^(Cell.Neg.κ_brug))
        κ_eff_Sep = κsep*(Cell.Sep.ϵ_e^(Cell.Sep.κ_brug))
        κ_eff_Pos = κpos*(Cell.Pos.ϵ_e^(Cell.Pos.κ_brug))
        σ_eff_Neg = σ_neg*Cell.Neg.ϵ_s^Cell.Neg.σ_brug #Effective Conductivity Neg
        σ_eff_Pos = σ_pos*Cell.Pos.ϵ_s^Cell.Pos.σ_brug #Effective Conductivity Pos

        #Resistances
        #Rtot_neg[i] = (Tk[i]*R)/(F^2*mean([abs(jeq_neg[i]) abs(javg_neg)]))+Cell.Neg.RFilm
        Rtot_neg[i+1] = (Tk[i+1]*R)/(F^2*sqrt(jeq_neg[i+1]^2+javg_neg^2/4))+Cell.Neg.RFilm
        #Rtot_pos[i] = (Tk[i]*R)/(F^2*mean([abs(jeq_pos[i]) abs(javg_pos)]))+Cell.Pos.RFilm
        Rtot_pos[i+1] = (Tk[i+1]*R)/(F^2*sqrt(jeq_pos[i+1]^2+javg_pos^2/4))+Cell.Pos.RFilm

        #Condensing Variable
        ν_neg = @. Cell.Neg.L*sqrt((Cell.Neg.as*(1/κ_eff_Neg+1/σ_eff_Neg))/Rtot_neg[i+1])
        ν_pos = @. Cell.Pos.L*sqrt((Cell.Pos.as*(1/κ_eff_Pos+1/σ_eff_Pos))/Rtot_pos[i+1])

        #Relinearise dependent on ν, σ, κ
        D = D_Linear(Cell, ν_neg, ν_pos, σ_eff_Neg, κ_eff_Neg, σ_eff_Pos, κ_eff_Pos, κ_eff_Sep)

        #Interpolate C & D Matrices
        C = interp(C0,SList,Cell_SOC[i+1])
        #D = interp(D0,SList,Cell_SOC[i+1])
        #SS Output
        y[i+1,:] = C*x[i+1,:] + D*Iapp[i+1]

        #Concentrations & Force electrode concentration maximum
        Cse_Neg[i+1,:] = (SOC_Neg.*Cell.Neg.cs_max .+ y[i+1,CseNegInd]) > ones(size(Cse_Neg,2))*Cell.Neg.cs_max ? ones(size(Cse_Neg,2))*Cell.Neg.cs_max : (SOC_Neg.*Cell.Neg.cs_max .+ y[i+1,CseNegInd])
        Cse_Pos[i+1,:] = (SOC_Pos.*Cell.Pos.cs_max .+ y[i+1,CsePosInd]) > ones(size(Cse_Pos,2))*Cell.Pos.cs_max ? ones(size(Cse_Pos,2))*Cell.Pos.cs_max : (SOC_Pos.*Cell.Pos.cs_max .+ y[i+1,CsePosInd])
        Ce[i+1,:] = @. Cell.Const.ce0 + y[i+1,CeInd]

        #Potentials
        Uocp_Neg[i+1] = Cell.Const.Uocp("Neg",Cse_Neg[i+1,1]/Cell.Neg.cs_max)
        Uocp_Pos[i+1] = Cell.Const.Uocp("Pos",Cse_Pos[i+1,1]/Cell.Pos.cs_max)
        ϕ_se_neg_0[i+1] = y[i+1, ϕ_seNegInd[1]] + Uocp_Neg[i+1] #Location 0
        ϕ_ẽ1[i+1,:] = y[i+1,ϕ_ẽInd]
        ϕ_ẽ2[i+1,:] = @. ((Tk[i+1]*2*R*(1-Cell.Const.tpf(Ce[i+1,:])))/F)*(log(Ce[i+1,:]/Ce[i+1,1]))
        ϕ_e[i+1,:] = @. [0; ϕ_ẽ1[i+1,:]]+ϕ_ẽ2[i+1,:]-ϕ_se_neg_0[i+1]

        #Flux
        jNeg[i+1,:] = y[i+1,FluxNegInd]
        j0[i+1] = y[i+1,FluxNegInd[1]]
        jPos[i+1,:] = y[i+1,FluxPosInd]
        jL[i+1] = y[i+1,FluxPosInd[1]]

        #j0_CC_neg[i+1] = findmax([eps(); (k_neg*(Cell.Neg.cs_max-Cse_Neg[i+1,1])^(1-Cell.Neg.α))*((Cse_Neg[i+1,1]^Cell.Neg.α)*(Ce[i+1,1]^(1-Cell.Neg.α)))])[1]
        j0_neg = findmax([ones(size(Cse_Neg,2))*eps() (k_neg.*(Cse_Neg[i+1,:].^Cell.Neg.α).*(Ce[i+1,1].^(1-Cell.Neg.α))).*(Cell.Neg.cs_max.-Cse_Neg[i+1,:]).^(1-Cell.Neg.α)], dims=2)[1]

        η0[i+1] = Tk[i+1]*2*R/F*asinh(j0[i+1]/(2*j0_neg[1]))
        η_neg[i+1,:] = @. (Tk[i+1]*2*R)/F*asinh((jNeg[i+1,:])/(2*j0_neg))

        #j0_CC_pos[i+1] = findmax([eps(); (k_pos*(Cell.Pos.cs_max-Cse_Pos[i+1,1])^(1-Cell.Pos.α))*((Cse_Pos[i+1,1]^Cell.Pos.α)*(Ce[i+1,end]^(1-Cell.Pos.α)))])[1]
        j0_pos = findmax([ones(size(Cse_Pos,2))*eps() (k_pos.*(Cse_Pos[i+1,:].^Cell.Pos.α).*(Ce[i+1,1].^(1-Cell.Pos.α))).*(Cell.Pos.cs_max.-Cse_Pos[i+1,:]).^(1-Cell.Pos.α)], dims=2)[1]
        ηL[i+1] = (Tk[i+1]*2*R)/F*asinh(jL[i+1]/(2*j0_pos[1]))
        η_pos[i+1,:] = @. (Tk[i+1]*2*R)/F*asinh(jPos[i+1,:]/(2*j0_pos))

        #Cell Voltage
        Cell_V[i+1] = @. (Uocp_Pos[i+1]-Uocp_Neg[i+1]) + (ηL[i+1]-η0[i+1]) + (ϕ_ẽ1[i+1,end]+ϕ_ẽ2[i+1,end]) + (Cell.Pos.RFilm*jL[i+1]-Cell.Neg.RFilm*j0[i+1])*F

        #ϕ_s
        ϕ_s_neg = y[i+1,ϕ_sNegInd]
        ϕ_s_pos = @. y[i+1,ϕ_sPosInd] + Cell_V[i+1]

        #Interpolate A Matrix
        A = interp(A0,SList,Cell_SOC[i+1])
        B = interp(B0,SList,Cell_SOC[i+1])

        #Update States
        x[i+2,:] = A*x[i+1,:] + B*Iapp[i+1]

        if Def =="Power"
            Iapp[i+2] = Input[i+1,1]/Cell_V[i+1]
        else
            Iapp[i+2] = Input[i+1,1]
        end

    end

return Cell_V, Ce, jNeg, jPos, Rtot_neg, Rtot_pos, η0, ηL, η_neg, η_pos, ϕ_ẽ1, ϕ_ẽ2, Uocp_Neg, Uocp_Pos, ϕ_e, Cse_Neg, Cse_Pos, Cell_SOC, jeq_neg, jeq_pos, j0, jL, t
end


function build_master(sub_problems::T where T <: Dict, tree::T where T <: AbstractTree, probabilities::Dict{AbstractTree,Float64}, solver, discount_factor::Float64,
  risk::Union{Risk,Array{Risk,1}}, sideconstraints)

	if typeof(solver) <: Tuple
		model = Model(solver[1])
	else
		model = Model(solver)
	end
	@objective(model,Min,0)

	model.ext[:columns] = Dict{AbstractTree,Array{Column,1}}()
	for node in collect(tree)
		model.ext[:columns][node]=Array{Column,1}()
	end
	leafs=Array{Leaf,1}()

	model.ext[:expansions] = Dict{AbstractTree,Dict{Symbol,Any}}()

	model.ext[:options]=Dict{Symbol,Tuple}()
	for (name,options) in sub_problems[tree].ext[:options]
		model.ext[:options][name]=options
	end

	form=nothing
	for (name,options) in model.ext[:options]
		if options[4]==:Bin
			if form==:continuous
				form=:mixed
				break
			else
				form=:binary
			end
		elseif options[4]==:Int
			form=:mixed
			break
		else
			if form==:binary
				form=:mixed
				break
			else
				form=:continuous
			end
		end
	end

	model.ext[:discrete_con]=Dict{AbstractTree,Dict{Int64,ConstraintRef}}()
	model.ext[:discrete_var]=Dict{AbstractTree,Dict{Int64,VariableRef}}()

	for (node,sp) in sub_problems
		if haskey(node.ext,:sum_max) || haskey(node.ext,:sum_min) || haskey(node.ext,:sum_set)
			sp.ext[:form]=:binarycolumns
		else
			sp.ext[:form]=form
		end
		if sp.ext[:form] ∉ [:binary,:binarycolumns]
			vars = all_variables(sp)
			if form==:continuous
				for var in vars
					if is_binary(var) || is_integer(var)
						sp.ext[:form]=:mixed
						break
					end
				end
			end

			if sp.ext[:form]==:mixed
				sp.ext[:discrete]=VariableRef[]
				sp.ext[:discrete_branch]=VariableRef[]
				ignore=VariableRef[]

				for (name,variable) in sp.ext[:expansions]
					if typeof(variable) <: AbstractArray
						for index in eachindex(variable)
							push!(ignore,variable[index])
						end
					else
						push!(ignore,variable)
					end
				end

				for var in vars
					if (is_binary(var) || is_integer(var)) && var ∉ ignore
						push!(sp.ext[:discrete],var)
						push!(sp.ext[:discrete_branch],var)
					end
				end
			end
		end
	end

	for (node,sp) in sub_problems
		if typeof(node)==Leaf
			push!(leafs,node)
		end
		model.ext[:expansions][node] = Dict{Symbol,Any}()

		for (name,variable) in sp.ext[:expansions]
			if sp.ext[:options][name][4]==:Bin
				model.ext[:expansions][node][name] = copy_variable!(model, variable, relaxbinary)
			elseif sp.ext[:options][name][4]==:Int
				model.ext[:expansions][node][name] = copy_variable!(model, variable, relaxinteger)
			else
				model.ext[:expansions][node][name] = copy_variable!(model, variable)
			end
			if model.ext[:options][name][5]!=nothing
				if typeof(variable) <: AbstractArray
					for i in eachindex(variable)
						set_lower_bound(model.ext[:expansions][node][name][i],model.ext[:options][name][5])
					end
				else
					set_lower_bound(model.ext[:expansions][node][name],model.ext[:options][name][5])
				end
			end
			if model.ext[:options][name][6]!=nothing
				if typeof(variable) <: AbstractArray
					for i in eachindex(variable)
						set_upper_bound(model.ext[:expansions][node][name][i],model.ext[:options][name][6])
					end
				else
					set_upper_bound(model.ext[:expansions][node][name],model.ext[:options][name][6])
				end
			end
		end
	end

	model.ext[:scenprofit_var] = Dict{Leaf,VariableRef}()
	model.ext[:scenprofit_con] = Dict{Leaf,ConstraintRef}()

	for leaf in leafs
		model.ext[:scenprofit_var][leaf] = @variable(model)
		model.ext[:scenprofit_con][leaf] = @constraint(model, 0 == model.ext[:scenprofit_var][leaf])

  		set_objective_coefficient(model, model.ext[:scenprofit_var][leaf], probabilities[leaf])
	end

	risk_objectives=AffExpr[]
	remain=1.0
	if typeof(risk)==Risk
		if (risk.α==1.0 || risk.λ==0.0) && risk.bound==nothing
			risk=[]
		else
			risk=[risk]
		end
	end
    for i in 1:length(risk)
		risk_objective=AffExpr(0.0)
		eta=@variable(model)
	    for leaf in leafs
			v = @variable(model)
			w = @variable(model)
			set_lower_bound(v,0.0)
			set_lower_bound(w,0.0)
			if risk[i].offset==nothing
				@constraint(model,v>=eta-model.ext[:scenprofit_var][leaf])
				@constraint(model,w>=model.ext[:scenprofit_var][leaf]-eta)
				add_to_expression!(risk_objective,(model.ext[:scenprofit_var][leaf])*probabilities[leaf])
			else
				@constraint(model,v>=eta-model.ext[:scenprofit_var][leaf]+risk[i].offset[leaf])
				@constraint(model,w>=model.ext[:scenprofit_var][leaf]-risk[i].offset[leaf]-eta)
				add_to_expression!(risk_objective,(model.ext[:scenprofit_var][leaf]-risk[i].offset[leaf])*probabilities[leaf])
			end
			add_to_expression!(risk_objective,probabilities[leaf]*(v+w/risk[i].α*(1-risk[i].α)))
		end
		remain-=risk[i].λ
		push!(risk_objectives,risk_objective)
	end

	for leaf in leafs
		nodes=history(leaf)
		for n in eachindex(nodes)
			node=nodes[n]
			sp=sub_problems[node]
			df=discount_factor^depth(node)
			for (name,variable) in sp.ext[:expansions]
				interval=max(1,n-sp.ext[:options][name][3]-sp.ext[:options][name][2]+1):n-sp.ext[:options][name][2]
				disc=Dict{Int64,Float64}()
				for i in interval
					disc[i]=df/discount_factor^(i-1)
				end
				if typeof(variable) <: AbstractArray
					for i in eachindex(variable)
						cost_coef=df*coef(sp.ext[:capitalcosts],variable[i])
						for j in interval
							cost_coef+=disc[j]*coef(sp.ext[:ongoingcosts],variable[i])
						end
						set_normalized_coefficient(model.ext[:scenprofit_con][leaf], model.ext[:expansions][node][name][i], cost_coef)
					end
				else
					cost_coef=df*coef(sp.ext[:capitalcosts],variable)
					for j in interval
						cost_coef+=disc[j]*coef(sp.ext[:ongoingcosts],variable)
					end
					set_normalized_coefficient(model.ext[:scenprofit_con][leaf], model.ext[:expansions][node][name], cost_coef)
				end
			end
		end
	end

	# create the cover constraints
	model.ext[:coverconstraint] = Dict{AbstractTree,Dict{Symbol,Any}}()
	for (node,sp) in sub_problems
		model.ext[:coverconstraint][node] = Dict{Symbol,Any}()
		past=history(node)
		for (name,variable) in sp.ext[:expansions]
			interval=sp.ext[:options][name][2]+1:min(sp.ext[:options][name][2]+sp.ext[:options][name][3],length(past))
			if typeof(variable) <: AbstractArray
				model.ext[:coverconstraint][node][name] = Dict()
				for i in eachindex(variable)
					if sp.ext[:options][name][1]==:shutdown
						model.ext[:coverconstraint][node][name][i] = @constraint(model, 0 >= sum(model.ext[:expansions][past[index]][name][i] for index in interval))
					elseif sp.ext[:options][name][1]==:expansion
						model.ext[:coverconstraint][node][name][i] = @constraint(model, 0 <= sum(model.ext[:expansions][past[index]][name][i] for index in interval))
					elseif sp.ext[:options][name][1]==:enforced
						model.ext[:coverconstraint][node][name][i] = @constraint(model, 0 == sum(model.ext[:expansions][past[index]][name][i] for index in interval))
					elseif sp.ext[:options][name][1]==:state
						if node.parent==nothing
							model.ext[:coverconstraint][node][name][i] = @constraint(model, 0 == -sp.ext[:options][name][7] + model.ext[:expansions][node][name][i])
						else
							model.ext[:coverconstraint][node][name][i] = @constraint(model, 0 == model.ext[:expansions][node][name][i] - model.ext[:expansions][node.parent][name][i])
						end
					end
				end
				# if typeof(node)==Leaf && !sp.ext[:options][name][4]
				# 	for i in eachindex(variable)
				# 		@constraint(model, sum(model.ext[:expansions][past][name][i] for past in history_function(node))<=1)
				# 	end
				# end
			else
				if sp.ext[:options][name][1]==:shutdown
					model.ext[:coverconstraint][node][name] = @constraint(model, 0 >= sum(model.ext[:expansions][past[index]][name] for index in interval))
				elseif sp.ext[:options][name][1]==:expansion
					model.ext[:coverconstraint][node][name] = @constraint(model, 0 <= sum(model.ext[:expansions][past[index]][name] for index in interval))
				elseif sp.ext[:options][name][1]==:enforced
					model.ext[:coverconstraint][node][name] = @constraint(model, 0 == sum(model.ext[:expansions][past[index]][name] for index in interval))
				elseif sp.ext[:options][name][1]==:state
					if node.parent==nothing
						model.ext[:coverconstraint][node][name] = @constraint(model, 0 == -sp.ext[:options][name][7] + model.ext[:expansions][node][name])
					else
						model.ext[:coverconstraint][node][name] = @constraint(model, 0 == model.ext[:expansions][node][name] - model.ext[:expansions][node.parent][name])
					end
				end
				# if typeof(node)==Leaf && !sp.ext[:options][name][4]
				# 	@constraint(model, sum(model.ext[:expansions][past][name] for past in history_function(node))<=1)
				# end
			end
		end
	end

	model.ext[:convexcombination] = Dict{AbstractTree,Array{ConstraintRef,1}}()
	model.ext[:mip]=false
	model.ext[:branch_cons]=ConstraintRef[]

	for node in keys(sub_problems)
		model.ext[:convexcombination][node]=ConstraintRef[]
		if haskey(node.ext,:sum_max)
			push!(model.ext[:convexcombination][node],@constraint(model, 0 <= node.ext[:sum_max]))
		end
		if haskey(node.ext,:sum_min)
			push!(model.ext[:convexcombination][node],@constraint(model, 0 >= node.ext[:sum_min]))
		end
		if haskey(node.ext,:sum_set)
			if length(model.ext[:convexcombination][node])>0
				error(":eq and :le/:ge specifed for node "*node.name)
			end
			push!(model.ext[:convexcombination][node],@constraint(model, 0 == node.ext[:sum_set]))
		end
		if length(model.ext[:convexcombination][node])==0
			push!(model.ext[:convexcombination][node],@constraint(model, 0 == 1))
		end
	end

	if typeof(sideconstraints) <: Function
		map(Main.eval,unpack_expansions(model.ext[:expansions])) #bring expansion variables into global scope
		sideconstraints(model,tree)
		map(Main.eval,clear_expansions(model.ext[:expansions]))
	end

	objective_fn=objective_function(model)*remain

	for i in 1:length(risk_objectives)
		objective_fn+=risk_objectives[i]*risk[i].λ
		if risk[i].bound!=nothing
			surplus=@variable(model)
			set_lower_bound(surplus,0)
			@constraint(model,risk_objectives[i]<=risk[i].bound+surplus)
			objective_fn+=risk[i].penalty*surplus
		end
	end

	set_objective_function(model,objective_fn)

	model
end

function solve_master_binary(judge::JuDGEModel,initial_time::Float64,termination::Termination,warm_starts::Bool,mp_callback,verbose)
	set=0
	current=nothing
	if (termination.allow_frac==:binary_solve || termination.allow_frac==:binary_solve_return_relaxation) #&& judge.log[end].int>0
		if verbose==2
			print("Solving master problem as MIP")
		end
		set=1
		if solve_binary(judge,termination.abstol,termination.reltol,warm_starts,mp_callback)
			current = ConvergenceState(judge.log[end].obj, judge.bounds.UB, judge.bounds.LB, time()-initial_time, judge.log[end].iter + 1, 0)
			if verbose>0
				if verbose==2
					overprint("")
				end
				display(current,relaxation=false)
			end
			push!(judge.log,current)
		else
			if verbose==2
				overprint("")
			end
		end
	end
	if termination.allow_frac==:binary_solve_return_relaxation && set==1
		remove_binary(judge)
		optimize!(judge.master_problem)
	end
	current
end

function solve_binary(judge::JuDGEModel,abstol::Float64,reltol::Float64,warm_starts::Bool,mp_callback)
	if typeof(judge.master_solver) <: Tuple
		JuMP.set_optimizer(judge.master_problem,judge.master_solver[2])
	end

	judge.master_problem.ext[:mip]=true

	for node in keys(judge.master_problem.ext[:expansions])
		if judge.sub_problems[node].ext[:form]==:binarycolumns
			for col in judge.master_problem.ext[:columns][node]
			 	if !haskey(col.node.ext,:col_max)
			 		set_binary(col.var)
				else
			 		set_integer(col.var)
			 	end
			 end
		 end

		for x in keys(judge.master_problem.ext[:expansions][node])
			var = judge.master_problem.ext[:expansions][node][x]
			if judge.sub_problems[judge.tree].ext[:options][x][4]==:Bin
				if typeof(var) <: JuMP.Containers.SparseAxisArray
					for key in keys(var.data)
						set_binary(var.data[key])
					end
				elseif typeof(var) <: AbstractArray
					for key in keys(var)
						set_binary(var[key])
					end
				else
					set_binary(var)
				end
			elseif judge.sub_problems[judge.tree].ext[:options][x][4]==:Int
				if typeof(var) <: JuMP.Containers.SparseAxisArray
					for key in keys(var.data)
						set_integer(var.data[key])
					end
				elseif typeof(var) <: AbstractArray
					for key in keys(var)
						set_integer(var[key])
					end
				else
					set_integer(var)
				end
			end
		end
		if judge.sub_problems[node].ext[:form]==:mixed
			for (i,var) in judge.master_problem.ext[:discrete_var][node]
				set_integer(var)
			end
		end
	end

	if typeof(mp_callback) <: Function
		mp_callback(judge,abstol,reltol)
		printright("Initialising")
	end
	optimize!(judge.master_problem)
	status=termination_status(judge.master_problem)
	if status!=MathOptInterface.INFEASIBLE_OR_UNBOUNDED && status!=MathOptInterface.INFEASIBLE && status!=MathOptInterface.DUAL_INFEASIBLE
		obj = objective_value(judge.master_problem)
		if obj<judge.bounds.UB
			judge.bounds.UB=obj
			if warm_starts
				vars=all_variables(judge.master_problem)
				set_start_value.(vars, JuMP.value.(vars))
			end
		end
		return true
	else
		return false
	end
end

function remove_binary(judge::JuDGEModel)
	# if judge.master_problem.ext[:binarycolumns]
	# 	for col in judge.master_problem.ext[:columns]
	# 		if !haskey(col.node.ext,:col_max)
	# 			unset_binary(col.var)
	# 		else
	# 			unset_integer(col.var)
	# 		end
	# 	end
	# end
	for node in keys(judge.master_problem.ext[:expansions])
		if judge.sub_problems[node].ext[:form]==:binarycolumns
			for col in judge.master_problem.ext[:columns][node]
			 	if !haskey(col.node.ext,:col_max)
			 		unset_binary(col.var)
				else
			 		unset_integer(col.var)
			 	end
			 end
		 end
		for x in keys(judge.master_problem.ext[:expansions][node])
			var = judge.master_problem.ext[:expansions][node][x]
			if judge.sub_problems[judge.tree].ext[:options][x][4]==:Bin
				if typeof(var) <: AbstractArray
					for key in eachindex(var)
						unset_binary(var[key])
					end
				else
					unset_binary(var)
				end
			elseif judge.sub_problems[judge.tree].ext[:options][x][4]==:Int
				if typeof(var) <: AbstractArray
					for key in eachindex(var)
						unset_integer(var[key])
					end
				else
					unset_integer(var)
				end
			end
		end
		if judge.sub_problems[node].ext[:form]==:mixed
			for (i,var) in judge.master_problem.ext[:discrete_var][node]
				unset_integer(var)
			end
		end
	end

	judge.master_problem.ext[:mip]=false

	if typeof(judge.master_solver) <: Tuple
		JuMP.set_optimizer(judge.master_problem,judge.master_solver[1])
	end
end


export extract_amp, extract_vis, extract_lcamp, extract_cphase

function getvisfield(obs)
    obsamps = obs.data::PyObject
    u = deepcopy((get(obsamps, Vector{Float64}, "u")))
    v = deepcopy((get(obsamps, Vector{Float64}, "v")))
    err = deepcopy((get(obsamps, Vector{Float64}, "sigma")))
    vis = deepcopy((get(obsamps, Vector{Complex{Float64}}, "vis")))
    t1 = Symbol.(deepcopy((get(obsamps, Vector{String}, "t1"))))
    t2 = Symbol.(deepcopy((get(obsamps, Vector{String}, "t2"))))
    baseline = tuple.(t1, t2)
    time = deepcopy((get(obsamps, Vector{Float64}, "time")))
    freq = zeros(length(time))
    bw = zeros(length(time))

    return  StructArray{Comrade.EHTVisibilityDatum{Float64}}(
        visr = real.(vis),
        visi = imag.(vis),
        u = u,
        v = v,
        error = err,
        time = time,
        frequency = freq,
        bandwidth = bw,
        baseline = baseline
    )
end

function getampfield(obs)
    obsamps = obs.amp::PyObject
    uamp = deepcopy(get(obsamps, Vector{Float64}, "u"))
    vamp = deepcopy(get(obsamps, Vector{Float64}, "v"))
    erramp = deepcopy(get(obsamps, Vector{Float64}, "sigma"))
    amps = deepcopy(get(obsamps, Vector{Float64}, "amp"))
    t1 = Symbol.(deepcopy(get(obsamps, Vector{String}, "t1")))
    t2 = Symbol.(deepcopy(get(obsamps, Vector{String}, "t2")))
    baseline = tuple.(t1, t2)
    time = deepcopy(get(obsamps, Vector{Float64}, "time"))
    freq = zeros(length(time))
    bw = zeros(length(time))

    return  StructArray{Comrade.EHTVisibilityAmplitudeDatum{Float64}}(
        amp = amps,
        u = uamp,
        v = vamp,
        error = erramp,
        time = time,
        frequency = freq,
        bandwidth = bw,
        baseline = baseline
    )
end

function getcpfield(obs)
    obscp = obs.cphase::PyObject
    u1 = deepcopy((get(obscp, Vector{Float64}, "u1")))
    v1 = deepcopy((get(obscp, Vector{Float64}, "v1")))
    u2 = deepcopy((get(obscp, Vector{Float64}, "u2")))
    v2 = deepcopy((get(obscp, Vector{Float64}, "v2")))
    u3 = deepcopy((get(obscp, Vector{Float64}, "u3")))
    v3 = deepcopy((get(obscp, Vector{Float64}, "v3")))
    cp = deg2rad.(deepcopy((get(obscp, Vector{Float64}, "cphase"))))
    errcp = deg2rad.(deepcopy((get(obscp, Vector{Float64}, "sigmacp"))))

    t1 = Symbol.(deepcopy((get(obscp, Vector{String}, "t1"))))
    t2 = Symbol.(deepcopy((get(obscp, Vector{String}, "t2"))))
    t3 = Symbol.(deepcopy((get(obscp, Vector{String}, "t3"))))
    baseline = tuple.(t1, t2, t3)
    time = deepcopy(get(obscp, Vector{Float64}, "time"))
    freq = zeros(length(time))
    bw = zeros(length(time))

    return StructArray{Comrade.EHTClosurePhaseDatum{Float64}}(
        phase = cp,
        u1 = u1,
        v1 = v1,
        u2 = u2,
        v2 = v2,
        u3 = u3,
        v3 = v3,
        error = errcp,
        time = time,
        frequency = freq,
        bandwidth = bw,
        triangle = baseline
    )

end

function getlcampfield(obs)
    obslcamp = obs.logcamp::PyObject
    u1 = get(obslcamp, Vector{Float64}, "u1")
    v1 = get(obslcamp, Vector{Float64}, "v1")
    u2 = get(obslcamp, Vector{Float64}, "u2")
    v2 = get(obslcamp, Vector{Float64}, "v2")
    u3 = get(obslcamp, Vector{Float64}, "u3")
    v3 = get(obslcamp, Vector{Float64}, "v3")
    u4 = get(obslcamp, Vector{Float64}, "u4")
    v4 = get(obslcamp, Vector{Float64}, "v4")
    camp = ((get(obslcamp, Vector{Float64}, "camp")))
    errcamp = ((get(obslcamp, Vector{Float64}, "sigmaca")))

    t1 = Symbol.(((get(obslcamp, Vector{String}, "t1"))))
    t2 = Symbol.(((get(obslcamp, Vector{String}, "t2"))))
    t3 = Symbol.(((get(obslcamp, Vector{String}, "t3"))))
    t4 = Symbol.(((get(obslcamp, Vector{String}, "t4"))))
    baseline = tuple.(t1, t2, t3, t4)
    time = (get(obslcamp, Vector{Float64}, "time"))
    freq = zeros(length(time))
    bw = zeros(length(time))

    return StructArray{Comrade.EHTLogClosureAmplitudeDatum{Float64}}(
        amp = camp,
        u1 = u1,
        v1 = v1,
        u2 = u2,
        v2 = v2,
        u3 = u3,
        v3 = v3,
        u4 = u4,
        v4 = v4,
        error = errcamp,
        time = time,
        frequency = freq,
        bandwidth = bw,
        quadrangle = baseline
    )

end

function getradec(obs)::Tuple{Float64, Float64}
    return (float(obs.ra), float(obs.dec))
end

function getmjd(obs)::Int
    return Int(obs.mjd)
end

function getsource(obs)::Symbol
    return Symbol(obs.source)
end

"""
    extract_amp(obs)
Extracts the visibility amplitudes from an ehtim observation object.

Any valid keyword arguments to `add_amp` in ehtim can be passed through extract_amp.

Returns an EHTObservation with visibility amplitude data
"""
function extract_amp(obsc; kwargs...)
    obs = obsc.copy()
    obs.reorder_tarr_snr()
    obs.add_amp(;kwargs...)
    data = getampfield(obs)
    ra, dec = getradec(obs)
    mjd = getmjd(obs)
    source = getsource(obs)
    bw = obs.bw
    rf = obs.rf
    ac = _arrayconfig(data, bw, rf)
    return Comrade.EHTObservation(data = data, mjd = mjd,
                   ra = ra, dec= dec,
                   config = ac,
                   bandwidth=bw, frequency=rf,
                   source = source,
    )
end

"""
    extract_vis(obs)
Extracts the complex visibilities from an ehtim observation object

This grabs the raw `data` object from the obs object. Any keyword arguments are ignored.

Returns an EHTObservation with complex visibility data
"""
function extract_vis(obsc; kwargs...)
    obs = obsc.copy()
    obs.reorder_tarr_snr()

    data = getvisfield(obs)
    ra, dec = getradec(obs)
    mjd = getmjd(obs)
    source = getsource(obs)
    bw = obs.bw
    rf = obs.rf
    ac = _arrayconfig(data, bw, rf)
    return Comrade.EHTObservation(
                   data = data, mjd = mjd,
                   config=ac,
                   ra = ra, dec= dec,
                   bandwidth=bw, frequency=rf,
                   source = source,
    )
end

function closure_designmat(scancp::Scan{A,B,C}, scanvis) where {A,B,C<:StructArray{<:EHTClosurePhaseDatum}}
    ant1, ant2, ant3 = baselines(scancp)
    antvis1, antvis2 = baselines(scanvis)
    design_mat = zeros(length(scancp), length(scanvis))
    #throw("here")
    # fill in each row of the design matrix
    for i in axes(design_mat, 2), j in axes(design_mat,1)
        a1 = ant1[j]
        a2 = ant2[j]
        a3 = ant3[j]
        # check leg 1
        ((antvis1[i] == a1) & (antvis2[i] == a2)) && (design_mat[j,i] = 1.0)
        ((antvis1[i] == a2) & (antvis2[i] == a1)) && (design_mat[j,i] = -1.0)

        #check leg 2
        ((antvis1[i] == a2) & (antvis2[i] == a3)) && (design_mat[j,i] = 1.0)
        ((antvis1[i] == a3) & (antvis2[i] == a2)) && (design_mat[j,i] = -1.0)

        #check leg 3
        ((antvis1[i] == a3) & (antvis2[i] == a1)) && (design_mat[j,i] = 1.0)
        ((antvis1[i] == a1) & (antvis2[i] == a3)) && (design_mat[j,i] = -1.0)
    end
    return design_mat
end

function closure_designmat(scanlca::Scan{A,B,C}, scanvis) where {A,B,C<:StructArray{<:EHTLogClosureAmplitudeDatum}}
    ant1, ant2, ant3, ant4 = baselines(scanlca)
    antvis1, antvis2 = baselines(scanvis)
    design_mat = zeros(length(scanlca), length(scanvis))
    # fill in each row of the design matrix
    for i in axes(design_mat, 2), j in axes(design_mat,1)
        a1 = ant1[j]
        a2 = ant2[j]
        a3 = ant3[j]
        a4 = ant4[j]

        av1 = antvis1[i]
        av2 = antvis2[i]
        # check leg 1
        (((av1 == a1)&(av2 == a2)) || (av1 == a2)&(av2 == a1))&& (design_mat[j,i] = 1.0)

        #check leg 2
        (((av1 == a3)&(av2 == a4)) || (av1 == a4)&(av2 == a3))&& (design_mat[j,i] = 1.0)

        #check leg 3
        (((av1 == a1)&(av2 == a4)) || (av1 == a4)&(av2 == a1))&& (design_mat[j,i] = -1.0)

        #check leg 4
        (((av1 == a2)&(av2 == a3)) || (av1 == a3)&(av2 == a2))&& (design_mat[j,i] = -1.0)

    end
    return design_mat
end

function minimal_lcamp(obsc; kwargs...)

    obs = obsc.copy()
    # reorder to maximize the snr
    obs.reorder_tarr_snr()

    lcamp = _ehtim_lcamp(obs; count="max", kwargs...)
    stlca = scantable(lcamp)

    #Now make the vis obs
    dvis = extract_vis(obsc)
    st = scantable(dvis)

    minset, dmat = _minimal_closure(stlca, st)

    # now create EHTObservation
    ra, dec = getradec(obs)
    mjd = getmjd(obs)
    source = getsource(obs)
    bw = obs.bw
    rf = obs.rf
    ac = arrayconfig(dvis)
    clac = ClosureConfig(ac, dmat)
    return EHTObservation(data = minset, mjd = mjd,
                          config=clac,
                          ra = ra, dec= dec,
                          bandwidth=bw, frequency=rf,
                          source = source,
                        )

end

function _ehtim_cphase(obsc; count="max", cut_trivial=false, uvmin=0.1e9, kwargs...)
    obs = obsc.copy()

    # cut 0 baselines since these are trivial triangles
    if cut_trivial
        obs = obs.flag_uvdist(uv_min=uvmin)
    end

    obs.reorder_tarr_snr()

    obs.add_cphase(;count=count, kwargs...)
    data = getcpfield(obs)
    ra, dec = getradec(obs)
    mjd = getmjd(obs)
    source = getsource(obs)
    bw = obs.bw
    rf = obs.rf

    cphase = EHTObservation(data = data, mjd = mjd,
                   config=nothing,
                   ra = ra, dec= dec,
                   bandwidth=bw, frequency=rf,
                   source = source,
    )

    stcp = scantable(cphase)

    #Now make the vis obs
    dvis = extract_vis(obsc)
    st = scantable(dvis)
    S = eltype(dvis[:visr])

    dmat = Matrix{S}[]
    for i in 1:length(st)
        scanvis = st[i]
        ind = findfirst(==(scanvis.time), stcp.times)
        if isnothing(ind)
            inow = length(dmat)
            # I want to construct block diagonal matrices for efficienty but we need
            # to be careful with scan that can't form closures. This hack moves these
            # scans into the preceding scan but fills it with zeros
            if i > 1
                dmat[inow] = hcat(dmat[inow], zeros(S, size(dmat[inow],1), length(scanvis.scan)))
            else
                push!(dmat, zeros(S, 1, length(scanvis.scan)))
            end
            continue
        end
        dmatscan = closure_designmat(stcp[ind], st[i])
        push!(dmat, dmatscan)
    end

    if iszero(dmat[1])
        dmat[2] = hcat(zeros(S, size(dmat[2],1), size(dmat[1],2)), dmat[2])
        dmat = dmat[2:end]
    end

    ac = arrayconfig(dvis)
    clac = ClosureConfig(ac, dmat)
    return  EHTObservation(data = data, mjd = mjd,
                           config=clac,
                           ra = ra, dec= dec,
                           bandwidth=bw, frequency=rf,
                           source = source
                          )
end

function _make_lcamp(obsc, count="max"; kwargs...)
    obs = obsc.copy()

    obs.reorder_tarr_snr()

    obs.add_logcamp(;count=count, kwargs...)
    data = getlcampfield(obs)
    ra, dec = getradec(obs)
    mjd = getmjd(obs)
    source = getsource(obs)
    bw = obs.bw
    rf = obs.rf

    return EHTObservation(data = data, mjd = mjd,
                   config=nothing,
                   ra = ra, dec= dec,
                   bandwidth=bw, frequency=rf,
                   source = source,
    )

end

function _ehtim_lcamp(obsc; count="max", kwargs...)

    lcamp = _make_lcamp(obsc; count, kwargs...)

    stlca = scantable(lcamp)

    #Now make the vis obs
    dvis = extract_vis(obsc)
    st = scantable(dvis)
    S = eltype(dvis[:visr])

    dmat = Matrix{S}[]
    for i in 1:length(st)
        scanvis = st[i]
        ind = findfirst(==(scanvis.time), stlca.times)
        if isnothing(ind)
            inow = length(dmat)
            # I want to construct block diagonal matrices for efficienty but we need
            # to be careful with scan that can't form closures. This hack moves these
            # scans into the preceding scan but fills it with zeros
            if i > 1
                dmat[inow] = hcat(dmat[inow], zeros(S, size(dmat[inow],1), length(scanvis.scan)))
            else
                push!(dmat, zeros(S, 1, length(scanvis.scan)))
            end
            continue
        end
        dmatscan = closure_designmat(stlca[ind], st[i])
        push!(dmat, dmatscan)
    end

    if iszero(dmat[1])
        dmat[2] = hcat(zeros(S, size(dmat[2],1), size(dmat[1],2)), dmat[2])
        dmat = dmat[2:end]
    end

    ac = arrayconfig(dvis)
    clac = ClosureConfig(ac, dmat)
    return  EHTObservation(data = lcamp.data, mjd = lcamp.mjd,
                           config=clac,
                           ra = lcamp.ra, dec= lcamp.dec,
                           bandwidth=lcamp.bandwidth, frequency=lcamp.frequency,
                           source = lcamp.source
                          )
end

function minimal_cphase(obsc; kwargs...)
    # compute a maximum set of closure phases
    obs = obsc.copy()

    # reorder to maximize the snr
    obs.reorder_tarr_snr()

    cphase = _ehtim_cphase(obs; count="max", kwargs...)
    stcp = scantable(cphase)

    #Now make the vis obs
    dvis = extract_vis(obsc)
    st = scantable(dvis)

    minset, dmat = _minimal_closure(stcp, st)
    # now create EHTObservation
    ra, dec = getradec(obs)
    mjd = getmjd(obs)
    source = getsource(obs)
    bw = obs.bw
    rf = obs.rf
    ac = arrayconfig(dvis)
    clac = ClosureConfig(ac, dmat)
    return EHTObservation(data = minset, mjd = mjd,
                          config=clac,
                          ra = ra, dec= dec,
                          bandwidth=bw, frequency=rf,
                          source = source,
                        )
end

function _minimal_closure(stcl, st)

    # Determine the number of timestamps containing a closure triangle
    T = typeof(stcl[1][1])
    S = typeof(stcl[1][1].u1)
    # loop over all timestamps
    minset = StructVector{T}(undef, 0)
    dmat = Matrix{S}[]
    for i in 1:length(st.times)

        scanvis = st[i]
        # find closure scan that matches time
        ind = findfirst(==(scanvis.time), stcl.times)
        if isnothing(ind)
            inow = length(dmat)
            # I want to construct block diagonal matrices for efficiency but we need
            # to be careful with scans that can't form closures. This hack moves these
            # scans into the preceding scan but fills it with zeros
            if i > 1
                dmat[inow] = hcat(dmat[inow], zeros(S, size(dmat[inow],1), length(scanvis.scan)))
            else
                push!(dmat, zeros(S, 1, length(scanvis.scan)))
            end
            continue
        end

        # Get our cphase scan
        scancl = stcl[ind]

        # sort by closure SNR so we form a nice minimal set
        snr = inv.(scancl.scan.error)
        ind_snr = sortperm(snr)
        scancl = scancl[ind_snr]
        snr = snr[ind_snr]

        # now find the visibility scan that matches the closure one based on timestamps

        # initialize the design matrix
        design_mat = closure_designmat(scancl, scanvis)

        # determine the expected size of the minimal set
        # this is needed to make sure we aren't killing too many triangles
        nmin = rank(design_mat)

        # print some info
        #println("For timestamp $(stcl.times[i]):")

        # get the current stations
        #println("Observing stations are $(stations(scancl))")

        #println("Size of maximal set of closure products = $(length(scancl))")
        #println("Size of minimal set of closure products = $(nmin)")
        #println("...")

        ##########################################################
        # start of loop to recover minimal set
        ##########################################################
        minset_scan, dmat_min = minimal_closure_scan(scancl, scanvis, nmin)
        minset = append!!(minset, minset_scan.scan)
        push!(dmat, dmat_min)
    end

    # Now if the first scan can't form a closure we will have an extra row of zeros
    # kill this
    if iszero(dmat[1])
        dmat[2] = hcat(zeros(S, size(dmat[2],1), size(dmat[1],2)), dmat[2])
        dmat = dmat[2:end]
    end
    return minset, dmat
end

function minimal_closure_scan(scancl0, scanvis, nmin::Int)
    # make a mask to keep track of which clhases will stick around
    scancl = deepcopy(scancl0)
    keep = fill(true, length(scancl))

    # remember the original minimal set size
    nmin0 = nmin

    # perform the loop
    count = 1
    keep = fill(true, length(scancl))
    good = true
    while good
        # recreate the mask each time
        scankeep = scancl[keep]

        design_mat = closure_designmat(scankeep, scanvis)

        # determine the size of the minimal set
        nmin = rank(design_mat)
        #println(nmin0, " ", nmin, " ", size(design_mat, 1))

        # Now decide whether to continue pruning
        if (sum(keep) == nmin0) && (nmin == nmin0)
            good = false
        else
            if nmin == nmin0
                keep[count] = false
            else
                keep[count-1] = true
                count -= 1
            end
            #(nmin == nmin0) && (keep[count] = false)
        end
        count += 1
        count+1 > length(scancl) && break
    end

    # print out the size of the recovered set for double-checking
    scankeep = scancl[keep]
    dmat = closure_designmat(scankeep, scanvis)
    if length(scankeep) != nmin
        @error "minimal set not found $(length(scankeep)) $(nmin)"
        throw("No minimal set found at time $(scancl.time)")
    end
    return scankeep, dmat
end

"""
    extract_cphase(obs)
Extracts the closure phases from an ehtim observation object

Any valid keyword arguments to `add_cphase` in ehtim can be passed through extract_cphase.

Returns an EHTObservation with closure phases datums

# Keyword arguments:
 - count: How the closures are formed, the available options are "min-correct", "min", "max"
 - cut_trivial: Cut the trivial triangles from the closures
 - uvmin: The flag to decide what are trivial triangles. Any baseline with ||(u,v)|| < uvmin
          are removed.

# Warning

The `count` keyword argument is treated specially in `Comrade`. The default option
is "min-correct" and should almost always be used.
This option construct a minimal set of closure phases that is valid even when
the array isn't fully connected. For testing and legacy reasons we `ehtim` other count
options are also included. However, the current `ehtim` count="min" option is broken
and does construct proper minimal sets of closure quantities if the array isn't fully connected.

"""
function extract_cphase(obs; count="min-correct", cut_trivial=false, uvmin=0.1e9,  kwargs...)
    if count == "min-correct"
        return minimal_cphase(obs; cut_trivial, uvmin, kwargs...)
    else
        return _ehtim_cphase(obs; count=count, cut_trivial=cut_trivial, uvmin=uvmin, kwargs...)
    end
end

"""
    extract_lcamp(obs)
Extracts the log-closure amp. from an ehtim observation object

Any valid keyword arguments to `add_logcamp` in ehtim can be passed through extract_lcamp.

Returns an EHTObservation with log-closure amp. datums

# Warning
The `count` keyword argument is treated specially in `Comrade`. The default option
is "min-correct" and should almost always be used.
This option construct a minimal set of closure phases that is valid even when
the array isn't fully connected. For testing and legacy reasons we `ehtim` other count
options are also included. However, the current `ehtim` count="min" option is broken
and does construct proper minimal sets of closure quantities if the array isn't fully connected.

"""
function extract_lcamp(obs; count="min-correct", kwargs...)
    if count == "min-correct"
        return minimal_lcamp(obs; kwargs...)
    else
        return _ehtim_lcamp(obs; count=count, kwargs...)
    end
end


export read_non_aligned_sequence, write_mapped_contacts, score_BM, PPV

##############################################################
"""
    function read_non_aligned_sequence(seq::String)

    (seq[might be a path]) --> sequence

    USE:
    Given the path of a non aligned sequence (with '.' and insertions)
    returns the sequence. Takes as input .fasta format, or sequence itself,
    as well as the sequence without the path.

 	INPUT:
    `seq`: amino acid sequence or path to it

    OUTPUT:
    String

"""

function read_non_aligned_sequence(seq::String)
    isfasta(file_path::String) = readdlm(file_path)[1][1] == '>'
    seq =   if ispath(seq)
                if isfasta(seq)
                    readdlm(seq)[2]
                else
                    readdlm(seq)
                end

            else
                seq
            end
    return seq
end

##############################################################
"""
    function afa2model(seq::String)

    (seq) --> dictionary

    USE:
    Given a .afa non aligned sequence coming from a HMM
    returns a dictionary that maps its positions to the matched states (model).

    Example: yAC returns the following dictionary

    position       -->  position
    in the model   -->  in the original sequence

    2              -->  1
    3              -->  2

 	INPUT:
    - `seq`: amino acid sequence with insertions

    OUTPUT:
    Dictionary

"""

function afa2model(seq::String)
    seq = read_non_aligned_sequence(seq)
    p_new =  Dict{Int64,Int64}()
    k_lower_and_upper = 0
    k_gap_and_upper = 0
    for lettr in seq
        if islowercase(lettr)
            k_lower_and_upper += 1
        elseif lettr == '-'
            k_gap_and_upper +=1
        elseif isuppercase(lettr)
            k_gap_and_upper +=1
            k_lower_and_upper += 1
            p_new[k_lower_and_upper] = k_gap_and_upper
        end
    end

    if length(p_new) != length([lettr for lettr in seq if isuppercase(lettr)])
        error("Error: there is some inconsistency in the dictionary, verify input.")
    end
    return p_new
end

##############################################################
"""
    function write_mapped_contacts(afa::String, old_distances_path:String; path_out::String = "")
    function write_mapped_contacts(afa::String, old_distances::Array{Float64, 2}; path_out::String = "")

    (afa, old_distances_path, path_out) --> Nothing

    USE:
    Given a .afa sequence with insertions aligned to the distances file
    writes a new distance file with mapped position numbers.

 	INPUT:
    - `afa`: amino acid sequence with insertions

    - `old_distance_path`: path of amino acid distances from crystal data
    - `old_distances`: array with pairs of positions and residue distances

    - `path_out`: output path, if not declared will be the same of "old_distances_path"

    OUTPUT:
    Nothing

"""

function write_mapped_contacts(afa::String, old_distances_path::String; path_out::String = "")
    println("MAKE SURE THAT THE .afa FILE AND THE DISTANCE FILE REFER TO THE SAME SEQUENCE OR ARE ALIGNED.\n")

    afa_sequence = read_non_aligned_sequence(afa)
    l_afa = length(afa_sequence)
    l_gaps = length([lettr for lettr in afa_sequence if lettr == '-'])
    l_upper = length([lettr for lettr in afa_sequence if isuppercase(lettr)])
    full_length = l_afa - l_gaps

    println(
"The length of the full sequence is          $(full_length).
The length of the model-aligned sequence is $(l_upper).
The gaps in the model-aligned sequence are  $(l_gaps).")

    pdb2wt = afa2model(afa)
    old_distances = readdlm(old_distances_path)

    # count how many pairs with distances there will be in the new file a cool julia oneliner
    n_distances = length([i for i in 1:size(old_distances, 1) if haskey(pdb2wt, old_distances[i, 1]) && haskey(pdb2wt, old_distances[i, 2]) ])

    # new file for amino acid distances
    new_distances = zeros(Float64, n_distances, 3)

    # fill new matrix with positions mapped to our models
    l = 0
    for i in 1:size(old_distances, 1)
        if haskey(pdb2wt, old_distances[i, 1]) && haskey(pdb2wt, old_distances[i, 2])
            l +=1
            new_distances[l, :] = [pdb2wt[old_distances[i, 1]], pdb2wt[old_distances[i, 2]] , old_distances[i, 3]]
        end
    end

    path_out =  path_out == "" ? join(split(old_distances_file, '.')[1:end-1], '.')*"_mapped_model.txt" : path_out
    writedlm(path_out, new_distances)
end

function write_mapped_contacts(afa::String, old_distances::Array{Float64, 2}; path_out::String = "")
    println("MAKE SURE THAT THE .afa FILE AND THE DISTANCE FILE REFER TO THE SAME SEQUENCE OR ARE ALIGNED.\n")

    afa_sequence = read_non_aligned_sequence(afa)
    l_afa = length(afa_sequence)
    l_gaps = length([lettr for lettr in afa_sequence if lettr == '-'])
    l_upper = length([lettr for lettr in afa_sequence if isuppercase(lettr)])
    full_length = l_afa - l_gaps

    println(
"The length of the full sequence is          $(full_length).
The length of the model-aligned sequence is $(l_upper).
The gaps in the model-aligned sequence are  $(l_gaps).")

    pdb2wt = afa2model(afa)

    # count how many pairs with distances there will be in the new file a cool julia oneliner
    n_distances = length([i for i in 1:size(old_distances, 1) if haskey(pdb2wt, old_distances[i, 1]) && haskey(pdb2wt, old_distances[i, 2]) ])

    # new file for amino acid distances
    new_distances = zeros(Float64, n_distances, 3)

    # fill new matrix with positions mapped to our models
    l = 0
    for i in 1:size(old_distances, 1)
        if haskey(pdb2wt, old_distances[i, 1]) && haskey(pdb2wt, old_distances[i, 2])
            l +=1
            new_distances[l, :] = [pdb2wt[old_distances[i, 1]], pdb2wt[old_distances[i, 2]] , old_distances[i, 3]]
        end
    end

    path_out =  path_out == "" ? join(split(old_distances_file, '.')[1:end-1], '.')*"_mapped_model.txt" : path_out
    writedlm(path_out, new_distances)
end
;

##############################################################
"""
    function PPV(scores::Array{Float64,2}, distances::Array{Float64,2}; amino_dist = 4, atom_dist = 8)

    (scores, distances) --> PPV

    USE:
    Given a list of residue pairs with a score (arbitrarly defined) and a list of experimental residue distances
    returns the fraction of pairs - out of the highest scoring ones - that are considered "contacts".
    The definition of contacts depends on linear amino acid distance and physical distance between residues.

    INPUT:
    - `scores`: list of protein positions with a score
    - `distances`: list of protein positions with physical distance
    - `amino_dist`: threshold in terms of linear amino acid distance
    - `atom_dist`: threshold in terms of physical distance

    OUTPUT:
    PPV vector

"""

function PPV(scores::Array{Float64,2}, distances::Array{Float64,2}; amino_dist = 4, atom_dist = 8)
    if size(scores[1,:],1)!=3 || size(distances[1,:],1)!=3
        error("contactprediction.jl - PPV: `scores` and `distances` should be in the format `i j value`.")
    end

    scores = sortslices(scores, dims=1, rev=true, by=x->x[3])
    d = Dict{Tuple{Float64,Float64},Float64}()
    for l in 1:size(distances,1)
        d[(distances[l,1], distances[l,2])] = distances[l,3]
    end
    tp = 0
    np = 0
    TP = []
    cc=0
    for r in 1:size(scores,1)
        if abs(scores[r,1] - scores[r,2]) > amino_dist && get(d, (scores[r,1], scores[r,2]), -1.)>=0.
            np +=1
            if d[(scores[r,1], scores[r,2])]<atom_dist
                tp+=1
            else
                # cc<10?println("$(scores[r,1]), $(scores[r,2])"):Void
                cc+=1
            end
            push!(TP,tp/np)
        end
    end

    return TP
end

##############################################################

"""
    score_BM(J::Array{Float64,2}; APC::Bool = true, gap::Bool = false)
    score_BM(path_param::AbstractString; APC::Bool = true, gap::Bool = false)

    (J) --> scores
    (path_param) --> scores

    USE:
    Compute Frobenius norm of `q x q` blocks in matrix `J` and adds APC correction.
    This is to use the DCA model as a contact prediction predictor.

    INPUT:
    - `J`: coupling matrix in `q x q x L x L` format
    - `path_param`: path of fields and couplings parameters
    - `APC`: apply the famous APC correction. Default to `true`.
    - `gap`: Remove the state `21` from the Frobenius norm. Default to `false`.

    OUTPUT:
    scores

"""

#-----------------------------------------------

function score_BM(path_param::AbstractString; APC::Bool = true, gap::Bool = false)
    h, J = extract_params(path_param)
    q = size(J, 1)
    L = size(J, 4)

    println("Size of input matrix is ",size(J))

    S = zeros(Float64, L,L)
    for i in 1:L-1
        for j in i+1 : L
            S[i,j] = gap ? sqrt(sum(J[:, : ,i, j].^2)) : sqrt(sum(J[1:20, 1:20, i, j].^2))
            S[j,i] = S[i,j]
        end
    end

    F = zeros(Float64, convert(Int64, L*(L-1)/2),3)
    pos = 1
    for i in 1:L
        for j in i+1:L
            F[pos,:] = [i,j,S[i,j] - ( sum(S[:,j])/L * sum( S[:,i])/L ) / (sum(S)/(L^2)) * APC]
            pos += 1
        end
    end

    return F
end

function score_BM(J::Array{Float64,4}; APC::Bool = true, gap::Bool = false)
    q = size(J, 1)
    L = size(J, 4)

    println("Size of input matrix is ",size(J))

    S = zeros(Float64, L,L)
    for i in 1:L-1
        for j in i+1 : L
            S[i,j] = gap ? sqrt(sum(J[:, : ,i, j].^2)) : sqrt(sum(J[1:20, 1:20, i, j].^2))
            S[j,i] = S[i,j]
        end
    end

    F = zeros(Float64, convert(Int64, L*(L-1)/2),3)
    pos = 1
    for i in 1:L
        for j in i+1:L
            F[pos,:] = [i,j,S[i,j] - ( sum(S[:,j])/L * sum( S[:,i])/L ) / (sum(S)/(L^2)) * APC]
            pos += 1
        end
    end

    return F
end


"""
    Calculators

This module exists to bridge the gap between the `Models` and the `Dynamics`.

Here we provide functions and types for evaluating and storing quantities obtained from the
`Models`.
In addition any further manipulation of those quantities, such as computing eigenvalues,
is included here.

This module is largely needed to facilitate seemless integration of both ring polymer and
classical dynamics, using the same models and functions for both.
Specific ring polymer types are provided that have the extra fields and methods needed
to evaluate the quantities for each bead.
"""
module Calculators

using LinearAlgebra: LinearAlgebra, Hermitian, I, Eigen, tr
using StaticArrays: SMatrix, SVector

using NQCModels: NQCModels, Model, nstates
using NQCModels.AdiabaticModels: AdiabaticModel
using NQCModels.DiabaticModels: DiabaticModel, DiabaticFrictionModel
using NQCModels.FrictionModels: AdiabaticFrictionModel

using NQCDynamics: RingPolymers, ndofs

"""
    AbstractCalculator{M<:Model}

Top-level type for all calculators.

Each concrete calculator contains the `Model` and the fields to store the quantities
obtained from the model.
"""
abstract type AbstractCalculator{T,M<:Model} end
abstract type AbstractAdiabaticCalculator{T,M<:AdiabaticModel} <: AbstractCalculator{T,M} end
abstract type AbstractDiabaticCalculator{T,M<:Union{DiabaticFrictionModel,DiabaticModel}} <: AbstractCalculator{T,M} end
abstract type AbstractStaticDiabaticCalculator{T,M} <: AbstractDiabaticCalculator{T,M} end
abstract type AbstractFrictionCalculator{T,M<:AdiabaticFrictionModel} <: AbstractCalculator{T,M} end

Base.eltype(::AbstractCalculator{T}) where {T} = T

mutable struct AdiabaticCalculator{T,M} <: AbstractAdiabaticCalculator{T,M}
    model::M
    potential::T
    derivative::Matrix{T}
    function AdiabaticCalculator{T}(model::M, atoms::Integer) where {T,M<:Model}
        new{T,M}(model, 0, zeros(ndofs(model), atoms))
    end
end

struct RingPolymerAdiabaticCalculator{T,M} <: AbstractAdiabaticCalculator{T,M}
    model::M
    potential::Vector{T}
    derivative::Array{T,3}
    function RingPolymerAdiabaticCalculator{T}(model::M, atoms::Integer, beads::Integer) where {T,M<:Model}
        new{T,M}(model, zeros(beads), zeros(ndofs(model), atoms, beads))
    end
end

mutable struct DiabaticCalculator{T,M,S,L} <: AbstractDiabaticCalculator{T,M}
    model::M
    potential::Hermitian{T,SMatrix{S,S,T,L}}
    derivative::Matrix{Hermitian{T,SMatrix{S,S,T,L}}}
    eigen::LinearAlgebra.Eigen{T,T,SMatrix{S,S,T,L},SVector{S,T}}
    adiabatic_derivative::Matrix{SMatrix{S,S,T,L}}
    nonadiabatic_coupling::Matrix{SMatrix{S,S,T,L}}
    tmp_position::Matrix{T}
    tmp_mat::Matrix{T}
    tmp_mat_complex1::Matrix{Complex{T}}
    tmp_mat_complex2::Matrix{Complex{T}}
    function DiabaticCalculator{T}(model::M, atoms::Integer) where {T,M<:Model}
        n = nstates(model)
        matrix_template = NQCModels.DiabaticModels.matrix_template(model, T)
        vector_template = NQCModels.DiabaticModels.vector_template(model, T)

        potential = Hermitian(matrix_template)
        derivative = [Hermitian(matrix_template) for _=1:ndofs(model), _=1:atoms]
        eigen = Eigen(vector_template, matrix_template + I)
        adiabatic_derivative = [matrix_template for _ in CartesianIndices(derivative)]
        nonadiabatic_coupling = [matrix_template for _ in CartesianIndices(derivative)]
        tmp_position = zeros(T, ndofs(model), atoms)
        tmp_mat = zeros(T, n, n)
        tmp_mat_complex1 = zeros(Complex{T}, n, n)
        tmp_mat_complex2 = zeros(Complex{T}, n, n)

        new{T,M,n,n^2}(model,
            potential, derivative, eigen,
            adiabatic_derivative, nonadiabatic_coupling,
            tmp_position, tmp_mat, tmp_mat_complex1, tmp_mat_complex2)
    end
end

mutable struct RingPolymerDiabaticCalculator{T,M,S,L} <: AbstractDiabaticCalculator{T,M}
    model::M
    potential::Vector{Hermitian{T,SMatrix{S,S,T,L}}}
    traceless_potential::Vector{Hermitian{T,SMatrix{S,S,T,L}}}
    V̄::Vector{T}
    derivative::Array{Hermitian{T,SMatrix{S,S,T,L}},3}
    traceless_derivative::Array{Hermitian{T,SMatrix{S,S,T,L}},3}
    D̄::Array{T,3}
    eigen::Vector{LinearAlgebra.Eigen{T,T,SMatrix{S,S,T,L},SVector{S,T}}}
    adiabatic_derivative::Array{SMatrix{S,S,T,L},3}
    traceless_adiabatic_derivative::Array{SMatrix{S,S,T,L},3}
    nonadiabatic_coupling::Array{SMatrix{S,S,T,L},3}

    centroid_potential::Hermitian{T,SMatrix{S,S,T,L}}
    centroid_derivative::Matrix{Hermitian{T,SMatrix{S,S,T,L}}}
    centroid_eigen::LinearAlgebra.Eigen{T,T,SMatrix{S,S,T,L},SVector{S,T}}
    centroid_adiabatic_derivative::Matrix{SMatrix{S,S,T,L}}
    centroid_nonadiabatic_coupling::Matrix{SMatrix{S,S,T,L}}
    tmp_centroid::Matrix{T}

    tmp_mat::Matrix{T}
    tmp_mat_complex1::Matrix{Complex{T}}
    tmp_mat_complex2::Matrix{Complex{T}}
    function RingPolymerDiabaticCalculator{T}(model::M, atoms::Integer, beads::Integer) where {T,M<:Model}
        n = nstates(model)
        matrix_template = NQCModels.DiabaticModels.matrix_template(model, T)
        vector_template = NQCModels.DiabaticModels.vector_template(model, T)

        potential = [Hermitian(matrix_template) for _=1:beads]
        traceless_potential = [Hermitian(matrix_template) for _=1:beads]
        V̄ = zeros(beads)

        derivative = [Hermitian(matrix_template) for _=1:ndofs(model), _=1:atoms, _=1:beads]
        traceless_derivative = [Hermitian(matrix_template) for _=1:ndofs(model), _=1:atoms, _=1:beads]
        D̄ = zeros(ndofs(model), atoms, beads)

        adiabatic_derivative = [matrix_template for _=1:ndofs(model), _=1:atoms, _=1:beads]
        traceless_adiabatic_derivative = [Hermitian(matrix_template) for _=1:ndofs(model), _=1:atoms, _=1:beads]

        eigen = [Eigen(vector_template, matrix_template + I) for _=1:beads]
        nonadiabatic_coupling = [matrix_template for _=1:ndofs(model), _=1:atoms, _=1:beads]

        centroid_potential = Hermitian(matrix_template)
        centroid_derivative = [Hermitian(matrix_template) for _=1:ndofs(model), _=1:atoms]
        centroid_eigen = Eigen(vector_template, matrix_template + I)
        centroid_adiabatic_derivative = [matrix_template for _ in CartesianIndices(centroid_derivative)]
        centroid_nonadiabatic_coupling = [matrix_template for _ in CartesianIndices(centroid_derivative)]
        tmp_centroid = zeros(T, ndofs(model), atoms)

        tmp_mat = zeros(T, n, n)
        tmp_mat_complex1 = zeros(Complex{T}, n, n)
        tmp_mat_complex2 = zeros(Complex{T}, n, n)
        new{T,M,n,n^2}(model,
            potential, traceless_potential, V̄,
            derivative, traceless_derivative, D̄,
            eigen, adiabatic_derivative, traceless_adiabatic_derivative, nonadiabatic_coupling,
            centroid_potential, centroid_derivative, centroid_eigen,
            centroid_adiabatic_derivative, centroid_nonadiabatic_coupling, tmp_centroid,
            tmp_mat, tmp_mat_complex1, tmp_mat_complex2)
    end
end

function Calculator(model::DiabaticModel, atoms::Integer, t::Type{T}) where {T}
    DiabaticCalculator{t}(model, atoms)
end
function Calculator(model::AdiabaticModel, atoms::Integer, t::Type{T}) where {T}
    AdiabaticCalculator{t}(model, atoms)
end
function Calculator(model::DiabaticModel, atoms::Integer, beads::Integer, t::Type{T}) where {T}
    RingPolymerDiabaticCalculator{t}(model, atoms, beads)
end
function Calculator(model::AdiabaticModel, atoms::Integer, beads::Integer, t::Type{T}) where {T}
    RingPolymerAdiabaticCalculator{t}(model, atoms, beads)
end

function evaluate_potential!(calc::AbstractCalculator, R)
    calc.potential = NQCModels.potential(calc.model, R)
end

function evaluate_potential!(calc::AbstractCalculator, R::AbstractArray{T,3}) where {T}
    @views for i in axes(R, 3)
        calc.potential[i] = NQCModels.potential(calc.model, R[:,:,i])
    end
end

function evaluate_V̄!(calc::RingPolymerDiabaticCalculator)
    for i in 1:length(calc.V̄)
        calc.V̄[i] = tr(calc.potential[i]) / nstates(calc.model)
    end
end

function evaluate_traceless_potential!(calc::RingPolymerDiabaticCalculator)
    n = nstates(calc.model)
    for I in 1:length(calc.traceless_potential)
        calc.traceless_potential[I] = Hermitian(SMatrix{n,n}(
            i != j ? calc.potential[I][j,i] : calc.potential[I][j,i] - calc.V̄[I] for j=1:n, i=1:n
        ))
    end
end

function evaluate_centroid_potential!(calc::AbstractCalculator, R::AbstractMatrix)
    calc.centroid_potential = NQCModels.potential(calc.model, R)
end

function evaluate_derivative!(calc::AbstractCalculator, R)
    NQCModels.derivative!(calc.model, calc.derivative, R)
end

function evaluate_derivative!(calc::AbstractCalculator, R::AbstractArray{T,3}) where {T}
    @views for i in axes(R, 3)
        NQCModels.derivative!(calc.model, calc.derivative[:,:,i], R[:,:,i])
    end
end

function evaluate_D̄!(calc::RingPolymerDiabaticCalculator)
    for I in eachindex(calc.D̄)
        calc.D̄[I] = tr(calc.derivative[I]) / nstates(calc.model)
    end
end

function evaluate_traceless_derivative!(calc::RingPolymerDiabaticCalculator)
    n = nstates(calc.model)
    for I in eachindex(calc.derivative)
        calc.traceless_derivative[I] = Hermitian(SMatrix{n,n}(
            i != j ? calc.derivative[I][j,i] : calc.derivative[I][j,i] - calc.D̄[I] for j=1:n, i=1:n
        ))
    end
end

function evaluate_traceless_adiabatic_derivative!(calc::RingPolymerDiabaticCalculator)
    n = nstates(calc.model)
    for I in eachindex(calc.adiabatic_derivative)
        calc.traceless_adiabatic_derivative[I] = SMatrix{n,n}(
            i != j ? calc.adiabatic_derivative[I][j,i] : calc.adiabatic_derivative[I][j,i] - calc.D̄[I] for j=1:n, i=1:n
        )
    end
end

function evaluate_centroid_derivative!(calc::AbstractCalculator, R::AbstractMatrix)
    NQCModels.derivative!(calc.model, calc.centroid_derivative, R)
end

function eigen!(calc::DiabaticCalculator)
    eig = LinearAlgebra.eigen(calc.potential)
    corrected_vectors = correct_phase(eig.vectors, calc.eigen.vectors)
    calc.eigen = Eigen(eig.values, corrected_vectors)
    return nothing
end

function centroid_eigen!(calc::RingPolymerDiabaticCalculator)
    eig = LinearAlgebra.eigen(calc.centroid_potential)
    corrected_vectors = correct_phase(eig.vectors, calc.centroid_eigen.vectors)
    calc.centroid_eigen = Eigen(eig.values, corrected_vectors)
    return nothing
end

function eigen!(calc::RingPolymerDiabaticCalculator)
    for i=1:length(calc.potential)
        eig = LinearAlgebra.eigen(calc.potential[i])
        corrected_vectors = correct_phase(eig.vectors, calc.eigen[i].vectors)
        calc.eigen[i] = Eigen(eig.values, corrected_vectors)
    end
    return nothing
end

function correct_phase(new_vectors::SMatrix, old_vectors::SMatrix)
    n = size(new_vectors, 1)
    vect = SVector{n}(sign(LinearAlgebra.dot(new_vectors[:,i], old_vectors[:,i])) for i=1:n)
    return new_vectors .* vect'
end

function transform_derivative!(calc::AbstractDiabaticCalculator)
    for I in eachindex(calc.derivative)
        calc.adiabatic_derivative[I] = calc.eigen.vectors' * calc.derivative[I] * calc.eigen.vectors
    end
end

function transform_centroid_derivative!(calc::RingPolymerDiabaticCalculator)
    for I in eachindex(calc.centroid_derivative)
        calc.centroid_adiabatic_derivative[I] = calc.centroid_eigen.vectors' * calc.centroid_derivative[I] * calc.centroid_eigen.vectors
    end
end

function transform_derivative!(calc::RingPolymerDiabaticCalculator)
    for i in axes(calc.derivative, 3) # Beads
        for j in axes(calc.derivative, 2) # Atoms
            for k in axes(calc.derivative, 1) # DoFs
                calc.adiabatic_derivative[k,j,i] = calc.eigen[i].vectors' * calc.derivative[k,j,i] * calc.eigen[i].vectors
            end
        end
    end
end

function evaluate_nonadiabatic_coupling!(calc::AbstractDiabaticCalculator)
    for I in eachindex(calc.adiabatic_derivative)
        calc.nonadiabatic_coupling[I] = evaluate_nonadiabatic_coupling(calc.adiabatic_derivative[I], calc.eigen.values)
    end
end

function evaluate_centroid_nonadiabatic_coupling!(calc::RingPolymerDiabaticCalculator)
    for I in eachindex(calc.centroid_adiabatic_derivative)
        calc.centroid_nonadiabatic_coupling[I] = evaluate_nonadiabatic_coupling(calc.centroid_adiabatic_derivative[I], calc.centroid_eigen.values)
    end
end

function evaluate_nonadiabatic_coupling!(calc::RingPolymerDiabaticCalculator)
    for i in axes(calc.nonadiabatic_coupling, 3) # Beads
        for I in CartesianIndices(size(calc.adiabatic_derivative)[1:2])
            calc.nonadiabatic_coupling[I,i] = evaluate_nonadiabatic_coupling(calc.adiabatic_derivative[I,i], calc.eigen[i].values)
        end
    end
end

"""
# References

- HammesSchifferTully_JChemPhys_101_4657_1994 Eq. (32)
- SubotnikBellonzi_AnnuRevPhyschem_67_387_2016, section 2.3
"""
function evaluate_nonadiabatic_coupling(adiabatic_derivative::SMatrix, eigenvalues::SVector)
    n = length(eigenvalues)
    SMatrix{n,n}(
        (i != j ? adiabatic_derivative[j,i] / (eigenvalues[i] - eigenvalues[j]) : 0
        for j=1:n, i=1:n))
end

"""
Evaluates all electronic properties for the current position `r`.

# Properties evaluated:
- Diabatic potential
- Diabatic derivative
- Eigenvalues and eigenvectors
- Adiabatic derivative
- Nonadiabatic coupling
"""
function update_electronics!(calculator::AbstractDiabaticCalculator, r::AbstractArray)
    evaluate_potential!(calculator, r)
    evaluate_derivative!(calculator, r)
    eigen!(calculator)
    transform_derivative!(calculator)
    evaluate_nonadiabatic_coupling!(calculator)
end

function update_electronics!(calculator::RingPolymerDiabaticCalculator, r::AbstractArray{T,3}) where {T}
    evaluate_potential!(calculator, r)
    evaluate_derivative!(calculator, r)
    eigen!(calculator)
    transform_derivative!(calculator)
    evaluate_nonadiabatic_coupling!(calculator)

    update_centroid_electronics!(calculator, r)
end

function update_centroid_electronics!(calculator::RingPolymerDiabaticCalculator, r::AbstractArray{T,3}) where {T}
    RingPolymers.get_centroid!(calculator.tmp_centroid, r)
    evaluate_centroid_potential!(calculator, calculator.tmp_centroid)
    evaluate_centroid_derivative!(calculator, calculator.tmp_centroid)
    centroid_eigen!(calculator)
    transform_centroid_derivative!(calculator)
    evaluate_centroid_nonadiabatic_coupling!(calculator)
end

include("large_diabatic.jl")
include("friction.jl")

position(calc) = calc.tmp_position

for field ∈ (:potential, :derivative)
    get_field = Symbol(:get, :_, field)
    evaluate_field = Symbol(:evaluate, :_, field, :!)

    eval(quote
        function $get_field(calc, r)
            if position(calc) ≈ r
                return calc.$field
            else
                copyto!(position(calc), r)
                return $evaluate_field(calc, r)
            end
        end
    end)
end

end # module


Materials = (
    tissueplastic = Material(
        "A-150 Tissue-Equivalent Plastic",
        0.54903,
        65.1eV,
        1.127g / cm^3,
        Dict(
            1 => 0.10133,
            6 => 0.775498,
            7 => 0.035057,
            8 => 0.052315,
            9 => 0.017423,
            20 => 0.018377,
        ),
    ),
    adipose = Material(
        "Adipose Tissue (ICRU-44)",
        0.55579,
        64.8eV,
        0.95g / cm^3,
        Dict(
            1 => 0.114,
            6 => 0.598,
            7 => 0.007,
            8 => 0.278,
            11 => 0.001,
            16 => 0.001,
            17 => 0.001,
        ),
    ),
    air = Material(
        "Air, Dry (near sea level)",
        0.49919,
        85.7eV,
        0.001205g / cm^3,
        Dict(6 => 0.000124, 7 => 0.755268, 8 => 0.231781, 18 => 0.012827),
    ),
    alanine = Material(
        "Alanine",
        0.53876,
        71.9eV,
        1.424g / cm^3,
        Dict(1 => 0.079192, 6 => 0.404437, 7 => 0.157213, 8 => 0.359157),
    ),
    boneplastic = Material(
        "B-100 Bone-Equivalent Plastic",
        0.5274,
        85.9eV,
        1.45g / cm^3,
        Dict(
            1 => 0.065473,
            6 => 0.536942,
            7 => 0.0215,
            8 => 0.032084,
            9 => 0.167415,
            20 => 0.176585,
        ),
    ),
    bakelite = Material(
        "Bakelite",
        0.52792,
        72.4eV,
        1.25g / cm^3,
        Dict(1 => 0.057444, 6 => 0.774589, 8 => 0.167968),
    ),
    wholeblood = Material(
        "Blood, Whole (ICRU-44)",
        0.54999,
        75.2eV,
        1.06g / cm^3,
        Dict(
            1 => 0.102,
            6 => 0.11,
            7 => 0.033,
            8 => 0.745,
            11 => 0.001,
            15 => 0.001,
            16 => 0.002,
            17 => 0.003,
            19 => 0.002,
            26 => 0.001,
        ),
    ),
    corticalbone = Material(
        "Bone, Cortical (ICRU-44)",
        0.51478,
        112.0eV,
        1.92g / cm^3,
        Dict(
            1 => 0.034,
            6 => 0.155,
            7 => 0.042,
            8 => 0.435,
            11 => 0.001,
            12 => 0.002,
            15 => 0.103,
            16 => 0.003,
            20 => 0.225,
        ),
    ),
    brain = Material(
        "Brain, Grey/White Matter (ICRU-44)",
        0.55239,
        73.9eV,
        1.04g / cm^3,
        Dict(
            1 => 0.107,
            6 => 0.145,
            7 => 0.022,
            8 => 0.712,
            11 => 0.002,
            15 => 0.004,
            16 => 0.002,
            17 => 0.003,
            19 => 0.003,
        ),
    ),
    breasttissue = Material(
        "Breast Tissue (ICRU-44)",
        0.55196,
        70.3eV,
        1.02g / cm^3,
        Dict(
            1 => 0.106,
            6 => 0.332,
            7 => 0.03,
            8 => 0.527,
            11 => 0.001,
            15 => 0.001,
            16 => 0.002,
            17 => 0.001,
        ),
    ),
    airplastic = Material(
        "C-552 Air-equivalent Plastic",
        0.49969,
        86.8eV,
        1.76g / cm^3,
        Dict(
            1 => 0.024681,
            6 => 0.50161,
            8 => 0.004527,
            9 => 0.465209,
            14 => 0.003973,
        ),
    ),
    CdTe = Material(
        "Cadmium Telluride",
        0.41665,
        539.3eV,
        6.2g / cm^3,
        Dict(48 => 0.468358, 52 => 0.531642),
    ),
    CaF = Material(
        "Calcium Fluoride",
        0.48671,
        166.0eV,
        3.18g / cm^3,
        Dict(9 => 0.486672, 20 => 0.513328),
    ),
    CaSO4 = Material(
        "Calcium Sulfate",
        0.49948,
        152.3eV,
        2.96g / cm^3,
        Dict(8 => 0.470081, 16 => 0.235534, 20 => 0.294385),
    ),
    ammoniumsulfate = Material(
        "15 mmol L-1 Ceric Ammonium Sulfate Solution",
        0.55282,
        76.7eV,
        1.030E+00g / cm^3,
        Dict(
            1 => 0.107694,
            7 => 0.000816,
            8 => 0.875172,
            16 => 0.01427958,
            58 => 0.00204,
        ),
    ),
    CsI = Material(
        "Cesium Iodide",
        0.41569,
        553.1eV,
        4.51g / cm^3,
        Dict(53 => 0.488451, 55 => 0.511549),
    ),
    concrete = Material(
        "Concrete, Ordinary",
        0.50932,
        124.5eV,
        2.3g / cm^3,
        Dict(
            1 => 0.0221,
            6 => 0.002484,
            8 => 0.57493,
            11 => 0.015208,
            12 => 0.001266,
            13 => 0.019953,
            14 => 0.304627,
            19 => 0.010045,
            20 => 0.042951,
            26 => 0.006435,
        ),
    ),
    concretebarite = Material(
        "Concrete, Barite (TYPE BA)",
        0.45714,
        248.2eV,
        3.35g / cm^3,
        Dict(
            1 => 0.003585,
            8 => 0.311622,
            12 => 0.001195,
            13 => 0.004183,
            14 => 0.010457,
            16 => 0.107858,
            20 => 0.050194,
            26 => 0.047505,
            56 => 0.4634,
        ),
    ),
    eyelens = Material(
        "Eye Lens (ICRU-44)",
        0.54709,
        74.3eV,
        1.07g / cm^3,
        Dict(
            1 => 0.096,
            6 => 0.195,
            7 => 0.057,
            8 => 0.646,
            11 => 0.001,
            15 => 0.001,
            16 => 0.003,
            17 => 0.001,
        ),
    ),
    ferroussulfate = Material(
        "Ferrous Sulfate Standard Fricke",
        0.55334,
        76.3eV,
        1.024g / cm^3,
        Dict(
            1 => 0.108376,
            8 => 0.878959,
            11 => 2.2e-05,
            16 => 0.012553,
            17 => 3.5e-05,
            26 => 5.5e-05,
        ),
    ),
    GOS = Material(
        "Gadolinium Oxysulfide",
        0.42265,
        493.3eV,
        7.44g / cm^3,
        Dict(8 => 0.084527, 16 => 0.084704, 64 => 0.830769),
    ),
    gafchromicsensor = Material(
        "Gafchromic Sensor",
        0.54384,
        67.2eV,
        1.3g / cm^3,
        Dict(1 => 0.0897, 6 => 0.6058, 7 => 0.1122, 8 => 0.1923),
    ),
    GaAs = Material(
        "Gallium Arsenide",
        0.44246,
        384.9eV,
        5.31g / cm^3,
        Dict(31 => 0.48203, 33 => 0.51797),
    ),
    pyrex = Material(
        "Glass, Borosilicate (Pyrex)",
        0.49707,
        134.0eV,
        2.23g / cm^3,
        Dict(
            5 => 0.040066,
            8 => 0.539559,
            11 => 0.028191,
            13 => 0.011644,
            14 => 0.37722,
            19 => 0.003321,
        ),
    ),
    leadglass = Material(
        "Glass, Lead",
        0.42101,
        526.4eV,
        6.22g / cm^3,
        Dict(
            8 => 0.156453,
            14 => 0.080866,
            22 => 0.008092,
            33 => 0.002651,
            82 => 0.751938,
        ),
    ),
    LiF = Material(
        "Lithium Fluride",
        0.46262,
        94.0eV,
        2.635g / cm^3,
        Dict(3 => 0.267585, 9 => 0.732415),
    ),
    LiB4O7 = Material(
        "Lithium Tetraborate",
        0.48485,
        94.6eV,
        2.44g / cm^3,
        Dict(3 => 0.082081, 5 => 0.255715, 8 => 0.662204),
    ),
    lung = Material(
        "Lung Tissue (ICRU-44)",
        0.55048,
        75.2eV,
        1.05g / cm^3,
        Dict(
            1 => 0.103,
            6 => 0.105,
            7 => 0.031,
            8 => 0.749,
            11 => 0.002,
            15 => 0.002,
            16 => 0.003,
            17 => 0.003,
            19 => 0.002,
        ),
    ),
    MgB4O7 = Material(
        "Magnesium Tetroborate",
        0.49012,
        108.3eV,
        2.53g / cm^3,
        Dict(5 => 0.24087, 8 => 0.623762, 12 => 0.135367),
    ),
    MgI2 = Material(
        "Mercuric Iodide",
        0.40933,
        684.5eV,
        6.36g / cm^3,
        Dict(53 => 0.55856, 80 => 0.44144),
    ),
    muscle = Material(
        "Muscle, Skeletal (ICRU-44)",
        0.55,
        74.6eV,
        1.05g / cm^3,
        Dict(
            1 => 0.102,
            6 => 0.143,
            7 => 0.034,
            8 => 0.71,
            11 => 0.001,
            15 => 0.002,
            16 => 0.003,
            17 => 0.001,
            19 => 0.004,
        ),
    ),
    ovary = Material(
        "Ovary (ICRU-44)",
        0.55149,
        75.0eV,
        1.05g / cm^3,
        Dict(
            1 => 0.105,
            6 => 0.093,
            7 => 0.024,
            8 => 0.768,
            11 => 0.002,
            15 => 0.002,
            16 => 0.002,
            17 => 0.002,
            19 => 0.002,
        ),
    ),
    emulsionkodak = Material(
        "Photographic Emulsion (Kodak Type AA)",
        0.48176,
        179.0eV,
        2.2g / cm^3,
        Dict(
            1 => 0.0305,
            6 => 0.2107,
            7 => 0.0721,
            8 => 0.1632,
            35 => 0.2228,
            47 => 0.3007,
        ),
    ),
    emulsionstandard = Material(
        "Photographic Emulsion (Standard Nuclear)",
        0.45453,
        331.0eV,
        3.815g / cm^3,
        Dict(
            1 => 0.0141,
            6 => 0.072261,
            7 => 0.01932,
            8 => 0.066101,
            16 => 0.00189,
            35 => 0.349104,
            47 => 0.474105,
            53 => 0.00312,
        ),
    ),
    vinyltoluene = Material(
        "Plastic Scintillator, Vinyltoluene",
        0.54141,
        64.7eV,
        1.032g / cm^3,
        Dict(1 => 0.085, 6 => 0.915),
    ),
    polyethylene = Material(
        "Polyethylene",
        0.57033,
        57.4eV,
        0.93g / cm^3,
        Dict(1 => 0.143716, 6 => 0.856284),
    ),
    mylar = Material(
        "Polyethylene Terephthalate, (Mylar)",
        0.52037,
        78.7eV,
        1.38g / cm^3,
        Dict(1 => 0.04196, 6 => 0.625016, 8 => 0.333024),
    ),
    PMMA = Material(
        "Polymethyl Methacrylate",
        0.53937,
        74.0eV,
        1.19g / cm^3,
        Dict(1 => 0.080541, 6 => 0.599846, 8 => 0.319613),
    ),
    polystyrene = Material(
        "Polystyrene",
        0.53768,
        68.7eV,
        1.06g / cm^3,
        Dict(1 => 0.077421, 6 => 0.922579),
    ),
    teflon = Material(
        "Polytetrafluoroethylene, (Teflon)",
        0.47993,
        99.1eV,
        2.25g / cm^3,
        Dict(6 => 0.240183, 9 => 0.759818),
    ),
    PVC = Material(
        "Polyvinyl Chloride",
        0.51201,
        108.2eV,
        1.406g / cm^3,
        Dict(1 => 0.048382, 6 => 0.384361, 17 => 0.567257),
    ),
    nylonfilm = Material(
        "Radiochromic Dye Film, Nylon Base",
        0.54987,
        64.5eV,
        1.08g / cm^3,
        Dict(1 => 0.101996, 6 => 0.654396, 7 => 0.098915, 8 => 0.144693),
    ),
    testis = Material(
        "Testis (ICRU-44)",
        0.552,
        74.7eV,
        1.04g / cm^3,
        Dict(
            1 => 0.106,
            6 => 0.099,
            7 => 0.02,
            8 => 0.766,
            11 => 0.002,
            15 => 0.001,
            16 => 0.002,
            17 => 0.002,
            19 => 0.002,
        ),
    ),
    softtissue = Material(
        "Tissue, Soft (ICRU-44)",
        0.54996,
        74.7eV,
        1.06g / cm^3,
        Dict(
            1 => 0.102,
            6 => 0.143,
            7 => 0.034,
            8 => 0.708,
            11 => 0.002,
            15 => 0.003,
            16 => 0.003,
            17 => 0.002,
            19 => 0.003,
        ),
    ),
    softtissue4 = Material(
        "Tissue, Soft (ICRU Four-Component)",
        0.54975,
        74.9eV,
        1.0g / cm^3,
        Dict(1 => 0.101174, 6 => 0.111, 7 => 0.026, 8 => 0.761826),
    ),
    methane = Material(
        "Tissue-Equivalent Gas, Methane Based",
        0.54992,
        61.2eV,
        0.001064g / cm^3,
        Dict(1 => 0.101873, 6 => 0.456177, 7 => 0.035172, 8 => 0.406778),
    ),
    propane = Material(
        "Tissue-Equivalent Gas, Propane Based",
        0.55027,
        59.5eV,
        0.001826g / cm^3,
        Dict(1 => 0.102676, 6 => 0.568937, 7 => 0.035022, 8 => 0.293365),
    ),
    water = Material(
        "Water, Liquid",
        0.55508,
        75.0eV,
        1.0g / cm^3,
        Dict(1 => 0.111898, 8 => 0.888102),
    ),
)


abstract type Wishart3x3{T} <: AbstractMatrix{T} end

struct InverseWishart{T} <: Wishart3x3{T}
    data::NTuple{8,VE{T}}
end
struct CholInvWishart{T} <: Wishart3x3{T}
    data::NTuple{8,T}
end
struct RevCholWishart{T} <: Wishart3x3{T}
    data::NTuple{8,T}
end

const WishartFactor{T} = Union{CholInvWishart{T},RevCholWishart{T}}

@inline InverseWishart(x::Vararg{T,8}   ) where T = InverseWishart{T}(VE.(x))
@inline InverseWishart{T}(x::Vararg{T,8}) where T = InverseWishart{T}(VE.(x))
# @generated function CholInvWishart(x::Vararg{T,8}) where T
#     quote
#         $(Expr(:meta,:inline))
#         CholInvWishart{$T}($(Expr(:tuple,[:(@inbounds x[$i]) for i ∈ 1:8]...)))
#     end
# end

@inline function LinearAlgebra.logdet(l::Union{CholInvWishart{T},RevCholWishart{T}}) where {T}
    @inbounds SLEEFPirates.log(l[1] * l[4] * l[6])
end
@inline function logdetw(l::Union{CholInvWishart{T},RevCholWishart{T}}) where {T}
    @inbounds SLEEFwrap.log(l[1] * l[4] * l[6])
end

@inline function CholInvWishart{T}(x_1::T,x_2::T,x_3::T,x_4::T,x_5::T,x_6::T) where T
    x_7 = x_8 = one(T)
    CholInvWishart{T}((Base.Cartesian.@ntuple 8 x))
end
@inline function CholInvWishart{T}(x_1::T,x_2::T,x_3::T,x_4::T,x_5::T,x_6::T,x_7::T,x_8::T) where T
    CholInvWishart{T}((Base.Cartesian.@ntuple 8 x))
end
@inline function CholInvWishart(x_1::T,x_2::T,x_3::T,x_4::T,x_5::T,x_6::T) where T
    x_7 = x_8 = one(T)
    CholInvWishart{T}((Base.Cartesian.@ntuple 8 x))
end
@inline function CholInvWishart(x_1::T,x_2::T,x_3::T,x_4::T,x_5::T,x_6::T,x_7::T,x_8::T) where T
    CholInvWishart{T}((Base.Cartesian.@ntuple 8 x))
end
@inline function RevCholWishart{T}(x_1::T,x_2::T,x_3::T,x_4::T,x_5::T,x_6::T) where T
    x_7 = x_8 = one(T)
    RevCholWishart{T}((Base.Cartesian.@ntuple 8 x))
end
@inline function RevCholWishart{T}(x_1::T,x_2::T,x_3::T,x_4::T,x_5::T,x_6::T,x_7::T,x_8::T) where T
    RevCholWishart{T}((Base.Cartesian.@ntuple 8 x))
end
@inline function RevCholWishart{T1}(
            x_1::SVec{W,T},x_2::SVec{W,T},x_3::SVec{W,T},x_4::SVec{W,T},x_5::SVec{W,T},x_6::SVec{W,T},x_7::SVec{W,T},x_8::SVec{W,T}
        ) where {W,T,T1}
    RevCholWishart{SVec{W,T}}((Base.Cartesian.@ntuple 8 x))
end
@inline function RevCholWishart(x_1::T,x_2::T,x_3::T,x_4::T,x_5::T,x_6::T) where T
    x_7 = x_8 = one(T)
    RevCholWishart{T}((Base.Cartesian.@ntuple 8 x))
end
@inline function RevCholWishart(x_1::T,x_2::T,x_3::T,x_4::T,x_5::T,x_6::T,x_7::T,x_8::T) where T
    RevCholWishart{T}((Base.Cartesian.@ntuple 8 x))
end
# @inline RevCholWishart{T}(x::Vararg{T,8}) where T = RevCholWishart{T}(x)
# @inline RevCholWishart(x::Vararg{T,6}   ) where T = RevCholWishart{T}(x,T(1),T(1))
# @inline RevCholWishart{T}(x::Vararg{T,6}) where T = RevCholWishart{T}(x,T(1),T(1))
@inline CholInvWishart{T}(x::Vararg{SVec{W,T},8}) where {W,T} = CholInvWishart{SVec{W,T}}(x)

inv_type(::Type{CholInvWishart{T}}) where T = RevCholWishart{T}
inv_type(::Type{RevCholWishart{T}}) where T = CholInvWishart{T}
# @inline function SArray{Tuple{S},T,N,L}(vs::Vararg{SIMDPirates.SVec{W,T},L}) where {S,T,N,L,W}
#     SArray{S}(vs...)
# end
# @inline Base.size(::Wishart2x2) = (2,2)
# @inline Base.length(::Wishart2x2) = 4
@inline Base.size(::Wishart3x3) = (3,3)
@inline Base.length(::Wishart3x3) = 9
# ScatteredArrays.type_length(::Type{Wishart2x2}) = 3
# ScatteredArrays.type_length(::Type{<:Wishart3x3}) = 6
# ScatteredArrays.type_length(::Type{CholInvWishart{T}}) where T = 6

@inline extractval(x::Core.VecElement{T}) where T = x.value
@inline extractval(x) = x
@inline Base.getindex(w::Wishart3x3, i::Integer) = extractval(@inbounds w.data[i])

@inline Base.getindex(A::Wishart3x3, ::LinearStorage, i::Integer) = @inbounds A.data[i]

@inline function Base.getindex(w::Wishart3x3, i::CartesianIndex{2})
    w[i[1], i[2]]
end
@inline function Base.getindex(iw::InverseWishart, i::Integer, j::Integer)
    i, j = minmax(i, j)
    @boundscheck (j > 3 || i < 1) && throwboundserror()
    im1 = i - 1
    @inbounds iw.data[3im1 - ( (im1*i) >> 1) + j].value
end
@inline function Base.getindex(w::WishartFactor{T}, j::Integer, i::Integer) where T
    j < i && return zero(T)
    @boundscheck (j > 3 || i < 1) && throwboundserror()
    im1 = i - 1
    @inbounds extractval(w.data[3im1 - ( (im1*i) >> 1) + j])
end
function Base.show(io::IO, ::MIME"text/plain", x::InverseWishart{T}) where T
    denom = x[7] > 2 ? x[7]-2 : one(T)
    println(io, "InverseWishart{$T} with ν = $(x[7]+2)")
    Base.print_matrix(io, x / denom)
end
function Base.show(io::IO, ::MIME"text/plain", x::CholInvWishart{T}) where T
    denom = x[7] > 2 ? sqrt(x[7]-2) : one(T)
    println(io, "Cholesky factor of InverseWishart{$T} with ν = $(x[7]+2)")
    Base.print_matrix(io, x / denom)
end

@inline function InverseWishart(x::T, y::T, z::T) where T
    InverseWishart(x * x, x * y, x * z, y * y, y * z, z * z, one(T), one(T))
end

@inline function Base.:+(a::W, b::W) where W <: Wishart3x3
    W(SIMDPirates.vadd(a.data, b.data))
end
@inline function Base.:*(L::WishartFactor{T}, x::SVector{3,T}) where {T}
    @fastmath @inbounds SVector(
        L[1]*x[1],
        L[2]*x[1] + L[4]*x[2],
        L[3]*x[1] + L[5]*x[2] + L[6]*x[3]
    )
end
@inline function Base.:*(L::WishartFactor{SVec{W,T}}, x::SVector{3,T}) where {W,T}
    @fastmath @inbounds SVector(
        L[1]*x[1],
        L[2]*x[1] + L[4]*x[2],
        L[3]*x[1] + L[5]*x[2] + L[6]*x[3]
    )
end
# L2 is treated as a lower triangular matrix
@inline function Base.:*(L1::W, L2::Union{W,NTuple{6,T}}) where {T, W <: WishartFactor{T}}
    @fastmath @inbounds W(
        L1[1]*L2[1],
        L1[2]*L2[1] + L1[4]*L2[2],
        L1[3]*L2[1] + L1[5]*L2[2] + L1[6]*L2[3],
        L1[4]*L2[4],
        L1[5]*L2[4] + L1[6]*L2[5],
        L1[6]*L2[6],
        L1[7],L1[8]
    )
end
@inline function Base.:*(L1::NTuple{6,T}, L2::W) where {T, W <: WishartFactor{T}}
    @fastmath @inbounds W(
        L1[1]*L2[1],
        L1[2]*L2[1] + L1[4]*L2[2],
        L1[3]*L2[1] + L1[5]*L2[2] + L1[6]*L2[3],
        L1[4]*L2[4],
        L1[5]*L2[4] + L1[6]*L2[5],
        L1[6]*L2[6],
        L2[7],L2[8]
    )
end
@inline function Base.:/(a::W, x::T) where {T <: Number, W <: Wishart3x3{T}}
    @fastmath @inbounds begin
        xinv = 1 / x
        W(
            a[1] * xinv, a[2] * xinv, a[3] * xinv, a[4] * xinv, a[5] * xinv, a[6] * xinv, a[7], a[8]
        )
    end
end

@inline extract_ν(iw::Wishart3x3) = iw[7]
@generated function extract_ν(iw::AbstractVector{<:Wishart3x3{T}}, ::Val{NG}) where {NG,T}
    quote
        $(Expr(:meta, :inline))
        ptr_iw = Base.unsafe_convert(Ptr{$T}, pointer(iw))
        $(Expr(:tuple, [:(unsafe_load(ptr_iw, $((ng-1)*8) + 7 ) ) for ng ∈ 1:NG]...))
    end
end
@inline extract_α(iw::Union{InverseWishart,CholInvWishart}) = iw[8]
@generated function extract_α(iw::AbstractVector{<:Union{InverseWishart{T},CholInvWishart{T}}}, ::Val{NG}) where {NG,T}
    quote
        $(Expr(:meta, :inline))
        ptr_iw = Base.unsafe_convert(Ptr{$T}, pointer(iw))
        $(Expr(:tuple, [:(unsafe_load(ptr_iw, $((ng-1)*8) + 8 ) ) for ng ∈ 1:NG]...))
    end
end

@inline function finite_sqrt(x::T, min_value = eps(T)) where {T}
    rx = @fastmath sqrt(x)
    # min_value = eps(T)
    # min_value = T(0.005) # 5
    isfinite(rx) ? rx : min_value
end

# @inline
function inv_and_cholesky!(riwv::AbstractVector{RevCholWishart{T}},
                                    ciwv::AbstractVector{CholInvWishart{T}},
                                    iwv::AbstractVector{InverseWishart{T}}) where T

    @inbounds for i ∈ eachindex(iwv)
        iw = iwv[i]
        @fastmath begin
            min_value = T(0.05) * sqrt(iw[7] - 2)
            L11 = finite_sqrt(iw[1], min_value)
            R11 = 1 / L11
            L21 = R11 * iw[2]
            L31 = R11 * iw[3]
            L22 = finite_sqrt(iw[4] - L21^2, min_value)
            R22 = 1 / L22
            L32 = R22 * (iw[5] - L21 * L31)
            L33 = finite_sqrt(iw[6] - L31^2 - L32^2, min_value)
            R33 = 1 / L33

            R21 = - R22 * L21 * R11
            R31 = - R33 * ( L31*R11 + L32*R21 )
            R32 = - R33 * L32 * R22
        end
        ciwv[i] = CholInvWishart(
            L11, L21, L31, L22, L32, L33, iw[7], iw[8]
        )
        riwv[i] = RevCholWishart(
            R11, R21, R31, R22, R32, R33, iw[7], iw[8]
        )
    end
end
@inline function triangle_inv(L::NTuple{8,T}) where T
    @inbounds @fastmath begin
        L11           = L[1]
        L21, L22      = L[2], L[4]
        L31, L32, L33 = L[3], L[5], L[6]
        R11 = 1 / L11
        R22 = 1 / L22
        R33 = 1 / L33
        R21 = - R22 * L21 * R11
        R31 = - R33 * ( L31*R11 + L32*R21 )
        R32 = - R33 * L32 * R22
    end
    @inbounds (R11, R21, R31, R22, R32, R33, L[7], L[8])
end
@inline function triangle_inv(L::NTuple{6,T}) where T
    @inbounds @fastmath begin
        L11           = L[1]
        L21, L22      = L[2], L[4]
        L31, L32, L33 = L[3], L[5], L[6]
        R11 = 1 / L11
        R22 = 1 / L22
        R33 = 1 / L33
        R21 = - R22 * L21 * R11
        R31 = - R33 * ( L31*R11 + L32*R21 )
        R32 = - R33 * L32 * R22
    end
    R11, R21, R31, R22, R32, R33
end
@inline function Base.inv(w::W) where {T,W<:WishartFactor{T}}
    inv_type(W)(triangle_inv(w.data))
end

@inline function UpperTriangle3(U::WishartFactor{T}) where T
    @inbounds UpperTriangle3(SVector{6,T}(
        U[1], U[2], U[4], U[3], U[5], U[6]
    ))
end

@inline function randinvwishartfactor(rng::AbstractRNG, rcw::RevCholWishart{T}) where T
    rcwU = UpperTriangle3(rcw)
    ν = extract_ν(rcw)
    U11 = randchi(rng, ν - 2)
    U12 = randn(rng)
    U22 = randchi(rng, ν - 1)
    U13 = randn(rng)
    U23 = randn(rng)
    U33 = randchi(rng, ν)
    U =  UpperTriangle3(
        U11, U12, U22, U13, U23, U33
    )
#    @show rcw
#    @show rcwU
#    @show U
    inv(rcwU * U)
end
function randinvwishart(rng::AbstractRNG, rcw::RevCholWishart{T}) where T
    xtx( randinvwishartfactor(rng, rcw) )
end