text:
using GeoData, Test, Statistics, Dates, CFTime, Plots
import ArchGDAL, NCDatasets
using GeoData: name, mode, span, sampling, val, Ordered, metadata, bounds,
               FileArray, FileStack, NCDfile
include(joinpath(dirname(pathof(GeoData)), "../test/test_utils.jl"))

ncexamples = "https://www.unidata.ucar.edu/software/netcdf/examples/"
ncsingle = maybedownload(joinpath(ncexamples, "tos_O1_2001-2002.nc"))
ncmulti = maybedownload(joinpath(ncexamples, "test_echam_spectral.nc"))

stackkeys = (
    :abso4, :aclcac, :aclcov, :ahfcon, :ahfice, :ahfl, :ahfliac, :ahfllac,
    :ahflwac, :ahfres, :ahfs, :ahfsiac, :ahfslac, :ahfswac, :albedo, :albedo_nir,
    :albedo_nir_dif, :albedo_nir_dir, :albedo_vis, :albedo_vis_dif, :albedo_vis_dir,
    :alsobs, :alsoi, :alsol, :alsom, :alsow, :ameltdepth, :ameltfrac, :amlcorac,
    :ao3, :apmeb, :apmegl, :aprc, :aprl, :aprs, :aps, :az0i, :az0l, :az0w,
    :barefrac, :dew2, :drain, :evap, :evapiac, :evaplac, :evapwac, :fage, :friac,
    :geosp, :glac, :gld, :hyai, :hyam, :hybi, :hybm, :lsp, :q, :qres, :qvi, :relhum,
    :runoff, :sd, :seaice, :siced, :sicepdi, :sicepdw, :sicepres, :slm, :sn, :snacl,
    :snc, :sni, :snifrac, :snmel, :sofliac, :sofllac, :soflwac, :srad0, :srad0d,
    :srad0u, :sradl, :srads, :sradsu, :sraf0, :srafl, :srafs, :st, :svo, :t2max,
    :t2min, :temp2, :thvsig, :topmax, :tpot, :trad0, :tradl, :trads, :tradsu,
    :traf0, :trafl, :trafs, :trfliac, :trfllac, :trflwac, :tropo, :tsi, :tsicepdi,
    :tslm1, :tsurf, :tsw, :u10, :ustr, :ustri, :ustrl, :ustrw, :v10, :vdis, :vdisgw,
    :vstr, :vstri, :vstrl, :vstrw, :wimax, :wind10, :wl, :ws, :wsmx, :xi, :xivi,
    :xl, :xlvi
)

@testset "GeoArray" begin
    @time ncarray = GeoArray(ncsingle)

    @testset "open" begin
        @test all(open(A -> A[Y=1], ncarray) .=== ncarray[:, 1, :])
    end

    # GeoData.create(tempfile, eltype(ncarray), dims(ncarray);
        # keys=keys(ncarray), missingval=missingval(gdalarray)
    # )

    @testset "read" begin
        @time A = read(ncarray);
        @test A isa GeoArray
        @test parent(A) isa Array
        A2 = zero(A)
        @time read!(ncarray, A2);
        A3 = zero(A)
        @time read!(ncsingle, A3)
        @test all(A .=== A2)
        @test all(A .=== A3)
    end

    @testset "array properties" begin
        @test size(ncarray) == (180, 170, 24)
        @test ncarray isa GeoArray
        @test index(ncarray, Ti) == DateTime360Day(2001, 1, 16):Month(1):DateTime360Day(2002, 12, 16)
        @test index(ncarray, Y) == -79.5:89.5
        @test index(ncarray, X) == 1.0:2:359.0
        @test bounds(ncarray) == (
            (0.0, 360.0),
            (-80.0, 90.0),
            (DateTime360Day(2001, 1, 1), DateTime360Day(2003, 1, 1)),
        )
    end

    @testset "dimensions" begin
        @test ndims(ncarray) == 3
        @test length.(dims(ncarray)) == (180, 170, 24)
        @test dims(ncarray) isa Tuple{<:X,<:Y,<:Ti}
        @test refdims(ncarray) == ()
        # TODO detect the time span, and make it Regular
        modes = (
             Mapped(Ordered(), Explicit(vcat((0.0:2.0:358.0)', (2.0:2.0:360.0)')), Intervals(Center()), EPSG(4326), EPSG(4326)),
             Mapped(Ordered(), Explicit(vcat((-80.0:89.0)', (-79.0:90.0)')), Intervals(Center()), EPSG(4326), EPSG(4326)),
             Sampled(Ordered(), Explicit(
                 vcat(permutedims(DateTime360Day(2001, 1, 1):Month(1):DateTime360Day(2002, 12, 1)),
                      permutedims(DateTime360Day(2001, 2, 1):Month(1):DateTime360Day(2003, 1, 1)))
                ), Intervals(Center())
            )
        )
        @test val.(span(ncarray)) == val.(span.(modes))
        @test typeof(mode(ncarray)) == typeof(modes)
        @test bounds(ncarray) == ((0.0, 360.0), (-80.0, 90.0), (DateTime360Day(2001, 1, 1), DateTime360Day(2003, 1, 1)))
    end
    tempfile = tempname() * ".nc"

    @testset "other fields" begin
        @test ismissing(missingval(ncarray))
        @test metadata(ncarray) isa Metadata{NCDfile}
        @test name(ncarray) == :tos
    end

    @testset "indexing" begin
        @test ncarray[Ti(1)] isa GeoArray{<:Any,2}
        @test ncarray[Y(1), Ti(1)] isa GeoArray{<:Any,1}
        @test ncarray[X(1), Ti(1)] isa GeoArray{<:Any,1}
        @test ncarray[X(1), Y(1), Ti(1)] isa Missing
        @test ncarray[X(30), Y(30), Ti(1)] isa Float32
        # Russia
        @test ncarray[X(50), Y(100), Ti(1)] isa Missing
        # Alaska
        @test ncarray[Y(Near(64.2008)), X(Near(149.4937)), Ti(1)] isa Missing
        @test ncarray[Ti(2), X(At(59.0)), Y(At(-50.5))] == ncarray[30, 30, 2] === 278.47168f0
    end

    @testset "methods" begin
        @testset "mean" begin
            @test all(mean(ncarray; dims=Y) .=== mean(parent(ncarray); dims=2))
        end
        @testset "trim, crop, extend" begin
            a = read(ncarray)
            a[X(1:20)] .= missingval(a)
            trimmed = trim(a)
            @test size(trimmed) == (160, 169, 24)
            cropped = crop(a; to=trimmed)
            @test size(cropped) == (160, 169, 24)
            @test all(collect(cropped .=== trimmed))
            extended = extend(cropped; to=a)
            @test all(collect(extended .=== a))
        end
        @testset "mask and mask!" begin
            msk = read(ncarray)
            msk[X(1:100), Y([1, 5, 95])] .= missingval(msk)
            @test !all(ncarray[X(1:100)] .=== missingval(msk))
            masked = mask(ncarray; to=msk)
            @test all(masked[X(1:100), Y([1, 5, 95])] .=== missingval(msk))
            tempfile = tempname() * ".nc"
            cp(ncsingle, tempfile)
            @test !all(GeoArray(tempfile)[X(1:100), Y([1, 5, 95])] .=== missing)
            open(GeoArray(tempfile); write=true) do A
                mask!(A; to=msk, missingval=missing)
                # TODO: replace the CFVariable with a FileArray{NCDfile} so this is not required
                nothing
            end
            @test all(GeoArray(tempfile)[X(1:100), Y([1, 5, 95])] .=== missing)
            rm(tempfile)
        end
        @testset "mosaic" begin
            @time ncarray = GeoArray(ncsingle)
            A1 = ncarray[X(1:80), Y(1:100)]
            A2 = ncarray[X(50:150), Y(90:150)]
            tempfile = tempname() * ".nc"
            Afile = mosaic(first, read(A1), read(A2); missingval=missing, atol=1e-7, filename=tempfile)
            Amem = mosaic(first, A1, A2; missingval=missing, atol=1e-7)
            Atest = ncarray[X(1:150), Y(1:150)]
            Atest[X(1:49), Y(101:150)] .= missing
            Atest[X(81:150), Y(1:89)] .= missing
            @test all(Atest .=== Afile .=== Amem)
        end
        @testset "chunk_series" begin
            @test GeoData.chunk_series(ncarray) isa GeoSeries
            @test size(GeoData.chunk_series(ncarray)) == (1, 1, 1)
        end
        @testset "slice" begin
            @test_throws ArgumentError GeoData.slice(ncarray, Z)
            ser = GeoData.slice(ncarray, Ti)
            @test ser isa GeoSeries
            @test size(ser) == (24,)
            @test index(ser, Ti) == DateTime360Day(2001, 1, 16):Month(1):DateTime360Day(2002, 12, 16)
            @test bounds(ser) == ((DateTime360Day(2001, 1, 1), DateTime360Day(2003, 1, 1)),)
            A = ser[1]
            @test index(A, Y) == -79.5:89.5
            @test index(A, X) == 1.0:2:359.0
            @test bounds(A) == ((0.0, 360.0), (-80.0, 90.0))
        end
    end

    @testset "indexing with reverse lat" begin
        if !haskey(ENV, "CI") # CI downloads fail. But run locally
            ncrevlat = maybedownload("ftp://ftp.cdc.noaa.gov/Datasets/noaa.ersst.v5/sst.mon.ltm.1981-2010.nc")
            ncrevlatarray = GeoArray(ncrevlat; name=:sst, missingval=-9.96921f36)
            @test order(dims(ncrevlatarray, Y)) == Ordered(ReverseIndex(), ReverseArray(), ForwardRelation())
            @test ncrevlatarray[Y(At(40)), X(At(100)), Ti(1)] == missingval(ncrevlatarray)
            @test ncrevlatarray[Y(At(-40)), X(At(100)), Ti(1)] == ncrevlatarray[51, 65, 1] == 14.5916605f0
            @test val(span(ncrevlatarray, Ti)) == Month(1)
            @test val(span(ncrevlatarray, Ti)) isa Month # Not CompoundPeriod
        end
    end

    @testset "selectors" begin
        a = ncarray[X(At(21.0)), Y(Between(50, 52)), Ti(Near(DateTime360Day(2002, 12)))]
        @test bounds(a) == ((50.0, 52.0),)
        x = ncarray[X(Near(150)), Y(Near(30)), Ti(1)]
        @test x isa Float32
        dimz = X(Between(0.0, 360)), Y(Between(-80, 90)),
               Ti(Between(DateTime360Day(2001, 1, 1), DateTime360Day(2003, 01, 02)))
        @test size(ncarray[dimz...]) == (180, 170, 24)
        @test index(ncarray[dimz...]) == index(ncarray)
        nca = ncarray[Y(Between(-80, -25)), X(Between(0, 180)), Ti(Contains(DateTime360Day(2002, 02, 20)))]
        @test size(nca) == (90, 55)
        @test index(nca, Y) == index(ncarray[1:90, 1:55, 2], Y)
        @test all(nca .=== ncarray[1:90, 1:55, 14])
    end

    @testset "conversion to GeoArray" begin
        geoA = ncarray[X(1:50), Y(20:20), Ti(1)]
        @test size(geoA) == (50, 1)
        @test eltype(geoA) <: Union{Missing,Float32}
        @test geoA isa GeoArray{Union{Missing,Float32},2}
        @test dims(geoA) isa Tuple{<:X,<:Y}
        @test refdims(geoA) isa Tuple{<:Ti}
        @test metadata(geoA) == metadata(ncarray)
        @test ismissing(missingval(geoA))
        @test name(geoA) == :tos
    end

    @testset "write" begin
        @testset "to netcdf" begin
            # TODO save and load subset
            geoA = read(ncarray)
            @test size(geoA) == size(ncarray)
            filename = tempname() * ".nc"
            write(filename, geoA)
            saved = read(GeoArray(filename))
            @test size(saved) == size(geoA)
            @test refdims(saved) == refdims(geoA)
            @test missingval(saved) === missingval(geoA)
            @test map(metadata.(dims(saved)), metadata.(dims(GeoArray))) do s, g
                all(s .== g)
            end |> all
            @test metadata(saved) == metadata(geoA)
            @test all(metadata.(dims(saved)) == metadata.(dims(geoA)))
            @test GeoData.name(saved) == GeoData.name(geoA)
            @test all(mode.(dims(saved)) .!= mode.(dims(geoA)))
            @test all(order.(dims(saved)) .== order.(dims(geoA)))
            @test all(typeof.(span.(dims(saved))) .== typeof.(span.(dims(geoA))))
            @test all(val.(span.(dims(saved))) .== val.(span.(dims(geoA))))
            @test all(sampling.(dims(saved)) .== sampling.(dims(geoA)))
            @test typeof(dims(saved)) <: typeof(dims(geoA))
            @test index(saved, 3) == index(geoA, 3)
            @test all(val.(dims(saved)) .== val.(dims(geoA)))
            @test all(data(saved) .=== data(geoA))
            @test saved isa typeof(geoA)
            # TODO test crs
        end
        @testset "to gdal" begin
            gdalfilename = tempname() * ".tif"
            nccleaned = replace_missing(ncarray[Ti(1)], -9999.0)
            write(gdalfilename, nccleaned)
            gdalarray = GeoArray(gdalfilename)
            # gdalarray WKT is missing one AUTHORITY
            # @test_broken crs(gdalarray) == convert(WellKnownText, EPSG(4326))
            # But the Proj representation is the same
            @test convert(ProjString, crs(gdalarray)) == convert(ProjString, EPSG(4326))
            @test bounds(gdalarray) == (bounds(nccleaned)..., (1, 1))
            # Tiff locus = Start, Netcdf locus = Center
            @test reverse(index(gdalarray, Y)) .+ 0.5 ≈ index(nccleaned, Y)
            @test index(gdalarray, X) .+ 1.0  ≈ index(nccleaned, X)
            @test reverse(GeoArray(gdalarray); dims=Y()) ≈ nccleaned
        end
        @testset "to grd" begin
            nccleaned = replace_missing(ncarray[Ti(1)], -9999.0)
            write("testgrd.gri", nccleaned)
            grdarray = GeoArray("testgrd.gri");
            @test crs(grdarray) == convert(ProjString, EPSG(4326))
            @test bounds(grdarray) == (bounds(nccleaned)..., (1, 1))
            @test index(grdarray, Y) ≈ index(nccleaned, Y) .- 0.5
            @test index(grdarray, X) ≈ index(nccleaned, X) .- 1.0
            @test GeoArray(grdarray) ≈ reverse(nccleaned; dims=Y)
        end
    end

    @testset "show" begin
        sh = sprint(show, MIME("text/plain"), ncarray)
        # Test but don't lock this down too much
        @test occursin("GeoArray", sh)
        @test occursin("Y", sh)
        @test occursin("X", sh)
        @test occursin("Time", sh)
    end

    @testset "plot" begin
        ncarray[Ti(1:3:12)] |> plot
        ncarray[Ti(1)] |> plot
        ncarray[Y(100), Ti(1)] |> plot
    end

end

@testset "Single file stack" begin
    @time ncstack = GeoStack(ncmulti)

    @testset "load ncstack" begin
        @test ncstack isa GeoStack
        @test ismissing(missingval(ncstack))
        @test metadata(ncstack) isa Metadata{NCDfile}
        @test dims(ncstack[:abso4]) == dims(ncstack, (X, Y, Ti))
        @test refdims(ncstack) == ()
        # Loads child as a regular GeoArray
        @test_throws ErrorException ncstack[:not_a_key]
        @test ncstack[:albedo] isa GeoArray{<:Any,3}
        @test ncstack[:albedo, 2, 3, 1] isa Float32
        @test ncstack[:albedo, :, 3, 1] isa GeoArray{<:Any,1}
        @test dims(ncstack[:albedo]) isa Tuple{<:X,<:Y,<:Ti}
        @test keys(ncstack) isa NTuple{131,Symbol}
        @test keys(ncstack) == stackkeys
        @test first(keys(ncstack)) == :abso4
        @test metadata(ncstack) isa Metadata{NCDfile}
        @test metadata(ncstack)["institution"] == "Max-Planck-Institute for Meteorology"
        @test metadata(ncstack, :albedo) isa Metadata{NCDfile}
        @test metadata(ncstack, :albedo)["long_name"] == "surface albedo"
        # Test some DimensionalData.jl tools work
        # Time dim should be reduced to length 1 by mean
        @test axes(mean(ncstack[:albedo, Y(1:20)] , dims=Ti)) ==
              (Base.OneTo(192), Base.OneTo(20), Base.OneTo(1))
        geoA = ncstack[:albedo][Ti(4:6), X(1), Y(2)]
        @test geoA == ncstack[:albedo, Ti(4:6), X(1), Y(2)]
        @test size(geoA) == (3,)
    end

    if VERSION > v"1.1-"
        @testset "copy" begin
            geoA = read(ncstack[:albedo]) .* 2
            copy!(geoA, ncstack, :albedo);
            # First wrap with GeoArray() here or == loads from disk for each cell.
            # we need a general way of avoiding this in all disk-based sources
            @test geoA == read(ncstack[:albedo])
        end
    end

    @testset "indexing" begin
        ncmultistack = GeoStack(ncsingle)
        @test dims(ncmultistack[:tos]) isa Tuple{<:X,<:Y,<:Ti}
        @test ncmultistack[:tos] isa GeoArray{<:Any,3}
        @test ncmultistack[:tos][Ti(1)] isa GeoArray{<:Any,2}
        @test ncmultistack[:tos, Y(1), Ti(1)] isa GeoArray{<:Any,1}
        @test ncmultistack[:tos, 8, 30, 10] isa Float32
    end

    @testset "Subsetting keys" begin
        smallstack = subset(ncstack, (:albedo, :evap, :runoff))
        @test keys(smallstack) == (:albedo, :evap, :runoff)
    end

    # This is slow. We combine read/save to reduce test time
    # And it seems the memory is not garbage collected??
    @testset "read and write" begin
        @time st = read(ncstack)
        @test st isa GeoStack
        @test st.data isa NamedTuple
        @test first(st.data) isa Array
        length(dims(st[:aclcac]))
        filename = tempname() * ".nc"
        write(filename, st);
        saved = GeoStack(GeoStack(filename))
        @test keys(saved) == keys(st)
        @test metadata(saved)["advection"] == "Lin & Rood"
        @test metadata(saved) == metadata(st) == metadata(ncstack)
        @test all(first(DimensionalData.layers(saved)) .== first(DimensionalData.layers(st)))
    end

    @testset "show" begin
        ncstack = view(GeoStack(ncmulti), X(7:99), Y(3:90));
        sh = sprint(show, MIME("text/plain"), ncstack)
        # Test but don't lock this down too much
        @test occursin("GeoStack", sh)
        @test occursin("Y", sh)
        @test occursin("X", sh)
        @test occursin("Ti", sh)
        @test occursin(":tropo", sh)
        @test occursin(":tsurf", sh)
        @test occursin(":aclcac", sh)
        @test occursin("test_echam_spectral.nc", sh)
    end

end

@testset "series" begin
    ncseries = GeoSeries([ncsingle, ncsingle], (Ti,); child=GeoStack)
    @testset "read" begin
        geoseries = read(ncseries)
        @test geoseries isa GeoSeries{<:GeoStack}
        @test geoseries.data isa Vector{<:GeoStack}
    end
    geoA = GeoArray(ncsingle; key=:tos)
    @test all(read(ncseries[Ti(1)][:tos]) .=== read(geoA))
end

nothing


text:
using SmartShockFinder

# Determine the type of Network
n_dims = 1
network_type = NNRH(n_dims)

# add mesh and polynomdegrees to then generate Traindata
meshes = SmartShockFinder.get_mesh_1d(5)
polydeg = collect(2:9)
train_setup = Traindata_Settings(polydeg, meshes)

# Parameters to train the Net
η = 0.002             # learning rate
β = 0.009             # regularization paramater
number_epochs = 100
Sb = 500                 # batch size
L = 10                   # early stopping parameter
hidden_units = [10, 10, 10, 10, 10]

model_setup = Model_Settings(η, β, number_epochs, Sb, L, hidden_units)

model = Network(network_type, train_setup, model_setup)

generate_traindataset1d(network_type, train_setup)

generate_validdataset1d(network_type, train_setup)

train_network(network_type, model_setup)

ext: jl


text:
using Metatheory
using Metatheory.Library
using Metatheory.EGraphs
using Metatheory.Classic
using Metatheory.Util
using Metatheory.EGraphs.Schedulers
using TermInterface

function customlt(x,y)
    if typeof(x) == Expr && typeof(y) == Expr
        false
    elseif typeof(x) == typeof(y)
        isless(x,y)
    elseif x isa Symbol && y isa Number
        false
    elseif x isa Expr && y isa Number
        false
    elseif x isa Expr && y isa Symbol
        false
    else true end
end

canonical_t = @theory begin
    # restore n-arity
    (x * x)             => x^2
    (x^n::Number * x)   => x^(n+1)
    (x * x^n::Number)   => x^(n+1)
    (x + (+)(ys...)) => +(x,ys...)
    ((+)(xs...) + y) => +(xs..., y)
    (x * (*)(ys...)) => *(x,ys...)
    ((*)(xs...) * y) => *(xs..., y)

    (*)(xs...)      |> Expr(:call, :*, sort!(xs; lt=customlt)...)
    (+)(xs...)      |> Expr(:call, :+, sort!(xs; lt=customlt)...)
end

# Metatheory.options[:verbose] = true
# Metatheory.options[:printiter] = true

function simplcost(n::ENode, g::EGraph, an::Type{<:AbstractAnalysis})
    cost = 0 + arity(n)
    if n.head == :∂
        cost += 20
    end
    for id ∈ n.args
        eclass = geteclass(g, id)
        !hasdata(eclass, an) && (cost += Inf; break)
        cost += last(getdata(eclass, an))
    end
    return cost
end

function simplify(ex; steps=4)
    params = SaturationParams(
        scheduler=ScoredScheduler,
        eclasslimit=5000,
        timeout=7,
        schedulerparams=(1000,5, Schedulers.exprsize),
        #stopwhen=stopwhen,
    )
    hist = UInt64[]
    push!(hist, hash(ex))
    for i ∈ 1:steps
        g = EGraph(ex)
        saturate!(g, cas, params; mod=@__MODULE__)
        ex = extract!(g, simplcost)
        ex = rewrite(ex, canonical_t; clean=false, m=@__MODULE__)
        if !TermInterface.isterm(typeof(ex))
            return ex
        end
        if hash(ex) ∈ hist
            println("loop detected $ex")
            return ex
        end
        println(ex)
        push!(hist, hash(ex))
    end
    # println(res)
    # for (id, ec) ∈ g.classes
    #     println(id, " => ", collect(ec.nodes))
    #     println("\t\t", getdata(ec, ExtractionAnalysis{astsize}))
    # end

end
macro simplify(ex)
    Meta.quot(simplify(ex))
end


text:
softmax!(x::A) where A<:AbstractVecOrMat{Float32} =
    nnp_softmax_output(x, x, threadpool = shared_threadpool[])

softmax!(y::A, x::A) where A<:AbstractVecOrMat{Float32} =
    nnp_softmax_output(x, y, threadpool = shared_threadpool[])

softmax(x::A) where A<:AbstractVecOrMat{Float32} =
    nnp_softmax_output(x, similar(x), threadpool = shared_threadpool[])

function maxpool(x::A, k; pad = map(_->0,k), stride = k) where A<:AbstractArray{Float32, 4}
    pad_, stride_= check_support(x, k, pad, stride)
    maxpool!(similar(x, pdims(size(x), k, pad_, stride_)), x, k, pad = pad_, stride = stride_)
end

maxpool!(y::A, x::A, k; pad = map(_->0,k), stride = k) where A<:AbstractArray{Float32, 4} =
    nnp_max_pooling_output(x, y, k, padding = expand(Val{length(k)}, pad), stride = expand(Val{length(k)}, stride), threadpool = shared_threadpool[])

function conv(x::A, w::A; pad = 0, stride = 1, dilation = 1, algo = UInt32(0)) where A<:AbstractArray{Float32, 4}
    pad_, stride_ = check_support(x, (size(w, 1), size(w, 2)), pad, stride, dilation)
    y = similar(x, cdims(size(x), dilation_dims(w, dilation), pad_, stride_))
    conv!(y, x, w, zeros(Float32, size(y, 3)), pad = pad_, stride = stride_, dilation = dilation, algo = UInt32(algo))
end

function conv(x::A1, w::A1, b::A2; pad = 0, stride = 1, dilation = 1, algo = UInt32(0)) where {A1<:AbstractArray{Float32, 4}, A2<:AbstractArray{Float32, 1}}
    pad_, stride_ = check_support(x, (size(w, 1), size(w, 2)), pad, stride, dilation)
    y = similar(x, cdims(size(x), dilation_dims(w, dilation), pad_, stride_))
    conv!(y, x, w, b, pad = pad_, stride = stride_, dilation = dilation, algo = UInt32(algo))
end

function conv!(y::A1, x::A1, w::A1, b::A2; pad = 0, stride = 1, dilation = 1, algo = UInt32(0), flipkernel = 0) where {A1<:AbstractArray{Float32, 4}, A2<:AbstractArray{Float32, 1}}
    flipkernel == 0 && (w = reverse(reverse(w, dims=1), dims=2))
    nnp_convolution_output(y, x, w, b, algo = algo, padding = pad, stride = stride, threadpool = shared_threadpool[])
end

function ∇conv_data(dy::A, x::A, w::A; pad = 0, stride = 1, dilation = 1, algo = UInt32(0)) where A<:AbstractArray{Float32, 4}
    pad_, stride_ = check_support(x, (size(w, 1), size(w, 2)), pad, stride, dilation)
    ∇conv_data!(zeros(Float32, size(x)), dy, x, w; pad = pad_, stride = stride_, dilation = dilation, algo = UInt32(algo))
end

function ∇conv_data!(dx::A, dy::A, x::A, w::A; pad = 0, stride = 1, dilation = 1, algo = UInt32(0), flipkernel = 0) where A<:AbstractArray{Float32, 4}
    flipkernel == 0 && (w = reverse(reverse(w, dims=1), dims=2))
    nnp_convolution_input_gradient(dx, x, dy, w, padding = pad, stride = stride, algo = algo, threadpool = shared_threadpool[])
end

function ∇conv_filter(dy::A, x::A, w::A; pad = 0, stride = 1, dilation = 1, algo = UInt32(0)) where A<:AbstractArray{Float32, 4}
    pad_, stride_ = check_support(x, (size(w, 1), size(w, 2)), pad, stride, dilation)
    ∇conv_filter!(zeros(Float32, size(w)), dy, x, w; pad = pad_, stride = stride_, dilation = dilation, algo = UInt32(algo))
end

function ∇conv_filter!(dw::A, dy::A, x::A, w::A; pad = 0, stride = 1, dilation = 1, algo = UInt32(0), flipkernel = 0) where A<:AbstractArray{Float32, 4}
    flipkernel == 0 && (w = reverse(reverse(w, dims=1), dims=2))
    dw .= nnp_convolution_kernel_gradient(dw, x, dy, w, padding = pad, stride = stride, algo = algo, threadpool = shared_threadpool[])
    flipkernel == 0 ? reverse(reverse(dw, dims=1), dims=2) : dw
end


text:
import Base: +, -, *, /, abs

+(u::Vec2D, v::Vec2D) = Vec2D(u.dx + v.dx, u.dy + v.dy)
-(u::Vec2D, v::Vec2D) = Vec2D(u.dx - v.dx, u.dy - v.dy)
*(factor::Number, v::Vec2D) = Vec2D(v.dx * factor, v.dy * factor)
*(v::Vec2D, factor::Number) = Vec2D(v.dx * factor, v.dy * factor)
/(v::Vec2D, factor::Number) = Vec2D(v.dx / factor, v.dy / factor)

abs(v::Vec2D) = Vec2D(abs(v.dx), abs(v.dy))


text:
module FirstFourierPlans

using FFTW
using LinearAlgebra: mul!

using ..PoincareInvariants: FirstPoincareInvariant, getpointnum, getform, getplan

import ..PoincareInvariants: compute!, getpoints, getpointspec

export FirstFourierPlan

struct FirstFourierPlan{T, D, FTP}
    θs::Matrix{T}
    input::Vector{Float64}
    θks::Vector{Complex{Float64}}
    zks::Vector{Complex{Float64}}
    fftplan::FTP
end

function FirstFourierPlan{T, D}(θ, N) where {T, D}
    θs = Matrix{T}(undef, N, D)
    input = Vector{Float64}(undef, N)
    θks = Vector{Complex{Float64}}(undef, N ÷ 2 + 1)
    zks = Vector{Complex{Float64}}(undef, N ÷ 2 + 1)

    # this makes use of FFTW internals, so it may easily break
    plan = FFTW.rFFTWPlan{Float64, FFTW.FORWARD, false, 1}(
        input, θks, 1:1, FFTW.ESTIMATE, FFTW.NO_TIMELIMIT
    )
    FirstFourierPlan{T, D, typeof(plan)}(θs, input, θks, zks, plan)
end

function compute!(
    pinv::FirstPoincareInvariant{T, D, <:Any, <:FirstFourierPlan}, t, p
) where {T, D}

    zs = pinv.points
    N = getpointnum(pinv)
    plan = getplan(pinv)

    θ = getform(pinv)
    for (i, z) in enumerate(eachrow(zs))
        plan.θs[i, :] .= θ(z, t, p)
    end

    I = zero(T)

    for d in 1:D
        copyto!(plan.input, view(plan.θs, :, d))
        mul!(plan.θks, plan.fftplan, plan.input)

        copyto!(plan.input, view(zs, :, d))
        mul!(plan.zks, plan.fftplan, plan.input)

        for k in 1:(N ÷ 2 + 1)
            I += T(real(plan.zks[k] * conj(plan.θks[k]) * im * 4π * (k-1) / N^2))
        end
    end

    return I
end

## getpoints and getpointspec ##

function getpointspec(N::Integer, ::Type{<:FirstFourierPlan})::Int
    N < 0 && error()
    return N
end

function getpoints(f, ::Type{T}, N::Integer, ::Type{<:FirstFourierPlan}) where T
    D = length(f(zero(T)))
    out = Matrix{T}(undef, N, D)

    for (i, x) in enumerate(range(0, 1, length=N+1)[1:end-1])
        out[i, :] .= f(x)
    end

    # Should return vector instead of matrix in 1D case
    return D == 1 ? vec(out) : out
end

end  # module FirstFourierPlans


text:
export LAPACKException,
       ARPACKException,
       SingularException,
       PosDefException,
       RankDeficientException

type LAPACKException <: Exception
    info::BlasInt
end

type ARPACKException <: Exception
    info::BlasInt
end

type SingularException <: Exception
    info::BlasInt
end

type PosDefException <: Exception
    info::BlasInt
end

type RankDeficientException <: Exception
    info::BlasInt
end

max_stars_repo_name: jakebolewski/julia

max_issues_repo_name: jakebolewski/julia

max_forks_repo_name: jakebolewski/julia


text:
# for AMIP we don't need regridding. WHen we do we re-introduce the ClimaCoreTempestRemap
function dummmy_remap!(target, source)  # TODO: bring back Tempest regrid
    parent(target) .= parent(source)
end

max_stars_repo_name: CliMA/ClimaCoupler


text:
module CO2SystemExample2

println(" ")
println("Here is another example:")
println(" ")
println("What is the pH on the Seawater Scale at 2 degrees, 4000 meters deep of a sample")
println("that has a pH of 7.8 on the Total Scale at 25 degrees at atmospheric pressure?")
println(" ")
println("(Addional info: alk=2400, si=50, po4=2, dissociation constats: Mehrbach Refit)")
println(" ")

par1type =   1    # The first parameter supplied is of type "1", which means "alkalinity"
par1    = 2400    # value of the first parameter
par2type =   3    # The first parameter supplied is of type "1", which means "pH"
par2    =  7.8    # value of the second parameter
sal     =   35    # Salinity of the sample
tempin  =   25    # Temperature at input conditions
tempout =    2    # Temperature at output conditions
presin  =    0    # Pressure    at input conditions
presout = 4000    # Pressure    at output conditions
sil     =   50    # Concentration of silicate  in the sample (in umol/kg)
po4     =    2    # Concentration of phosphate in the sample (in umol/kg)
pHscale =    1    # pH scale at which the input pH is reported ("1" means "Total Scale")
k1k2c   =    4    # Choice of H2CO3 and HCO3- dissociation constants K1 and K2 ("4" means "Mehrbach refit")
kso4c   =    1    # Choice of HSO4- dissociation constants KSO4 ("1" means "Dickson")

# Import CO2SYS
using CO2System

# Do the calculation. See CO2SYS's help for syntax and output format
A = CO2SYS(par1,par2,par1type,par2type,sal,tempin,tempout,presin,presout,
    sil,po4,pHscale,k1k2c,kso4c)[1]

println("The answer:") # It is the 38th element of the output of CO2SYS (pHoutSWS)
println(A[38])
println(" ")
println("See CO2SystemExample2.jl to see what the syntax for this calculation was.")
println(" ")

end # module CO2SystemExample2


text:
using RRIFT
using Test

@testset "in-vivo" begin
    test_study = RRIFT.gbm_study_uids[1]
    # Run download functions twice to check for overwrite protection
    dicom_folders = download_invivo_studies(test_study, destination = "./data/tcga-gbm")
    download_invivo_studies(test_study, destination = "./data/tcga-gbm")
    mask_folder = download_invivo_masks(destination = "./data/tcga-gbm-masks")
    download_invivo_masks(destination = "./data/tcga-gbm-masks")

    mat_file_test = preprocess_dicom_to_mat(destination = "./data/tcga-gbm-test", dicom_folders = dicom_folders, mask_folder = mask_folder)

    mat_folder = "./data/tcga-gbm-mat"
    download_invivo_preprocessed(destination = mat_folder)

    test_fits_a = process_patients(mat_file_test)
    test_fits_b = process_patients(joinpath(mat_folder, test_study * ".mat"))

    thresh = 1e-9
    println("Printing Δ between fits using local file versus downloaded file:")
    for param in (:kt, :kep, :ve, :vp, :kt_rr, :kep_rr, :ve_rr), method in (:tofts, :rrift)
        Δ = maximum(abs.(test_fits_a[param][method][1] .- test_fits_b[param][method][1]))
        println(string(param) * " " * string(method) * ": " * string(Δ))
        @test thresh > Δ
    end

    fits = process_patients(mat_folder)

    ccc_kt = ccc(fits[:kt], lim=(0, 0.2))
    ccc_ve = ccc(fits[:ve], lim=(0, 0.5))
    ccc_vp = ccc(fits[:vp], lim=(0, 0.05))

    ccc_ktrr = ccc(fits[:kt_rr])
    ccc_verr = ccc(fits[:ve_rr])
    ccc_keprr = ccc(fits[:kep_rr])

    println("======")
    println("Printing CCC:")
    println("CCC [kt, ve vp]:")
    println([ccc_kt[:rrift][:tofts], ccc_ve[:rrift][:tofts], ccc_vp[:rrift][:tofts]])
    println("CCC [kt_rr, ve_rr, kep_rr]:")
    println([ccc_ktrr[:rrift][:tofts], ccc_verr[:rrift][:tofts], ccc_keprr[:rrift][:tofts]])
end


text:
using Random, CSV, Dates, Statistics
using DataFrames: DataFrame, DataFrameRow, DataFrameRows, DataFrameColumns, nrow, ncol, eachcol
using Dates: DateTime, Date,Dates, Day,Year, Month, Week, Minute, Hour, year, day, month, hour, minute, week, dayofweek, dayofquarter
using Random: seed!
using CSV: read, write
using Statistics: mean, median, middle, quantile, std, var, cor, cov

export seed!
export read, write
export mean, median, middle, quantile, std, var, cor, cov
export Random, CSV, Dates, Statistics
export DataFrame, DataFrameRow, DataFrameRows, DataFrameColumns, nrow, ncol, eachcol
export DateTime, Date,Dates, Day,Year, Month, Week, Minute, Hour, year, day, month, hour, minute, week, dayofweek, dayofquarter


text:
### A Pluto.jl notebook ###
# v0.12.21

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 3f2e9c20-974c-11eb-3e97-757b3e0255f4
begin
	init_notebook = true
	using OpticSim, OpticSim.Geometry, OpticSim.Emitters
	import OpticSim.NotebooksUtils as NB

	NB.InitNotebook()
end

# ╔═╡ d6c83680-879e-11eb-31d4-7dbda7e93e48
begin
	init_notebook

	import PlutoUI

	using Makie

	defs = OpticSim.NotebooksUtils.Defs("ran")

	NB.DefsClearHTML(defs)

	# this function is needed to allow the visualization scene to be displayed inside a pluto notebook
	function Makie.display(obj)
		#@info "RG: $obj"
		return obj
	end

	#text of this cell to appear in the notebook
	NB.DefsAddHTML(defs, NB.HTMLFromObj(md"Basic Initialization (code)"))

	PlutoUI.Show(MIME"text/html"(), NB.DefsHTML(defs))
end

# ╔═╡ 576fb1c0-8ba3-11eb-1c10-fdb51172e2c9
md"# Basic CSG (more to come soon)"

# ╔═╡ aa47f750-8c1f-11eb-21d7-d3a969fb1f6d
begin
	csg1_top_surface_caption = "Top Surface"
	csg1_bottom_surface_caption = "Bottom Surface"
	csg1_cylinder_surface_caption = "Cylinder"
	csg1_result_surface_caption = "Intersection Result"
	csg1_surfaces = [
		csg1_top_surface_caption,
		csg1_bottom_surface_caption,
		csg1_cylinder_surface_caption,
		csg1_result_surface_caption,
	]

	NB.DefsClearHTML(defs)

	csg1_selected_surfaces_info = NB.GetVarInfo(@bind csg1_selected_surfaces PlutoUI.MultiSelect(csg1_surfaces, default=csg1_surfaces))

	cyl_rot_x_info = NB.GetVarInfo(@bind cyl_rot_x NB.UISlider(0:30, 0))
	cyl_rot_y_info = NB.GetVarInfo(@bind cyl_rot_y NB.UISlider(0:30, 0))
	cyl_rot_z_info = NB.GetVarInfo(@bind cyl_rot_z NB.UISlider(0:30, 0))

	NB.DefsAddHTML(defs,
		"
" * csg1_selected_surfaces_info.html *"
" *
		"Rotation: (x, y, z):" * cyl_rot_x_info.html *
		cyl_rot_y_info.html *
		cyl_rot_z_info.html
	)

	PlutoUI.Show(MIME"text/html"(), NB.DefsHTML(defs))
end

# ╔═╡ 5cc45e60-8c1f-11eb-3e4a-17c26a3214f6
begin
	function beziersurface()
		points = map(
			x -> collect(x),
			[
				(0.0, 0.0, 0.0) (0.0, 0.33, 0.0) (0.0, 0.66, 0.0) (0.0, 1.0, 0.0)
				(0.33, 0.0, 0.0) (0.33, 0.33, 1.0) (0.33, 0.66, 1.0) (0.33, 1.0, 0.0)
				(0.66, 0.0, 0.0) (0.66, 0.33, 1.0) (0.66, 0.66, 1.0) (0.66, 1.0, 0.0)
				(1.0, 0.0, 0.0) (1.0, 0.33, 0.0) (1.0, 0.66, 0.0) (1.0, 1.0, 0.0)
			],
		)
		return BezierSurface{OpticSim.Euclidean,Float64,3,3}(points)
	end

	# canonic bezier surface
	csg1_surf1 = AcceleratedParametricSurface(beziersurface(), 25);

	# two transformed copies of the canonic bezier surface
	csg1_surf2 = leaf(csg1_surf1, OpticSim.translation(-0.5, -0.5, 0.0))
	csg1_surf3 = leaf(csg1_surf1, Transform(0.0, Float64(π), 0.0, 0.5, -0.5, 0.0))

	# transformed cilinder
	csg1_surf4_1 = leaf(Cylinder(0.3, 1.0), OpticSim.translation(0.0, 0.0, 0.0))
	csg1_surf4 = leaf(csg1_surf4_1, OpticSim.rotation(deg2rad(cyl_rot_x), deg2rad(cyl_rot_y), deg2rad(cyl_rot_z)))

	# intersection result
	csg1_surf5 = (csg1_surf2 ∩ csg1_surf4 ∩ csg1_surf3)()

	md"## Define Surfaces and Perform CSG Optrations (code)"
end

# ╔═╡ 5f77cfd0-8854-11eb-377c-ef975f9abf63
md"## Initialization Stuff - Can be ignored"

# ╔═╡ 68e9b210-87ad-11eb-0f3a-5bb2dbf7d86c
begin
	NB.DefsClearHTML(defs)

	default_resolution_info = NB.GetVarInfo(
		@bind default_resolution PlutoUI.Select(["Small", "Medium", "Large"], default="Medium")
	)
	makie_backend_info = NB.GetVarInfo(
		@bind makie_backend PlutoUI.Select(["Static", "Web"], default="Static")
	)
	show_toc_info = NB.GetVarInfo(
		@bind show_toc PlutoUI.CheckBox(default=true)
	)

	fb3_items = [
		"Options",
		"@  Drawings Size $(default_resolution_info.html)",
		"Document Options",
		"@  $(show_toc_info.html) Show Table Of Content",
		"@  Makie Backend $(makie_backend_info.html) (Leave as Static for now)",
		]

	NB.DefsAddHTML(defs,
		NB.HTMLFloatingBox(
			fb3_items,
			name="plutoui-docmenu",
			header="Document Menu",
			width="25%",
			top="45%",
		)
	)

	# wide document layout
	NB.DefsAddHTML(defs, NB.HTMLNewDocLayout())

	# add TOC
	NB.DefsAddHTML(defs,
		NB.HTMLFromObj(
			PlutoUI.TableOfContents(title = "Document Table of Content", depth = 4)
		)*
		NB.HTMLFixTOC()
	)

	#text of this cell to appear in the notebook
	NB.DefsAddHTML(defs, NB.HTMLFromObj(md"Define Floating Menu, TOC and New Layout (code)"))

	PlutoUI.Show(MIME"text/html"(), NB.DefsHTML(defs))
end

# ╔═╡ 3a5d3ba0-87ae-11eb-1717-93be0b802cab
begin
	drawing = 1
	OpticSim.NotebooksUtils.SetBackend(defs, makie_backend)
	md"Set Makie backend"
end

# ╔═╡ 8ba2f700-8c1f-11eb-209f-b76b9713b576
begin
	drawing	# notebook only - create dependency on drawing backend - comment if running in REPL

	Vis.draw(OpticSim.SVector(0.0, 0.0, 0.0), markersize=0)

	if (csg1_top_surface_caption in csg1_selected_surfaces)
		Vis.draw!(
			csg1_surf2;
			wireframe=true,
			linewidth=1,
			color=:orange,
			shaded=false,
			normals=false,
			numdivisions=50
		)
	end
	if (csg1_bottom_surface_caption in csg1_selected_surfaces)
		Vis.draw!(
			csg1_surf3;
			wireframe=true,
			linewidth=1,
			color=:blue,
			shaded=false,
			normals=false,
			numdivisions=50
		)
	end
	if (csg1_cylinder_surface_caption in csg1_selected_surfaces)
		Vis.draw!(
			csg1_surf4;
			wireframe=true,
			linewidth=1,
			color=:red,
			shaded=false,
			normals=false,
			numdivisions=50
		)
	end
	if (csg1_result_surface_caption in csg1_selected_surfaces)
		Vis.draw!(
			csg1_surf5;
			wireframe=true,
			linewidth=1,
			color=:green,
			shaded=true,
			normals=false,
			numdivisions=50
		)
	end
	Vis.current_main_scene
end

# ╔═╡ 96e423a0-885a-11eb-02a3-8704e8dbdab6
begin
	function Vis.scene(resolution = (1000, 1000))
		# @info "RG: Vis.Scene Replacement"
		scene, layout = Makie.layoutscene(resolution = resolution)
		Vis.set_current_main_scene(scene)
		lscene = layout[1, 1] = Makie.LScene(scene, scenekw = (camera = Makie.cam3d_cad!, axis_type = Makie.axis3d!, raw = false))
		Vis.set_current_3d_scene(lscene)
		return scene, lscene
	end

	function resolution()
		if (default_resolution == "Small")
			return (300, 300)
		elseif (default_resolution == "Medium")
			return (500, 500)
		else
			return (1000, 1000)
		end
	end
end

# ╔═╡ Cell order:
# ╟─576fb1c0-8ba3-11eb-1c10-fdb51172e2c9
# ╟─5cc45e60-8c1f-11eb-3e4a-17c26a3214f6
# ╟─aa47f750-8c1f-11eb-21d7-d3a969fb1f6d
# ╠═8ba2f700-8c1f-11eb-209f-b76b9713b576
# ╟─5f77cfd0-8854-11eb-377c-ef975f9abf63
# ╠═3f2e9c20-974c-11eb-3e97-757b3e0255f4
# ╠═d6c83680-879e-11eb-31d4-7dbda7e93e48
# ╟─68e9b210-87ad-11eb-0f3a-5bb2dbf7d86c
# ╟─3a5d3ba0-87ae-11eb-1717-93be0b802cab
# ╟─96e423a0-885a-11eb-02a3-8704e8dbdab6


text:
using Turing, MCMCChains
using Distributions
using ParetoSmooth

samples = randn(100, 1, 1)
data = randn(50)
chain = Chains(samples)

compute_loglike(μ, data) = logpdf(Normal(μ, 1), data)
compute_loglike(μ, σ, data) = logpdf(Normal(μ, σ), data)

pll1 = pointwise_log_likelihoods(compute_loglike, chain, data)
size(pll1) |> display


text:
###############################################################################
##### Gibbs and Metropolis-Hastings Proposals
###############################################################################

## Probability a Node is a Branch
function probgrow(d::Int64, bm::BartModel)
  bm.hypers.α * (1.0 + d) ^ (-bm.hypers.β)
end

## Prior probability of a Tree
function log_tree_prior(tree::Tree, bm::BartModel)
  if isa(tree.root, Leaf)
    d = depth(tree.root, tree)
    return log(1 - probgrow(d, bm))
  end
  lp = log(probgrow(depth(tree.root, tree), bm))
  log_tree_prior(tree.root.left, tree, bm, lp) +
    log_tree_prior(tree.root.right, tree, bm, lp)
end

function log_tree_prior(trees::Vector{Tree}, bm::BartModel)
  ltp = 0
  for tree in trees
    ltp += log_tree_prior(tree, bm)
  end
  ltp
end

function log_tree_prior(branch::Branch, tree::Tree, bm::BartModel, lp::Float64)
  d = depth(branch, tree)
  lp += log(probgrow(d, bm))
  log_tree_prior(branch.left, tree, bm, lp) +
    log_tree_prior(branch.right, tree, bm, lp)
end

function log_tree_prior(leaf::Leaf, tree::Tree, bm::BartModel, lp::Float64)
  d = depth(leaf, tree)
  lp + log(1 - probgrow(d, bm))
end

## Probability that an observtion goes left at a given Branch
function probleft(x::Vector{Float64}, branch::Branch, tree::Tree)
  1 / (1 + exp((x[branch.var] - branch.cut) / tree.λ))
end

function probleft(X::Matrix{Float64}, branch::Branch, tree::Tree)
  1 ./ (1 .+ exp.((X[:,branch.var] .- branch.cut) ./ tree.λ))
end

## Probability that observations end up in each Leaf
function leafprob(x::Vector{Float64}, tree::Tree)
  if isa(tree.root, Leaf)
    return 1.0
  end
  S = Float64[]
  goesleft = probleft(x, tree.root, tree)
  goesright = 1 - goesleft
  leafprob(x, tree.root.left, tree, goesleft, S)
  leafprob(x, tree.root.right, tree, goesright, S)
end

function leafprob(x::Vector{Float64}, branch::Branch, tree::Tree, psi::Float64, S::Vector{Float64})
  goesleft = psi * probleft(x, branch, tree)
  goesright = psi - goesleft
  leafprob(x, branch.left, tree, goesleft, S)
  leafprob(x, branch.right, tree, goesright, S)
end

function leafprob(x::Vector{Float64}, leaf::Leaf, tree::Tree, psi::Float64, S::Vector{Float64})
  push!(S, psi)
end

function leafprob(X::Matrix{Float64}, bt::BartTree, bm::BartModel)
  S = zeros(bm.td.n, bt.ss.Lt)
  for i in 1:bm.td.n
    S[i,:] .= leafprob(X[i,:], bt.tree)
  end
  S
end

function leafprob(X::Matrix{Float64}, tree::Tree)
  n = size(X, 1)
  S = zeros(n, length(leafnodes(tree.root)))
  for i in 1:n
    S[i,:] .= leafprob(X[i,:], tree)
  end
  S
end

## Draw a new cut value for the proposed split
function drawcut(leaf::Leaf, var::Int64, tree::Tree, bm::BartModel)
  branch = leaf
  lower = [bm.td.xmin[:,var][1]]
  upper = [bm.td.xmax[:,var][1]]
  check = branch == tree.root ? false : true
  while check
    left = isleft(branch, tree)
    branch = parent(branch, tree)
    check = branch == tree.root ? false : true
    if (branch.var == var)
      check
      if (left)
        upper = push!(upper, branch.cut)
      else
        lower = push!(lower, branch.cut)
      end
    end
  end
  lower = maximum(lower)
  upper = minimum(upper)
  rand(Uniform(lower, upper))
end

## If the birthproposal is accepted
## grow the Leaf into a Branch with two new Leafs
function birthleaf!(leaf::Leaf, tree::Tree, branch::Branch)
  if isa(tree.root, Leaf)
    tree.root = branch
  else
    parentnode = parent(leaf, tree)
    if parentnode.left == leaf
      parentnode.left = branch
    else
      parentnode.right = branch
    end
  end
end

## Probability of making a birth proposal
## If the tree only has a root node, a birth proposal has to be made
## Otherwise randomly choose a Birth or Death proposal with probability 0.5
function birthprob(tree::Tree)
  isa(tree.root, Leaf) ? 1.0 : 0.5
end

function birthprob(S::Matrix{Float64})
  size(S, 2) == 1 ? 1.0 : 0.5
end

## The log ratio of the transition probabilities for a birth proposal
function log_birth_trans(bt::BartTree, S_prime::Matrix{Float64})
  # Probability of transitioning from proposed Tree back to the current Tree
  numr = (1 - birthprob(S_prime)) / (length(onlyparents(bt.tree)) + 1)
  # Probability of transitioning from the current Tree to the proposed Tree
  denomr = birthprob(bt.tree) / bt.ss.Lt
  log(numr) - log(denomr)
end

## The log ratio of the tree probabilities
function log_birth_tree(node::Node, tree::Tree, bm::BartModel)
  d = depth(node, tree)
  numr = probgrow(d, bm) * (1 - probgrow(d + 1, bm)) ^ 2
  denomr = 1 - probgrow(d, bm)
  log(numr) - log(denomr)
end

## Compute the mariginal log likelihood of the current Tree residuals
function mll(rt::Vector{Float64}, ss::SuffStats, bs::BartState, bm::BartModel)
  mll = 0.5 * logdet(2 * pi * ss.Ω)
  mll -= 0.5 * bm.td.n * log(2 * pi * bs.σ^2)
  mll -= 0.5 * ss.Lt * log(2 * pi * bm.hypers.τ)
  mll -= 0.5 * dot(rt, rt) / bs.σ^2
  mll += 0.5 * dot(ss.rhat, ss.Ω * ss.rhat)
  mll
end

function sample_var(bs::BartState, bm::BartModel)
  sample(1:bm.td.p, weights(exp.(bs.s)))
end

## Conduct a Birth proposal
function birthproposal!(bt::BartTree, rt::Vector{Float64}, bs::BartState, bm::BartModel)
  leaves = leafnodes(bt.tree.root)
  index = rand(1:length(leaves))
  leaf = leaves[index]

  newvar = sample_var(bs, bm)
  newcut = drawcut(leaf, newvar, bt.tree, bm)
  branch = Branch(newvar, newcut, Leaf(0.0), Leaf(0.0))
  goesleft = bt.S[:,index] .* probleft(bm.td.X, branch, bt.tree)
  goesright = bt.S[:,index] .- goesleft
  if size(bt.S, 2) == 1
    S_prime = hcat(goesleft, goesright)
  else
    indices = [index, index + 1]
    S_prime = zeros(bm.td.n, length(leaves) + 1)
    S_prime[:,indices] = hcat(goesleft, goesright)
    S_prime[:,setdiff(1:end, indices)] = bt.S[:,setdiff(1:end, index)]
  end
  bt.ss = suffstats(rt, bt.S, bs, bm)
  ss_prime = suffstats(rt, S_prime, bs, bm)
  mloglikratio = mll(rt, ss_prime, bs, bm) - mll(rt, bt.ss, bs, bm)
  treeratio = log_birth_tree(leaf, bt.tree, bm)
  transratio = log_birth_trans(bt, S_prime)
  logr = mloglikratio + treeratio + transratio
  if log(rand()) < logr
    birthleaf!(leaf, bt.tree, branch)
    bt.S = S_prime
    bt.ss = ss_prime
  end
end

## The log ratio of the transition probabilities for a Death proposal
function log_death_trans(bt::BartTree, S_prime::Matrix{Float64})
  denomr = birthprob(S_prime) / length(onlyparents(bt.tree))
  numr = (1 - birthprob(bt.tree)) / (bt.ss.Lt - 1)
  log(numr) - log(denomr)
end

## The log ratio of the Tree probabilities
function log_death_tree(node::Node, tree::Tree, bm::BartModel)
  -1.0*log_birth_tree(node, tree, bm)
end

## If doing a Death proposal, kill two Leaf Nodes of
## a randomly selected Branch and turn the Branch into a Leaf
function deathbranch!(branch::Branch, tree::Tree)
  if tree.root == branch
    tree.root = Leaf(0.0)
  else
    parentnode = parent(branch, tree)
    if parentnode.left == branch
      parentnode.left = Leaf(0.0)
    else
      parentnode.right = Leaf(0.0)
    end
  end

end

## Conduct a Death proposal
function deathproposal!(bt::BartTree, rt::Vector{Float64}, bs::BartState, bm::BartModel)
  branch = rand(onlyparents(bt.tree))
  indexes = findall(x -> (x == branch.left) || (x == branch.right), leafnodes(bt.tree.root))
  S_prime = copy(bt.S)
  S_prime[:,indexes[1]] = sum(bt.S[:,indexes], dims = 2)
  S_prime = S_prime[:,1:end .!= indexes[2]]
  bt.ss = suffstats(rt, bt.S, bs, bm)
  ss_prime = suffstats(rt, S_prime, bs, bm)
  mloglikratio = mll(rt, ss_prime, bs, bm) - mll(rt, bt.ss, bs, bm)
  treeratio = log_death_tree(branch, bt.tree, bm)
  transratio = log_death_trans(bt, S_prime)
  logr = mloglikratio + treeratio + transratio
  if log(rand()) < logr
    deathbranch!(branch, bt.tree)
    bt.S = S_prime
    bt.ss = ss_prime
  end
end

## Function to perform Metropolis-Hastings step for a single Tree update
function drawT!(bt::BartTree, rt::Vector{Float64}, bs::BartState, bm::BartModel)
  if rand() < birthprob(bt.tree)
    birthproposal!(bt, rt, bs, bm)
  else
    deathproposal!(bt, rt, bs, bm)
  end
end

## Function to perform Metropolis-Hastings step to update λ on a single Tree
function drawλ!(bt::BartTree, rt::Vector{Float64}, bs::BartState, bm::BartModel)
  lp_λ = logpdf(Exponential(bm.hypers.λmean), bt.tree.λ)
  log_λ = log(bt.tree.λ)
  log_λprime = log_λ + rand(Uniform(-1, 1))
  bt.tree.λ = exp(log_λprime)
  lp_λprime = logpdf(Exponential(bm.hypers.λmean), bt.tree.λ)
  S_prime = leafprob(bm.td.X, bt, bm)
  ss_prime = suffstats(rt, S_prime, bs, bm)
  logr = mll(rt, ss_prime, bs, bm) + lp_λprime + log_λprime -
    (mll(rt, bt.ss, bs, bm) + lp_λ + log_λ)
  if log(rand()) < logr
    bt.S = S_prime
    bt.ss = ss_prime
  else
    bt.tree.λ = exp(log_λ)
  end
end

## Gibbs step to update leaf parameters for a single Tree
function drawμ!(bt::BartTree, rt::Vector{Float64}, bs::BartState, bm::BartModel)
  newμ = rand(MvNormal(bt.ss.Ω * bt.ss.rhat, Symmetric(bt.ss.Ω)))
  leaves = leafnodes(bt.tree.root)
  for l in 1:length(leaves)
    leaves[l].μ = newμ[l]
  end
end

## Draw latent from truncated normal
function drawz!(bs::BartState, bm::BartModel)
  for i in 1:bm.td.n
    if bm.td.y[i] == 1
      bs.z[i] = rand(truncated(Normal(bs.fhat[i]), 0, Inf))
    else
      bs.z[i] = rand(truncated(Normal(bs.fhat[i]), -Inf, 0))
    end
  end
end

function drawtrees!(bs::RegBartState, bm::BartModel)
  for bt in bs.ensemble.trees
    fhat_t = bs.fhat .- predict(bt)
    rt = bm.td.y .- fhat_t
    drawT!(bt, rt, bs, bm)
    bm.hypers.λfix ? nothing : drawλ!(bt, rt, bs, bm)
    drawμ!(bt, rt, bs, bm)
    bs.fhat = fhat_t .+ predict(bt)
  end
end

function drawtrees!(bs::ProbitBartState, bm::BartModel)
  drawz!(bs, bm)
  for bt in bs.ensemble.trees
    fhat_t = bs.fhat .- predict(bt)
    rt = bs.z .- fhat_t
    drawT!(bt, rt, bs, bm)
    bm.hypers.λfix ? nothing : drawλ!(bt, rt, bs, bm)
    drawμ!(bt, rt, bs, bm)
    bs.fhat = fhat_t .+ predict(bt)
  end
end

## Gibbs step to update error variance
function drawσ!(bs::RegBartState, bm::BartModel)
  a = 0.5 * (bm.hypers.ν + bm.td.n)
  b = 0.5 * (bm.hypers.ν * bm.hypers.δ + sum((bm.td.y - bs.fhat).^2))
  bs.σ = sqrt(rand(InverseGamma(a, b)))
end

## Draw sparsity per Linero & Yang (2018)
varcount!(leaf::Leaf, counts::Vector) = nothing

function varcount!(branch::Branch, counts::Vector)
  counts[branch.var] += 1
  varcount!(branch.left, counts)
  varcount!(branch.right, counts)
end

function varcount(tree::Tree, bm::BartModel)
  counts = zeros(bm.td.p)
  varcount!(tree.root, counts)
  counts
end

function varcounts(trees::Vector{BartTree}, bm::BartModel)
  vec(sum(reduce(hcat, [varcount(bt.tree, bm) for bt in trees]), dims = 2))
end

# function rlgam(shape)
#   if shape >= 0.1
#     return log(rand(Gamma(shape, 1)))
#   else
#     a = shape
#     L = 1.0/a- 1.0;
#     w = exp(-1.0) * a / (1.0 - a)
#     ww = 1.0 / (1.0 + w)
#     z = 0.0
#     cond = true
#     while cond
#       U = rand()
#       if (U <= ww)
#         z = -log(U / ww);
#       else
#         z = log(rand()) / L
#       end
#       eta = z >= 0 ? -z : log(w)  + log(L) + L * z
#       h = -z - exp(-z / a);
#       cond = h - eta > log(rand())
#     end
#     return -z/a
#   end
# end

function log_sum_exp(x)
  m = maximum(x)
  m + log(sum(exp.(x .- m)))
end

function draws!(bs::BartState, bm::BartModel)
  vc = varcounts(bs.ensemble.trees, bm)
  # gc = [sum(counts[findall(bm.hypers.group_idx .== g)]) for g in bm.hypers.groups]
  shapes = bs.shape / bm.td.p .+ vc
  # bs.s = rand(Dirichlet(shapes))
  # logs = rlgam.(shapes)
  # logs = logs .- log_sum_exp(logs)
  # bs.s = exp.(logs)
  y = log.(rand.(Gamma.(shapes .+ 1, 1)))
  z = log.(rand(length(shapes))) ./ shapes
  logs = y + z
  logs = logs .- log_sum_exp(logs)
  bs.s = logs
end

function loglikω(ω, log_s, a, b, p)
  aω = p * (ω/(1-ω))
  loggamma(aω) - p*loggamma(aω/p) + aω*mean(log_s) + logpdf(Beta(a, b), ω)
end

function drawα!(bs::BartState, bm::BartModel)
  ω = bs.shape / (bs.shape + bm.td.p)
  grid = collect(1:1000) ./ 1001
  lωg = [loglikω(g, bs.s, bm.hypers.a, bm.hypers.b, bm.td.p) for g in grid]
  lse = log_sum_exp(lωg)
  prob = exp.(lωg .- lse)
  new_ω = sample(grid, weights(prob))
  bs.shape = (new_ω / (1-new_ω)) * bm.td.p
  # bs.shape = bm.hypers.shape
  # println(bs.shape)
end

## Log tree posterior: continuous respose
function ptd(c::Array)
  bm, trees, σ, y = c
  S = reduce(hcat, [leafprob(bm.td.X, tree) for tree in trees])
  Σ = Symmetric(σ^2*I + bm.hypers.τ*S*S')
  mll = loglikelihood(MvNormal(zeros(bm.td.n), Σ), reshape(y, bm.td.n, 1))
  ltp = sum([log_tree_prior(tree, bm) for tree in trees])
  llp = sum([logpdf(Exponential(bm.hypers.λmean), tree.λ) for tree in trees])
  lσp = logpdf(InverseGamma(bm.hypers.ν/2, bm.hypers.ν*bm.hypers.δ/2), σ^2)
  -2*(mll + ltp + llp + lσp)
end

function ptd(bc::RegBartChain)
  S = size(bc.treedraws, 1)*size(bc.treedraws, 3)
  treedraws = reshape(bc.treedraws, S)
  σdraws = reshape(bc.σdraws, S)
  c = [[bc.bm, treedraws[l], σdraws[l], bc.bm.td.y] for l in 1:S]
  pmap(ptd, c)
end

## Log tree posterior: binary response
# function log_tree_post(bc::ProbitBartChain)
#   S = size(bc.treedraws, 1)*size(bc.treedraws, 3)
#   treedraws = reshape(bc.treedraws, S)
#   zdraws = reshape(bc.zdraws, bc.bm.td.n, S)
#   c = [[bc.bm, treedraws[l], 1.0, zdraws[:,l]] for l in 1:S]
#   pmap(log_tree_post, c)
# end

function _ptd(c)
  trees, bm = c
  M = reduce(hcat, [leafprob(bm.td.X, tree) for tree in trees])
  sig = [1 + bm.hypers.τ*dot(M[i,:], M[i,:]) for i in 1:size(M, 1)]
  ym = sum(log.(cdf.(Normal.(0, sig), 0)))
  ltp = sum([log_tree_prior(tree, bm) for tree in trees])
  -2*(ym + ltp)
end

function ptd(bc::BART.ProbitBartChain)
  treedraws = reshape(bc.treedraws, size(bc.treedraws, 1)*size(bc.treedraws, 3))
  pmap(_ptd, [[trees, bc.bm] for trees in treedraws])
end


text:
module MDPs

using Compat

if VERSION < v"0.4-dev"
  using Docile
end

export # MDP types
       AbstractMDP,
       MDP,

       # Q-function types
       AbstractQFunction,
       ArrayQFunction,
       QFunction,  # constructor helper
       VectorQFunction,

       # transition probability types
       AbstractTransitionProbability,
       AbstractArrayTransitionProbability,
       ArrayTransitionProbability,
       FunctionTransitionProbability,
       SparseArrayTransitionProbability,
       TransitionProbability,  # constructor helper

       # reward types
       AbstractReward,
       AbstractArrayReward,
       ArrayReward,
       Reward,  # constructor helper
       SparseReward,

       # functions
       bellman,
       bellman!,
       getvalue,
       is_square_stochastic,
       ismdp,
       num_actions,
       num_states,
       policy,
       policy!,
       probability,
       reward,
       setvalue!,
       value,
       value!,
       valuetype,
       value_iteration,
       value_iteration!

include("transition.jl")
include("reward.jl")
include("qfunction.jl")
include("bellman.jl")
include("mdp.jl")
include("examples.jl")
include("utils.jl")

end # module


text:
## recognizer.jl.  Some routines for old-fashioned GMM-based (speaker) recognizers
## (c) 2013--2014 David A. van Leeuwen

## This function computes the `dotscoring' linear appoximation of a GMM/UBM log likelihood ratio
## of test data y using MAP adapted model for x.
## We can compute this with just the stats:
function dotscore(x::CSstats, y::CSstats, r::Real=1.)
    sum(broadcast(/, x.f, x.n + r) .* y.f)
end
## or directly from the UBM and the data x and y
dotscore{T<:Real}(gmm::GMM, x::Matrix{T}, y::Matrix{T}, r::Real=1.) =
    dotscore(CSstats(gmm, x), CSstats(gmm, y), r)

## or produce a score matrix
function dotscore{T}(x::Vector{CSstats{T}}, y::Vector{CSstats{T}}, r::Real=1.)
    xx = hcat([vec(x.f ./ (x.n+r)) for x in x]...)
    yy = hcat([vec(y.f) for y in y]...)
    return xx' * yy
end

## Maximum A Posteriori adapt a gmm
function maxapost{T<:AbstractFloat}(gmm::GMM, x::Matrix{T}, r::Real=16.; means::Bool=true, weights::Bool=false, covars::Bool=false)
    nₓ, ll, N, F, S = stats(gmm, x)
    α = N ./ (N+r)
    if weights
        w = α .* N / sum(N) + (1-α) .* gmm.w
        w ./= sum(g.w)
    else
        w = gmm.w
    end
    if means
        μ = broadcast(*, α./N, F) + broadcast(*, 1-α, gmm.μ)
    else
        μ = gmm.μ
    end
    if covars
        kind(gmm) == :diag || error("Sorry, can't MAP adapt full covariance matrix GMM yet")
        Σ = broadcast(*, α./N, S) + broadcast(*, 1-α, gmm.Σ .^2 + gmm.μ .^2) - g.μ .^2
    else
        Σ = gmm.Σ
    end
    hist = vcat(gmm.hist, History(@sprintf "MAP adapted with %d data points relevance %3.1f %s %s %s" size(x,1) r means ? "means" : ""  weights ? "weights" : "" covars ? "covars" : ""))
    return GMM(w, μ, Σ, hist, nₓ)
end

## compute a supervector from a MAP adapted utterance.
function Base.vec(gmm::GMM, x::Matrix, r=16.)
    kind(gmm) == :diag || error("Sorry, can't compute MAP adapted supervector for full covariance matrix GMM yet")
    nₓ, ll, N, F, S = stats(gmm, x)
    α = N ./ (N+r)
    Δμ = broadcast(*, α./N, F) - broadcast(*, α, gmm.μ)
    v = Δμ .* sqrt(broadcast(/, gmm.w, gmm.Σ))
    return vec(v)
end


text:
# This functionality could be moved to a separate package which carries the
# Plots dependency. An alternative would be to use @require here, but we don't.
# Refer to https://github.com/JuliaLang/Pkg.jl/issues/1285 - the functionality
# might become part of Julia.
import Plots
import Plots: plot, plot!, default, px
using RecipesBase
"Similar to other 'place_image', but not included in MechanicalSketch
Consider moving to MechGlueCode, provided Plots and MechanicalSketch are loaded."
function place_image_37(pos::Point, plott::Plots.Plot;
    width = missing, height = missing, scalefactor = missing,
    centered = false)
    ioc = IOContext(IOBuffer(), :color=>true)
    show(ioc, MIME("image/svg+xml"), plott)
    stsvg = String(take!(ioc.io))
    stsvg = replace(stsvg, "fill-opacity=\"1\"" => "fill-opacity=\"0.3\""; count = 2)
    place_image(pos, readsvg(stsvg); width, height, scalefactor, centered)
end


text:
# # plot recipe for plotting an Observable object
@recipe function f(
    observable::Observable,
    params::Any;
    titlewidth=50
)
    x_nt, x_range, x_i = range_to_namedtuples(params)
    x_label = keys(params)[x_i]
    title_string = get_obs_title(params, x_i, titlewidth=titlewidth)
    y = observable.func.(x_nt)

    y_label = string(observable.func)

    @series begin
        seriestype --> :path
        label --> y_label*"("*string(x_label)*")"
        xguide --> string(x_label)
        yguide --> y_label
        title --> title_string
        x_range, y
    end

end

@recipe function f(
    observable::Observable,
    param::NamedTuple,
    default_params::Any;
    titlewidth=50
)
    params = merge(default_params, param)

    @series begin
        titlewidth --> titlewidth
        observable, params
    end
end

# plot recipe for plotting a particular Observable from a NamedTuple
@recipe function f(
    observables::NamedTuple{<:Any, <:Tuple{Vararg{Observable}}},
    name::Symbol,
    params::Any;
    titlewidth=50
)
    @series begin
        titlewidth --> titlewidth
        observables[name], params
    end
end

# plot recipe for plotting a particular Observable from a NamedTuple
@recipe function f(
    observables::NamedTuple{<:Any, <:Tuple{Vararg{Observable}}},
    name::Symbol,
    params::Any,
    default_params::Any;
    titlewidth=50
)
    @series begin
        titlewidth --> titlewidth
        observables[name], params, default_params
    end
end

# plot recipe for plotting an Array{Observable}
@recipe function f(
    observables::AbstractVector{Observable},
    params::Any;
    titlewidth=50,
    bin_names = ["bin$i" for i in 1:length(observables)]
)
    nbins = length(observables)
    ys = [obs.func(params) for obs in observables]

    @series begin
        seriestype --> :steppre
        label --> "prediction"
        xticks --> (collect(1:nbins).+0.5, bin_names)
        linewidth --> 1.5
        1:nbins+1, [ys[1], ys...]
    end
end

# internal: convert a NamedTuple that contains a range into a vector of NamedTuples
function range_to_namedtuples(params)
    ks = keys(params)
    vals = collect(values(params))

    i = 0
    for j in 1:length(ks)
        if !isa(vals[j], Real)
            i = j
        end
    end

    xs = collect(vals[i])
    nt = Vector{NamedTuple}(undef, length(xs))

    for j in 1:length(xs)
        vals[i] = xs[j]
        nt[j] = namedtuple(collect(ks), vals)
    end

    return nt, xs, i
end

function get_obs_title(params, i; titlewidth=50)
    ks = keys(params)
    vals = collect(values(params))

    title_string = ""
    for j in 1:length(ks)
        if j != i
            title_string *= string(ks[j], "=", vals[j], ", ")
        end
    end
    title_string = title_string[1:end-2]

    return replace(title_string, Regex("(.{$(titlewidth)})") => Base.SubstitutionString("\\1\n"))
end


text:
using Statistics

"""

Convenience interface for discrete distributions .

If you want to do S{T} <: FiniteSupportDistribution{T}, you need to implement:

- `supp(s::S)` returning a Vector{T} with the support points

  This will fall back to doing `s.support` if you don't define a separate
  method.

- `pmf(s::S, i)` returning the mass of the i-th support element

"""
abstract type FiniteSupportDistribution{T} <: AbstractArray{T, 1} end

Base.size(fsd::FiniteSupportDistribution)=(length(fsd), )
Base.IndexStyle(::Type{<:FiniteSupportDistribution}) = IndexLinear()
Base.getindex(fsd::FiniteSupportDistribution, i)=supp(fsd)[i]

supp(fsd::FiniteSupportDistribution)=fsd.support
Base.length(fsd::FiniteSupportDistribution)=length(supp(fsd))

" Mean of the finite supp. distribution (not the underlying generator) "
Statistics.mean(fsd::FiniteSupportDistribution)=begin
    m = 0.
    for i in 1:length(fsd)
        m += fsd[i] * pmf(fsd, i)
    end
    return m
end

" Divide support points by `mean(fsd)`"
normalized(fsd::FiniteSupportDistribution)=(fsd2 = deepcopy(fsd); normalized!(fsd2); fsd2)
normalized!(fsd::FiniteSupportDistribution)=(rmul!(supp(fsd), 1/mean(fsd)))

struct BasicFiniteSupportDistribution{T} <: FiniteSupportDistribution{T}
    support::Vector{T}
    probvec::Vector{Float64}
end

pmf(fsd::BasicFiniteSupportDistribution)=fsd.probvec
pmf(fsd::BasicFiniteSupportDistribution, i::Int64)=fsd.probvec[i]

struct EquiprobDistribution{T} <: FiniteSupportDistribution{T}
    support::Vector{T}
    equiprob::Float64 # Avoids (super minor) cost of computing 1/n every time
end
EquiprobDistribution(support)=EquiprobDistribution([x for x in support],
                                                   1/length(support))

pmf(epd::EquiprobDistribution)=fill(epd.equiprob,length(epd))
pmf(epd::EquiprobDistribution, i::Int64)=epd.equiprob
Statistics.mean(epd::EquiprobDistribution)=sum(supp(epd))*pmf(epd,1)

struct Discretized1dDistribution{D, S, T <: FiniteSupportDistribution{S}} <: FiniteSupportDistribution{S}
    fsd::T
    distrib::D
end

supp(dd::Discretized1dDistribution)=supp(dd.fsd)
pmf(dd::Discretized1dDistribution)=pmf(dd.fsd)
pmf(dd::Discretized1dDistribution, i)=pmf(dd.fsd,i)
recover(dd::Discretized1dDistribution)=dd.distrib

function eqprob_discretize(dd, npts)
    # Approximate distribution `dd` with a finite support distribution with
    # equally probable points.
    # --------------------------------------------------------------------------
    #
    # Idea of aproximation -----------------------------------------------------
    #     ^
    #     |
    #     |
    # 1   -                                                  o
    #     |
    #     x                                       o
    #     |
    #     -                               o
    #     |
    #     x                           o
    #     |
    #     -                      o
    #     |
    #     x              o
    #     |
    #     -        o
    #     |
    #     x  o
    # 0 _ |__________________________________________________________\
    #                                                                /
    #  Assign F(-) - F(prev(-)) to `x`s
    # --------------------------------------------------------------------------

    ran1 = range(0, (npts-1)/npts, length=npts)
    ran2 = range(1/npts, 1, length=npts)
    ran = (ran1 + ran2)/2

    zsupport = quantile.(dd, ran)

    Discretized1dDistribution(EquiprobDistribution(zsupport), dd)
end

max_stars_repo_name: pereiragc/rouwenhorst

max_issues_repo_name: pereiragc/rouwenhorst

max_forks_repo_name: pereiragc/rouwenhorst


text:
# # Perfect baroclinic adjustment calibration with Ensemble Kalman Inversion
#
# This example showcases a "perfect model calibration" of the two-dimensional baroclinic adjustement
# problem (depth-latitude) with eddies parametrized using Gent-McWilliams--Redi isoneutral diffusion
# closure. We use output for buoyancy (``b``) and a passive-tracer concentration (``c``) to calibrate
# the parametrization.

# ## Install dependencies
#
# First let's make sure we have all required packages installed.

# ```julia
# using Pkg
# pkg"add Oceananigans, Distributions, CairoMakie, ParameterEstimocean"
# ```

# First we load few things

using ParameterEstimocean

using Oceananigans
using Oceananigans.Units
using Oceananigans.TurbulenceClosures: FluxTapering
using Oceananigans.Models.HydrostaticFreeSurfaceModels: SliceEnsembleSize
using Distributions
using Printf
using LinearAlgebra: norm

# ## Set up the problem and generate observations

# Define the  "true" skew and symmetric diffusivity coefficients. These are the parameter values that we
# use to generate the data. Then, we'll see if the EKI calibration can recover these values.

κ_skew = 1000.0       # [m² s⁻¹] skew diffusivity
κ_symmetric = 900.0   # [m² s⁻¹] symmetric diffusivity
nothing # hide

# We gather the "true" parameters in a named tuple ``θ_*``:

θ★ = (; κ_skew, κ_symmetric)

# The experiment name and where the synthetic observations will be saved.
experiment_name = "baroclinic_adjustment"
data_path = experiment_name * ".jld2"

# The domain, number of grid points, and other parameters.
architecture = CPU()      # CPU or GPU?
Ly = 1000kilometers       # north-south extent [m]
Lz = 1kilometers          # depth [m]
Ny = 64                   # grid points in north-south direction
Nz = 16                   # grid points in the vertical
Δt = 10minute             # time-step
stop_time = 1days         # length of run
save_interval = 0.25days  # save observation every so often

force_generate_observations = false
nothing # hide

horizontal_diffusivity = HorizontalScalarDiffusivity(κ=100)
vertical_diffusivity = VerticalScalarDiffusivity(κ=1e-2)

# The isopycnal skew-symmetric diffusivity closure.
gerdes_koberle_willebrand_tapering = FluxTapering(1e-2)
gent_mcwilliams_diffusivity = IsopycnalSkewSymmetricDiffusivity(κ_skew = κ_skew,
                                                                κ_symmetric = κ_symmetric,
                                                                slope_limiter = gerdes_koberle_willebrand_tapering)

# ## Generate synthetic observations

if force_generate_observations || !(isfile(data_path))
    grid = RectilinearGrid(architecture,
                           topology = (Flat, Bounded, Bounded),
                           size = (Ny, Nz),
                           y = (-Ly/2, Ly/2),
                           z = (-Lz, 0),
                           halo = (3, 3))

    closures = (gent_mcwilliams_diffusivity, horizontal_diffusivity, vertical_diffusivity)

    model = HydrostaticFreeSurfaceModel(grid = grid,
                                        tracers = (:b, :c),
                                        buoyancy = BuoyancyTracer(),
                                        coriolis = BetaPlane(latitude=-45),
                                        closure = closures,
                                        free_surface = ImplicitFreeSurface())

    @info "Built $model."

    ##### Initial conditions of an unstable buoyancy front

    """
    Linear ramp from 0 to 1 between -Δy/2 and +Δy/2.

    For example:

    y < y₀           => ramp = 0
    y₀ < y < y₀ + Δy => ramp = y / Δy
    y > y₀ + Δy      => ramp = 1
    """
    ramp(y, Δy) = min(max(0, y/Δy + 1/2), 1)

    N² = 4e-6             # [s⁻²] buoyancy frequency / stratification
    M² = 8e-8             # [s⁻²] horizontal buoyancy gradient

    Δy = 50kilometers     # horizontal extent of the font

    Δc_y = 2Δy            # horizontal extent of initial tracer concentration
    Δc_z = 50             # [m] vertical extent of initial tracer concentration

    Δb = Δy * M²          # inital buoyancy jump

    bᵢ(x, y, z) = N² * z + Δb * ramp(y, Δy)
    cᵢ(x, y, z) = exp(-y^2 / 2Δc_y^2) * exp(-(z + Lz/2)^2 / (2Δc_z^2))

    set!(model, b=bᵢ, c=cᵢ)

    simulation = Simulation(model, Δt=Δt, stop_time=stop_time)

    simulation.output_writers[:fields] = JLD2OutputWriter(model, merge(model.velocities, model.tracers),
                                                          schedule = TimeInterval(save_interval),
                                                          filename = experiment_name,
                                                          array_type = Array{Float64},
                                                          with_halos = true,
                                                          overwrite_existing = true)

    run!(simulation)
end

# ## Load truth data as observations

# We use here the `Transformation` functionality to slice up the observation data a bit.
# In particular, we choose to exclude the 3 grid points on either side of the `y` dimension,
# and 3 grid points from the bottom of the domain. Also, we only use the last 3 snapshots of
# the observations.
#
# We use `SpaceIndices` and `TimeIndices` to denote which space-time indices we would like to
# keep in observations.

transformation = Transformation(space=SpaceIndices(y=4:Ny-3, z=4:Nz), time=TimeIndices(3:5), normalization=ZScore())
observations = SyntheticObservations(data_path; field_names=(:b, :c), transformation)

# ## Calibration with Ensemble Kalman Inversion

# ### Ensemble model

# First we set up an ensemble model,
ensemble_size = 20

slice_ensemble_size = SliceEnsembleSize(size=(Ny, Nz), ensemble=ensemble_size)

@show ensemble_grid = RectilinearGrid(architecture,
                                      size=slice_ensemble_size,
                                      topology = (Flat, Bounded, Bounded),
                                      y = (-Ly/2, Ly/2),
                                      z = (-Lz, 0),
                                      halo=(3, 3))

gm_ensemble = [deepcopy(gent_mcwilliams_diffusivity) for i = 1:ensemble_size]
closures = (gm_ensemble, horizontal_diffusivity, vertical_diffusivity)

@show ensemble_model = HydrostaticFreeSurfaceModel(grid = ensemble_grid,
                                                   tracers = (:b, :c),
                                                   buoyancy = BuoyancyTracer(),
                                                   coriolis = BetaPlane(latitude=-45),
                                                   closure = closures,
                                                   free_surface = ImplicitFreeSurface())

# and then we create an ensemble simulation:

ensemble_simulation = Simulation(ensemble_model; Δt, stop_time)

ensemble_simulation

# ### Free parameters
#
# We construct some prior distributions for our free parameters. We found that it often helps to
# constrain the prior distributions so that neither very high nor very low values for diffusivities
# can be drawn out of the distribution.

priors = (κ_skew = ScaledLogitNormal(bounds=(400.0, 1300.0)),
          κ_symmetric = ScaledLogitNormal(bounds=(700.0, 1700.0)))

free_parameters = FreeParameters(priors)

# To visualize the prior distributions we randomly sample out values from then and plot the p.d.f.

using CairoMakie
using ParameterEstimocean.Parameters: unconstrained_prior, transform_to_constrained

samples(prior) = [transform_to_constrained(prior, x) for x in rand(unconstrained_prior(prior), 10000000)]

samples_κ_skew = samples(priors.κ_skew)
samples_κ_symmetric = samples(priors.κ_symmetric)

fig = Figure()
ax = Axis(fig[1, 1], xlabel = "Diffusivities [m² s⁻¹]", ylabel = "PDF")
densities = []
push!(densities, density!(ax, samples_κ_skew))
push!(densities, density!(ax, samples_κ_symmetric))
Legend(fig[1, 2], densities, ["κ_skew", "κ_symmetric"], valign = :bottom, halign = :left)

save("visualize_prior_diffusivities_baroclinic_adjustment.svg", fig); nothing # hide

# ![](visualize_prior_diffusivities_baroclinic_adjustment.svg)

# ### The inverse problem

# We can construct the inverse problem ``y = G(θ) + η``. Here, ``y`` are the `observations` and ``G`` is the
# `ensemble_model`.
calibration = InverseProblem(observations, ensemble_simulation, free_parameters)

# ### Assert that ``G(θ_*) ≈ y``
#
# As a sanity check we apply the `forward_map` on the calibration after we initialize all ensemble
# members with the true parameter values. We then confirm that the output of the `forward_map` matches
# the observations to machine precision.

G = forward_map(calibration, θ★)
y = observation_map(calibration)
nothing #hide

# The `forward_map` output `G` is a two-dimensional matrix whose first dimension is the size of the state
# space. Here, after the transformation we applied to the observations, we have that the state space size
# is `` 2 \times (N_y - 6) \times (N_z - 3) \times 3``; the 2 comes from the two tracers we used as observations
# and the 3 comes from only using the last three snapshots of the observations. The second dimension of
# the `forward_map` output is the `ensemble_size`.

@show size(G) == (2 * (Ny-6) * (Nz-3) * 3, ensemble_size)

# Since above we computed `G` using the true parameters ``θ_*``, all columns of the forward map output should
# be the same as the observations:

mean(G, dims=2) ≈ y

# Next, we construct an `EnsembleKalmanInversion` (EKI) object,

eki = EnsembleKalmanInversion(calibration; pseudo_stepping=ConstantConvergence(0.4))

# and perform few iterations to see if we can converge to the true parameter values.

params = iterate!(eki; iterations = 5)

@show params

# Last, we visualize few metrics regarding how the EKI calibration went about.

θ̅(iteration) = [eki.iteration_summaries[iteration].ensemble_mean...]
varθ(iteration) = eki.iteration_summaries[iteration].ensemble_var

weight_distances = [norm(θ̅(iter) - [θ★[1], θ★[2]]) for iter in 1:eki.iteration]
output_distances = [norm(forward_map(calibration, θ̅(iter))[:, 1] - y) for iter in 1:eki.iteration]
ensemble_variances = [varθ(iter) for iter in 1:eki.iteration]

f = Figure()
lines(f[1, 1], 1:eki.iteration, weight_distances, color = :red, linewidth = 2,
      axis = (title = "Parameter distance",
              xlabel = "Iteration",
              ylabel="|θ̅ₙ - θ⋆|",
              yscale = log10))
lines(f[1, 2], 1:eki.iteration, output_distances, color = :blue, linewidth = 2,
      axis = (title = "Output distance",
              xlabel = "Iteration",
              ylabel="|G(θ̅ₙ) - y|",
              yscale = log10))
ax3 = Axis(f[2, 1:2], title = "Parameter convergence",
           xlabel = "Iteration",
           ylabel = "Ensemble variance",
           yscale = log10)

for (i, pname) in enumerate(free_parameters.names)
    ev = getindex.(ensemble_variances, i)
    lines!(ax3, 1:eki.iteration, ev / ev[1], label = String(pname), linewidth = 2)
end

axislegend(ax3, valign = :top, halign = :right)
save("summary_baroclinic_adjustment.svg", f); nothing #hide

# ![](summary_baroclinic_adjustment.svg)

# And also we plot the the distributions of the various model ensembles for few EKI iterations to see
# if and how well they converge to the true diffusivity values.

f = Figure()

axtop = Axis(f[1, 1])

axmain = Axis(f[2, 1],
              xlabel = "κ_skew [m² s⁻¹]",
              ylabel = "κ_symmetric [m² s⁻¹]")

axright = Axis(f[2, 2])
scatters = []
labels = String[]

for iteration in [0, 1, 2, 5]
    ## Make parameter matrix
    parameters = eki.iteration_summaries[iteration].parameters
    Nensemble = length(parameters)
    Nparameters = length(first(parameters))
    parameter_ensemble_matrix = [parameters[i][j] for i=1:Nensemble, j=1:Nparameters]

    label = iteration == 0 ? "Initial ensemble" : "Iteration $iteration"
    push!(labels, label)
    push!(scatters, scatter!(axmain, parameter_ensemble_matrix))
    density!(axtop, parameter_ensemble_matrix[:, 1])
    density!(axright, parameter_ensemble_matrix[:, 2], direction = :y)
end

vlines!(axmain, [κ_skew], color = :red)
vlines!(axtop, [κ_skew], color = :red)

hlines!(axmain, [κ_symmetric], color = :red)
hlines!(axright, [κ_symmetric], color = :red)

colsize!(f.layout, 1, Fixed(300))
colsize!(f.layout, 2, Fixed(200))

rowsize!(f.layout, 1, Fixed(200))
rowsize!(f.layout, 2, Fixed(300))

Legend(f[1, 2], scatters, labels, valign = :bottom, halign = :left)

hidedecorations!(axtop, grid = false)
hidedecorations!(axright, grid = false)

xlims!(axmain, 350, 1350)
xlims!(axtop, 350, 1350)
ylims!(axmain, 650, 1750)
ylims!(axright, 650, 1750)
xlims!(axright, 0, 0.015)
ylims!(axtop, 0, 0.015)

save("distributions_baroclinic_adjustment.svg", f); nothing #hide

# ![](distributions_baroclinic_adjustment.svg)


text:
@testset "save" begin

    @test_nowarn begin
        m = @agent(0,
            l::Local,
            g::Global,
            ga::GlobalArray
        )
        model = compile(m,save="RAM",debug=false)

        com = Community(model,N=5)
        com.l .= 0.1 .*1:5
        com.id .= 1:5
        com.g = 6
        com.ga = [1 2;3 4]

        saveCSV(com,"testCSV")

        com2 = loadCommunityFromCSV(model,"testCSV")
    end

    m = @agent(0,
        l::Local,
        g::Global,
        ga::GlobalArray
    )
    model = compile(m,save="RAM",debug=false)
    com = Community(model,N=5)
    com.l .= 0.1 .*1:5
    com.id .= 1:5
    com.g = 6
    com.ga = [1 2;3 4]
    saveCSV(com,"testCSV")
    com2 = loadCommunityFromCSV(model,"testCSV")
    @test com.declaredSymbols_ == com2.declaredSymbols_
    @test com.local_ == com2.local_
    @test com.identity_ == com2.identity_
    @test com.global_ == com2.global_
    @test com.globalArray_ == com2.globalArray_
    @test com.t == com2.t
    @test com.N == com2.N

    for platform in testplatforms

        m = @agent(0,
            l::GlobalArray,

            UpdateGlobal = begin
                    l[1,1] += 1
                end
            )
        model = compile(m,save="RAM",debug=false,platform=platform)
        # println(model.program)
        com = Community(model)
        com.l = zeros(2,2)
        comt = model.evolve(com,dt=0.1,tMax=10.05,dtSave=0.2)

        @test length(comt)==51
        @test comt[end].l[1,2] == 0.
        @test comt[1].l[1,1] == 0
    end

    for platform in testplatforms

        m = @agent(0,
            l::GlobalArray,

            UpdateGlobal = begin
                    l[1,1] += 1
                end
            )
        model = compile(m,save="CSV",debug=false,platform=platform)
        # println(model.program)
        com = Community(model)
        com.l = zeros(2,2)
        model.evolve(com, dt=0.1, tMax=10.05, dtSave=0.2, saveFile="testCSV")

        @test_nowarn loadCommunityInTimeFromCSV(model,"testCSV")
    end

    dir = readdir("./")
    for i in dir
        if occursin(".csv",i)
            rm(i)
        end
    end

end


text:
function mark_centers(r)
    num_steps = through_rectangles_into_angle(r,(Sud,West))
    num_steps_to_ost = sum(num_steps[1:2:end])
    num_steps_to_nord = sum(num_steps[2:2:end])

    movements!(r,Nord,num_steps_to_nord)
    putmarker!(r)
    num_steps_to_sud = movements!(r,Nord)

    movements!(r,Ost,num_steps_to_ost)
    putmarker!(r)
    num_steps_to_west = movements!(r,Ost)

    movements!(r,Sud,num_steps_to_sud)
    putmarker!(r)
    movements!(r,Sud)

    movements!(r,Sud,num_steps_to_west)
    putmarker!(r)
    movements!(r,Sud)

    movements!(r,(Ost,Nord),num_steps)

end

max_stars_repo_path: zadacha 11.jl

max_stars_repo_name: Anpilogova/-

max_issues_repo_path: zadacha 11.jl

max_issues_repo_name: Anpilogova/-

max_forks_repo_path: zadacha 11.jl

max_forks_repo_name: Anpilogova/-


text:
using Turkie
using Test
using CairoMakie
CairoMakie.activate!()
using OnlineStats
using Turing

@testset "Turkie.jl" begin
    Turing.@model function demo(x) # Some random Turing model
        m0 ~ Normal(0, 2)
        s ~ InverseGamma(2, 3)
        m ~ Normal(m0, √s)
        for i in eachindex(x)
            x[i] ~ Normal(m, √s)
        end
    end
    N = 10
    xs = randn(N) .+ 1
    model = demo(xs)
    vars = [:m0, :s, :m]
    @testset "Interface" begin
        @test Turkie.std_colors == Turkie.ColorSchemes.seaborn_colorblind
        @test Turkie.name(:blah) == "blah"
        @test Turkie.name(OnlineStats.Mean(Float32)) == "Mean"

        cb = TurkieCallback(model; blah=2.0)
        @test cb.figure isa Figure
        @test sort(collect(keys(cb.data))) == sort(vcat(vars, :iter))
        @test cb.data[:m] isa MovingWindow{Float32}
        @test sort(collect(keys(cb.vars))) == sort(vars)
        @test cb.vars[:m][1] == :histkde
        @test cb.vars[:m][2] == Mean(Float32)
        @test cb.vars[:m][3] == Variance(Float32)
        @test cb.vars[:m][4] == AutoCov(20, Float32)
        @test length(cb.axis_dict) == 15
        @test cb.params[:blah] == 2.0
        @test cb.params[:window] == 1000
        @test cb.iter[] == 0
        @test cb.iter isa Observable{Int}
    end
    @testset "Testing all possible plots" begin
        @testset "Vector of symbols" begin
            for stat in [:histkde, :kde, :hist, :mean, :var, :trace, :autocov]
                cb = TurkieCallback(Dict(:m => [stat]))
                sample(model, MH(), 50; progress=false, callback=cb)
            end
        end
        @testset "Series" begin
            for stat in [Mean(Float32), Variance(Float32)]
                cb = TurkieCallback(model, OnlineStats.Series(stat))
                sample(model, MH(), 50; progress=false, callback=cb)
            end
        end
    end

end


text:
function e002(N)
    return sum_evens(N)
end

function fib(channel::Channel, num0::Integer, num1::Integer)
    put!(channel, num0)
    put!(channel, num1)
    while true
        sum = num0 + num1
        num0 = num1
        num1 = sum
        put!(channel, sum)
    end
end

function sum_evens(N)
    channel = Channel(c -> fib(c,1,2))
    i = take!(channel)
    sum = 0
    while i < N
        if i % 2 == 0
            sum += i
        end
        i = take!(channel)
    end
    close(channel)
    return sum
end

function solution002()
    N = 4000000
    answer = sum_evens(N)
    println("The answer to problem 002 is $answer")
end

max_stars_repo_name: mithodin/ttpe

max_issues_repo_name: mithodin/ttpe

max_forks_repo_name: mithodin/ttpe


text:
"""
Mig() |>

	check_sql(\"EXISTS TABLE rnd600.t01\") |>

	migrate_sql(\"""
	CREATE TABLE rnd600.t01
	...
	\""") |>

	rollback_sql(\"DROP TABLE rnd600.t01\") |>

        cache_text( s\"""(m)-> (KH, server) = "chouse_conf.jl"|>include \""" )|>

        check_text( s\""" function (m)
         (KH, server) = cached(m)
         dbexists = any( _->ismatch( r"RND600", _), KH.READONLY( server, "SHOW DATABASES")|>readlines )
        end \""") |>

        migrate_text( s\""" function(m)
         (KH, server) = cached(m)
         KH.MODIFY( server, "CREATE DATABASE RND600")
        end \""", danger=false) |>

        rollback_text( s\""" function(m)
         (KH, server) = cached(m)
         KH.MODIFY( server, "DROP DATABASE RND600" )
        end \""") |> migrate

        # or ... |> rollback( DANGER=true)

        # or ... |> check
"""
module Migrations

typealias Str AbstractString

abstract Action

immutable NotImplemented<:Action end
typealias NotImp NotImplemented

immutable Implemented<:Action
 expr::Expr
 func::Function
 danger::Bool

 function Implemented( t::Str, applicable_to_args::Tuple; danger=false)
  try ex = parse(t)

    if isa( ex, Expr)
        ex.head == :incomplete && error("Incomplete expression:\n$t")
    else
        error("Bad type: $(ex|>typeof). Must be Expr. \nText: $t")
    end
    f = eval(ex)

    !isa( f, Function) && error("Bad type: $(f|>typeof). Must be a Function. \nExpr: $ex")

    if !applicable( f, applicable_to_args... )
        types = map( typeof, applicable_to_args)
        mets = methods(f)
        warn("Must be applicable to argument types: $types.\nNow only applicable to: $mets.")
        error("Creation of Implemented() was stopped.")
    end

    !danger && ismatch( r"\bDROP\b", t) &&
        "Found smth dangerous in text:\n$t\nYou must mark action as danger=true"|>error

    new( ex, f, danger)

  catch e error("$e\nParsed text:\n$t") end
 end
end

typealias Imp Implemented

abstract Cache

immutable EmptyCache<:Cache end

immutable FullCache{T}<:Cache
 value::T
end

"""
using Migrations

m = Migration()"""
immutable Migration
 check_sql::AbstractString
 migrate_sql::AbstractString
 rollback_sql::AbstractString
 cache::Action
 cached::Cache
 check::Action
 migrate::Action
 rollback::Action
end

"typealias Mig == Migration"
typealias Mig Migration

export Mig, Migration

Mig(; check_sql::AbstractString="",
      migrate_sql::AbstractString="",
      rollback_sql::AbstractString="",
      cache::Action=NotImp(),
      cached::Cache=EmptyCache(),
      check::Action=NotImp(),
      migrate::Action=NotImp(),
      rollback::Action=NotImp() ) = Mig( check_sql, migrate_sql, rollback_sql, cache, cached, check, migrate, rollback )

"""
Migrations.Migration() |> check_sql(\"""EXISTS TABLE db1.table1\""")
"""
check_sql( t::Str) = (m::Mig)->check_sql( m, t)
export check_sql

"""
m = Migrations.Migration();

check_sql( m, \"""(m)->info(\"EXISTS TABLE db1.table1\")\""" )
"""
check_sql( m::Mig, t::Str) = Mig(
    check_sql=t, migrate_sql=m.migrate_sql, rollback_sql=m.rollback_sql,
    cache=m.cache, cached=m.cached, check=m.check, migrate=m.migrate, rollback=m.rollback)

"""
Migrations.Migration() |> migrate_sql(\"""EXISTS TABLE db1.table1\""")
"""
migrate_sql( t::Str) = (m::Mig)->migrate_sql( m, t)
export migrate_sql

"""
m = Migrations.Migration();

migrate_sql( m, \"""(m)->info(\"CREATE TABLE db1.table1 ...  \")\""" )
"""
migrate_sql( m::Mig, t::Str) = Mig(
    check_sql=m.check_sql, migrate_sql=t, rollback_sql=m.rollback_sql,
    cache=m.cache, cached=m.cached, check=m.check, migrate=m.migrate, rollback=m.rollback)

"""
Migrations.Migration() |> rollback_sql(\"""EXISTS TABLE db1.table1\""")
"""
rollback_sql( t::Str) = (m::Mig)->rollback_sql( m, t)
export rollback_sql

"""
m = Migrations.Migration();

rollback_sql( m, \"""(m)->info(\"DROP TABLE db1.table1  \")\""" )
"""
rollback_sql( m::Mig, t::Str) = Mig(
    check_sql=m.check_sql, migrate_sql=m.migrate_sql, rollback_sql=t,
    cache=m.cache, cached=m.cached, check=m.check, migrate=m.migrate, rollback=m.rollback)

"""
Migrations.Migration() |> cache_text(\"""(m)->true\""")
"""
cache_text( t::Str) = (m::Mig)->cache_text( m, t)
export cache_text

"""
m = Migrations.Migration();

cache_text( m, \"""(m)->info(\"hello from migration\")\""" )
"""
cache_text( m::Mig, t::Str) = Mig(
    check_sql=m.check_sql, migrate_sql=m.migrate_sql, rollback_sql=m.rollback_sql,
    cache=Imp( t, (m,)), cached=m.cached, check=m.check, migrate=m.migrate, rollback=m.rollback)

"""
expr::Expr = m|>cache_text # return 'cache' action from migration
"""
cache_text( m::Mig) = expr( m.cache)

"""
Migrations.Migration() |> check_text(\"""(m)->true\""")
"""
check_text( t::Str) = (m::Mig)->check_text( m, t)
export check_text

"""
m = Migrations.Migration();

check_text( m, \"""(m)->info(\"hello from migration\")\""" )
"""
check_text( m::Mig, t::Str) = Mig(
    check_sql=m.check_sql, migrate_sql=m.migrate_sql, rollback_sql=m.rollback_sql,
    cache=m.cache, cached=m.cached, check=Imp( t, (m,)), migrate=m.migrate, rollback=m.rollback)

"""
expr::Expr = m|>check_text
"""
check_text( m::Mig) = expr( m.check)

"""
Migrations.Migration() |> migrate_text(\"""(m)->info(\"hello from migration\")\""")
"""
migrate_text( t::Str; danger::Bool=true ) =
    (m::Mig)->migrate_text( m, t, danger=danger)
export migrate_text

"""
m = Migrations.Migration();

migrate_text( m, \"""(m)->info(\"hello from migration\")\""" )
"""
migrate_text( m::Mig, t::Str; danger::Bool=true) = Mig(
    check_sql=m.check_sql, migrate_sql=m.migrate_sql, rollback_sql=m.rollback_sql,
    cache=m.cache, cached=m.cached, check=m.check, migrate=Imp( t, (m,), danger=danger), rollback=m.rollback)

"""
expr::Expr = m|>migrate_text
"""
migrate_text( m::Mig) = expr( m.migrate)

expr( ni::NotImp) = ni
expr( imp::Imp) = imp.expr

"""
expr::Expr = m|>rollback_text
"""
rollback_text( m::Mig) = expr( m.rollback)
export rollback_text

"""
Migrations.Migration() |> rollback_text(\"""(m)->info(\"hello from rollback\")\""")
"""
rollback_text( t::Str; danger::Bool=true) = (m::Mig)->rollback_text( m, t, danger=danger)

"""
m = Migrations.Migration();

rollback_text( m, \"""(m)->info(\"hello from migration\")\""" )
"""
rollback_text( m::Mig, t::Str; danger::Bool=true) = Mig(
    check_sql=m.check_sql, migrate_sql=m.migrate_sql, rollback_sql=m.rollback_sql,
    cache=m.cache, cached=m.cached, check=m.check, migrate=m.migrate, rollback=Imp( t, (m,), danger=danger))

"""
Returns old m::Migration if cache() not implemented, and return new with cached value otherwise

m|>cache
"""
function cache( m::Mig; debug::Bool=false)::Mig
 if isa( m.cached, EmptyCache )
    rv = doit( m.cache, m, debug=debug, prefix="Run cache()")
    isa( rv, NotImp ) && return m
    isa( rv, Void ) && error("cache() not return any value! Must return something.\n$(m.cache.expr)")
    Mig(
	check_sql=m.check_sql, migrate_sql=m.migrate_sql, rollback_sql=m.rollback_sql,
	cache=m.cache, cached=FullCache(rv), check=m.check, migrate=m.migrate, rollback=m.rollback )
 else
    m
 end
end
export cache

"Returns migrations cached value (returned from cache())"
cached(m::Mig) = isa( m.cached, FullCache) ?
    m.cached.value :
        isa( m.cache, NotImp ) ?
            error("Migration has not cached value, because cache() not implemented.") :
            error("Migration has not cached value. Forgot to call cache()? cache expression:\n $(m.cache.expr)")
export cached

"""Checks is migration successed.

m|>check"""
function check( m::Mig; debug::Bool=false, prefix::Str="check()")
 isa( m.cached, EmptyCache) && ( m = cache( m, debug=debug))
 rv = doit( m.check, m, debug=debug, prefix=prefix)
 if isa( rv, Bool)
    debug && info("Return: $rv")
    return rv
 end
 isa( rv, NotImp ) && error("Not implemented check() for migration $m")
 error("Bad type $(rv|>typeof) of returned value ($rv). Must be a Bool. check():\n$(m.check.expr)")
end
export check

check(;debug::Bool=false, prefix::Str="check()") = (m::Mig)->check( m, debug=debug, prefix=prefix)

function migrate( m::Mig; checking::Bool=true, DANGER::Bool=false, debug::Bool=true )
    isa( m.cached, EmptyCache) && (m = cache( m, debug=debug))
    if checking
        chk1 = check( m, debug=true, prefix="Check before migrate()")::Bool
        chk1 && return Dict(:check_before=>chk1, :migrate=>:not_called, :resume=>:not_need)
        mgr = doit( m.migrate, m, DANGER=DANGER, debug=debug, prefix="migrate()")
        chk2 = check( m, debug=true, prefix="Check after migrate()")::Bool
        !chk2 && warn( "Check return $chk2 after migration!")
        Dict( :check_before=>chk1, :migrate=>mgr, :check_after=>chk2, :resume=>(chk2? :success : :fail) )
    else
        mgr = doit( m.migrate, m, DANGER=DANGER, debug=debug, prefix="migrate() without checking" )
        Dict( :check_before=>:not_called, :migrate=>mgr, :check_after=>:not_called, :resume=>:not_checked_success)
    end
end
export migrate

"""
m|>migrate

m|>migrate( check=false) # without check()

does migrate if check returns false """
migrate( ;checking::Bool=true, DANGER::Bool=false, debug::Bool=true) = (m::Mig)->migrate( m, checking=checking, DANGER=DANGER, debug=debug)

"""rollback( m)

rollback( m, check=false) # without check()

does rollback if check() returns true """
function rollback( m::Mig; checking::Bool=true, DANGER::Bool=false, debug::Bool=true)
    isa( m.cached, EmptyCache) && ( m = cache( m, debug=debug))
    if checking
        ch1 = check( m, debug=true, debug=debug, prefix="Checking before rollback")::Bool
        !ch1 && return Dict( :check_before=>ch1, :rollback=>:not_called, :resume=>:not_need)
        rbk = doit( m.rollback, m, DANGER=DANGER, debug=debug, prefix="rollback()")
        ch2 = check( m, debug=true, prefix="Checking after rollback")::Bool
        ch2 && warn("Check after rollback return $ch2!")
        Dict( :check_before=>ch1, :rollback=>rbk, :check_after=>ch2, :resume=>(ch2? :fail : :success))
    else
        rbk = doit( m.rollback, m, DANGER=DANGER, debug=debug, prefix="rollback()")
        Dict( :check_before=>:skipped, :rollback=>rbk, :check_after=>:skipped, :resume=>:not_checked_success)
    end
end
export rollback

"""m|>rollback

m|>rollback( check=false)
"""
rollback( ;checking::Bool=true, DANGER::Bool=false, debug::Bool=true ) = (m::Mig)->rollback( m, checking=checking, DANGER=DANGER, debug=debug)

function doit( i::Imp, args... ; DANGER=false, debug::Bool=false, prefix::Str="" )
    !DANGER && i.danger && error("Canceled dangerous action. Must be called with DANGER=true.\n$(i.expr)")
    debug && info("$prefix call:\n$(i.expr)")
    try rv = i.func( args...)
        debug && info("$prefix return: $rv")
        rv
    catch e
        error("$prefix: $e\n$(i.expr)\nArguments:\n$(args...)")
    end
end

doit( ni::NotImp, args...; DANGER=false, debug::Bool=false, prefix::Str="" )::NotImp = ni

end # module

max_stars_repo_name: closescreen/Migrations

max_issues_repo_name: closescreen/Migrations

max_forks_repo_name: closescreen/Migrations