text:
function verifymeasure(;
    measure::DataFrame,
    measureinfo::NamedTuple,
    lpfilterperiods::Vector{<:Union{Day, Nothing}} = PARAM[:analysislpfilterperiods],
    benchmarkfunds::Bool = PARAM[:analysisbenchmarkfunds],
    benchmarkcomomentum::Bool=PARAM[:analysisbenchmarkcomomentum],
    comomcols::Union{Vector{Symbol}, Nothing} = PARAM[:analysiscomomcols],
    hfrname::String = PARAM[:hfrname],
    mfname::String = PARAM[:mfname],
    plname::String = PARAM[:plname],
    weeklyaggregatename = PARAM[:analysisweeklyaggregatename]
  )

  @unpack inputpath, outputpath, resultid, returntype, datefrequency = measureinfo

  hf = CSV.File("$inputpath\\$(returntype)_$(hfrname).csv") |> DataFrame
  mf = CSV.File("$inputpath\\$(returntype)_$(mfname).csv") |> DataFrame

    #compare the measure to particular benchmarks
    if datefrequency === :month
      validatedates(measure.date, frequency=:month, validateexactfrequency=true)

      if benchmarkfunds
        hfresults = comparetofundassets(measure, measureinfo, hf;
            fundslabel=:hf,
            datefrequency=:month,
            measurecols = findcols(measure, :loggrowth),
            benchmarkcols = findcols(hf, :loggrowth),
            usefocalcolheuristics=true,
            matchonlpfilter=true)

        mfresults = comparetofundassets(measure, measureinfo, mf;
            fundslabel=:mf,
            datefrequency=:month,
            measurecols = findcols(measure, :loggrowth),
            benchmarkcols = findcols(mf, :loggrowth),
            usefocalcolheuristics=true,
            matchonlpfilter=true)
        fundresults = [hfresults.results; mfresults.results;]
        sort!(fundresults, :cor)
        hfresults.data |> CSV.write("$outputpath\\$(resultid)_$(hfrname)_series.csv")
        mfresults.data |> CSV.write("$outputpath\\$(resultid)_$(mfname)_series.csv")
        fundresults |> CSV.write("$outputpath\\$(resultid)_fund_corresults.csv")

    elseif datefrequency === :week
      validatedates(measure.date, frequency=:week, validateexactfrequency=true)
      benchmarkfunds && throw("Only monthly data is supported when comparetofunds=true")
    end

    else

      @info "Note: fund benchmark analysis not run as analysisbenchmarkfunds=false"

    end

    if benchmarkcomomentum
      throw("Momentum not set up!!! Need to re-run everything before disabling this error.
        To set up, re-run/refresh momentum soup-2-nuts, then disable this error
        (can also disable the load warning in the analyze file)")

      comom = CSV.File("$inputpath\\$(returntype)_$(plname).csv") |> DataFrame

      benchmarkcols = isnothing(comomcols) ? setdiff(propertynames(comom), [:date]) : comomcols
      measurecols = [findcols(measure, :loggrowth); findcols(measure, :loggrowthpressure);
        findcols(measure, :assetmarketcap); findcols(measure, :loggrowthvolumepressure);
        findcols(measure, :growthvolumepressure)]

      println("benchmarkcols: $benchmarkcols")
      measurecols = [measurecols; measurecols .|> differencecol]
      comomresults = comparetofundassets(measure, measureinfo, comom,
        fundslabel = :pl,
        matchonlpfilter=false,
        usefocalcolheuristics=false;
        datefrequency, benchmarkcols, measurecols)
      sort!(comomresults.results, :cor)
      comomresults.data |> CSV.write("$outputpath\\$(resultid)_$(plname)_series.csv")
      comomresults.results |> CSV.write("$outputpath\\$(resultid)_comom_corresults.csv")
    else
      @info "Note: comomentum not run as analysisbenchmarkcomomentum=false"
    end

end

#useful function to translate monthyl dates between dataframes
function conformmonthlydates(old::AbstractVector{Date}, valid::AbstractVector{Date};
  datejoin = leftjoin)

  validatedates(old, frequency=:month, validateexactfrequency=true)
  validatedates(valid, frequency=:month, validateexactfrequency=true)

  #form a lookup table out of the available dates
  olddf = DataFrame(idx=1:length(old) |> collect, old=old)
  olddf.year = year.(old)
  olddf.month = month.(old)

  validdf = DataFrame(valid=valid)
  validdf.year = year.(valid)
  validdf.month = month.(valid)
  @assert sum(nonunique(validdf, [:year, :month])) == 0

  joint = datejoin(olddf, validdf, on=[:year, :month], validate=(false, true))
  @assert sum(nonunique(joint, [:year, :month])) == 0
  sort!(joint, :idx)

  return joint.valid
end

#useful function to translate weekly dates between dataframes
function conformweeklydates(old::AbstractVector{Date}, valid::AbstractVector{Date};
  datejoin = leftjoin)

  validatedates(old, frequency=:week, validateexactfrequency=true)
  validatedates(valid, frequency=:week, validateexactfrequency=true)

  olddf = DataFrame(idx=1:length(old) |> collect, old=old, lastdow = lastdayofweek.(old))
  validdf = DataFrame(valid=valid, lastdow = lastdayofweek.(valid))
  @assert sum(nonunique(validdf, [:lastdow])) == 0

  joint = datejoin(olddf, validdf, on=:lastdow, validate=(false, true))
  @assert sum(nonunique(joint, :lastdow)) == 0
  sort!(joint, :idx)

  return joint.valid
end

#compute correlations and statistics
function comparetofundassets(measure, measureinfo, funds;
    fundslabel::Symbol = throw("results label is required"),
    datefrequency::Symbol = throw("datefrequency is required"),
    measurecols = setdiff(propertynames(measure), [:date]),
    matchonlpfilter=true,
    benchmarkcols = setdiff(propertynames(funds), [:date]), #benchmarkcols
    extrabetaweights = PARAM[:fundextrabetaweights],
    extrabetaweightsidx = PARAM[:fundextrabetaweightsidx],
    minσ = eps(Float64)^0.5,
    usefocalcolheuristics = true
    )

  function bsstatfunc(sim)
    b,m = sim[1], sim[2]

    ρ = cor(b,m)
    return Dict(
      :cor  => ρ,
      :kend =>corkendall(b,m),
      :spear =>corspearman(b,m),
      :betabm => ρ * std(b)/std(m),
      :betamb => ρ * std(m)/std(b))
  end

  nullar1statfunc(b,m) = bsstatfunc((b,m))

  @unpack (neweylags, nwΣ,
    bsaggfunc, bsmeasurecolsregex, bsbenchmarkcolsregex, bssamples, bswidth,
      nullar1aggfunc, nullar1simulations) = measureinfo

  if datefrequency ≡ :month
    funds.measuredate = conformmonthlydates(funds.date, measure.date)
  elseif datefrequency ≡ :week
    funds.measuredate = conformweeklydates(funds.date, measure.date)
  else
    throw("unrecognized date frequency")
  end

  funds = funds[completecases(funds, :measuredate),:]
  select!(funds, Not(:date))
  joint = innerjoin(measure, funds,
    on=:date=>:measuredate, validate=(true,true))
  sort!(joint, :date)

  #compute the statistics pairwise in a semi-long format
  results = DataFrame(
    label=Symbol[], #label in case we concatenate this df with others
    measure=Symbol[], #measure field that we are benchmarking
    benchmark=Symbol[], #benchmark field

    cor=MFloat64[],
    cort=MFloat64[],
    corhac=MFloat64[],
    cormbbz = MFloat64[],
    cormbbp = MFloat64[],
    corar1h0z = MFloat64[],
    corar1h0p = MFloat64[],

    spear=MFloat64[],
    speart=MFloat64[],
    spearmbbz=MFloat64[],
    spearmbbp=MFloat64[],
    spearar1h0z=MFloat64[],
    spearar1h0p=MFloat64[],
    #spearz=MFloat64[],
    kend=MFloat64[],
    kendz=MFloat64[],
    kendmbbz=MFloat64[],
    kendmbbp=MFloat64[],
    kendar1h0z=MFloat64[],
    kendar1h0p=MFloat64[],

    betabm=MFloat64[],
    betabmmw=MFloat64[],
    betabmhac=MFloat64[],
    betabmmbbp=MFloat64[],
    betabmmbbz=MFloat64[],
    betabmar1h0z=MFloat64[],
    betabmar1h0p=MFloat64[],

    betamb=MFloat64[],
    betambmw=MFloat64[],
    betambhac=MFloat64[],
    betambmbbp=MFloat64[],
    betambmbbz=MFloat64[],
    betambar1h0z=MFloat64[],
    betambar1h0p=MFloat64[],
    N=MInt[])

  resultsbythread = [deepcopy(results) for t ∈ 1:Threads.nthreads()]

  #this subfunction determines which measure cols will be tested against a bcol
  function formfocalcols(bcol)
    matchedcols = matchonlpfilter ? matchcolsonlpsuffix(bcol, measurecols) : measurecols
    (!usefocalcolheuristics) && return measurecols

    focalmeasurecols = Symbol[]

    #=if (strippedbcol |> Symbol) ∉ extrabetaweights
      #along with maybe the s
      strippedbcol = stripcolprefix(bcol)=#

    for mcol ∈ matchedcols
      local isfocal::Bool
      strippedmcol = stripcolprefix(mcol) |> string

      if occursin(strippedmcol, bcol |> string)
        isfocal=true

      #see if we explicitly want to include it as a focal col
      else
        isfocal = false
        for k ∈ extrabetaweights
          #in this case, the specified extra column is embeded in the name
          # of the bcol
          if occursin(k|> string, bcol |> string) && (
            any((m->
              occursin(m|>string, strippedmcol)).(extrabetaweightsidx[k])))

            isfocal=true
          end
        end
      end

      if isfocal
        push!(focalmeasurecols, mcol)
      end
    end

    return focalmeasurecols
  end
  for bcol ∈ benchmarkcols
    tresults = resultsbythread[Threads.threadid()]

    focalmeasurecols = formfocalcols(bcol)
    #focalmeasurecols = matchonlpfilter ? matchcolsonlpsuffix(bcol, measurecols) : measurecols
    #we want to match the smoothing/lp filter level between the measure and the benchmark

    for mcol ∈ focalmeasurecols

      jointnomissing = view(joint, completecases(joint, [mcol,bcol]), [:date, mcol, bcol])
      m::Vector{Float64} = jointnomissing[:,mcol]
      b::Vector{Float64} = jointnomissing[:,bcol]
      N=nrow(jointnomissing)

      #try
        #stats all from wikipedia, so probably want better sources if used
        pear = cor(m,b) #pearson correlation
        peart = pear*((N-2)/(1-pear^2))^(0.5)
        spear = StatsBase.corspearman(m, b) #spearman correlation
        #spearz = atanh(spear)*((N-3)/1.06)^0.5
        speart = spear*((N-2)/(1-spear^2))^(0.5)
        kend = StatsBase.corkendall(m, b) #kendall's tau
        kendz = kend/((2*(2*N+5))/(9*N*(N-1)))^0.5 #not sure about this- investigate if using

        #=Xb = hcat(b, ones(N))
        lm = FMLM(Xb, m)
        beta=lm.β[1]
        @assert lm.β[1] ≈ pear * std(m)/std(b)=#
        Xm = hcat(m, ones(N))
        lmbm = FMLM(Xm, b)
        betabm=lmbm.β[1]
        @assert (lmbm.β[1] ≈ pear * std(b)/std(m)) || (lmbm.β[1] ≡ pear * std(b)/std(m)) "
          lmbm.β[1]= $(lmbm.β[1]) while pear * std(b)/std(m) = $(pear * std(b)/std(m))!!"
        betabmmw = betabm/modifiedwhiteΣ!(lmbm)[1,1]^0.5
        betabmhac = betabm/nwΣ(lmbm)[1,1]^0.5

        Xb = hcat(b, ones(N))
        lmmb = FMLM(Xb, m)
        betamb=lmmb.β[1]
        @assert ((lmmb.β[1] ≈ pear * std(m)/std(b)) || (lmmb.β[1] ≡ pear * std(m)/std(b)))
        betambmw = betamb/modifiedwhiteΣ!(lmmb)[1,1]^0.5
        betambhac = betamb/nwΣ(lmmb)[1,1]^0.5

        #compute the pearson correlation as a regression
        #this is kinda hacky- might want to replace with a block bootstrap
        #=Xbcor = hcat(b .* std(m) ./ std(b), ones(N))
        lmcor = FMLM(Xbcor, m)=#
        Xmcor = hcat(m .* std(b) ./ std(m), ones(N))
        lmcor = FMLM(Xmcor, b)
        @assert ((lmcor.β[1] ≈ pear) || (lmcor.β[1] ≡ pear))
        pearhacz = pear/nwΣ(lmcor)[1,1]^0.5

        #perform moving block bootstrap if desired
        if (occursin(bsmeasurecolsregex, string(mcol)) &&
            occursin(bsbenchmarkcolsregex, string(bcol)) &&
            (std(b) > minσ) && (std(m) > minσ))
          #@info "bcol:$bcol, mcol:$mcol"
          bs = mbbootstrap([b,m], bsstatfunc, bsaggfunc, B=bssamples, w=bswidth)
          cormbbz = pear/bs.cor.se
          spearmbbz = spear/bs.spear.se
          kendmbbz =  kend/bs.kend.se
          betabmmbbz =  betabm/bs.betabm.se
          betambmbbz =  betamb/bs.betamb.se

          cormbbp = bs.cor.p
          spearmbbp = bs.spear.p
          kendmbbp =  bs.kend.p
          betabmmbbp =  bs.betabm.p
          betambmbbp =  bs.betamb.p

          ar1h0teststats = Dict(:cor=>pear, :kend=>kend, :spear=>spear,
            :betabm => betabm, :betamb => betamb)
          ar1h0 = comparear1undernull(b,m, nullar1statfunc,
            (v,s)->nullar1aggfunc(v,s,ar1h0teststats),
            Nsimulations=nullar1simulations)
          #these are the alternate hypotheses to test

          corar1h0z =  pear/ar1h0.cor.se
          spearar1h0z = spear/ar1h0.spear.se
          kendar1h0z =  kend/ar1h0.kend.se
          betabmar1h0z =  betabm/ar1h0.betabm.se
          betambar1h0z =  betamb/ar1h0.betamb.se

          corar1h0p = ar1h0.cor.p
          spearar1h0p = ar1h0.spear.p
          kendar1h0p =  ar1h0.kend.p
          betabmar1h0p =  ar1h0.betabm.p
          betambar1h0p =  ar1h0.betamb.p
        else
          cormbbp=missing; spearmbbp=missing; kendmbbp=missing; betabmmbbp=missing; betambmbbp=missing;
          cormbbz=missing; spearmbbz=missing; kendmbbz=missing; betabmmbbz=missing; betambmbbz=missing;

          corar1h0p=missing; spearar1h0p=missing; kendar1h0p=missing; betabmar1h0p=missing; betambar1h0p=missing;
          corar1h0z=missing; spearar1h0z=missing; kendar1h0z=missing; betabmar1h0z=missing; betambar1h0z=missing;
        end

        push!(tresults,(;
          label=fundslabel,
          measure=mcol,
          benchmark=bcol,
          cor=pear,
          cort=peart,
          corhac=pearhacz,
          cormbbz, cormbbp, corar1h0z, corar1h0p,
          spear, speart, spearmbbz, spearmbbp, spearar1h0z, spearar1h0p,
          kend, kendz, kendmbbp, kendmbbz, kendar1h0p, kendar1h0z,
          betabm, betabmmw, betabmhac, betabmmbbz,  betabmmbbp, betabmar1h0z,  betabmar1h0p,
          betamb, betambmw, betambhac, betambmbbz,  betambmbbp, betambar1h0z,  betambar1h0p,
          N))

    end
  end

  joint |> CSV.write(
    "$(PARAM[:testpath])\\$(PARAM[:analysismeasurefilename])_$(fundslabel)_jointcor.csv")

  #collect the results in a single df
  append!(results, reduce(append!, resultsbythread))
  sort!(results, [:benchmark, :measure])

  return (;results, data=joint)
end

#normalizedassetcols(measure) = propertynames(measure[!, r"^NA_"])

max_stars_repo_name: clintonTE/CCA

max_issues_repo_name: clintonTE/CCA

max_forks_repo_name: clintonTE/CCA


text:
@testset "Formatting Options" begin
    @testset "remove extra newlines" begin
        str_ = """
        a = 10

        # foo1
        # ooo

        # aooo

        # aaaa
        b = 20

        # hello
        """
        str = """
        a = 10

        # foo1
        # ooo

        # aooo

        # aaaa
        b = 20

        # hello
        """
        @test fmt(str_, remove_extra_newlines = true) == str
        @test fmt(str_, remove_extra_newlines = false) == str_

        str_ = """
        module M

        function foo(bar::Int64, baz::Int64)

            return bar + baz
        end

        function foo(bar::In64, baz::Int64)
            return bar + baz

        end

        end
        """
        str = """
        module M

        function foo(bar::Int64, baz::Int64)
            return bar + baz
        end

        function foo(bar::In64, baz::Int64)
            return bar + baz
        end

        end
        """
        @test fmt(str_, remove_extra_newlines = true) == str
        @test fmt(str_, remove_extra_newlines = false) == str_
    end

    @testset "whitespace in typedefs" begin
        str_ = "Foo{A,B,C}"
        str = "Foo{A, B, C}"
        @test fmt(str_, whitespace_typedefs = true) == str

        str_ = """
        struct Foo{A<:Bar,Union{B<:Fizz,C<:Buzz},<:Any}
            a::A
        end"""
        str = """
        struct Foo{A <: Bar, Union{B <: Fizz, C <: Buzz}, <:Any}
            a::A
        end"""
        @test fmt(str_, whitespace_typedefs = true) == str

        str_ = """
        function foo() where {A,B,C{D,E,F{G,H,I},J,K},L,M<:N,Y>:Z}
            body
        end
        """
        str = """
        function foo() where {A, B, C{D, E, F{G, H, I}, J, K}, L, M <: N, Y >: Z}
            body
        end
        """
        @test fmt(str_, whitespace_typedefs = true) == str

        str_ = "foo() where {A,B,C{D,E,F{G,H,I},J,K},L,M<:N,Y>:Z} = body"
        str = "foo() where {A, B, C{D, E, F{G, H, I}, J, K}, L, M <: N, Y >: Z} = body"
        @test fmt(str_, whitespace_typedefs = true) == str
    end

    @testset "whitespace ops in indices" begin
        str = "arr[1 + 2]"
        @test fmt("arr[1+2]", m = 1, whitespace_ops_in_indices = true) == str

        str = "arr[(1 + 2)]"
        @test fmt("arr[(1+2)]", m = 1, whitespace_ops_in_indices = true) == str

        str_ = "arr[1:2*num_source*num_dump-1]"
        str = "arr[1:(2 * num_source * num_dump - 1)]"
        @test fmt(str_, m = 1, whitespace_ops_in_indices = true) == str

        str_ = "arr[2*num_source*num_dump-1:1]"
        str = "arr[(2 * num_source * num_dump - 1):1]"
        @test fmt(str_, m = 1, whitespace_ops_in_indices = true) == str

        str = "arr[(a + b):c]"
        @test fmt("arr[(a+b):c]", m = 1, whitespace_ops_in_indices = true) == str

        str = "arr[a in b]"
        @test fmt(str, m = 1, whitespace_ops_in_indices = true) == str

        str_ = "a:b+c:d-e"
        str = "a:(b + c):(d - e)"
        @test fmt(str_, m = 1, whitespace_ops_in_indices = true) == str

        # issue 180
        str_ = "s[m+i+1]"
        str = "s[m+i+1]"
        @test fmt(str, m = 1) == str

        str = "s[m + i + 1]"
        @test fmt(str_, m = 1, whitespace_ops_in_indices = true) == str
    end

    @testset "rewrite import to using" begin
        str_ = "import A"
        str = "using A: A"
        @test fmt(str_, import_to_using = true) == str

        str_ = """
        import A,

        B, C"""
        str = """
        using A: A
        using B: B
        using C: C"""
        @test_broken fmt(str_, import_to_using = true) == str

        str_ = """
        import A,
               # comment
        B, C"""
        str = """
        using A: A
        # comment
        using B: B
        using C: C"""
        @test fmt(str_, import_to_using = true) == str

        str_ = """
        import A, # inline
               # comment
        B, C # inline"""
        str = """
        using A: A # inline
        # comment
        using B: B
        using C: C # inline"""
        @test fmt(str_, import_to_using = true) == str

        str_ = """
        import ..A, .B, ...C"""
        str = """
        using ..A: A
        using .B: B
        using ...C: C"""
        @test fmt(str_, import_to_using = true) == str
        t = run_pretty(str_, opts = Options(margin = 80, import_to_using = true))
        @test t.len == 13

        # issue 232
        str = """import A.b"""
        @test fmt(str, import_to_using = true) == str

        str = """import A.b: c"""
        @test fmt(str, import_to_using = true) == str

        str = """import A.b.c"""
        @test fmt(str, import_to_using = true) == str

        str = """import A.b.c: d"""
        @test fmt(str, import_to_using = true) == str
    end

    @testset "always convert `=` to `in` (for loops)" begin
        str_ = """
        for i = 1:n
            println(i)
        end"""
        str = """
        for i in 1:n
            println(i)
        end"""
        @test fmt(str_, always_for_in = true) == str
        @test fmt(str, always_for_in = true) == str

        str_ = """
        for i = I1, j in I2
            println(i, j)
        end"""
        str = """
        for i in I1, j in I2
            println(i, j)
        end"""
        @test fmt(str_, always_for_in = true) == str
        @test fmt(str, always_for_in = true) == str

        str_ = """
        for i = 1:30, j = 100:-2:1
            println(i, j)
        end"""
        str = """
        for i in 1:30, j in 100:-2:1
            println(i, j)
        end"""
        @test fmt(str_, always_for_in = true) == str
        @test fmt(str, always_for_in = true) == str

        str_ = "[(i,j) for i=I1,j=I2]"
        str = "[(i, j) for i in I1, j in I2]"
        @test fmt(str_, always_for_in = true) == str
        @test fmt(str, always_for_in = true) == str

        str_ = "((i,j) for i=I1,j=I2)"
        str = "((i, j) for i in I1, j in I2)"
        @test fmt(str_, always_for_in = true) == str
        @test fmt(str, always_for_in = true) == str

        str_ = "[(i, j) for i = 1:2:10, j = 100:-1:10]"
        str = "[(i, j) for i in 1:2:10, j in 100:-1:10]"
        @test fmt(str_, always_for_in = true) == str
        @test fmt(str, always_for_in = true) == str
    end

    @testset "rewrite x |> f to f(x)" begin
        @test fmt("x |> f", pipe_to_function_call = true) == "f(x)"

        str_ = "var = func1(arg1) |> func2 |> func3 |> func4 |> func5"
        str = "var = func5(func4(func3(func2(func1(arg1)))))"
        @test fmt(str_, pipe_to_function_call = true) == str
        @test fmt(str_, pipe_to_function_call = true, margin = 1) == fmt(str)
    end

    @testset "function shortdef to longdef" begin
        str_ = "foo(a) = bodybodybody"
        str = """
        function foo(a)
            bodybodybody
        end"""
        @test fmt(str_, 4, length(str_), short_to_long_function_def = true) == str_
        @test fmt(str_, 4, length(str_) - 1, short_to_long_function_def = true) == str

        str_ = "foo(a::T) where {T} = bodybodybodybodybodybodyb"
        str = """
        function foo(a::T) where {T}
            bodybodybodybodybodybodyb
        end"""
        @test fmt(str_, 4, length(str_), short_to_long_function_def = true) == str_
        @test fmt(str_, 4, length(str_) - 1, short_to_long_function_def = true) == str

        str_ = "foo(a::T)::R where {T} = bodybodybodybodybodybodybody"
        str = """
        function foo(a::T)::R where {T}
            bodybodybodybodybodybodybody
        end"""
        @test fmt(str_, 4, length(str_), short_to_long_function_def = true) == str_
        @test fmt(str_, 4, length(str_) - 1, short_to_long_function_def = true) == str
    end

    @testset "always use return" begin
        str_ = "foo(a) = bodybodybody"
        str = """
        function foo(a)
            return bodybodybody
        end"""
        @test fmt(
            str_,
            4,
            length(str_) - 1,
            short_to_long_function_def = true,
            always_use_return = true,
        ) == str

        str_ = """
        function foo()
            expr1
            expr2
        end"""
        str = """
        function foo()
            expr1
            return expr2
        end"""
        @test fmt(str_, 4, length(str_) - 1, always_use_return = true) == str

        str_ = """
        macro foo()
            expr1
            expr2
        end"""
        str = """
        macro foo()
            expr1
            return expr2
        end"""
        @test fmt(str_, 4, length(str_) - 1, always_use_return = true) == str

        str_ = """
        map(arg1, arg2) do x, y
            expr1
            expr2
        end"""
        str = """
        map(arg1, arg2) do x, y
            expr1
            return expr2
        end"""
        @test fmt(str_, 4, length(str_) - 1, always_use_return = true) == str

        str = """
        function foo()
            @macrocall(expr2)
        end"""
        @test fmt(str, 4, 92, always_use_return = true) == str

        str = """
        function foo()
            @macroblock expr2
        end"""
        @test fmt(str, 4, 92, always_use_return = true) == str

        str = """
        function foo()
            for i = 1:10
                println(i)
            end
        end"""
        @test fmt(str, 4, 92, always_use_return = true) == str

        str = """
        function f(a)
            if a > 0
                return -1
            else
                return 1
            end
        end"""
        @test fmt(str, 4, 92, always_use_return = true) == str
    end

    @testset "whitespace in keyword arguments" begin
        str_ = "f(; a = b)"
        str = "f(; a=b)"
        @test fmt(str_, 4, 92, whitespace_in_kwargs = false) == str

        str = "f(; a!) = a!"
        @test fmt(str, 4, 92, whitespace_in_kwargs = false) == str

        # issue 242
        str_ = "f(a, b! = 1; c! = 2, d = 3, e! = 4)"
        str = "f(a, (b!)=1; (c!)=2, d=3, (e!)=4)"
        @test fmt(str_, 4, 92, whitespace_in_kwargs = false) == str

        str_ = "( k1 =v1,  k2! = v2)"
        str = "(k1=v1, (k2!)=v2)"
        @test fmt(str_, 4, 80, style = YASStyle(), whitespace_in_kwargs = false) == str
        @test fmt(str_, 4, 80, style = DefaultStyle(), whitespace_in_kwargs = false) == str

        str_ = "( k1 =v1,  k2! = v2)"
        str = "(k1 = v1, k2! = v2)"
        @test fmt(str_, 4, 80, style = YASStyle(), whitespace_in_kwargs = true) == str
        @test fmt(str_, 4, 80, style = DefaultStyle(), whitespace_in_kwargs = true) == str

        str_ = "(; g = >=(1))"
        str = "(; g=(>=(1)))"
        @test fmt(str_, 4, 92, whitespace_in_kwargs = false) == str
    end

    @testset "annotate untyped fields with `Any`" begin
        str = """
        struct name
            arg::Any
        end"""

        str_ = """
        struct name
            arg
        end"""
        @test fmt(str_) == str

        str_ = """
        struct name
        arg
        end"""
        @test fmt(str_) == str

        str_ = """
        struct name
                arg
            end"""
        @test fmt(str_) == str

        t = run_pretty(str_, 80)
        @test length(t) == 12

        str = """
        mutable struct name
            reallylongfieldname::Any
        end"""

        str_ = """
        mutable struct name
            reallylongfieldname
        end"""
        @test fmt(str_) == str

        str_ = """
        mutable struct name
        reallylongfieldname
        end"""
        @test fmt(str_) == str

        str_ = """
        mutable struct name
                reallylongfieldname
            end"""
        @test fmt(str_) == str

        t = run_pretty(str_, 80)
        @test length(t) == 28

        str = """
        struct name
            arg
        end"""

        str_ = """
        struct name
            arg
        end"""
        @test fmt(str_, annotate_untyped_fields_with_any = false) == str

        str_ = """
        struct name
        arg
        end"""
        @test fmt(str_, annotate_untyped_fields_with_any = false) == str

        str_ = """
        struct name
                arg
            end"""
        @test fmt(str_, annotate_untyped_fields_with_any = false) == str

        t = run_pretty(
            str_,
            opts = Options(margin = 80, annotate_untyped_fields_with_any = false),
        )
        @test length(t) == 11

        str = """
        mutable struct name
            reallylongfieldname
        end"""

        str_ = """
        mutable struct name
            reallylongfieldname
        end"""
        @test fmt(str_, annotate_untyped_fields_with_any = false) == str

        str_ = """
        mutable struct name
        reallylongfieldname
        end"""
        @test fmt(str_, annotate_untyped_fields_with_any = false) == str

        str_ = """
        mutable struct name
                reallylongfieldname
            end"""
        @test fmt(str_, annotate_untyped_fields_with_any = false) == str

        t = run_pretty(
            str_,
            opts = Options(margin = 80, annotate_untyped_fields_with_any = false),
        )
        @test length(t) == 23
    end

    @testset "format docstrings - basic" begin
        str = """
        \"""
        doc
        \"""
        function f()
            20
        end"""
        t = run_pretty(str, 80)
        @test length(t) == 12

        normalized = """
        \"""
        doc
        \"""
        function f()
            20
        end"""

        str = """
        \"""doc
        \"""
        function f()
            20
        end"""
        @test fmt(str) == str
        @test fmt(str, format_docstrings = true) == normalized

        str = """
        \"""
        doc\"""
        function f()
            20
        end"""
        @test fmt(str) == str
        @test fmt(str, format_docstrings = true) == normalized

        str = """
        \"""doc\"""
        function f()
            20
        end"""
        @test fmt(str) == str
        @test fmt(str, format_docstrings = true) == normalized

        str = """
        "doc
        "
        function f()
            20
        end"""
        @test fmt(str) == str
        @test fmt(str, format_docstrings = true) == normalized

        str = """
        "
        doc"
        function f()
            20
        end"""
        @test fmt(str) == str
        @test fmt(str, format_docstrings = true) == normalized

        str = """
        "doc"
        function f()
            20
        end"""
        @test fmt(str) == str
        @test fmt(str, format_docstrings = true) == normalized

        # test aligning to function identation
        str_ = """
            "doc"
        function f()
            20
        end"""
        str = """
        "doc"
        function f()
            20
        end"""
        @test fmt(str_) == str
        @test fmt(str_, format_docstrings = true) == normalized

        str = """\"""
                 doc for Foo
                 \"""
                 Foo"""
        @test fmt(str) == str
        t = run_pretty(str, 80)
        @test length(t) == 11

        str = """
        \"""
        doc
        \"""
        function f()    #  comment
            20
        end"""
        @test fmt(str) == str
        @test fmt(str, format_docstrings = true) == str

        # Issue 157
        str = raw"""
        @doc \"""
           foo()
        \"""
        foo() = bar()"""
        @test fmt(str) == str
        @test fmt(str, format_docstrings = true) == str

        str = raw"""
        @doc docϵ\"""
           foo()
        \"""
        foo() = bar()"""
        @test fmt(str) == str
        @test fmt(str, format_docstrings = true) == str

        str = raw"""@doc "doc for foo" foo"""
        @test fmt(str) == str
        @test fmt(str, format_docstrings = true) == str

        str = raw"""@doc \"""doc for foo\""" foo"""
        @test fmt(str) == str
        @test fmt(str, format_docstrings = true) == str

        str = raw"""@doc doc\"""doc for foo\""" foo()"""
        @test fmt(str) == str
        @test fmt(str, format_docstrings = true) == str

        str = raw"""@doc foo"""
        @test fmt(str) == str
        @test fmt(str, format_docstrings = true) == str

        # issue 160
        str = raw"""
        module MyModule

        import Markdown: @doc_str

        @doc doc\"""
            foo()
        \"""
        foo() = bar()

        end # module"""
        @test fmt(str) == str
        @test fmt(str, format_docstrings = true) == str
    end

    @testset "format docstrings - with code" begin
        unformatted = """
        \"""
        This is a docstring

        ```@example
        a =  1
         b  = 2
         a + b
        ```

        ```jldoctest
        a =  1
        b  = 2
        a + b

        # output

        3
        ```

        ```jldoctest
        julia> a =  1
        1

        julia> b  = 2;

        julia>  a + b
        3

        julia> function test(x)
               x + 1
               x + 2
               end;
        ```
        \"""
        function test(x) x end"""

        formatted = """
        \"""
        This is a docstring

        ```@example
        a = 1
        b = 2
        a + b
        ```

        ```jldoctest
        a = 1
        b = 2
        a + b

        # output

        3
        ```

        ```jldoctest
        julia> a = 1
        1

        julia> b = 2;

        julia> a + b
        3

        julia> function test(x)
                   x + 1
                   x + 2
               end;

        ```
        \"""
        function test(x)
            x
        end"""
        @test fmt(unformatted, format_docstrings = true) == formatted
    end

    @testset "format docstrings - Multi-line indented code-blocks" begin
        unformatted = """
        \"""
            fmt(
            )
        \"""
        function fmt() end"""

        formatted = """
        \"""
            fmt(
            )
        \"""
        function fmt() end"""
        @test fmt(unformatted, format_docstrings = true) == formatted
    end

    @testset "format docstrings - Empty line in docstring" begin
        unformatted = """
        \"""

        \"""
        function test() end"""

        formatted = """
        \"""
        \"""
        function test() end"""
        @test fmt(unformatted, format_docstrings = true) == formatted
    end

    @testset "format docstrings - Indented docstring" begin
        unformatted = """
        begin
            \"""
            Indented docstring

            with multiple paragraphs
            \"""
            indented_item
        end"""
        formatted = """
        begin
            \"""
            Indented docstring

            with multiple paragraphs
            \"""
            indented_item
        end"""
        @test fmt(unformatted, format_docstrings = true) == formatted

        short = """
        begin
            "Short docstring"
            item
        end
        """
        short_formatted = """
        begin
            \"""
            Short docstring
            \"""
            item
        end
        """
        @test fmt(short, format_docstrings = true) == short_formatted
    end

    @testset "align struct fields" begin
        str_ = """
        struct Foo
            a::T
        end"""
        str = """
        struct Foo
            a::T
        end"""
        @test fmt(str_, align_struct_field = true) == str

        str = """
        struct Foo
            a             :: T
            longfieldname :: B
        end"""
        str_ = """
        struct Foo
            a::T
            longfieldname::B
        end"""
        @test fmt(str, align_struct_field = true) == str
        @test fmt(str, align_struct_field = false) == str_

        str_ = """
        Base.@kwdef struct Options
            indent_size::Int                       = 4
            margin::Int                            = 92
            always_for_in::Bool                 = false
            whitespace_typedefs::Bool          = false
            whitespace_ops_in_indices::Bool        = false
            remove_extra_newlines::Bool            = false
            import_to_using::Bool                  = false
            pipe_to_function_call::Bool            = false
            short_to_long_function_def::Bool      = false
            always_use_return::Bool           = false
            whitespace_in_kwargs::Bool          = true
            annotate_untyped_fields_with_any::Bool = true
            format_docstrings::Bool             = false
            align_struct_fields::Bool           = false

            another_field1::BlahBlahBlah = 10
            field2::Foo                          = 10

            Options() = new()
        end"""
        str = """
        Base.@kwdef struct Options
            indent_size::Int                       = 4
            margin::Int                            = 92
            always_for_in::Bool                    = false
            whitespace_typedefs::Bool              = false
            whitespace_ops_in_indices::Bool        = false
            remove_extra_newlines::Bool            = false
            import_to_using::Bool                  = false
            pipe_to_function_call::Bool            = false
            short_to_long_function_def::Bool       = false
            always_use_return::Bool                = false
            whitespace_in_kwargs::Bool             = true
            annotate_untyped_fields_with_any::Bool = true
            format_docstrings::Bool                = false
            align_struct_fields::Bool              = false

            another_field1::BlahBlahBlah = 10
            field2::Foo = 10

            Options() = new()
        end"""
        @test fmt(str_, align_struct_field = true) == str

        str_ = """
        Base.@kwdef struct Options
            indent_size::Int = 4
            margin::Int = 92
            always_for_in::Bool = false
            whitespace_typedefs::Bool = false
            whitespace_ops_in_indices::Bool = false
            remove_extra_newlines::Bool = false
            import_to_using::Bool = false
            pipe_to_function_call::Bool = false
            short_to_long_function_def::Bool = false
            always_use_return::Bool = false
            whitespace_in_kwargs::Bool = true
            annotate_untyped_fields_with_any::Bool = true
            format_docstrings::Bool = false
            align_struct_fields::Bool = false

            another_field1::BlahBlahBlah = 10
            field2::Foo = 10

            Options() = new()
        end"""
        @test fmt(str, align_struct_field = true) == str
        @test fmt(str, align_struct_field = false) == str_

        str = """
        Base.@kwdef struct Options
            indent_size::Int                       = 4
            margin::Int                            = 92
            always_for_in::Bool                    = false
            whitespace_typedefs::Bool              = false
            whitespace_ops_in_indices::Bool        = false
            remove_extra_newlines::Bool            = false
            import_to_using::Bool                  = false
            pipe_to_function_call::Bool            = false
            short_to_long_function_def::Bool       = false
            always_use_return::Bool                = false
            whitespace_in_kwargs::Bool             = true
            annotate_untyped_fields_with_any::Bool = true
            format_docstrings::Bool                = false
            align_struct_fields::Bool              = false

            another_field1::BlahBlahBlah =
                10
            field2::Foo =
                10

            Options() =
                new()
        end"""
        @test fmt(str, 4, 1, align_struct_field = true) == str
    end

    @testset "align assignment" begin
        str_ = """
        const variable1 = 1
        const var2      = 2
        const var3 = 3
        const var4 = 4"""
        str = """
        const variable1 = 1
        const var2      = 2
        const var3      = 3
        const var4      = 4"""
        @test fmt(str_, align_assignment = true) == str

        str = """
        const variable1 = 1
        const variable2 = 2
        const var3 = 3
        const var4 = 4"""
        @test fmt(str, align_assignment = true) == str

        str_ = """
        module Foo

        const UTF8PROC_STABLE    = (1<<1)
        const UTF8PROC_COMPAT    = (1<<2)
        const UTF8PROC_COMPOSE   = (1<<3)
        const UTF8PROC_DECOMPOSE = (1<<4)
        const UTF8PROC_IGNORE    = (1<<5)
        const UTF8PROC_REJECTNA  = (1<<6)
        const UTF8PROC_NLF2LS    = (1<<7)
        const UTF8PROC_NLF2PS    = (1<<8)
        const UTF8PROC_NLF2LF    = (UTF8PROC_NLF2LS | UTF8PROC_NLF2PS)
        const UTF8PROC_STRIPCC   = (1<<9)
        const UTF8PROC_CASEFOLD  = (1<<10)
        const UTF8PROC_CHARBOUND = (1<<11)
        const UTF8PROC_LUMP=(1<<12)
        const UTF8PROC_STRIP         = (1<<13) # align this

        const FOOBAR = 0
        const FOO = 1

        end"""

        str = """
        module Foo

        const UTF8PROC_STABLE    = (1 << 1)
        const UTF8PROC_COMPAT    = (1 << 2)
        const UTF8PROC_COMPOSE   = (1 << 3)
        const UTF8PROC_DECOMPOSE = (1 << 4)
        const UTF8PROC_IGNORE    = (1 << 5)
        const UTF8PROC_REJECTNA  = (1 << 6)
        const UTF8PROC_NLF2LS    = (1 << 7)
        const UTF8PROC_NLF2PS    = (1 << 8)
        const UTF8PROC_NLF2LF    = (UTF8PROC_NLF2LS | UTF8PROC_NLF2PS)
        const UTF8PROC_STRIPCC   = (1 << 9)
        const UTF8PROC_CASEFOLD  = (1 << 10)
        const UTF8PROC_CHARBOUND = (1 << 11)
        const UTF8PROC_LUMP      = (1 << 12)
        const UTF8PROC_STRIP     = (1 << 13) # align this

        const FOOBAR = 0
        const FOO = 1

        end"""
        # @test fmt(str_, align_assignment = false) == str
        @test fmt(str_, align_assignment = true) == str

        # the aligned consts will NOT be nestable
        str = """
        module Foo

        const UTF8PROC_STABLE    = (1 << 1)
        const UTF8PROC_COMPAT    = (1 << 2)
        const UTF8PROC_COMPOSE   = (1 << 3)
        const UTF8PROC_DECOMPOSE = (1 << 4)
        const UTF8PROC_IGNORE    = (1 << 5)
        const UTF8PROC_REJECTNA  = (1 << 6)
        const UTF8PROC_NLF2LS    = (1 << 7)
        const UTF8PROC_NLF2PS    = (1 << 8)
        const UTF8PROC_NLF2LF    = (UTF8PROC_NLF2LS | UTF8PROC_NLF2PS)
        const UTF8PROC_STRIPCC   = (1 << 9)
        const UTF8PROC_CASEFOLD  = (1 << 10)
        const UTF8PROC_CHARBOUND = (1 << 11)
        const UTF8PROC_LUMP      = (1 << 12)
        const UTF8PROC_STRIP     = (1 << 13) # align this

        const FOOBAR =
            0
        const FOO =
            1

        end"""
        @test fmt(str_, 4, 1, align_assignment = true) == str

        str = """
        a  = 1
        bc = 2

        long_variable = 1
        other_var     = 2
        """
        @test fmt(str, 4, 1, align_assignment = true) == str

        str = """
        vcat(X::T...) where {T}         = T[X[i] for i = 1:length(X)]
        vcat(X::T...) where {T<:Number} = T[X[i] for i = 1:length(X)]
        hcat(X::T...) where {T}         = T[X[j] for i = 1:1, j = 1:length(X)]
        hcat(X::T...) where {T<:Number} = T[X[j] for i = 1:1, j = 1:length(X)]
        """
        @test fmt(str, 4, 1, align_assignment = true) == str

        str = """
        μs, ns = divrem(ns, 1000)
        ms, μs = divrem(μs, 1000)
        s, ms = divrem(ms, 1000)
        """
        @test fmt(str, 4, 100, align_assignment = true) == str

        str = """
        run = wandb.init(
            name      = name,
            project   = project,
            config    = config,
            notes     = notes,
            tags      = tags,
            dir       = dir,
            job_type  = job_type,
            entity    = entity,
            group     = group,
            id        = id,
            reinit    = reinit,
            resume    = resume,
            anonymous = anonymous ? "allow" : "never",
        )
        """
        @test fmt(str, 4, 100, align_assignment = true, whitespace_in_kwargs = false) == str

        str_ = """
        s           = model.sys
        @unpack A,K = s
        C           = s.C
        poles       = eigvals(A - K * C)
        """
        str = """
        s            = model.sys
        @unpack A, K = s
        C            = s.C
        poles        = eigvals(A - K * C)
        """
        @test fmt(str_, 4, 100, align_assignment = true) == str

        str_ = """
        s             = model.sys
        @unpack A,K   = s
        C             = s.C
        const polesss = eigvals(A - K * C)
        """
        str = """
        s             = model.sys
        @unpack A, K  = s
        C             = s.C
        const polesss = eigvals(A - K * C)
        """
        @test fmt(str_, 4, 100, align_assignment = true) == str

        str = """
        function stabilize(model)
            s            = model.sys
            @unpack A, K = s
            C            = s.C
            poles        = eigvals(A - K * C)
            newpoles     = map(poles) do p
                ap = abs(p)
                ap <= 1 && (return p)
                p / (ap + sqrt(eps()))
            end
            K2           = ControlSystems.acker(A', C', newpoles)' .|> real
            all(abs(p) <= 1 for p in eigvals(A - K * C)) || @warn("Failed to stabilize predictor")
            s.K .= K2
            model
        end
        """
        @test fmt(str, 4, 100, align_assignment = true) == str
    end

    @testset "align conditionals" begin
        str_ = """
        index = zeros(n <= typemax(Int8)  ? Int8  :
                      n <= typemax(Int16) ? Int16 :
                      n <= typemax(Int32) ? Int32 : Int64, n)
        """

        str = """
        index = zeros(
            n <= typemax(Int8)  ? Int8  :
            n <= typemax(Int16) ? Int16 :
            n <= typemax(Int32) ? Int32 : Int64,
            n,
        )
        """
        @test fmt(str_, align_conditional = true) == str

        str = """
        index =
            zeros(
                n <= typemax(Int8)  ? Int8  :
                n <= typemax(Int16) ? Int16 :
                n <= typemax(Int32) ? Int32 : Int64,
                n,
            )
        """
        @test fmt(str_, 4, 1, align_conditional = true) == str

        str_ = """
        index = zeros(n <= typemax(Int8)  ? Int8 :   # inline
                        #comment 1
                      n <= typemax(Int16) ? Int16 :   # inline 2
                              # comment 2
                      n <= typemax(Int32) ? Int32 : # inline 3
                      Int64, n)
        """
        str = """
        index =
            zeros(
                n <= typemax(Int8)  ? Int8 :   # inline
                #comment 1
                n <= typemax(Int16) ? Int16 :   # inline 2
                # comment 2
                n <= typemax(Int32) ? Int32 : # inline 3
                Int64,
                n,
            )
        """
        @test fmt(str_, 4, 1, align_conditional = true) == str

        str_ = """
        index = zeros(n <= typemax(Int8)  ? Int8  :    # inline
                      n <= typemax(Int16) ? Int16 : n <= typemax(Int32) ? Int32 : Int64, n)
        """

        str = """
        index =
            zeros(
                n <= typemax(Int8)  ? Int8  :    # inline
                n <= typemax(Int16) ? Int16 : n <= typemax(Int32) ? Int32 : Int64,
                n,
            )
        """
        @test fmt(str_, 4, 1, align_conditional = true) == str

        str_ = """
        index =
            zeros(
                n <= typemax(Int8)     ? Int8  :
                n <= typemax(Int16A) ? Int16  :
                n <= typemax(Int32)  ? Int322 : Int64,
                n,
            )
        """
        str = """
        index =
            zeros(
                n <= typemax(Int8)   ? Int8   :
                n <= typemax(Int16A) ? Int16  :
                n <= typemax(Int32)  ? Int322 : Int64,
                n,
            )
        """
        @test fmt(str_, 4, 1, align_conditional = true) == str

        str_ = """
        val = cst.kind === Tokens.ABSTRACT ? "abstract" :
            cst.kind === Tokens.BAREMODULE ? "baremodule" : ""
        """
        str = """
        val = cst.kind === Tokens.ABSTRACT ? "abstract" : cst.kind === Tokens.BAREMODULE ? "baremodule" : ""
        """
        @test fmt(str_, 4, 100, align_conditional = true) == str

        str_ = """
        val = cst.kind === Tokens.ABSTRACT ? "abstract" :
            cst.kind === Tokens.BAREMODUL  ? "baremodule" : ""
        """
        str = """
        val = cst.kind === Tokens.ABSTRACT  ? "abstract" :
              cst.kind === Tokens.BAREMODUL ? "baremodule" : ""
        """
        @test fmt(str_, 4, 100, align_conditional = true) == str

        str = """
        val =
            cst.kind === Tokens.ABSTRACT  ? "abstract" :
            cst.kind === Tokens.BAREMODUL ? "baremodule" : ""
        """
        @test fmt(str_, 4, 1, align_conditional = true) == str
    end

    @testset "align pair arrow `=>`" begin
        str_ = """
        pages = [
            "Introduction" => "index.md",
            "How It Works" => "how_it_works.md",
            "Code Style"          => "style.md",
            "Skipping Formatting" => "skipping_formatting.md",
            "Syntax Transforms" => "transforms.md",
            "Custom Alignment" => "custom_alignment.md",
            "Custom Styles" => "custom_styles.md",
            "YAS Style" => "yas_style.md",
            "Configuration File" => "config.md",
            "API Reference" => "api.md",
        ]
        """
        str = """
        pages = [
            "Introduction"        => "index.md",
            "How It Works"        => "how_it_works.md",
            "Code Style"          => "style.md",
            "Skipping Formatting" => "skipping_formatting.md",
            "Syntax Transforms"   => "transforms.md",
            "Custom Alignment"    => "custom_alignment.md",
            "Custom Styles"       => "custom_styles.md",
            "YAS Style"           => "yas_style.md",
            "Configuration File"  => "config.md",
            "API Reference"       => "api.md",
        ]
        """
        @test fmt(str_, 4, 100, align_pair_arrow = true) == str

        str = """
        pages =
            [
                "Introduction"        => "index.md",
                "How It Works"        => "how_it_works.md",
                "Code Style"          => "style.md",
                "Skipping Formatting" => "skipping_formatting.md",
                "Syntax Transforms"   => "transforms.md",
                "Custom Alignment"    => "custom_alignment.md",
                "Custom Styles"       => "custom_styles.md",
                "YAS Style"           => "yas_style.md",
                "Configuration File"  => "config.md",
                "API Reference"       => "api.md",
            ]
        """
        @test fmt(str_, 4, 1, align_pair_arrow = true) == str
    end

    @testset "conditional to `if` block" begin
        str_ = """
        E ? A : B
        """
        @test fmt(str_, 2, 9, conditional_to_if = true) == str_

        str = """
        if E
          A
        else
          B
        end
        """
        @test fmt(str_, 2, 8, conditional_to_if = true) == str

        str_ = """
        begin
            E1 ? A : E2 ? B : foo(E333, E444) ? D : E
        end
        """
        @test fmt(str_, 4, 45, conditional_to_if = true) == str_

        str = """
        begin
            if E1
                A
            elseif E2
                B
            elseif foo(E333, E444)
                D
            else
                E
            end
        end
        """
        @test fmt(str_, 4, 44, conditional_to_if = true) == str
        @test fmt(str_, 4, 26, conditional_to_if = true) == str

        str = """
        begin
            if E1
                A
            elseif E2
                B
            elseif foo(
                E333,
                E444,
            )
                D
            else
                E
            end
        end
        """
        @test fmt(str_, 4, 25, conditional_to_if = true) == str

        str_ = """
        foobar = some_big_long_thing * 10_000 == 2 ?
            #comment
            bar :
            #comment
            another_big_long_thing * 10^300 / this_things_here
        """

        str = """
        foobar = if some_big_long_thing * 10_000 == 2
            #comment
            bar
        else
            #comment
            another_big_long_thing * 10^300 / this_things_here
        end
        """
        @test fmt(str_, conditional_to_if = true) == str
    end

    @testset "normalize_line_endings" begin
        windows_str = "a\r\nb\r\nc\r\nd"
        unix_str = "a\nb\nc\nd"
        mixed_windows_str = "a\r\nb\r\nc\nd"
        mixed_unix_str = "a\r\nb\nc\nd"

        @test fmt(windows_str, normalize_line_endings = "auto") == windows_str
        @test fmt(unix_str, normalize_line_endings = "auto") == unix_str
        @test fmt(mixed_windows_str, normalize_line_endings = "auto") == windows_str
        @test fmt(mixed_unix_str, normalize_line_endings = "auto") == unix_str

        @test fmt(windows_str, normalize_line_endings = "unix") == unix_str
        @test fmt(unix_str, normalize_line_endings = "unix") == unix_str
        @test fmt(mixed_windows_str, normalize_line_endings = "unix") == unix_str
        @test fmt(mixed_unix_str, normalize_line_endings = "unix") == unix_str

        @test fmt(windows_str, normalize_line_endings = "windows") == windows_str
        @test fmt(unix_str, normalize_line_endings = "windows") == windows_str
        @test fmt(mixed_windows_str, normalize_line_endings = "windows") == windows_str
        @test fmt(mixed_unix_str, normalize_line_endings = "windows") == windows_str
    end
end

ext: jl


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

@stochastic_model simple begin
    @stage 1 begin
        @decision(simple, x₁ >= 40)
        @decision(simple, x₂ >= 20)
        objective = @expression(simple, 100*x₁ + 150*x₂)
        @objective(simple, Min, objective)
        @constraint(simple, x₁ + x₂ <= 120)
    end
    @stage 2 begin
        @uncertain q₁ q₂ d₁ d₂
        @recourse(simple, 0 <= y₁ <= d₁)
        @recourse(simple, 0 <= y₂ <= d₂)
        @objective(simple, Max, q₁*y₁ + q₂*y₂)
        @constraint(simple, 6*y₁ + 10*y₂ <= 60*x₁)
        @constraint(simple, 8*y₁ + 5*y₂ <= 80*x₂)
    end
end

ξ₁ = @scenario q₁ = 24.0 q₂ = 28.0 d₁ = 500.0 d₂ = 100.0 probability = 0.4
ξ₂ = @scenario q₁ = 28.0 q₂ = 32.0 d₁ = 300.0 d₂ = 300.0 probability = 0.6

simple_res = SPResult([46.67, 36.25], Dict(1 => [300., 100.], 2 => [300., 100.]), -855.83, -1518.75, 662.92, 286.92, -1445.92, -568.92)
push!(problems, (simple, [ξ₁,ξ₂], simple_res, "Simple"))


text:
using CUDA, Flux, MLDatasets, Statistics, Images, Parameters, Printf, Random
using Base.Iterators: partition
using Flux.Optimise: update!
using Flux.Losses: logitbinarycrossentropy
using Flux.Data: DataLoader
using Flux.Optimise: train!, @epochs
using Parameters, Printf, Random
using BSON: @save
using Distributions
using Zygote

include("spectral_norm.jl")

@with_kw struct Settings
	G = nothing
	D = nothing
	loss = nothing
	img_fun = nothing
	rand_input = nothing
	batch_size::Int = 128
	latent_dim::Int = 100
	nclasses::Int = 2
	epochs::Int = 120
	verbose_freq::Int = 5
	output_x::Int = 6
	output_y::Int = 6
	optD = ADAM(0.0002, (0.5, 0.99))
	optG = ADAM(0.0002, (0.5, 0.99))
	output_dir = "output" # save dir
end

## Regular DCGAN
struct DCGANLoss end

function Lᴰ(t::DCGANLoss, G, D, z, ny, x, y)
	real_loss = logitbinarycrossentropy(D(x, y), 1f0, agg=mean)
	fake_loss = logitbinarycrossentropy(D(G(z, ny), ny), 0f0, agg=mean)
	return real_loss + fake_loss
end

Lᴳ(t::DCGANLoss, G, D, z, ny) = logitbinarycrossentropy(D(G(z, ny), ny), 1f0, agg=mean)

## Least Squares Loss
struct LSLoss end

function Lᴰ(t::LSLoss, G, D, z, ny, x, y)
	Flux.mse(D(x, y), 1f0) + Flux.mse(D(G(z, ny), ny), 0f0)
end

Lᴳ(t::LSLoss, G, D, z, ny) = Flux.mse(D(G(z, ny), ny), 1f0)

## WGAN-GP
@with_kw struct WLossGP
	λ = 10f0 # gp penalty
	gp_type = :exact #:approx or :exact
	noise_range = 1 # Range of the ϵ parameter that picks a point to evalutate the gradient (smaller is closer to "real" data)
end

function gradient_penalty(D, x, y)
	B = size(y, 2)
	l, b = Flux.pullback(() -> D(x, y), Flux.params(x, y))
	grads = b(ones(Float32, 1, size(x, 4)) |> gpu)
	Flux.mean((sqrt.(sum(reshape(grads[x], :, B).^2, dims = 1) .+ sum(grads[y].^2, dims = 1)) .- 1f0).^2)
end

function approx_penalty(D, x, xhat, y, yhat)
	B = size(x, 4) # batch size
	Δx, Δy = reshape(x .- xhat, :, B), y .- yhat # difference
	mag = sqrt.(sum(Δx.^2, dims = 1) + sum(Δy.^2, dims = 1))
	xdir, ydir = reshape(Δx ./ mag, size(xhat)), Δy ./ mag
	δ = 0.01f0
	ΔD = abs.(D(xhat, yhat) .- D(xhat .+ δ .* xdir, yhat .+ δ .* ydir))
	Flux.mean((ΔD ./ δ .- 1f0).^2)
	Flux.mean(max.((ΔD ./ δ .- 1f0), 0f0).^2)
end

function Lᴰ(t::WLossGP, G, D, z, ny, x, y)
	# ϵ = Zygote.ignore(() -> Float32.(rand(Uniform(0, t.noise_range), 1, size(y,2))) |> gpu)
	xtilde = G(z, ny)
	# ϵx = reshape(ϵ, 1, 1, 1, length(ϵ))
	# xhat = ϵx .* xtilde + (1f0 .- ϵx) .* x
	# yhat = ϵ .* ny + (1f0 .- ϵ) .* y
	loss = mean(D(xtilde, ny) .- D(x, y))
	# if t.gp_type == :approx
	# 	loss += t.λ*approx_penalty(D, x, xhat, y, yhat)
	# else
	# 	loss += t.λ*gradient_penalty(D, xhat, yhat)
	# end
end

Lᴳ(t::WLossGP, G, D, z, ny) = -mean(D(G(z, ny), ny))

## hinge
struct HingeLoss end

function Lᴰ(t::HingeLoss, G, D, z, ny, x, y)
	real_loss = mean(relu.(1f0 .- D(x, y)))
	fake_loss = mean(relu.(1f0 .+ D(G(z, ny), ny)))
	return real_loss + fake_loss
end

Lᴳ(t::HingeLoss, G, D, z, ny) = -mean(D(G(z, ny), ny))

## Training Code
function train_discriminator!(loss, G, D, z, ny, x, y, optD)
    θ = Flux.params(D)
    loss, back = Flux.pullback(() -> Lᴰ(loss, G, D, z, ny, x, y), θ)
    update!(optD, θ, back(1f0))
    loss
end

function orthogonal_regularization(G, β=1f-4)
    reg = 0f0
    for l in Zygote.ignore(()->filter((l)->hasproperty(l, :weight), vcat([[l...] for l in Flux.trainable(G)]...)))
		W = reshape(l.weight, :, size(l.weight)[end])
        prod = W' * W
        mat = Zygote.ignore(() -> ones(Float32, size(prod)...) .- Matrix{Float32}(I, size(prod)...) |> gpu)
        reg += norm(prod .* mat)^2
    end
    β*reg
end

function train_generator!(loss, G, D, z, ny, optG)
	θ = Flux.params(G)
	loss, back = Flux.pullback(() -> Lᴳ(loss, G, D, z, ny) + orthogonal_regularization(G), θ)
	update!(optG, θ, back(1f0))
	loss
end

function to_image(G, fixed_noise, fixed_labels, s)
    fake_images = cpu.(G.(fixed_noise, fixed_labels))
    image_array = permutedims(dropdims(reduce(vcat, reduce.(hcat, partition(fake_images, s.output_y))); dims=(3, 4)), (2, 1))
    image_array = map(clamp01nan, Gray.(image_array .+ 1f0) ./ 2f0)
    return image_array
end

function train(s::Settings)
	G, D = s.G(s), s.D(s) # Build the models
	data, fixed_noise, fixed_labels = s.img_fun(s) # Load the data

    # Training
	Dhist, Ghist = [], []
	for e=1:s.epochs
		@info("Epoch $(e)!")
		avg_loss_D, avg_loss_G, steps = 0, 0, 0
		for (x, y) in data
			loss_D = train_discriminator!(s.loss, G, D, s.rand_input(s)..., x, y, s.optD)
			loss_G = train_generator!(s.loss, G, D, s.rand_input(s)..., s.optG)

			push!(Dhist, loss_D)
			push!(Ghist, loss_G)
			avg_loss_D += loss_D
			avg_loss_G += loss_G
			steps += 1
		end
		avg_loss_D /= steps
		avg_loss_G /= steps
		if e % s.verbose_freq == 0
			@info("Discriminator loss = $avg_loss_D, Generator loss = $avg_loss_G)")
			name = @sprintf("cgan_epochs_%06d.png", e)
			save(string(s.output_dir, "/", name), to_image(G, fixed_noise, fixed_labels, s))
		end
    end
	G, D, Ghist, Dhist
end

max_stars_repo_name: sisl/VerifyGAN

max_issues_repo_name: sisl/VerifyGAN

max_forks_repo_name: sisl/VerifyGAN


text:
using DataVisualization
using Test

include("preprocessing.jl")


text:
function simulate_data(m::AbstractDSGEModel,
                       para_init::OrderedDict{Symbol, Vector{Float64}};
                       filepath::String = "sim_data.h5",
                       burnin::Int = 1, n_periods::Int = 200,
                       s_0::Vector{Float64} = Vector{Float64}(undef, 0),
                       P_0::Matrix{Float64} = Matrix{Float64}(undef, 0, 0),
                       meas_error::Float64 = 0.)

    defined_keys = para_init.keys[[isassigned(para_init.keys, i) for i in 1:length(para_init.keys)]]
    defined_vals = [para_init[key] for key in defined_keys]

    # Ensure that para_init is well-formed
    n_para      = length(defined_keys)
    n_para_sets = length(defined_vals[1])
    @assert all(n_para_sets .== map(length, defined_vals[2:end])) "Must give provide the same number of values for each parameter in para_init"

    for i in 1:n_para_sets
        para_init_single = Dict{Symbol, Float64}(zip(defined_keys, [defined_vals[j][i] for j in 1:n_para]))
        simulate_data(m; para_init = para_init_single, filepath = filepath,
                      burnin = burnin, n_periods = n_periods,
                      s_0 = s_0, P_0 = P_0, meas_error = meas_error)
    end
end

function simulate_data(m::AbstractDSGEModel;
                       para_init::OrderedDict{Symbol, Float64} = OrderedDict{Symbol, Float64}(),
                       filepath::String = "sim_data.h5",
                       burnin::Int = 1, n_periods::Int = 200,
                       s_0::Vector{Float64} = Vector{Float64}(undef, 0),
                       P_0::Matrix{Float64} = Matrix{Float64}(undef, 0, 0),
                       meas_error::Float64 = 0.)

    # Set initial parameters
    if !isempty(para_init)
        for (key, value) in para_init
            m[key] = value
        end
    end

    answer = ""
    if isfile(filepath)
        invalid_answer = true
        while invalid_answer
            println("WARNING: $filepath already exists. Do you want to overwrite it with a new dataset? (y/n)")
            answer = readline(stdin)
            if (answer == "y" || answer == "n")
                invalid_answer = false
            else
                println("Invalid answer! Must be `y` or `n`")
            end
        end
    end

    # If you want to over-write, or if there is no data currently existing in the $filepath location
    if answer == "y" || !isfile(filepath)
        data = simulate_observables(m, burnin = burnin, n_periods = n_periods, s_0 = s_0, P_0 = P_0, meas_error = meas_error)

        if answer == "y"
            rm(filepath)
            h5write(filepath, "data", data)
        else
            h5write(filepath, "data", data)
        end
        println("Wrote $filepath")

        return data
    elseif answer == "n"
        println("Aborting data simulation because the user chose not to overwrite the existing dataset.")
    end
end

function simulate_observables(m::AbstractDSGEModel;
                              burnin::Int = 1, n_periods::Int = 200,
                              s_0::Vector{Float64} = Vector{Float64}(undef, 0),
                              P_0::Matrix{Float64} = Matrix{Float64}(undef, 0, 0),
                              meas_error::Float64 = 1e-3)
    system = compute_system(m; tvis = haskey(get_settings(m), :tvis_information_set))

    n_states = size(system[:TTT], 1)
    n_obs    = length(system[:DD])

    s, ϵ = simulate_states(system[:TTT], system[:RRR], system[:CCC], system[:QQ],
                           burnin = burnin, n_periods = n_periods,
                           s_0 = s_0, P_0 = P_0)

    meas_error_dist = DegenerateMvNormal(zeros(length(system[:DD])),
                                         meas_error*Matrix{Float64}(I, length(system[:DD]), length(system[:DD])))

    y = Matrix{Float64}(undef, n_obs, n_periods)
    for i in 1:n_periods
        y[:, i] = system[:DD] + system[:ZZ]*s[:, i + burnin] + rand(meas_error_dist)
    end

    return y
end

function simulate_states(TTT::Matrix{Float64},
                         RRR::Matrix{Float64},
                         CCC::Vector{Float64},
                         QQ::Matrix{Float64};
                         burnin::Int = 5000,
                         n_periods::Int = 200,
                         ϵ::Matrix{Float64} = Matrix{Float64}(undef, 0, 0),
                         s_0::Vector{Float64} = Vector{Float64}(undef, 0),
                         P_0::Matrix{Float64} = Matrix{Float64}(undef, 0, 0))

    n_states = size(TTT, 1)
    n_shocks = size(RRR, 2)

    # Initializing shocks
    if isempty(ϵ)
        # Because we parameterize DegenerateMvNormal distributions
        # with the σ not the var
        ϵ = rand(DegenerateMvNormal(zeros(n_shocks), sqrt.(QQ)), burnin + n_periods)
    end

    # Only need shocks starting from t = 1
    # since the states will be 0 indexed, starting from the initial state s_0
    # but will generate one extra shock for the 0 index period and not use
    # it just for convenience of notation
    @assert size(ϵ) == (n_shocks, burnin + n_periods)

    # Initializing states/state covariances
    # Either specify both initial states and state covariances or neither
    @assert (isempty(s_0) && isempty(P_0)) || (!isempty(s_0) && !isempty(P_0))

    if isempty(s_0) && isempty(P_0)
        s_0, P_0 = init_stationary_states(TTT, RRR, CCC, QQ)
    end
    s = Matrix{Float64}(undef, n_states, burnin + n_periods)
    s[:, 1] = s_0

    for i in 2:(burnin + n_periods)
        s[:, i] = iterate_state(s[:, i-1], TTT, RRR, CCC, ϵ[:, i])
    end

    return s, ϵ
end

# Iterates states from s_t to s_{t+1}
function iterate_state(s::Vector{Float64}, TTT::Matrix{Float64}, RRR::Matrix{Float64},
                       CCC::Vector{Float64}, ϵ::Vector{Float64})
    return CCC + TTT*s + RRR*ϵ
end


text:
include("plotter_setup.jl")

using LightGraphs
using Expokit

##

n = 101
t = 22

g = PathGraph(n)
a = adjacency_matrix(g)
l = laplacian_matrix(g)

init_state = zeros(Float64, n)
init_state[div((n+1),2)] = 1.
##

classical_dist = expmv(t, -l, init_state)
@assert sum(classical_dist) ≈ 1.
quantum_dist = abs2.(expmv(t, 1im*a, init_state))
@assert sum(quantum_dist) ≈ 1.

##
using PyPlot
rc("text", usetex=true)
rc("font", family="serif")

plot_style= (markersize=1)

f = figure(figsize=(2.5,2))
bound = div(n, 2)
plot(-bound:bound, classical_dist,"ko-", markersize=2,linewidth=.4)
ylim(-0.005,0.085)
ylabel("probability")
xlabel("position")
savefig("plots/introduction/classical.pdf", bbox_inches="tight", metadata = Dict("CreationDate" => nothing))

f = figure(figsize=(2.5,2))
plot(-bound:bound, quantum_dist,"ko-", markersize=2,linewidth=.4)
ylim(-0.005,0.085)
yticks(visible=false)
xlabel("position")
savefig("plots/introduction/quantum.pdf", bbox_inches="tight", metadata = Dict("CreationDate" => nothing))


text:
abstract type Hamiltonian{D,N,O<:AbstractOperatorSampler} end

localdim(::Hamiltonian{D}) where {D} = D
dim(::Hamiltonian{D,N}) where {D,N} = N

struct TFIM{N,O} <: Hamiltonian{2,N,O}
    op_sampler::O
    h::Float64
    J::Float64
    P_normalization::Float64
    Ns::Int
    Nb::Int
end

function TFIM(bond_spin, Dim::Int, Ns::Int, Nb::Int, h::Float64, J::Float64)
    ops, p = make_prob_vector(bond_spin, Ns, J, h)
    op_sampler = HierarchicalOperatorSampler(ops, p)
    return TFIM{Dim, typeof(op_sampler)}(op_sampler, h, J, sum(p), Ns, Nb)
end

zero(H::Hamiltonian{2}) = falses(nspins(H))
zero(H::Hamiltonian) = zeros(nspins(H))
one(H::Hamiltonian{2}) = trues(nspins(H))
one(H::Hamiltonian) = ones(nspins(H))

nspins(H::TFIM) = H.Ns
nbonds(H::TFIM) = H.Nb


text:
abstract type DescentDirectionAllocs{T<:Real} end

mutable struct DescentDirectionAllocsPC{T<:Real} <: DescentDirectionAllocs{T}
    Δxy_aff          :: Vector{T} # affine-step solution of the augmented system [Δx_aff; Δy_aff], size nvar + ncon
    Δs_l_aff         :: Vector{T} # size nlow
    Δs_u_aff         :: Vector{T} # size nupp
    x_m_l_αΔ_aff     :: Vector{T} # x + α_aff * Δxy_aff - lvar , size nlow
    u_m_x_αΔ_aff     :: Vector{T} # uvar - (x + α_aff * Δxy_aff) , size nupp
    s_l_αΔ_aff       :: Vector{T} # s_l + α_aff * Δs_l_aff , size nlow
    s_u_αΔ_aff       :: Vector{T} # s_u + α_aff * Δs_u_aff , size nupp
    rxs_l            :: Vector{T} # - σ * μ * e + ΔX_aff * Δ_S_l_aff , size nlow
    rxs_u            :: Vector{T} # σ * μ * e + ΔX_aff * Δ_S_u_aff , size nupp
end

DescentDirectionAllocsPC(id :: QM_IntData, T :: DataType) =
    DescentDirectionAllocsPC{T}(zeros(T, id.nvar+id.ncon), # Δxy_aff
                                zeros(T, id.nlow), # Δs_l_aff
                                zeros(T, id.nupp), # Δs_u_aff
                                zeros(T, id.nlow), # x_m_l_αΔ_aff
                                zeros(T, id.nupp), # u_m_x_αΔ_aff
                                zeros(T, id.nlow), # s_l_αΔ_aff
                                zeros(T, id.nupp), # s_u_αΔ_aff
                                zeros(T, id.nlow), # rxs_l
                                zeros(T, id.nupp)  # rxs_u
                                )

convert(::Type{<:DescentDirectionAllocs{T}}, dda :: DescentDirectionAllocsPC{T0}) where {T<:Real, T0<:Real} =
    DescentDirectionAllocsPC(convert(Array{T}, dda.Δxy_aff),
                             convert(Array{T}, dda.Δs_l_aff),
                             convert(Array{T}, dda.Δs_u_aff),
                             convert(Array{T}, dda.x_m_l_αΔ_aff),
                             convert(Array{T}, dda.u_m_x_αΔ_aff),
                             convert(Array{T}, dda.s_l_αΔ_aff),
                             convert(Array{T}, dda.s_u_αΔ_aff),
                             convert(Array{T}, dda.rxs_l),
                             convert(Array{T}, dda.rxs_u))

function update_pt_aff!(x_m_l_αΔ_aff, u_m_x_αΔ_aff, s_l_αΔ_aff, s_u_αΔ_aff, Δxy_aff, Δs_l_aff, Δs_u_aff, x_m_lvar, uvar_m_x,
                        s_l, s_u, α_aff_pri, α_aff_dual, ilow, iupp)

    x_m_l_αΔ_aff .= @views x_m_lvar .+ α_aff_pri .* Δxy_aff[ilow]
    u_m_x_αΔ_aff .= @views uvar_m_x .- α_aff_pri .* Δxy_aff[iupp]
    s_l_αΔ_aff .= s_l .+ α_aff_dual .* Δs_l_aff
    s_u_αΔ_aff .= s_u .+ α_aff_dual .* Δs_u_aff
end

# Mehrotra's Predictor-Corrector algorithm
function update_dd!(dda :: DescentDirectionAllocsPC{T}, pt :: Point{T}, itd :: IterData{T}, fd :: Abstract_QM_FloatData{T},
                    id :: QM_IntData, res :: Residuals{T}, pad :: PreallocatedData{T}, cnts :: Counters, T0 :: DataType) where {T<:Real}

    # solve system aff
    dda.Δxy_aff[1:id.nvar] .= .- res.rc
    dda.Δxy_aff[id.nvar+1:end] .= .-res.rb
    dda.Δxy_aff[id.ilow] .+= pt.s_l
    dda.Δxy_aff[id.iupp] .-= pt.s_u
    out = solver!(pad, dda, pt, itd, fd, id, res, cnts, T0, :aff)
    out == 1 && return out
    dda.Δs_l_aff .= @views .-pt.s_l .- pt.s_l .* dda.Δxy_aff[id.ilow] ./ itd.x_m_lvar
    dda.Δs_u_aff .= @views .-pt.s_u .+ pt.s_u .* dda.Δxy_aff[id.iupp] ./ itd.uvar_m_x

    α_aff_pri, α_aff_dual = compute_αs(pt.x, pt.s_l, pt.s_u, fd.lvar, fd.uvar, dda.Δxy_aff, dda.Δs_l_aff,
                                       dda.Δs_u_aff, id.nvar)

    # (x-lvar, uvar-x, s_l, s_u) .+= α_aff * Δ_aff
    update_pt_aff!(dda.x_m_l_αΔ_aff, dda.u_m_x_αΔ_aff, dda.s_l_αΔ_aff, dda.s_u_αΔ_aff, dda.Δxy_aff,
                   dda.Δs_l_aff, dda.Δs_u_aff, itd.x_m_lvar, itd.uvar_m_x, pt.s_l, pt.s_u,
                   α_aff_pri, α_aff_dual, id.ilow, id.iupp)

    μ_aff = compute_μ(dda.x_m_l_αΔ_aff, dda.u_m_x_αΔ_aff, dda.s_l_αΔ_aff, dda.s_u_αΔ_aff, id.nlow, id.nupp)
    σ = (μ_aff / itd.μ)^3

    # corrector-centering step
    dda.rxs_l .= @views -σ * itd.μ .+ dda.Δxy_aff[id.ilow] .* dda.Δs_l_aff
    dda.rxs_u .= @views σ * itd.μ .+ dda.Δxy_aff[id.iupp] .* dda.Δs_u_aff
    itd.Δxy .= 0
    itd.Δxy[id.ilow] .+= dda.rxs_l ./ itd.x_m_lvar
    itd.Δxy[id.iupp] .+= dda.rxs_u ./ itd.uvar_m_x
    out = solver!(pad, dda, pt, itd, fd, id, res, cnts, T0, :cc)
    out == 1 && return out
    itd.Δs_l .= @views .-(dda.rxs_l .+ pt.s_l .* itd.Δxy[id.ilow]) ./ itd.x_m_lvar
    itd.Δs_u .= @views (dda.rxs_u .+ pt.s_u .* itd.Δxy[id.iupp]) ./ itd.uvar_m_x

    # final direction
    itd.Δxy .+= dda.Δxy_aff
    itd.Δs_l .+= dda.Δs_l_aff
    itd.Δs_u .+= dda.Δs_u_aff

    return out
end

mutable struct DescentDirectionAllocsIPF{T<:Real} <: DescentDirectionAllocs{T}
    compl :: Vector{T} # complementarity s_lᵀ(x-lvar) + s_uᵀ(uvar-x)
end

DescentDirectionAllocsIPF(id :: QM_IntData, T :: DataType) = DescentDirectionAllocsIPF{T}(zeros(T, id.nvar))

convert(::Type{<:DescentDirectionAllocs{T}}, dda :: DescentDirectionAllocsIPF{T0}) where {T<:Real, T0<:Real} =
    DescentDirectionAllocsIPF(convert(Array{T}, dda.compl))

function update_dd!(dda :: DescentDirectionAllocsIPF{T}, pt :: Point{T}, itd :: IterData{T}, fd :: Abstract_QM_FloatData{T},
                    id :: QM_IntData, res :: Residuals{T}, pad :: PreallocatedData{T}, cnts :: Counters, T0 :: DataType) where {T<:Real}

    r, γ =  T(0.999), T(0.05)
    # D = [s_l (x-lvar) + s_u (uvar-x)]
    dda.compl .= 0
    dda.compl[id.ilow] .+= pt.s_l .* itd.x_m_lvar
    dda.compl[id.iupp] .+= pt.s_u .* itd.uvar_m_x
    dda.compl[id.ifree] .= one(T)
    ξ = minimum(dda.compl) / itd.μ
    σ = γ * min((one(T) - r) * (one(T) - ξ) / ξ, T(2))^3

    itd.Δxy[1:id.nvar] .= .-res.rc
    itd.Δxy[id.nvar+1:end] .= .-res.rb
    itd.Δxy[id.ilow] .+= pt.s_l - σ * itd.μ ./ itd.x_m_lvar
    itd.Δxy[id.iupp] .-= pt.s_u - σ * itd.μ ./ itd.uvar_m_x

    out = solver!(pad, dda, pt, itd, fd, id, res, cnts, T0, :IPF)
    out == 1 && return out
    itd.Δs_l .= @views (σ * itd.μ .- pt.s_l .* itd.Δxy[id.ilow]) ./ itd.x_m_lvar .- pt.s_l
    itd.Δs_u .= @views (σ * itd.μ .+ pt.s_u .* itd.Δxy[id.iupp]) ./ itd.uvar_m_x .- pt.s_u
end


text:
# table.go
const VtableMetadataFields = 2

const TableOrBuilder = Union{Table,Builder}

const Bytes2Type = Dict{Int, DataType}(1=>UInt8, 2=>UInt16, 4=>UInt32, 8=>UInt64)

Base.get(t::TableOrBuilder, pos, ::Type{T}) where {T} = read(IOBuffer(view(t.bytes, (pos+1):length(t.bytes))), T)
readbuffer(t::AbstractVector{UInt8}, pos::Int, ::Type{T}) where {T} = read(IOBuffer(view(t, (pos+1):length(t))), T)
Base.get(t::TableOrBuilder, pos, ::Type{T}) where {T <: Enum} = T(read(IOBuffer(view(t.bytes, (pos+1):length(t.bytes))), Bytes2Type[sizeof(T)]))

"""
`offset` provides access into the Table's vtable.

Deprecated fields are ignored by checking against the vtable's length.
"""
function offset(t::Table, vtableoffset)
	vtable = t.pos - get(t, t.pos, Int32)
	return vtableoffset < get(t, vtable, Int16) ? get(t, vtable + vtableoffset, Int16) : 0
end

"`indirect` retrieves the relative offset stored at `offset`."
indirect(t::Table, off) = off + get(t, off, Int32)

"""
`vectorlen` retrieves the length of the vector whose offset is stored at
`off` in this object.
"""
function vectorlen(t::Table, off)
	off += t.pos
	off += get(t, off, Int32)
	return get(t, off, Int32)
end

"""
`vector` retrieves the start of data of the vector whose offset is stored
at `off` in this object.
"""
function vector(t::Table, off)
	off += t.pos
	off += get(t, off, Int32)
	# data starts after metadata containing the vector length
	return off + sizeof(Int32)
end

"""
GetVOffsetTSlot retrieves the VOffsetT that the given vtable location
points to. If the vtable value is zero, the default value `d`
will be returned.
"""
function getoffsetslot(t::Table, slot, d)
	off = offset(t, slot)
	if off == 0
		return d
	end
	return off
end

"""
`getslot` retrieves the `T` that the given vtable location
points to. If the vtable value is zero, the default value `d`
will be returned.
"""
function getslot(t::Table, slot, d::T) where {T}
	off = offset(t, slot)
	if off == 0
		return d
	end

	return get(t, t.pos + off, T)
end

# builder.go
value(x::T) where {T <: Enum} = length(T.types) == 0 ? Int(x) : getfield(x,1)

Base.write(sink::Builder, o, x::Union{Bool,UInt8}) = sink.bytes[o+1] = UInt8(x)
function Base.write(sink::Builder, off, x::T) where {T}
	off += 1
	for (i,ind) = enumerate(off:(off + sizeof(T) - 1))
		sink.bytes[ind] = (x >> ((i-1) * 8)) % UInt8
	end
end
Base.write(b::Builder, o, x::Float32) = write(b, o, reinterpret(UInt32, x))
Base.write(b::Builder, o, x::Float64) = write(b, o, reinterpret(UInt64, x))
Base.write(b::Builder, o, x::Enum) = write(b, o, reinterpret(Bytes2Type[sizeof(x)], value(x)))

"Offset relative to the end of the buffer."
offset(b::Builder) = length(b.bytes) - b.head

pad!(b::Builder, n) = foreach(x->place!(b, 0x00), 1:n)

"""
`finishedbytes` returns a pointer to the written data in the byte buffer.
Panics if the builder is not in a finished state (which is caused by calling
`finish!()`).
"""
function finishedbytes(b::Builder)
	assertfinished(b)
	return b.bytes[b.head+1:end]
end

function startobject(b::Builder, numslots)
	assertnotnested(b)
	b.nested = true
	b.vtable = zeros(Int, numslots)
	b.objectend = offset(b)
	b.minalign = 1
	return b
end

"""
`endobject` writes data necessary to finish object construction.
"""
function endobject(b::Builder{T}) where {T}
	assertnested(b)
	n = writevtable!(b)
	b.nested = false
	return n
end

"""
`prep!` prepares to write an element of `size` after `additionalbytes`
have been written, e.g. if you write a string, you need to align such
the int length field is aligned to sizeof(Int32), and the string data follows it
directly.
If all you need to do is align, `additionalbytes` will be 0.
"""
function prep!(b::Builder, size, additionalbytes)
	# Track the biggest thing we've ever aligned to.
	if size > b.minalign
		b.minalign = size
	end
	# Find the amount of alignment needed such that `size` is properly
	# aligned after `additionalBytes`:
	alignsize = xor(Int(-1), (length(b.bytes) - b.head) + additionalbytes) + 1
	alignsize &= (size - 1)

	# Reallocate the buffer if needed:
	totalsize = alignsize + size + additionalbytes
	if b.head <= totalsize
		len = length(b.bytes)
		prepend!(b.bytes, zeros(UInt8, totalsize))
		b.head += length(b.bytes) - len
	end
	pad!(b, alignsize)
	return
end

"""
`prepend!` prepends a `T` to the Builder buffer.
Aligns and checks for space.
"""
function Base.prepend!(b::Builder, x::T) where {T}
	prep!(b, sizeof(T), 0)
	place!(b, x)
	return
end

"""
`place!` prepends a `T` to the Builder, without checking for space.
"""
function place!(b::Builder, x::T) where {T}
	b.head -= sizeof(T)
	write(b, b.head, x)
	return
end

"""
`startvector` initializes bookkeeping for writing a new vector.

A vector has the following format:
<UOffsetT: number of elements in this vector>
<T: data>+, where T is the type of elements of this vector.
"""
function startvector(b::Builder, elemSize, numElems, alignment)
	assertnotnested(b)
	b.nested = true
	prep!(b, sizeof(UInt32), elemSize * numElems)
	prep!(b, alignment, elemSize * numElems)
	return offset(b)
end

"""
`endvector` writes data necessary to finish vector construction.
"""
function endvector(b::Builder, vectorNumElems)
	assertnested(b)
	place!(b, UInt32(vectorNumElems))
	b.nested = false
	return offset(b)
end

"""
`createstring` writes a null-terminated string as a vector.
"""
function createstring(b::Builder, s::AbstractString)
	assertnotnested(b)
	b.nested = true
	s = codeunits(s)
	prep!(b, sizeof(UInt32), length(s) + 1)
	place!(b, UInt8(0))

	l = length(s)

	b.head -= l
	copyto!(b.bytes, b.head+1, s, 1, l)
	return endvector(b, length(s))
end

"""
`createbytevector` writes a byte vector
"""
function createbytevector(b::Builder, v::AbstractVector{UInt8})
	assertnotnested(b)
	b.nested = true

	prep!(b, sizeof(UInt32), length(v))

	l = length(v)

	b.head -= l
	copyto!(b.bytes, b.head+1, v, 1, l)

	return endvector(b, length(v))
end

"""
`prependoffset!` prepends an Int32, relative to where it will be written.
"""
function prependoffset!(b::Builder, off)
	prep!(b, sizeof(Int32), 0) # Ensure alignment is already done.
	if !(off <= offset(b))
		throw(ArgumentError("unreachable: $off <= $(offset(b))"))
	end
	off2 = offset(b) - off + sizeof(Int32)
	place!(b, Int32(off2))
	return
end

function prependoffsetslot!(b::Builder, o::Int, x::T, d) where {T}
	if x != T(d)
		prependoffset!(b, x)
		slot!(b, o)
	end
	return
end

"""
`prependslot!` prepends a `T` onto the object at vtable slot `o`.
If value `x` equals default `d`, then the slot will be set to zero and no
other data will be written.
"""
function prependslot!(b::Builder, o::Int, x::T, d) where {T}
	if x != T(d)
		prepend!(b, x)
		slot!(b, o)
	end
	return
end

"""
`prependstructslot!` prepends a struct onto the object at vtable slot `o`.
Structs are stored inline, so nothing additional is being added.
In generated code, `d` is always 0.
"""
function prependstructslot!(b::Builder, voffset, x, d)
	if x != d
		assertnested(b)
		if x != offset(b)
			throw(ArgumentError("inline data write outside of object"))
		end
		slot!(b, voffset)
	end
	return
end

"""
`slot!` sets the vtable key `voffset` to the current location in the buffer.
"""
function slot!(b::Builder, slotnum)
	b.vtable[slotnum] = offset(b)
end

"""
`finish!` finalizes a buffer, pointing to the given `rootTable`.
"""
function finish!(b::Builder{T}, rootTable) where {T}
	assertnotnested(b)
	identifier = file_identifier(T)
	n = length(identifier)
	prep!(b, b.minalign, sizeof(UInt32))
	for i = 0:(n-1)
		prepend!(b, UInt8(identifier[n - i]))
	end
	prependoffset!(b, Int32(rootTable))
	b.finished = true
	return
end

function assertnested(b::Builder)
	# If you get this assert, you're in an object while trying to write
	# data that belongs outside of an object.
	# To fix this, write non-inline data (like vectors) before creating
	# objects.
	if !b.nested
		throw(ArgumentError("Incorrect creation order: must be inside object."))
	end
	return
end

function assertnotnested(b::Builder)
	# If you hit this, you're trying to construct a Table/Vector/String
	# during the construction of its parent table (between the MyTableBuilder
	# and builder.Finish()).
	# Move the creation of these view-objects to above the MyTableBuilder to
	# not get this assert.
	# Ignoring this assert may appear to work in simple cases, but the reason
	# it is here is that storing objects in-line may cause vtable offsets
	# to not fit anymore. It also leads to vtable duplication.
	if b.nested
		throw(ArgumentError("Incorrect creation order: object must not be nested."))
	end
	return
end

function assertfinished(b::Builder)
	# If you get this assert, you're attempting to get access a buffer
	# which hasn't been finished yet. Be sure to call builder.Finish()
	# with your root table.
	# If you really need to access an unfinished buffer, use the bytes
	# buffer directly.
	if !b.finished
		throw(ArgumentError("Incorrect use of FinishedBytes(): must call 'Finish' first."))
	end
end

"""
WriteVtable serializes the vtable for the current object, if applicable.

Before writing out the vtable, this checks pre-existing vtables for equality
to this one. If an equal vtable is found, point the object to the existing
vtable and return.

Because vtable values are sensitive to alignment of object data, not all
logically-equal vtables will be deduplicated.

A vtable has the following format:
<VOffsetT: size of the vtable in bytes, including this value>
<VOffsetT: size of the object in bytes, including the vtable offset>
<VOffsetT: offset for a field> * N, where N is the number of fields in
the schema for this type. Includes deprecated fields.
Thus, a vtable is made of 2 + N elements, each SizeVOffsetT bytes wide.

An object has the following format:
<SOffsetT: offset to this object's vtable (may be negative)>
<byte: data>+
"""
function writevtable!(b::Builder{T}) where {T}
	# Prepend a zero scalar to the object. Later in this function we'll
	# write an offset here that points to the object's vtable:
	prepend!(b, Int32(0))

	objectOffset = offset(b)
	existingVtable = 0

	# Search backwards through existing vtables, because similar vtables
	# are likely to have been recently appended. See
	# BenchmarkVtableDeduplication for a case in which this heuristic
	# saves about 30% of the time used in writing objects with duplicate
	# tables.
	for i = length(b.vtables):-1:1
		# Find the other vtable, which is associated with `i`:
		vt2Offset = b.vtables[i]
		vt2Start = length(b.bytes) - vt2Offset
		vt2Len = readbuffer(b.bytes, vt2Start, Int16)

		metadata = VtableMetadataFields * sizeof(Int16)
		vt2End = vt2Start + vt2Len
		vt2 = view(b.bytes, (vt2Start + metadata + 1):vt2End) #TODO: might need a +1 on the start of range here

		# Compare the other vtable to the one under consideration.
		# If they are equal, store the offset and break:
		if vtableEqual(b.vtable, objectOffset, vt2)
			existingVtable = vt2Offset
			break
		end
	end

	if existingVtable == 0
		# Did not find a vtable, so write this one to the buffer.

		# Write out the current vtable in reverse , because
		# serialization occurs in last-first order:
		for i = length(b.vtable):-1:1
			off::Int16 = 0
			if b.vtable[i] != 0
				# Forward reference to field;
				# use 32bit number to assert no overflow:
				off = objectOffset - b.vtable[i]
			end
			prepend!(b, Int16(off))
		end

		# The two metadata fields are written last.

		# First, store the object bytesize:
		objectSize::Int16 = objectOffset - b.objectend
		prepend!(b, objectSize)

		# Second, store the vtable bytesize:
		vbytes::Int16 = (length(b.vtable) + VtableMetadataFields) * sizeof(Int16)
		prepend!(b, vbytes)

		# Next, write the offset to the new vtable in the
		# already-allocated SOffsetT at the beginning of this object:
		objectStart::Int32 = length(b.bytes) - objectOffset
		write(b, objectStart, Int32(offset(b) - objectOffset))

		# Finally, store this vtable in memory for future
		# deduplication:
		push!(b.vtables, offset(b))
	else
		# Found a duplicate vtable.

		objectStart = length(b.bytes) - objectOffset
		b.head = objectStart

		# Write the offset to the found vtable in the
		# already-allocated SOffsetT at the beginning of this object:
		write(b, b.head, Int32(existingVtable - objectOffset))
	end

	empty!(b.vtable)
	return objectOffset
end

"vtableEqual compares an unwritten vtable to a written vtable."
function vtableEqual(a::Vector{Int}, objectStart, b::AbstractVector{UInt8})
	if length(a) * sizeof(Int16) != length(b)
		return false
	end

	for i = 0:(length(a)-1)
		x = read(IOBuffer(view(b, (i * sizeof(Int16) + 1):length(b))), Int16)

		# Skip vtable entries that indicate a default value.
		x == 0 && a[i+1] == 0 && continue

		y = objectStart - a[i+1]
		x != y && return false
	end
	return true
end


text:
using BinaryProvider
using Compat

const verbose = "--verbose" in ARGS
const prefix = Prefix(joinpath(@__DIR__, "usr"))

# BinaryProvider v0.3.2 has libdir(prefix) bug on windows
@static if Compat.Sys.iswindows()
    file_path = joinpath(prefix.path, "lib", string("pa_ringbuffer_", Sys.ARCH, "-w64-mingw32.dll"))
    product = FileProduct(file_path, :libpa_ringbuffer)
else
    product = LibraryProduct(prefix, "pa_ringbuffer", :libpa_ringbuffer)
end
satisfied(product; verbose=verbose) && write_deps_file(joinpath(@__DIR__, "deps.jl"), [product])

max_stars_repo_name: wookay/TestRingBuffers

max_issues_repo_name: wookay/TestRingBuffers

max_forks_repo_name: wookay/TestRingBuffers


text:
function iterate!(mc::MarkovChain{T},
                  mcmc::MCMC{T},
                  n::Int64,
                  Σ::Array{Float64, 2},
                  σ::Float64;
                  condition_on_network::Bool=false,
                  event_batches::Int64=1,
                  adapt_cov::Int64=100) where T <: EpidemicModel
  if adapt_cov < 0
    @warn "`adapt_cov` argument indicates the increment in iterations in which the covariance matrix is updated and must be ≧ 0. Setting to 0 for non-adaptive Metropolis-Hastings."
    adapt_cov = 0
  end
  pmeter = Progress(n, "MCMC progress")
  use_adapted_cov = false
  local adapted_cov
  if (adapt_cov != 0 & mc.iterations > adapt_cov) && LinearAlgebra.isposdef(value(mc.cov))
    use_adapted_cov = true
    adapted_cov = OnlineStats.value(mc.cov) * 2.38^2 / length(mc.risk_parameters[1])
  end
  tr_cache = TransmissionRateCache(mcmc.population.individuals)
  for i = 1:n
    if use_adapted_cov
      update!(mc, mcmc, adapted_cov, σ, condition_on_network = condition_on_network, event_batches = event_batches, transmission_rate_cache=tr_cache)
    else
      update!(mc, mcmc, Σ, σ, condition_on_network = condition_on_network, event_batches = event_batches, transmission_rate_cache=tr_cache)
    end
    next!(pmeter)
    @logmsg LogLevel(-5000) "MCMC progress" progress = i/n
    if adapt_cov > 0 && mod(i, adapt_cov) == 0
      OnlineStats.fit!(mc.cov, eachrow(convert(Array{Float64, 2}, mc.risk_parameters[end-adapt_cov:end])))
      if LinearAlgebra.isposdef(value(mc.cov))
        if !use_adapted_cov
          use_adapted_cov = true
          @debug "Now using adapted covariance matrix for adaptive MCMC on core $(Distributed.myid()) (updated every $adapt_cov iterations)"
        end
        adapted_cov = OnlineStats.value(mc.cov) * 2.38^2 / length(mc.risk_parameters[1])
        @debug "Covariance matrix updated for Adaptive Metropolis-Hastings MCMC" Covariance = adapted_cov
      end
    end
  end
  return mc
end

function iterate!(mc::MarkovChain{T},
                  mcmc::MCMC{T},
                  n::Int64,
                  Σ::Array{Float64, 2},
                  σ::Float64,
                  progress_channel::RemoteChannel;
                  condition_on_network::Bool=false,
                  event_batches::Int64=1,
                  adapt_cov::Int64=100) where T <: EpidemicModel
  use_adapted_cov = false
  local adapted_cov
  if (adapt_cov != 0 & mc.cov.n >= adapt_cov) && LinearAlgebra.isposdef(value(mc.cov))
    use_adapted_cov = true
    adapted_cov = OnlineStats.value(mc.cov) * 2.38^2 / length(mc.risk_parameters[1])
  end
  tr_cache = TransmissionRateCache(mcmc.population.individuals)
  for i = 1:n
    if use_adapted_cov
      update!(mc, mcmc, adapted_cov, σ, condition_on_network = condition_on_network, event_batches = event_batches, transmission_rate_cache=tr_cache)
    else
      update!(mc, mcmc, Σ, σ, condition_on_network = condition_on_network, event_batches = event_batches, transmission_rate_cache=tr_cache)
    end
    put!(progress_channel, true)
    if adapt_cov != 0 && mod(i, adapt_cov) == 0
      OnlineStats.fit!(mc.cov, eachrow(convert(Array{Float64, 2}, mc.risk_parameters[end-adapt_cov:end])))
      if LinearAlgebra.isposdef(value(mc.cov))
        if !use_adapted_cov
          use_adapted_cov = true
          @debug "Now using adapted covariance matrix for adaptive MCMC on core $(Distributed.myid()) (updated every $adapt_cov iterations)"
        end
        adapted_cov = OnlineStats.value(mc.cov) * 2.38^2 / length(mc.risk_parameters[1])
        @debug "Covariance matrix updated for Adaptive Metropolis-Hastings MCMC" Covariance = adapted_cov
      end
    end
  end
  return mc
end

function iterate!(mcmc::MCMC{T},
                  n::Int64,
                  Σ::Array{Float64, 2},
                  σ::Float64;
                  condition_on_network::Bool=false,
                  event_batches::Int64=1,
                  adapt_cov::Int64=100) where T <: EpidemicModel
  if length(mcmc.markov_chains) == 1
    iterate!(mcmc.markov_chains[1], mcmc, n, Σ, σ,
                    condition_on_network = condition_on_network,
                    event_batches = event_batches,
                    adapt_cov = adapt_cov)
    return mcmc
  else
    if adapt_cov < 0
      @warn "`adapt_cov` argument indicates the increment in iterations in which the covariance matrix is updated and must be ≧ 0. Setting to 0 for non-adaptive Metropolis-Hastings."
      adapt_cov = 0
    end
    pmeter = Progress(n*length(mcmc.markov_chains), "MCMC progress")
    pchannel = RemoteChannel(()->Channel{Bool}(10), 1)
    mc_Futures = Future[]
    for mc in mcmc.markov_chains
      @debug "Starting MCMC..."
      push!(mc_Futures, @spawn iterate!(mc, mcmc, n, Σ, σ, pchannel, condition_on_network = condition_on_network, event_batches = event_batches, adapt_cov = adapt_cov))
    end
    @debug "MCMC in progress..."
    while !all(isready.(mc_Futures))
      take!(pchannel) && next!(pmeter)
    end
    close(pchannel)
    finish!(pmeter)
    mcmc.markov_chains = [fetch(i) for i in mc_Futures]
    @debug "MCMC complete..."
    return mcmc
  end
end

function iterate!(mc::MarkovChain{T},
                  mcmc::MCMC{T},
                  n::Int64,
                  σ::Float64;
                  condition_on_network::Bool=false,
                  event_batches::Int64=1,
                  adapt_cov::Int64=100) where T <: EpidemicModel
  return iterate!(mc,
                  mcmc,
                  n,
                  Matrix(LinearAlgebra.Diagonal([var(mcmc.risk_priors[i]) for i in 1:length(mcmc.risk_priors)])),
                  σ,
                  condition_on_network = condition_on_network,
                  event_batches = event_batches,
                  adapt_cov = adapt_cov)
end

function iterate!(mcmc::MCMC{T},
                  n::Int64,
                  σ::Float64;
                  condition_on_network::Bool=false,
                  event_batches::Int64=1,
                  adapt_cov::Int64=100) where T <: EpidemicModel
  return iterate!(mcmc,
                  n,
                  Matrix(LinearAlgebra.Diagonal([var(mcmc.risk_priors[i]) for i in 1:length(mcmc.risk_priors)]))/(10*length(mcmc.risk_priors)),
                  σ,
                  condition_on_network = condition_on_network,
                  event_batches = event_batches,
                  adapt_cov = adapt_cov)
end


text:
# Code for a generic orthogonal polynomial, defined in terms of its recurrence coefficients

"""
A generic orthogonal polynomial sequence is determined by its recurrence
coefficients. The `GenericOPS` type stores the coefficients `A_n`, `B_n` and
`C_n` from the recurrence relation in the following form:
```
p_{n+1}(x) = (A_n x + B_n) * p_n(x) - C_n * p_{n-1}(x).
p_{-1} = 0, p_0 = p0
```
"""
struct GenericOPS{T} <: OPS{T}
    moment  ::  T
    p0      ::  T
    rec_a   ::  Vector{T}
    rec_b   ::  Vector{T}
    rec_c   ::  Vector{T}
    support
    weight

    function GenericOPS{T}(moment, rec_a, rec_b, rec_c, support, p0=one(T), weight=nothing) where {T}
        @assert length(rec_a) == length(rec_b) == length(rec_c)
        new(moment, p0, rec_a, rec_b, rec_c, support, weight)
    end
end

GenericOPS(moment::T, rec_a::Vector{A}, rec_b::Vector{B}, rec_c::Vector{C}, support, p0=one(T), weight=nothing) where {T,A,B,C} =
    GenericOPS{promote_type(T,A,B,C)}(moment, rec_a, rec_b, rec_c, support, p0, weight)

function MonicOPSfromQuadrature(n, my_quadrature_rule, other...; options...)
    α, β = adaptive_stieltjes(n,my_quadrature_rule; options...)
    MonicOPSfromMonicCoefficients(α, β, other...)
end

function OrthonormalOPSfromQuadrature(n, my_quadrature_rule, other...; options...)
    α, β = adaptive_stieltjes(n+1,my_quadrature_rule; options...)
    ONPSfromMonicCoefficients(α, β, other...)
end

MonicOPSfromMonicCoefficients(α::Vector{A}, β::Vector{B}, support::Domain{T}, other...) where {T,A,B} =
    GenericOPS{promote_type(T,A,B)}(β[1], ones(T,length(α)), -α, β, support, one(T), other...)

function ONPSfromMonicCoefficients(α::Vector{A}, β::Vector{B}, support::Domain{T}, other...) where {T,A,B}
    a,b,c = monic_to_orthonormal_recurrence_coefficients(α,β)
    GenericOPS{promote_type(T,A,B)}(β[1], a, b, c, support, 1/sqrt(β[1]), other...)
end

support(b::GenericOPS) = b.support
support(b::GenericOPS, idx) = support(b)

size(b::GenericOPS) = (length(b.rec_a),)

name(b::GenericOPS) = "Generic OPS"

weightfun(b::GenericOPS, x) = (b.weight==nothing) ? error("weight not defined for this Generic OPS") : b.weight(x)

function similar(b::GenericOPS, ::Type{T}, n::Int) where {T}
    @assert n <= length(b)
    GenericOPS{T}(b.moment, b.rec_a[1:n], b.rec_b[1:n], b.rec_c[1:n], b.support, b.p0, b.weight)
end

first_moment(b::GenericOPS) = b.moment

rec_An(b::GenericOPS, n::Int) = b.rec_a[n+1]

rec_Bn(b::GenericOPS, n::Int) = b.rec_b[n+1]

rec_Cn(b::GenericOPS, n::Int) = b.rec_c[n+1]

p0(b::GenericOPS) = b.p0


text:
const default_warn_order = 14

const warn_order =
  Ref{Union{Int, Nothing}}(default_warn_order)

"""
    get_warn_order()

Return the threshold for the order of an ITensor above which
ITensors will emit a warning.

You can set the threshold with the function `set_warn_order!(N::Int)`.
"""
get_warn_order() = warn_order[]

"""
    set_warn_order!(N::Int)

After this is called, ITensor will warn about ITensor contractions
that result in ITensors above the order `N`.

This function returns the initial warning threshold (what it was
set to before this function was called).

You can get the current threshold with the function `get_warn_order(N::Int)`. You can reset to the default value with
`reset_warn_order!()`.
"""
function set_warn_order!(N::Union{Int, Nothing})
  N_init = get_warn_order()
  warn_order[] = N
  return N_init
end

"""
    reset_warn_order!()

After this is called, ITensor will warn about ITensor contractions
that result in ITensors above the default order
$default_warn_order.

This function returns the initial warning threshold (what it was
set to before this function was called).
"""
reset_warn_order!() =
  set_warn_order!(default_warn_order)

"""
    disable_warn_order!()

After this is called, ITensor will not warn about ITensor
contractions that result in large ITensor orders.

This function returns the initial warning threshold (what it was
set to before this function was called).
"""
disable_warn_order!() =
  set_warn_order!(nothing)

"""
    @disable_warn_order

Disable warning about the ITensor order in a block of code.

# Examples
```julia
@disable_warn_order begin
  A * B
end
```
"""
macro disable_warn_order(block)
  quote
    old_order = disable_warn_order!()
    r = $(esc(block))
    set_warn_order!(old_order)
    return r
  end
end

"""
    @set_warn_order

Temporarily set the order threshold for warning about the ITensor
order in a block of code.

# Examples
```julia
@set_warn_order 12 begin
  A * B
end
```
"""
macro set_warn_order(new_order, block)
  quote
    old_order = set_warn_order!($(esc(new_order)))
    r = $(esc(block))
    set_warn_order!(old_order)
    return r
  end
end

"""
    @reset_warn_order

Temporarily reset the order threshold for warning about the ITensor
order in a block of code to the default value $default_warn_order.

# Examples
```julia
@reset_warn_order begin
  A * B
end
```
"""
macro reset_warn_order(block)
  quote
    old_order = reset_warn_order!()
    r = $(esc(block))
    set_warn_order!(old_order)
    return r
  end
end

const GLOBAL_TIMER = TimerOutput()


text:
# read the output voltage
export psp_readvoltage

const READ_VOLTAGE = [0x56, 0x0d]

"""
    psp_readvoltage(io)

Read the output voltage.
"""
function psp_readvoltage(io_psp::IO)
  write(io_psp, READ_VOLTAGE)
  voltagestring = bytestring(read(io_psp,UInt8, 9))
  if  voltagestring[1] != 'V' ||
      voltagestring[4] != '.' ||
      voltagestring[7:9] != "\r\r\n"
        error("PSP Invalid Response: ",voltagestring)
  end
  parse(Float64,voltagestring[2:6])
end


text:
using MultivariateStats
using GLMakie
using Statistics
using Random
using Distributions
using LinearAlgebra

include("utils.jl")
include("decompositions.jl")

n_latent = 1
n_out = 1
n_samples = 10000
s_noise = 0.00001
sigma = 0.1

# Atest = rand(n_out,n_latent)
Atest = Matrix(I,n_out,n_latent)
data = Atest*rand(Laplace(0,sigma),n_latent,n_samples) + randn(n_out,n_samples)*s_noise

ica = fit(ICA, data, n_latent; maxiter=100000, tol=0.001)

max_stars_repo_name: comp-physics/tensor-modal-decomp

max_issues_repo_name: comp-physics/tensor-modal-decomp

max_forks_repo_name: comp-physics/tensor-modal-decomp


text:
using Requires

# TODO: GPUArrays
# TODO: MPI


text: # Chapter 7. Iteration