text:
function f_min(iter, factor)
    return  (iter * factor) -(factor-1)
end

function f_max(iter, factor)
    return (iter * factor)
end

function projection()

    println("Predefined vars for projections:")
    println("------------------------------------------------")
    println("=====================[gas]:=====================")
    println("       -all the non derived hydro vars-")
    println(":cpu, :level, :rho, :cx, :cy, :cz, :vx, :vy, :vz, :p, var6,...")
    println("further possibilities: :rho, :density, :ρ")
    println("              -derived hydro vars-")
    println(":x, :y, :z")
    println(":sd or :Σ or :surfacedensity")
    println(":mass, :cellsize, :freefall_time")
    println(":cs, :mach, :jeanslength, :jeansnumber")
    println()
    println("==================[particles]:==================")
    println("        all the non derived  vars:")
    println(":cpu, :level, :id, :family, :tag ")
    println(":x, :y, :z, :vx, :vy, :vz, :mass, :birth, :metal....")
    println()
    println("              -derived particle vars-")
    println(":age")
    println()
    println("==============[gas or particles]:===============")
    println(":v, :ekin")
    println("squared => :vx2, :vy2, :vz2")
    println("velocity dispersion => σx, σy, σz, σ")
    println()
    println("related to a given center:")
    println("---------------------------")
    println(":vr_cylinder, vr_sphere (radial components)")
    println(":vϕ_cylinder, :vθ")
    println("squared => :vr_cylinder2, :vϕ_cylinder2")
    println("velocity dispersion => σr_cylinder, σϕ_cylinder ")
    #println(":l, :lx, :ly, :lz :lr, :lϕ, :lθ")
    println()
    println("2d maps (not projected):")
    println(":r_cylinder") #, :r_sphere")
    println(":ϕ") # :θ
    println("------------------------------------------------")
    println()
    return
end

"""
#### Remap projected data of DataMapsType onto a coarser grid:
- select the level of the coarser grid
- select the weighting
- toggle verbose mode

```julia
remap(dataobject::DataMapsType, lmax::Real; weighting::Symbol=:volume, verbose::Bool=verbose_mode)

return DataMapsType
```
#### Arguments
##### Required:
- **`dataobject`:** needs to be of type: "DataMapsType", e.g. created by the function *projection*
- **`lmax`:** the level of the coarser grid to be created
##### Predefined/Optional Keywords:
- **`weighting`:** choose between: :volume, :mass
- **`verbose`:** print the dimensions/pixel size of the provided and created maps on screen; default: set by the variable `verbose_mode`
"""
function remap(dataobject::DataMapsType, lmax::Real; weighting::Symbol=:volume, verbose::Bool=verbose_mode)

    # checks to use maps instead of mass weighting
    rcheck = [:r_cylinder, :r_sphere]
    anglecheck = [:ϕ]
    ranglecheck = [rcheck..., anglecheck...]

    mcheck = [:sd, :Σ, :surfacedensity,:density, :rho, :ρ]
    mvar = :no
    if weighting == :mass
        notonly_ranglecheck_vars = checkformaps(dataobject, ranglecheck)
        if notonly_ranglecheck_vars
            no_mass_map = false
            for ivar in keys(dataobject.maps)
                if in(ivar, mcheck)
                    mvar = ivar
                    no_mass_map = false
                    break
                else
                    no_mass_map = true
                end
            end
            if no_mass_map
                error("""[Mera]: For mass weighting a map with "rho" or "sd" is necessary.""")
            end
        end
    end

    simlmax = dataobject.lmax
    lmax_projected = lmax

    maps = dataobject.maps
    selected_vars= Symbol[]
    for k in keys(maps)
        push!(selected_vars, k)
    end

    maps_unit = dataobject.maps_unit
    units= Symbol[]
    for k in selected_vars
        kunit = maps_unit[ Symbol(string(k) ) ]
        push!(units, kunit)
    end
    σcheck = [:σx, :σy, :σz, :σ, :σr_cylinder, :σϕ_cylinder]
    σ_to_v = SortedDict(  :σx => [:vx, :vx2],
                          :σy => [:vy, :vy2],
                          :σz => [:vz, :vz2],
                          :σ  => [:v,  :v2],
                          :σr_cylinder => [:vr_cylinder, :vr_cylinder2],
                          :σϕ_cylinder => [:vϕ_cylinder, :vϕ_cylinder2] )

    wcheck = [:r_cylinder, :r_sphere, :ϕ, :sd, :Σ, :surfacedensity,:density, :rho, :ρ] # 2d maps (no projections)

    maps_lmax = SortedDict()
    # remap onto lmax grid
    if simlmax > lmax
        if verbose
            println()
            println("remap from:")
            println("level ", simlmax, " => ", lmax)

            min_cellsize, min_unit  = humanize(dataobject.info.boxlen / 2^lmax, dataobject.info.scale, 2, "length")
            max_cellsize, max_unit  = humanize(dataobject.info.boxlen / 2^simlmax, dataobject.info.scale, 2, "length")
            println("cellsize ", max_cellsize," [$max_unit] => ", min_cellsize ," [$min_unit]")
        end

        first_time  =1
        for ivar in selected_vars
            if !in(ivar, σcheck)
                if mvar == :no || in(ivar, wcheck)
                    maps_buffer = maps[Symbol(ivar)]
                else # use mass-weighting
                    maps_buffer = maps[Symbol(ivar)] .* maps[Symbol(mvar)]
                end

                s = size(maps_buffer)
                lmax_ratio = 2^simlmax / 2^lmax
                scaled_length1 = round(Int, s[1] / lmax_ratio)
                scaled_length2 = round(Int, s[2] / lmax_ratio)
                if first_time == 1
                    if verbose println("pixels ", s, " => ($scaled_length1, $scaled_length2)" ) end
                    first_time = 0
                end
                maps_lmax[Symbol(ivar)] = [maps_buffer[floor(Int,x),floor(Int,y)]  for x in range(1, stop=s[1], length=scaled_length1), y in range(1, stop=s[2], length=scaled_length2)]
            end

        end
    else
        error("lmax of simulation is =< given lmax.")
    end

    if mvar != :no # reverse mass-weighting
        for ivar in selected_vars
            if  !in(ivar, wcheck) && !in(ivar, σcheck) # exclude 2d maps (no projections) and velocity dispersion maps
                maps_lmax[Symbol(ivar)] = maps_lmax[Symbol(ivar)] ./ maps_lmax[Symbol(mvar)]
            end
        end
    end

    # create velocity dispersion maps from rebinned maps
    if simlmax > lmax
        for ivar in selected_vars

            if in(ivar, σcheck)
                    selected_unit, unit_name= getunit(dataobject, ivar, selected_vars, units, uname=true)

                        selected_v = σ_to_v[ivar]
                        iv  = maps_lmax[selected_v[1]]
                        iv_unit = maps_unit[Symbol( string(selected_v[1])  )]
                        iv2 = maps_lmax[selected_v[2]]
                        iv2_unit = maps_unit[Symbol( string(selected_v[2])  )]
                        if iv_unit == iv2_unit
                            diff_iv = iv2 .- iv .^2
                            if typeof(dataobject) == PartMapsType
                                diff_iv[isnan.(diff_iv)] .= 0
                                diff_iv[diff_iv .< 0.] .= 0
                            end

                            if iv_unit == unit_name
                                maps_lmax[Symbol(ivar)] = sqrt.( diff_iv )
                            elseif iv_unit == :standard
                                maps_lmax[Symbol(ivar)] = sqrt.( diff_iv )  .* selected_unit
                            elseif iv_unit == :km_s
                                maps_lmax[Symbol(ivar)] = sqrt.( diff_iv )  ./ dataobject.info.scale.km_s
                            end
                        elseif iv_unit != iv2_unit
                            if iv_unit == :km_s && unit_name == :standard
                                iv = iv ./ dataobject.info.scale.km_s
                            elseif iv_unit == :standard && unit_name == :km_s
                                iv = iv .* dataobject.info.scale.km_s
                            end
                            if iv2_unit == :km_s && unit_name == :standard
                                iv2 = iv2 ./ dataobject.info.scale.km_s.^2
                            elseif iv2_unit == :standard && unit_name == :km_s
                                iv2 = iv2 .* dataobject.info.scale.km_s.^2
                            end

                            diff_iv = iv2 .- iv .^2
                            if typeof(dataobject) == PartMapsType
                                diff_iv[isnan.(diff_iv)] .= 0
                                diff_iv[diff_iv .< 0.] .= 0
                            end

                            maps_lmax[Symbol(ivar)]  = sqrt.( diff_iv )
                        end

            end
        end
    end

    if typeof(dataobject) == HydroMapsType
        return HydroMapsType(maps, maps_unit, maps_lmax, SortedDict(),
                        lmax_projected,
                        dataobject.lmin,
                        simlmax,
                        dataobject.ranges,
                        dataobject.extent,
                        dataobject.cextent,
                        dataobject.ratio,
                        dataobject.boxlen,
                        dataobject.smallr, dataobject.smallc, dataobject.scale,
                        dataobject.info)
    elseif typeof(dataobject) == PartMapsType
        return PartMapsType(maps, maps_unit, maps_lmax, SortedDict(),
                            lmax_projected,
                            dataobject.lmin,
                            simlmax,
                            dataobject.ref_time,
                            dataobject.ranges,
                            dataobject.extent,
                            dataobject.cextent,
                            dataobject.ratio,
                            dataobject.boxlen,
                            dataobject.scale,
                            dataobject.info)
    end

end

# check if only variables from ranglecheck are selected
function checkformaps(selected_vars::Array{Symbol,1}, reference_vars::Array{Symbol,1})
    Nvars = length(selected_vars)
    cw = 0
    for iw in selected_vars
        if in(iw,reference_vars)
            cw +=1
        end
    end
    Ndiff = Nvars-cw
    return Ndiff != 0
end

function checkformaps(dataobject::DataMapsType, reference_vars::Array{Symbol,1})
    Nvars =0
    cw = 0
    for iw in keys(dataobject.maps)
        Nvars +=1
        if in(iw,reference_vars)
            cw +=1
        end
    end
    Ndiff = Nvars-cw
    return Ndiff != 0
end


text:
using LoopVectorization, SnoopCompile
tinf = @snoopi_deep include(joinpath(pkgdir(LoopVectorization), "test", "runtests.jl"))

ttot, pcs = SnoopCompile.parcel(tinf);

lv, pcslv = last(pcs);
@assert lv === LoopVectorization

blacklist = (
  :_turbo_!, :vmaterialize!, :vmaterialize, Symbol("vreduce##kw"), :_vreduce_dims!, :vreduce, :vmapreduce, :SIMDMapBack, :launch_thread_vmap!, :_turbo_loopset_debug, :all_dense,
  :sigmoid_fast, :rrule, :add_broadcast!, :create_mrefs!, :avx_config_val, :subsetview, :ifelsepartial, :tanh_fast, :check_args, :relu, :init_dual
)
filteredmethods = filter(m -> !Base.sym_in(m[2].def.name, blacklist), last(pcslv)); length(filteredmethods)

SnoopCompile.write("/tmp/precompile_loopvec", [LoopVectorization => (sum(first,filteredmethods),filteredmethods)])

# pc = SnoopCompile.parcel(tinf; blacklist=["vmaterialize", "vmaterialize!", "vreduce", "Base.Broadcast.materialize", "_vreduce_dims!", "vmapreduce"])
# pcs = pc[:LoopVectorization]
# open(joinpath(pkgdir, "src", "precompile.jl"), "w") do io
#     println(io, """
#     function _precompile_()
#         ccall(:jl_generating_output, Cint, ()) == 1 || return nothing
#     """)
#     for stmt in sort(pcs)
#         println(io, "    ", stmt)
#     end
#     println(io, "end")
# end

ext: jl


text:
using FileManager
using Test

# NN
include("las.jl")
include("txt.jl")
include("hierarchy.jl")


text:
function rr_vle_vapor_fraction(K,z)
    #Solution of the Rachford Rice equation using perturbation analysis
    #https://doi.org/10.1016/j.fluid.2021.112981
    #strategy: find 4 strongest compositions
    #calculate exact β for those 4
    #real_β = β_ℵ1 + ε*β_ℵ2(ε), where ℵ1 is the strong group of component and
    #ℵ2 is the weak group of components.
    #on ε = 0, we have the aprox solution
    #on ε = 1 we have the real solution
    #ε*β_ℵ2(ε) ≈ ∑ci*ε^i
    #we obtain the coefficients and evaluate ε = 1.
    #for n = length(z) <= 4, we return exact solutions.

    if length(z) <= 4
        return rr_vle_vapor_fraction_exact(K,z)
    end

    #z_ℵ1 is normalized
    ℵ1,k_ℵ1,z_ℵ1 = Clapeyron.rr_find_strongest(K,z)
    β0 = rr_vle_vapor_fraction_exact(k_ℵ1,z_ℵ1)
    b0 = β0
    _0 = zero(first(z)+first(K))
    _1 = one(_0)
    M1,M2 = _0,_0
    N1,N2 = _0,_0
    Ξ1,Ξ2 = _0,_0
    P1 = _0
    Λ2 = _0
    for i in 1:length(z)
        βi = _1/(_1-K[i])
        γi = βi - b0
        γiinv = _1/γi
        λi = z[i]*γiinv
        μi = λi*γiinv
        νi = μi*γiinv
        ξi = νi*γiinv
        if i in ℵ1
            M1 += μi
            N1 += νi
            Ξ1 += ξi
            P1 += ξi*γiinv
        else
            M2 += μi
            N2 += νi
            Ξ2 += ξi
            Λ2 += λi
        end
    end
    b1 = -Λ2/M1
    b2 = -(b1*b1*N1 + b1*M2)/M1
    b11 = b1*b1
    b111 = b11*b1
    b12 = b1*b2
    b22 = b2*b2
    b3 = -((b111*Ξ1+2*b12*N1) + (b11*N2+b2*M2))/M1
    t1 = b11*b11*P1 + 3*b11*b2*Ξ1 + (b22+2*b1*b3)*N1
    t2 = b111*Ξ2 + 2*b12*N2 + b3*M2
    b4 = -(t1+t2)/M1
    #@show b0,b1,b2,b3,b4
    βx = (b0+b1+b2+b3+b4)
    #now we have 5 points to select.
    r0 = rr_flash_eval(K,z,b0)
    β1 = b0+b1
    r1 = rr_flash_eval(K,z,β1)
    β2 = β1+b2
    r2 = rr_flash_eval(K,z,β2)
    β3 = β2+b3
    r3 = rr_flash_eval(K,z,β3)
    β4 = β3+b4
    r4 = rr_flash_eval(K,z,β4)
    bx = SA[r0,r1,r2,r3,r4]
    all_β = (β0,β1,β2,β3,β4)
    idx_b = sortperm(bx)
    sort_β = map(x->getindex(all_β,x),idx_b)
    #TODO: refine
    return β4
end
function rr_find_strongest(K,z)
    _0 = zero(first(z)+first(K))
    _1 = one(_0)
    (_,idmin) =  findmin(K)
    (_,idmax) =  findmax(K)
    id1 = 0
    id2 = 0

    idx = 0
    strong_zβ = (-_1,-_1)
    #find the strongest components, by zi/|βi|, ignoring βmin and βmax
    for i in 1:length(z)
        Ki = K[i]
        zi = z[i]
        idx+=1
        idx in (idmin,idmax) && continue
        zβi =abs(zi*(_1-Ki))
        zβ1,zβ2 = strong_zβ
        if zβ2 <= zβi
            strong_zβ = (zβ2,zβi)
            id1 = id2
            id2 = idx
        elseif zβ1 <= zβi
            strong_zβ = (zβi,zβ2)
            id1 = idx
        end
    end

    #indices of the strongest components
    idxs = (idmin,idmax,id1,id2)
    sumz = z[idmin] + z[idmax] +z[id1] +z[id2]
    invsumz = _1/sumz
    #normalize z
    zs = (invsumz*z[idmin],invsumz*z[idmax],invsumz*z[id1],invsumz*z[id2])
    ks = (K[idmin],K[idmax],K[id1],K[id2])
    return idxs,ks,zs
    end
    #=here are some tests, from the paper.
    the paper has errors, z2 has 10 components, 9 and 10 repeated.
    and after that, they dont include the bmax,bmin in their selected pair.
    z1 = [0.30,0.15, 0.05, 0.02, 0.01, 0.02, 0.02, 0.03, 0.07, 0.33]
    k1 = [3.0E+00, 2.0E+00, 1.1E+00, 8.0E-01, 4.0E-01, 1.0E-01, 5.0E-02, 2.0E-02, 1.0E-02, 1.0E-04]
    idxs1,ks1,zs1 = Clapeyron.rr_find_strongest(k1,z1)
    @test all(in.(idxs1,Ref((1,2,9,10))))

    z2 = [0.00034825,0.01376300,0.13084000,0.10925000,0.00001000,0.51009000,0.23564000,0.00006000]
    k2 = [5.2663E+02, 5.0400E+01, 1.6463E+00, 8.7450E-01, 1.5589E-01, 3.6588E-02, 2.6625E-02, 4.8918E-06]
    idxs2,ks2,zs2 = Clapeyron.rr_find_strongest(k2,z2)
    @test_broken all(in.(idxs2,Ref((1,2,6,7))))
    #the paper is wrong here.
    z3 = [0.0187002, 0.0243002, 0.5419054, 0.0999010, 0.0969010, 0.0400004, 0.0212002, 0.0148001, 0.0741507, 0.0350404, 0.0173602, 0.0157402]
    k3 = [1.32420, 1.12778, 1.22222, 1.11760, 9.88047E-01, 8.94344E-01, 7.87440E-01, 7.43687E-01, 8.11797E-01, 6.93279E-01, 5.09443E-01, 2.28721E-01]
    idxs3,ks3,zs3 = Clapeyron.rr_find_strongest(k3,z3)
    @test all(in.(idxs3,Ref((1,3,9,12))))

    =#
    function rr_vle_vapor_fraction_exact(K,z)
    n = length(z)
    if n == 2
        #0 = a1b -a0
        z1,z2 = z
        k1,k2 = K
        b1,b2 = 1/(1-k1),1/(1-k2)
        a1 = (z1*b2 + z2*b1)
        a0 = z1+z2
        return a1/a0
    elseif n == 3
        #0 = a0 + a1β + a2β^2
        z1,z2,z3 = z
        k1,k2,k3 = K
        b1,b2,b3 = 1/(1-k1),1/(1-k2),1/(1-k3)
        bmin = min(b1,b2,b3)
        bmax = max(b1,b2,b3)
        a2 =(z1 + z2 + z3)
        a1 = -b1*(z2 + z3) - b2*(z1 + z3) - b3*(z1 + z2)
        a0 = b1*b2*z3 + b1*b3*z2 + b2*b3*z1
        Δ = a1*a1 - 4*a0*a2
        inva2 = 1/(2*a2)
        Δ2 = sqrt(Δ)*inva2
        x = -a1*inva2
        β1 = x-Δ2
        β2 = x+Δ2
        #@show β1,β2
        #@show bmin,bmax
        if bmin < β1 < bmax
            return β1
        else
            return β2
        end
    elseif n == 4
        #0 = a0 + a1β + a2β^2 + a3β^3
        z1,z2,z3,z4 = z
        k1,k2,k3,k4 = K
        b1,b2,b3,b4 = 1/(1-k1),1/(1-k2),1/(1-k3),1/(1-k4)
        bmin = min(b1,b2,b3,b4)
        bmax = max(b1,b2,b3,b4)
        a3 =(z1 + z2 + z3 + z4)
        a2 = -b1*(z2 + z3 + z4) - b2*(z1 + z3 + z4) - b3*(z1 + z2 + z4)
        a1 = b1*b2*(z3+z4) + b1*b3*(z2+z4) + b2*b3*(z1+z4) + b1*b4*(z2+z3) + b2*b4*(z1+z3) + b3*b4*(z1+z2)
        a0 = -b1*b2*b3*z4 - b1*b2*b4*z3 - b1*b3*b4*z2 - b2*b3*b4*z1
        res1 =  Solvers.roots3(a0,a1,a2,a3)
        clx1,clx2,clx3 = res1
        r1,r2,r3 = real(clx1),real(clx2),real(clx3)
        rsum = r1+r2+r3
        rmax,rmin = extrema((r1,r2,r3))
        rmid = rsum - rmax - rmin
        return rmid
    else
        _0 = zero(first(K)+first(z))
        return _0/_0
    end
    end

function rr_flash_eval(K,z,β,normalize=true)
        _0 = zero(first(z)+first(K)+first(β))
        _1 = one(_0)
        if normalize
            sumz = sum(z)
            invsumz = _1/sumz
        else
            invsumz = _1
        end
        res = _0
        for i in 1:length(z)
            Kim1 = K[i] - _1
            res += invsumz*z[i]*Kim1/(1+β*Kim1)
        end
        #@show β,res

        return res
    end
     #porting those functions from my ex pkg.

"""
    rr_flash_vapor(k,z,β)
Returns the gas phase composition, given k-values `k`, the initial molar composition `z` and the molar vapor fraction ´β´.

Each gas phase composition is calculated acording to:

    xvᵢ = kᵢzᵢ/(1+ β(kᵢ-1))
"""
function rr_flash_vapor(k,z,β)
    function f(ki,zi)
    _1 = one(ki)
        return ki*zi/(_1+β*(ki-_1))
    end

    return map(f,k,z)
end

"""
    rr_flash_liquid(k,z,β)
Returns the liquid phase composition, given k-values `k`, the initial molar composition `z` and the molar vapor fraction ´β´.

Each gas phase composition is calculated acording to:

    xlᵢ = zᵢ/(1+ β(kᵢ-1))
"""
function rr_flash_liquid(k,z,β)
    function f(ki,zi)
        _1 = one(ki)
        return zi/(_1+β*(ki-_1))
    end
    return map(f,k,z)
end


text:
module MayanNumerals

export covert

convtable = Dict([
									(0,"\U1D2E0"),
									(1,"\U1D2E1"),
									(2,"\U1D2E2"),
									(3,"\U1D2E3"),
									(4,"\U1D2E4"),
									(5,"\U1D2E5"),
									(6,"\U1D2E6"),
									(7,"\U1D2E7"),
									(8,"\U1D2E8"),
									(9,"\U1D2E9"),
									(10,"\U1D2EA"),
									(11,"\U1D2EB"),
									(12,"\U1D2EC"),
									(13,"\U1D2ED"),
									(14,"\U1D2EE"),
									(15,"\U1D2EF"),
									(16,"\U1D2F0"),
									(17,"\U1D2F1"),
									(18,"\U1D2F2"),
									(19,"\U1D2F3")
									])

function convert(num::Integer; trailingstring::String=" ", blankzeros::Bool=false)
	if num < 0 || num > 19
		return "0-19 please"
	elseif blankzeros && num==0
		return " " * trailingstring
	else
		return get(convtable, num, "wrong for some other reason") * trailingstring
	end
end

function convert(num::AbstractFloat; trailingstring::String=" ", blankzeros::Bool=false)
	num = Int(round(num))
	if num < 0 || num > 19
		return "0-19 please"
	elseif blankzeros && num==0
		return " " * trailingstring
	else
		return get(convtable, num, "wrong for some other reason") * trailingstring
	end
end

end # module


text:
mutable struct BaseDelay{T <: Integer}
    shift_reg::T
    time_delay::Int64

    # Input
    stream::AbstractBitStream

    function BaseDelay{T}(stream::AbstractBitStream) where {T <: Unsigned}
        o = new()
        o.shift_reg = 0
        o.time_delay = 0
        o.stream = stream
        o
    end
end

function next!(base::BaseDelay{T}) where {T <: Integer}
    step!(base.stream)
    # Take stream's output shift into the delay position base of T
    base.shift_reg = base.shift_reg | (T(output(base.stream)) << base.time_delay)
end

function shift!(base::BaseDelay{T}) where {T <: Integer}
    # Shift register right
    base.shift_reg = base.shift_reg >> 1
end

# ----------------------------------------------------------------
# Abstract
# ----------------------------------------------------------------
function pre!(delay::AbstractDelay)
    next!(delay.base)
end

function post!(delay::AbstractDelay)
    shift!(delay.base)
end

# ----------------------------------------------------------------
# Tiny
# ----------------------------------------------------------------
# MSB = input bit
# LSB = output bit
#
#  /--- input from stream
# v
# 00000000
#        ^
#         \--- output
mutable struct TinyDelay <: AbstractDelay
    base::BaseDelay{UInt8}

    function TinyDelay(stream::AbstractBitStream, time_delay::Int64 = 0)
        @assert time_delay < 8 "TinyDelay time_delay must be < 8"
        o = new()
        o.base = BaseDelay{UInt8}(stream)
        o.base.time_delay = time_delay
        o
    end
end

function output(delay::TinyDelay)
    # Output is LSB
    delay.base.shift_reg & UInt8(1)
end

# ----------------------------------------------------------------
# Small
# ----------------------------------------------------------------
# Note: if the global time step is 100us then the maximum time_delay
# will be: 100us * 64 = 6400us = 6.4ms
mutable struct SmallDelay <: AbstractDelay
    base::BaseDelay{UInt64}

    function SmallDelay(stream::AbstractBitStream, time_delay::Int64 = 0)
        @assert time_delay < 64 "SmallDelay time_delay must be < 64"
        o = new()
        o.base = BaseDelay{UInt64}(stream)
        o.base.time_delay = time_delay
        o
    end
end

function output(delay::SmallDelay)
    # Output is LSB
    delay.base.shift_reg & UInt64(1)
end

# ----------------------------------------------------------------
# Medium
# ----------------------------------------------------------------
# Note: if the global time step is 100us then the maximum time_delay
# will be: 100us * 128 = 12800us = 12.8ms
mutable struct MediumDelay <: AbstractDelay
    base::BaseDelay{UInt128}

    function MediumDelay(stream::AbstractBitStream, time_delay::Int64 = 0)
        @assert time_delay < 128 "MediumDelay time_delay must be < 128"
        o = new()
        o.base = BaseDelay{UInt128}(stream)
        o.base.time_delay = time_delay
        o
    end
end

function output(delay::MediumDelay)
    # Output is LSB
    delay.base.shift_reg & UInt128(1)
end


text:
@testset "Path Normalization" begin
    mixed_path = ".\\\\path///to\\a\\place/..\\///"
    norm_path = "path/to/a"
    @test Zarr.normalize_path(mixed_path) == norm_path
    @test Zarr.DirectoryStore(mixed_path).folder == norm_path
    @test Zarr.normalize_path("/") == "/"
    @test Zarr.normalize_path("/a/") == "/a"
    @test Zarr.normalize_path("/path/to/a") == "/path/to/a"
end

"""
Function to test the interface of AbstractStore. Every complete implementation should pass this test.
"""
function test_store_common(ds)
  @test !Zarr.is_zgroup(ds)
  ds[".zgroup"]=rand(UInt8,50)
  @test Zarr.is_zgroup(ds)
  @test !Zarr.is_zarray(ds)

  @test Zarr.zname(ds)=="foo"
  @test isempty(Zarr.subdirs(ds))
  @test sort(collect(Zarr.keys(ds)))==[".zgroup"]

  #Create a subgroup
  snew = Zarr.newsub(ds,"bar")
  @test !Zarr.is_zarray(ds)
  snew[".zarray"] = rand(UInt8,50)
  @test Zarr.is_zgroup(ds)
  #Test getindex and setindex
  data = rand(UInt8,50)
  snew["0.0.0"] = data
  @test snew["0.0.0"]==data
  @test Zarr.storagesize(snew)==50
  @test Zarr.isinitialized(snew,"0.0.0")
  @test !Zarr.isinitialized(snew,"0.0.1")
  Zarr.writeattrs(snew,Dict("a"=>"b"))
  @test Zarr.getattrs(snew)==Dict("a"=>"b")
  snew2 = Zarr.getsub(ds,"bar")
  @test Zarr.getattrs(snew2)==Dict("a"=>"b")
  @test snew2["0.0.0"]==data
  delete!(snew2,"0.0.0")
  @test !Zarr.isinitialized(snew2,"0.0.0")
  snew["0.0.0"] = data
end

@testset "DirectoryStore" begin
  A = fill(1.0, 30, 20)
  chunks = (5,10)
  metadata = Zarr.Metadata(A, chunks; fill_value=-1.5)
  p = tempname()
  mkpath(joinpath(p,"foo"))
  ds = Zarr.DirectoryStore(joinpath(p,"foo"))
  test_store_common(ds)
  @test isdir(joinpath(p,"foo"))
  @test isfile(joinpath(p,"foo",".zgroup"))
  @test isdir(joinpath(p,"foo","bar"))
  @test isfile(joinpath(p,"foo","bar","0.0.0"))
  @test isfile(joinpath(p,"foo","bar",".zarray"))
  @test Zarr.path(ds)==joinpath(p,"foo")
end

@testset "DictStore" begin
  A = fill(1.0, 30, 20)
  chunks = (5,10)
  metadata = Zarr.Metadata(A, chunks; fill_value=-1.5)
  ds = Zarr.DictStore("foo")
  test_store_common(ds)
  @test ds.name == "foo"
  @test haskey(ds.a,".zgroup")
  @test ds.subdirs["bar"].a[".zattrs"]==UInt8[0x7b, 0x22, 0x61, 0x22, 0x3a, 0x22, 0x62, 0x22, 0x7d]
  @test sort(collect(keys(ds.subdirs["bar"].a)))==[".zarray", ".zattrs", "0.0.0"]
  @test isempty(keys(ds.subdirs["bar"].subdirs))
end

import AWSCore: aws_config

@testset "AWS S3 Storage" begin
  # These tests work locally but not on Travis, not idea why, will skip them for now
  # TODO fix
  #if get(ENV,"TRAVIS","") != "true"
    bucket = "zarr-demo"
    store = "store/foo"
    region = "eu-west-2"
    S3 = S3Store(bucket, store, aws = aws_config(creds=nothing, region = region))
    @test storagesize(S3) == 0
    @test Zarr.zname(S3) == "foo"
    @test Zarr.is_zgroup(S3) == true
    S3group = zopen(S3)
    @test Zarr.zname(S3group) == "foo"
    S3Array = S3group.groups["bar"].arrays["baz"]
    @test Zarr.zname(S3Array) == "baz"
    @test eltype(S3Array) == Zarr.ASCIIChar
    @test storagesize(S3Array) == 69
    @test String(S3Array[:]) == "Hello from the cloud!"
  #end
end

@testset "GCS S3 Storage" begin
  bucket = "cmip6"
  store = "ScenarioMIP/DKRZ/MPI-ESM1-2-HR/ssp370/r4i1p1f1/Amon/tasmax/gn"
  region = ""
  aws_google = aws_config(creds=nothing, region="", service_host="googleapis.com", service_name="storage")
  cmip6 = S3Store(bucket,store,aws = aws_google, listversion=1)
  @test storagesize(cmip6) == 7557
  @test Zarr.zname(cmip6) == "gn"
  g = zopen(cmip6)
  arr = g["tasmax"]
  @test size(arr) == (384,192,1032)
  @test eltype(arr) == Union{Missing, Float32}
  @test all(isapprox.(arr[1:2,1:2,2], [237.519 239.618; 237.536 239.667]))
  g2 = zopen("gs://cmip6/ScenarioMIP/DKRZ/MPI-ESM1-2-HR/ssp370/r4i1p1f1/Amon/tasmax/gn")
  arr = g["tasmax"]
  @test size(arr) == (384,192,1032)
  @test eltype(arr) == Union{Missing, Float32}
  @test all(isapprox.(arr[1:2,1:2,2], [237.519 239.618; 237.536 239.667]))
end

@testset "HTTP Storage" begin
  s = Zarr.DictStore()
  g = zgroup(s, attrs = Dict("groupatt"=>5))
  a = zcreate(Int,g,"a1",10,20,chunks=(5,5),attrs=Dict("arratt"=>2.5))
  a .= reshape(1:200,10,20)
  using Zarr.HTTP, Sockets
  server = Sockets.listen(0)
  ip,port = getsockname(server)
  @async HTTP.serve(g,ip,port,server=server)
  g2 = zopen("http://$ip:$port")
  @test g2.attrs == Dict("groupatt"=>5)
  @test g2["a1"].attrs == Dict("arratt"=>2.5)
  @test g2["a1"][:,:] == reshape(1:200,10,20)
  close(server)
end


text:
"""
    SubstitutionModel

Abstract type for substitution models,
using a continous time Markov model on a phylogeny.
Adapted from [SubstitutionModels.jl](https://github.com/BioJulia/SubstitutionModels.jl/)
in BioJulia.

For variable rates, see [`RateVariationAcrossSites`](@ref)

For sub types, see [`NucleicAcidSubstitutionModel`](@ref), [`TraitSubstitutionModel`](@ref)

All these models are supposed to have fields `rate` and `eigeninfo`.
"""
abstract type SubstitutionModel end #ideally, we'd like this to be SubstitutionModels.SubstitionModel
const SM = SubstitutionModel
const Qmatrix = StaticArrays.SMatrix{4, 4, Float64}
const Pmatrix = StaticArrays.MMatrix{4, 4, Float64}
const Bmatrix = StaticArrays.MMatrix{2, 2, Float64}

"""
    TraitSubstitutionModel

For subtypes, see [`BinaryTraitSubstitutionModel`](@ref),
[`EqualRatesSubstitutionModel`](@ref),
[`TwoBinaryTraitSubstitutionModel`](@ref)
"""
abstract type TraitSubstitutionModel{T} <: SubstitutionModel end #this accepts labels
const TSM = TraitSubstitutionModel{T} where T #T is type of labels

"""
    NucleicAcidSubstitutionModel

Adapted from [SubstitutionModels.jl](https://github.com/BioJulia/SubstitutionModels.jl/)
in BioJulia. The same [`Q`](@ref) and [`P`](@ref) function names are used for the
transition rates and probabilities.

For subtypes, see [`JC69`](@ref), [`HKY85`](@ref)
"""
abstract type NucleicAcidSubstitutionModel <: SubstitutionModel end
const NASM = NucleicAcidSubstitutionModel
function Base.show(io::IO, obj::SM)
    error("show not defined for $(typeof(obj)).")
end

"""
    nparams(model)

Number of parameters for a given trait evolution model
(length of field `model.rate`).
"""
nparams(obj::SM) = error("nparams not defined for $(typeof(obj)).")

"""
    setrates!(model, rates)

update rates then call [`seteigeninfo!`](@ref) to update a model's eigeninfo
"""
function setrates!(obj::SM, rates::AbstractVector)
    obj.rate[:] = rates
    seteigeninfo!(obj)
end

"""
    seteigeninfo!(obj)

Calculate eigenvalue & eigenfector information for a substitution model (SM) object
(as needed to calculate transition rate matrices) and store this info within the object.
"""
function seteigeninfo!(obj::SM)
    error("seteigeninfo! not finalized yet for generic substitution models")
    # eig_vals, eig_vecs = eigen(Q(obj)) #this assumes that every SM has an eigeninfo and has
    # obj.eigeninfo[:] = eig_vals, eig_vecs, inv(eig_vecs)
    # then make P! use seteigeninfo! by for a generic SM
end

"""
    getlabels(model)

State labels of a substitution model.
"""
function getlabels(obj::SM)
    error("Model must be of type TraitSubstitutionModel or NucleicAcidSubstitutionModel. Got $(typeof(obj))")
end

"""
    nstates(model)

Number of character states for a given evolution model.
"""
nstates(obj::SM) = error("nstates not defined for $(typeof(obj)).")

"""
For example, this is 4 for a `NucleicAcidSubstitutionModel`.

```jldoctest
julia> nstates(JC69([0.03], false))
4

julia> nstates(HKY85([.5], [0.25, 0.25, 0.25, 0.25]))
4
```
"""
function nstates(obj::NASM)
    return 4::Int
end

function getlabels(obj::TSM)
    return obj.label
end

"""
for a given [`NucleicAcidSubstitutionModel`](@ref), labels are symbols
from [BioSymbols](https://github.com/BioJulia/BioSymbols.jl).
For now, only ACGTs are allowed. (When fitting data, any ambiguity code
in the data would be treated as missing value).

# examples

```jldoctest
julia> getlabels(JC69([0.03], false))
4-element Array{BioSymbols.DNA,1}:
 DNA_A
 DNA_C
 DNA_G
 DNA_T

julia> getlabels(HKY85([.5], repeat([0.25], 4)))
4-element Array{BioSymbols.DNA,1}:
 DNA_A
 DNA_C
 DNA_G
 DNA_T

```
"""
function getlabels(::NASM)
    return [BioSymbols.DNA_A, BioSymbols.DNA_C, BioSymbols.DNA_G, BioSymbols.DNA_T]
end

"""
    Q(model)

Substitution rate matrix for a given substitution model:
Q[i,j] is the rate of transitioning from state i to state j.
"""
Q(obj::SM) = error("rate matrix Q not defined for $(typeof(obj)).")

"""
    showQ(IO, model)

Print the Q matrix to the screen, with trait states as labels on rows and columns.
adapted from prettyprint function by mcreel, found 2017/10 at
https://discourse.julialang.org/t/display-of-arrays-with-row-and-column-names/1961/6
"""
function showQ(io::IO, obj::SM)
    M = Q(obj)
    pad = max(8,maximum(length(getlabels(obj))+1))
    for i = 1:size(M,2) # print the header
        print(io, lpad(getlabels(obj)[i],(i==1 ? 2*pad : pad), " "))
    end
    print(io, "\n")
    for i = 1:size(M,1) # print one row per state
        if getlabels(obj) != ""
            print(io, lpad(getlabels(obj)[i],pad," "))
        end
        for j = 1:size(M,2)
            if j == i
                print(io, lpad("*",pad," "))
            else
                fmt = "%$(pad).4f"
                @eval(@printf($io,$fmt,$(M[i,j])))
            end
        end
        print(io, "\n")
    end
end

"""
    P(obj, t)

Probability transition matrix for a [`TraitSubstitutionModel`](@ref), of the form

    P[1,1] ... P[1,k]
       .          .
       .          .
    P[k,1] ... P[k,k]

where P[i,j] is the probability of ending in state j after time t,
given that the process started in state i.
see also: [`P!`](@ref).

HKY example:
```jldoctest
julia> m1 = HKY85([0.5], [0.20, 0.30, 0.30, 0.20])
HKY85 Substitution Model base frequencies: [0.2, 0.3, 0.3, 0.2]
relative rate version with transition/tranversion ratio kappa = 0.5,
 scaled so that there is one substitution per unit time
rate matrix Q:
               A       C       G       T
       A       *  0.4839  0.2419  0.3226
       C  0.3226       *  0.4839  0.1613
       G  0.1613  0.4839       *  0.3226
       T  0.3226  0.2419  0.4839       *

julia> PhyloNetworks.P(m1, 3.0)
4×4 StaticArrays.MArray{Tuple{4,4},Float64,2,16} with indices SOneTo(4)×SOneTo(4):
 0.217509  0.198417  0.190967  0.198417
 0.297625  0.312992  0.297625  0.28645
 0.28645   0.297625  0.312992  0.297625
 0.198417  0.190967  0.198417  0.217509
```

Juke-Cantor example:
```jldoctest
julia> m1 = JC69([1.]);

julia> PhyloNetworks.P(m1, 0.2)
4×4 StaticArrays.MArray{Tuple{4,4},Float64,2,16} with indices SOneTo(4)×SOneTo(4):
 0.824446   0.0585179  0.0585179  0.0585179
 0.0585179  0.824446   0.0585179  0.0585179
 0.0585179  0.0585179  0.824446   0.0585179
 0.0585179  0.0585179  0.0585179  0.824446
```
"""
@inline function P(obj::SM, t::Float64)
    t >= 0.0 || error("substitution model: >=0 branch lengths are needed")
    k = nstates(obj)
    Pmat = MMatrix{k,k}(repeat([0.0], k*k))
    return P!(Pmat, obj, t)
end

"""
    P!(Pmat::AbstractMatrix, obj::SM, t::Float64)

Fill in the input matrix `Pmat` with the transition rates
to go from each state to another in time `t`, according to rates in `Q`.
see also: [`P`](@ref).

```jldoctest
julia> m1 = BinaryTraitSubstitutionModel([1.0,2.0], ["low","high"])
Binary Trait Substitution Model:
rate low→high α=1.0
rate high→low β=2.0

julia> PhyloNetworks.P!(Matrix{Float64}(undef,2,2), m1, 3.0) # fills an uninitialized 2x2 matrix of floats
2×2 Array{Float64,2}:
 0.666708  0.333292
 0.666584  0.333416

julia> m2 = JC69([1.]);

julia> PhyloNetworks.P!(Matrix{Float64}(undef,4,4), m2, 0.2)
4×4 Array{Float64,2}:
 0.824446   0.0585179  0.0585179  0.0585179
 0.0585179  0.824446   0.0585179  0.0585179
 0.0585179  0.0585179  0.824446   0.0585179
 0.0585179  0.0585179  0.0585179  0.824446

```
"""
@inline function P!(Pmat::AbstractMatrix, obj::SM, t::Float64)
    Pmat[:] = exp(Q(obj) * t)
    return Pmat
end

"""
    BinaryTraitSubstitutionModel(α, β [, label])

Model for binary traits, that is, with 2 states. Default labels are "0" and "1".
α is the rate of transition from "0" to "1", and β from "1" to "0".
"""
struct BinaryTraitSubstitutionModel{T} <: TraitSubstitutionModel{T} # fixit: back to mutable struct?
    rate::Vector{Float64}
    label::Vector{T} # most often: T = String, but could be BioSymbols.DNA
    eigeninfo::Vector{Float64}
    function BinaryTraitSubstitutionModel{T}(rate::Vector{Float64}, label::Vector{T}, eigeninfo::Vector{Float64}) where T
        #Warning: this constructor should not be used directly. Use it with the constructors below.
        @assert length(rate) == 2 "binary state: need 2 rates"
        rate[1] >= 0. || error("parameter α must be non-negative")
        rate[2] >= 0. || error("parameter β must be non-negative")
        @assert length(label) == 2 "need 2 labels exactly"
        new(rate, label, eigeninfo)
    end
end
const BTSM = BinaryTraitSubstitutionModel{T} where T
function BinaryTraitSubstitutionModel(r::AbstractVector, label::AbstractVector)
    obj = BinaryTraitSubstitutionModel{eltype(label)}(r::AbstractVector, label::AbstractVector, zeros(3)) # Vector{Float64}(undef,3) for julia v1.0
    seteigeninfo!(obj)
    return obj
end
BinaryTraitSubstitutionModel(α::Float64, β::Float64, label::AbstractVector) = BinaryTraitSubstitutionModel([α,β], label)
BinaryTraitSubstitutionModel(α::Float64, β::Float64) = BinaryTraitSubstitutionModel(α, β, ["0", "1"])

"""
for a [`BinaryTraitSubstitutionModel`]: store eigenvalue (q_01+q_10) and stationary distribution
"""
function seteigeninfo!(obj::BTSM)
    ab = obj.rate[1] + obj.rate[2] #eigenvalue = -(a+b)
    ab > 0. || error("α+β must be positive")
    p0 = obj.rate[2]/ab # asymptotic frequency of state "0"
    p1 = obj.rate[1]/ab # asymptotic frequency of state "1"
    obj.eigeninfo[1] = ab
    obj.eigeninfo[2] = p0
    obj.eigeninfo[3] = p1
end

"""
for a `BinaryTraitSubstitutionModel`, this is 2:

```jldoctest
julia> m1 = BinaryTraitSubstitutionModel([1.0,2.0], ["low","high"])
Binary Trait Substitution Model:
rate low→high α=1.0
rate high→low β=2.0

julia> nstates(m1)
2
```
"""
function nstates(::BTSM)
    return 2::Int
end

nparams(::BTSM) = 2::Int

"""
For a BinaryTraitSubstitutionModel, the rate matrix Q is of the form:

    -α  α
     β -β
"""
@inline function Q(obj::BTSM)
    return Bmatrix(-obj.rate[1], obj.rate[2], obj.rate[1], -obj.rate[2])
end

function Base.show(io::IO, obj::BTSM)
    str = "Binary Trait Substitution Model:\n"
    str *= "rate $(obj.label[1])→$(obj.label[2]) α=$(round(obj.rate[1], digits=5))\n"
    str *= "rate $(obj.label[2])→$(obj.label[1]) β=$(round(obj.rate[2], digits=5))\n"
    print(io, str)
end

function P!(Pmat::AbstractMatrix, obj::BTSM, t::Float64)
    e1 = exp(-obj.eigeninfo[1]*t)
    a0= obj.eigeninfo[2]*e1
    a1= obj.eigeninfo[3]*e1
    Pmat[1,1] = obj.eigeninfo[2]+a1
    Pmat[2,1] = obj.eigeninfo[2]-a0
    Pmat[1,2] = obj.eigeninfo[3]-a1
    Pmat[2,2] = obj.eigeninfo[3]+a0
    return Pmat
end

"""
    TwoBinaryTraitSubstitutionModel(rate [, label])

[`TraitSubstitutionModel`](@ref) for two binary traits, possibly correlated.
Default labels are "x0", "x1" for trait 1, and "y0", "y1" for trait 2.
If provided, `label` should be a vector of size 4, listing labels for
trait 1 first then labels for trait 2.
`rate` should be a vector of substitution rates of size 8.
rate[1],...,rate[4] describe rates of changes in trait 1.
rate[5],...,rate[8] describe rates of changes in trait 2.
In the transition matrix, trait combinations are listed in the following order:
x0-y0, x0-y1, x1-y0, x1-y1.

# example

```julia
model = TwoBinaryTraitSubstitutionModel([2.0,1.2,1.1,2.2,1.0,3.1,2.0,1.1],
        ["carnivory", "noncarnivory", "wet", "dry"]);
model
using PhyloPlots
plot(model) # to visualize states and rates
```
"""
struct TwoBinaryTraitSubstitutionModel <: TraitSubstitutionModel{String}
    rate::Vector{Float64}
    label::Vector{String}
    function TwoBinaryTraitSubstitutionModel(α, label)
        all( x -> x >= 0., α) || error("rates must be non-negative")
        @assert length(α)==8 "need 8 rates"
        @assert length(label)==4 "need 4 labels for all combinations of 2 binary traits"
        new(α, [string(label[1], "-", label[3]), # warning: original type of 'label' lost here
                string(label[1], "-", label[4]),
                string(label[2], "-", label[3]),
                string(label[2], "-", label[4])])
    end
end
const TBTSM = TwoBinaryTraitSubstitutionModel
TwoBinaryTraitSubstitutionModel(α::Vector{Float64}) = TwoBinaryTraitSubstitutionModel(α, ["x0", "x1", "y0", "y1"])

nstates(::TBTSM) = 4::Int
nparams(::TBTSM) = 8::Int

function Q(obj::TBTSM)
    M = fill(0.0,(4,4))
    a = obj.rate
    M[1,3] = a[1]
    M[3,1] = a[2]
    M[2,4] = a[3]
    M[4,2] = a[4]
    M[1,2] = a[5]
    M[2,1] = a[6]
    M[3,4] = a[7]
    M[4,3] = a[8]
    M[1,1] = -M[1,2] - M[1,3]
    M[2,2] = -M[2,1] - M[2,4]
    M[3,3] = -M[3,4] - M[3,1]
    M[4,4] = -M[4,3] - M[4,2]
    return M
end

function Base.show(io::IO, obj::TBTSM)
    print(io, "Substitution model for 2 binary traits, with rate matrix:\n")
    showQ(io, obj)
end

"""
    EqualRatesSubstitutionModel(numberStates, α, labels)

[`TraitSubstitutionModel`](@ref) for traits with any number of states
and equal substitution rates α between all states.
Default labels are "1","2",...

# example

```jldoctest
julia> m1 = EqualRatesSubstitutionModel(2, [1.0], ["low","high"])
Equal Rates Substitution Model with k=2,
all rates equal to α=1.0.
rate matrix Q:
             low    high
     low       *  1.0000
    high  1.0000       *
```
"""
struct EqualRatesSubstitutionModel{T} <: TraitSubstitutionModel{T}
    k::Int
    rate::Vector{Float64}
    label::Vector{T}
    eigeninfo::Vector{Float64}
    function EqualRatesSubstitutionModel{T}(k::Int, rate::Vector{Float64},
        label::Vector{T}, eigeninfo::Vector{Float64}) where T
        k >= 2 || error("parameter k must be greater than or equal to 2")
        @assert length(rate)==1 "rate must be a vector of length 1"
        rate[1] > 0 || error("parameter α (rate) must be positive")
        @assert length(label)==k "label vector of incorrect length"
        new(k, rate, label, eigeninfo)
    end
end
const ERSM = EqualRatesSubstitutionModel{T} where T
function EqualRatesSubstitutionModel(k::Int, rate::Vector{Float64}, label::AbstractVector)
    obj = EqualRatesSubstitutionModel{eltype(label)}(k::Int, rate::Vector{Float64}, label::AbstractVector, zeros(1)) # Vector{Float64}(undef,3) for julia v1.0
    seteigeninfo!(obj)
    return obj
end
EqualRatesSubstitutionModel(k::Int, α::Float64, label::AbstractVector) = EqualRatesSubstitutionModel(k,[α],label)
EqualRatesSubstitutionModel(k::Int, α::Float64) = EqualRatesSubstitutionModel(k, [α], string.(1:k))

"""
for a [`EqualRatesSubstitutionModel`]: store lambda = k/(k-1), where k is the number of states
"""
function seteigeninfo!(obj::ERSM)
    q = (obj.k/(obj.k-1.0))
    obj.eigeninfo[1] = q
end

function nstates(obj::ERSM)
    return obj.k
end
nparams(::ERSM) = 1::Int

function Base.show(io::IO, obj::ERSM)
    str = "Equal Rates Substitution Model with k=$(obj.k),\n"
    str *= "all rates equal to α=$(round(obj.rate[1], digits=5)).\n"
    str *= "rate matrix Q:\n"
    print(io, str)
    showQ(io, obj)
end
function Q(obj::ERSM)
    #this might be wrong, doesnt match ERSM proof
    α = obj.rate[1]
    M = fill(α, (obj.k,obj.k))
    d = -(obj.k-1) * α
    for i in 1:obj.k
        M[i,i] = d
    end
    return M
end

"""
    randomTrait(model, t, start)
    randomTrait!(end, model, t, start)

Simulate traits along one edge of length t.
`start` must be a vector of integers, each representing the starting value of one trait.
The bang version (ending with !) uses the vector `end` to store the simulated values.

# Examples
```jldoctest
julia> m1 = BinaryTraitSubstitutionModel(1.0, 2.0)
Binary Trait Substitution Model:
rate 0→1 α=1.0
rate 1→0 β=2.0

julia> using Random; Random.seed!(12345);

julia> randomTrait(m1, 0.2, [1,2,1,2,2])
5-element Array{Int64,1}:
 1
 2
 1
 1
 2
```
"""
function randomTrait(obj::TSM, t::Float64, start::AbstractVector{Int})
    res = Vector{Int}(undef, length(start))
    randomTrait!(res, obj, t, start)
end

@doc (@doc randomTrait) randomTrait!
function randomTrait!(endTrait::AbstractVector{Int}, obj::TSM, t::Float64, start::AbstractVector{Int})
    Pt = P(obj, t)
    k = size(Pt, 1) # number of states
    w = [aweights(Pt[i,:]) for i in 1:k]
    for i in 1:length(start)
        endTrait[i] =sample(1:k, w[start[i]])
    end
    return endTrait
end

"""
    randomTrait(model, net; ntraits=1, keepInternal=true, checkPreorder=true)

Simulate evolution of discrete traits on a rooted evolutionary network based on
the supplied evolutionary model. Trait sampling is uniform at the root.

optional arguments:

- `ntraits`: number of traits to be simulated (default: 1 trait).
- `keepInternal`: if true, export character states at all nodes, including
  internal nodes. if false, export character states at tips only.

output:

- matrix of character states with one row per trait, one column per node;
  these states are *indices* in `model.label`, not the trait labels themselves.
- vector of node labels (for tips) or node numbers (for internal nodes)
  in the same order as columns in the character state matrix

# examples

```jldoctest
julia> m1 = BinaryTraitSubstitutionModel(1.0, 2.0, ["low","high"]);

julia> net = readTopology("(((A:4.0,(B:1.0)#H1:1.1::0.9):0.5,(C:0.6,#H1:1.0::0.1):1.0):3.0,D:5.0);");

julia> using Random; Random.seed!(1234);

julia> trait, lab = randomTrait(m1, net)
([1 2 … 1 1], ["-2", "D", "-3", "-6", "C", "-4", "H1", "B", "A"])

julia> trait
1×9 Array{Int64,2}:
 1  2  1  1  2  2  1  1  1

julia> lab
9-element Array{String,1}:
 "-2"
 "D"
 "-3"
 "-6"
 "C"
 "-4"
 "H1"
 "B"
 "A"
```
"""
function randomTrait(obj::TSM, net::HybridNetwork;
    ntraits=1::Int, keepInternal=true::Bool, checkPreorder=true::Bool)
    net.isRooted || error("net needs to be rooted for preorder recursion")
    if(checkPreorder)
        preorder!(net)
    end
    nnodes = net.numNodes
    M = Matrix{Int}(undef, ntraits, nnodes) # M[i,j]= trait i for node j
    randomTrait!(M,obj,net)
    if !keepInternal
        M = getTipSubmatrix(M, net, indexation=:cols) # subset columns only. rows=traits
        nodeLabels = [n.name for n in net.nodes_changed if n.leaf]
    else
        nodeLabels = [n.name == "" ? string(n.number) : n.name for n in net.nodes_changed]
    end
    return M, nodeLabels
end

function randomTrait!(M::Matrix, obj::TSM, net::HybridNetwork)
    recursionPreOrder!(net.nodes_changed, M, # updates M in place
            updateRootRandomTrait!,
            updateTreeRandomTrait!,
            updateHybridRandomTrait!,
            obj)
end

function updateRootRandomTrait!(V::AbstractArray, i::Int, obj)
    sample!(1:nstates(obj), view(V, :, i)) # uniform at the root
    return
end

function updateTreeRandomTrait!(V::Matrix,
    i::Int,parentIndex::Int,edge::Edge,
    obj)
    randomTrait!(view(V, :, i), obj, edge.length, view(V, :, parentIndex))
end

function updateHybridRandomTrait!(V::Matrix,
        i::Int, parentIndex1::Int, parentIndex2::Int,
        edge1::Edge, edge2::Edge, obj)
    randomTrait!(view(V, :, i), obj, edge1.length, view(V, :, parentIndex1))
    tmp = randomTrait(obj, edge2.length, view(V, :, parentIndex2))
    for j in 1:size(V,1) # loop over traits
        if V[j,i] == tmp[j] # both parents of the hybrid node have the same trait
            continue # skip the rest: go to next trait
        end
        if rand() > edge1.gamma
            V[j,i] = tmp[j] # switch to inherit trait of parent 2
        end
    end
end

"""
    JC69(rate, relative)

Jukes Cantor (1969) nucleic acid substitution model, which has a single rate parameter.
`rate` corresponds to the absolute diagonal elements, that is, the rate of change
(to any of the other 2 states). Individual rates are `rate`/3.
If `relative` is true (default), the transition matrix [`Q`](@ref) is normalized
to an average of 1 transition per unit of time: in which case `rate` is set to 1.0.

# examples

```jldoctest
julia> m1 = JC69([0.25], false)
Jukes and Cantor 69 Substitution Model,
absolute rate version
off-diagonal rates equal to 0.25/3.
rate matrix Q:
               A       C       G       T
       A       *  0.0833  0.0833  0.0833
       C  0.0833       *  0.0833  0.0833
       G  0.0833  0.0833       *  0.0833
       T  0.0833  0.0833  0.0833       *

julia> nstates(m1)
4

julia> nparams(m1)
1

julia> m2 = JC69([0.5])
Jukes and Cantor 69 Substitution Model,
relative rate version
off-diagonal rates equal to 1/3
rate matrix Q:
               A       C       G       T
       A       *  0.3333  0.3333  0.3333
       C  0.3333       *  0.3333  0.3333
       G  0.3333  0.3333       *  0.3333
       T  0.3333  0.3333  0.3333       *

julia> nparams(m2)
0
```
"""
struct JC69 <: NucleicAcidSubstitutionModel
    rate::Vector{Float64}
    relative::Bool
    eigeninfo::Vector{Float64} #TODO change to MVector, see if its faster
    function JC69(rate::Vector{Float64}, relative::Bool, eigeninfo::Vector{Float64})
        #Warning: this constructor should not be used directly. Use the constructor below
        #   which will call this.
        0 <= length(rate) <= 1 || error("rate not a valid length for a JC69 model")
        all(rate .> 0.0) || error("All elements of rate must be positive for a JC69 model")
        new(rate, relative, eigeninfo)
    end
end
function JC69(rate::AbstractVector, relative=true::Bool)
    obj = JC69(rate, relative, zeros(1))
    seteigeninfo!(obj)
    return obj
end
JC69(rate::Float64, relative=true::Bool) = JC69([rate], relative)

"""
for [`JC69`](@ref): store lambda = 4/3 (if relative) or rate * 4/3 (absolute).
"""
function seteigeninfo!(obj::JC69)
    if obj.relative
        lambda = (4.0/3.0) #eigen value of the Q matrix
    else
        lambda = (4.0/3.0)*obj.rate[1] #eigen value of the Q matrix
    end
    obj.eigeninfo[1] = lambda
end

function Base.show(io::IO, obj::JC69)
    str = "Jukes and Cantor 69 Substitution Model,\n"
    if obj.relative == true
        str *= "relative rate version\n"
        str *= "off-diagonal rates equal to 1/3\n"
    else
        str *= "absolute rate version\n"
        str *= "off-diagonal rates equal to $(round(obj.rate[1], digits=5))/3.\n"
    end
    str *= "rate matrix Q:\n"
    print(io, str)
    showQ(io, obj)
end

"""
    HKY85(rate, pi, relative)

A nucleic acid substitution model based on Hasegawa et al. 1985 substitution model.
`rate` should be a vector of 1 or 2 rates, and `pi` a vector of 4 probabilities summing to 1.

If `relative` is false, the 2 rates represent the transition rate and the transversion rate,
α and β. If `relative` is true (default), only the first rate is used and represents the
transition/transversion ratio: κ=α/β. The rate transition matrix Q is normalized to have
1 change / unit of time on average, i.e. the absolute version of Q is divided by
`2(piT*piC + piA*piG)α + 2(piY*piR)β`.

`nparams` returns 1 or 2.
In other words: the stationary distribution is not counted in the number of parameters
(and `fitdiscrete` does not optimize the pi values at the moment).

# examples

```jldoctest
julia> m1 = HKY85([.5], [0.20, 0.30, 0.30, 0.20])
HKY85 Substitution Model base frequencies: [0.2, 0.3, 0.3, 0.2]
relative rate version with transition/tranversion ratio kappa = 0.5,
 scaled so that there is one substitution per unit time
rate matrix Q:
               A       C       G       T
       A       *  0.4839  0.2419  0.3226
       C  0.3226       *  0.4839  0.1613
       G  0.1613  0.4839       *  0.3226
       T  0.3226  0.2419  0.4839       *

julia> nstates(m1)
4

julia> m2 = HKY85([0.5, 0.5], [0.20, 0.30, 0.30, 0.20], false)
HKY85 Substitution Model base frequencies: [0.2, 0.3, 0.3, 0.2]
absolute rate version with transition/transversion ratio kappa = a/b = 1.0
 with rates a = 0.5 and b = 0.5
rate matrix Q:
               A       C       G       T
       A       *  0.1500  0.1500  0.1000
       C  0.1000       *  0.1500  0.1000
       G  0.1000  0.1500       *  0.1000
       T  0.1000  0.1500  0.1500       *

```
"""
struct HKY85 <: NucleicAcidSubstitutionModel
    rate::Vector{Float64}
    pi::Vector{Float64}
    relative::Bool
    eigeninfo::Vector{Float64}
    function HKY85(rate::Vector{Float64}, pi::Vector{Float64}, relative::Bool, eigeninfo::Vector{Float64})
        #Warning: this constructor should not be used directly. Use the constructor below,
        #   which will call this.
        all(rate .> 0.) || error("All elements of rate must be positive")
        1 <= length(rate) <= 2 || error("rate has invalid length for HKY85 model")
        if relative length(rate) == 1 || error("the relative version of HKY85 takes a single rate")
        else        length(rate) == 2 || error("the absolute version of HKY85 takes 2 rates")
        end
        length(pi) == 4 || error("pi must be of length 4")
        all(0. .< pi.< 1.) || error("All base proportions must be between 0 and 1")
        sum(pi) == 1. || error("Base proportions must sum to 1")
        new(rate, pi, relative, eigeninfo)
    end
end
function HKY85(rate::AbstractVector, pi::Vector{Float64}, relative=true::Bool)
    obj = HKY85(rate, pi, relative, zeros(5))
    seteigeninfo!(obj)
    return obj
end
HKY85(rate::Float64, pi::Vector{Float64}) = HKY85([rate], pi, true)

"""
for [`HKY85`](@ref): store piR, piY, the 2 non-zero eigenvalues and a scaling factor
"""
function seteigeninfo!(obj::HKY85)
    piA = obj.pi[1]; piC = obj.pi[2]; piG = obj.pi[3]; piT = obj.pi[4]
    piR = piA + piG
    piY = piT + piC
    obj.eigeninfo[4] = piR
    obj.eigeninfo[5] = piY

    a = obj.rate[1]
    if obj.relative
        lambda = 2*a*(piT*piC + piA*piG) + 2*(piY*piR) #for sub/year interpretation
        b = 1.0/lambda #scaling factor so that branch lengths can be interpreted as substitutions
        obj.eigeninfo[2] = -((piR * a*b) + (piY * b))
        obj.eigeninfo[3] = -((piY * a*b) + (piR * b))
    else
        b = obj.rate[2]
        obj.eigeninfo[2] = -((piR * a) + (piY * b))
        obj.eigeninfo[3] = -((piY * a) + (piR * b))
    end
    obj.eigeninfo[1] = -b
end

function Base.show(io::IO, obj::HKY85)
    str = "HKY85 Substitution Model base frequencies: $(obj.pi)\n"
    if obj.relative
        str *= "relative rate version with transition/tranversion ratio kappa = $(round(obj.rate[1], digits=5)),"
        str *= "\n scaled so that there is one substitution per unit time\n"
    else
        str *= "absolute rate version with transition/transversion ratio kappa = a/b = "
        str *= "$(round(obj.rate[1]/obj.rate[2], digits=5))"
        str *= "\n with rates a = $(round(obj.rate[1], digits=5)) and b = $(round(obj.rate[2], digits=5))\n"
    end
    str *= "rate matrix Q:\n"
    print(io, str)
    showQ(io, obj)
end

"""
for `JC69` model: 0 if relative, 1 if absolute
"""
function nparams(obj::JC69)
    return (obj.relative ? 0 : 1)
end

"""
for `HKY85` model: 1 if relative, 2 if absolute
"""
function nparams(obj::HKY85)
    return (obj.relative ? 1 : 2)
end

@inline function Q(obj::JC69)
    if obj.relative
        Q0 = 1.0/3.0
        #we multiply by 1/3 to make branch lengths interpretable as avg(#substitutions) (pi-weighted avg of diagonal is -1)
    else
        Q0 = (1.0/3.0)*obj.rate[1]
    end

    return Qmatrix(-3*Q0, Q0, Q0, Q0,
                    Q0, -3*Q0, Q0, Q0,
                    Q0, Q0, -3*Q0, Q0,
                    Q0, Q0, Q0, -3*Q0)
end

@inline function Q(obj::HKY85)
    piA = obj.pi[1]; piC = obj.pi[2]; piG = obj.pi[3]; piT = obj.pi[4]
    piR = piA + piG
    piY = piT + piC
    if obj.relative
        k = obj.rate[1]
        lambda = (2*k*(piT*piC + piA*piG) + 2*(piY*piR)) # for sub/time interpretation. see HKY docstring
        Q₁  = piA/lambda
        Q₂  = k * Q₁
        Q₃  = piC/lambda
        Q₄  = k * Q₃
        Q₆  = piG/lambda
        Q₅  = k * Q₆
        Q₇  = piT/lambda
        Q₈  = k * Q₇
        Q₉  = -(Q₃ + Q₅ + Q₇)
        Q₁₀ = -(Q₁ + Q₆ + Q₈)
        Q₁₁ = -(Q₂ + Q₃ + Q₇)
        Q₁₂ = -(Q₁ + Q₄ + Q₆)
    else #GT, AT, CG, and AC are less likely (transversions)
        a = obj.rate[1]; b = obj.rate[2]
        Q₁  = b * piA
        Q₂  = a * piA
        Q₃  = b * piC
        Q₄  = a * piC
        Q₅  = a * piG
        Q₆  = b * piG
        Q₇  = b * piT
        Q₈  = a * piT
        Q₉  = -(Q₃ + Q₅ + Q₇) #AA
        Q₁₀ = -(Q₁ + Q₆ + Q₈) #CC
        Q₁₁ = -(Q₂ + Q₃ + Q₇) #GG
        Q₁₂ = -(Q₁ + Q₄ + Q₆) #TT
    end
    return Qmatrix(Q₉,  Q₁,  Q₂,  Q₁, #ACGT
                    Q₃,  Q₁₀, Q₃,  Q₄,
                    Q₅,  Q₆,  Q₁₁, Q₆,
                    Q₇,  Q₈,  Q₇,  Q₁₂)
end

function P!(Pmat::AbstractMatrix, obj::JC69, t::Float64)
    P0 = 0.25 + 0.75 * exp(-t * obj.eigeninfo[1]) #lambda
    P1 = 0.25 - 0.25 * exp(-t * obj.eigeninfo[1])
    Pmat[:,:] .= P1 # puts P1 on the off-diagonals (in particular)
    for i in 1:4 Pmat[i,i]=P0; end # diagonal
    return Pmat
end

function P!(Pmat::AbstractMatrix, obj::HKY85, t::Float64)
    #returns ACGT
    piA = obj.pi[1]; piC = obj.pi[2]; piG = obj.pi[3]; piT = obj.pi[4]
    piR = obj.eigeninfo[4]
    piY = obj.eigeninfo[5]

    e2 = exp(obj.eigeninfo[1] * t)
    e3 = exp(obj.eigeninfo[2] * t)
    e4 = exp(obj.eigeninfo[3] * t)

    P_AA = piA + ((piA * piY)/piR) * e2 + (piG/piR) * e3
    P_CC = piC + ((piC * piR)/piY) * e2 + (piT/piY) * e4
    P_GG = piG + ((piG * piY)/piR) * e2 + (piA/piR) * e3
    P_TT = piT + ((piT * piR)/piY) * e2 + (piC/piY) * e4

    P_TA_CA = piA * (1 - e2) #double name because TA and CA are equal
    P_GA = piA + (piA*piY/piR)*e2 - (piA/piR)*e3 #P6

    P_AC_GC = piC * (1 - e2) #double name because AC and GC are equal #p7
    P_TC = piC + ((piC * piR)/piY) * e2 - (piC/piY) * e4

    P_AG = piG + ((piG * piY)/piR) * e2 - (piG/piR) * e3 #p9
    P_TG_CG = piG * (1 - e2) #double name because TG and CG are equal #p10

    P_AT_GT = piT * (1 - e2) #double name because AT and GT are equal #P11
    P_CT = piT + ((piT * piR)/piY) * e2 - (piT/piY) * e4

    Pmat[1,1] = P_AA
    Pmat[2,2] = P_CC
    Pmat[3,3] = P_GG
    Pmat[4,4] = P_TT

    Pmat[1,2] = P_TA_CA
    Pmat[1,3] = P_GA
    Pmat[1,4] = P_TA_CA

    Pmat[2,1] = P_AC_GC
    Pmat[2,3] = P_AC_GC
    Pmat[2,4] = P_TC

    Pmat[3,1] = P_AG
    Pmat[3,2] = P_TG_CG
    Pmat[3,4] = P_TG_CG

    Pmat[4,1] = P_AT_GT
    Pmat[4,2] = P_CT
    Pmat[4,3] = P_AT_GT
    return Pmat
end

"""
    RateVariationAcrossSites(α=1.0, ncat=4)

Model for variable substitution rates across sites (or across traits) using
the discrete Gamma model (Yang 1994, Journal of Molecular Evolution).
Turn any SM model to SM + Gamma.
Using this model increases the number of parameters by one.

Because the mean of the desired continuous Gamma distribution is 1, we use
shape α and scale θ=1/α (e.g. rate β=α).
The shape parameter is referred to as alpha here.
The Gamma distribution is discretized into `ncat` categories.
In each category, the category's rate multiplier is a normalized quantile of the gamma distribution.

```jldoctest
julia> rv = RateVariationAcrossSites()
Rate Variation Across Sites using Discretized Gamma Model
alpha: 1.0
categories for Gamma discretization: 4
ratemultiplier: [0.14578, 0.51313, 1.07083, 2.27025]

julia> PhyloNetworks.setalpha!(rv, 2.0)

julia> rv
Rate Variation Across Sites using Discretized Gamma Model
alpha: 2.0
categories for Gamma discretization: 4
ratemultiplier: [0.31907, 0.68336, 1.10898, 1.8886]

julia> RateVariationAcrossSites(2.0, 4)
Rate Variation Across Sites using Discretized Gamma Model
alpha: 2.0
categories for Gamma discretization: 4
ratemultiplier: [0.31907, 0.68336, 1.10898, 1.8886]
```
"""
mutable struct RateVariationAcrossSites
    alpha::Float64
    ncat::Int
    ratemultiplier::Array{Float64}
    function RateVariationAcrossSites(alpha = 1.0::Float64, ncat = 4::Int)
        @assert alpha >= 0.0 "alpha must be >= 0"
        @assert ncat > 0 "ncat must be 1 or greater"
        if ncat == 1
            ratemultiplier = [1.0]
        else
            cuts = Vector((0:(ncat-1))/ncat) + repeat([1/(2ncat)], ncat)
            ratemultiplier = quantile.(Distributions.Gamma(alpha, 1/alpha), cuts)
            ratemultiplier ./= mean(ratemultiplier)
        end
        new(alpha, ncat, ratemultiplier)
    end
end
const RVAS = RateVariationAcrossSites

"""
    setalpha!(obj, alpha)

Set the shape parameter `alpha` in a RateVariationAcrossSites model `obj`,
and update the rate multipliers accordingly.
"""
function setalpha!(obj::RateVariationAcrossSites, alpha::Float64)
    @assert alpha >= 0 "alpha must be a float >= 0"
    obj.alpha = alpha
    if obj.ncat > 1
        rv = obj.ratemultiplier
        cuts = Vector((0:(obj.ncat-1))/obj.ncat) + repeat([1/(2*obj.ncat)], obj.ncat)
        rv[:] = quantile.(Distributions.Gamma(alpha, 1/alpha), cuts)
        rv ./= mean(rv)
    end
    return nothing
end

function Base.show(io::IO, obj::RateVariationAcrossSites)
    str = "Rate Variation Across Sites using Discretized Gamma Model\n"
    str *= "alpha: $(round(obj.alpha, digits=5))\n"
    str *= "categories for Gamma discretization: $(obj.ncat)\n"
    str *= "ratemultiplier: $(round.(obj.ratemultiplier, digits=5))\n"
    print(io, str)
end

function nparams(obj::RateVariationAcrossSites)
    return (obj.ncat == 1 ? 0 : 1)
end

"""
    empiricalDNAfrequencies(DNAdata::AbstractDataFrame, DNAweights,
                            correction=true, useambiguous=true)

Estimate base frequencies in DNA data `DNAdata`, ordered ACGT.

- `DNAdata`: data frame. All columns are used. If the first column
  gives species names, find a way to ignore it before calculating empirical
  frequencies, e.g. `empiricalDNAfrequencies(view(DNAdata, :, 2:ncol(DNAdata)))`.
  Data type must be `BioSymbols.DNA` or `Char` or `String`.
  WARNING: this is checked on the first column only.
- `DNAweights`: vector of weights, to weigh each column in `DNAdata`.
- `correction`: if `true`, add 1 to each count and 4 to the denominator
  for a more stable estimator, similar to Bayes prior of 1/4 and
  the Agresti-Coull interval in binomial estimation.
- `useambiguous`: if `true`, ambiguous bases are used (except gaps and Ns).
  For example, `Y` adds 0.5 weight to `C` and 0.5 weight to `T`.
"""
function empiricalDNAfrequencies(dnaDat::AbstractDataFrame, dnaWeights::Vector,
    correctedestimate=true::Bool, useambiguous=true::Bool)

    # warning: checking first column and first row only
    eltypes(dnaDat)[1] == BioSymbols.DNA || eltypes(dnaDat)[1] == Char ||
      dnaDat[1,1] ∈ string.(BioSymbols.alphabet(DNA)) ||
        error("empiricalDNAfrequencies requires data of type String, Char, or BioSymbols.DNA")

    # initialize counts: keys same as BioSymbols.ACGT (which are ordered)
    prior = correctedestimate ? 1.0 : 0.0
    dnacounts = Dict(DNA_A=>prior, DNA_C=>prior, DNA_G=>prior, DNA_T=>prior)

    convert2dna = eltypes(dnaDat)[1] != BioSymbols.DNA
    for j in 1:ncol(dnaDat) # for each column
        col = dnaDat[!,j]
        wt = dnaWeights[j]
        for nuc in col      # for each row
            if convert2dna
                nuc = convert(DNA, nuc[1]) # if nuc is a string, nuc[1] = 1st character
            end
            if nuc ∈ BioSymbols.ACGT
                dnacounts[nuc] += wt
            elseif nuc == DNA_Gap || nuc == DNA_N || !useambiguous
                continue # to next row
            # else: ambiguity, uses BioSequences definitions
            elseif nuc == DNA_M # A or C
                dnacounts[DNA_A] += 0.5*wt
                dnacounts[DNA_C] += 0.5*wt
            elseif nuc == DNA_R # A or G
                dnacounts[DNA_A] += 0.5*wt
                dnacounts[DNA_G] += 0.5*wt
            elseif nuc == DNA_W # A or T/U
                dnacounts[DNA_A] += 0.5*wt
                dnacounts[DNA_T] += 0.5*wt
            elseif nuc == DNA_S # C or G
                dnacounts[DNA_C] += 0.5*wt
                dnacounts[DNA_G] += 0.5*wt
            elseif nuc == DNA_Y # C or T/U
                dnacounts[DNA_C] += 0.5*wt
                dnacounts[DNA_T] += 0.5*wt
            elseif nuc == DNA_K # G or T/U
                dnacounts[DNA_G] += 0.5*wt
                dnacounts[DNA_T] += 0.5*wt
            elseif nuc == DNA_V # A or C or G
                dnacounts[DNA_A] += wt/3
                dnacounts[DNA_C] += wt/3
                dnacounts[DNA_G] += wt/3
            elseif nuc == DNA_H # A or C or T
                dnacounts[DNA_A] += wt/3
                dnacounts[DNA_C] += wt/3
                dnacounts[DNA_T] += wt/3
            elseif nuc == DNA_D # A or G or T/U
                dnacounts[DNA_A] += wt/3
                dnacounts[DNA_G] += wt/3
                dnacounts[DNA_T] += wt/3
            elseif nuc == DNA_B # C or G or T/U
                dnacounts[DNA_C] += wt/3
                dnacounts[DNA_G] += wt/3
                dnacounts[DNA_T] += wt/3
            end
        end
    end
    totalweight = sum(values(dnacounts))
    res = [dnacounts[key]/totalweight for key in BioSymbols.ACGT] # to control the order
    all(0. .<= res .<= 1.) || error("""weird: empirical base frequency < 0 or > 1""")
    return res
end

"""
    stationary(substitutionmodel)

Stationary distribution of a Markov model
"""
stationary(mod::SM) = error("stationary not defined for $(typeof(mod)).")

function stationary(mod::JC69)
    return [0.25,0.25,0.25,0.25]
end

function stationary(mod::HKY85)
    return mod.pi
end

function stationary(mod::ERSM) #for ERSM, uniform = stationary
    return [1/mod.k for i in 1:mod.k]
end

function stationary(mod::BTSM)
    return [mod.eigeninfo[2], mod.eigeninfo[3]]
end


text:
module CDLagrangianFESpacesTests

using Test
using Gridap.Geometry
using Gridap.ReferenceFEs
using Gridap.FESpaces
using Gridap.CellData
using Gridap.TensorValues
using Gridap.Fields

domain =(0,1,0,1)
partition = (3,3)
model = CartesianDiscreteModel(domain,partition)
trian = get_triangulation(model)

# order = 1
orders = (2,1)
order = 2*max(orders...)

T = VectorValue{2,Float64}
reffe = LagrangianRefFE(T,QUAD,orders)
V = FESpace(model=model,reffe=reffe,conformity=CDConformity((DISC,CONT)))#),dirichlet_tags = [1,6])
test_single_field_fe_space(V)

u(x) = x

U = TrialFESpace(V,u)
uh = interpolate(u,U)
e = u - uh
trian = Triangulation(model)
quad = CellQuadrature(trian,order)
el2 = sqrt(sum(integrate(inner(e,e),trian,quad)))
@test el2 < 1.0e-10

reffe = LagrangianRefFE(T,QUAD,2)
V = FESpace(model=model,reffe=reffe,conformity=CDConformity((CONT,DISC)))
U = TrialFESpace(V,u)
uh = interpolate(u,U)
e = u - uh
el2 = sqrt(sum(integrate(inner(e,e),trian,quad)))
@test el2 < 1.0e-10

reffe = LagrangianRefFE(T,QUAD,(2,1))
V = FESpace(model=model,reffe=reffe,conformity=CDConformity((DISC,CONT)))
U = TrialFESpace(V,u)
uh = interpolate(u,U)
e = u - uh
el2 = sqrt(sum(integrate(inner(e,e),trian,quad)))
@test el2 < 1.0e-10

reffe = LagrangianRefFE(T,QUAD,(2,0))
V = FESpace(model=model,reffe=reffe,conformity=CDConformity((CONT,DISC)))
u(x) = VectorValue(x[1],0.0)
U = TrialFESpace(V,u)
uh = interpolate(u,U)
e = u - uh
el2 = sqrt(sum(integrate(inner(e,e),trian,quad)))
@test el2 < 1.0e-10

reffe = LagrangianRefFE(T,QUAD,(2,0))
V = FESpace(model=model,reffe=reffe,conformity=CDConformity((DISC,DISC)))
u(x) = VectorValue(x[1],0.0)
U = TrialFESpace(V,u)
uh = interpolate(u,U)
e = u - uh
el2 = sqrt(sum(integrate(inner(e,e),trian,quad)))
@test el2 < 1.0e-10

# # Skeleton triangulation
# face_own_dofs = get_face_own_dofs(reffe)
# strian = SkeletonTriangulation(model,reffe,face_own_dofs)
# ns = get_normal_vector(strian)
# writevtk(strian,"strian",cellfields=["normal"=>ns])
#
# # Random function for visualization purposes
# model = CartesianDiscreteModel((0,1,0,1),(10,5))
# V = FESpace(model=model,reffe=reffe,conformity=true)
# trian = Triangulation(model)
# vh = FEFunction(V,rand(num_free_dofs(V)))
# writevtk(trian,"trian",nsubcells=20,cellfields=["vh"=>vh])

#using Gridap.Visualization
#
#writevtk(trian,"trian",nsubcells=10,cellfields=["uh"=>uh])

end # module


text:
function logpriors(rparams::RiskParameters{T}, rpriors::RiskPriors{T}) where T <: EpidemicModel
  lprior = 0.0
  for i in 1:length(rpriors.sparks)
    lprior == -Inf && break
    lprior += logpdf(rpriors.sparks[i], rparams.sparks[i])
  end
  for i in 1:length(rpriors.susceptibility)
    lprior == -Inf && break
    lprior += logpdf(rpriors.susceptibility[i], rparams.susceptibility[i])
  end
  for i in 1:length(rpriors.infectivity)
    lprior == -Inf && break
    lprior += logpdf(rpriors.infectivity[i], rparams.infectivity[i])
  end
  for i in 1:length(rpriors.transmissibility)
    lprior == -Inf && break
    lprior += logpdf(rpriors.transmissibility[i], rparams.transmissibility[i])
  end
  if T in [SEIR; SEI]
    for i in 1:length(rpriors.latency)
      lprior == -Inf && break
      lprior += logpdf(rpriors.latency[i], rparams.latency[i])
    end
  end
  if T in [SEIR; SIR]
    for i in 1:length(rpriors.removal)
      lprior == -Inf && break
      lprior += logpdf(rpriors.removal[i], rparams.removal[i])
    end
  end
  return lprior
end


text:
const Min_Max = NamedTuple{(:min, :max),Tuple{Float64,Float64}}
const From_To_Float = NamedTuple{(:from, :to),Tuple{Float64,Float64}}
const From_To_Bus =  NamedTuple{(:from, :to),Tuple{Bus,Bus}}
const FromTo_ToFrom_Float = NamedTuple{(:from_to, :to_from),Tuple{Float64,Float64}}

struct Line <: Branch
    name::String
    available::Bool
    connectionpoints::From_To_Bus
    r::Float64 #[pu]
    x::Float64 #[pu]
    b::From_To_Float # [pu]
    rate::Float64
    anglelimits::Min_Max #Degrees
end

function Line(name::String,
              available::Bool,
              connectionpoints::From_To_Bus,
              r::Float64,
              x::Float64,
              b::From_To_Float,
              rate::Float64,
              anglelimits::Float64)
    return Line(name, available, connectionpoints, r, x, b, rate,
                (min=-anglelimits, max=anglelimits))
end

function Line(; name="init",
              available=false,
              connectionpoints=From_To_Bus((from=Bus(), to=Bus())),
              r=0.0,
              x=0.0,
              b=(from=0.0, to=0.0),
              rate=0.0,
              anglelimits=(min=-1.57, max=1.57))
    return Line(name, available, connectionpoints, r, x, b, rate, anglelimits)
end

struct MonitoredLine <: Branch
    name::String
    available::Bool
    connectionpoints::From_To_Bus
    r::Float64 #[pu]
    x::Float64 #[pu]
    b::From_To_Float # [pu]
    flowlimits::FromTo_ToFrom_Float #MW
    rate::Float64
    anglelimits::Min_Max #Degrees
end

function MonitoredLine(; name="init",
                       available=false,
                       connectionpoints=From_To_Bus((from=Bus(), to=Bus())),
                       r=0.0,
                       x=0.0,
                       b=(from=0.0, to=0.0),
                       rate=0.0,
                       anglelimits=(min=-90.0, max=90.0))
    return Line(name, available, connectionpoints, r, x, b, rate, anglelimits)
end


text:
using Images
using FileIO

# load our mandel functions
include("mandel.jl")

# set parameters
h = 2500
w = 5000
iter = 10000

# time measurements
print("starting...\n")
tStart=time()

m_set = mandelbrot_set(-2.5,1.5,-1.25,1.25,w,h,iter)

tStop = time()

# write the image-file
img = colorview(RGB, m_set)
save("mandel.jpg", img)

print("done. took ", tStop-tStart, " seconds\n");

max_stars_repo_name: rabarar/juliaFractals

max_issues_repo_name: rabarar/juliaFractals

max_forks_repo_name: rabarar/juliaFractals


text:
# test.jl

include("./infrastructure/infrastructure.jl")

using DataFrames
# see: https://github.com/JuliaDebug/Debugger.jl
#using Debugger # use REPL for debugging

# activating logging
# see: https://discourse.julialang.org/t/how-to-save-logging-output-to-a-log-file/14004/5
#using Logging
#io = open("log_invoicing.txt", "w+")
#logger = SimpleLogger(io)
#global_logger(logger)

@info("$(now()) - Test program started.")

const PATH_CSV = "./bank.csv"

# get orders
orders = AppliSales.process()
@info("$(now()) - Orders received.");
invnbr = 1000
invoices = [create(order, "A" * string(global invnbr += 1)) for order in orders]
@info("$(now()) - Unpaid invoices created.")

# first invoice
invoice_1 = invoices[1]

# meta data
@info("$(now())", meta(invoice_1))

# header
@info("$(now())", header(invoice_1))

# body
@info("$(now())", body(invoice_1))

# id
@info("$(now())", id(invoice_1))

# process orders
journal_entries_1 = process(orders)
@info("$(now()) - Journal entries for unpaid invoices created.")

# get Bank statements and the unpaid invoices
stms = read_bank_statements(PATH_CSV)
@info("$(now()) - Get bankstatement entries for testing.")

unpaid_invoices = retrieve_unpaid_invoices()
@info("$(now()) - Unpaid invoices retrieved from database.")

# process the unpaid invoices and bank statements
journal_entries_2 = process(unpaid_invoices, stms)
@info("$(now()) - Journal entries for paid invoices created.")

paid_invoices = retrieve_paid_invoices()
@info("$(now()) - Paid invoices retrieved from database.")
#flush(io) # needed when logging is enabled
#@info("Test program finished")

# print an invoice example
#include("./print_invoice.jl")
#invoices = retrieve_unpaid_invoices(PATH_DB_TEST)
#basic(invoices[2])

# cleanup
cmd = `rm ./test_invoicing.sqlite`
run(cmd)

@info("$(now()) - databaseT removed")


text:
using RecordStores
using Test

@testset "RecordStores.jl" begin
    records = [
        Dict(:a => 1, :b => [1, 2, 3]),
        Dict(:c => [1.0, 2.0, 3im]),
    ]

    @testset for archiver in [:zip, :dir]
        mktempdir() do dir
            store = recordstore(joinpath(dir, "store"); archiver = archiver)
            open(store, "w") do w
                for r in records
                    write(w, r)
                end
            end

            @test read(store) == records
            @test open(collect, store) == records
        end
    end
end

@testset "guess_archiver" begin
    if Sys.isunix()
        @test recordstore("/store.zip").writer == RecordStores.ZipFile.Writer
        @test recordstore("/store/").writer == RecordStores.DirWriter
    end
end


text:
using Pkg

cd("Julia")

Pkg.activate(".")
Pkg.instantiate()

using Revise

include("01_wrangling.jl")
include("../src/tweet_helpers_refac_dateticks.jl")
include("../src/tweet_graphing.jl")

using Chain
using CairoMakie
using Colors
using AlgebraOfGraphics
using DataFrames
using DataFramesMeta
using Dates

set_aog_theme!()

normalize(vec) = (vec .- minimum(vec)) / (maximum(vec) - minimum(vec))

function cumulative_sorting(df,column;for_comparison::Bool = false, scandal::String)
    nonprofit = ["ComicRelief","WWP1","WWP2","Oxfam","RedCrossAU","MIT1","MIT2","MIT3"]

    df = @chain df begin
        @subset(:true_type .!= "retweeted")
        groupby(:author_id)
        combine(column => sum => :tot)
        @orderby(-:tot)
        @subset(:tot .> 0)
        @transform(_,
        :cumulative = cumsum(:tot),
        :rownumber = rownumber.(eachrow(_)))
        @select :rownumber :cumulative
        @rtransform :type = scandal ∈ nonprofit ? "Non Profit" : "Profit"
        @rtransform :scandal = scandal
        rename(:rownumber => :x, :cumulative => :y)
    end

    if for_comparison

        df = @chain df begin
            @transform :y = normalize(:y)
            @transform :x = normalize(:x)
        end

    end

    return df
end

scandals = string.(read_names("Data"));
dfs = [unique(df,[:tweet_id,:author_id]) for df in read_files("Data")]

get_date(x) = Date(x)

function cumulative_comparison_plot(dfs, scandals, what;for_comparison=false)
    cumulatives = [
        cumulative_sorting(df, what;
            for_comparison=for_comparison, scandal)
        for (df, scandal) in zip(dfs, scandals)]
    cumulatives = vcat(cumulatives...)

    xtitle = for_comparison ? "Proportion" : "Number"

    cumulatives_plot = data(cumulatives) *
                       mapping(:x => xtitle * " of authors considered",
                               :y => "Cumulative " * string(what),
                               group=:scandal,
                               color=:type) *
                       visual(Lines)

    draw(cumulatives_plot)

    return (cumulatives_plot)
end

colors = cgrad(:thermal, 19, categorical=true)

resolution = (800, 600)
fig = Figure(; resolution);

draw!(fig[1,1],
    cumulative_comparison_plot(dfs,scandals,:retweet_count);
    axis=(yscale=log10,
          xscale=log10)
    #palettes=(color=colors,)
)

draw!(fig[1,2],
    cumulative_comparison_plot(dfs,scandals,:reply_count);
    axis=(yscale=log10,
          xscale=log10)
    #palettes=(color=colors,)
)

draw!(fig[2,1],
    cumulative_comparison_plot(dfs,scandals,:quote_count);
    axis=(yscale=log10,
          xscale=log10)
    #palettes=(color=colors,)
)

mkpath("PlotsDate/comparative_plots/")
save("PlotsDate/comparative_plots/comparative_cumulatives_raw_numbers.png", fig, px_per_unit = 3)

resolution = (800, 600)
fig = Figure(; resolution);

draw!(fig[1,1],
    cumulative_comparison_plot(dfs,scandals,:retweet_count;for_comparison = true);
    #palettes=(color=colors,)
)

draw!(fig[1,2],
    cumulative_comparison_plot(dfs,scandals,:reply_count;for_comparison = true);
    #palettes=(color=colors,)
)

draw!(fig[2,1],
    cumulative_comparison_plot(dfs,scandals,:quote_count;for_comparison = true);
    #palettes=(color=colors,)
)

fig

save("PlotsDate/comparative_plots/comparative_cumulatives_normalized.png", fig, px_per_unit = 3)


text:
loadbutton = filepicker()
columnbuttons = Observable{Any}(dom"div"())
data = Observable{Any}(Array{Sequence})
map!(f->begin
        print("Loading... $f\n")
        JLD2.load(FileIO.File{FileIO.DataFormat{:JLD2}}(f),"seq")
    end
    , data, loadbutton)

hh, ww = 450, 150
l = PlotlyJS.Layout(;title="Acquired signal", yaxis_title="Signal (a.u.)",
    xaxis_title="samples",
    xaxis=attr(range=[length(signal)/2-250,length(signal)/2+250],
			    rangeslider=attr(visible=true),
                gridcolor="gray"),
    yaxis=attr(gridcolor="gray"),
    plot_bgcolor="black",
    modebar=attr(orientation="v"),legend=false,height=400)
absS = PlotlyJS.scatter(y=abs.(signal),name="|S(t)|")
reS = PlotlyJS.scatter(y=real.(signal),name="Re{S(t)}")
imS = PlotlyJS.scatter(y=imag.(signal),name="Im{S(t)}")
p = PlotlyJS.plot([absS,reS,imS],l)
plt = Observable{Any}(p)

function makebuttons(df)
    namesdf = 1:length(df)
    buttons = button.(string.())
    for (btn, name) in zip(buttons, namesdf)
        map!(t -> begin
            p = PlotlyJS.plot(df[name][:])
            PlotlyJS.relayout!(p,height=hh)
            p
            end
            , plt, btn)
    end
    dom"div"(hbox(buttons))
end

map!(makebuttons, columnbuttons, data)
pulseseq = dom"div"(loadbutton, columnbuttons, plt)
content!(w, "div#content", pulseseq)


text:
"""
    abstract type AbstractLegendreNorm end

Abstract trait supertype for normalization conditions of the associated Legendre functions.
"""
abstract type AbstractLegendreNorm end

# Define element types as a way to conveniently promote, even for the trait types.
# Union{} will always lose out in type promotion.
Base.eltype(::Type{<:AbstractLegendreNorm}) = Union{}
# Normalizations are scalars in a broadcasting context
Base.broadcastable(x::AbstractLegendreNorm) = Ref(x)

## Below are overloadable non-exported functions to implement the AbstractLegendreNorm
## interface.

"""
    initcond(::N, ::Type{T}) where {N<:AbstractLegendreNorm, T}

Returns the initial condition ``P_0^0(x)`` for the associated Legendre recursions based
on the normalization choice `N` for numeric type `T`.
"""
function initcond end

"""
    coeff_μ(norm::N, ::Type{T}, l::Integer) where {N<:AbstractLegendreNorm, T}

Returns the coefficient ``\\mu_\\ell`` for the single-term recursion relation
```math
    P_\\ell^\\ell(x) = -\\mu_\\ell \\sqrt{1-x^2} P_{\\ell-1}^{\\ell-1}(x)
```
where ``\\mu_\\ell`` is appropriate for the choice of normalization `N`.
"""
function coeff_μ end

"""
    coeff_ν(norm::N, ::Type{T}, l::Integer) where {N<:AbstractLegendreNorm, T}

Returns the coefficient ``\\nu_\\ell`` for the single-term recursion relation
```math
    P_\\ell^{\\ell-1}(x) = \\nu_\\ell x P_{\\ell-1}^{\\ell-1}(x)
```
where ``\\nu_\\ell`` is appropriate for the choice of normalization `N`.
"""
function coeff_ν end

"""
    coeff_α(norm::N, ::Type{T}, l::Integer, m::Integer) where {N<:AbstractLegendreNorm, T}

Returns the coefficient ``\\alpha_\\ell^m`` for the two-term recursion relation
```math
    P_\\ell^m(x) = \\alpha_\\ell^m x P_{\\ell-1}^m(x) - \\beta_\\ell^m P_{\\ell-2}^m(x)
```
where ``\\alpha_\\ell^m`` is appropriate for the choice of normalization `N`.
"""
function coeff_α end

"""
    coeff_β(norm::N, ::Type{T}, l::Integer, m::Integer) where {N<:AbstractLegendreNorm, T}

Returns the coefficient ``\\beta_\\ell^m`` for the two-term recursion relation
```math
    P_\\ell^m(x) = \\alpha_\\ell^m x P_{\\ell-1}^m(x) - \\beta_\\ell^m P_{\\ell-2}^m(x)
```
where ``\\beta_\\ell^m`` is appropriate for the choice of normalization `N`.
"""
function coeff_β end

"""
    boundscheck_hook(norm::AbstractLegendreNorm, lmax, mmax)

A bounds-checking hook executed at the beginning of each [`legendre!`](@ref) call to
permit a normalization `norm` to validate that the given maximum ``(\\ell,m)`` will be
within the ability to satisfy. The default case always returns `nothing`. A custom
normalization should throw an error if `lmax` or `mmax` is out of bounds or return `nothing`
otherwise.

For example, the precomputed coefficients of [`LegendreNormCoeff`](@ref) are limited to
a given domain at time of construction and cannot be used to calculate terms to arbitrary
orders/degrees.
"""
function boundscheck_hook end

boundscheck_hook(norm::AbstractLegendreNorm, lmax, mmax) = nothing


text:
using Metal

@show devices()
dev = MtlDevice(1)

vecC=rand(Float32, 128)

bufferSize = 128
bufferA = MtlBuffer(Float32, dev, bufferSize, storage=Shared)
bufferB = MtlBuffer(Float32, dev, bufferSize, storage=Shared)
bufferC = MtlBuffer(Float32, dev, bufferSize, storage=Shared)

vecA = unsafe_wrap(Vector{Float32}, bufferA, (bufferSize,))
vecB = unsafe_wrap(Vector{Float32}, bufferB, (bufferSize,))
vecC = unsafe_wrap(Vector{Float32}, bufferC, (bufferSize,))

using Random
rand!(vecA)

queue = global_queue(dev)

vecB .= 0
cmdBuffer = commit!(queue) do buffer
    MTL.MtlBlitCommandEncoder(buffer) do enc
        MTL.append_copy!(enc, bufferA, 1, bufferC, 1, 128*4)
    end
end

Base.unsafe_copyto!(dev, bufferA, 1, bufferB, 1, 128)

vecD = rand(Float32, 128)
ptrD = pointer(vecD)
Base.unsafe_copyto!(dev, ptrD, bufferB, 0, 128)

###
arr2 = MTL.MtlArray{Float32,1}(undef, (bufferSize,), storage=MTL.MtResourceStorageModeShared)
arrptr = MTL.content(arr2.buffer)
arrvec = unsafe_wrap(Vector{Float32}, arrptr, bufferSize)

Base.unsafe_copyto!(dev, arr2.buffer, 1, bufferB, 1, 128)

MTL.waitUntilCompleted(cmdBuffer)


text:
module ExclusionProcesses

# See http://www.stochasticlifestyle.com/finalizing-julia-package-documentation-testing-coverage-publishing/
export EPState, SingleSpinState, MultiSpinState, spin_configs, PeriodicBC, OpenBC, SEP,
        update!, Simulation

#Import your packages

# import Base: func1 #Any function you add dispatches to need to be imported directly

include("epstate.jl")
include("boundarycondition.jl")
include("exclusionprocess.jl")
include("sep.jl")
include("runsim.jl")

#export coolfunc, coolfunc2 #Export the functions you want users to use

using .epstate
using .boundarycondition
using .exclusionprocess
using .sep
using .runsim

end

max_stars_repo_name: AustenLamacraft/ExclusionProcesses

max_issues_repo_name: AustenLamacraft/ExclusionProcesses

max_forks_repo_name: AustenLamacraft/ExclusionProcesses


text:
using PkgSRA
using Plots
using Random
#Random, Distributions, Interpolations
using StatsBase, StatsPlots
pyplot()
Random.seed!(11)
using BSON: @save
include("../scripts/paper_settings.jl");
include("../../utils/sindy_turing_utils.jl")
# include("../../utils/sindy_statistics_utils.jl")
# include("../../utils/main_algorithm_utils.jl")

#####
##### Define the controlled dynamical system
#####
# Load example problem
include(EXAMPLE_FOLDERNAME*"example_spiking_neuron.jl")
U_starts = 5.0
    amplitude = 40.0;
function U_func_current(t)
    ctr = zeros(2,1)
    ctr[1] = (t>U_starts ? amplitude : 0.0)
    return ctr
end

# Solve the system
dyn_with_ctr = solve_neuron_system(U_func_time=U_func_current)
plot(dyn_with_ctr)
dat = Array(dyn_with_ctr)

numerical_grad = numerical_derivative(dat, ts)
grad_true = dyn_with_ctr(ts, Val{1})

# Get just the control signals
U_true = hcat(U_func_current.(ts)...)

## Add a spiking control signal
num_ctr = 50;
    U_starts2 = rand(1, num_ctr) .* tspan[2]
    U_widths2 = [0.05];
    amplitudes = [200.0]
my_U_func_time2(t) = U_func_time_multivariate(t, zeros(2,1),
                            U_widths2, U_starts2,
                            F_dim=[1],
                            amplitudes=amplitudes)
U_func_total(t) = my_U_func_time2(t) + U_func_current(t)

dyn_with_spikes = solve_neuron_system(U_func_time=U_func_total)
plot(dyn_with_spikes)
dat2 = Array(dyn_with_spikes)

grad_true2 = dyn_with_spikes(ts, Val{1})
numerical_grad2 = numerical_derivative(dat2, ts)

U_true2 = hcat(U_func_total.(ts)...)

#####
##### First panels: Remove the control signal and fit a controlled model
#####
sindy_library = Dict("cross_terms"=>2, "constant"=>nothing);
chain_unctr, best_sindy_unctr = calc_distribution_of_models(
    dat, numerical_grad, sindy_library,
    # val_list = calc_permutations(3,2)
    # val_list = combinations(1:3, 2)
    val_list = Iterators.product(1:3,1:3) # DEBUG
)

plot(chain_unctr)
## Get the posterior distribution
(residual_unctr, sample_gradients, sample_noise, dat_grad) =
        calc_distribution_of_residuals(
        dat, numerical_grad, chain_unctr, 1:length(ts), best_sindy_unctr)
ctr_guess = process_residual(residual_unctr, sample_noise)

ind = 1:5000
i = 1
    plot(residual_unctr[i,ind], ribbon=500*sample_noise)
    plot!(U_true[:,ind]', label="True")
    title!("Residual and true control for channel $i")
i = 2
    plot(sample_gradients[i,ind], label="model", lw=2)
    plot!(numerical_grad[i,ind], label="data")
    title!("Data and model for channel $i")

start_of_subsample = 1001
accepted_ind = subsample_using_residual(
            residual_unctr[:,start_of_subsample:end], sample_noise,
            noise_factor=500,
            min_length=10, shorten_length=2)
num_pts = 2000; start_ind = 101
subsample_ind = accepted_ind[start_ind:num_pts+start_ind-1] .+ start_of_subsample

U_sub = U_true[:,subsample_ind]
plot(U_sub')
    # plot!(residual[:,subsample_ind]', ribbon=sample_trajectory_noise)
    title!("Control signals in the subsampled dataset")
i = 2
    plot(subsample_ind,sample_gradients[i,subsample_ind], label="model", lw=2)
    plot!(subsample_ind,numerical_grad[i,subsample_ind], label="data")
    title!("Data and model for channel $i")

# val_list = calc_permutations(5,2)
# val_list = combinations(1:5, 2)
val_list = Iterators.product(1:3,1:3) # DEBUG
(final_sindy_model,best_criterion,all_criteria,_, _) =
     sindyc_ensemble(dat[:,subsample_ind],
                     numerical_grad[:, subsample_ind],
                     sindy_library, val_list,
                     selection_criterion=my_aicc,
                     sparsification_mode="num_terms",
                     selection_dist=Normal(0,10),
                     use_clustering_minimization=true)
print_equations(final_sindy_model)
scatter(sum.(val_list), all_criteria)

chain_ctr, best_sindy_ctr = calc_distribution_of_models(
    dat[:,subsample_ind],
    numerical_grad[:,subsample_ind],
    sindy_library,
    # val_list = calc_permutations(5,2),
    # val_list = combinations(1:5, 2),
    val_list = Iterators.product(1:3,1:3), # DEBUG
    chain_opt = (iterations=200, train_ind=1:num_pts)
)

print_equations(best_sindy_ctr)
sindy_sample =  sindy_from_chain(best_sindy_ctr, chain_ctr)
print_equations(sindy_sample)
density(chain_ctr)

# Calculate updated control signal
(residual_ctr, _, noise_ctr, _) =
        calc_distribution_of_residuals(
                dat, grad_true, chain_ctr,
                1:length(ts), best_sindy_ctr)
ctr_final = process_residual(residual_ctr, 100*noise_ctr)

plot(ctr_final', label="Learned", lw=3, color=:black)
    plot!(U_true', label="True")

#####
##### Second panels: nonlinearity as well as control
#####

num_pts = 200;
sindy_library = Dict("cross_terms"=>2, "constant"=>nothing);
best_initial_subsample = calc_best_random_subsample(
    dat2, numerical_grad2, sindy_library;
    num_pts=num_pts,
    num_subsamples=10,
    # val_list=calc_permutations(5,2)
    # val_list = combinations(1:5, 2)
    val_list = Iterators.product(1:3,1:3) # DEBUG
)[1]

# Now start the regularly scheduled Turing
chain_unctr2, best_sindy_unctr2 = calc_distribution_of_models(
    dat2[:,best_initial_subsample],
    numerical_grad2[:,best_initial_subsample], sindy_library,
    # val_list = calc_permutations(5,2),
    # val_list = combinations(1:5, 2),
    val_list = Iterators.product(1:3,1:3), # DEBUG
    chain_opt = (iterations=200, train_ind=1:num_pts)
)

plot(chain_unctr2[:noise])
## Get the posterior distribution
(residual_unctr2, sample_gradients2, sample_noise2, dat_grad2) =
        calc_distribution_of_residuals(
        dat2, grad_true2, chain_unctr2, 1:length(ts), best_sindy_unctr2)
# ctr_guess2 = process_residual(residual_unctr2, sample_noise2)

ind = 3200:4200
i = 1
    plot(residual_unctr2[i,ind], ribbon=1*sample_noise2)
    plot!(U_true2[:,ind]', label="True")
    ylims!(-400, 400)
    title!("Residual and true control for channel $i")
i = 1
    plot(sample_gradients2[i,ind], label="model", lw=2)
    plot!(numerical_grad2[i,ind], label="data")
    title!("Data and model for channel $i")

start_of_subsample = 1001
accepted_ind = subsample_using_residual(
            residual_unctr2[:,start_of_subsample:end], sample_noise2,
            noise_factor=1,
            min_length=10, shorten_length=2)
num_pts = 2000; start_ind = 101
subsample_ind = accepted_ind[start_ind:num_pts+start_ind-1] .+ start_of_subsample

U_sub2 = U_true2[:,subsample_ind]
plot(U_sub2')
    # plot!(residual[:,subsample_ind]', ribbon=sample_trajectory_noise)
    title!("Control signals in the subsampled dataset")
i = 2
    plot(subsample_ind,sample_gradients2[i,subsample_ind], label="model", lw=2)
    plot!(subsample_ind,numerical_grad2[i,subsample_ind], label="data")
    title!("Data and model for channel $i")

# val_list = calc_permutations(5,2)
# val_list = combinations(1:5, 2)
val_list = Iterators.product(1:3,1:3) # DEBUG
(final_sindy_model2,best_criterion2,all_criteria2,_, _) =
     sindyc_ensemble(dat2[:,subsample_ind],
                     numerical_grad2[:, subsample_ind],
                     sindy_library, val_list,
                     selection_criterion=my_aicc,
                     sparsification_mode="num_terms",
                     selection_dist=Normal(0,10),
                     use_clustering_minimization=true)
println("Model")
print_equations(final_sindy_model2)
println("True equations")
print_equations(core_dyn_true)
scatter(sum.(val_list), all_criteria)

chain_ctr2, best_sindy_ctr2 = calc_distribution_of_models(
    dat2[:,subsample_ind],
    numerical_grad2[:,subsample_ind],
    sindy_library,
    # val_list = calc_permutations(4,2),
    # val_list = combinations(1:4, 2),
    val_list = Iterators.product(1:3,1:3) # DEBUG
    chain_opt = (iterations=200, train_ind=1:num_pts)
)

print_equations(best_sindy_ctr2)
# sindy_sample =  sindy_from_chain(best_sindy_ctr, chain_ctr)
plot(chain_ctr2)

# Calculate updated control signal
(residual_ctr2, _, noise_ctr2, _) =
        calc_distribution_of_residuals(
                dat2, grad_true2, chain_ctr2,
                1:length(ts), best_sindy_ctr2)
ctr_final2 = process_residual(residual_ctr2, 1*noise_ctr2)

plot(ctr_final2', label="Learned", lw=3, color=:black)
    plot!(U_true2', label="True")

#####
##### Save
#####
this_dat_name = DAT_FOLDERNAME*"dat_neuron_"

## Basic spiking model
# Raw data
fname = this_dat_name*"raw.bson";
@save fname dat grad_true numerical_grad dyn_with_ctr U_true chain_unctr

# Controlled model
fname = this_dat_name*"controlled_model.bson";
@save fname chain_ctr ctr_final #best_sindy_ctr

## Spikes with varying input
# Raw data
fname = this_dat_name*"raw2.bson";
@save fname dat2 grad_true2 numerical_grad2 dyn_with_spikes U_true2 chain_unctr2

# Controlled model
fname = this_dat_name*"controlled_model2.bson";
@save fname chain_ctr2 ctr_final2 #best_sindy_ctr2


text:
__precompile__(true)
module Magnus

using LinearMaps

abstract type Exponentiator end
abstract type KrylovExponentiator <: Exponentiator end

include("scalar.jl")
# include("pade.jl")
include("sub_exp.jl")
include("lanczos.jl")
include("propagators.jl")
include("midpoint.jl")
include("cfet.jl")
include("integrate.jl")
include("convergence.jl")

end # module


text:
struct Layers
    layers::Vector{Layer}
end

Base.convert(::Type{Layers}, s::Layer) = Layers([s])
Base.convert(::Type{Layers}, s::Layers) = s

Base.getindex(v::Layers, i::Int) = v.layers[i]
Base.length(v::Layers) = length(v.layers)
Base.eltype(::Type{Layers}) = Layer
Base.iterate(v::Layers, args...) = iterate(v.layers, args...)

const OneOrMoreLayers = Union{Layers, Layer}

function Base.:+(s1::OneOrMoreLayers, s2::OneOrMoreLayers)
    l1::Layers, l2::Layers = s1, s2
    return Layers(vcat(l1.layers, l2.layers))
end

function Base.:*(s1::OneOrMoreLayers, s2::OneOrMoreLayers)
    l1::Layers, l2::Layers = s1, s2
    return Layers([el1 * el2 for el1 in l1 for el2 in l2])
end

uniquevalues(v::ArrayLike) = collect(uniquesorted(vec(v)))

function uniquevalues(e::Entry)
    uv = Dict{KeyType, Any}()
    for (key, val) in pairs(e.primary)
        uv[key] = uniquevalues(val)
    end
    for (key, val) in pairs(e.positional)
        all(iscontinuous, val) && continue
        all(isgeometry, val) && continue
        uv[key] = mapreduce(uniquevalues, mergesorted, val)
    end
    return uv
end

function extremas(e::Entry)
    es = Dict{KeyType, Any}()
    for tup in (e.positional, e.named)
        for (key, val) in pairs(tup)
            all(iscontinuous, val) || continue
            es[key] = mapreduce(Makie.extrema_nan, extend_extrema, val)
        end
    end
    return es
end

function compute_grid_positions(scales, primary=(;))
    return map((:row, :col), (first, last)) do sym, f
        scale = get(scales, sym, nothing)
        lscale = get(scales, :layout, nothing)
        return if !isnothing(scale)
            rg = Base.OneTo(maximum(scale.plot))
            haskey(primary, sym) ? rescale(fill(primary[sym]), scale) : rg
        elseif !isnothing(lscale)
            rg = Base.OneTo(maximum(f, lscale.plot))
            haskey(primary, :layout) ? map(f, rescale(fill(primary[:layout]), lscale)) : rg
        else
            Base.OneTo(1)
        end
    end
end

function compute_axes_grid(fig, s::OneOrMoreLayers;
                           axis=NamedTuple(), palettes=NamedTuple())
    layers::Layers = s
    entries = map(process, layers)

    uv = mapreduce(uniquevalues, mergewith!(mergesorted), entries)
    es = mapreduce(extremas, mergewith!(extend_extrema), entries)

    theme_palettes = NamedTuple(Makie.current_default_theme()[:palette])
    palettes = merge((layout=wrap,), map(to_value, theme_palettes), palettes)
    scales = default_scales(merge(uv, es), palettes)

    axs = compute_grid_positions(scales)

    axes_grid = map(CartesianIndices(axs)) do c
        type = get(axis, :type, Axis)
        options = Base.structdiff(axis, (; type))
        ax = type(fig[Tuple(c)...]; options...)
        return AxisEntries(ax, Entry[], scales, Dict{KeyType, Any}())
    end

    for e in entries
        for c in CartesianIndices(shape(e))
            primary, positional, named = map((e.primary, e.positional, e.named)) do tup
                return map(v -> getnewindex(v, c), tup)
            end
            rows, cols = compute_grid_positions(scales, primary)
            entry = Entry(e; primary, positional, named)
            labels = copy(e.labels)
            map!(l -> getnewindex(l, c), values(labels))
            for i in rows, j in cols
                ae = axes_grid[i, j]
                push!(ae.entries, entry)
                merge!(assert_equal, ae.labels, labels)
            end
        end
    end

    # Link colors
    labeledcolorbar = getlabeledcolorbar(axes_grid)
    if !isnothing(labeledcolorbar)
        _, colorbar = labeledcolorbar
        colorrange = colorbar.extrema
        for entry in AlgebraOfGraphics.entries(axes_grid)
            entry.attributes[:colorrange] = colorrange
        end
    end

    # Axis labels and ticks
    for ae in axes_grid
        # TODO: support log colorscale
        ndims = isaxis2d(ae) ? 2 : 3
        for (i, var) in zip(1:ndims, (:x, :y, :z))
            label, scale = get(ae.labels, i, nothing), get(ae.scales, i, nothing)
            any(isnothing, (label, scale)) && continue
            for (k′, v) in pairs((label=string(label), ticks=ticks(scale)))
                k = Symbol(var, k′)
                k in keys(axis) || (getproperty(Axis(ae), k)[] = v)
            end
        end
    end

    return axes_grid

end

function Makie.plot!(fig, s::OneOrMoreLayers;
                     axis=NamedTuple(), palettes=NamedTuple())
    grid = compute_axes_grid(fig, s; axis, palettes)
    foreach(plot!, grid)
    return grid
end

function Makie.plot(s::OneOrMoreLayers;
                    axis=NamedTuple(), figure=NamedTuple(), palettes=NamedTuple())
    fig = Figure(; figure...)
    grid = plot!(fig, s; axis, palettes)
    return FigureGrid(fig, grid)
end

# Convenience function, which may become superfluous if `plot` also calls `facet!`
function draw(s::OneOrMoreLayers;
              axis = NamedTuple(), figure=NamedTuple(), palettes=NamedTuple())
    fg = plot(s; axis, figure, palettes)
    facet!(fg)
    Legend(fg)
    resizetocontent!(fg)
    return fg
end

function draw!(fig, s::OneOrMoreLayers;
               axis=NamedTuple(), palettes=NamedTuple())
    ag = plot!(fig, s; axis, palettes)
    facet!(fig, ag)
    return ag
end


text:
# [test/runtests.jl]
using Bioceananigans

# Test scripts
include("NNP.jl")

max_stars_repo_name: iuryt/biocenanigans

max_issues_repo_name: iuryt/biocenanigans