text:
using Winston
import Winston.plot
import Visualizer.visualize

visualize{T<:GroupedCountEntropy}(::Type{T}, fnames) = visualize(T, fnames, "Grouped Count Entropy")
visualize{T<:GroupedTemporalEntropy}(::Type{T}, fnames) = visualize(T, fnames, "Grouped Temporal Entropy")
visualize{T<:GroupedPopulationEntropy}(::Type{T}, fnames) = visualize(T, fnames, "Grouped Population Entropy")

function visualize{T<:String,T2<:GroupedShannonEntropy}(::Type{T2}, fnames::Array{T,1},title::String)
	GTE = T2(fnames[1])
	ta = plot(GTE)

	function func3(_ta,Q,i)
		if i > 0 && i <= length(Q)
			GTE = T2(Q[i])
			plot(_ta, GTE)
		end
	end
	visualize(fnames,800,600,title, func3, ta)
end

function plot(H::GroupedShannonEntropy)
	nrows = int(ceil(sqrt(maximum(H.group_labels))))
	ncols = nrows
	ta = Winston.Table(nrows,ncols)
	for i=1:ncols
		for j=1:nrows
		   l = ((i-1)*nrows+j)
		   if (l in H.group_labels) || (i==div(ncols,2)+1) && (j==div(nrows,2)+1)
				ta[j,i] = Winston.FramedPlot()
			end
		end
	end
	ta = plot(ta, H)
	ta
end

function plot(ta::Winston.Table,H::GroupedShannonEntropy)
	nrows = ta.rows
	ncols = ta.cols
	c = (div(nrows,2)+1 , div(ncols,2)+1)
	ta.attr[:align_interiors] = true
	ta.attr[:cellspacing] = 6.0 #tweaks
	bins = H.bins[1:size(H.H,1)]

	for i=1:ncols
		for j=1:nrows
		   l = ((i-1)*nrows+j)
			if !(l in H.group_labels)
				continue
			end
			idx = find(H.group_labels.==l)
			k = 1
			bidx = find(!isfinite(H.Hc[:,idx]))
			if !isempty(bidx)
				for bi in bidx
					if !(k in bidx)
						Winston.add(ta[j,i],Winston.FillBetween(bins[k:bi-1], H.Hc[:,idx][k:bi-1]-H.dHc[:,idx][k:bi-1],
																bins[k:bi-1], H.Hc[:,idx][k:bi-1] + H.dHc[:,idx][k:bi-1]))
					end
					k = bi + 1

				end
			end
			Winston.add(ta[j,i],Winston.FillBetween(bins[k:end], H.Hc[:,idx][k:end]-H.dHc[:,idx][k:end],
													bins[k:end], H.Hc[:,idx][k:end] + H.dHc[:,idx][k:end]))

			Winston.add(ta[j,i],Winston.Curve(bins, H.Hc[:,idx][:]))
		    Winston.setattr(ta[j,i],"title","location $l")
		    Winston.setattr(ta[j,i].x2,"draw_axis",false)
		    Winston.setattr(ta[j,i].y2,"draw_axis",false)
		    Winston.setattr(ta[j,i].x1,"tickdir",1)
		    Winston.setattr(ta[j,i].y1,"tickdir",1)
		end
	end
	#plot the information in the center plot
	p = ta[c...]
	k = 1
	bidx = find(!isfinite(H.I))
	if !isempty(bidx)
		for bi in bidx
			if !(k in bidx)
				Winston.add(p,Winston.FillBetween(bins[k:bi-1], H.I[k:bi-1]-H.dI[k:bi-1],
														bins[k:bi-1], H.I[k:bi-1] + H.dI[k:bi-1]))
			end
			k = bi + 1

		end
	end
	Winston.add(p,Winston.FillBetween(bins[k:end], H.I[k:end]-H.dI[k:end],
											bins[k:end], H.I[k:end] + H.dI[k:end]))
	Winston.add(p, Winston.Curve(bins, H.I))
	Winston.setattr(p.x2,"draw_axis",false)
	Winston.setattr(p.y2,"draw_axis",false)
	Winston.setattr(p.x1,"tickdir",1)
	Winston.setattr(p.y1,"tickdir",1)
	Winston.setattr(p, "xlabel","Time [ms]")
	Winston.setattr(p, "ylabel", "Information [bits]")
	ta
end

function plot(H::TemporalEntropy)
	p = Winston.FramedPlot()
	plot(p, H)
	p
end

function plot(p::Winston.FramedPlot,H::TemporalEntropy)
	Winston.add(p,Winston.Curve(H.bins, H.H))
	Winston.setattr(p.x2,"draw_axis",false)
	Winston.setattr(p.y2,"draw_axis",false)
	Winston.setattr(p.x1,"tickdir",1)
	Winston.setattr(p.y1,"tickdir",1)
	Winston.setattr(p, "xlabel", "Time [ms]")
	Winston.setattr(p, "ylabel", "Entropy [bits]")
	p
end


text:
function tsQuickScatter(o::tsParm,pngImageFile::AbstractString)

  plt=PyPlot
  plt.ioff()
  plt.clf()
  nPlots=size(o.Data,2)
  #nPlots=4
  f=plt.figure("tsQuickScatter",figsize=(4*nPlots,4*nPlots))
  subplots_adjust(wspace=0.3)
  k=0
  xBinCentres=zeros(19)
  p5=zeros(19)
  p50=zeros(19)
  p95=zeros(19)
  for i=1:nPlots
    for j=1:nPlots
       k+=1
       ax=plt.subplot(nPlots,nPlots,k)
       ax[:set_ylabel](o.Name[j],fontsize=10)
       ax[:tick_params]("both",labelsize=10)
       ax[:set_xlabel](o.Name[i],fontsize=10)
       ax[:tick_params]("both",labelsize=10)

       x=o.Data[:,i]
       y=o.Data[:,j]

       xMin,xMax,xTicks=ascale(minimum(x),maximum(x),"s")
       # Calculate binned percentiles
       xBinEdges=linspace(xTicks[1],xTicks[end],20)
       mask=falses(length(x))
       for ik=2:length(xBinEdges)
          mask = (x.>xBinEdges[ik-1]) & (x.<xBinEdges[ik]) & (!isnan(y))
	  if all(!mask)
	    p5[ik-1]=NaN
	    p50[ik-1]=NaN
	    p95[ik-1]=NaN
	  else
            p5[ik-1]=StatsBase.percentile(y[mask],5)
            p50[ik-1]=StatsBase.percentile(y[mask],50)
            p95[ik-1]=StatsBase.percentile(y[mask],95)
	  end
          xBinCentres[ik-1]=0.5*(xBinEdges[ik-1]+xBinEdges[ik])
       end

       ax[:set_xlim]((xMin,xMax))
       ax[:set_xticks]((xTicks))
       plot(x,y,"b.",color="#a0a060")
       plot(xBinCentres,p5,"k-o",color="#3030f0")
       plot(xBinCentres,p50,"r-o")
       plot(xBinCentres,p95,"k-o",color="#3030f0")
       grid("on")
    end
  end
  savefig(pngImageFile)
  println(*("File saved to:",pngImageFile))
  return true
  end

function tsQuickPlot(o::tsParm,pngImageFile::AbstractString)

   plt=PyPlot
   #plt.ioff()
   plt.clf()
   nPlots=size(o.Data,2)
   f=plt.figure("tsQuickPlot",figsize=(10,4*nPlots))
   subplots_adjust(hspace=0.25)
   #majorformatter = matplotlib[:dates][:DateFormatter]("%d.%m.%Y")
   #majorformatter = matplotlib[:dates][:DateFormatter]("%Y")
   #minorformatter = matplotlib[:dates][:DateFormatter]("%H:%M")
   #majorlocator = matplotlib[:dates][:MonthLocator](interval=12*5)
   #minorlocator = matplotlib[:dates][:HourLocator](byhour=(8, 16))
   jDates=Dates.datetime2julian(o.sDate)
   ind=[x for x=1:length(o.sDate)]
   # If running in julia not IJulia
   #pygui(true)
   dateTicks(o.sDate)

   for i=1:nPlots
      ax=plt.subplot(nPlots,1,i)
      ax[:set_ylabel](o.Name[i],fontsize=10)
      ax[:tick_params]("both",labelsize=10)
      #ax[:xaxis][:set_major_formatter](majorformatter)
      #ax[:xaxis][:set_minor_formatter](minorformatter)
      #ax[:xaxis][:set_major_locator](majorlocator)
      #ax[:xaxis][:set_minor_locator](minorlocator)
      #plt.plot_date(o.sDate,o.Data[:,i],markersize=1,fmt=".",color="red", linewidth=0.25)
      #plt.plot_date(jDates,o.Data[:,i],markersize=1,fmt=".",color="red", linewidth=0.25)
      #plt.plot_date(ind,o.Data[:,i],markersize=1,fmt=".",color="red", linewidth=0.25)
      #plt.plot(ind,o.Data[:,i],"r.")#,markersize=1,fmt=".",color="red", linewidth=0.25)
      plt.plot(jDates-jDates[1],o.Data[:,i],"r.",markersize=1)#,fmt=".",color="red", linewidth=0.25)
      grid("on")
   end

   plt.show()
   savefig(pngImageFile)
   println(*("File saved to:",pngImageFile))

   return true
end

ext: jl

max_stars_repo_name: jadamc/MetoceanTools

max_issues_repo_name: jadamc/MetoceanTools

max_forks_repo_name: jadamc/MetoceanTools


text:
"helper function to calculate a run-length encoding of a tuple type"
Base.@pure function runlength(::Type{T}) where {T <: Tuple}
    rle = Tuple{Type, Int}[]
    fieldcount(T) == 0 && return rle
    curT = fieldtype(T, 1)
    prevT = curT
    len = 1
    for i = 2:fieldcount(T)
        @inbounds curT = fieldtype(T, i)
        if curT === prevT
            len += 1
        else
            push!(rle, (prevT, len))
            prevT = curT
            len = 1
        end
    end
    push!(rle, (curT, len))
    return rle
end

"""
    Tables.eachcolumn(f, sch::Tables.Schema{names, types}, x::Union{Tables.AbstractRow, Tables.AbstractColumns})
    Tables.eachcolumn(f, sch::Tables.Schema{names, nothing}, x::Union{Tables.AbstractRow, Tables.AbstractColumns})

Takes a function `f`, table schema `sch`, `x`, which is an object that satisfies the `AbstractRow` or `AbstractColumns` interfaces;
it generates calls to get the value for each column (`Tables.getcolumn(x, nm)`) and then calls `f(val, index, name)`, where `f` is the
user-provided function, `val` is the column value (`AbstractRow`) or entire column (`AbstractColumns`), `index` is the column index as an `Int`,
and `name` is the column name as a `Symbol`.

An example using `Tables.eachcolumn` is:
```julia
rows = Tables.rows(tbl)
sch = Tables.schema(rows)
if sch === nothing
    state = iterate(rows)
    state === nothing && return
    row, st = state
    sch = Tables.schema(Tables.columnnames(row), nothing)
    while state !== nothing
        Tables.eachcolumn(sch, row) do val, i, nm
            bind!(stmt, i, val)
        end
        state = iterate(rows, st)
        state === nothing && return
        row, st = state
    end
else
    for row in rows
        Tables.eachcolumn(sch, row) do val, i, nm
            bind!(stmt, i, val)
        end
    end
end
```
Note in this example we account for the input table potentially returning `nothing` from `Tables.schema(rows)`; in that case, we start
iterating the rows, and build a partial schema using the column names from the first row `sch = Tables.schema(Tables.columnnames(row), nothing)`,
which is valid to pass to `Tables.eachcolumn`.
"""
function eachcolumn end

quot(s::Symbol) = Meta.QuoteNode(s)
quot(x::Int) = x

@inline function eachcolumn(f::F, sch::Schema{names, types}, row::T) where {F, names, types, T}
    N = fieldcount(types)
    if N <= SPECIALIZATION_THRESHOLD
        Base.@nexprs 100 i -> begin
            if i <= N
                f(getcolumn(row, fieldtype(types, i), i, names[i]), i, names[i])
            end
        end
    else
        for (i, nm) in enumerate(names)
            f(getcolumn(row, fieldtype(types, i), i, nm), i, nm)
        end
    end
    return
end

@inline function eachcolumn(f::F, sch::Schema{names, nothing}, row::T) where {F, names, T}
    N = length(names)
    if N <= SPECIALIZATION_THRESHOLD
        Base.@nexprs 100 i -> begin
            if i <= N
                f(getcolumn(row, names[i]), i, names[i])
            end
        end
    else
        for (i, nm) in enumerate(names)
            f(getcolumn(row, nm), i, nm)
        end
    end
    return
end

# these are specialized `eachcolumn`s where we also want
# the indexing of `columns` to be constant propagated, so it needs to be returned from the generated function
@inline function eachcolumns(f::F, sch::Schema{names, types}, row::T, columns::S, args...) where {F, names, types, T, S}
    N = fieldcount(types)
    if N <= SPECIALIZATION_THRESHOLD
        Base.@nexprs 100 i -> begin
            if i <= N
                f(getcolumn(row, fieldtype(types, i), i, names[i]), i, names[i], columns[i], args...)
            end
        end
    else
        for (i, nm) in enumerate(names)
            f(getcolumn(row, fieldtype(types, i), i, nm), i, nm, columns[i], args...)
        end
    end
    return
end

@inline function eachcolumns(f::F, sch::Schema{names, nothing}, row::T, columns::S, args...) where {F, names, T, S}
    N = length(names)
    if N <= SPECIALIZATION_THRESHOLD
        Base.@nexprs 100 i -> begin
            if i <= N
                f(getcolumn(row, names[i]), i, names[i], columns[i], args...)
            end
        end
    else
        for (i, nm) in enumerate(names)
            f(getcolumn(row, nm), i, nm, columns[i], args...)
        end
    end
    return
end

"""
    rowmerge(row, other_rows...)
    rowmerge(row; fields_to_merge...)

Return a `NamedTuple` by merging `row` (an `AbstractRow`-compliant value) with `other_rows`
(one or more `AbstractRow`-compliant values) via `Base.merge`. This function is similar to
`Base.merge(::NamedTuple, ::NamedTuple...)`, but accepts `AbstractRow`-compliant values
instead of `NamedTuple`s.

A convenience method `rowmerge(row; fields_to_merge...) = rowmerge(row, fields_to_merge)`
is defined that enables the `fields_to_merge` to be specified as keyword arguments.
"""
rowmerge(row, other) = merge(_row_to_named_tuple(row), _row_to_named_tuple(other))
rowmerge(row, other, more...) = merge(_row_to_named_tuple(row), rowmerge(other, more...))
rowmerge(row; fields_to_merge...) = rowmerge(row, fields_to_merge.data)

_row_to_named_tuple(row::NamedTuple) = row
_row_to_named_tuple(row) = NamedTuple(Row(row))


text:
using Printf, Test, Plots, LinearAlgebra, BenchmarkTools

include("../../hybmc/mathutils/Helpers.jl")

include("../../hybmc/termstructures/YieldCurve.jl")
include("../../hybmc/models/HullWhiteModel.jl")
include("../../hybmc/models/AssetModel.jl")
include("../../hybmc/models/DeterministicModel.jl")
include("../../hybmc/models/QuasiGaussianModel.jl")
include("../../hybmc/models/HybridModel.jl")

include("../../hybmc/simulations/McSimulation.jl")
include("../../hybmc/simulations/Payoffs.jl")

function HwModel(rate=0.01,vol=0.0050,mean=0.03)
    curve = YieldCurve(rate)
    times = [ 10.0 ]
    vols  = [ vol  ]
    return HullWhiteModel(curve, mean, times, vols)
end

function HybridHullWhiteModel4Ccy()
    domAlias = "EUR"
    domModel = HwModel(0.01, 0.0050, 0.03)
    forAliases = [ "USD", "GBP", "JPY" ]
    spotS0     = [   1.0,   2.0,   3.0 ]
    spotVol    = [   0.3,   0.3,   0.3 ]
    rates      = [  0.02,  0.03,  0.04 ]
    ratesVols  = [ 0.006, 0.007, 0.008 ]
    mean       = [  0.01,  0.01,  0.01 ]
    #
    forAssetModels = [
        AssetModel(S0, vol) for (S0, vol) in zip(spotS0,spotVol) ]
    forRatesModels = [
        HwModel(r,v,m) for (r,v,m) in zip(rates,ratesVols,mean) ]
    #
    nCorrs = 2 * size(forAliases)[1] + 1
    corr = Matrix{Float64}(I, nCorrs, nCorrs)
    corr[2,3] = 0.6  # USD quanto
    corr[3,2] = 0.6
    corr[4,5] = 0.8  # GBP quanto
    corr[5,4] = 0.8
    corr[6,7] = -0.7  # JPY quanto
    corr[7,6] = -0.7
    #
    return HybridModel(domAlias,domModel,forAliases,forAssetModels,forRatesModels,corr)
end

function HybridQuasiaussianModel3Ccy()
    ### full smile/skew model
    # domestic rates
    domAlias = "EUR"
    eurCurve = YieldCurve(0.03)
    d     = 2
    times = [  10.0     ]
    sigma = [ 0.0060 ;
              0.0040 ]
    slope = [ 0.10  ;
              0.15  ]
    curve = [ 0.05  ;
              0.10  ]
    delta = [ 1.0, 10.0  ]
    chi   = [ 0.01, 0.15 ]
    Gamma = [ 1.0   0.6  ;
              0.6   1.0  ]
    eurRatesModel = QuasiGaussianModel(eurCurve,d,times,sigma,slope,curve,delta,chi,Gamma)
    #
    forAliases = [ "USD", "GBP" ]
    spotS0     = [   1.0,   2.0 ]
    spotVol    = [   0.3,   0.2 ]
    forAssetModels = [
        AssetModel(S0, vol) for (S0, vol) in zip(spotS0,spotVol) ]
    # USD rates
    usdCurve = YieldCurve(0.02)
    d     = 3
    times = [  10.0   ]
    sigma = [  0.0060 ;
               0.0050 ;
               0.0040 ]
    slope = [  0.10   ;
               0.20   ;
               0.30   ]
    curve = [  0.05   ;
               0.10   ;
               0.20   ]
    delta = [ 1.0, 5.0, 20.0 ]
    chi   = [ 0.01, 0.05, 0.15 ]
    Gamma = [ 1.0  0.8  0.6 ;
              0.8  1.0  0.8 ;
              0.6  0.8  1.0 ]
    usdRatesModel = QuasiGaussianModel(usdCurve,d,times,sigma,slope,curve,delta,chi,Gamma)
    #
    gbpRatesModel = HwModel()
    #
           # 'EUR_x_0', 'EUR_x_1', 'USD_logS', 'USD_x_0', 'USD_x_1', 'USD_x_2', 'GBP_logS', 'GBP_x'
    corr = [  1.0        0.0         0.5       -0.5        0.0        0.0        -0.5       -0.5   ;   # EUR_x_0
              0.0        1.0         0.0        0.0       -0.5        0.0        -0.5        0.0   ;   # EUR_x_1
              0.5        0.0         1.0       -0.5       -0.5       -0.5         0.0        0.0   ;   # USD_logS
             -0.5        0.0        -0.5        1.0        0.0        0.0         0.0        0.0   ;   # USD_x_0
              0.0       -0.5        -0.5        0.0        1.0        0.0         0.0        0.0   ;   # USD_x_1
              0.0        0.0        -0.5        0.0        0.0        1.0         0.0        0.0   ;   # USD_x_2
             -0.5       -0.5         0.0        0.0        0.0        0.0         1.0        0.5   ;   # GBP_logS
             -0.5        0.0         0.0        0.0        0.0        0.0         0.5        1.0   ]   # GBP_x
    model = HybridModel(domAlias,eurRatesModel,forAliases,forAssetModels,[usdRatesModel,gbpRatesModel],corr)
    return model
end

function HybridHwDcfModel2Ccy()
    curve0 = YieldCurve(0.03)
    hwModel = HwModel(0.01,0.0050,0.03)
    asModel = AssetModel(1.0,0.30)
    dcfModel0 = DcfModel(curve0,"USD")
    corr =  Matrix{Float64}(I, 2, 2)
    model = HybridModel("EUR",hwModel,["USD"],[asModel],[dcfModel0],corr)
    return model
end

function fwd(mcSim::McSimulation,p::Payoff)
    samples = discountedAt(p,mcSim)
    fwd = mean(samples) /
        discount(mcSim.model.domRatesModel.yieldCurve,obsTime(p))
    err = std(samples) / sqrt(size(samples)[1]) /
        discount(mcSim.model.domRatesModel.yieldCurve,obsTime(p))
    return (fwd, err)
end

function npv(mcSim::McSimulation,leg)
    samples = transpose(discountedAt(leg,mcSim))
    npv_ = mean(sum(samples, dims=2), dims=1)
    err = std(sum(samples, dims=2), dims=1)
    return (npv_,err)
end

yieldCurve(model) = model.yieldCurve

yieldCurve(model::DeterministicModel) = model.domCurve

yieldCurve(model::HybridModel) = yieldCurve(model.domRatesModel)

yieldCurve(sim::McSimulation) = yieldCurve(sim.model)

# We run the test

function run_test(model, modelName)
    times = append!(Array(range(0,stop=10.0,length=11)), [10.5])
    nPaths = 2^13
    seed = 141592653593
    @printf("\n")
    @printf("Run performance test for model %s: nPaths = %d  seed = %d\n", modelName, nPaths, seed)
    sim = McSimulation(model,times,2,seed,true)  #  warm-up
    @printf("Simulation:\n")
    @time sim = McSimulation(model,times,nPaths,seed,true)
    # domestic numeraire
    T = 10.0
    P = Pay(Fixed(1.0),T)
    (fw, err) = fwd(sim,P)
    @printf("1.0   @ %4.1lfy %8.6lf - mc_err = %8.6lf\n", T,fw,err)
    @time fwd(sim,P)
    # foreign asset
    alias = "USD"
    k = 1
    p = Asset(T,alias)
    xT = model.forAssetModels[k].X0 *
        discount(yieldCurve(model.forRatesModels[k]),T) /
        discount(yieldCurve(model.domRatesModel),T)
    (fw, err) = fwd(sim,p)
    @printf("%s   @ %4.1lfy %8.6lf vs %8.6lf (curve) - mc_err = %8.6lf\n", alias,T,fw,xT,err)
    @time fwd(sim,p)
    # domestic Libor rate
    Tstart = 10.0
    Tend = 10.5
    L = Pay(LiborRate(T,Tstart,Tend,alias="EUR"),Tend)
    (fw, err) = fwd(sim,L)
    Lref = (discount(yieldCurve(sim),Tstart) /
            discount(yieldCurve(sim), Tend) - 1) /
            (Tend - Tstart)
    @printf("L_EUR @ %4.1lfy %8.6lf vs %8.6lf (curve) - mc_err = %8.6lf\n", T,fw,Lref,err)
    @time fwd(sim,L)
    # foreign Lbor rates
    alias = "USD"
    k = 1
    L = Pay(LiborRate(T,Tstart,Tend,alias=alias)*Asset(Tend,alias),Tend)
    (fw, err) = fwd(sim,L)
    fw *= discount(yieldCurve(sim.model.domRatesModel), Tend) /
          discount(yieldCurve(sim.model.forRatesModels[k]), Tend) /
          sim.model.forAssetModels[k].X0
    err *= discount(yieldCurve(sim.model.domRatesModel), Tend) /
           discount(yieldCurve(sim.model.forRatesModels[k]), Tend) /
           sim.model.forAssetModels[k].X0
    Lref = (discount(yieldCurve(sim.model.forRatesModels[k]),Tstart) /
            discount(yieldCurve(sim.model.forRatesModels[k]), Tend) - 1) /
            (Tend - Tstart)
    @printf("L_%s @ %4.1lfy %8.6lf vs %8.6lf (curve) - mc_err = %8.6lf\n", alias,T,fw,Lref,err)
    @time fwd(sim,L)
    # Libor leg
    alias = "EUR"
    schedule = Array(range(0,stop=T,length=21))
    Leg = [ (0.5 * LiborRate(Tstart,Tstart,Tend,alias="EUR")) % Tend
            for (Tstart,Tend) in zip(schedule[begin:end-1], schedule[begin+1:end]) ]
    (npv_,err) = npv(sim,Leg)
    npv_ref = 1.0 - discount(yieldCurve(sim), T)
    @printf("F_%s @ %4.1lfy %8.6lf vs %8.6lf (curve) - mc_err = %8.6lf\n", alias,T,npv_[1,1],npv_ref,err[1,1])
    @time npv(sim,Leg)
end

# actually run the tests...

run_test(HybridHwDcfModel2Ccy(), "HybridHwDcfModel2Ccy")

run_test(HybridHullWhiteModel4Ccy(), "HybridHullWhiteModel4Ccy")

run_test(HybridQuasiaussianModel3Ccy(), "HybridQuasiaussianModel3Ccy")

max_stars_repo_name: sschlenkrich/HybridMonteCarlo

max_issues_repo_name: sschlenkrich/HybridMonteCarlo

max_forks_repo_name: sschlenkrich/HybridMonteCarlo


text:
###############################################################################
#
#  NfRelOrd
#
###############################################################################

# I don't like that we have to drag S around
# Still hoping for julia/#18466

mutable struct NfRelOrdSet{T}
  nf::NumField{T}

  function NfRelOrdSet{T}(K::NumField{T}) where {T}
    a = new(K)
    return a
  end
end

mutable struct NfRelOrd{T, S, U} <: Ring
  nf::NumField{T}
  basis_nf::Vector{U}
  basis_matrix::Generic.MatSpaceElem{T}
  basis_mat_inv::Generic.MatSpaceElem{T}
  basis_pmatrix::PMat{T, S}
  pseudo_basis::Vector{Tuple{U, S}}

  disc_abs::NfOrdIdl # used if T == nf_elem
  disc_rel#::NfRelOrdIdl{T} # used otherwise; is a forward declaration
  parent::NfRelOrdSet{T}

  isequation_order::Bool

  ismaximal::Int                   # 0 Not known
                                   # 1 Known to be maximal
                                   # 2 Known to not be maximal

  trace_mat::Generic.MatSpaceElem{T}

  index::NfOrdIdl #Only if the base field is AnticNumberField

  inv_coeff_ideals::Vector{S}
  index_div
  @declare_other

  function NfRelOrd{T, S, U}(K::NumField{T}) where {T, S, U}
    z = new{T, S, U}()
    z.nf = K
    z.parent = NfRelOrdSet{T}(K)
    z.isequation_order = false
    z.ismaximal = 0
    z.index_div = Dict{ideal_type(order_type(base_field(K))), Vector{Tuple{ideal_type(z), Int}}}()
    return z
  end

  function NfRelOrd{T, S, U}(K::NumField{T}, M::PMat{T, S}) where {T, S, U}
    z = NfRelOrd{T, S, U}(K)
    z.nf = K
    z.parent = NfRelOrdSet{T}(K)
    z.basis_pmatrix = M
    z.basis_matrix = M.matrix
    z.index_div = Dict{ideal_type(order_type(base_field(K))), Vector{Tuple{ideal_type(z), Int}}}()
    return z
  end

  function NfRelOrd{T, S, U}(K::NumField{T}, M::Generic.MatSpaceElem{T}) where {T, S, U}
    z = NfRelOrd{T, S, U}(K)
    z.nf = K
    z.parent = NfRelOrdSet{T}(K)
    z.basis_matrix = M
    z.basis_pmatrix = pseudo_matrix(M)
    z.index_div = Dict{ideal_type(order_type(base_field(K))), Vector{Tuple{ideal_type(z), Int}}}()
    return z
  end
end

###############################################################################
#
#  NfRelOrdElem
#
###############################################################################

mutable struct NfRelOrdElem{T, U} <: RingElem
  parent#::NfRelOrd{T, S} # I don't want to drag the S around
  elem_in_nf::U
  coordinates::Vector{T}
  has_coord::Bool

  function NfRelOrdElem{T, U}(O::NfRelOrd{T}) where {T, U}
    z = new{T, U}()
    z.parent = O
    z.elem_in_nf = zero(nf(O))
    z.coordinates = Vector{T}(undef, degree(O))
    z.has_coord = false
    return z
  end

  function NfRelOrdElem{T, U}(O::NfRelOrd{T}, a::U) where {T, U}
    z = new{T, U}()
    z.parent = O
    z.elem_in_nf = a
    z.coordinates = Vector{T}(undef, degree(O))
    z.has_coord = false
    return z
  end

  function NfRelOrdElem{T, U}(O::NfRelOrd{T}, a::U, arr::Vector{T}) where {T, U}
    z = new{T, U}()
    z.parent = O
    z.elem_in_nf = a
    z.coordinates = arr
    z.has_coord = true
    return z
  end
end

###############################################################################
#
#  NfRelOrdFracIdl
#
###############################################################################

mutable struct NfRelOrdFracIdlSet{T, S, U}
  order::NfRelOrd{T, S, U}

  function NfRelOrdFracIdlSet{T, S, U}(O::NfRelOrd{T, S, U}) where {T, S, U}
    a = new(O)
    return a
  end
end

mutable struct NfRelOrdFracIdl{T, S, U}
  order::NfRelOrd{T, S, U}
  parent::NfRelOrdFracIdlSet{T, S, U}
  basis_pmatrix::PMat{T, S}
  pseudo_basis::Vector{Tuple{U, S}}
  basis_matrix::Generic.MatSpaceElem{T}
  basis_mat_inv::Generic.MatSpaceElem{T}
  den::fmpz

  norm
  has_norm::Bool

  function NfRelOrdFracIdl{T, S, U}(O::NfRelOrd{T, S, U}) where {T, S, U}
    z = new{T, S, U}()
    z.order = O
    z.parent = NfRelOrdFracIdlSet{T, S, U}(O)
    z.has_norm = false
    return z
  end

  function NfRelOrdFracIdl{T, S, U}(O::NfRelOrd{T, S, U}, M::PMat{T, S}) where {T, S, U}
    z = NfRelOrdFracIdl{T, S, U}(O)
    z.basis_pmatrix = M
    z.basis_matrix = M.matrix
    return z
  end
end

###############################################################################
#
#  NfRelOrdIdl
#
###############################################################################

mutable struct NfRelOrdIdlSet{T, S, U}
  order::NfRelOrd{T, S, U}

  function NfRelOrdIdlSet{T, S, U}(O::NfRelOrd{T, S, U}) where {T, S, U}
    a = new(O)
    return a
  end
end

mutable struct NfRelOrdIdl{T, S, U}
  order::NfRelOrd{T, S, U}
  parent::NfRelOrdIdlSet{T, S, U}
  basis_pmatrix::PMat{T, S}
  pseudo_basis::Vector{Tuple{U, S}}
  basis_matrix::Generic.MatSpaceElem{T}
  basis_mat_inv::Generic.MatSpaceElem{T}

  norm
  has_norm::Bool
  is_prime::Int            # 0: don't know
                           # 1 known to be prime
                           # 2 known to be not prime
  splitting_type::Tuple{Int, Int}

  minimum
  non_index_div_poly::fq_poly # only used if the ideal is a prime ideal not dividing the index
  p_uniformizer::NfRelOrdElem{T, U}
  anti_uniformizer::U

  function NfRelOrdIdl{T, S, U}(O::NfRelOrd{T, S, U}) where {T, S, U}
    z = new{T, S, U}()
    z.order = O
    z.parent = NfRelOrdIdlSet{T, S, U}(O)
    z.has_norm = false
    z.is_prime = 0
    z.splitting_type = (0,0)
    return z
  end

  function NfRelOrdIdl{T, S, U}(O::NfRelOrd{T, S, U}, M::PMat{T, S}) where {T, S, U}
    z = NfRelOrdIdl{T, S, U}(O)
    z.basis_pmatrix = M
    z.basis_matrix = M.matrix
    return z
  end
end

################################################################################
#
#  NfRelNS / NfRelNSElem
#
################################################################################

mutable struct NfRelNS{T} <: NonSimpleNumField{T}
  base_ring::Nemo.Field
  pol::Array{Nemo.Generic.MPoly{T}, 1}
  abs_pol::Array{Generic.Poly{T}, 1}
  S::Array{Symbol, 1}
  basis_traces::Vector{Vector{T}}
  auxilliary_data::Array{Any, 1}
  @declare_other

  function NfRelNS(abs_pol::Array{Generic.Poly{T}}, f::Array{Nemo.Generic.MPoly{T}, 1}, S::Array{Symbol, 1}; cached::Bool = false) where T
    r = new{T}()
    r.pol = f
    r.abs_pol = abs_pol
    r.base_ring = base_ring(f[1])
    r.S = S
    r.auxilliary_data = Array{Any}(undef, 5)
    return r
  end
end

mutable struct NfRelNSElem{T} <: NonSimpleNumFieldElem{T}
  data::Nemo.Generic.MPoly{T}
  parent::NfRelNS{T}

  NfRelNSElem{T}(g::Generic.MPoly{T}) where {T} = new{T}(g)
end

################################################################################
#
#  Quotient rings of orders
#
################################################################################

# T1 is the type of the order, T2 the type of the ideal,
# T3 the type of the basis pseudo matrix of the ideal
mutable struct RelOrdQuoRing{T1, T2, T3} <: Ring
  base_ring::T1
  ideal::T2
  basis_pmatrix::T3

  function RelOrdQuoRing{T1, T2, T3}(O::T1, I::T2) where { T1, T2, T3 }
    z = new{T1, T2, T3}()
    z.base_ring = O
    z.ideal = I
    z.basis_pmatrix = basis_pmatrix_wrt(I, O)
    return z
  end
end

function RelOrdQuoRing(O::T1, I::T2) where { T1, T2 }
  @assert T2 == ideal_type(O)
  return RelOrdQuoRing{T1, T2, typeof(basis_pmatrix(I, copy = false))}(O, I)
end

# T1, T2 and T3 as for RelOrdQuoRing, S is the elem_type of the order
mutable struct RelOrdQuoRingElem{T1, T2, T3, S} <: RingElem
  elem::S
  parent::RelOrdQuoRing{T1, T2, T3}

  function RelOrdQuoRingElem{T1, T2, T3, S}(Q::RelOrdQuoRing{T1, T2, T3}, x::S) where { T1, T2, T3, S }
    z = new{T1, T2, T3, S}()
    z.elem = mod(x, Q)
    z.parent = Q
    return z
  end
end

function RelOrdQuoRingElem(Q::RelOrdQuoRing{T1, T2, T3}, x::S) where { T1, T2, T3, S }
  return RelOrdQuoRingElem{T1, T2, T3, S}(Q, x)
end


text:
@everywhere function remove_sigs(j)

    exposuresOutput, accr, frob_rel_div, norm_one_dif = removeAllSingleSignatures(parallelExposure[:,j], parallelProcesses , originalGenomes[:, j])

    return j, exposuresOutput, accr, frob_rel_div, norm_one_dif
end

function deconvoluteMutCounts(
        processes,
        iterations,
        sort_idx,
        n_types_to_remove,
    )

    totalReplicates = 100
    removeLastPercentage = 0.07

    Wall, Hall, genomeErrors, genomesReconstructed = extractSignatures(genomes, processes, iterations)
    Wall, Hall, genomeErrors, genomesReconstructed = filterOutIterations( Wall, Hall, genomeErrors, processes, genomesReconstructed, removeLastPercentage )

    processes, processesStd, exposure, exposureStd, idx, idxS, processStab, processStabAvg, clusterCompactness = evaluateStability(
        Wall, Hall, processes, totalReplicates)

    processes, processesStd, Wall, genomeErrors, genomesReconstructed = addWeak(
        sort_idx, n_types_to_remove, processes, processesStd, Wall, genomeErrors, genomesReconstructed)

    # signature fitting
    exposure_copied = deepcopy(exposure)

    @eval @everywhere parallelExposure=$exposure_copied
    @eval @everywhere parallelProcesses=$processes

    for (j, exposuresOutput, accr, frob_rel_div, norm_one_dif) in pmap(remove_sigs, 1:size(parallelExposure,2))
        exposure[:, j] = exposuresOutput
    end

    # calculate stabilities and errors in reconstruction
    avgReconstructionError, avgReconstructionErrorPercentage, avgStability, reconError, reconErrorPercentage, simError = calculate_stabilities(processes, exposure, processStabAvg)

    return processes, exposure_copied, processStab, processStabAvg, exposure, avgReconstructionError, avgReconstructionErrorPercentage, avgStability, reconError, reconErrorPercentage, simError
end

max_stars_repo_name: ypradat/sigprofilerjulia

max_issues_repo_name: ypradat/sigprofilerjulia

max_forks_repo_name: ypradat/sigprofilerjulia


text:
# Julia wrapper for header: /Volumes/BOOTCAMP/programs/compa_libs/iup/include/iupole.h
# Automatically generated using Clang.jl wrap_c, version 0.0.0

function IupOleControl()
  ccall( (:IupOleControl, iup), Ptr{Cint}, (), )
end
function IupOleControlOpen()
  ccall( (:IupOleControlOpen, iup), Cint, (), )
end


text:
module My-Julia-ML-Package

# Write your package code here.
println("hello, my first line for my Julia ML package hope this goes well")

function sum_of_three(x,y,z)
    return x+y+z
end

sum_of_three(3,4,5)
end


text:
using Amphitheatre
using Test

@testset "Amphitheatre.jl" begin

    @testset "Examples" begin
        @info "Testing if examples run"
        include("../example/extendingAmphi.jl")
    end
end

max_stars_repo_name: Affie/Amphitheatre

max_issues_repo_name: Affie/Amphitheatre


text:
using  Logging, LoggingExtras

function msg()
    @warn "warning 1\nanother line\nlast one"
    @error "error 1"
    @info "info 1"
    @debug "debug 1"
end

msg()

with_logger(FileLogger("salida.log")) do
    msg()
end


text:
# SPDX-License-Identifier: X11
# 2020-11-19
# Water Heater (400pt)

const TIME_END = 2 * 10^5 + 5

function main()
  N, W = parse.(Int, (split ∘ readline)())
  plans = zeros(Int, TIME_END)
  for i ∈ 1:N
    S,T,P = parse.(Int, (split ∘ readline)())
    plans[S + 1] += P
    plans[T + 1] -= P
  end

  for i ∈ 2:TIME_END
    plans[i] += plans[i - 1]
    if plans[i - 1] > W
      println("No")
      return
    end
  end

  println("Yes")
end

main()

max_stars_repo_name: jthvai/competitive-programming

max_issues_repo_name: jthvai/competitive-programming

max_forks_repo_name: jthvai/competitive-programming


text:
      function percentile(perc::Float64,p::Array{Float64,2})

      sig::Float64=0.0
      sig0::Float64=0.0
      sig1::Float64=0.0
      dsig::Float64=0.0
      pmin::Float64=0.0
      delta::Float64=0.0
      nvec::Int64= size(p,1)*size(p,2)
      nsig::Int64= Int64(floor(nvec/2))
      itnwt::Int64=0

      phi::Float64=0.0
      phi0::Float64=0.0
      phi1::Float64=0.0
      dphi::Float64=0.0
      TOL::Float64 = 0.000001
      #
      pmin=p[1,1]
      delta=pmin
      for i2=1:size(p,2)
          for i1=1:size(p,1)
              delta=max(delta,p[i1,i2])
              pmin=min(pmin,p[i1,i2])
          end
      end
      delta=delta-pmin
      if delta <= 0.0
          println("Null interval ")
          return nothing
      end

      sig0=0.0
      phi0=0.0
      for i2=1:size(p,2)
          for i1=1:size(p,1)
              if (p[i1,i2]-pmin) >= sig0
                 phi0 += 1.0
              end
          end
      end
      phi0 = perc - phi0/Float64(nvec)
      sig1=delta
      phi1=0.0
      for i2=1:size(p,2)
          for i1=1:size(p,1)
              if (p[i1,i2]-pmin) >= sig1
                 phi1+=1.0
              end
          end
      end

      phi1=perc - phi1/Float64(nvec)
      dsig=-phi0*(sig1-sig0)/(phi1-phi0)
      sig = sig0+dsig

      for itnwt=1:99
      phi=0.0
      for i2=1:size(p,2)
          for i1=1:size(p,1)
              if (p[i1,i2]-pmin) >= sig
                 phi=phi+1.0
              end
          end
      end
      phi = perc - phi/Float64(nvec)

      if phi*phi0 < 0.0
         sig1 = sig
         phi1 = phi
      else
         sig0 = sig
         phi0 = phi
      end
      dsig=-phi0*(sig1-sig0)/(phi1-phi0)
      sig = sig0+dsig

      if  abs(phi) < TOL || abs(dsig) < 1.e-16*delta
          println("iter = ",itnwt)
          break
      end
      end

      sig=sig+pmin
      println("sig = ",sig, " phi = ",phi, " dsig = ",dsig)

      return nothing
      end # function


text:
"""
Stochastic Scheme for Uncertainty Quantification
Copyright (c) 2020-2022 Tianbai Xiao
"""

module Langevin

using Reexport
@reexport using PolyChaos
@reexport using KitBase
using Base.Threads: @threads
using KitBase.FiniteMesh.DocStringExtensions
using KitBase.JLD2
using KitBase.OffsetArrays

export AbstractUQ, UQ1D, UQ2D, ran_chaos, chaos_ran, lambda_tchaos, t_lambdachaos
export uq_moments_conserve,
       uq_maxwellian,
       uq_energy_distribution,
       uq_prim_conserve,
       uq_conserve_prim,
       uq_conserve_prim!,
       uq_prim_conserve!,
       uq_sound_speed,
       uq_vhs_collision_time,
       uq_aap_hs_collision_time,
       uq_aap_hs_prim

include("struct.jl")
include("uq.jl")
include("filter.jl")
include("kinetic.jl")
include("initialize.jl")
include("flux_ctr.jl")
include("flux_sol.jl")
include("out.jl")
include("solver.jl")
include("step.jl")

end


text:
import Base.LinAlg.BlasFloat

export eigvals_lanczos

function lanczos!{T}(K::KrylovSubspace{T})
    m = K.n
    αs = Array(T, m)
    βs = Array(T, m-1)
    for j=1:m-1
        w = nextvec(K)
        if j>1 w -= βs[j-1]*K.v[1] end
        w, y = orthogonalize(w, K, 1)
        αs[j] = y[1]
        βs[j] = convert(T, norm(w))
        append!(K, w/βs[j])
    end
    αs[m]= dot(nextvec(K), lastvec(K))
    SymTridiagonal(αs, βs)
end

function eigvals_lanczos(A, neigs::Int=size(A,1); tol::Real=size(A,1)^3*eps(), maxiter::Int=size(A,1))
    K = KrylovSubspace(A, 2) #In Lanczos, only remember the last two vectors
    initrand!(K)
    resnorms = zeros(maxiter)
    e1 = eigvals(lanczos!(K), 1:neigs)
    for iter=1:maxiter
        e0, e1 = e1, eigvals(lanczos!(K), 1:neigs)
        resnorms[iter] = norm(e1-e0)
        if resnorms[iter] < tol
            resnorms = resnorms[1:iter]
            break
        end
    end
    e1, ConvergenceHistory(0<=resnorms[end]<tol, tol, K.mvps, resnorms)
end


text:
## First example block
using Manopt, Manifolds, LinearAlgebra, Random
Random.seed!(42)
M = Sphere(2)
n = 40
p = 1 / sqrt(3) .* ones(3)
B = DefaultOrthonormalBasis()
pts = [exp(M, p, get_vector(M, p, 0.425 * randn(2), B)) for _ in 1:n]

F(M, y) = sum(1 / (2 * n) * distance.(Ref(M), pts, Ref(y)) .^ 2)
gradF(M, y) = sum(1 / n * grad_distance.(Ref(M), pts, Ref(y)))

x_mean = gradient_descent(M, F, gradF, pts[1])

euclidean_mean = mean(pts)
print("Norm of Euclieadn mean:", norm(euclidean_mean), "\n\n")
euclidean_mean_normed = euclidean_mean / norm(euclidean_mean)

## Second example block
o = gradient_descent(
    M,
    F,
    gradF,
    pts[1];
    debug=[:Iteration, " | ", :Change, " | ", :Cost, "\n", :Stop],
    record=[:x, :Change, :Cost],
    return_options=true,
)
x_mean2 = get_solver_result(o)
values = get_record(o) # vector with a tuple per iteration
iterates = get_record(o, :Iteration, :x) # from each iteration get the recorded iterate x.

## Export with Paul Tols colors https://personal.sron.nl/~pault/
using Colors
TolVibrantOrange = RGBA{Float64}(colorant"#EE7733")
TolVibrantTeal = RGBA{Float64}(colorant"#009988")
TolVibrantCyan = RGBA{Float64}(colorant"#33BBEE")
asymptote_export_S2_signals(
    joinpath(@__DIR__, "img/MeanIllustr.asy");
    points=[[x_mean], [euclidean_mean_normed], pts],
    colors=Dict(:points => [TolVibrantOrange, TolVibrantCyan, TolVibrantTeal]),
    dot_sizes=[3.5, 2.5, 3.5],
    camera_position=(0.7, 0.7, 0.5),
)
render_asymptote(joinpath(@__DIR__, "img/MeanIllustr.asy"); render=4)  #src


text:
module Nierika

# using NullableArrays

export  DB, Stmt, Cursor, Table,

        # methods
        colnames,
        query,
        reset!,
        results,
        step!,
        where

include("../deps/deps.jl")
include("SQLite_consts.jl")
include("SQLite_utils.jl")
include("SQLite_api.jl")
include("db.jl")
include("cursor.jl")
include("table.jl")
include("show.jl")

end # module


text:
#Dot product using LoopVectorization (real ⋅ real)
"""
    l2inner(a,b)

Computes the inner product in the l2 vector space (dot product) between Real or Complex vectors `a` and `b`. The lenghts of `a` and `b` must match.

See also [`l2norm`](@ref), [`l2dist`](@ref), [`sql2norm`](@ref), [`sql2dist`](@ref).

# Example
```jldoctest
a = [1,2,3]
b = [2,0,4]
l2inner(a,b)

# output

14
```
"""
function l2inner(a::AbstractVector{T}, b::AbstractVector{T}) where {T}
    s = zero(T)
    @turbo for i ∈ eachindex(a,b)
        s += a[i] * b[i]
    end
    return s
end

#Dot product using LoopVectorization (real ⋅ complex)
function l2inner(a::AbstractVector{T}, cb::AbstractVector{Complex{T}}) where {T}
    re = zero(T)
    im = zero(T)
    b = reinterpret(reshape, T, cb)
    @turbo for i ∈ eachindex(a)
        re += a[i] * b[1,i]
        im += a[i] * b[2,i]
    end
    return Complex(re, im)
end

#Dot product using LoopVectorization (real ⋅ complex)
function l2inner(ca::AbstractVector{Complex{T}}, b::AbstractVector{T}) where {T}
    re = zero(T)
    im = zero(T)
    a = reinterpret(reshape, T, ca)
    @turbo for i ∈ eachindex(b)
        re += a[1,i] * b[i]
        im += - (a[2,i] * b[i])
    end
    return Complex(re, im)
end

#Dot product using LoopVectorization (complex ⋅ complex)
function l2inner(ca::AbstractVector{Complex{T}}, cb::AbstractVector{Complex{T}}) where {T}
    re = zero(T)
    im = zero(T)
    a = reinterpret(reshape, T, ca)
    b = reinterpret(reshape, T, cb)
    @turbo for i ∈ axes(a,2) #Conjugate(a) × b
        re += (a[1,i] * b[1,i]) + (a[2,i] * b[2,i])
        im += (a[1,i] * b[2,i]) - (a[2,i] * b[1,i])
    end
    return Complex(re, im)
end

#L2 Norm square using LoopVectorization (real)
"""
    l2inner(a)

Computes the inner product in the l2 vector space (dot product) of the Real or Complex vectors `a` with itself. i.e., the Squared l2 norm of `a`. This is identical to calling `sql2norm(a)`.
"""
function l2inner(a::AbstractVector{T}) where {T}
    s = zero(T)
    @turbo for i ∈ eachindex(a)
        s += a[i] * a[i]
    end
    return s
end

"""
    sql2norm(a)

Compute the square of the l2 norm of a real valued vector `a`. This is identical to calling `l2inner(a)`.
"""
sql2norm(a::AbstractVector{T}) where {T} = l2inner(a)

"""
    l2norm(a)

Compute the l2 norm of a real valued vector `a`.
"""
l2norm(a::AbstractVector{T}) where {T} = sqrt(l2inner(a))

#L2 Distance square using LoopVectorization (real)
"""
    sql2dist(a,b)

Compute the square of the l2 distance between two real vectors `a` and `b`. The lengths of `a` and `b` must match.
"""
function sql2dist(a::AbstractVector{T}, b::AbstractVector{T}) where {T}
    s = zero(T)
    @turbo for i ∈ eachindex(a,b)
        t = a[i] - b[i]
        s += t * t
    end
    return s
end

"""
    l2dist(a,b)

Compute the l2 distance between two real vectors `a` and `b`. The lengths of `a` and `b` must match.
"""
l2dist(a::AbstractVector{T}, b::AbstractVector{T}) where {T} = sqrt(sql2dist(a,b))

#L1 Norm of real vectors
"""
    l1norm(a)

Compute the l1 norm of a real valued vector `a`.
"""
function l1norm(a::AbstractVector{T}) where {T}
    s = zero(T)
    @turbo for i ∈ eachindex(a)
        s += abs(a[i])
    end
    return s
end

#L1 distance between real vectors
"""
    l1dist(a,b)

Compute the l1 distance between two real vectors `a` and `b`. The lengths of `a` and `b` must match.
"""
 function l1dist(a::AbstractVector{T}, b::AbstractVector{T}) where {T}
    s = zero(T)
    @turbo for i ∈ eachindex(a,b)
        t = a[i] - b[i]
        s += abs(t)
    end
    return s
 end


text:
function gauss(f)
    0.2369268851*f(-0.9061798459)+
    0.2369268851*f(0.9061798459)+
    0.5688888889*f(0)+
    0.4786286705*f(0.5384693101)+
    0.4786286705*f(-0.5384693101)
end

f(x) = 1/sqrt(x^2+6x+13)
F(x) = log(x+3+sqrt(x^2+6x+13))

println("Gauss-Legendre=", gauss(f))
println("F(1)-F(-1)=", F(1)-F(-1))


text:
using RobotDynamics: InPlace, StaticReturn, ForwardAD, FiniteDifference
using BenchmarkTools

struct MyFun{JCH} <: RobotDynamics.ScalarFunction
    gradcfg::ForwardDiff.GradientConfig{Nothing, Float64, JCH, Vector{ForwardDiff.Dual{Nothing, Float64, JCH}}}
    hesscfg::ForwardDiff.HessianConfig{Nothing, Float64, JCH, Vector{ForwardDiff.Dual{Nothing, ForwardDiff.Dual{Nothing, Float64, JCH}, JCH}}, Vector{ForwardDiff.Dual{Nothing, Float64, JCH}}}
    gradcache::FiniteDiff.GradientCache{Nothing, Nothing, Nothing, Vector{Float64}, Val{:forward}(), Float64, Val{true}()}
    hesscache::FiniteDiff.HessianCache{Vector{Float64}, Val{:hcentral}(), Val{true}()}
    function MyFun()
        n,m = 4,2
        gradcfg = ForwardDiff.GradientConfig(nothing, zeros(n+m))
        hesscfg = ForwardDiff.HessianConfig(nothing, zeros(n+m))
        gradcache = FiniteDiff.GradientCache(zeros(n+m), zeros(n+m), Val(:forward))
        hesscache = FiniteDiff.HessianCache(zeros(n+m), zeros(n+m), zeros(n+m), zeros(n+m), Val(:hcentral), Val(true))
        new{length(gradcfg.seeds)}(gradcfg, hesscfg, gradcache, hesscache)
    end
end
RD.state_dim(::MyFun) = 4
RD.control_dim(::MyFun) = 2

function RD.evaluate(::MyFun, x, u)
    x[1]^2 + cos(x[2] + 2u[2]) + u[1]*x[4] + x[3]^2
end

function RD.gradient!(::MyFun, grad, x, u)
    grad[1] = 2x[1]
    grad[2] = -sin(x[2] + 2u[2])
    grad[3] = 2x[3]
    grad[4] = u[1]
    grad[5] = x[4]
    grad[6] = -sin(x[2] + 2u[2]) * 2
    return nothing
end

function RD.hessian!(::MyFun, H, x, u)
    H .= 0
    H[1,1] = 2.0
    H[2,2] = -cos(x[2] + 2u[2])
    H[2,6] = -cos(x[2] + 2u[2]) * 2
    H[3,3] = 2
    H[4,5] = 1.0
    H[5,4] = 1.0
    H[6,2] = -cos(x[2] + 2u[2]) * 2
    H[6,6] = -cos(x[2] + 2u[2]) * 4
    return
end

function RD.gradient!(::ForwardAD, fun::MyFun, grad, z)
    f(_z) = RD.evaluate(fun, RD.getstate(z, _z), RD.getcontrol(z, _z), RD.getparams(z))
    ForwardDiff.gradient!(grad, f, RD.getdata(z), fun.gradcfg)
    return nothing
end

function RD.hessian!(::ForwardAD, fun::MyFun, hess, z, cache=fun.hesscfg)
    f(_z) = RD.evaluate(fun, RD.getstate(z, _z), RD.getcontrol(z, _z), RD.getparams(z))
    ForwardDiff.hessian!(hess, f, RD.getdata(z), fun.hesscfg)
    return nothing
end

function hessian(::StaticReturn, ::ForwardAD, fun::MyFun, hess, z, cache=fun.hesscfg)
    f(_z) = RD.evaluate(fun, RD.getstate(z, _z), RD.getcontrol(z, _z), RD.getparams(z))
    hess .= ForwardDiff.hessian(f, RD.getdata(z), fun.hesscfg)
end

function RD.gradient!(::FiniteDifference, fun::MyFun, grad, z)
    f(_z) = RD.evaluate(fun, RD.getstate(z, _z), RD.getcontrol(z, _z), RD.getparams(z))
    FiniteDiff.finite_difference_gradient!(grad, f, RD.getdata(z), fun.gradcache)
    return nothing
end

function RD.hessian!(::FiniteDifference, fun::MyFun, hess, z, cache=fun.hesscache)
    cache.xmm .= z
    cache.xmp .= z
    cache.xpm .= z
    cache.xpp .= z
    f(_z) = RD.evaluate(fun, RD.getstate(z, _z), RD.getcontrol(z, _z), RD.getparams(z))
    FiniteDiff.finite_difference_hessian!(hess, f, RD.getdata(z), cache)
    return nothing
end

function test_scalar_fun(fun)
    n,m = RD.dims(fun)
    x = @SVector rand(n)
    u = @SVector rand(m)
    t,h = 1.0, 0.1
    grad = zeros(n+m)
    grad0 = zeros(n+m)
    hess = zeros(n+m, n+m)
    hess0 = zeros(n+m, n+m)
    zs = RD.KnotPoint(x, u, t, h)
    z = RD.KnotPoint{n,m}(Vector(x), Vector(u), t, h)

    RD.gradient!(fun, grad0, x, u)
    RD.gradient!(ForwardAD(), fun, grad, z)
    @test grad ≈ grad0
    RD.gradient!(FiniteDifference(), fun, grad, z)
    @test grad ≈ grad0 atol = 1e-6
    RD.gradient!(ForwardAD(), fun, grad, zs)
    @test grad ≈ grad0
    RD.gradient!(FiniteDifference(), fun, grad, zs)
    @test grad ≈ grad0 atol = 1e-6

    if run_alloc_tests
        allocs = 0
        allocs += @allocated RD.gradient!(ForwardAD(), fun, grad, z)
        @test allocs == 0
        allocs += @allocated RD.gradient!(ForwardAD(), fun, grad, zs)
        @test allocs == 0
        allocs += @allocated RD.gradient!(FiniteDifference(), fun, grad, z)
        @test allocs == 0
        allocs += @allocated RD.gradient!(FiniteDifference(), fun, grad, zs)
        @test allocs == 0
    end

    RD.hessian!(fun, hess0, x, u)
    RD.hessian!(ForwardAD(), fun, hess, z)
    @test hess ≈ hess0
    RD.hessian!(FiniteDifference(), fun, hess, z)
    @test hess ≈ hess0 atol = 1e-6
    RD.hessian!(ForwardAD(), fun, hess, zs)
    @test hess ≈ hess0
    RD.hessian!(FiniteDifference(), fun, hess, zs)
    @test hess ≈ hess0 atol = 1e-6

    if run_alloc_tests
        allocs += @allocated RD.hessian!(RD.UserDefined(), fun, hess, z)
        # allocs += @allocated RD.hessian!(ForwardAD(), fun, hess, z)  # this has some un-avoidable allocations
        allocs += @allocated RD.hessian!(ForwardAD(), fun, hess, zs)
        allocs += @allocated RD.hessian!(FiniteDifference(), fun, hess, z)
        allocs += @allocated RD.hessian!(FiniteDifference(), fun, hess, zs)
        @test allocs == 0
    end
end

##
fun = MyFun()
test_scalar_fun(fun)

##
@macroexpand RD.@autodiff struct MyFunAuto <: RobotDynamics.ScalarFunction end
RD.@autodiff struct MyFunAuto <: RobotDynamics.ScalarFunction end
RD.state_dim(::MyFunAuto) = 4
RD.control_dim(::MyFunAuto) = 2

function RD.evaluate(::MyFunAuto, x, u)
    x[1]^2 + cos(x[2] + 2u[2]) + u[1]*x[4] + x[3]^2
end

function RD.gradient!(::MyFunAuto, grad, x, u)
    grad[1] = 2x[1]
    grad[2] = -sin(x[2] + 2u[2])
    grad[3] = 2x[3]
    grad[4] = u[1]
    grad[5] = x[4]
    grad[6] = -sin(x[2] + 2u[2]) * 2
    return nothing
end

function RD.hessian!(::MyFunAuto, H, x, u)
    H .= 0
    H[1,1] = 2.0
    H[2,2] = -cos(x[2] + 2u[2])
    H[2,6] = -cos(x[2] + 2u[2]) * 2
    H[3,3] = 2
    H[4,5] = 1.0
    H[5,4] = 1.0
    H[6,2] = -cos(x[2] + 2u[2]) * 2
    H[6,6] = -cos(x[2] + 2u[2]) * 4
    return
end

fun = MyFunAuto()
test_scalar_fun(fun)


text:
using CLEARSWI
using Test

@time begin
    @time @testset "Utils Tests" begin include("utility_test.jl") end
    @time @testset "Functions Test" begin include("functions_test.jl") end
    @time @testset "With FFTW Test" begin include("fftw_test.jl") end
end


text:
using OptimalTransport

using Distances
using ForwardDiff
using ReverseDiff
using LogExpFunctions
using PythonOT: PythonOT
using LinearAlgebra
using Random
using Test

const POT = PythonOT

Random.seed!(100)

@testset "sinkhorn_divergence.jl" begin
    @testset "fixed_support" begin
        # size of problem
        N = 250
        # number of target measures
        x = range(-1, 1; length=N)
        C = pairwise(SqEuclidean(), x)
        f(x; μ, σ) = exp(-((x - μ) / σ)^2)
        # regularization parameter
        ε = 0.05
        @testset "basic" begin
            μ = normalize!(f.(x; μ=0, σ=0.5), 1)
            M = 100

            ν_all = [normalize!(f.(x; μ=y, σ=0.5), 1) for y in range(-1, 1; length=M)]

            for reg in (true, false)
                loss = map(ν -> sinkhorn_divergence(μ, ν, C, ε; regularization=reg), ν_all)
                loss_ = map(
                    ν ->
                        sinkhorn2(μ, ν, C, ε; regularization=reg) -
                        (
                            sinkhorn2(μ, μ, C, ε; regularization=reg) +
                            sinkhorn2(ν, ν, C, ε; regularization=reg)
                        ) / 2,
                    ν_all,
                )

                @test loss ≈ loss_ rtol = 1e-6
                @test all(loss .≥ 0)
                @test sinkhorn_divergence(μ, μ, C, ε) ≈ 0 atol = 1e-9
            end
        end
        @testset "batch" begin
            M = 10
            μ = hcat([normalize!(f.(x; μ=randn(), σ=0.5), 1) for _ in 1:M]...)
            ν = hcat([normalize!(f.(x; μ=randn(), σ=0.5), 1) for _ in 1:M]...)
            for reg in (true, false)
                loss_batch = sinkhorn_divergence(μ, ν, C, ε; regularization=reg)
                @test loss_batch ≈ [
                    sinkhorn_divergence(μ[:, i], ν[:, i], C, ε; regularization=reg) for
                    i in 1:M
                ]
                loss_batch_μ = sinkhorn_divergence(μ, ν[:, 1], C, ε; regularization=reg)
                @test loss_batch_μ ≈ [
                    sinkhorn_divergence(μ[:, i], ν[:, 1], C, ε; regularization=reg) for
                    i in 1:M
                ]
                loss_batch_ν = sinkhorn_divergence(μ[:, 1], ν, C, ε; regularization=reg)
                @test loss_batch_ν ≈ [
                    sinkhorn_divergence(μ[:, 1], ν[:, i], C, ε; regularization=reg) for
                    i in 1:M
                ]
            end
        end
        @testset "AD" begin
            ε = 0.05
            μ = normalize!(f.(x; μ=-0.5, σ=0.5), 1)
            ν = normalize!(f.(x; μ=0.5, σ=0.5), 1)
            for Diff in [ForwardDiff, ReverseDiff]
                for reg in (true, false)
                    ∇ = Diff.gradient(log.(ν)) do xs
                        sinkhorn_divergence(μ, softmax(xs), C, ε; regularization=reg)
                    end
                    @test size(∇) == size(ν)
                    ∇ = Diff.gradient(log.(μ)) do xs
                        sinkhorn_divergence(μ, softmax(xs), C, ε; regularization=reg)
                    end
                    @test norm(∇, Inf) ≈ 0 atol = 1e-9 # Sinkhorn divergence has minimum at SD(μ, μ)
                end
            end
        end
    end
    @testset "empirical" begin
        N = 50
        M = 64
        d = 2
        μ_spt = randn(N, d)
        ν_spt = 1.5randn(M, d)
        μ = fill(1 / N, N)
        ν = fill(1 / M, M)
        Cμν = pairwise(SqEuclidean(), μ_spt', ν_spt'; dims=2)
        Cμ = pairwise(SqEuclidean(), μ_spt'; dims=2)
        Cν = pairwise(SqEuclidean(), ν_spt'; dims=2)
        ε = 1.0

        @testset "basic" begin
            for reg in (true, false)
                @test sinkhorn_divergence(μ, ν, Cμν, Cμ, Cν, ε; regularization=reg) ≥ 0
                @test sinkhorn_divergence(μ, μ, Cμ, Cμ, Cμ, ε; regularization=reg) ≈ 0 atol =
                    1e-6
            end
        end

        @testset "AD" begin
            for Diff in [ForwardDiff, ReverseDiff]
                for reg in (true, false)
                    ∇ = Diff.gradient(ν_spt) do xs
                        Cμν = pairwise(SqEuclidean(), μ_spt', xs'; dims=2)
                        Cν = pairwise(SqEuclidean(), xs'; dims=2)
                        sinkhorn_divergence(μ, ν, Cμν, Cμ, Cν, ε)
                    end
                    @test size(∇) == size(ν_spt)
                    ∇ = Diff.gradient(μ_spt) do xs
                        Cμν = pairwise(SqEuclidean(), μ_spt', xs'; dims=2)
                        Cν = pairwise(SqEuclidean(), xs'; dims=2)
                        sinkhorn_divergence(μ, μ, Cμν, Cμ, Cν, ε)
                    end
                    @test norm(∇, Inf) ≈ 0 atol = 1e-6
                end
            end
        end
    end
end


text:
function show_versions(res_name::AbstractString;
        start = 1, window = 10,
        cpreview_l = 200)
    res_name = extract_basename(res_name);
    versions = Core.find_versions(get_working_dir(), res_name);
    println("Found $(length(versions)) version(s) for $(res_name), showing from $(start)");
    println();
    c = 1;
    for (i, v) in enumerate(versions)
        i < start && continue;
        println(i, ": ", basename(v.res_dir), " tt: ", Core.strf_time_tag(v.time_tag));
        comment = Core.get_desc_comment(Core.load_res_desc(v));
        print("   ")
        length(comment) > cpreview_l && println(comment[1:cpreview_l], " ...");
        length(comment) <= cpreview_l && println(comment);
        println();
        c == window && return;
        c += 1;
    end
end

function show_repo(res, ver::Union{Void, Int, String} = nothing; log_lw = 20)
    repo = get_loadable_repo(res, ver);
    println(basename(repo.res_dir), " tt: ", Core.strf_time_tag(repo.time_tag));
    println();
    println("Comment ----------------------------------------------");
    println(Core.get_desc_comment(Core.load_res_desc(repo)));
    println();
    println("Log Frgment----------------------------------------------");
    println("(...)")
    log = split(Core.get_desc_log(Core.load_res_desc(repo)),"\n");
    length(log) <= log_lw && for logline in log println(logline); end
    length(log) > log_lw && for logline in log[end - log_lw:end] println(logline); end
end

function show_repos(; cpreview_l = 200)
    repos = Core.find_repos(get_working_dir());

    for (i, repo) in enumerate(repos)
        println(i, ": ", basename(repo.res_dir), " tt: ", Core.strf_time_tag(repo.time_tag));
        comment = Core.get_desc_comment(Core.load_res_desc(repo));
        print("   ")
        length(comment) > cpreview_l && println(comment[1:cpreview_l], " ...");
        length(comment) <= cpreview_l && println(comment);
        println();
    end
end


text:
# Process-local dictionary mapping from job IDs to instances of `Job`
global jobs = Dict()

# TODO: Allow for different threads to use different jobs by making this
# thread-local. For now, we only allow a single `Job` for each process
# and no sharing between threads; i.e., the burden is on the user to make
# sure they are synchronizing access to the `Job` if using the same one from
# different threads.
# TODO: Allow for different threads to use the same job by wrapping each
# `Job` in `jobs` in a mutex to allow only one to use it at a time. Further
# modifications would be required to make sharing a job between threads
# ergonomic.
global current_job_id = nothing

function set_job(job_id::Union{JobId,Nothing})
    global current_job_id
    current_job_id = job_id
end

function get_job_id()::JobId
    global current_job_id
    if isnothing(current_job_id)
        error(
            "No job selected using `create_job` or `with_job` or `set_job`. The current job may have been destroyed or no job created yet.",
        )
    end
    current_job_id
end

function get_job(job_id=get_job_id())::Job
    global jobs
    if !haskey(jobs, job_id)
        error("The selected job does not have any information; if it was created by this process, it has either failed or been destroyed.")
    end
    jobs[job_id]
end

get_cluster_name() = get_job().cluster_name

function create_job(;
    cluster_name::Union{String,Nothing} = nothing,
    nworkers::Union{Integer,Nothing} = 16,
    print_logs::Union{Bool,Nothing} = false,
    store_logs_in_s3::Union{Bool,Nothing} = true,
    store_logs_on_cluster::Union{Bool,Nothing} = false,
    sample_rate::Union{Integer,Nothing} = nworkers,
    job_name::Union{String,Nothing} = nothing,
    files::Union{Vector,Nothing} = [],
    code_files::Union{Vector,Nothing} = [],
    force_update_files::Union{Bool,Nothing} = false,
    pf_dispatch_table::Union{String,Nothing} = "",
    using_modules::Union{Vector,Nothing} = [],
    url::Union{String,Nothing} = nothing,
    branch::Union{String,Nothing} = nothing,
    directory::Union{String,Nothing} = nothing,
    dev_paths::Union{Vector,Nothing} = [],
    force_clone::Union{Bool,Nothing} = false,
    force_pull::Union{Bool,Nothing} = false,
    force_install::Union{Bool,Nothing} = false,
    nowait::Bool=false,
    kwargs...,
)
    global jobs
    global current_job_id

    if cluster_name == ""
        cluster_name = nothing
    end

    # Configure
    configure(; kwargs...)

    # Construct parameters for creating job
    cluster_name = if isnothing(cluster_name)
        running_clusters = get_running_clusters()
        if length(running_clusters) == 0
            error("Failed to create job: you don't have any clusters created")
        end
        first(keys(running_clusters))
    else
        cluster_name
    end

    julia_version = get_julia_version()

    job_configuration = Dict{String,Any}(
        "cluster_name" => cluster_name,
        "num_workers" => nworkers,
    	"return_logs" => print_logs,
	    "store_logs_in_s3" => store_logs_in_s3,
        "store_logs_on_cluster" => store_logs_on_cluster,
        "julia_version" => julia_version,
        "nowait" => nowait,
        "benchmark" => get(ENV, "BANYAN_BENCHMARK", "0") == "1",
        "main_modules" => get_loaded_packages(),
        "using_modules" => using_modules,
    )
    if !isnothing(job_name)
        job_configuration["job_name"] = job_name
    end

    s3_bucket_name = get_cluster_s3_bucket_name(cluster_name)

    environment_info = Dict{String,Any}()
    # If a url is not provided, then use the local environment
    if isnothing(url)
        # TODO: Optimize to not have to send tomls on every call
        local_environment_dir = get_julia_environment_dir()
        project_toml = load_file("file://$(local_environment_dir)Project.toml")
        if !isfile("$(local_environment_dir)Manifest.toml")
            manifest_toml = ""
            @warn "Manifest file not present for this environment"
        else
            manifest_toml = load_file("file://" * local_environment_dir * "Manifest.toml")
        end
        environment_hash = get_hash(project_toml * manifest_toml)
        environment_info["environment_hash"] = environment_hash
        environment_info["project_toml"] = "$(environment_hash)/Project.toml"
        if !isfile(S3Path("s3://$(s3_bucket_name)/$(environment_hash)/Project.toml", config=get_aws_config()))
            s3_put(get_aws_config(), s3_bucket_name, "$(environment_hash)/Project.toml", project_toml)
        end
        if manifest_toml != ""
            environment_info["manifest_toml"] = "$(environment_hash)/Manifest.toml"
            if !isfile(S3Path("s3://$(s3_bucket_name)/$(environment_hash)/Manifest.toml", config=get_aws_config()))
                s3_put(get_aws_config(), s3_bucket_name, "$(environment_hash)/Manifest.toml", manifest_toml)
            end
        end
    else
        # Otherwise, use url and optionally a particular branch
        environment_info["url"] = url
        if isnothing(directory)
            error("Directory must be provided for given URL $url")
        end
        environment_info["directory"] = directory
        if !isnothing(branch)
            environment_info["branch"] = branch
        end
        environment_info["dev_paths"] = dev_paths
        environment_info["force_clone"] = force_clone
        environment_info["force_pull"] = force_pull
        environment_info["force_install"] = force_install
        environment_info["environment_hash"] = get_hash(
            url * (if isnothing(branch) "" else branch end)
        )
    end
    job_configuration["environment_info"] = environment_info

    # Upload files to S3
    for f in vcat(files, code_files)
        s3_path = S3Path("s3://$(s3_bucket_name)/$(basename(f))", config=get_aws_config())
        if !isfile(s3_path) || force_update_files
            s3_put(get_aws_config(), s3_bucket_name, basename(f), load_file(f))
        end
    end
    # TODO: Optimize so that we only upload (and download onto cluster) the files if the filename doesn't already exist
    job_configuration["files"] = [basename(f) for f in files]
    job_configuration["code_files"] = [basename(f) for f in code_files]

    if pf_dispatch_table == ""
        pf_dispatch_table = "https://raw.githubusercontent.com/banyan-team/banyan-julia/v0.1.3/Banyan/res/pf_dispatch_table.json"
    end
    job_configuration["pf_dispatch_table"] = load_json(pf_dispatch_table)

    # Create the job
    @debug "Sending request for job creation"
    job_response = send_request_get_response(:create_job, job_configuration)
    sleep(30)
    job_id = job_response["job_id"]
    @info "Started creating job with ID $job_id on cluster named \"$cluster_name\""

    # Store in global state
    current_job_id = job_id
    jobs[current_job_id] = Job(cluster_name, current_job_id, nworkers, sample_rate)

    wait_for_cluster(cluster_name)

    if !nowait
        wait_for_job(job_id)
    end

    @debug "Finished creating job $job_id"
    job_id
end

function destroy_job(job_id::JobId = get_job_id(); failed = false, force = false, kwargs...)
    global jobs
    global current_job_id

    @info "Destroying job with ID $job_id"
    send_request_get_response(
        :destroy_job,
        Dict{String,Any}("job_id" => job_id, "failed" => failed),
    )

    # Remove from global state
    # TODO: Maybe `job_id`must always has be in `jobs`. In that case, we should
    # not be needing this `if` statement here.
    if !isnothing(current_job_id) && get_job_id() == job_id
        set_job(nothing)
    end
    delete!(jobs, job_id)
    job_id
end

function get_jobs(cluster_name = nothing; status = nothing, kwargs...)
    @debug "Downloading description of all jobs in cluster named $cluster_name"
    configure(; kwargs...)
    filters = Dict()
    if !isnothing(cluster_name)
        filters["cluster_name"] = cluster_name
    end
    if !isnothing(status)
        filters["status"] = status
    end

    finished = false
    indiv_response = send_request_get_response(:describe_jobs, Dict{String,Any}("filters"=>filters))
    response = indiv_response
    if  isnothing(indiv_response["last_eval"])
        finished = true
    else
        curr_last_eval = indiv_response["last_eval"]
        while !finished
            indiv_response = send_request_get_response(:describe_jobs, Dict{String,Any}("filters"=>filters, "this_start_key"=>curr_last_eval))
            response["jobs"] = merge!(response["jobs"], indiv_response["jobs"])
            if isnothing(indiv_response["last_eval"])
                finished = true
            else
                curr_last_eval = indiv_response["last_eval"]
            end
        end
    end

    for (id, j) in response["jobs"]
        if response["jobs"][id]["ended"] == ""
            response["jobs"][id]["ended"] = nothing
        else
            response["jobs"][id]["ended"] = parse_time(response["jobs"][id]["ended"])
        end
        response["jobs"][id]["created"] = parse_time(response["jobs"][id]["created"])
    end
    response["jobs"]
end

get_running_jobs(args...; kwargs...) = get_jobs(args...; status="running", kwargs...)

function download_job_logs(job_id::JobId, cluster_name::String, filename::String=nothing; kwargs...)
    @debug "Downloading logs for job"
    configure(; kwargs...)
    s3_bucket_name = get_cluster_s3_bucket_name(cluster_name)
    log_file_name = "banyan-log-for-job-$(job_id)"
    filename = !isnothing(filename) ? filename : joinpath(homedir(), ".banyan", "logs")
    s3_get_file(get_aws_config(), s3_bucket_name, log_file_name, filename)
    @info "Downloaded logs for job with ID $job_id to $filename"
end

function destroy_all_jobs(cluster_name::String; kwargs...)
    @info "Destroying all running jobs for cluster named $cluster_name"
    configure(; kwargs...)
    jobs = get_jobs(cluster_name, status=["creating", "running"])
    for (job_id, job) in jobs
        if job["status"] == "running"
            destroy_job(job_id, kwargs...)
        end
    end
end

function get_job_status(job_id::String=get_job_id(); kwargs...)
    configure(; kwargs...)
    filters = Dict("job_id" => job_id)
    response = send_request_get_response(:describe_jobs, Dict{String,Any}("filters"=>filters))
    job_status = response["jobs"][job_id]["status"]
    if job_status == "failed"
        # We don't immediately fail - we're just explaining. It's only later on
        # where it's like we're actually using this job do we set the status.
        @error response["jobs"][job_id]["status_explanation"]
    end
    job_status
end

function wait_for_job(job_id::JobId=get_job_id(), kwargs...)
    t = 5
    job_status = get_job_status(job_id; kwargs)
    p = ProgressUnknown("Preparing job with ID $job_id", spinner=true)
    while job_status == "creating"
        sleep(t)
        next!(p)
        if t < 80
            t *= 2
        end
        job_status = get_job_status(job_id; kwargs)
    end
    finish!(p, spinner = job_status == "running" ? '✓' : '✗')
    if job_status == "running"
        @debug "Job with ID $job_id is ready"
    elseif job_status == "completed"
        error("Job with ID $job_id has already completed")
    elseif job_status == "failed"
        error("Job with ID $job_id has failed")
    else
        error("Unknown job status $job_status")
    end
end

function with_job(f::Function; kwargs...)
    # This is not a constructor; this is just a function that ensures that
    # every job is always destroyed even in the case of an error
    use_existing_job = :job in keys(kwargs)
    destroy_job_on_error = get(kwargs, :destroy_job_on_error, true)
    destroy_job_on_exit = get(kwargs, :destroy_job_on_exit, true)
    j = use_existing_job ? kwargs[:job] : create_job(; kwargs...)
    destroyed = false # because of weird catch/finally stuff
    try
        set_job(j)
        f(j)
    catch
        # If there is an error we definitely destroy the job
        # TODO: Cache the job so that even if there is a failure we can still
        # reuse it
        if destroy_job_on_error
            destroy_job(j)
            destroyed = true
        end
        rethrow()
    finally
        # We only destroy the job if it hasn't already been destroyed because
        # of an error and if we don't intend to reuse a job
        if destroy_job_on_exit && !destroyed
            destroy_job(j)
        end
    end
end

# `create_job` creates a new job. `use_job` will create the job or reuse an
# existing job created by this process with the same configuration if possible.
# `with_job`

# Every job is owned by some process that created it. Until that process
# calls destroy_job, it maintains ownership.

# create_job should be optimized to try to re-use resources if possible. If a
# job was created by the current process with the same configuration, that job
# will be reused (unless force=true).

# In an interactive use-case, re-running a notebook and creating a job should
# probably anyway create a new job since it is running again. And the way a
# notebook is used is that you only re-run the cell that you edit. So you will
# keep re-using the job. If anything, we should have an option to automatically
# destroy the job after some time or make it so that if a job is used but it is
# destroyed, then the job should be reinstated. Actually, never mind. We
# shouldn't reinstate the job. If something fails and it is a hard failure, the
# job should be restarted. The only thing we should do for that is make sure
# that recoverable exceptions on the backend are propagated while errors that
# cause the job to crash result in the job being ended. Then we should reuse
# destroyed jobs that haven't crashed.

max_stars_repo_name: banyan-team/banyan-julia

max_issues_repo_name: banyan-team/banyan-julia

max_forks_repo_name: banyan-team/banyan-julia


text:
module Docker
    using Requests
    using JSON

    immutable DockerError
	status::Int
	msg::ByteString
    end

    const headers = Dict{}("Content-Type" => "application/json")

    docker_uri(host) = URI("http://$host/v1.21")
    docker_uri(host,endpoint) = URI("http://$host/v1.21/$endpoint")
    parse(data) = JSON.parse(join(map(Char,data)))

    function create_container(host, image;
                            cmd::Cmd     = ``,
                            entryPoint   = "",
                            tty          = true,
                            attachStdin  = false,
                            openStdin    = false,
                            attachStdout = true,
                            attachStderr = true,
                            memory       = 0,
                            cpuSets      = "",
                            volumeDriver = "",
                            portBindings = ["",""], # [ContainerPort,HostPort]
                            ports        = [],
                            pwd          = "")

		url = docker_uri(host)

		params = Dict{}("Image" => image,
						"Cmd" => collect(cmd.exec),
						"Tty" => tty,
						"AttachStdin" 	=> attachStdin,
						"OpenStdin" 	=> openStdin,
						"AttachStdout" 	=> attachStdout,
						"AttachStderr" 	=> attachStderr,
						"ExposedPorts" 	=> [string(dec(p),"/tcp")=>Dict{}() for p in ports],
						"HostConfig" 	=> Dict{}(
													"Memory"       => memory,
													"CpusetCpus"   => cpuSets,
													"VolumeDriver" => volumeDriver,
													"PortBindings" => Dict{}( string(portBindings[1],"/tcp") => [Dict{}( "HostPort" => string(portBindings[2]))]
																			)
												)
						)

		if !isempty(entryPoint)
			params["Entrypoint"] = entryPoint
		end

		if !isempty(cmd.exec)
			params["Cmd"] = cmd
		end

		if !isempty(pwd)
			params["WorkingDir"] = pwd
		end

		resp = post(URI("$url/containers/create"),json=params,headers=headers)
		if resp.status != 201
			throw(DockerError(resp.status,resp.data))
		end
		parse(resp.data)
	end

	function inspect_container(host,id)
		resp = get(docker_uri(host,"containers/$id/json"))
		if resp.status != 200
			throw(DockerError(resp.status,resp.data))
		end
		parse(resp.data)
	end

	function start_container(host, id)
		resp = post(docker_uri(host,"containers/$id/start"))
		if resp.status != 204
			throw(DockerError(resp.status,resp.data))
		end
		id
	end

	function restart_container(host, id)
		resp = post(docker_uri(host,"containers/$id/restart"))
		if resp.status != 204
			throw(DockerError(resp.status,resp.data))
		end
		resp
	end

	function stop_container(host, id)
		resp = post(docker_uri(host,"containers/$id/stop"))
		if resp.status != 204
			throw(DockerError(resp.status,resp.data))
		end
		resp
	end

	function pause_container(host, id)
		resp = post(docker_uri(host,"containers/$id/pause"))
		if resp.status != 204
			throw(DockerError(resp.status,resp.data))
		end
		resp
	end

	function unpause_container(host, id)
		resp = post(docker_uri(host,"containers/$id/unpause"))
		if resp.status != 204
			throw(DockerError(resp.status,resp.data))
		end
		id
	end

	function kill_container(host, id)
		resp = post(docker_uri(host,"containers/$id/kill"),"")
		if resp.status != 204
			throw(DockerError(resp.status,resp.data))
		end
		resp
	end

	function remove_container(host, id)
		resp = Requests.delete(docker_uri(host,"containers/$id?force=1"))
		if resp.status != 204
			throw(DockerError(resp.status,resp.data))
		end
		resp
	end

	function processes_container(host, id)
		resp = get(docker_uri(host,"containers/$id/top"))
		println(resp.status)
		if resp.status != 200
			throw(DockerError(resp.status,resp.data))
		end
		parse(resp.data)
	end

	function list_containers(host)
		resp = get(docker_uri(host,"containers/json"))
		if resp.status != 200
			throw(DockerError(resp.status,resp.data))
		end
		parse(resp.data)
	end

	function stats_container(host,id)
		resp = get(docker_uri(host,"containers/$id/stats?stream=0"))
		if resp.status != 200
			throw(DockerError(resp.status,resp.data))
		end
		parse(resp.data)
	end

	function open_logs_stream(host, id; history=false)
		path = "containers/$id/attach?logs&follow=1&stdout=1"
		if history
			path *=  "&logs=1"
		end
		url = docker_uri(host,path)
		Requests.open_stream(url,[Dict{}("Content-Type"=>"plain/text")],"","POST")
	end

	function cleanse!(host)
		resp = get(docker_uri(host,"containers/json?all=true"))
		if resp.status != 200
			throw(DockerError(resp.status,resp.data))
		end
		data = parse(resp.data)
		for c in data
			remove_container(host,c["Id"])
			kill_container(host,c["Id"])
		end
		nothing
	end

end


text:
function executaCalculoCabidas(dcp, dcn, dca, dcc, dcu, dcf, dcr, fpe, conjuntoTemplates)

    numTiposDepto = length(dcc.SUPDEPTOUTIL);

    # Posiciona el origen del predio (esq. inferior izquierda) en el punto 0,0. Obtiene estructura con las matrices de rotación
    V_predio, R = infoPredio(dcp.x, dcp.y);

    # Calcula matriz V_areaEdif asociada a los vértices del area de edificación
    V_areaEdif = copy(V_predio);
    numLadosPredio = size(V_areaEdif, 1);
    conjuntoLados = 1:numLadosPredio;
    conjuntoLadosCalle = dcp.ladosConCalle;
    conjuntoLadosVecinos = setdiff(conjuntoLados, conjuntoLadosCalle);
    anchoLado = dca.ANCHOMAX
    sepVecinos = dcn.SEPMIN;
    rasante = dcn.RASANTE;
    alturaMax = dcn.ALTURAMAX
    anchoEspacioPublico = dcp.ANCHOESPACIOPUBLICO
    antejardin = dcn.ANTEJARDIN
    ps_areaEdif = PolyShape([V_predio], 1)
    ps_areaEdif = polyShape.polyExpandSides_v2(ps_areaEdif, -ones(length(conjuntoLadosVecinos),1)*sepVecinos, conjuntoLadosVecinos)
    ps_areaEdif = polyShape.polyExpandSides_v2(ps_areaEdif, -ones(length(conjuntoLadosCalle),1)*antejardin, conjuntoLadosCalle)
    V_areaEdif = ps_areaEdif.Vertices[1]

    ps_predio = PolyShape([V_predio], 1)
    superficieTerreno = polyShape.polyArea_v2(ps_predio);
    ps_bruto, ps_publico = generaSupBruta(ps_predio, conjuntoLadosCalle, anchoEspacioPublico)
    ps_calles = generaCalles(ps_predio, ps_publico, anchoEspacioPublico)
    V_bruto = ps_bruto.Vertices[1]
    V_publico = ps_publico.Vertices[1]

    # global altCorteVecinos, areaSombra_p, areaSombra_o, areaSombra_s
    ps_bruto = PolyShape([V_bruto], 1)
    superficieTerrenoBruta = polyShape.polyArea_v2(ps_bruto);
    ps_publico = PolyShape([V_publico], 1)
    ps_areaEdif = PolyShape([V_areaEdif], 1)

    # Calcula el volumen y sombra teórica
    rasante_ss = dcn.RASANTE;
    matConexionVertices_ss, vecVertices_ss, ps_volteor = generaVol3d_v4(V_predio, V_bruto, rasante_ss, dcn, dcp)
    V_volteor = ps_volteor.Vertices[1]
    ps_SombraVolTeor_p, ps_SombraVolTeor_o, ps_SombraVolTeor_s = generaSombraTeor_v3(ps_volteor, matConexionVertices_ss, vecVertices_ss, ps_publico, ps_calles)
    rasante_cs = 5;
    matConexionVertices_cs, vecVertices_cs, ps_restSombra = generaVol3d_v4(V_predio, V_bruto, rasante_cs, dcn, dcp)
    V_restSombra = ps_restSombra.Vertices[1]

    areaSombra_p = polyShape.polyArea_v2(ps_SombraVolTeor_p)
    areaSombra_o = polyShape.polyArea_v2(ps_SombraVolTeor_o)
    areaSombra_s = polyShape.polyArea_v2(ps_SombraVolTeor_s)

    vecAlturas_ss = sort(unique(V_volteor[:,end]))
    vecAlturas_cs = sort(unique(V_restSombra[:,end]))
    altCorteVecinos = dcn.RASANTE * dcn.SEPMIN
    numVertices = size(V_areaEdif, 1);

    function fitness_ss(x)  # Función de Fitness Sin Sombra
        alt, areaBasal, ps_base, ps_baseSeparada, psCorte = resultConverter_v2(x, V_restSombra, anchoLado, matConexionVertices_cs, vecVertices_cs, vecAlturas_cs)

        total_fit = 0

        ps_r = polyShape.polyDifference_v3(ps_base, psCorte) #Sector de la base del edificio que sobrepasa el areaEdif
        area_r = polyShape.polyArea_v2(ps_r) #Area del sector que sobrepasa
        penalizacion_r = area_r^1.1
        penalizacionCoefOcup = max(0, areaBasal - dcn.COEFOCUPACION * superficieTerreno)

        total_fit = alt*(areaBasal - 5*(penalizacionCoefOcup + penalizacion_r))

        return -total_fit
    end

    function fitness_cs(x)  # Función de Fitness Con Sombra

        alt, areaBasal, ps_base, ps_baseSeparada, psCorte = resultConverter_v2(x, V_restSombra, anchoLado, matConexionVertices_cs, vecVertices_cs, vecAlturas_cs)

        total_fit = 0
        ps_sombraEdif_p, ps_sombraEdif_o, ps_sombraEdif_s = generaSombraEdificio(ps_baseSeparada, alt, ps_publico, ps_calles)

        areaSombraEdif_p = polyShape.polyArea_v2(ps_sombraEdif_p)
        areaSombraEdif_o = polyShape.polyArea_v2(ps_sombraEdif_o)
        areaSombraEdif_s = polyShape.polyArea_v2(ps_sombraEdif_s)
        penalizacionSombra_p = max(0, areaSombraEdif_p - areaSombra_p)
        penalizacionSombra_o = max(0, areaSombraEdif_o - areaSombra_o)
        penalizacionSombra_s = max(0, areaSombraEdif_s - areaSombra_s)

        ps_r = polyShape.polyDifference_v3(ps_base, psCorte) #Sector de la base del edificio que sobrepasa el areaEdif
        area_r = polyShape.polyArea_v2(ps_r) #Area del sector que sobrepasa
        penalizacion_r = area_r^1.1

        penalizacionCoefOcup = max(0, areaBasal - dcn.COEFOCUPACION * superficieTerreno)

        #V_ch = poly2D.convHull(ps_base.Vertices[1])
        #area_ch = poly2D.polyArea(V_ch)

        total_fit = alt*(areaBasal - 5*(penalizacion_r + penalizacionCoefOcup + penalizacionSombra_p + penalizacionSombra_o + penalizacionSombra_s))

        return -total_fit
    end

    numParticles = 4000# 2000;
    maxIterations = 300# 300;

    resultados = fill(ResultadoCabida(nothing, nothing, nothing, nothing, nothing, nothing, []), length(conjuntoTemplates), 1)

    @time begin
        cont = 0

        t = conjuntoTemplates[1] # for t in conjuntoTemplates

        cont += 1

        display("Inicio de Cálculo")
        min_alt = min(maximum(vecAlturas_ss), dcn.RASANTE * dcn.SEPMIN)
        max_alt = maximum(vecAlturas_ss)
        min_theta = -pi;
        max_theta =  pi;

        min_largo = anchoLado;
        max_largo = 100;
        min_largo1 = anchoLado;
        max_largo1 = 100;
        min_largo2 = anchoLado;
        max_largo2 = 100;

        xmin = minimum(V_areaEdif[:,1]);  xmax = maximum(V_areaEdif[:,1]);
        ymin = minimum(V_areaEdif[:,2]);  ymax = maximum(V_areaEdif[:,2]);

        if t == 1
            min_alfa = 0;
            max_alfa = pi / 2;

            lb = [min_alt, min_theta, min_alfa, xmin, ymin, min_largo1, min_largo2,t];
            ub = [max_alt, max_theta, max_alfa, xmax, ymax, max_largo1, max_largo2,t];

            xopt_cs, fopt_cs = evol(fitness_cs, lb, ub, numParticles, maxIterations, false)

        elseif t == 2
            largos, angulosExt, angulosInt, largosDiag =  polyShape.extraeInfoPoly(ps_areaEdif)
            maxDiagonal = maximum(largosDiag)

            min_phi1 = 0; max_phi1 =  pi / 2;
            min_phi2 = 0; max_phi2 =  pi / 2;
            min_largo0 = 3 * anchoLado; max_largo0 = maxDiagonal

            lb = [min_alt, min_theta, min_phi1, min_phi2, xmin, ymin, min_largo0, min_largo1, min_largo2, t];
            ub = [max_alt, max_theta, max_phi1, max_phi2, xmax, ymax, max_largo0, max_largo1, max_largo2, t];

            xopt_cs, fopt_cs = evol(fitness_cs, lb, ub, numParticles, maxIterations, false)

        elseif t == 3

        elseif t == 4
            largos, angulosExt, angulosInt, largosDiag =  polyShape.extraeInfoPoly(ps_areaEdif)
            maxDiagonal = maximum(largosDiag)

            min_alfa = 0;
            max_alfa =  pi / 2;

            lb = [min_alt, min_theta, min_alfa, xmin, ymin, min_largo1, min_largo2, t];
            ub = [max_alt, max_theta, max_alfa, xmax, ymax, max_largo1, max_largo2, t];

            xopt_cs, fopt_cs = evol(fitness_cs, lb, ub, numParticles, maxIterations, false)

        end

        alt, areaBasal, ps_base, ps_baseSeparada, psCorte = resultConverter_v2(xopt_cs, V_restSombra, anchoLado, matConexionVertices_cs, vecVertices_cs, vecAlturas_cs)
        numPisos = Int(floor(alt / dca.ALTURAPISO))
        alturaEdif = numPisos * dca.ALTURAPISO
        sn, sa, si, st, sm, sf = optiEdificio(dcn, dca, dcp, dcc, dcu, dcf, dcr, alturaEdif, ps_base, superficieTerreno, superficieTerrenoBruta)
        xopt_cs[1] = sa.altura
        numPisos = sa.numPisos
        resultados[cont] = ResultadoCabida(sn, sa, si, st, sm, sf, [xopt_cs])
        resultados_ = [sn, sa, si, st, sm, sf, [xopt_cs]]

        ps_sombraEdif_p, ps_sombraEdif_o, ps_sombraEdif_s = generaSombraEdificio(ps_baseSeparada, alt, ps_publico, ps_calles)

        #fig = plotBaseEdificio3d_v2(fpe, resultados[cont].Xopt[1], dca.ALTURAPISO, ps_predio,
        #                            ps_volteor, matConexionVertices_ss, vecVertices_ss,
        #                            ps_restSombra, matConexionVertices_cs, vecVertices_cs,
        #                            ps_publico, ps_calles, ps_base, ps_baseSeparada);

        displayResults(resultados_)
        println(" ")
        println(" ")
        println(" ")

    end

    return resultados, ps_predio, ps_base, xopt_cs, fopt_cs, psCorte,
            ps_SombraVolTeor_p, ps_sombraEdif_p, ps_SombraVolTeor_s, ps_sombraEdif_s, ps_SombraVolTeor_o, ps_sombraEdif_o
end


text:
module HeatVectorEquationTests

using Gridap
using ForwardDiff
using LinearAlgebra
using Test
using GridapODEs.ODETools
using GridapODEs.TransientFETools
using Gridap.FESpaces: get_algebraic_operator

import Gridap: ∇
import GridapODEs.TransientFETools: ∂t

θ = 0.5

u(x,t) = VectorValue(x[1],x[2])*t
u(t::Real) = x -> u(x,t)

f(t) = x -> ∂t(u)(x,t)-Δ(u(t))(x)

domain = (0,1,0,1)
partition = (1,1)
model = CartesianDiscreteModel(domain,partition)

order = 2

reffe = ReferenceFE(lagrangian,VectorValue{2,Float64},order)
V0 = FESpace(
  model,
  reffe,
  conformity=:H1,
  dirichlet_tags="boundary"
)

U = TransientTrialFESpace(V0,u)

Ω = Triangulation(model)
degree = 2*order
dΩ = Measure(Ω,degree)

a(u,v) = ∫(∇(v)⊙∇(u))dΩ
m(u,v) = ∫(u⋅v)dΩ
b(v,t) = ∫(v⋅f(t))dΩ

res(t,u,ut,v) = a(u,v) + m(ut,v) - b(v,t)
jac(t,u,ut,du,v) = a(du,v)
jac_t(t,u,ut,dut,v) = m(dut,v)

op = TransientFEOperator(res,jac,jac_t,U,V0)

t0 = 0.0
tF = 1.0
dt = 0.1

U0 = U(0.0)
uh0 = interpolate_everywhere(u(0.0),U0)

ls = LUSolver()
odes = ThetaMethod(ls,dt,θ)
solver = TransientFESolver(odes)

sol_t = solve(solver,op,uh0,t0,tF)

l2(w) = w⋅w

tol = 1.0e-6
_t_n = t0

rv, _ = Base.iterate(sol_t)
uh_tn, tn = rv
uh_tn.free_values

_t_n = t0
for (uh_tn, tn) in sol_t
  global _t_n
  _t_n += dt
  @test tn≈_t_n
  e = u(tn) - uh_tn
  el2 = sqrt(sum( ∫(l2(e))dΩ ))
  @test el2 < tol
end

u0 = get_free_values(uh0)
uf = copy(get_free_values(uh0))

odeop = get_algebraic_operator(op)

ode_cache = allocate_cache(odeop)
vθ = similar(u0)
nl_cache = nothing

odes.θ == 0.0 ? dtθ = dt : dtθ = dt*odes.θ
tθ = t0+dtθ
ode_cache = update_cache!(ode_cache,odeop,tθ)

using GridapODEs.ODETools: ThetaMethodNonlinearOperator
nlop = ThetaMethodNonlinearOperator(odeop,tθ,dtθ,u0,ode_cache,vθ)

nl_cache = solve!(uf,odes.nls,nlop,nl_cache)
uf

K = nl_cache.A
h = nl_cache.b

# Steady version of the problem to extract the Laplacian and mass matrices
tf = tθ
Utf = U(tf)
fst(x) = f(tf)(x)
a(u,v) = ∫(∇(v)⊙∇(u))dΩ

function extract_matrix_vector(a,fst)
  btf(v) = ∫(v⋅fst)dΩ
  op = AffineFEOperator(a,btf,Utf,V0)
  ls = LUSolver()
  solver = LinearFESolver(ls)
  uh = solve(solver,op)

  tol = 1.0e-6
  e = uh-u(tf)
  l2(e) = e⋅e
  l2e = sqrt(sum( ∫(l2(e))dΩ ))
  # @test l2e < tol

  Ast = op.op.matrix
  bst = op.op.vector

  @test uh.free_values ≈ Ast \ bst

  return Ast, bst
end

A,rhs = extract_matrix_vector(a,fst)

gst(x) = -1.0*u(0.0)(x)
m(u,v) = (1/(θ*dt))*∫(u⋅v)dΩ

M,rhs2 = extract_matrix_vector(m,gst)

@test rhs + rhs2 ≈ h
@test A+M ≈ K
@test K \ h ≈ uf
@test K \ h ≈ uf

end #module


text:
"""
```
augment_states_with_shocks(regime_indices, TTTs, RRRs, CCCs, QQs, ZZs, z0, P0)
```

Go from original state space:

```
z_{t+1} = CCC + TTT*z_t + RRR*ϵ_t    (transition equation)
y_t     = DD  + ZZ*z_t  + η_t        (measurement equation)

ϵ_t ∼ N(0, QQ)
η_t ∼ N(0, EE)
Cov(ϵ_t, η_t) = 0
```

to augmented state space:

```
|z_{t+1}| = |CCC| + |TTT 0| |z_t| + |RRR| ϵ_t    (transition equation)
|ϵ_{t+1}|   | 0 |   | 0  0| |ϵ_t|   | I |

y_t = DD + |ZZ| |z_t| + η_t                      (measurement equation)
           | 0| |ϵ_t|

ϵ_t ∼ N(0, QQ)
η_t ∼ N(0, EE)
Cov(ϵ_t, η_t) = 0
```

with initial state and covariance:

```
|z0| and |P0  0|
| 0|     | 0 QQ|
```

Returns the augmented `TTTs`, `RRRs`, CCCs`, `ZZs`, `z0`, and `P0`.
"""
function augment_states_with_shocks(regime_indices::Vector{UnitRange{Int64}},
    TTTs::Vector{Matrix{S}}, RRRs::Vector{Matrix{S}}, CCCs::Vector{Vector{S}},
    QQs::Vector{Matrix{S}}, ZZs::Vector{Matrix{S}}, z0::Vector{S}, P0::Matrix{S}) where S <: AbstractFloat

    n_regimes = length(regime_indices)

    # Dimensions
    Nz = size(TTTs[1], 1) # number of states
    Ne = size(RRRs[1], 2) # number of shocks
    Ny = size(ZZs[1],  1) # number of observables

    # Augment initial conditions
    z0 = vcat(z0, zeros(Ne))
    P0 = vcat(hcat(P0, zeros(Nz, Ne)), hcat(zeros(Ne, Nz), QQs[1]))

    for i = 1:n_regimes
        # Get state-space system matrices for this regime
        TTT, RRR, CCC = TTTs[i], RRRs[i], CCCs[i]
        QQ,  ZZ       = QQs[i],  ZZs[i]

        # Augment regime-specific matrices
        TTTs[i] = vcat(hcat(TTT, zeros(Nz, Ne)), hcat(zeros(Ne, Nz), zeros(Ne, Ne)))
        RRRs[i] = vcat(RRR, Matrix{Float64}(I, Ne, Ne))
        CCCs[i] = vcat(CCC, zeros(Ne))
        ZZs[i]  = hcat(ZZ, zeros(Ny, Ne))
    end

    return TTTs, RRRs, CCCs, ZZs, z0, P0
end


text:
using Utils
using Chain

function parse_file(f, combiner)
    map(split_blocks(f)) do block
        @chain block begin
            split
            map(Set, _)
            reduce(combiner, _)
        end
    end
end

function load(combiner)
    open(joinpath(@__DIR__, "06_customs-declarations.txt"), "r") do f
        parse_file(f, combiner)
    end
end

function sumcount(sets)
    sum(length, sets)
end

println("Misunderstood sum of counts is: ", union |> load |> sumcount)
println("Correct       sum of counts is: ", intersect |> load |> sumcount)

using Test

@testset "Custom Customs" begin
    input = """
    abc

    a
    b
    c

    ab
    ac

    a
    a
    a
    a

    b
    """

    @testset "example 1" begin
        example = parse_file(IOBuffer(input), union)

        @test length(example) == 5
        @test example[1] == example[2] == example[3]
        @test sumcount(example) == 11
    end

    @testset "example 2" begin
        example = parse_file(IOBuffer(input), intersect)

        @test length(example) == 5
        @test sumcount(example) == 6
    end

    @testset "results" begin
        @test union |> load |> sumcount == 6947
        @test intersect |> load |> sumcount == 3398
    end
end

max_stars_repo_name: CmdQ/AoC