text:
# SF6 defined as "linear gas" or gas 4 in PAGE 2009; equations are the same as for gas 3 (CH4) in PAGE2002
@defcomp LGemissions begin
    region=Index()

    # global emissions
    e_globalLGemissions=Variable(index=[time],unit="Mtonne/year")
    # baseline emissions
    e0_baselineLGemissions=Parameter(index=[region],unit="Mtonne/year")
    # regional emissions
    e_regionalLGemissions=Variable(index=[time,region],unit="Mtonne/year")
    # growth rate by region
    er_LGemissionsgrowth=Parameter(index=[time,region],unit="%")

    function run_timestep(p, v, d, t)

        #eq.4 in Hope (2006) - regional LG emissions as % change from baseline
        for r in d.region
            v.e_regionalLGemissions[t,r]=p.er_LGemissionsgrowth[t,r]*p.e0_baselineLGemissions[r]/100
        end

        #eq. 5 in Hope (2006) - global LG emissions are sum of regional emissions
        v.e_globalLGemissions[t]=sum(v.e_regionalLGemissions[t,:])
    end
end

max_stars_repo_name: lrennels/PAGE-ICE

max_issues_repo_name: lrennels/PAGE-ICE

max_forks_repo_name: lrennels/PAGE-ICE


text:
export number_neighbors, all_neighbors, random_neighbor,
    neutral_neighbors, fitter_neighbors, fitness_range_neighbors,
    fittest_neighbors, fittest_neighbor

@doc """neighbors(g::Genotype)

Returns a list of neighbors of Genotype g.
"""
# TODO: Implement all four of Nowak and Krug's methods.

function neighbors(g::Genotype)
  result = Genotype[]
  for i = 0:(g.landscape.n - 1)
    locusmask = AlleleMask(1) << i
    alleles::AlleleString = xor(g.alleles, locusmask)
    Base.push!(result,Genotype(alleles, g.landscape))
  end
  result
end

@doc """number_neighbors(g::Genotype)

Returns the number of one-mutant neighbors of a genotype.
"""
function number_neighbors(g::Genotype)
  (g.landscape.a - 1) * g.landscape.n
end

@doc """all_neighbors(g::Genotype)

Returns a vector that contains all neighbors of the given `Genotype`.
"""
function all_neighbors(g::Genotype; sorted::Bool=true)
  genotypes = map(neighbors(g)) do gt
    return gt
  end
  return if sorted
    sort!(genotypes, by=(g) -> fitness(g))
  else
    genotypes
  end
end

@doc """random_neighbor(g::Genotype)

Returns a random neighbor of the given genotype.
"""
function random_neighbor(g::Genotype)
  i = rand(0:(g.landscape.n - 1))
  locusmask = AlleleMask(1) << i
  alleles::AlleleString = xor(g.alleles, locusmask)
  return Genotype(alleles, g.landscape)
end

@doc """neutral_neighbors(g::Genotype)

Returns a vector of genotypes with fitness approximately equal to the given
genotype.
"""
function neutral_neighbors(g::Genotype)
  f0 = fitness(g)
  return filter(neighbors(g)) do gt
    isapprox(fitness(gt), f0)
  end |> collect
end

@doc """fitter_neighbors(g::Genotype; sorted::Bool=true, orequal::Bool=false)

Returns a vector of genotypes with fitness greater than the given genotype. If
`sorted` is `true`, the genotypes returned will be sorted from lowest fitness
to highest fitness. If `orequal` is `true` (the default is `false`), then
genotypes of approximately equal fitness will be included.
"""
function fitter_neighbors(g::Genotype; sorted::Bool=true, orequal::Bool=false)
  f0 = fitness(g)
  if !orequal
    genotypes = filter(neighbors(g)) do gt
      fitness(gt) > f0
    end |> collect
  else
    genotypes = filter(neighbors(g)) do gt
      fitness(gt) >= f0
    end |> collect
  end
  return if sorted
    sort!(genotypes, by=(g) -> fitness(g))
  else
    genotypes
  end
end

@doc """fitness_range_neighbors(g::Genotype, lb::Float64, ub::Float64; sorted::Bool=true)

Returns all neighbors with fitness in the half open interval [lb, ub) as
columns of a matrix, sorted from lowest fitness (left) to highest fitness
(right).
"""
function fitness_range_neighbors(g::Genotype, lb::Float64, ub::Float64; sorted::Bool=true)
  genotypes = filter(neighbors(g)) do gt
    lb <= fitness(gt) < ub
  end |> collect
  return if sorted
    sort!(genotypes, by=(g) -> fitness(g))
  else
    genotypes
  end
end

@doc """fittest_neighbors(g::Genotype, count::Int64; sorted::Bool=true)

Returns the fittest neighbors, some of which might be less fit than the given
genotype, sorted from lowest fitness to highest fitness.
"""
function fittest_neighbors(g::Genotype, count::Int64)
  @assert count > 0
  @assert count <= number_neighbors(g)

  return all_neighbors(g)[(end - count + 1):end]
end

@doc """fittest_neighbor(g::Genotype)

Returns the single fittest neighbor, even if it is less fit than the given
genotype.
"""
function fittest_neighbor(g::Genotype)
  return all_neighbors(g)[end]
end

ext: jl


text:
# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.7.0-beta3
#     language: julia
#     name: julia-1.7
# ---

# %%
#Base.findmin(f, X) = mapfoldl(x -> (f(x), x), min, X)

# %%
#ts = range(0, 2; length=201)
#findmin(sinpi, ts)

# %%
#rosenbrock2d((x, y),) = (1 - x)^2 + 100(y - x^2)^2
#xs = ys = range(-5, 5; length=1001)
#findmin(rosenbrock2d, Iterators.product(xs, ys))

# %%
struct Foo{F} f::F end
Base.findmin(foo::Foo, X) = mapfoldl(x -> (foo.f(x), x), min, X)

rosenbrock2d((x, y),) = (1 - x)^2 + 100(y - x^2)^2
xs = ys = range(-5, 5; length=1001)
findmin(Foo(rosenbrock2d), Iterators.product(xs, ys))

# %%
module My
findmin(f, X) = mapfoldl(x -> (f(x), x), min, X)
end

rosenbrock2d((x, y),) = (1 - x)^2 + 100(y - x^2)^2
xs = ys = range(-5, 5; length=1001)
My.findmin(rosenbrock2d, Iterators.product(xs, ys))

# %%
rosenbrock2d((x, y),) = (1 - x)^2 + 100(y - x^2)^2
xs = ys = range(-5, 5; length=1001)
findmin(rosenbrock2d, Iterators.product(xs, ys))

# %%
Base.keys(itr::Iterators.ProductIterator) = Iterators.product(keys.(itr.iterators)...)
findmin(rosenbrock2d, Iterators.product(xs, ys))

# %%
argmin(rosenbrock2d, Iterators.product(xs, ys))

# %%
valargmin(f, X) = (arg = argmin(f, X); (f(arg), arg))
valargmin(rosenbrock2d, Iterators.product(xs, ys))

# %%
Base.keys(itr::Iterators.ProductIterator) = Iterators.product(keys.(itr.iterators)...)
argmin(rosenbrock2d∘last, pairs(Iterators.product(xs, ys)))

# %%
Base.keys(itr::Iterators.ProductIterator) = Iterators.product(keys.(itr.iterators)...)
valargindmin(f, X) = ((ind, arg) = argmin(f∘last, pairs(X)); (f(arg), arg, ind))
valargindmin(rosenbrock2d, Iterators.product(xs, ys))

# %%

max_stars_repo_name: genkuroki/public

max_issues_repo_name: genkuroki/public

max_forks_repo_name: genkuroki/public


text:
export projGNexplicit

"""
function projGNexplicit(gc,pMis,pInv,sig,dsig,d2F,d2R,Active)

explicitly builds and solves projected normal equation.

Inputs:

	gc        - gradient
	pMis      - misfit params
	pInv      - inverse param
	sig,dsig  - current model and derivative
	d2F,d2R   - Hessians os misfit and regularizer
	Active    - indicator of active set, default false

Outputs:

	dm        - search direction
	times     - array containing time to build and solve Hessian
"""
function projGNexplicit(gc::Vector,pMis,pInv::InverseParam,sig::Vector,dsig,d2F,d2R,Active=falses(length(gc)))
	if all(Active)
		return 0*gc,[0.0;0.0]
	end
	# get Hessian of misfit
	timeBuild = @elapsed begin
	Hm = getHessian(sig,pMis,d2F)
	end;

	# build overall Hessian
	H  = dsig'*Hm*dsig + d2R

	# remove Active constraints
	timeSolve = @elapsed begin
	Hr = H[.!(Active),.!(Active)]
	gr = gc[.!(Active)]
	dm = 0*gc
	dr = -(Hr\gr)
	if any(isinf.(dr)) || any(isnan.(dr))
		dr = -(pinv(Hr)*gr)
	end
	dm[.!(Active)] = dr
	end

	# solve and return
	return dm,[timeBuild;timeSolve]
end


text:
@testset "CheckData" begin
    @testset "checkdata!" begin
        hrs = exampledata("hrs")
        nt = (data=hrs, subset=nothing, weightname=nothing)
        ret = checkdata!(nt...)
        @test ret.esample == trues(size(hrs,1))

        ismale = hrs.male.==1
        nt = merge(nt, (weightname=:rwthh, subset=ismale))
        ret = checkdata!(nt...)
        @test ret.esample == ismale
        @test ret.esample !== ismale

        df = DataFrame(hrs)
        allowmissing!(df, :rwthh)
        df.rwthh[1] = missing
        nt = merge(nt, (data=df, subset=nothing))
        ret = checkdata!(nt...)
        @test ret.esample == (hrs.rwthh.>0) .& (.!ismissing.(df.rwthh))

        nt = merge(nt, (data=rand(10,10),))
        @test_throws ArgumentError checkdata!(nt...)
        nt = merge(nt, (data=[(a=1, b=2), (a=1, b=2)],))
        @test_throws ArgumentError checkdata!(nt...)
        nt = merge(nt, (data=hrs, subset=BitArray(hrs.male[1:100])))
        @test_throws DimensionMismatch checkdata!(nt...)
        nt = merge(nt, (subset=falses(size(hrs,1)),))
        @test_throws ErrorException checkdata!(nt...)
    end

    @testset "StatsStep" begin
        @test sprint(show, CheckData()) == "CheckData"
        @test sprint(show, MIME("text/plain"), CheckData()) ==
            "CheckData (StatsStep that calls DiffinDiffsBase.checkdata!)"

        @test _f(CheckData()) == checkdata!

        hrs = exampledata("hrs")
        nt = (data=hrs, subset=nothing, weightname=nothing)
        ret = CheckData()(nt)
        @test ret.esample == trues(size(hrs,1))
        @test ret.aux isa BitVector
        @test_throws ErrorException CheckData()()
    end
end

@testset "GroupTreatintterms" begin
    @testset "grouptreatintterms" begin
        nt = (treatintterms=TermSet(),)
        @test grouptreatintterms(nt...) == nt
    end

    @testset "StatsStep" begin
        @test sprint(show, GroupTreatintterms()) == "GroupTreatintterms"
        @test sprint(show, MIME("text/plain"), GroupTreatintterms()) ==
            "GroupTreatintterms (StatsStep that calls DiffinDiffsBase.grouptreatintterms)"
        @test _byid(GroupTreatintterms()) == false
        nt = (treatintterms=TermSet(),)
        @test GroupTreatintterms()() == nt
    end
end

@testset "GroupXterms" begin
    @testset "groupxterms" begin
        nt = (xterms=TermSet(term(:x)),)
        @test groupxterms(nt...) == nt
    end

    @testset "StatsStep" begin
        @test sprint(show, GroupXterms()) == "GroupXterms"
        @test sprint(show, MIME("text/plain"), GroupXterms()) ==
            "GroupXterms (StatsStep that calls DiffinDiffsBase.groupxterms)"
        @test _byid(GroupXterms()) == false
        nt = (xterms=TermSet(),)
        @test GroupXterms()() == nt
    end
end

@testset "GroupContrasts" begin
    @testset "groupcontrasts" begin
        nt = (contrasts=Dict{Symbol,Any}(),)
        @test groupcontrasts(nt...) == nt
    end

    @testset "StatsStep" begin
        @test sprint(show, GroupContrasts()) == "GroupContrasts"
        @test sprint(show, MIME("text/plain"), GroupContrasts()) ==
            "GroupContrasts (StatsStep that calls DiffinDiffsBase.groupcontrasts)"
        @test _byid(GroupContrasts()) == false
        nt = (contrasts=nothing,)
        @test GroupContrasts()() == nt
    end
end

@testset "CheckVars" begin
    @testset "checkvars!" begin
        hrs = exampledata("hrs")
        N = size(hrs,1)
        us = unspecifiedpr()
        tr = dynamic(:wave, -1)
        nt = (data=hrs, pr=us, yterm=term(:oop_spend),
            treatname=:wave_hosp, esample=trues(N), aux=BitVector(undef, N),
            treatintterms=TermSet(), xterms=TermSet(),
            tytr=typeof(tr), trvars=(termvars(tr)...,), tr=tr)
        @test checkvars!(nt...) == (esample=trues(N), tr_rows=trues(N))

        nt = merge(nt, (pr=nevertreated(11),))
        @test checkvars!(nt...) == (esample=trues(N), tr_rows=hrs.wave_hosp.!=11)

        nt = merge(nt, (pr=notyettreated(11),))
        @test checkvars!(nt...) == (esample=hrs.wave.!=11,
            tr_rows=(hrs.wave_hosp.!=11).&(hrs.wave.!=11))

        nt = merge(nt, (pr=notyettreated(11, 10), esample=trues(N)))
        @test checkvars!(nt...) ==
            (esample=.!(hrs.wave_hosp.∈(10,)).& .!(hrs.wave.∈((10,11),)),
            tr_rows=(.!(hrs.wave_hosp.∈((10,11),)).& .!(hrs.wave.∈((10,11),))))

        nt = merge(nt, (pr=notyettreated(10),
            esample=.!((hrs.wave_hosp.==10).&(hrs.wave.==9))))
        @test checkvars!(nt...) == (esample=(hrs.wave.<9).&(hrs.wave_hosp.<=10),
            tr_rows=(hrs.wave.<9).&(hrs.wave_hosp.<=9))

        nt = merge(nt, (pr=us, treatintterms=TermSet(term(:male)),
            xterms=TermSet(term(:white)), esample=trues(N)))
        @test checkvars!(nt...) == (esample=trues(N), tr_rows=trues(N))

        nt = merge(nt, (pr=nevertreated(11), treatintterms=TermSet(term(:male)),
            xterms=TermSet(term(:white)), esample=trues(N)))
        @test checkvars!(nt...) == (esample=trues(N), tr_rows=hrs.wave_hosp.!=11)

        df = DataFrame(hrs)
        allowmissing!(df, :male)
        df.male .= ifelse.(hrs.wave_hosp.==11, missing, df.male)

        nt = merge(nt, (data=df, pr=us, esample=trues(N)))
        @test checkvars!(nt...) == (esample=hrs.wave_hosp.!=11, tr_rows=hrs.wave_hosp.!=11)
        df.male .= ifelse.(hrs.wave_hosp.==10, missing, hrs.male)
        nt = merge(nt, (esample=trues(N),))
        @test checkvars!(nt...) == (esample=hrs.wave_hosp.!=10, tr_rows=hrs.wave_hosp.!=10)

        df.male .= ifelse.(hrs.wave_hosp.==11, missing, hrs.male)
        nt = merge(nt, (pr=nevertreated(11), esample=trues(N)))
        @test checkvars!(nt...) == (esample=trues(N), tr_rows=hrs.wave_hosp.!=11)
        df.male .= ifelse.(hrs.wave_hosp.==10, missing, df.male)
        @test checkvars!(nt...) == (esample=hrs.wave_hosp.!=10, tr_rows=hrs.wave_hosp.∈((8,9),))
        df.white = ifelse.(hrs.wave_hosp.==9, missing, df.white.+0.5)
        ret = (esample=hrs.wave_hosp.∈((8,11),), tr_rows=hrs.wave_hosp.==8)
        @test checkvars!(nt...) == ret

        df.wave_hosp = Date.(df.wave_hosp)
        @test_throws ArgumentError checkvars!(nt...)
        df.wave = Date.(df.wave)
        @test_throws ArgumentError checkvars!(nt...)
        nt = merge(nt, (pr=nevertreated(Date(11)),))
        @test_throws ArgumentError checkvars!(nt...)
        df.wave = settime(df.wave, Year(1))
        df.wave_hosp = settime(df.wave_hosp, Year(1))
        @test_throws ArgumentError checkvars!(nt...)
        df.wave_hosp = aligntime(df, :wave_hosp, :wave)
        @test checkvars!(nt...) == ret
        nt = merge(nt, (pr=notyettreated(Date(11)), esample=trues(N)))
        ret = checkvars!(nt...)
        @test ret == (esample=(df.wave_hosp.∈((Date(8),Date(11)),)).&(df.wave.!=Date(11)),
            tr_rows=(df.wave_hosp.==Date(8)).&(df.wave.!=Date(11)))

        allowmissing!(df, :wave)
        df.wave .= ifelse.((df.wave_hosp.==Date(8)).&(df.wave.∈((Date(7), Date(8)),)),
            missing, df.wave)
        nt = merge(nt, (esample=trues(N), pr=nevertreated(Date(11))))
        ret = checkvars!(nt...)
        @test ret.esample == (df.wave_hosp.∈((Date(8),Date(11)),)).& (.!(hrs.wave.∈((7,8),)))
        @test ret.tr_rows == ret.esample.&(df.wave_hosp.!=Date(11))
        nt = merge(nt, (esample=trues(N), pr=notyettreated(Date(11))))
        ret = checkvars!(nt...)
        @test ret.esample == (df.wave_hosp.∈((Date(8),Date(11)),)).& (hrs.wave.∈((9,10),))
        @test ret.tr_rows == ret.esample.&(df.wave_hosp.!=Date(11))

        df = DataFrame(hrs)
        rot = ifelse.(isodd.(df.hhidpn), 1, 2)
        df.wave_hosp = rotatingtime(rot, df.wave_hosp)
        df.wave = rotatingtime(rot, df.wave)
        e = rotatingtime((1,2), 11)
        nt = merge(nt, (data=df, pr=nevertreated(e), treatintterms=TermSet(),
            xterms=TermSet(), esample=trues(N)))
        # Check RotatingTimeArray
        @test_throws ArgumentError checkvars!(nt...)
        df.wave_hosp = settime(hrs.wave_hosp, rotation=rot)
        df.wave = settime(hrs.wave, rotation=rot)
        # Check aligntime
        @test_throws ArgumentError checkvars!(nt...)
        df.wave_hosp = aligntime(df.wave_hosp, df.wave)
        ret1 = checkvars!(nt...)
        @test ret1 == (esample=trues(N), tr_rows=hrs.wave_hosp.!=11)
        df.wave_hosp = settime(hrs.wave_hosp, start=7, rotation=ones(N))
        df.wave = settime(hrs.wave, rotation=ones(N))
        # Check the match of field type of pr.e
        @test_throws ArgumentError checkvars!(nt...)
        df.wave_hosp = settime(hrs.wave_hosp, start=7, rotation=ones(Int, N))
        df.wave = settime(hrs.wave, rotation=ones(Int, N))
        nt = merge(nt, (esample=trues(N), pr=notyettreated(e),))
        ret2 = checkvars!(nt...)
        @test ret2 == (esample=hrs.wave.!=11, tr_rows=(hrs.wave_hosp.!=11).&(hrs.wave.!=11))

        # RotatingTimeArray with time field of type Array
        df.wave_hosp = RotatingTimeArray(rot, hrs.wave_hosp)
        df.wave = RotatingTimeArray(rot, hrs.wave)
        e = rotatingtime((1,2), (10,11))
        nt = merge(nt, (esample=trues(N), pr=notyettreated(e)))
        ret3 = checkvars!(nt...)
        ret3e = (hrs.wave.!=11).&(.!((hrs.wave.==10).&(rot.==1))).&
            (.!((hrs.wave_hosp.==11).&(rot.==1)))
        @test ret3 == (esample=ret3e, tr_rows=ret3e.&
            (hrs.wave_hosp.!=11).&(.!((hrs.wave_hosp.==10).&(rot.==1))))

        df.wave = settime(Date.(hrs.wave), Year(1), rotation=rot)
        df.wave_hosp = settime(Date.(hrs.wave_hosp), Year(1), start=Date(7), rotation=rot)
        e = rotatingtime((1,2), Date(11))
        nt = merge(nt, (esample=trues(N), pr=nevertreated(e)))
        @test checkvars!(nt...) == ret1
        nt = merge(nt, (esample=trues(N), pr=notyettreated(e),))
        @test checkvars!(nt...) == ret2
    end

    @testset "StatsStep" begin
        @test sprint(show, CheckVars()) == "CheckVars"
        @test sprint(show, MIME("text/plain"), CheckVars()) ==
            "CheckVars (StatsStep that calls DiffinDiffsBase.checkvars!)"

        @test _f(CheckVars()) == checkvars!

        hrs = exampledata("hrs")
        N = size(hrs,1)
        nt = (data=hrs, tr=dynamic(:wave, -1), pr=nevertreated(11), yterm=term(:oop_spend),
            treatname=:wave_hosp, treatintterms=TermSet(), xterms=TermSet(),
            esample=trues(N), aux=BitVector(undef, N))
        gargs = groupargs(CheckVars(), nt)
        @test gargs[copyargs(CheckVars())...] == nt.esample

        @test CheckVars()(nt) ==
            merge(nt, (esample=trues(N), tr_rows=hrs.wave_hosp.!=11))
        nt = (data=hrs, tr=dynamic(:wave, -1), pr=nevertreated(11), yterm=term(:oop_spend),
            treatname=:wave_hosp, esample=trues(N), aux=BitVector(undef, N),
            treatintterms=TermSet(), xterms=TermSet())
        @test CheckVars()(nt) ==
            merge(nt, (esample=trues(N), tr_rows=hrs.wave_hosp.!=11))
        @test_throws ErrorException CheckVars()()
    end
end

@testset "GroupSample" begin
    @testset "groupsample" begin
        nt = (esample=trues(3),)
        @test groupsample(nt...) == nt
    end

    @testset "StatsStep" begin
        @test sprint(show, GroupSample()) == "GroupSample"
        @test sprint(show, MIME("text/plain"), GroupSample()) ==
            "GroupSample (StatsStep that calls DiffinDiffsBase.groupsample)"
        @test _byid(GroupSample()) == false
        nt = (esample=trues(3),)
        @test GroupSample()(nt) == nt
    end
end

@testset "MakeWeights" begin
    @testset "makeweights" begin
        hrs = exampledata("hrs")
        nt = (data=hrs, esample=trues(size(hrs,1)), weightname=nothing)
        r = makeweights(nt...)
        @test r.weights isa UnitWeights && sum(r.weights) == size(hrs,1)

        nt = merge(nt, (weightname=:rwthh,))
        r = makeweights(nt...)
        @test r.weights isa Weights && sum(r.weights) == sum(hrs.rwthh)
    end

    @testset "StatsStep" begin
        @test sprint(show, MakeWeights()) == "MakeWeights"
        @test sprint(show, MIME("text/plain"), MakeWeights()) ==
            "MakeWeights (StatsStep that calls DiffinDiffsBase.makeweights)"

        @test _f(MakeWeights()) == makeweights

        hrs = exampledata("hrs")
        nt = (data=hrs, esample=trues(size(hrs,1)))
        r = MakeWeights()(nt)
        @test r.weights isa UnitWeights && sum(r.weights) == size(hrs,1)
    end
end


text:
int_rules_4_7_3 = @theory begin

    #= ::Subsection::Closed:: =#
    #= 4.7.3*Secant*normalization*rules =#
    @apply_utils Antiderivative(~u * (~(c') * sin(~(a') + ~(b') * ~x)) ^ ~(m') * (~(d') * csc(~(a') + ~(b') * ~x)) ^ ~(n'), ~x) => (~c * sin(~a + ~b * ~x)) ^ ~m * (~d * csc(~a + ~b * ~x)) ^ ~m * Antiderivative(ActivateTrig(~u) * (~d * csc(~a + ~b * ~x)) ^ (~n - ~m), ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~m, ~n], ~x) && KnownSecantIntegrandQ(~u, ~x)
    @apply_utils Antiderivative(~u * (~(c') * cos(~(a') + ~(b') * ~x)) ^ ~(m') * (~(d') * sec(~(a') + ~(b') * ~x)) ^ ~(n'), ~x) => (~c * cos(~a + ~b * ~x)) ^ ~m * (~d * sec(~a + ~b * ~x)) ^ ~m * Antiderivative(ActivateTrig(~u) * (~d * sec(~a + ~b * ~x)) ^ (~n - ~m), ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~m, ~n], ~x) && KnownSecantIntegrandQ(~u, ~x)
    @apply_utils Antiderivative(~u * (~(c') * tan(~(a') + ~(b') * ~x)) ^ ~(m') * (~(d') * sec(~(a') + ~(b') * ~x)) ^ ~(n'), ~x) => (((~c * tan(~a + ~b * ~x)) ^ ~m * (~d * csc(~a + ~b * ~x)) ^ ~m) / (~d * sec(~a + ~b * ~x)) ^ ~m) * Antiderivative((ActivateTrig(~u) * (~d * sec(~a + ~b * ~x)) ^ (~m + ~n)) / (~d * csc(~a + ~b * ~x)) ^ ~m, ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~m, ~n], ~x) && (KnownSecantIntegrandQ(~u, ~x) && Not(IntegerQ(~m)))
    @apply_utils Antiderivative(~u * (~(c') * tan(~(a') + ~(b') * ~x)) ^ ~(m') * (~(d') * csc(~(a') + ~(b') * ~x)) ^ ~(n'), ~x) => (((~c * tan(~a + ~b * ~x)) ^ ~m * (~d * csc(~a + ~b * ~x)) ^ ~m) / (~d * sec(~a + ~b * ~x)) ^ ~m) * Antiderivative((ActivateTrig(~u) * (~d * sec(~a + ~b * ~x)) ^ ~m) / (~d * csc(~a + ~b * ~x)) ^ (~m - ~n), ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~m, ~n], ~x) && (KnownSecantIntegrandQ(~u, ~x) && Not(IntegerQ(~m)))
    @apply_utils Antiderivative(~u * (~(c') * cot(~(a') + ~(b') * ~x)) ^ ~(m') * (~(d') * sec(~(a') + ~(b') * ~x)) ^ ~(n'), ~x) => (((~c * cot(~a + ~b * ~x)) ^ ~m * (~d * sec(~a + ~b * ~x)) ^ ~m) / (~d * csc(~a + ~b * ~x)) ^ ~m) * Antiderivative((ActivateTrig(~u) * (~d * csc(~a + ~b * ~x)) ^ ~m) / (~d * sec(~a + ~b * ~x)) ^ (~m - ~n), ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~m, ~n], ~x) && (KnownSecantIntegrandQ(~u, ~x) && Not(IntegerQ(~m)))
    @apply_utils Antiderivative(~u * (~(c') * cot(~(a') + ~(b') * ~x)) ^ ~(m') * (~(d') * csc(~(a') + ~(b') * ~x)) ^ ~(n'), ~x) => (((~c * cot(~a + ~b * ~x)) ^ ~m * (~d * sec(~a + ~b * ~x)) ^ ~m) / (~d * csc(~a + ~b * ~x)) ^ ~m) * Antiderivative((ActivateTrig(~u) * (~d * csc(~a + ~b * ~x)) ^ (~m + ~n)) / (~d * sec(~a + ~b * ~x)) ^ ~m, ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~m, ~n], ~x) && (KnownSecantIntegrandQ(~u, ~x) && Not(IntegerQ(~m)))
    @apply_utils Antiderivative(~u * (~(c') * sin(~(a') + ~(b') * ~x)) ^ ~(m'), ~x) => (~c * csc(~a + ~b * ~x)) ^ ~m * (~c * sin(~a + ~b * ~x)) ^ ~m * Antiderivative(ActivateTrig(~u) / (~c * csc(~a + ~b * ~x)) ^ ~m, ~x) <-- FreeQ([~a, ~b, ~c, ~m], ~x) && (Not(IntegerQ(~m)) && KnownSecantIntegrandQ(~u, ~x))
    @apply_utils Antiderivative(~u * (~(c') * cos(~(a') + ~(b') * ~x)) ^ ~(m'), ~x) => (~c * cos(~a + ~b * ~x)) ^ ~m * (~c * sec(~a + ~b * ~x)) ^ ~m * Antiderivative(ActivateTrig(~u) / (~c * sec(~a + ~b * ~x)) ^ ~m, ~x) <-- FreeQ([~a, ~b, ~c, ~m], ~x) && (Not(IntegerQ(~m)) && KnownSecantIntegrandQ(~u, ~x))
    @apply_utils Antiderivative(~u * (~(c') * tan(~(a') + ~(b') * ~x)) ^ ~(m'), ~x) => (((~c * tan(~a + ~b * ~x)) ^ ~m * (~c * csc(~a + ~b * ~x)) ^ ~m) / (~c * sec(~a + ~b * ~x)) ^ ~m) * Antiderivative((ActivateTrig(~u) * (~c * sec(~a + ~b * ~x)) ^ ~m) / (~c * csc(~a + ~b * ~x)) ^ ~m, ~x) <-- FreeQ([~a, ~b, ~c, ~m], ~x) && (Not(IntegerQ(~m)) && KnownSecantIntegrandQ(~u, ~x))
    @apply_utils Antiderivative(~u * (~(c') * cot(~(a') + ~(b') * ~x)) ^ ~(m'), ~x) => (((~c * cot(~a + ~b * ~x)) ^ ~m * (~c * sec(~a + ~b * ~x)) ^ ~m) / (~c * csc(~a + ~b * ~x)) ^ ~m) * Antiderivative((ActivateTrig(~u) * (~c * csc(~a + ~b * ~x)) ^ ~m) / (~c * sec(~a + ~b * ~x)) ^ ~m, ~x) <-- FreeQ([~a, ~b, ~c, ~m], ~x) && (Not(IntegerQ(~m)) && KnownSecantIntegrandQ(~u, ~x))
    @apply_utils Antiderivative(~u * (~(c') * sec(~(a') + ~(b') * ~x)) ^ ~(n') * (~A + ~(B') * cos(~(a') + ~(b') * ~x)), ~x) => ~c * Antiderivative(ActivateTrig(~u) * (~c * sec(~a + ~b * ~x)) ^ (~n - 1) * (~B + ~A * sec(~a + ~b * ~x)), ~x) <-- FreeQ([~a, ~b, ~c, ~A, ~B, ~n], ~x) && KnownSecantIntegrandQ(~u, ~x)
    @apply_utils Antiderivative(~u * (~(c') * csc(~(a') + ~(b') * ~x)) ^ ~(n') * (~A + ~(B') * sin(~(a') + ~(b') * ~x)), ~x) => ~c * Antiderivative(ActivateTrig(~u) * (~c * csc(~a + ~b * ~x)) ^ (~n - 1) * (~B + ~A * csc(~a + ~b * ~x)), ~x) <-- FreeQ([~a, ~b, ~c, ~A, ~B, ~n], ~x) && KnownSecantIntegrandQ(~u, ~x)
    @apply_utils Antiderivative(~u * (~A + ~(B') * cos(~(a') + ~(b') * ~x)), ~x) => Antiderivative((ActivateTrig(~u) * (~B + ~A * sec(~a + ~b * ~x))) / sec(~a + ~b * ~x), ~x) <-- FreeQ([~a, ~b, ~A, ~B], ~x) && KnownSecantIntegrandQ(~u, ~x)
    @apply_utils Antiderivative(~u * (~A + ~(B') * sin(~(a') + ~(b') * ~x)), ~x) => Antiderivative((ActivateTrig(~u) * (~B + ~A * csc(~a + ~b * ~x))) / csc(~a + ~b * ~x), ~x) <-- FreeQ([~a, ~b, ~A, ~B], ~x) && KnownSecantIntegrandQ(~u, ~x)
    @apply_utils Antiderivative(~(u') * (~(c') * sec(~(a') + ~(b') * ~x)) ^ ~(n') * (~(A') + ~(B') * cos(~(a') + ~(b') * ~x) + ~(C') * cos(~(a') + ~(b') * ~x) ^ 2), ~x) => (~c) ^ 2 * Antiderivative(ActivateTrig(~u) * (~c * sec(~a + ~b * ~x)) ^ (~n - 2) * (~C + ~B * sec(~a + ~b * ~x) + ~A * sec(~a + ~b * ~x) ^ 2), ~x) <-- FreeQ([~a, ~b, ~c, ~A, ~B, ~C, ~n], ~x) && KnownSecantIntegrandQ(~u, ~x)
    @apply_utils Antiderivative(~(u') * (~(c') * csc(~(a') + ~(b') * ~x)) ^ ~(n') * (~(A') + ~(B') * sin(~(a') + ~(b') * ~x) + ~(C') * sin(~(a') + ~(b') * ~x) ^ 2), ~x) => (~c) ^ 2 * Antiderivative(ActivateTrig(~u) * (~c * csc(~a + ~b * ~x)) ^ (~n - 2) * (~C + ~B * csc(~a + ~b * ~x) + ~A * csc(~a + ~b * ~x) ^ 2), ~x) <-- FreeQ([~a, ~b, ~c, ~A, ~B, ~C, ~n], ~x) && KnownSecantIntegrandQ(~u, ~x)
    @apply_utils Antiderivative(~(u') * (~(c') * sec(~(a') + ~(b') * ~x)) ^ ~(n') * (~A + ~(C') * cos(~(a') + ~(b') * ~x) ^ 2), ~x) => (~c) ^ 2 * Antiderivative(ActivateTrig(~u) * (~c * sec(~a + ~b * ~x)) ^ (~n - 2) * (~C + ~A * sec(~a + ~b * ~x) ^ 2), ~x) <-- FreeQ([~a, ~b, ~c, ~A, ~C, ~n], ~x) && KnownSecantIntegrandQ(~u, ~x)
    @apply_utils Antiderivative(~(u') * (~(c') * csc(~(a') + ~(b') * ~x)) ^ ~(n') * (~A + ~(C') * sin(~(a') + ~(b') * ~x) ^ 2), ~x) => (~c) ^ 2 * Antiderivative(ActivateTrig(~u) * (~c * csc(~a + ~b * ~x)) ^ (~n - 2) * (~C + ~A * csc(~a + ~b * ~x) ^ 2), ~x) <-- FreeQ([~a, ~b, ~c, ~A, ~C, ~n], ~x) && KnownSecantIntegrandQ(~u, ~x)
    @apply_utils Antiderivative(~u * (~(A') + ~(B') * cos(~(a') + ~(b') * ~x) + ~(C') * cos(~(a') + ~(b') * ~x) ^ 2), ~x) => Antiderivative((ActivateTrig(~u) * (~C + ~B * sec(~a + ~b * ~x) + ~A * sec(~a + ~b * ~x) ^ 2)) / sec(~a + ~b * ~x) ^ 2, ~x) <-- FreeQ([~a, ~b, ~A, ~B, ~C], ~x) && KnownSecantIntegrandQ(~u, ~x)
    @apply_utils Antiderivative(~u * (~(A') + ~(B') * sin(~(a') + ~(b') * ~x) + ~(C') * sin(~(a') + ~(b') * ~x) ^ 2), ~x) => Antiderivative((ActivateTrig(~u) * (~C + ~B * csc(~a + ~b * ~x) + ~A * csc(~a + ~b * ~x) ^ 2)) / csc(~a + ~b * ~x) ^ 2, ~x) <-- FreeQ([~a, ~b, ~A, ~B, ~C], ~x) && KnownSecantIntegrandQ(~u, ~x)
    @apply_utils Antiderivative(~u * (~A + ~(C') * cos(~(a') + ~(b') * ~x) ^ 2), ~x) => Antiderivative((ActivateTrig(~u) * (~C + ~A * sec(~a + ~b * ~x) ^ 2)) / sec(~a + ~b * ~x) ^ 2, ~x) <-- FreeQ([~a, ~b, ~A, ~C], ~x) && KnownSecantIntegrandQ(~u, ~x)
    @apply_utils Antiderivative(~u * (~A + ~(C') * sin(~(a') + ~(b') * ~x) ^ 2), ~x) => Antiderivative((ActivateTrig(~u) * (~C + ~A * csc(~a + ~b * ~x) ^ 2)) / csc(~a + ~b * ~x) ^ 2, ~x) <-- FreeQ([~a, ~b, ~A, ~C], ~x) && KnownSecantIntegrandQ(~u, ~x)
    @apply_utils Antiderivative(~u * (~(A') * sec(~(a') + ~(b') * ~x) ^ ~(n') + ~(B') * sec(~(a') + ~(b') * ~x) ^ ~n1 + ~(C') * sec(~(a') + ~(b') * ~x) ^ ~n2), ~x) => Antiderivative(ActivateTrig(~u) * sec(~a + ~b * ~x) ^ ~n * (~A + ~B * sec(~a + ~b * ~x) + ~C * sec(~a + ~b * ~x) ^ 2), ~x) <-- FreeQ([~a, ~b, ~A, ~B, ~C, ~n], ~x) && (EqQ(~n1, ~n + 1) && EqQ(~n2, ~n + 2))
    @apply_utils Antiderivative(~u * (~(A') * csc(~(a') + ~(b') * ~x) ^ ~(n') + ~(B') * csc(~(a') + ~(b') * ~x) ^ ~n1 + ~(C') * csc(~(a') + ~(b') * ~x) ^ ~n2), ~x) => Antiderivative(ActivateTrig(~u) * csc(~a + ~b * ~x) ^ ~n * (~A + ~B * csc(~a + ~b * ~x) + ~C * csc(~a + ~b * ~x) ^ 2), ~x) <-- FreeQ([~a, ~b, ~A, ~B, ~C, ~n], ~x) && (EqQ(~n1, ~n + 1) && EqQ(~n2, ~n + 2))
end

max_stars_repo_path: src/rules/4 Trig functions/4.7 Miscellaneous/4.7.3 Secant normalization rules.jl


text:
function instantiate_stokes_drift()
    ∂t_uˢ(z, t) = exp(z/20) * cos(t)
    ∂t_vˢ(z, t) = exp(z/20) * cos(t)
    ∂z_uˢ(z, t) = exp(z/20) * cos(t)
    ∂z_vˢ(z, t) = exp(z/20) * cos(t)
    stokes_drift = UniformStokesDrift(∂t_uˢ=∂t_uˢ, ∂t_vˢ=∂t_vˢ,
                                      ∂z_uˢ=∂z_uˢ, ∂z_vˢ=∂z_vˢ)

    return true
end

@testset "Stokes drift" begin
    @info "Testing Stokes drift..."

    @testset "Stokes drift" begin
        @test instantiate_stokes_drift()
    end
end


text:
"""
    Blocking(data_collection)

A very minimalistic implementation of blocking collection. It is used in
`TaskGroup`.

(TODO: use channel in `TaskGroup`?)
"""
struct Blocking{T,Data,Send,Receive}
    eltype::Val{T}
    data::Data          # holds value of type T
    send::Send          # swaps value::T -> nothing
    receive::Receive    # swaps nothing -> value::T
end

function Blocking(data)
    send, receive = Reagents.channel(eltype(data), Nothing)
    return Blocking(Val(eltype(data)), data, send, receive)
end

Base.eltype(::Type{<:Blocking{T}}) where {T} = T

# Note: these are not quite enough for ordering guarantee; see `block_if_nonempty`
putting(b::Blocking) = b.send | putting(b.data)
taking(b::Blocking) = b.receive | taking(b.data)

Base.put!(b::Blocking, x) = putting(b)(convert(eltype(b), x))
Base.take!(b::Blocking) = taking(b)()

blocknothing() = Map(x -> x === nothing ? Block() : something(x))

putting(c::TreiberStack) = pushing(c)
taking(c::TreiberStack) = trypopping(c) ⨟ blocknothing()

putting(c::MSQueue) = pushing(c)
taking(c::MSQueue) = trypoppingfirst(c) ⨟ blocknothing()

blocking_treiberstack(T) = Blocking(TreiberStack{T}())


text:
using LinearAlgebra
using Scapin
using Scapin.Elasticity
using StaticArrays
using Test

function isotropic_tensors(T, d)
    s = (d * (d + 1)) ÷ 2
    I₄ = Matrix{T}(I, s, s)
    I₂ = zeros(T, s)
    I₂[begin:begin+d-1] .= one(T)
    J₄ = I₂ * transpose(I₂) / d
    K₄ = I₄ - J₄
    return I₄, J₄, K₄
end

for d = 2:3
    @testset "Hooke $(d)D" begin
        T = Float64
        (I₄, J₄, K₄) = isotropic_tensors(T, d)
        C = Hooke{d,T}(1.0, 0.3)
        C_exp = d * bulk_modulus(C) * J₄ + 2C.μ * K₄
        C_act = convert(Array, C)
        @test isapprox(C_act, C_exp, rtol = 1e-15)
    end
end

function fourier_matrix_ref(Γ::GreenOperatorHooke{d,T}, k) where {d,T}
    if d == 2
        sym = 3
        ij2i = [1, 2, 1]
        ij2j = [1, 2, 2]
    elseif d == 3
        sym = 6
        ij2i = [1, 2, 3, 2, 3, 1]
        ij2j = [1, 2, 3, 3, 1, 2]
    else
        throw(ArgumentError("d must be 2 or 3 (was $d)"))
    end

    mat = zeros(T, sym, sym)
    n = k / norm(k)
    for ij = 1:sym
        i = ij2i[ij]
        j = ij2j[ij]
        w_ij = ij <= d ? one(T) : sqrt(2 * one(T))
        for kl = 1:sym
            k = ij2i[kl]
            l = ij2j[kl]
            w_kl = kl <= d ? one(T) : sqrt(2 * one(T))
            δik_nj_nl = i == k ? n[j] * n[l] : zero(T)
            δil_nj_nk = i == l ? n[j] * n[k] : zero(T)
            δjk_ni_nl = j == k ? n[i] * n[l] : zero(T)
            δjl_ni_nk = j == l ? n[i] * n[k] : zero(T)
            aux1 = (δik_nj_nl + δil_nj_nk + δjk_ni_nl + δjl_ni_nk) / 4
            aux2 = n[i] * n[j] * n[k] * n[l] / (2 * (1 - Γ.C.ν))
            mat[ij, kl] = w_ij * w_kl * (aux1 - aux2) / Γ.C.μ
        end
    end
    mat
end

@testset "Green operator, Hooke 2D" begin
    C = Hooke{2,Float64}(5.6, 0.3)
    Γ = GreenOperatorHooke(C)
    for k_norm ∈ [0.12, 2.3, 14.5]
        for θ ∈ LinRange(0.0, 2 * π, 21)[1:end-1]
            k = @SVector [k_norm * cos(θ), k_norm * sin(θ)]
            act = fourier_matrix(Γ, k)
            exp = fourier_matrix_ref(Γ, k)

            @test all(isapprox.(act, exp, atol = 1e-15))
        end
    end
end

@testset "Green operator, Hooke 3D" begin
    C = Hooke{3,Float64}(5.6, 0.3)
    Γ = GreenOperatorHooke(C)
    for k_norm ∈ [0.12, 2.3, 14.5]
        for φ ∈ LinRange(0.0, 2 * π, 21)[1:end-1]
            for θ ∈ LinRange(0.0, π, 11)
                k = k_norm * (@SVector [sin(θ) * cos(φ), sin(θ) * sin(φ), cos(θ)])
                act = fourier_matrix(Γ, k)
                exp = fourier_matrix_ref(Γ, k)

                @test all(isapprox.(act, exp, atol = 1e-15))
            end
        end
    end
end

# @testset "Discrete Green operator [MS94], Hooke 2D" begin
#     T = Float64
#     DIM = 2
#     Γ = Hooke{T,DIM}(1.2, 0.3)
#     N = @SVector [5, 6]
#     L = @SVector [3.4, 5.6]
#     Γ_h = TruncatedGreenOperator{T,DIM}(Γ, N, L)
#     k = [(2π / L[1]) * [zero(T), 1, 2, -2, -1], (2π / L[2]) * [zero(T), 1, 2, -3, -2, -1]]
#     @test Γ_h.N == N
#     @test Γ_h.L == L
#     @test Γ_h.k == k
# end

# @testset "Discrete Green operator [MS94], Hooke 3D" begin
#     T = Float64
#     DIM = 3
#     Γ = Hooke{T,DIM}(1.2, 0.3)
#     N = @SVector [5, 6, 7]
#     L = @SVector [3.4, 5.6, 7.8]
#     Γ_h = TruncatedGreenOperator{T,DIM}(Γ, N, L)
#     k = [
#         (2π / L[1]) * [zero(T), 1, 2, -2, -1],
#         (2π / L[2]) * [zero(T), 1, 2, -3, -2, -1],
#         (2π / L[3]) * [zero(T), 1, 2, 3, -3, -2, -1],
#     ]
#     @test Γ_h.N == N
#     @test Γ_h.L == L
#     @test Γ_h.k == k
# end

include("test_grid.jl")
include("test_brick.jl")
include("test_bri17.jl")


text:
# Lazy structure for an array whose real and fourier part can be accessed by A.real and A.fourier
# Used to represent potentials and densities
# For normalization conventions, see PlaneWaveBasis.jl

import Base.getproperty, Base.setproperty!, Base.propertynames
import Base.*, Base.+, Base./
import Base.eltype, Base.size, Base.length

## TODO these functions are minimally optimized for the moment
## TODO this always allocates both real and fourier parts, which might be a bit wasteful
## TODO implement broadcasting. This seems to be non-trivial, so I just implement a few simple methods
## TODO implement compressed representation in Fourier space of real arrays

@enum RFA_state RFA_only_real RFA_only_fourier RFA_both

"""
A structure to facilitate manipulations of an array of real-space type T, in both real
and fourier space. Create with `from_real` or `from_fourier`, and access
with `A.real` and `A.fourier`.
"""
mutable struct RealFourierArray{T <: Real}
    basis::PlaneWaveBasis{T}
    _real::AbstractArray{T, 3}
    _fourier::AbstractArray{Complex{T}, 3}
    _state::RFA_state
end
# Type of the real part
Base.eltype(::RealFourierArray{T}) where {T} = T
Base.size(array::RealFourierArray) = size(array._real)
Base.size(array::RealFourierArray, i) = size(array._real, i)
Base.length(array::RealFourierArray) = length(array._real)

# Constructors
RealFourierArray(basis, real, fourier) = RealFourierArray(basis, real, fourier, RFA_both)
function RealFourierArray(basis::PlaneWaveBasis{T}) where {T}
    from_real(basis, zeros(T, basis.fft_size...))
end

function from_real(basis, real_part::AbstractArray{T}) where {T <: Real}
    RealFourierArray(basis, real_part, similar(real_part, complex(T)), RFA_only_real)
end
function from_fourier(basis, fourier_part::AbstractArray{T}; check_real=false) where {T <: Complex}
    if check_real
        error("Not implemented yet")
        # Go through G vectors and check c_{-G} = (c_G)' (if both G and -G are in the grid)
        # Should be reasonably fast so we can make it the default
    end
    RealFourierArray(basis, similar(fourier_part, real(T)), fourier_part, RFA_only_fourier)
end

function Base.propertynames(array::RealFourierArray, private=false)
    ret = [:basis, :real, :fourier]
    private ? append!(ret, fieldnames(RealFourierArray)) : ret
end

function Base.getproperty(A::RealFourierArray, x::Symbol)
    if x == :real
        if A._state == RFA_only_fourier
            r = G_to_r(A.basis, A._fourier)
            if eltype(A) <: Real
                A._real .= real(r)
            else
                A._real .= r
            end
            setfield!(A, :_state, RFA_both)
        end
        return A._real
    elseif x == :fourier
        if A._state == RFA_only_real
            r_to_G!(A._fourier, A.basis, complex.(A._real))
            setfield!(A, :_state, RFA_both)
        end
        return A._fourier
    else
        getfield(A, x)
    end
end
function Base.setproperty!(A::RealFourierArray, x)
    error("RealFourierArray is intended to be read-only." *
          "(This can be bypassed by `setfield!` if you really want to).")
end

# Algebraic operations
function Base.:+(A::RealFourierArray, B::RealFourierArray) where T
    from_real(A.basis, A.real + B.real)
end
function Base.:-(A::RealFourierArray, B::RealFourierArray) where T
    from_real(A.basis, A.real - B.real)
end
function Base.:*(α::Number, A::RealFourierArray)
    from_real(A.basis, α * A.real)
end
Base.:*(A::RealFourierArray, α::Number) = α * A
Base.:/(A::RealFourierArray, α::Number) = A * inv(α)


text:
export Manager, Node
export CUDD_UNIQUE_SLOTS, CUDD_CACHE_SLOTS
export initialize_cudd
export add_var, add_ith_var, add_level_var, add_const, add_constraint
export ref, deref, recursive_deref
export add_constraint, add_restrict, evaluate
export get_value, is_const, is_nonconst, read_index, node_count
export output_dot, output_stats

const Manager = DdManager
const Node = DdNode
mutable struct FILE end

function initialize_cudd()
    dd_manager = Cudd_Init(0, 0, CUDD_UNIQUE_SLOTS, CUDD_CACHE_SLOTS, 0)
    if dd_manager == C_NULL # Could not allocate memory
        throw(OutOfMemoryError())
    end
    return dd_manager
end

function add_var(mgr::Ptr{Manager})
    new_var = Cudd_addNewVar(mgr)
    if new_var == C_NULL # Could not allocate memory
        throw(OutOfMemoryError())
    end
    return new_var
end

function add_ith_var(mgr::Ptr{Manager}, i::Int)
    new_var = Cudd_addIthVar(mgr, i)
    if new_var == C_NULL # Could not allocate memory
        throw(OutOfMemoryError())
    end
    return new_var
end

function add_level_var(mgr::Ptr{Manager}, level::Int)
    new_var = Cudd_addNewVarAtLevel(mgr, level)
    if new_var == C_NULL
        throw(OutOfMemoryError())
    end
    return new_var
end

function add_const(mgr::Ptr{Manager}, c::Real)
    const_node = Cudd_addConst(mgr, c)
    if const_node == C_NULL
        throw(OutOfMemoryError())
    end
    return const_node
end

function add_constraint(mgr::Ptr{Manager}, f::Ptr{Node}, c::Ptr{Node})
    res = Cudd_addConstrain(mgr, f, c)
    if res == C_NULL
        throw(OutOfMemoryError())
    end
    return res
end

function add_restrict(mgr::Ptr{Manager}, f::Ptr{Node}, c::Ptr{Node})
    res = Cudd_addRestrict(mgr, f, c)
    if res == C_NULL
        throw(OutOfMemoryError())
    end
    return res
end

# WARNING: the assignment array starts with the variable indexed 0;
# hence, be sure to create the ith variable from index 0.
function evaluate(mgr::Ptr{Manager}, f::Ptr{Node}, assignment::Array{Cint})
    res = Cudd_Eval(mgr, f, assignment)
    if res == C_NULL
        throw(OutOfMemoryError())
    end
    return res
end

function ref(n::Ptr{Node})
    Cudd_Ref(n)
end

function deref(n::Ptr{Node})
    Cudd_Deref(n)
end

function recursive_deref(table::Ptr{Manager}, n::Ptr{Node})
    Cudd_RecursiveDeref(table, n)
end

function get_value(c::Ptr{Node})
    res = Cudd_V(c)
    if res == C_NULL
        throw(OutOfMemoryError())
    end
    return res
end

function is_const(n::Ptr{Node})
    return Cudd_IsConstant(n)
end

function is_nonconst(n::Ptr{Node})
    return Cudd_IsNonConstant(n)
end

function read_index(n::Ptr{Node})
    return Cudd_NodeReadIndex(n)
end

function node_count(mgr::Ptr{Manager})
    return Cudd_ReadNodeCount(mgr)
end

function output_dot(mgr::Ptr{Manager}, f::Ptr{Node}, filename::String)
    outfile = ccall(:fopen, Ptr{FILE}, (Cstring, Cstring), filename, "w")
    res = Cudd_DumpDot(mgr, 1, Ref(f), C_NULL, C_NULL, outfile)
    ccall(:fclose, Cint, (Ptr{FILE},), outfile)  # here returns 0 suggesting the file is successfully closed
    return res
end

function output_stats(mgr::Ptr{Manager}, filename::String)
    outfile = ccall(:fopen, Ptr{FILE}, (Cstring, Cstring), filename, "w")
    res = Cudd_PrintInfo(mgr, outfile)
    ccall(:fclose, Cint, (Ptr{FILE},), outfile)
    return res
end


text:
# */ --------------------------------------------------------------------------------------------- /* #

function print0(b::B1_type)

    println(" ")
    println(" */"*"-"^ 32*"/*")
    println(" algorithm          : " * "stand_bend")
    println(" case               : $(b.data.case)")
    println(" investment  nodes  : $(b.data.ms.I0[end])")
    println(" operational nodes  : $(b.data.ms.I[end])")
    println(" workers            : $(b.data.wrks)")
    println(" */"*"-"^ 32*"/*")

end

function print0(b::B2_type)

    println(" ")
    println(" */"*"-"^ 32*"/*")
    println(" algorithm          : " * "adapt_bend")
    println(" case               : $(b.data.case)")
    println(" investment  nodes  : $(b.data.ms.I0[end])")
    println(" operational nodes  : $(b.data.ms.I[end])")
    println(" workers            : $(b.data.wrks)")
    println(" */"*"-"^ 32*"/*")

end

# */ --------------------------------------------------------------------------------------------- /* #

function print_info(b::B1_type)

    a1 = "$(round(b.temp.Δ;digits=3))" * "0" ^ (5 - length("$(round(b.temp.Δ%1;digits=3))"))
    a2 = "$(round(sum(b.hist.T[b.hist.k,:]);digits=2))"
    str = " k =" * (" " ^ (4-length("$(b.hist.k)"))) * "$(b.hist.k), "
    str *= "δ =" * (" " ^ (7-length(a1))) * "$(a1) %, "
    str *= "t =" * (" " ^ (8-length(a2))) * "$(a2) s"
    println(str)

end

function print_info(b::B2_type)

    a1 = "$(round(b.temp.Δ;digits=3))" * "0" ^ (5 - length("$(round(b.temp.Δ%1;digits=3))"))
    a2 = "$(round(sum(b.hist.T[b.hist.k,:]);digits=2))"
    str = " k =" * (" " ^ (4-length("$(b.hist.k)"))) * "$(b.hist.k), "
    str *= "δ =" * (" " ^ (7-length(a1))) * "$(a1) %, "
    str *= "t =" * (" " ^ (8-length(a2))) * "$(a2) s"
    println(str)

end

# */ --------------------------------------------------------------------------------------------- /* #

function str_fcn1(x::Array{Float64,1},s::String)::String

    str = " "*s*" "^(9-length(s))
    for i in 1:length(x)
        str *= (" "^(7-length(string(x[i])))*string(x[i]))
    end

    return str
end

function print_summary(b)

    print_summary_a(b)
    print_summary_b(b)

end

function print_summary_a(b)

    println(" */"*"-"^ 32*"/*")
    println(" ")
    println(" */"*"-"^ 68*"/*")
    x0 = value.(b.rmp.m[:x0])
    x1 = round.(x0[:,1];digits=1)
    x2 = round.(x0[:,2:end];digits=1)
    str_tech = ["coal","coalccs","OCGT","CCGT","diesel","nuclear","pumpL","pumpH","lithium","onwind","offwind","solar"]
    int1 = " tech."*" "^9*"i1"
    println(" ")
    println(" co2 emission limit : $( b.data.case <= 1 ? "known" : "uncertain" )")
    println(" co2 emission cost  : $( b.data.case <= 2 ? "known" : "uncertain" )")
    println(" uranium cost       : $( b.data.case <= 3 ? "known" : "uncertain" )")
    println(" investment nodes   : $( b.data.ms.I0[end] )")
    println(" operational nodes  : $( b.data.ms.I[ end] )")
    println(" optimal objective  : $( round(b.hist.U[b.hist.k]*exp10(-11);digits=3) ) x 10^11 £")
    println(" ")
    println(" */"*"-"^ 68*"/*")
    println(" ")
    println(" optimal investments (GW) @ 0 years")
    println(" " * "-" ^ (length(int1)-1))
    println(int1)
    println(" " * "-" ^ (length(int1)-1))
    for p in b.data.ms.P
        println(str_fcn1(x1[p,:],str_tech[p]))
    end
    println(" "*"-"^(length(int1)-1))
    println(" ")
    println(" optimal investments (GW) @ 5 years")
    if (b.data.case<=3)
        int2 = " tech.    "
        for n in 1:size(x2)[2]
            int2 *= " "^(7-length("i$n"))*"i$n"
        end
        println(" " * "-" ^ (length(int2)-1))
        println(int2)
        println(" " * "-" ^ (length(int2)-1))
        for p in b.data.ms.P
            println(str_fcn1(x2[p,:],str_tech[p]))
        end
        println(" "*"-"^(length(int2)-1))
    else
        int2 = [" tech.    "," tech.    "," tech.    "]
        for j in 1:3
            for i in 1:Int64(size(x2)[2]/3)
                int2[j] *= " "^(7-length("i$(9*(j-1)+i)"))*"i$(9*(j-1)+i)"
            end
        end
        for j in 1:3
            println(" " * "-" ^ (length(int2[j])-1))
            println(int2[j])
            println(" " * "-" ^ (length(int2[j])-1))
            for p in b.data.ms.P
                println(str_fcn1(x2[p,9*(j-1)+1:9*j],str_tech[p]))
            end
            println(" " * "-" ^ (length(int2[j])-1))
        end
    end
    println(" ")
    println(" */"*"-"^ 68*"/*")
end

function print_summary_b(b)

    println(" ")
    println(" computational results:")
    println(" " * "-" ^ 67)
    println(" ϵ-target (%)" * " | " * "iters     time (s)" * "   " * "RMP (%)   SP (%)   Oracles (%)")
    println(" " * "-" ^ 67)
    δ = exp10(2)*(b.hist.U[1:b.hist.k].-b.hist.L[1:b.hist.k])./b.hist.U[1:b.hist.k]
    for e in [1.00, 0.10, 0.01]
        n =findmin(max.(δ.-e,0))[2]
        t  = sum(b.hist.T[1:n,:])
        tm = sum(b.hist.T[1:n,1])
        ts = sum(b.hist.T[1:n,3])
        to = sum(b.hist.T[1:n,4])
        rm = round(exp10(2)*tm/t;digits=2)
        rs = round(exp10(2)*ts/t;digits=2)
        ro = round(exp10(2)*to/t;digits=2)
        str_e = " $(round(e;digits=2))"
        str_n = "$(n)"
        str_t = "$(Int64(round(t;digits=0)))"
        str_e *= "0" ^ (5-length(str_e))
        str_n *= " " ^ (10-length(str_n))
        str_t *= " " ^ (9-length(str_t))
        str_m = "$(rm)"
        str_s = "$(rs)"
        str_o = "$(ro)"
        str_m = " " ^ (3-findfirst(isequal('.'),str_m)) * str_m
        str_s = " " ^ (4-findfirst(isequal('.'),str_s)) * str_s
        str_o = " " ^ (3-findfirst(isequal('.'),str_o)) * str_o
        str_m *= "0" ^ (5-length(str_m))
        str_s *= "0" ^ (6-length(str_s))
        str_o *= "0" ^ (5-length(str_o))
        str_m *= " " ^ (10-length(str_m))
        str_s *= " " ^ (9-length(str_s))
        println(str_e * "         | " * str_n * str_t * "  " * str_m * str_s * str_o)
    end
    println(" " * "-" ^ 67)
    println(" ")
    println(" */"*"-"^ 68*"/*")

end


text:
#dsdemo4_audio.jl: A GUI-based demonstration of delta-sigma with sound (Currently GUI-less).
using RSDeltaSigmaPort

@warn("dsdemo4_audio requires additional libraries:\n - MAT.jl\n - WAV.jl")

if !@isdefined(DSAudioDemo) #Hack to avoid continuously appending docs
@doc """# `DSAudioDemo`: Delta-Sigma Audio Demo Module

 - Set source, modulator and decimation filter paramters,
 - then click "Go."
 - Click on input or output waveforms to listen to them.
 - Click on the output spectrum to listen to the amplified error.

## TODO
 - Implement GUI, as originally intended.
""" DSAudioDemo
end

module DSAudioDemo

using RSDeltaSigmaPort
using RSDeltaSigmaPort.EasyPlot #set, cons
using RSDeltaSigmaPort: fft
using RSDeltaSigmaPort: linlin, loglin
import Printf: @sprintf
import MAT, WAV

#==Constants
===============================================================================#
const 𝑓ₛout = 8192 #Output sample rate is fixed
const j=im

"ABCD matrices of pre-defined modulators:"
const MODABCD_PRESETS = Dict{Int, Array}(
	1 => [1 1 -1; 1 0 0],
	2 => [1 0 1 -1; 1 1 1 -2; 0 1 0 0],
)

#==Types
===============================================================================#
abstract type AbstractSource; end

struct SrcRamp <: AbstractSource; end

mutable struct SrcRaw <: AbstractSource
	u::Vector #{Float64} #Oversampled signal
	u0::Vector #{Float64} #Decimated signal
end

mutable struct SrcSine <: AbstractSource
	f::Int #Frequency (Hz)
	A::Float64
end
SrcSine(;f=500, A=0.5) = SrcSine(Int(f), A)

struct Modulator
	sys
	fsout::Int #Output sample rate
	OSR::Int
	sincorder::Int
end

#==Constructors
===============================================================================#
function _Modulator(sys, 𝑓ₒₛ::Int, OSR::Nothing, sincorder::Int)
	#Round 𝑓ₒₛ to a multiple of 𝑓ₛout:
	OSR = round(Int, 𝑓ₒₛ/𝑓ₛout)
	OSR = max(1, OSR)
	𝑓ₒₛnew = 𝑓ₛout*OSR #𝑓ₛout is always hardcoded in dsdemo4

	if 𝑓ₒₛ < 𝑓ₛout
		@warn("Requested (over)-sampling frequency, 𝑓ₒₛ < 𝑓ₛout." *
			"\nSetting 𝑓ₒₛ = 𝑓ₛout = $𝑓ₛout Hz."
		)
	elseif 𝑓ₒₛ != 𝑓ₒₛnew
		@warn("Adjusting (over)-sampling frequency\nfrom 𝑓ₒₛ = $𝑓ₒₛ" *
			"\nto 𝑓ₒₛ = $𝑓ₒₛnew (OSR = $OSR)."
		)
	end

	return Modulator(sys, 𝑓ₛout, OSR, sincorder)
end
_Modulator(sys, 𝑓ₒₛ::Nothing, OSR::Int, sincorder::Int) =
	Modulator(sys, 𝑓ₛout, OSR, sincorder) #𝑓ₛout is always hardcoded in dsdemo4
_Modulator(sys, 𝑓ₒₛ::Nothing, OSR::Nothing, sincorder::Int) =
	_Modulator(sys, 𝑓ₒₛ, 32, sincorder) #Default OSR
_Modulator(sys, 𝑓ₒₛ, OSR, sincorder::Int) =
	throw(ArgumentError("Modulator: Invalid type for 𝑓ₒₛ or OSR (or overspecified)."))

Modulator(preset; fos=nothing, OSR=nothing, sincorder::Int=2) =
	_Modulator(MODABCD_PRESETS[preset], fos, OSR, sincorder)

#==Helper functions/ data generators
===============================================================================#
function validate(m::Modulator)
	(m.OSR < 1) &&
		throw("Modulator: Invalid OSR value: $(m.OSR).")
	return
end

function generate(src::SrcSine, m::Modulator; N=10)
	(src.f > m.fsout/2) &&
		throw("Source 𝑓 = $(src.f) too high. Must be <= 𝑓ₛout/2.")
	𝑓ₒₛ = m.fsout*m.OSR
	𝜔₀ = 2π*src.f
	tidx = (0:N-1); t = tidx/𝑓ₒₛ
	u = src.A * sin.((𝜔₀/𝑓ₒₛ) * tidx) .* ds_hann(N)
	u0 = u[1:m.OSR:end]
	return (t, u, u0)
end

function generate(src::SrcRamp, m::Modulator; N=10)
	𝑓ₒₛ = m.fsout*m.OSR
	tidx = (0:N-1); t = tidx/𝑓ₒₛ
	u = collect(range(-0.7, stop=0.7, length=N))
	u0 = u[1:m.OSR:end]
	return (t, u, u0)
end

function generate(src::SrcRaw, m::Modulator; N=10)
	𝑓ₒₛ = m.fsout*m.OSR
	N = length(src.u)
	tidx = (0:N-1); t = tidx/𝑓ₒₛ
	return (t, src.u, src.u0)
end

function play(data::Vector)
	WAV.wavplay(data, 𝑓ₛout)
	return :PLAYBACK_COMLETE
end

#==Main algorithms
===============================================================================#
function run(m::Modulator; Tsim::Float64=2.0, input=SrcSine(f=500))
	validate(m)
	(m.fsout != 𝑓ₛout) && throw("dsdemo4 ony supports modulation schemes where 𝑓ₛout = $𝑓ₛout.")
	fos = 𝑓ₛout*m.OSR #(over)-sampling frequency

	#Compute signal `u` depending on desired source:
	N = round(Int, Tsim*fos) #Unless overwritten
	(t, u, u0) = generate(input, m, N=round(Int, Tsim*fos))
	OSR = m.OSR #Copy for string interpolation

	#Plot signal `u0`
	plot_tdec = cons(:plot, linlin, title="Decimated Signal (dec=$OSR)", legend=false,
		xyaxes=set(ymin=-1, ymax=1),
		labels=set(xaxis="Time [s]", yaxis="u₀(t), w(t)"),
	)

	N = length(u0)
	t = (0:N-1)/𝑓ₛout
	u0w = waveform(t, u0)
	push!(plot_tdec,
		cons(:wfrm, u0w, line=set(style=:solid, color=:blue, width=2), label="u₀"),
	)

	#Plot U(𝑓), from 0 to 𝑓ₛout/2
	plot_𝑓dec = cons(:plot, linlin, title="Spectrum (dec=$OSR)", legend=false,
		xyaxes=set(xmin=0, xmax=𝑓ₛout/2, ymin=-160, ymax=0),
		labels=set(xaxis="Frequency [Hz]", yaxis="U₀(𝑓), W(𝑓)"),
	)

	if typeof(input) in [SrcSine, SrcRaw]
		N = length(u0)
		local U
		if isa(input, SrcSine)
			U = fft(u0)/(N/4)
		else #SrcRaw
			U = fft(applywnd(u0, ds_hann(N)))/(N/4)
		end
		𝑓 = range(0, stop=𝑓ₛout, length=N+1); 𝑓 = 𝑓[1:div(N,2)+1]
		UdB = dbv.(U[keys(𝑓)])
		UdBw = waveform(𝑓, UdB)
		push!(plot_𝑓dec,
			cons(:wfrm, UdBw, line=set(style=:solid, color=:blue, width=2), label="U₀"),
		)
	end

	#Plot ΔΣ signals (time domain):
	Nplot = 300 #Number of samples to plot
	simresult = simulateDSM(u, m.sys)
	v = simresult.v; y = simresult.y
	q = v - y #Quantization error; assumes quantizer gain = 1.
	N = length(v)
	n = 1:N
	if N>Nplot #Sample values from the middle:
		n = floor.(Int, N/2-Nplot/2:N/2+Nplot/2-1)
	end

	plot_tos = cons(:plot, linlin, title="Signals @ Modulator Input/Output", legend=false,
		xyaxes=set(ymin=-1.2, ymax=1.2),
		labels=set(xaxis="samples", yaxis="v(t)"),
	)

	_x = collect(0:length(n)-1)
	vw = wfrm_stairs(_x, v[n])
	uw = waveform(_x, u[n])

	push!(plot_tos,
		cons(:wfrm, vw, line=set(style=:solid, color=:blue, width=2), label="v"),
		cons(:wfrm, uw, line=set(style=:solid, color=:green, width=2), label="u"),
	)

	#Plot V(𝑓), from 0 to 𝑓ₒₛ/2. Use the middle Nfft points of v
	N = length(v)
	Nfft = min(N, 16*8192)
	n = array_round((N-Nfft)/2+1):array_round((N+Nfft)/2)
	V = fft(applywnd(v[n], ds_hann(Nfft)))/(Nfft/4)
	inBin = ceil(Int, Nfft/1000)
	if isa(input, SrcSine)
		inBin = round(Int, input.f/fos*Nfft+1) #Bin of tone
	end
	(𝑓nrm, Vp) = logsmooth(V, inBin)

	plot_𝑓os = cons(:plot, loglin, title="Spectrum (OSR=$OSR)", legend=false,
		xyaxes=set(xmin=100, xmax=fos/2, ymin=-160, ymax=0),
		labels=set(xaxis="Frequency [Hz]", yaxis="V(𝑓)"),
	)

	Vpw = waveform(𝑓nrm*fos, Vp)
	nbw_str = @sprintf("NBW = %.1f Hz", fos*1.5/Nfft) #orig. coords: (Fs/2, -90); :cr
	push!(plot_𝑓os,
		cons(:wfrm, Vpw, line=set(style=:solid, color=:blue, width=2), label="V"),
		cons(:atext, nbw_str, y=-90, reloffset=set(x=0.95), align=:cr),
	)

	#Compute w
	w = sinc_decimate(v, m.sincorder, m.OSR)
	filtered_q = sinc_decimate(q, m.sincorder, m.OSR)
	N = length(w)
	t = collect(0:N-1)/𝑓ₛout
	ww = wfrm_stairs(t, w)
	push!(plot_tdec,
		cons(:wfrm, ww, line=set(style=:solid, color=:red, width=2), label="w"),
	)

	#Plot W(𝑓), from 0 to 𝑓ₛout/2
	if typeof(input) in [SrcSine, SrcRaw]
		Nfft = length(w)
		local W
		if isa(input, SrcSine)
			W = fft(w)/(N/4)
		else
			W = fft(applywnd(w, ds_hann(N)))/(N/4)
		end
		𝑓 = range(0, stop=𝑓ₛout, length=Nfft+1); 𝑓 = 𝑓[1:div(Nfft,2)+1]
		WdB = dbv.(W[keys(𝑓)])
		WdBw = waveform(𝑓, WdB)
		nbw_str = @sprintf("NBW = %.1f Hz", 𝑓ₛout*1.5/Nfft) #orig. coords: (10, -90); :cl
		push!(plot_𝑓dec,
			cons(:wfrm, WdBw, line=set(style=:solid, color=:red, width=2), label="W"),
			cons(:atext, nbw_str, y=-90, reloffset=set(x=0.05), align=:cl),
		)
	end

	pcoll = push!(cons(:plot_collection, ncolumns=2),
		plot_tdec, plot_𝑓dec, plot_tos, plot_𝑓os,
	)

	return (plot=pcoll, u0=u0, w=w, input=input)
end

end #module DSAudioDemo

#==Auto-run test code:
===============================================================================#
println()
display(@doc(DSAudioDemo)) #Show user how to use DSAudioDemo

function load_demo4_audio_data(m::DSAudioDemo.Modulator)
	srcpath = dirname(pathof(RSDeltaSigmaPort))
	fpath = joinpath(srcpath, "..", "original_source", "delsig", "dsdemo4.mat")
	alldata = DSAudioDemo.MAT.matread(fpath)
	ds = alldata["ds"][:]
	sd = alldata["sd"][:]
	u0 = ds
	u = interp(sd, m.OSR)
	return DSAudioDemo.SrcRaw(u, u0)
end

function playresults(results)
	if isa(results.input, DSAudioDemo.SrcRamp)
		@warn("Will not playback ramp signal.")
		return
	end
	println()
	@info("Listening to ideally sampled/decimated signal...")
		flush(stdout); flush(stderr)
		@show DSAudioDemo.play(results.u0)
	println()
	@info("Listening to modulator output...")
		flush(stdout); flush(stderr)
		@show DSAudioDemo.play(results.w)

	println("\nReplay results with:")
	println("\tplayresults(results)")
	return
end

#Inputs
dsm = DSAudioDemo.Modulator(1, OSR=32, sincorder=2)
#sig = load_demo4_audio_data(dsm)
sig = DSAudioDemo.SrcSine(f=500) #500, 4000, 4200
#sig = DSAudioDemo.SrcRamp()

println("\nUsing modulator:")
@show dsm

println()
@info("Performing ΔΣ audio simulation..."); flush(stdout); flush(stderr)
results = DSAudioDemo.run(dsm, Tsim=3.0, input=sig)
println("\tdone."); flush(stdout); flush(stderr)

println()
@info("Displaying results..."); flush(stdout); flush(stderr)
displaygui(results.plot)

playresults(results)
println()

:END_OF_DEMO


text:
struct Kosarkasi
    ime::String
    prezime::String
    visina::Int
    pozicija::String
end

function zad41(A,visina)
    m = Int(floor(length(A)/2))
    rez = Array{Kosarkasi}(undef,0)
    if A[m].visina == visina
        k = length(A)
        i = 1
        while(i<m && k > m)
            if A[i].visina == visina
                push!(rez,A[i])
            end
            if A[k].visina == visina
                push!(rez,A[k])
            end
            i+=1
            k-=1
        end
    elseif A[m].visina > visina
        for j = m+1:length(A)
            if A[j].visina == visina
                push!(rez,A[j])
            end
            if(A[j] != key && br !=0)
                    break
            end
        end
    else
        for k = 1:m-1
            if A[k].visina == visina
                push!(rez,A[k])
            end
            if(A[k] != key && br !=0)
                    break
            end
        end
    end
    return br
end


text:
var_count = 0

function variables(::Type{T}; n::Int = 1, unique::Bool = true) where {T}
    if unique
        global var_count += n
        varcnt = var_count
    else
        varcnt = n
    end
    if n == 1
        return var(T, "ω$(varcnt)")
    end
    return [var(T, "ω$i") for i in varcnt-n+1:varcnt]
end

var(::Type{Symbol}, s::String) = Symbol(s)

function pascal(n; alt = false)
    f = alt ? -1 : 1
    entries = [1]
    for k in 0:n-1
        append!(entries, f * entries[k+1] * (n-k) / (k+1))
    end
    entries
end

initvar(v::T, i::Union{T, Int64}=0) where {T} = T("$(string(v))$(i)$(i)")


text:
using .FFTW
laplacianunwrap(ϕ) = laplacianunwrap!(copy(ϕ))
function laplacianunwrap!(ϕ::AbstractArray)
    FFTW.set_num_threads(Threads.nthreads())
    ϕ .+= 2π .* k(ϕ) # rounding k as suggested in the paper does not work
end

# Schofield and Zhu 2003, https://doi.org/10.1364/OL.28.001194
k(ϕw) = 1 / 2π .* ∇⁻²(∇²_nw(ϕw) - ∇²(ϕw))  # (1)

∇²(x) = -(2π)^ndims(x) / length(x) .* idct(pqterm(size(x)) .* dct(x))  # (2)

∇⁻²(x) = -length(x) / (2π)^ndims(x) .* idct(dct(x) ./ pqterm(size(x)))  # (3)

∇²_nw(ϕw) = cos.(ϕw) .* ∇²(sin.(ϕw)) .- sin.(ϕw) .* ∇²(cos.(ϕw))  # (in text)

pqterm(sz::NTuple{1}) = (1:sz[1]).^2  # 1D case
pqterm(sz::NTuple{2}) = [p^2 + q^2 for p in 1:sz[1], q in 1:sz[2]]  # 2D case
pqterm(sz::NTuple{3}) = [p^2 + q^2 + t^2 for p in 1:sz[1], q in 1:sz[2], t in 1:sz[3]]  # 3D case
pqterm(sz::NTuple{4}) = [p^2 + q^2 + t^2 + r^2 for p in 1:sz[1], q in 1:sz[2], t in 1:sz[3], r in 1:sz[4]]  # 4D case

"""
    laplacianunwrap(ϕ::AbstractArray)

Performs laplacian unwrapping on the input phase. (1D - 4D)
The phase has to be scaled to radians.
The implementation is close to the original publication: Schofield and Zhu 2003, https://doi.org/10.1364/OL.28.001194.
It is not the fastest implementation of laplacian unwrapping (doesn't use discrete laplacian).
"""
laplacianunwrap, laplacianunwrap!

# FFT variant
using OffsetArrays, ImageFiltering
function laplacianunwrap_fft(ϕ::AbstractArray, z_weight=1)
    FFTW.set_num_threads(min(4, Threads.nthreads()))

    kernel = float.(convert(AbstractArray, Kernel.Laplacian((true,true,true))))
    kernel[0,0,1] = kernel[0,0,-1] = z_weight
    kernel[0,0,0] += 2 * (1 - z_weight)

    ∇²(x) = imfilter(x, kernel)

    kernel_full = centered(zeros(size(ϕ)))
    kernel_full[CartesianIndices(kernel)] .= kernel

    del_op = fft(kernel_full)

    del_inv = 1 ./ del_op
    del_inv[.!isfinite.(del_inv)] .= 0

    del_phase = cos.(ϕ) .* ∇²(sin.(ϕ)) .- sin.(ϕ) .* ∇²(cos.(ϕ))

    unwrapped = real.(ifft( fft(del_phase) .* del_inv ))

    return unwrapped
end

function laplacianunwrap_mixed(ϕ::AbstractArray)
    FFTW.set_num_threads(min(4, Threads.nthreads()))

    kernel = Kernel.Laplacian((true,true,true))
    ∇²(x) = imfilter(x, kernel)

    del_phase = cos.(ϕ) .* ∇²(sin.(ϕ)) .- sin.(ϕ) .* ∇²(cos.(ϕ))

    unwrapped = ∇⁻²(del_phase)

    return unwrapped
end


text:
using Plots
using NPZ
using Random

include("../HCOPE.jl")
include("../HCOVE.jl")
include("../UnO_utils.jl")
include("../UnO_params.jl")
include("../UnO_CDF.jl")

function compute_trend(props,
                        data::Array{Float64, 3},
                        pi_data::Array{Float64, 1},
                        trials::Int,
                        mini::Float64,
                        maxi::Float64,
                        intervals::Int)

    ######################################################
    # Pre-allocate arrays for storing all the estimates
    n_points = 1000      # Number of support points for plotting CDFs
    α_quantile = 0.5     # For inverse CDF based stuff
    α_cvar = 0.25    # For inverse CDF based stuff
    δ = 0.05             # Failure rates

    CDF_estimates = fill(0.0, trials, intervals, n_points)
    mean_estimates = fill(0.0, trials, intervals)
    variance_estimates = fill(0.0, trials, intervals)
    quantile_estimates = fill(0.0, trials, intervals)
    cvar_estimates = fill(0.0, trials, intervals)

    CDF_bounds = fill(0.0, trials, intervals, 2, n_points)
    mean_bounds = fill(0.0, trials, intervals, 2)
    mean_baseline_bounds = fill(0.0, trials, intervals, 4)
    variance_bounds = fill(0.0, trials, intervals, 2)
    variance_baseline_bounds = fill(0.0, trials, intervals, 4)
    quantile_bounds = fill(0.0, trials, intervals, 2)
    cvar_bounds = fill(0.0, trials, intervals, 2)

    CDF_boot_bounds = fill(0.0, trials, intervals, 2, n_points)
    mean_boot_bounds = fill(0.0, trials, intervals, 2)
    variance_boot_bounds = fill(0.0, trials, intervals, 2)
    quantile_boot_bounds = fill(0.0, trials, intervals, 2)
    cvar_boot_bounds = fill(0.0, trials, intervals, 2)

    # True CDF and parameters
    true_CDF = fill(0.0, n_points)
    CDF_plot_xs = fill(0.0, n_points)
    true_params = fill(0.0, 5)      # Order: Mean, quantile, CVaR, variance, entropy

    ###################################################
    # Compute estimates and bounds for CDF and params for different trials
    ###################################################

    step = (maxi - mini) / (intervals - 1)
    full_size, _, _ = size(data)
    trend_xs = fill(0.0, intervals)

    # Run in CLI before executing code
    # Be careful, this changes arrays asynchronously
    # export JULIA_NUM_THREADS=10
    #
    for index in 1:1:intervals

        n = Int(floor( 10^(mini + (index - 1)*step) ))
        trend_xs[index] = n

        perm_idxs = shuffle(collect(1:full_size))   # Shuffle indices
        Threads.@threads for trial in 1:1:trials
            println(n, ":", trial)

            ## Without replacement sampling
            ## Ensure unique data for each trial
            # start = (trial-1)*n + 1
            # stop = n*trial
            # idxs = perm_idxs[start:stop]

            # With replacement sampling
            perm2 = shuffle(perm_idxs)
            idxs = perm2[1:n]

            fragment = data[idxs, :, :]

            ###########################################
            # Estimators for CDF and other parameters
            ###########################################

            xs, ys = IS_CDF(fragment, props)
            wis_xs, wis_ys = WIS_CDF(fragment, props)     # WIS version

            CDF_estimates[trial, index, :] = get_equispaced_CDF(xs, ys, props, n_points)[2]
            mean_estimates[trial, index] = CDF_mean(xs, ys)
            variance_estimates[trial, index] = CDF_variance(xs, ys)
            quantile_estimates[trial, index] = CDF_quantile(xs, ys, α_quantile)
            cvar_estimates[trial, index] = CDF_CVaR(xs, ys, α_cvar)

            ############################################################
            # Guaranteed coverage bounds for CDF and other parameters
            ###########################################################

            # bound_xs, lb_ys, ub_ys = CDF_band(xs, ys, props, δ)
            bound_xs, lb_ys, ub_ys = CDF_band_optimized(xs, ys, props, δ)
            CDF_bounds[trial, index, 1, :] = get_equispaced_CDF(bound_xs, lb_ys, props, n_points)[2]
            CDF_bounds[trial, index, 2, :] = get_equispaced_CDF(bound_xs, ub_ys, props, n_points)[2]

            lb, ub = CDF_mean_bound(bound_xs, lb_ys, ub_ys)
            mean_bounds[trial, index, 1] = lb
            mean_bounds[trial, index, 2] = ub

            # Divide by 8 because of two-sided bounds for 4 statistics
            mean_baseline_bounds[trial, index, 1] = CI_lower(fragment, props, δ/8)
            mean_baseline_bounds[trial, index, 2] = CI_upper(fragment, props, δ/8)
            # mean_baseline_bounds[trial, index, 3] = CI_lower(fragment, props, δ/2)  # w/o union
            # mean_baseline_bounds[trial, index, 4] = CI_upper(fragment, props, δ/2)  # w/o union

            lb, ub = CDF_variance_bound(bound_xs, lb_ys, ub_ys)
            variance_bounds[trial, index, 1] = lb
            variance_bounds[trial, index, 2] = ub

            # Divide by 8 because of two-sided bounds for 4 statistics
            variance_baseline_bounds[trial, index, 1] = HCOVE_lower(fragment, props, δ/8)
            variance_baseline_bounds[trial, index, 2] = HCOVE_upper(fragment, props, δ/8)
            # variance_baseline_bounds[trial, index, 3] = HCOVE_lower(fragment, props, δ/2)   # w/o union
            # variance_baseline_bounds[trial, index, 4] = HCOVE_upper(fragment, props, δ/2)   # w/o union

            lb, ub = CDF_quantile_bound(bound_xs, lb_ys, ub_ys, α_quantile)
            quantile_bounds[trial, index, 1] = lb
            quantile_bounds[trial, index, 2] = ub

            lb, ub = CDF_CVaR_bound(bound_xs, lb_ys, ub_ys, α_cvar)
            cvar_bounds[trial, index, 1] = lb
            cvar_bounds[trial, index, 2] = ub

            ############################
            # Bootstrap based bounds
            ############################
            lb, ub = CDF_mean_boot_bound(wis_xs, wis_ys, δ/4)
            mean_boot_bounds[trial, index, 1] = lb
            mean_boot_bounds[trial, index, 2] = ub

            lb, ub = CDF_variance_boot_bound(wis_xs, wis_ys, δ/4)
            variance_boot_bounds[trial, index, 1] = lb
            variance_boot_bounds[trial, index, 2] = ub

            lb, ub = CDF_quantile_boot_bound(wis_xs, wis_ys, α_quantile, δ/4)
            quantile_boot_bounds[trial, index, 1] = lb
            quantile_boot_bounds[trial, index, 2] = ub

            lb, ub = CDF_CVaR_boot_bound(wis_xs, wis_ys, α_cvar, δ/4)
            cvar_boot_bounds[trial, index, 1] = lb
            cvar_boot_bounds[trial, index, 2] = ub

        end;
    end;

    #######################
    # True CDF and parameters
    ########################

    xs, ys = on_policy_CDF(pi_data)
    CDF_plot_xs, true_CDF = get_equispaced_CDF(xs, ys, props, n_points)
    true_params[1] = CDF_mean(xs, ys)
    true_params[2] = CDF_quantile(xs, ys, α_quantile)
    true_params[3] = CDF_CVaR(xs, ys, α_cvar)
    true_params[4] = CDF_variance(xs, ys)

    ##############################################
    ## Save all the results from the experiments
    #################################################

    npzwrite(props["name"]*"_CDF.npy", CDF_estimates)
    npzwrite(props["name"]*"_mean.npy", mean_estimates)
    npzwrite(props["name"]*"_variance.npy", variance_estimates)
    npzwrite(props["name"]*"_quantile.npy", quantile_estimates)
    npzwrite(props["name"]*"_cvar.npy", cvar_estimates)

    npzwrite(props["name"]*"_CDF_bounds.npy", CDF_bounds)
    npzwrite(props["name"]*"_mean_bounds.npy", mean_bounds)
    npzwrite(props["name"]*"_mean_baseline_bounds.npy", mean_baseline_bounds)
    npzwrite(props["name"]*"_variance_bounds.npy", variance_bounds)
    npzwrite(props["name"]*"_variance_baseline_bounds.npy", variance_baseline_bounds)
    npzwrite(props["name"]*"_quantile_bounds.npy", quantile_bounds)
    npzwrite(props["name"]*"_cvar_bounds.npy", cvar_bounds)

    npzwrite(props["name"]*"_CDF_boot_bounds.npy", CDF_boot_bounds)
    npzwrite(props["name"]*"_mean_boot_bounds.npy", mean_boot_bounds)
    npzwrite(props["name"]*"_variance_boot_bounds.npy", variance_boot_bounds)
    npzwrite(props["name"]*"_quantile_boot_bounds.npy", quantile_boot_bounds)
    npzwrite(props["name"]*"_cvar_boot_bounds.npy", cvar_boot_bounds)

    npzwrite(props["name"]*"_true_CDF.npy", true_CDF)
    npzwrite(props["name"]*"_CDF_plot_xs.npy", CDF_plot_xs)
    npzwrite(props["name"]*"_true_params.npy", true_params)
    npzwrite(props["name"]*"_trend_xs.npy", trend_xs)

    #######################
    # Plots
    ########################

    function average_plot!(xs, ys, label)
        μs = dropdims(mean(ys, dims=1), dims=1)
        σs = dropdims(std(ys, dims=1), dims=1) ./ sqrt(trials)
        plot!(xs,μs, ribbon=σs, fillalpha=.3, label=label, legend=:left)
    end

    ########## CDF Plots ############

    # selected_int = ceil(Int32, intervals/2 + 1)         # Take the middle number of samples
    selected_int = intervals         # Take the middle number of samples
    plotCDF = plot(CDF_plot_xs, true_CDF, label="Fπ")                   # true CDF
    average_plot!(CDF_plot_xs, CDF_estimates[:, selected_int, :], "̂̂F")              # esimtated CDF
    average_plot!(CDF_plot_xs, CDF_bounds[:, selected_int, 1, :], "̂̂F-")             # lower bound CDF
    average_plot!(CDF_plot_xs, CDF_bounds[:, selected_int, 2, :], "̂̂F+")             # upper bound CDF
    average_plot!(CDF_plot_xs, CDF_boot_bounds[:, selected_int, 1, :], "Boot-")     # boot lower bound CDF
    average_plot!(CDF_plot_xs, CDF_boot_bounds[:, selected_int, 2, :], "Boot+")     # boot upper bound CDF

    ######## Mean Plots ###############

    plotMean = plot([true_params[1]], linetype=:hline, label="μπ", xaxis=:log)   # true mean
    average_plot!(trend_xs, mean_estimates[:, :], "̂μ")                              # esimtated mean
    average_plot!(trend_xs, mean_bounds[:, :, 1], "̂μ-")                             # lower bound mean
    average_plot!(trend_xs, mean_bounds[:, :, 2], "̂μ+")                             # upper bound mean
    average_plot!(trend_xs, mean_baseline_bounds[:, :, 1], "Base-")                 # Baseline lower bound mean
    average_plot!(trend_xs, mean_baseline_bounds[:, :, 2], "Base+")                 # Baseline upper bound mean
    average_plot!(trend_xs, mean_boot_bounds[:, :, 1], "Boot-")                     # boot lower bound mean
    average_plot!(trend_xs, mean_boot_bounds[:, :, 2], "Boot+")                     # boot upper bound mean

    ######## Variance Plots ###############

    plotVariance = plot([true_params[4]], linetype=:hline, label="σπ", xaxis=:log)   # true mean
    average_plot!(trend_xs, variance_estimates[:, :], "̂μ")                              # esimtated mean
    average_plot!(trend_xs, variance_bounds[:, :, 1], "̂μ-")                             # lower bound mean
    average_plot!(trend_xs, variance_bounds[:, :, 2], "̂μ+")                             # upper bound mean
    average_plot!(trend_xs, variance_baseline_bounds[:, :, 1], "Base-")                 # Baseline lower bound mean
    average_plot!(trend_xs, variance_baseline_bounds[:, :, 2], "Base+")                 # Baseline upper bound mean
    average_plot!(trend_xs, variance_boot_bounds[:, :, 1], "Boot-")                     # boot lower bound mean
    average_plot!(trend_xs, variance_boot_bounds[:, :, 2], "Boot+")                     # boot upper bound mean

    ######## Quantile Plots ###############

    plotQuantile = plot([true_params[2]],linetype=:hline, label="Qπ", xaxis=:log)   # true quantile
    average_plot!(trend_xs, quantile_estimates[:, :], "̂Qπ")                             # esimtated quantile
    average_plot!(trend_xs, quantile_bounds[:, :, 1], "̂Q-")                             # lower bound quantile
    average_plot!(trend_xs, quantile_bounds[:, :, 2], "̂Q+")                             # upper bound quantile
    average_plot!(trend_xs, quantile_boot_bounds[:, :, 1], "Boot-")                     # boot lower bound quantile
    average_plot!(trend_xs, quantile_boot_bounds[:, :, 2], "Boot+")                     # boot upper bound quantile

    ######## CVaR Plots ###############

    plotCVaR = plot([true_params[3]],linetype=:hline, label="Cπ", xaxis=:log)   # true cvar
    average_plot!(trend_xs, cvar_estimates[:, :], "̂Cπ")                             # esimtated cvar
    average_plot!(trend_xs, cvar_bounds[:, :, 1], "̂C-")                             # lower bound cvar
    average_plot!(trend_xs, cvar_bounds[:, :, 2], "̂C+")                             # upper bound cvar
    average_plot!(trend_xs, cvar_boot_bounds[:, :, 1], "Boot-")                     # boot lower bound cvar
    average_plot!(trend_xs, cvar_boot_bounds[:, :, 2], "Boot+")                     # boot upper bound cvar

    ####### Failure rates ##############

    plotFailure = plot([δ], linetype=:hline, label="δ", xaxis=:log)

    # Simultaneous guaranteed coverage bounds failure
    f1l = mean_bounds[:, :, 1] .> true_params[1]         # Lower mean bound failed
    f2l = quantile_bounds[:, :, 1] .> true_params[2]     # Lower quantile bound failed
    f3l = cvar_bounds[:, :, 1] .> true_params[3]         # Lower cvar bound failed
    f4l = variance_bounds[:, :, 1] .> true_params[4]     # Lower variance bound failed

    f1u = mean_bounds[:, :, 2] .< true_params[1]         # Upper mean bound failed
    f2u = quantile_bounds[:, :, 2] .< true_params[2]     # Upper quantile bound failed
    f3u = cvar_bounds[:, :, 2] .< true_params[3]         # Upper cvar bound failed
    f4u = variance_bounds[:, :, 2] .< true_params[4]     # Upper variance bound failed

    ci_failure = f1l .| f2l .| f3l .| f4l .| f5l .| f1u .| f2u .| f3u .| f4u .| f5u       # Bitwise OR
    average_plot!(trend_xs, ci_failure, "CI")

    # Simultaneous Bootstrap bounds failure
    f1l = mean_boot_bounds[:, :, 1] .> true_params[1]         # Lower mean boot bound failed
    f2l = quantile_boot_bounds[:, :, 1] .> true_params[2]     # Lower quantile boot bound failed
    f3l = cvar_boot_bounds[:, :, 1] .> true_params[3]         # Lower cvar boot bound failed
    f4l = variance_boot_bounds[:, :, 1] .> true_params[4]     # Lower variance boot bound failed

    f1u = mean_boot_bounds[:, :, 2] .< true_params[1]         # Upper mean boot bound failed
    f2u = quantile_boot_bounds[:, :, 2] .< true_params[2]     # Upper quantile boot bound failed
    f3u = cvar_boot_bounds[:, :, 2] .< true_params[3]         # Upper cvar boot bound failed
    f4u = variance_boot_bounds[:, :, 2] .< true_params[4]     # Upper variance boot bound failed

    boot_failure = f1l .| f2l .| f3l .| f4l .| f1u .| f2u .| f3u .| f4u          # Bitwise OR
    average_plot!(trend_xs, boot_failure, "Boot")

    npzwrite(props["name"]*"_ci_failure.npy", ci_failure)
    npzwrite(props["name"]*"_boot_failure.npy", boot_failure)

    ###### Plot everything in a single plot ##########
    plot(plotCDF, plotMean, plotVariance, plotQuantile, plotCVaR, plotFailure, layout=(7,1))

end

function compute_NS_trend(props,
    data::Array{Float64, 5},  # speed, trial, number of eps, horizon, 2
    pi_data::Array{Float64, 2}, # speed, trials
    trials::Int)

    ######################################################
    # Pre-allocate arrays for storing all the estimates
    n_points = 1000      # Number of support points for plotting CDFs
    δ = 0.05             # Failure rates
    speeds = [0,1,2]     # Speeds of non-stationarity

    CDF_bounds = fill(0.0, length(speeds), trials, 2, n_points)

    # True CDF and parameters
    true_CDF = fill(0.0, length(speeds), n_points)
    CDF_plot_xs = fill(0.0, n_points)

    ###################################################
    # Compute estimates and bounds for CDF and params for different trials
    ###################################################

    for speed in speeds
        Threads.@threads for trial in 1:1:trials
            println(speed, ":", trial)

            fragment = data[speed+1, trial, :, :, :]

            # Guaranteed coverage bounds for CDF and other parameters
            # bound_xs, lb_ys, ub_ys = CDF_NS_band(fragment, props, δ)
            bound_xs, lb_ys, ub_ys = CDF_NS_band_optimized(fragment, props, δ)
            CDF_bounds[speed+1, trial, 1, :] = get_equispaced_CDF(bound_xs, lb_ys, props, n_points)[2]
            CDF_bounds[speed+1, trial, 2, :] = get_equispaced_CDF(bound_xs, ub_ys, props, n_points)[2]

        end;

        #######################
        # True CDF and parameters
        ########################
        xs, ys = on_policy_CDF(pi_data[speed+1, :])
        CDF_plot_xs, true_CDF[speed+1, :] = get_equispaced_CDF(xs, ys, props, n_points)
    end;

    ##############################################
    ## Save all the results from the experiments
    #################################################

    npzwrite(props["name"]*"_CDF_bounds.npy", CDF_bounds)
    npzwrite(props["name"]*"_true_CDF.npy", true_CDF)
    npzwrite(props["name"]*"_CDF_plot_xs.npy", CDF_plot_xs)

    #######################
    # Plots
    ########################

    function average_plot!(xs, ys, label)
        μs = dropdims(mean(ys, dims=1), dims=1)
        σs = dropdims(std(ys, dims=1), dims=1) ./ sqrt(trials)
        plot!(xs,μs, ribbon=σs, fillalpha=.3, label=label, legend=:left)
    end

    ########## CDF Plots ############

    plotCDF1 = plot(CDF_plot_xs, true_CDF[1, :], label="Fπ")                   # true CDF
    average_plot!(CDF_plot_xs, CDF_bounds[1, :, 1, :], "̂̂F-")             # lower bound CDF
    average_plot!(CDF_plot_xs, CDF_bounds[1, :, 2, :], "̂̂F+")             # upper bound CDF

    plotCDF2 = plot(CDF_plot_xs, true_CDF[2, :], label="Fπ")                   # true CDF
    average_plot!(CDF_plot_xs, CDF_bounds[2, :, 1, :], "̂̂F-")             # lower bound CDF
    average_plot!(CDF_plot_xs, CDF_bounds[2, :, 2, :], "̂̂F+")             # upper bound CDF

    plotCDF3 = plot(CDF_plot_xs, true_CDF[3, :], label="Fπ")                   # true CDF
    average_plot!(CDF_plot_xs, CDF_bounds[3, :, 1, :], "̂̂F-")             # lower bound CDF
    average_plot!(CDF_plot_xs, CDF_bounds[3, :, 2, :], "̂̂F+")             # upper bound CDF

    ###### Plot everything in a single plot ##########
    plot(plotCDF1, plotCDF2, plotCDF3, layout=(3, 1))

end

####
# For all of the domains, R \in [rmin, rmax] is always required
# If the range for return [Gmin, Gmax] is also known,
# then the code can use that to make bounds tighter as well.

domain_props = Dict(
    "Reco" => Dict(
                    "beta_path" => "../../Experiments/Reco/ActorCritic/DataDefault/0/Results/beta_trajectories_0.5.npy",
                    "pi_path" => "../../Experiments/Reco/ActorCritic/DataDefault/0/Results/eval_returns_0.5.npy",
                    "rmin" => 0.0,
                    "rmax" => 10.0,
                    "Gmin" => 0.0,
                    "Gmax" => 10.0,
                    "horizon" => 1,
                    "gamma" => 0.99,
                    "name" => "Reco"
    ),
    "NS_Reco" => Dict(
                    "beta_path" => "../../Experiments/NS_Reco/beta_trajectories_1000.npy",
                    "pi_path" => "../../Experiments/NS_Reco/eval_returns_1000.npy",
                    "rmin" => -1.0,
                    "rmax" => +1.0,
                    "Gmin" => -1.0,
                    "Gmax" => +1.0,
                    "horizon" => 1,
                    "gamma" => 1.0,
                    "name" => "NS_Reco"
    ),
    "Maze_POMDP" => Dict(
                    # POMDP + multiple behavior policies
                    "beta_path" => "../../Experiments/Maze/ActorCritic/DataDefault/0/Results/beta_trajectories_2.npy",
                    "pi_path" => "../../Experiments/Maze/ActorCritic/DataDefault/0/Results/eval_returns_0.75.npy",
                    "rmin" => -1.0,
                    "rmax" => +1.0,
                    "Gmin" => -6.0,
                    "Gmax" => -1.0,
                    "horizon" => 6,
                    "gamma" => 0.99,
                    "name" => "Maze_POMDP"
    ),
    "Diabetes" => Dict(
                    "beta_path" => "../../Experiments/SimGlucosediscrete-v0/ActorCritic/DataDefault/0/Results/beta_trajectories_0.5.npy",
                    "pi_path" => "../../Experiments/SimGlucosediscrete-v0/ActorCritic/DataDefault/0/Results/eval_returns_0.5.npy",
                    "rmin" => -15.0,
                    "rmax" => +15.0,
                    "Gmin" => -15.0,
                    "Gmax" => +15.0,
                    "horizon" => 1,
                    "gamma" => 0.99,
                    "name" => "Diabetes"
    )

)

function run_UnO(domain)
    props = domain_props[domain]
    beta_data = npzread(joinpath(@__DIR__, props["beta_path"]))
    pi_data = npzread(joinpath(@__DIR__, props["pi_path"]))
    compute_trend(props, beta_data, pi_data, 30, 2.0, 4.5, 6)
    # compute_trend(props, beta_data, pi_data, 3, 2.0, 4.0, 5)
end;

function run_NS_UnO(domain)
    props = domain_props[domain]
    beta_data = npzread(joinpath(@__DIR__, props["beta_path"]))
    pi_data = npzread(joinpath(@__DIR__, props["pi_path"]))
    compute_NS_trend(props, beta_data, pi_data, 30)
end;

#######################
# IMP: To make use threading
#
# Run in CLI before executing code
# export JULIA_NUM_THREADS=10
###########################

# run_UnO("Maze")
# run_UnO("Reco")
# run_UnO("Diabetes")

# run_UnO("Maze_POMDP")
# run_NS_UnO("NS_Reco")

max_stars_repo_name: yashchandak/UnO

max_issues_repo_name: yashchandak/UnO

max_forks_repo_name: yashchandak/UnO


text:
@timedtestset "zygote_rules" begin
    @timedtestset "Diagonal" begin
        rng, N = MersenneTwister(123456), 11
        adjoint_test(Diagonal, rand(rng, N, N), randn(rng, N))
        adjoint_test(x->Diagonal(x).diag, randn(rng, N), randn(rng, N))
    end
    @timedtestset "broadcast" begin
        @timedtestset "exp" begin
            rng, N = MersenneTwister(123456), 11
            adjoint_test(x->exp.(x), randn(rng, N), randn(rng, N))
        end
        @timedtestset "-" begin
            rng, N = MersenneTwister(123456), 11
            adjoint_test(x->.-x, randn(rng, N), randn(rng, N))
        end
    end
    @timedtestset "ldiv(::Diagonal, ::Matrix)" begin
        rng, P, Q = MersenneTwister(123456), 13, 15
        Ȳ = randn(rng, P, Q)
        d = randn(rng, P)
        X = randn(rng, P, Q)
        adjoint_test((d, X)->Diagonal(d) \ X, Ȳ, d, X)
    end
end


text:
function SampledSignals.unsafe_read!(source::SndFileSource, buf::Array, frameoffset, framecount)
    total = min(framecount, nframes(source) - source.pos + 1)
    nread = 0
    readbuf = source.readbuf
    while nread < total
        n = min(size(readbuf, 2), total - nread)
        # transpose! needs the ranges to all use Ints, which on 32-bit systems
        # is an Int32, but sf_writef returns Int64 on both platforms, so we
        # convert to a platform-native Int. This also avoids a
        # type-inferrability problem where `nw` would otherwise change type.
        nr::Int = sf_readf(source.filePtr, readbuf, n)
        # the data comes in interleaved, so we need to transpose
        transpose!(view(buf, (1:nr) .+ frameoffset .+ nread, :),
                   view(readbuf, :, 1:nr))
        source.pos += nr
        nread += nr
        nr == n || break
    end

    nread
end

# returns the number of samples written
function SampledSignals.unsafe_write(sink::SndFileSink, buf::Array, frameoffset, framecount)
    nwritten = 0
    writebuf = sink.writebuf
    while nwritten < framecount
        n = min(size(writebuf, 2), framecount - nwritten)
        # the data needs to be interleaved, so we need to transpose
        transpose!(view(writebuf, :, 1:n),
                   view(buf, (1:n) .+ frameoffset .+ nwritten, :))
        # transpose! needs the ranges to all use Ints, which on 32-bit systems
        # is an Int32, but sf_writef returns Int64 on both platforms, so we
        # convert to a platform-native Int. This also avoids a
        # type-inferrability problem where `nw` would otherwise change type.
        nw::Int = sf_writef(sink.filePtr, writebuf, n)
        sink.nframes += nw
        nwritten += nw
        nw == n || break
    end

    nwritten
end


text:
using ConvexHull, JuMP, Base.Test

# A simple square
let
    m = Model()
    @defVar(m, 0 ≤ x[1:2] ≤ 1)

    v, r = get_extrema(m)

    @test length(v) == 4
    @test length(r) == 0
    @test [0,0] in v
    @test [1,0] in v
    @test [0,1] in v
    @test [1,1] in v
end

# the hypercube
for N in 2:11
    let
        m = Model()
        @defVar(m, 0 ≤ x[1:N] ≤ 1)

        v, r = get_extrema(m)

        @test length(v) == 2^N
        @test length(r) == 0

        for k in 0:N
            for p in combinations(1:N, k)
                ex = zeros(N)
                ex[p] = ones(length(p))
                @test ex in v
            end
        end
    end
end

# The 6D simplex
for N in 1:3:51
    let
        N = 6
        m = Model()
        @defVar(m, x[1:N] >= 0)
        @addConstraint(m, sum{x[i], i=1:N} == 1)

        v, r = get_extrema(m)

        @test length(v) == N
        @test length(r) == 0

        for k in 1:N
            ex = zeros(N)
            ex[k] = 1
            @test ex in v
        end
    end
end

# The 6D simplex, plus the origin
for N in 1:3:39
    let
        m = Model()
        @defVar(m, x[1:N] >= 0)
        @addConstraint(m, sum{x[i], i=1:N} ≤ 1)

        v, r = get_extrema(m)

        @test length(v) == N+1
        @test length(r) == 0

        @test zeros(N) in v
        for k in 1:N
            ex = zeros(N)
            ex[k] = 1
            @test ex in v
        end
    end
end

# The ex1 example
let
    m = Model()
    @defVar(m, x)
    @defVar(m, y)
    @addConstraints(m, begin
        12 + 2x - y ≥ 0
        -6 - x + 2y ≥ 0
        -3 + x + y ≥ 0
        1 + x ≥ 0
    end)

    v, r = get_extrema(m)

    @test length(r) == 2
    @test length(v) == 3

    for ex in Vector[[ 0, 3],
                     [-1, 4],
                     [-1,10]]
        @test is_approx_included(v, ex)
    end

    for ex in Vector[[1,2],
                     [2,1]]
        @test is_approx_included(r, ex)
    end
end

# infeasible
for N in 2:8
    let
        m = Model()
        @defVar(m, 0 ≤ x[1:N] ≤ 1)
        @addConstraint(m, x[1] ≥ 2)

        r, v = get_extrema(m)

        @test isempty(r)
        @test isempty(v)
    end
end

# non full-dimensional
let
    m = Model()
    @defVar(m, x[1:3] ≥ 1)
    @addConstraints(m, begin
        x[1] == 2
        x[2] ≤ 2
    end)

    v, r = get_extrema(m)

    @test is_approx_included(v, [2,2,1])
    @test is_approx_included(v, [2,1,1])
    @test is_approx_included(r, [0,0,1])
end

# cross polytope
for N in 2:9
    let
        m = Model()

        @defVar(m, x[1:N])

        for k in 0:N
            for S in combinations(1:N, k)
                Sᶜ = setdiff(1:N, S)
                @addConstraint(m, sum{x[i], i in S} - sum{x[i], i in Sᶜ} ≤ 1)
            end
        end

        v, r = get_extrema(m)

        @test length(v) == 2N
        @test length(r) == 0

        for i in 1:N
            ex = zeros(N)
            ex[i] = 1
            @test is_approx_included(v, ex)
            ex[i] = -1
            @test is_approx_included(v, ex)
        end
    end
end


text:
# x, y, rho and Phi(x, y, rho) computed using 10000 Gauss points
using StaticArrays

data = [
    1.4643071301680395 0.2608711310869101 -0.39340549998271723 0.5412822677479813
    1.1341188715556492 -1.5745869848370844 0.047636263318690336 0.05137668451285231
    -0.5638279972380127 -1.0279660845549863 -0.9024285135658796 6.758922980430737e-6
    0.40752339592449827 1.4121749577696892 0.599862776662886 0.6426102791769261
    -0.3112378143828986 0.09830439448059536 -0.27232272454795603 0.16195548268402457
    0.21327813956606173 -0.949087921746004 0.014741864645310599 0.10156803307865547
    0.7772983205313361 0.05747447390909049 0.2452968774301893 0.4377541290132034
    1.1675220009018015 0.28224129574653145 0.9163748809032705 0.6104378198196154
    0.10817469836005027 0.14109232234154198 0.22726503215215224 0.3379740085030564
    -1.8380592107548481 -1.5555742025386252 -0.31576708068649745 0.0002959449027418672
    0.643694737206525 1.2972628288018058 0.28647938907273673 0.6859654918890726
    -0.12308284078752839 0.3187970252733201 0.44882000014473206 0.35077850917848047
    0.4421427803274385 -1.0104910883715064 -0.3192501356898158 0.07502632133595051
    0.07200507807828677 -1.2444991825818057 -0.08607973357657019 0.050070937587043035
    -1.1034372169774 1.9233445877702353 0.5130398911643251 0.1347986750919063
    -1.867527226614093 0.47882320039568355 -0.06002767119617492 0.01961640082984873
    0.44972418526181324 0.05118485079801709 0.8076998575057721 0.4817208100419016
    0.24537526050312564 -1.2026725186704708 -0.9331611608750189 0.00015601307309313917
    -0.7361552384181024 -0.9948592466621384 -0.44632484102400966 0.009749883973644399
    -0.4450002990337082 0.9776749351851451 0.475920120611815 0.3118895797012318
    Inf -0.6733567529686127 0.9543028899648932 0.25036017750415995
    -Inf 0.12537172573685507 0.6490753182220962 0.0
    -Inf Inf 0.4691937784801415 0.0
    -Inf 1.6557106424466033 -0.21725994167383422 0.0
    Inf -1.4682072138174136 -0.02027841919955494 0.07102397267717517
    -0.2363119428322761 -Inf 0.5737718491966176 0.0
    Inf 0.21544610030974956 0.3517368660664412 0.58529023452831
    0.9993603424079786 Inf 0.17674031246464647 0.8411899181549679
#    -2.1081948831537214 0.03380065156887355 -0.9488037280958662 4.774982492738999e-13
    -Inf -0.08625166314905748 -0.30604233868164277 0.0
    1.4859295914108603 -Inf 0.199939784570541 0.0
    0.7475058996740315 -Inf 0.07531069470545042 0.0
    -Inf Inf -0.44055484032035086 0.0
    Inf -Inf -0.5163347434074138 0.0
    -0.6246291628524274 -1.6712393716866896 -0.2882232466347343 0.0048098181070054605
    -Inf -Inf -0.6626521561983009 0.0
    -Inf -Inf -0.21276586850580753 0.0
    Inf -Inf -0.22060273983396822 0.0
    Inf 1.1679218940355824 0.8628552904703347 0.8785808640718042
    Inf Inf 0.2047007663963864 1.0
    1.4643071301680395 0.2608711310869101 0.9018650879232706 0.602781040176274
    -1.5745869848370844 -0.8998931491797961 0.8968484585837353 0.05514962202849921
    -0.06462944484077995 0.40752339592449827 -0.9097039001685411 0.1522552330577691
    -0.3112378143828986 0.09830439448059536 -0.9511622546807661 0.018122797835366328
#    -0.949087921746004 -0.41839491618414215 -0.9832714583771583 3.3306690738754696e-16
    -0.5206684964371484 1.1675220009018015 -0.8937184403584206 0.18451550710692455
    0.10817469836005027 0.14109232234154198 0.9241414860485208 0.48746907008890855
    -1.5555742025386252 1.0980377024685246 -0.8922680775583764 0.0064844180972321586
    -1.6191589467053888 -0.12308284078752839 -0.9173323232250136 2.409223393735893e-7
    0.4421427803274385 -1.0104910883715064 -0.9723685127796805 0.00018629380796320039
    -1.2444991825818057 0.7321681826266581 0.9189212215747404 0.10665793625632225
    1.3775612444949075 -1.867527226614093 -0.9430317005654634 0.0011605534014627501
    0.44972418526181324 0.05118485079801709 -0.9793209397632535 0.19413558048329926
    -1.2026725186704708 -0.050779830784901565 0.9209700880204509 0.11447955295814552
    -0.27768412415298127 -0.4450002990337082 -0.9159636302721306 0.0025951526377828993
    0.8486637089024585 -1.4446783561804906 0.8965098114508927 0.07427411843427388
    1.9444589221380262 -0.7083489702163178 0.9799915616744654 0.2393642867830409
    0.1450805244965497 0.8437651107489547 0.9696991409001233 0.5576179126178168
    -0.35191886502552666 0.6718101538102055 0.9680571279920187 0.3624489827303736
    -0.5410899301286479 1.6557106424466033 0.9509727243269613 0.29422279933864653
]
for i in 1:size(data, 1)
    local x, y, ρ, pr = data[i, :]
    @test GaussianDistributions.Phi(x, y, ρ) ≈ pr
end

P = GaussianDistributions.BiNormal(0.6)
@test cdf(P, [0,0]) ≈ 1/4 + asin(P.ρ)/2pi

x, y = data[2, 1:2]
dx = dy = 1e-5
@test isapprox(((cdf(P, [x + dx, y + dx]) - cdf(P, [x+dx, y]))/dy + (cdf(P, [x, y])-cdf(P, [x, y+dx]))/dy)/dx,  pdf(P, [x,y]), atol=1e-6)

L = rand(2,2)
μ = rand(2)
P = Gaussian(μ, L*L')
σ1 = sqrt(P.Σ[1,1])
σ2 = sqrt(P.Σ[2,2])
ρ = P.Σ[1,2]/(σ1*σ2)
P2 = GaussianDistributions.BiNormal(ρ)
z = [rand(P2)...]

@test cdf(P2, z) ≈ cdf(P, μ + z.*[σ1, σ2])


text:
# the following creates doc-strings for the aliases (`instances`) of each measure:

for m in measures()
    name = m.name
    for instance in m.instances
        alias = Symbol(instance)
        quote
            @doc "An instance of type [`$($name)`](@ref). "*
                "Query the [`$($name)`](@ref) doc-string for details. " $alias
        end |> eval
    end
end


text:
using Base: sigatomic_begin, sigatomic_end
using PyCall: @pycheckz, TypeTuple

"""
Low-level version of `pycall(o, ...)` that always returns `PyObject`.
"""
function _pycall_legacy(o::Union{PyObject,PyPtr}, args...; kwargs...)
    oargs = map(PyObject, args)
    nargs = length(args)
    sigatomic_begin()
    try
        arg = PyObject(@pycheckn ccall((@pysym :PyTuple_New), PyPtr, (Int,),
                                       nargs))
        for i = 1:nargs
            @pycheckz ccall((@pysym :PyTuple_SetItem), Cint,
                             (PyPtr,Int,PyPtr), arg, i-1, oargs[i])
            pyincref(oargs[i]) # PyTuple_SetItem steals the reference
        end
        if isempty(kwargs)
            ret = PyObject(@pycheckn ccall((@pysym :PyObject_Call), PyPtr,
                                          (PyPtr,PyPtr,PyPtr), o, arg, C_NULL))
        else
            #kw = PyObject((AbstractString=>Any)[string(k) => v for (k, v) in kwargs])
            kw = PyObject(Dict{AbstractString, Any}([Pair(string(k), v) for (k, v) in kwargs]))
            ret = PyObject(@pycheckn ccall((@pysym :PyObject_Call), PyPtr,
                                            (PyPtr,PyPtr,PyPtr), o, arg, kw))
        end
        return ret::PyObject
    finally
        sigatomic_end()
    end
end

"""
    pycall(o::Union{PyObject,PyPtr}, returntype::TypeTuple, args...; kwargs...)

Call the given Python function (typically looked up from a module) with the given args... (of standard Julia types which are converted automatically to the corresponding Python types if possible), converting the return value to returntype (use a returntype of PyObject to return the unconverted Python object reference, or of PyAny to request an automated conversion)
"""
pycall_legacy(o::Union{PyObject,PyPtr}, returntype::TypeTuple, args...; kwargs...) =
    return convert(returntype, _pycall_legacy(o, args...; kwargs...)) #::returntype

pycall_legacy(o::Union{PyObject,PyPtr}, ::Type{PyAny}, args...; kwargs...) =
    return convert(PyAny, _pycall_legacy(o, args...; kwargs...))


text:
# This file is a part of BAT.jl, licensed under the MIT License (MIT).

using BAT
using Test

using ArraysOfArrays, Distributions, StatsBase, IntervalSets, ValueShapes

@testset "truncated_density" begin
    prior_dist = NamedTupleDist(
        a = truncated(Normal(), -2, 2),
        b = Exponential(),
        c = [1 2; 3 4],
        d = [-3..3, -4..4]
    )
    prior = convert(AbstractDensity, prior_dist)

    likelihood = v -> (logval = 0,)

    posterior = PosteriorDensity(likelihood, prior)

    intervals = [-1..3, 1..4, -2..2, 0..4]
    bounds = BAT.HyperRectBounds(intervals)

    @test (@inferred BAT.truncate_dist_hard(Exponential(), -1..5)).dist isa Truncated{<:Exponential}

    let
        orig_dist = Exponential()
        trunc_dist, logscalecorr = @inferred BAT.truncate_dist_hard(orig_dist, -6..6)
        @test trunc_dist.lower == 0 && trunc_dist.upper == 6
        @test logpdf(trunc_dist, 3) + logscalecorr ≈ logpdf(orig_dist, 3)
    end

    let
        orig_dist = truncated(Exponential(), 1, 5)
        trunc_dist, logscalecorr = @inferred BAT.truncate_dist_hard(orig_dist, 0..4)
        @test trunc_dist.lower == 1 && trunc_dist.upper == 4
        @test logpdf(trunc_dist, 3) + logscalecorr ≈ logpdf(orig_dist, 3)
    end

    let
        orig_dist = truncated(Exponential(), 1, 5)
        trunc_dist, logscalecorr = @inferred BAT.truncate_dist_hard(orig_dist, 2..6)
        @test trunc_dist.lower == 2 && trunc_dist.upper == 5
        @test logpdf(trunc_dist, 3) + logscalecorr ≈ logpdf(orig_dist, 3)
    end

    @test @inferred(BAT.truncate_dist_hard(prior_dist, intervals)).dist isa NamedTupleDist
    let
        trunc_dist, logscalecorr = @inferred BAT.truncate_dist_hard(prior_dist, intervals)
        @test logpdf(unshaped(trunc_dist), [1, 2, 0, 3]) + logscalecorr ≈ logpdf(unshaped(prior_dist), [1, 2, 0, 3])
    end

    @test @inferred(BAT.truncate_density(prior, bounds)) isa BAT.TruncatedDensity

    @test BAT.eval_logval(unshaped(BAT.truncate_density(prior, bounds)), [1, 2, 0, 3], BAT.default_dlt()) ≈ BAT.eval_logval(unshaped(prior), [1, 2, 0, 3], BAT.default_dlt())
    @test BAT.eval_logval(BAT.truncate_density(prior, bounds), varshape(prior)([1, 2, 0, 3]), BAT.default_dlt()) ≈ BAT.eval_logval(prior, varshape(prior)([1, 2, 0, 3]), BAT.default_dlt())
    @test varshape(BAT.truncate_density(prior, bounds)) == varshape(prior)

    @test @inferred(BAT.truncate_density(posterior, bounds)) isa PosteriorDensity

    trunc_pstr = BAT.truncate_density(posterior, bounds)
    @test @inferred(BAT.eval_logval(unshaped(trunc_pstr), [1, 2, 0, 3], BAT.default_dlt())) ≈ BAT.eval_logval(unshaped(posterior), [1, 2, 0, 3], BAT.default_dlt())
    @test @inferred(BAT.eval_logval(unshaped(trunc_pstr), [-1, -1, -1, -1], BAT.default_dlt())) ≈ -Inf
    @test varshape(trunc_pstr) == varshape(posterior)

    let
        trunc_prior_dist = parent(BAT.getprior(trunc_pstr)).dist
        s = bat_sample(trunc_pstr, MCMCSampling(mcalg = MetropolisHastings(), trafo = NoDensityTransform(), nsteps = 10^5)).result
        s_flat = flatview(unshaped.(s))
        @test all(minimum.(intervals) .< minimum(s_flat))
        @test all(maximum.(intervals) .> maximum(s_flat))
        cov_est = cov(unshaped.(s))
        @test isapprox(cov_est[1,1], var(trunc_prior_dist.a), rtol = 0.05)
        @test isapprox(cov_est[4,4], var(trunc_prior_dist.d.v[2]), rtol = 0.05)
    end

    @test @inferred(BAT.truncate_density(BAT.ConstDensity(LogDVal(0)), bounds)) == BAT.ConstDensity(LogDVal(0), bounds)
end