text:
export alias_elimination, flatten

using SymbolicUtils: Rewriters

function fixpoint_sub(x, dict)
    y = substitute(x, dict)
    while !isequal(x, y)
        y = x
        x = substitute(y, dict)
    end

    return x
end

function substitute_aliases(diffeqs, dict)
    lhss(diffeqs) .~ fixpoint_sub.(rhss(diffeqs), (dict,))
end

# Note that we reduce parameters, too
# i.e. `2param = 3` will be reduced away
isvar(s) = s isa Sym ? true :
           istree(s) ? isvar(operation(s)) :
                       false

function get_α_x(αx)
    if isvar(αx)
        return 1, αx
    elseif istree(αx) && operation(αx) === (*)
        args = arguments(αx)
        nums = []
        syms = []
        for arg in args
            isvar(arg) ? push!(syms, arg) : push!(nums, arg)
        end

        if length(syms) == 1
            return prod(nums), syms[1]
        end
    else
        return nothing
    end
end

function is_univariate_expr(ex, iv)
    count = 0
    for var in vars(ex)
        if !isequal(iv, var) && !isparameter(var)
            count += 1
            count > 1 && return false
        end
    end
    return count <= 1
end

function is_sub_candidate(ex, iv, conservative)
    conservative || return true
    isvar(ex) || ex isa Number || is_univariate_expr(ex, iv)
end

function maybe_alias(lhs, rhs, diff_vars, iv, conservative)
    is_sub_candidate(rhs, iv, conservative) || return false, nothing

    res_left = get_α_x(lhs)
    if res_left !== nothing && !(res_left[2] in diff_vars)
        α, x = res_left
        sub = x => _isone(α) ? rhs : rhs / α
        return true, sub
    else
        return false, nothing
    end
end

function alias_elimination(sys; conservative=true)
    sys = flatten(sys)
    iv = independent_variable(sys)
    eqs = equations(sys)
    diff_vars = filter(!isnothing, map(eqs) do eq
            if isdiffeq(eq)
                arguments(eq.lhs)[1]
            else
                nothing
            end
        end) |> Set

    deps = Set()
    subs = Pair[]
    neweqs = Equation[]; sizehint!(neweqs, length(eqs))

    for (i, eq) in enumerate(eqs)
        # only substitute when the variable is algebraic
        if isdiffeq(eq)
            push!(neweqs, eq)
            continue
        end

        # `α x = rhs` => `x = rhs / α`
        ma, sub = maybe_alias(eq.lhs, eq.rhs, diff_vars, iv, conservative)
        if !ma
            # `lhs = β y` => `y = lhs / β`
            ma, sub = maybe_alias(eq.rhs, eq.lhs, diff_vars, iv, conservative)
        end

        isalias = false
        if ma
            l, r = sub
            # alias equations shouldn't introduce cycles
            if !(l in deps) && isempty(intersect(deps, vars(r)))
                push!(deps, l)
                push!(subs, sub)
                isalias = true
            end
        end

        if !isalias
            neweq = _iszero(eq.lhs) ? eq : 0 ~ eq.rhs - eq.lhs
            push!(neweqs, neweq)
        end
    end

    alias_vars = first.(subs)
    sts = states(sys)
    fullsts = vcat(map(eq->eq.lhs, observed(sys)), sts, parameters(sys))
    alias_eqs = topsort_equations(alias_vars .~ last.(subs), fullsts)
    newstates = setdiff(sts, alias_vars)

    @set! sys.eqs = substitute_aliases(neweqs, Dict(subs))
    @set! sys.states = newstates
    @set! sys.observed = [observed(sys); alias_eqs]
    return initialize_system_structure(sys)
end

"""
$(SIGNATURES)

Use Kahn's algorithm to topologically sort observed equations.

Example:
```julia
julia> @variables t x(t) y(t) z(t) k(t)
(t, x(t), y(t), z(t), k(t))

julia> eqs = [
           x ~ y + z
           z ~ 2
           y ~ 2z + k
       ];

julia> ModelingToolkit.topsort_equations(eqs, [x, y, z, k])
3-element Vector{Equation}:
 Equation(z(t), 2)
 Equation(y(t), k(t) + 2z(t))
 Equation(x(t), y(t) + z(t))
```
"""
function topsort_equations(eqs, states; check=true)
    graph, assigns = observed2graph(eqs, states)
    neqs = length(eqs)
    degrees = zeros(Int, neqs)

    for 𝑠eq in 1:length(eqs); var = assigns[𝑠eq]
        for 𝑑eq in 𝑑neighbors(graph, var)
            # 𝑠eq => 𝑑eq
            degrees[𝑑eq] += 1
        end
    end

    q = Queue{Int}(neqs)
    for (i, d) in enumerate(degrees)
        d == 0 && enqueue!(q, i)
    end

    idx = 0
    ordered_eqs = similar(eqs, 0); sizehint!(ordered_eqs, neqs)
    while !isempty(q)
        𝑠eq = dequeue!(q)
        idx+=1
        push!(ordered_eqs, eqs[𝑠eq])
        var = assigns[𝑠eq]
        for 𝑑eq in 𝑑neighbors(graph, var)
            degree = degrees[𝑑eq] = degrees[𝑑eq] - 1
            degree == 0 && enqueue!(q, 𝑑eq)
        end
    end

    (check && idx != neqs) && throw(ArgumentError("The equations have at least one cycle."))

    return ordered_eqs
end

function observed2graph(eqs, states)
    graph = BipartiteGraph(length(eqs), length(states))
    v2j = Dict(states .=> 1:length(states))

    # `assigns: eq -> var`, `eq` defines `var`
    assigns = similar(eqs, Int)

    for (i, eq) in enumerate(eqs)
        lhs_j = get(v2j, eq.lhs, nothing)
        lhs_j === nothing && throw(ArgumentError("The lhs $(eq.lhs) of $eq, doesn't appear in states."))
        assigns[i] = lhs_j
        vs = vars(eq.rhs)
        for v in vs
            j = get(v2j, v, nothing)
            j !== nothing && add_edge!(graph, i, j)
        end
    end

    return graph, assigns
end


text:
using Base.Test
using LowRankModels

# Check that sparse algorithm converges to the same solution
# from the same initial conditions for simple pca.
m,n,k = 100,100,3
A = randn(m,k)*randn(k,n)
loss = QuadLoss()
r = ZeroReg()

glrm_1 = GLRM(A,loss,r,r,k) # solve with prox algorithm
glrm_2 = deepcopy(glrm_1)   # solve with sparse prox algorithm

X1,Y1,ch1 = fit!(glrm_1,ProxGradParams())
X2,Y2,ch2 = fit!(glrm_2,SparseProxGradParams())

@test_approx_eq X1 X2
@test_approx_eq Y1 Y2

# Check that the sparsity pattern in the data is correctly identified.
A = sprand(m,n,0.5)
I,J = findn(A)

glrm = GLRM(A,loss,r,r,k) # create glrm from sparse matrix

@test length(glrm.observed_features) == m
@test length(glrm.observed_examples) == n

for i = 1:m
	for j = 1:n
		if j in glrm.observed_features[i]
			@test A[i,j] != 0.0
		else
			@test A[i,j] == 0.0
		end
	end
end

for j = 1:n
	for i = 1:m
		if i in glrm.observed_examples[j]
			@test A[i,j] != 0.0
		else
			@test A[i,j] == 0.0
		end
	end
end

ext: jl


text:
"""
Helper functions for streaming lm

Author: Wayne Wang
Email: wayneyw@umich.edu
Last modified: 08/02/2021
"""

function plot_partitions(partitions::AbstractArray, x::AbstractArray, y::AbstractArray)
    """
    Plot the estimated regression curve on top of the point clouds
        for each partition
    """
    plt = scatter(x, y, legend = false)

    for i = 1:length(partitions)
        par = partitions[i]
        rng = par[1]
        coefs = par[2][2:3]
        f(x) = coefs[1] + coefs[2]*x
        plot!(f, rng[1], rng[2], color = :red, legend = false)
    end

    display(plt)
end

max_stars_repo_name: ayanbiswas/PRISM

max_issues_repo_name: ayanbiswas/PRISM

max_forks_repo_name: ayanbiswas/PRISM


text:
export Parallel

"""
    struct Parallel{N}

A parallel of `N` things.

# Fields
- `xs`: The things that are in parallel.
"""
struct Parallel{N}
    xs
end

Flux.functor(::Type{<:Parallel}, p) = p.xs, xs -> Parallel(xs...)

Parallel(xs...) = Parallel{length(xs)}(collect(xs))

Base.eltype(p::Parallel) = eltype(p.xs)
Base.getindex(p::Parallel, i) = p.xs[i]
Base.iterate(p::Parallel, state...) = iterate(p.xs, state...)
Base.length(p::Parallel) = length(p.xs)
Base.map(f, p::Parallel) = Parallel([map(f, x) for x in p.xs]...)

"""
    flatten(p::Parallel)

Flatten a recursive `Parallel` structure.

# Arguments
- `p::Parallel`: Recursive structure to flatten.

# Returns
- `Vector`: Flattened version of `p`.
"""
flatten(p::Parallel) = vcat(flatten.(p.xs)...)
flatten(x) = [x]

"""
    sample(p::Parallel; num_samples::Integer)

Sample a parallel of distributions.

# Arguments
- `p::Parallel`: Parallel of distributions.

# Keywords
- `num_samples::Integer`: Number of samples to take.

# Returns
- `Parallel`: A parallel of samples.
"""
sample(p::Parallel; num_samples::Integer) =
    Parallel(sample.(p.xs, num_samples=num_samples)...)

"""
    materialise(p::Parallel, f=xs -> repeat_cat(xs..., dims=2))

Turn a parallel of tensors into a tensor.

# Arguments
- `sample::Parallel`: Sample to turn into a tensor.
- `f`: Function that takes in a `Vector` of tensors and turns them into a tensor.

# Returns
- `AA`: Tensor corresponding to parallel.
"""
materialise(p::Parallel, f=xs -> repeat_cat(xs..., dims=2)) = f(materialise.(p.xs, f))
materialise(sample, f) = sample

_add(x, y) = x .+ y

"""
    kl(p::Parallel{N}, q::Parallel{N}) where N

Compute the KL divergence between two parallels of distributions.

# Arguments
- `p::Parallel{N}`: `p`.
- `q::Parallel{N}`: `q`.

# Returns
- `AA`: `KL(p, q)`.
"""
kl(p::Parallel{N}, q::Parallel{N}) where N =
    reduce(_add, [sum(kl(pᵢ, qᵢ), dims=(1, 2)) for (pᵢ, qᵢ) in zip(p, q)])

"""
    logpdf(d::Parallel{N}, x::Parallel{N}) where N

Compute log-pdf of a parallel of distributions at a parallel of values.

# Arguments
- `d::Parallel{N}`: Parallel of distributions.
- `x::Parallel{N}`: Parallel of values.

# Returns
- `AA`: Log-pdfs of `d` as `x`.
"""
logpdf(d::Parallel{N}, x::Parallel{N}) where N =
    reduce(_add, [sum(logpdf(dᵢ, xᵢ), dims=(1, 2)) for (dᵢ, xᵢ) in zip(d, x)])


text:
# functions related to binomial distribution

# R implementations
# For pdf and logpdf we use the Julia implementation
using .RFunctions:
    binomcdf,
    binomccdf,
    binomlogcdf,
    binomlogccdf,
    binominvcdf,
    binominvccdf,
    binominvlogcdf,
    binominvlogccdf

# Julia implementations
binompdf(n::Real, p::Real, k::Real) = exp(binomlogpdf(n, p, k))

binomlogpdf(n::Real, p::Real, k::Real) = binomlogpdf(promote(n, p, k)...)
function binomlogpdf(n::T, p::T, k::T) where {T<:Real}
    m = clamp(k, 0, n)
    val = min(0, betalogpdf(m + 1, n - m + 1, p) - log(n + 1))
    return 0 <= k <= n && isinteger(k) ? val : oftype(val, -Inf)
end


text:
export
    RefLine,
    value,
    reference_coordinates,
    faces


text:
using Test

include("makingtypes.jl")


text:
#======================================#
#        Functions for the user        #
#======================================#

function print_starting_cards(_Cards)
  println("\nThese are your cards mate,")

  for i in 1:(trunc(Int64, (length(_Cards) * 0.5 )))
    print(
      """\n\t
      \t#==============#\t#==============#
      \t    Card $(i)    \t\t    Card $(i + 2)
      \t              \t\t
      \t Health: $(_Cards[i]["Health"])   \t\t Health: $(_Cards[i + 2]["Health"])
      \t Attack: $(_Cards[i]["Attack"])   \t\t Attack: $(_Cards[i + 2]["Attack"])
      \t Defence: $(_Cards[i]["Defence"])  \t\t Defence: $(_Cards[i + 2]["Defence"])
      \t Speed: $(_Cards[i]["Speed"])    \t\t Speed: $(_Cards[i + 2]["Speed"])
      \t Crit: $(_Cards[i]["Critical"])     \t\t Crit: $(_Cards[i + 2]["Critical"])
      \t#==============#\t#==============#\n\n
      """)
  end
end

function pick_card()
  print("Choose the number card you want: ")
  @label start # @label will mark checkpoint
  userCardChoice = readline()
  userCardChoice = try; parse(Int, userCardChoice); catch e; print("Please only enter numbers: "); @goto start; end
  if userCardChoice ∉ 1:length(Cards) # ∉ means "not in"
    print("Please pick valid card: ")
    @goto start # go back to @label start
  else
    println("You picked card $userCardChoice. Good luck!")
  end
  return userCardChoice
end

function user_pick_move()
  print("Select the move that you want to use: attack, defend, dodge")
end

max_stars_repo_name: NarutoXY/austroporean-cards

max_issues_repo_name: NarutoXY/austroporean-cards

max_forks_repo_name: NarutoXY/austroporean-cards


text:
module RayTracingInOneJulia

using LinearAlgebra
using StaticArrays
using Images
using CUDA, KernelAbstractions, CUDAKernels
using Tullio

CUDA.allowscalar(false)

export Vec3, Point3

const Vec3 = SVector{3, T} where T
const Point3 = Vec3

include("ray.jl")
include("hittable.jl")
include("material.jl")
include("sphere.jl")
include("camera.jl")

export render!

function ray_color(r::Ray{T}, world, max_depth) where T
    depth = max_depth
    attenuation = RGB{T}(1)
    while true
        if depth <= 0
            return RGB{T}(0)
        end
        rec = hit(world, r, T(0.001), T(Inf))
        if !isnothing(rec)
            scat_rec = scatter(r, rec)
            if !isnothing(scat_rec)
                attenuation = attenuation ⊙ scat_rec.attenuation
                r = scat_rec.scattered
                depth -= one(depth)
            else
                return RGB{T}(0)
            end
        else
            unit_direction = normalize(r.dir)
            t = T(0.5) * (unit_direction.y + T(1))
            return attenuation ⊙ ((T(1.0)-t) * RGB{T}(1, 1, 1) + t * RGB{T}(0.5, 0.7, 1))
        end
    end
end

function write_color(pixel_color, samples_per_pixel)
    pixel_color = (1/samples_per_pixel) * pixel_color
    pixel_color = RGB(sqrt(pixel_color.r), sqrt(pixel_color.g), sqrt(pixel_color.b))
    return clamp01(pixel_color)
end

function render_pixel(::Type{T}, i, j, world, cam, image_width, image_height, samples_per_pixel, max_depth) where T
    (i, j, image_width, image_height) = T.((i, j, image_width, image_height))
    pixel_color = RGB{T}(0)
    for _ in Base.OneTo(samples_per_pixel)
        u = (j-1+rand(T)) / (image_width-1)
        v = (image_height-i+rand(T)) / (image_height-1)
        r = get_ray(cam, u, v)
        pixel_color += ray_color(r, world, max_depth)
    end
    return write_color(pixel_color, samples_per_pixel)
end

function render!(image::AbstractArray{RGB{T}}, world, cam, image_width, image_height, samples_per_pixel, max_depth) where T
    @tullio image[i,j] = render_pixel(T, i, j, $world, cam, image_width, image_height, Int32(samples_per_pixel), Int32(max_depth))
end

end # module RayTracingInOneJulia


text:
module S2Geometry

using Reexport

include("R1/R1.jl")

@reexport using .R1

end


text:
function hierNet(df; nfolds::Integer = 5, nlam::Integer = 10, strong::Bool = false,
   max_nest_size = 0, step_size = 1.01)
# --------------------------------------------------------------------------
# Options:
# nfolds - number of folds for cross-validation
# nlam - number of values of regularization parameter to try
# strong -  whether strong hierarchy is imposed or not
# --------------------------------------------------------------------------

# Unpack DataFrame df
s = Matrix(df[:, r"shares"]);
p = Matrix(df[:, r"prices"]);
x = Matrix(df[:, r"x"]);
iv = Matrix(df[:, r"demand_instruments"]);

J = size(s,2);

@rput s
@rput p
@rput iv
@rput strong
@rput nlam
@rput nfolds
@rput J

if max_nest_size == 0
  R"
  library(hierNet)
  J <-  ncol(s)

  rownames(s) <- seq(length=nrow(s))
  rownames(p) <- seq(length=nrow(p))
  rownames(iv) <- seq(length=nrow(iv))

  # Instrument for market shares
  s_hat <- matrix(nrow = nrow(s),ncol = ncol(s))
  one_mat <- matrix(1,nrow(s), 1)
  iv <- as.matrix(iv)
  for (j in 1:J) {
    iv_temp <- as.matrix(cbind(matrix(1,nrow(s),1),iv))
    beta<- solve(t(iv_temp)%*%iv_temp) %*% t(iv_temp) %*% s[,j]
    s_hat[,j] <-iv_temp%*% beta
  }

  # Estimate Hierarchical Lasso
  for (j in 1:J) {
     # Estimate path of regularization parameters
    structured_fit <- hierNet.path(as.matrix(s_hat), as.vector(p[,j]),
                                  lamlist = NULL, delta=1e-8, minlam = NULL, maxlam = NULL, nlam=nlam, flmin=.05,
                                  diagonal = TRUE, strong = strong, aa = NULL, zz = NULL,
                                  stand.main = TRUE, stand.int = TRUE,
                                  rho = nrow(s), niter = 100, sym.eps = 0.001,
                                  step = 1, maxiter = 2000, backtrack = 0.2, tol = 1e-05, trace = 0)
    if (j==1) {
      matsize <- nrow(structured_fit$bp)
      coefs <- matrix(nrow = 0 ,ncol = matsize)
    }

    # Select parameter by CV, append to matrix
    a<- hierNet.cv(structured_fit, as.matrix(s_hat), as.vector(p[,j]), nfolds=nfolds,folds=NULL,trace=0)
    lambdahat <- a['lamhat.1se']
    # lamhat lamhat.1se

    # Final Estimates
    structured_fit<- hierNet(as.matrix(s_hat), as.vector(p[,j]), as.numeric(lambdahat), delta=1e-8, strong=strong, diagonal=TRUE, aa=NULL,
                               zz=NULL,center=TRUE, stand.main=TRUE, stand.int=TRUE, rho=nrow(s), niter=100,
                               sym.eps=1e-3,step=1, maxiter=5000, backtrack=0.2, tol=1e-5, trace=0)
    coefs <- rbind(coefs,t(structured_fit$bp - structured_fit$bn))
  }
  included = (abs(coefs)>0)

  for (j in 1:J) {
    included[j,j] = 1
  }

  included_symmetric <- included
  for (j1 in 1:J) {
    for (j2 in 1:J) {
      if (included_symmetric[j1,j2] != included_symmetric[j2,j1]){
        included_symmetric[j1,j2] = 1
        included_symmetric[j2,j1] = 1
      }
    }
  }
  "
  @rget included
  @rget included_symmetric
else
  @rput max_nest_size
  @rput step_size
  R"
  library(hierNet)
  J <-  ncol(s)

  rownames(s) <- seq(length=nrow(s))
  rownames(p) <- seq(length=nrow(p))
  rownames(iv) <- seq(length=nrow(iv))

  # Instrument for market shares
  s_hat <- matrix(nrow = nrow(s),ncol = ncol(s))
  one_mat <- matrix(1,nrow(s), 1)
  iv <- as.matrix(iv)
  for (j in 1:J) {
    iv_temp <- as.matrix(cbind(matrix(1,nrow(s),1),iv))
    beta<- solve(t(iv_temp)%*%iv_temp) %*% t(iv_temp) %*% s[,j]
    s_hat[,j] <-iv_temp%*% beta
  }

  # Estimate Hierarchical Lasso
  for (j in 1:J) {
    nest_size = J+1
    lam = 0.1
    while (nest_size > max_nest_size) {
      lam = lam * step_size;
      fit <- hierNet(as.matrix(s_hat), as.vector(p[,j]), lam, delta=1e-8, strong=strong, diagonal=TRUE, aa=NULL, zz=NULL,
        center=TRUE, stand.main=TRUE, stand.int=FALSE, niter=100, sym.eps=1e-3,
        step=1, maxiter=2000, backtrack=0.2, tol=1e-5, trace=0)
      abs_coef <- abs(t(fit$bp - fit$bn))
      abs_coef[j] <- 1
      nest_size <- sum((abs_coef > 0))
    }

    if (j==1) {
      matsize <- J
      coefs <- matrix(nrow = 0 ,ncol = matsize)
    }

    # Final Estimates
    structured_fit<- hierNet(as.matrix(s_hat), as.vector(p[,j]), as.numeric(lam), delta=1e-8, strong=strong, diagonal=TRUE, aa=NULL,
                               zz=NULL,center=TRUE, stand.main=TRUE, stand.int=TRUE, rho=nrow(s), niter=100,
                               sym.eps=1e-3,step=1, maxiter=5000, backtrack=0.2, tol=1e-5, trace=0)
    coefs <- rbind(coefs,t(structured_fit$bp - structured_fit$bn))
  }
  included = (abs(coefs)>0)

  for (j in 1:J) {
    included[j,j] = 1
  }

  "
  @rget included
  included_symmetric = [];
end

return included, included_symmetric
end


text:
using Printf
using JLD2
using DataDeps
using OceanTurb
using CairoMakie
using WindMixing
using OceanParameterizations
using Images, FileIO
using ImageTransformations

include("modified_pacalowski_philander_model.jl")

train_files = ["wind_-5.5e-4_new"]
𝒟 = WindMixing.data(train_files, scale_type=ZeroMeanUnitVarianceScaling, enforce_surface_fluxes=false)

NDE_profile_path = joinpath(pwd(), "final_results", "18sim_old", "test_$(train_files[1])", "profiles_fluxes_oceananigans.jld2")
file = jldopen(NDE_profile_path)
data = file["NDE_profile"]
close(file)

Nz = 32
Lz = 256
u₀ = 𝒟.u.coarse[:,1]
v₀ = 𝒟.v.coarse[:,1]
T₀ = 𝒟.T.coarse[:,1]
Fu = 𝒟.uw.coarse[end,1]
Fθ = 𝒟.wT.coarse[end,1]
f₀ = 1f-4
g = 9.80655f0
α = 2f-4
∂T₀∂z_bottom = (T₀[end-1] - T₀[end]) / (Lz / Nz)

times = 𝒟.t

# constants_kpp = (f=f₀, α=α, g=g, Nz=Nz, H=Lz)
# BCs_kpp = (uw=(top=𝒟.uw.coarse[end, 1], bottom=𝒟.uw.coarse[1, 1]),
#             vw=(top=𝒟.vw.coarse[end, 1], bottom=𝒟.uw.coarse[1, 1]),
#             wT=(top=𝒟.wT.coarse[end, 1], bottom=𝒟.wT.coarse[1, 1]))
# ICs_kpp = (u=u₀, v=v₀, T=T₀)

# column_model_1D_kpp(constants_kpp, BCs_kpp, ICs_kpp, times, OceanTurb.KPP.Parameters())

constants_pp = OceanTurb.Constants(Float64, f=f₀)
parameters_pp = PacanowskiPhilander.Parameters()
model_pp = PacanowskiPhilander.Model(N=Nz, L=Lz, stepper=:BackwardEuler, constants=constants_pp, parameters=parameters_pp)

model_pp.bcs[1].top = OceanTurb.FluxBoundaryCondition(Fu)
model_pp.bcs[3].top = OceanTurb.FluxBoundaryCondition(Fθ)

model_pp.solution[1].data[1:Nz] .= u₀
model_pp.solution[2].data[1:Nz] .= v₀
model_pp.solution[3].data[1:Nz] .= T₀

Δt = 1.0
Nt = length(times)

u_solution_pp = zeros(Nz, Nt)
v_solution_pp = zeros(Nz, Nt)
T_solution_pp = zeros(Nz, Nt)

uw_solution_pp = zeros(Nz+1, Nt)
vw_solution_pp = zeros(Nz+1, Nt)
wT_solution_pp = zeros(Nz+1, Nt)

Ri_solution_pp = zeros(Nz+1, Nt)

function get_diffusive_flux(field_index, model)
    flux = FaceField(model.grid)
    field = model.solution[field_index]
    K = model.timestepper.eqn.K[field_index]
    for i in interiorindices(flux)
        @inbounds flux[i] = - K(model, i) * ∂z(field, i)
    end
    return flux
end

function get_richardson_number_profile(model)
    Ri = FaceField(model.grid)
    for i in interiorindices(Ri)
        @inbounds Ri[i] = local_richardson(model, i)
    end
    return Ri
end

for n in 1:Nt
    OceanTurb.run_until!(model_pp, Δt, times[n])
    @info "Time = $(times[n])"

    u_solution_pp[:, n] .= model_pp.solution[1][1:Nz]
    v_solution_pp[:, n] .= model_pp.solution[2][1:Nz]
    T_solution_pp[:, n] .= model_pp.solution[3][1:Nz]

    uw_solution_pp[:, n] .= get_diffusive_flux(1, model_pp)[1:Nz+1]
    vw_solution_pp[:, n] .= get_diffusive_flux(2, model_pp)[1:Nz+1]
    wT_solution_pp[:, n] .= get_diffusive_flux(3, model_pp)[1:Nz+1]

    uw_solution_pp[Nz+1, n] = Fu
    wT_solution_pp[Nz+1, n] = Fθ

    # Ri_solution_pp[:, n] = get_richardson_number_profile(model_pp)[1:Nz+1]
end

D_face = Dᶠ(32, 256 / 32)

function local_richardson_profile(u, v, T, g, α)
    ∂u∂z = D_face * u
    ∂v∂z = D_face * v
    ∂T∂z = D_face * T
    Bz = g .* α .* ∂T∂z
    S² = ∂u∂z .^ 2 .+ ∂v∂z .^ 2
    return Bz ./ S²
end

for i in 1:size(Ri_solution_pp, 2)
    Ri_solution_pp[:,i] = local_richardson_profile(u_solution_pp[:,i], v_solution_pp[:,i], T_solution_pp[:,i], g, α)
end

# Get rid of ∞ and super large values.
Ri_solution_pp = clamp.(Ri_solution_pp, -1, 2)

u_img = rotr90(load(assetpath(joinpath(pwd(), "images_plot", "u.png"))))
v_img = rotr90(load(assetpath(joinpath(pwd(), "images_plot", "v.png"))))
T_img = rotr90(load(assetpath(joinpath(pwd(), "images_plot", "T.png"))))
uw_img = rotr90(load(assetpath(joinpath(pwd(), "images_plot", "uw.png"))))
vw_img = rotr90(load(assetpath(joinpath(pwd(), "images_plot", "vw.png"))))
wT_img = rotr90(load(assetpath(joinpath(pwd(), "images_plot", "wT.png"))))
Ri_img = rotr90(load(assetpath(joinpath(pwd(), "images_plot", "Ri.png"))))
z_img = load(assetpath(joinpath(pwd(), "images_plot", "z.png")))

axis_images = (
    u = u_img,
    v = v_img,
    T = T_img,
    z = z_img,
    uw = uw_img,
    vw = vw_img,
    wT = wT_img,
    Ri = Ri_img,
)

FILE_PATH = joinpath(pwd(), "final_results", "base_kpp_pp")
fps=60

times_days = data["t"] ./ 86400

frame = Node(1)

time_point = @lift [times_days[$frame]]

u_data = [
    data["truth_u"],
    # data["test_u_modified_pacanowski_philander"],
    data["test_u_kpp"],
    # data["test_u"],
    u_solution_pp,
]

v_data = [
    data["truth_v"],
    # data["test_v_modified_pacanowski_philander"],
    data["test_v_kpp"],
    # data["test_v"],
    v_solution_pp,
]

T_data = [
    data["truth_T"],
    # data["test_T_modified_pacanowski_philander"],
    data["test_T_kpp"],
    # data["test_T"],
    T_solution_pp,
]

uw_data = [
    data["truth_uw"],
    # data["test_uw_modified_pacanowski_philander"],
    data["test_uw_kpp"],
    # data["test_uw"],
    uw_solution_pp,
]

vw_data = [
    data["truth_vw"],
    # data["test_vw_modified_pacanowski_philander"],
    data["test_vw_kpp"],
    # data["test_vw"],
    vw_solution_pp,
]

wT_data = [
    data["truth_wT"],
    # data["test_wT_modified_pacanowski_philander"],
    data["test_wT_kpp"],
    # data["test_wT"],
    wT_solution_pp,
]

uw_data .*= 1f4
vw_data .*= 1f4
wT_data .*= 1f5

Ri_data = [
    clamp.(data["truth_Ri"], -1, 2),
    # clamp.(data["test_Ri_modified_pacanowski_philander"], -1, 2),
    clamp.(data["test_Ri_kpp"], -1, 2),
    # clamp.(data["test_Ri"], -1, 2),
    Ri_solution_pp
]

u_frames = [@lift data[:,$frame] for data in u_data]
v_frames = [@lift data[:,$frame] for data in v_data]
T_frames = [@lift data[:,$frame] for data in T_data]

uw_frames = [@lift data[:,$frame] for data in uw_data]
vw_frames = [@lift data[:,$frame] for data in vw_data]
wT_frames = [@lift data[:,$frame] for data in wT_data]

Ri_frames = [@lift data[:,$frame] for data in Ri_data]

# losses_point_frames = [@lift [data[$frame]] for data in losses_data]

@inline function find_lims(datasets)
    return maximum(maximum.(datasets)), minimum(minimum.(datasets))
end

u_max, u_min = find_lims(u_data)
v_max, v_min = find_lims(v_data)
T_max, T_min = find_lims(T_data)

uw_max, uw_min = find_lims(uw_data)
vw_max, vw_min = find_lims(vw_data)
wT_max, wT_min = find_lims(wT_data)

# losses_max, losses_min = find_lims(losses_data)

train_parameters = data["train_parameters"]
ν₀ = train_parameters.ν₀
ν₋ = train_parameters.ν₋
ΔRi = train_parameters.ΔRi
Riᶜ = train_parameters.Riᶜ
Pr = train_parameters.Pr
loss_scalings = train_parameters.loss_scalings

BC_str = @sprintf "Momentum Flux = %.1e m² s⁻², Temperature Flux = %.1e m s⁻¹ °C" data["truth_uw"][end, 1] maximum(data["truth_wT"][end, :])

plot_title = @lift "Traditional Parameterisations: $BC_str, Time = $(round(times_days[$frame], digits=2)) days"

# fig = Figure(resolution=(1920, 1080))
fig = Figure(resolution=(1920, 960))

# colors = distinguishable_colors(length(uw_data)+1, [RGB(1,1,1), RGB(0,0,0)], dropseed=true)
colors = distinguishable_colors(length(uw_data), [RGB(1,1,1), RGB(0,0,0)], dropseed=true)

u_img = axis_images.u
v_img = axis_images.v
T_img = axis_images.T
uw_img = axis_images.uw
vw_img = axis_images.uw
wT_img = axis_images.uw
Ri_img = axis_images.uw
z_img = axis_images.z

zc = data["depth_profile"]
zf = data["depth_flux"]
zf_interior = zf[2:end-1]

rel_size = 30
# aspect = 1920 / 1080
aspect = 2

ax_u = fig[1, 2] = CairoMakie.Axis(fig)
ax_v = fig[1, 4] = CairoMakie.Axis(fig)

T_layout = fig[1:4, 5] = GridLayout()
colsize!(fig.layout, 5, CairoMakie.Relative(0.4))

ax_T = T_layout[1, 2] = CairoMakie.Axis(fig)
y_ax_T = T_layout[1,1] = CairoMakie.Axis(fig, aspect=DataAspect())
x_ax_T = T_layout[2,2] = CairoMakie.Axis(fig, aspect=DataAspect())

ax_Ri = fig[1, 7] = CairoMakie.Axis(fig)
ax_uw = fig[3, 2] = CairoMakie.Axis(fig)
ax_vw = fig[3, 4] = CairoMakie.Axis(fig)
ax_wT = fig[3, 7] = CairoMakie.Axis(fig)

y_ax_u = CairoMakie.Axis(fig[1,1], aspect=DataAspect())
y_ax_v = CairoMakie.Axis(fig[1,3], aspect=DataAspect())
y_ax_Ri = CairoMakie.Axis(fig[1,6], aspect=DataAspect())
y_ax_uw = CairoMakie.Axis(fig[3,1], aspect=DataAspect())
y_ax_vw = CairoMakie.Axis(fig[3,3], aspect=DataAspect())
y_ax_wT = CairoMakie.Axis(fig[3,6], aspect=DataAspect())

x_ax_u = CairoMakie.Axis(fig[2,2], aspect=DataAspect())
x_ax_v = CairoMakie.Axis(fig[2,4], aspect=DataAspect())
x_ax_Ri = CairoMakie.Axis(fig[2,7], aspect=DataAspect())
x_ax_uw = CairoMakie.Axis(fig[4,2], aspect=DataAspect())
x_ax_vw = CairoMakie.Axis(fig[4,4], aspect=DataAspect())
x_ax_wT = CairoMakie.Axis(fig[4,7], aspect=DataAspect())

hidedecorations!(y_ax_u)
hidedecorations!(y_ax_v)
hidedecorations!(y_ax_Ri)
hidedecorations!(y_ax_uw)
hidedecorations!(y_ax_vw)
hidedecorations!(y_ax_wT)

hidedecorations!(x_ax_u)
hidedecorations!(x_ax_v)
hidedecorations!(x_ax_Ri)
hidedecorations!(x_ax_uw)
hidedecorations!(x_ax_vw)
hidedecorations!(x_ax_wT)

hidespines!(y_ax_u)
hidespines!(y_ax_v)
hidespines!(y_ax_T)
hidespines!(y_ax_Ri)
hidespines!(y_ax_uw)
hidespines!(y_ax_vw)
hidespines!(y_ax_wT)

hidespines!(x_ax_u)
hidespines!(x_ax_v)
hidespines!(x_ax_T)
hidespines!(x_ax_Ri)
hidespines!(x_ax_uw)
hidespines!(x_ax_vw)
hidespines!(x_ax_wT)

image!(x_ax_u, axis_images.u)
image!(x_ax_v, axis_images.v)
image!(x_ax_Ri, axis_images.Ri)
image!(x_ax_uw, axis_images.uw)
image!(x_ax_vw, axis_images.vw)
image!(x_ax_wT, axis_images.wT)

image!(y_ax_u, axis_images.z)
image!(y_ax_v, axis_images.z)
image!(y_ax_Ri, axis_images.z)
image!(y_ax_uw, axis_images.z)
image!(y_ax_vw, axis_images.z)
image!(y_ax_wT, axis_images.z)

hidedecorations!(y_ax_T)
hidedecorations!(x_ax_T)
image!(x_ax_T, axis_images.T)
image!(y_ax_T, axis_images.z)

colsize!(T_layout, 1, CairoMakie.Relative(1 / rel_size * 1.5))
rowsize!(T_layout, 2, CairoMakie.Relative(1 / rel_size / aspect * 1.75))
colgap!(T_layout, Relative(1 / rel_size / aspect / 2))
rowgap!(T_layout, Relative(1 / rel_size / aspect))

rowsize!(fig.layout, 2, CairoMakie.Relative(1 / rel_size))
rowsize!(fig.layout, 4, CairoMakie.Relative(1 / rel_size))
colsize!(fig.layout, 1, CairoMakie.Relative(1 / rel_size / aspect))
colsize!(fig.layout, 3, CairoMakie.Relative(1 / rel_size / aspect))
colsize!(fig.layout, 6, CairoMakie.Relative(1 / rel_size / aspect))

colgap!(fig.layout, Relative(1 / rel_size / aspect / 2))
rowgap!(fig.layout, Relative(1 / rel_size / aspect))

alpha=0.5
truth_linewidth = 7
linewidth = 3

CairoMakie.xlims!(ax_u, u_min, u_max)
CairoMakie.xlims!(ax_v, v_min, v_max)
CairoMakie.xlims!(ax_T, T_min, T_max)
CairoMakie.xlims!(ax_uw, uw_min, uw_max)
CairoMakie.xlims!(ax_vw, vw_min, vw_max)
CairoMakie.xlims!(ax_wT, wT_min, wT_max)
CairoMakie.xlims!(ax_Ri, -1, 2)
# CairoMakie.xlims!(ax_losses, times[1], times[end])

CairoMakie.ylims!(ax_u, minimum(zc), 0)
CairoMakie.ylims!(ax_v, minimum(zc), 0)
CairoMakie.ylims!(ax_T, minimum(zc), 0)
CairoMakie.ylims!(ax_uw, minimum(zf), 0)
CairoMakie.ylims!(ax_vw, minimum(zf), 0)
CairoMakie.ylims!(ax_wT, minimum(zf), 0)
CairoMakie.ylims!(ax_Ri, minimum(zf), 0)
# CairoMakie.ylims!(ax_losses, losses_min, losses_max)

u_lines = [
     lines!(ax_u, u_frames[1], zc, linewidth=truth_linewidth, color=(colors[1], alpha));
    [lines!(ax_u, u_frames[i], zc, linewidth=linewidth, color=colors[i]) for i in 2:length(u_data)]
]

v_lines = [
     lines!(ax_v, v_frames[1], zc, linewidth=truth_linewidth, color=(colors[1], alpha));
    [lines!(ax_v, v_frames[i], zc, linewidth=linewidth, color=colors[i]) for i in 2:length(v_data)]
]

T_lines = [
     lines!(ax_T, T_frames[1], zc, linewidth=truth_linewidth, color=(colors[1], alpha));
    [lines!(ax_T, T_frames[i], zc, linewidth=linewidth, color=colors[i]) for i in 2:length(T_data)]
]

uw_lines = [
     lines!(ax_uw, uw_frames[1], zf, linewidth=truth_linewidth, color=(colors[1], alpha));
    [lines!(ax_uw, uw_frames[i], zf, linewidth=linewidth, color=colors[i]) for i in 2:length(uw_data)]
]

vw_lines = [
     lines!(ax_vw, vw_frames[1], zf, linewidth=truth_linewidth, color=(colors[1], alpha));
    [lines!(ax_vw, vw_frames[i], zf, linewidth=linewidth, color=colors[i]) for i in 2:length(vw_data)]
]

wT_lines = [
    lines!(ax_wT, wT_frames[1], zf, linewidth=truth_linewidth, color=(colors[1], alpha));
   [lines!(ax_wT, wT_frames[i], zf, linewidth=linewidth, color=colors[i]) for i in 2:length(wT_data)]
]

Ri_lines = [
     lines!(ax_Ri, Ri_frames[1], zf, linewidth=truth_linewidth, color=(colors[1], alpha));
    [lines!(ax_Ri, Ri_frames[i], zf, linewidth=linewidth, color=colors[i]) for i in 2:length(Ri_data)]
]

axislegend(ax_T, T_lines, ["Oceananigans.jl Large Eddy Simulation", "K-Profile Parameterisation", "Pacanowski-Philander Model"], "Data Type", position = :rb)

supertitle = fig[0, :] = Label(fig, plot_title, textsize=25)

trim!(fig.layout)

print_frame = maximum([1, Int(floor(length(times)/20))])

function print_progress(n, n_total, print_frame, type)
    if n % print_frame == 0
        @info "Animating $(type) frame $n/$n_total"
    end
end

@info "Starting Animation"

CairoMakie.record(fig, "$FILE_PATH.mp4", 1:length(times_days), framerate=fps, compression=1) do n
    print_progress(n, length(times_days), print_frame, "mp4")
    frame[] = n
end


text:
export layer_norm, batched_mlp, Splitter, MeanPooling, SumPooling

"""
    struct LayerNorm

Layer normalisation.

# Fields
- `w`: Weights.
- `b`: Biases.
- `dims`: Dimensions to apply the normalisation to.
"""
struct LayerNorm
    w
    b
    dims::Tuple
end

@Flux.functor LayerNorm

"""
    layer_norm(shape::Integer...)

Construct a `LayerNorm` layer.

# Arguments
- `shape...`: A tuple containing one integer per dimension. Set a dimension to `1` to not
    normalise. Set a dimension to the size of that dimension to do normalise.

# Returns
- `LayerNorm`: Corresponding layer.
"""
function layer_norm(shape::Integer...)
    return LayerNorm(
        ones(Float32, shape...),
        zeros(Float32, shape...),
        Tuple(findall(x -> x > 1, shape))
    )
end

function (layer::LayerNorm)(x::AA)
    x = x .- mean(x, dims=layer.dims)
    x = x ./ sqrt.(mean(x .* x, dims=layer.dims) .+ 1f-8)
    return x .* layer.w .+ layer.b
end

"""
    struct BatchedMLP

# Fields
- `mlp`: MLP.
- `dim_out::Integer`: Dimensionality of the output.
"""
struct BatchedMLP
    mlp
    dim_out::Integer
end

@Flux.functor BatchedMLP

function (layer::BatchedMLP)(x::AA)
    x, back = to_rank(3, x)  # Compress all batch dimensions.

    n, _, batch_size = size(x)

    # Merge data point and batch dimension.
    x = permutedims(x, (2, 1, 3))
    x = reshape(x, :, n * batch_size)

    x = layer.mlp(x)

    # Unmerge data point and batch dimension.
    x = reshape(x, :, n, batch_size)
    x = permutedims(x, (2, 1, 3))

    return back(x)  # Uncompress batch dimensions.
end

"""
    batched_mlp(;
        dim_in::Integer,
        dim_hidden::Integer=dim_in,
        dim_out::Integer,
        num_layers::Integer,
        act=x -> leakyrelu(x, 0.1f0)
    )

Construct a batched MLP.

# Keywords
- `dim_in::Integer`: Dimensionality of the input.
- `dim_hidden::Integer`: Dimensionality of the hidden layers.
- `dim_out::Integer`: Dimensionality of the output.
- `num_layers::Integer`: Number of layers.
- `act=x -> leakyrelu(x, 0.1f0)`: Activation function to use.

# Returns
- `BatchedMLP`: Corresponding batched MLP.
"""
function batched_mlp(;
    dim_in::Integer,
    dim_hidden::Integer,
    dim_out::Integer,
    num_layers::Integer,
    act=x -> leakyrelu(x, 0.1f0)
)
    if num_layers == 1
        return BatchedMLP(Chain(Dense(dim_in, dim_out)), dim_out)
    else
        layers = Any[Dense(dim_in, dim_hidden, act)]
        for i = 1:num_layers - 2
            push!(layers, Dense(dim_hidden, dim_hidden, act))
        end
        push!(layers, Dense(dim_hidden, dim_out))
        return BatchedMLP(Chain(layers...), dim_out)
    end
end

"""
    struct Splitter

# Fields
- `num_channels::Integer`: Number of channels to split off.
"""
struct Splitter
    num_channels::Integer
end

@Flux.functor Splitter

function (layer::Splitter)(x::AA)
    num_remaining_channels = size(x, 2) - layer.num_channels
    x₁ = slice_at(x, 2, 1:num_remaining_channels)
    x₂ = slice_at(x, 2, (num_remaining_channels + 1):size(x, 2))
    return Parallel(x₁, x₂)
end

"""
    abstract type Pooling end

Poolings over dimension one.
"""
abstract type Pooling end

"""
    struct MeanPooling

Mean pooling over dimension one.

# Fields
- `ln`: Layer normalisation to apply after pooling.
"""
struct MeanPooling <: Pooling
    ln
end

@Flux.functor MeanPooling

(layer::MeanPooling)(x::AA) = layer.ln(mean(x, dims=1))

"""
    struct SumPooling

Sum pooling over dimension one.

# Fields
- `factor::Integer`: Factor to divide by after pooling to help initialisation.
"""
struct SumPooling <: Pooling
    factor::Integer
end

@Flux.functor SumPooling

(layer::SumPooling)(x::AA) = sum(x, dims=1) ./ layer.factor


text:
using Dates, CompoundPeriods, TimeZones, TimesDates
using Test

str1 = "2011-05-08T11:15:00"
str2 = "2015-08-17T21:08:11.125"
str3 = "2018-03-09T18:29:00.04296875"

dt1 = DateTime(str1)
dt2 = DateTime(str2)

td1 = TimeDate(str1)
td2 = TimeDate(str2)
td3 = TimeDate(str3)

dtz1 = ZonedDateTime(dt1, tz"UTC")
dtz2 = ZonedDateTime(dt1, tz"America/New_York")
tdz1 = TimeDateZone(td1, tz"UTC")
tdz2 = TimeDateZone(td1, tz"America/New_York")

@test string(td1) == str1
@test string(td2) == str2
@test string(td3) == str3

@test DateTime(td1) == dt1
@test DateTime(td2) == dt2

@test Year(td1) == Year(dt1)
@test Minute(td2) == Minute(dt2)

@test Minute(td1 + Minute(1)) == Minute(td1) + Minute(1)
@test Microsecond(td3 - Microsecond(1)) == Microsecond(td3) - Microsecond(1)

@test TimeDateZone(dtz1) == tdz1
@test TimeDateZone(dtz2) == tdz2

@test yearmonthday(td1) == yearmonthday(dt1)
@test lastdayofmonth(td1) == lastdayofmonth(dt1)

str2000ms    = "2000-01-01T00:00:00.123"
str2000ms_ut = string(str2000ms, "+00:00")
str2000ms_ny = string(str2000ms, "-05:00")

str2000ns    = "2000-01-01T00:00:00.123456789"
str2000ns_ut = string(str2000ns, "+00:00")
str2000ns_ny = string(str2000ns, "-05:00")

zdt2000_ut = ZonedDateTime(str2000ms_ut)
zdt2000_ny = ZonedDateTime(str2000ms_ny)

tdz2000_ut = TimeDateZone(str2000ns_ut)
tdz2000_ny = TimeDateZone(str2000ns_ny)

@test  zdt2000_ut == ZonedDateTime(tdz2000_ut)
@test  zdt2000_ny == ZonedDateTime(tdz2000_ny)

zdt = ZonedDateTime(DateTime("2011-05-08T12:11:15.050"), tz"Australia/Sydney")
tdz = TimeDateZone(zdt)
tdz += Microsecond(11)

# recasted tests from TimeZones.jl/test/

warsaw = tz"Europe/Warsaw"

normal = DateTime(2015, 1, 1, 0)   # a 24 hour day in warsaw
spring = DateTime(2015, 3, 29, 0)  # a 23 hour day in warsaw
fall = DateTime(2015, 10, 25, 0)   # a 25 hour day in warsaw

# Unary plus
@test +TimeDateZone(normal, warsaw) == TimeDateZone(normal, warsaw)

# Period arithmetic
@test TimeDateZone(normal, warsaw) + Day(1) == TimeDateZone(normal + Day(1), warsaw)
@test TimeDateZone(spring, warsaw) + Day(1) == TimeDateZone(spring + Day(1), warsaw)
@test TimeDateZone(fall, warsaw) + Day(1) == TimeDateZone(fall + Day(1), warsaw)

@test TimeDateZone(normal, warsaw) + Hour(24) == TimeDateZone(normal + Hour(24), warsaw)
@test TimeDateZone(spring, warsaw) + Hour(24) == TimeDateZone(spring + Hour(25), warsaw)
@test TimeDateZone(fall, warsaw) + Hour(24) == TimeDateZone(fall + Hour(23), warsaw)

# Do the same calculations but backwards over the transitions.
@test TimeDateZone(normal + Day(1), warsaw) - Day(1) == TimeDateZone(normal, warsaw)
@test TimeDateZone(spring + Day(1), warsaw) - Day(1) == TimeDateZone(spring, warsaw)
@test TimeDateZone(fall + Day(1), warsaw) - Day(1) == TimeDateZone(fall, warsaw)

@test TimeDateZone(normal + Day(1), warsaw) - Hour(24) == TimeDateZone(normal, warsaw)
@test TimeDateZone(spring + Day(1), warsaw) - Hour(23) == TimeDateZone(spring, warsaw)
@test TimeDateZone(fall + Day(1), warsaw) - Hour(25) == TimeDateZone(fall, warsaw)

# Ensure that arithmetic around transitions works.
@test TimeDateZone(spring, warsaw) + Hour(1) == TimeDateZone(spring + Hour(1), warsaw)
@test TimeDateZone(spring, warsaw) + Hour(2) == TimeDateZone(spring + Hour(3), warsaw)
# ambiguous
# @test TimeDateZone(fall, warsaw) + Hour(2) == TimeDateZone(fall + Hour(2), warsaw, 1)
# @test TimeDateZone(fall, warsaw) + Hour(3) == TimeDateZone(fall + Hour(2), warsaw, 2)

# CompoundPeriod canonicalization interacting with period arithmetic. Since `spring_zdt` is
# a 23 hour day this means adding `Day(1)` and `Hour(23)` are equivalent.
spring_zdt = TimeDateZone(spring, warsaw)
@test spring_zdt + Day(1) + Minute(1) == spring_zdt + Hour(23) + Minute(1)

# Non-Associativity

zdt_explicit_hourday  = ZonedDateTime(2015, 3, 31, 1, warsaw)
tdz_explicit_hourday  = TimeDateZone(zdt_explicit_hourday)
tdz_explicit_hour_day = (TimeDateZone(spring, warsaw) + Hour(24)) + Day(1)
@test zdt_explicit_hourday == ZonedDateTime(tdz_explicit_hourday)
@test tdz_explicit_hourday == tdz_explicit_hour_day

zdt_explicit_dayhour = ZonedDateTime(2015, 3, 31, 0, warsaw)
tdz_explicit_dayhour  = TimeDateZone(zdt_explicit_dayhour)
tdz_explicit_day_hour = (TimeDateZone(spring, warsaw) + Day(1)) + Hour(24)
@test zdt_explicit_dayhour == ZonedDateTime(tdz_explicit_dayhour)
@test tdz_explicit_dayhour == tdz_explicit_day_hour

zdt_implicit_hourday  = ZonedDateTime(2015, 3, 31, 0, warsaw)
tdz_implicit_hourday = TimeDateZone(zdt_implicit_hourday)
tdz_implicit_hour_day = TimeDateZone(spring, warsaw) + (Hour(24) + Day(1))
@test zdt_implicit_hourday == ZonedDateTime(tdz_implicit_hourday)
@test tdz_implicit_hourday == tdz_implicit_hour_day
tdz_implicit_hour_day = TimeDateZone(spring, warsaw) + Hour(24) + Day(1)
@test tdz_implicit_hourday == tdz_implicit_hour_day

zdt_implicit_dayhour = ZonedDateTime(2015, 3, 31, 0, warsaw)
tdz_implicit_dayhour  = TimeDateZone(zdt_implicit_dayhour)
tdz_implicit_day_hour = TimeDateZone(spring, warsaw) + Day(1) + Hour(24)
@test zdt_implicit_dayhour == ZonedDateTime(tdz_implicit_dayhour)
@test tdz_implicit_dayhour == tdz_implicit_day_hour

Δms = DateTime(str2)-DateTime(1900)
@test TimeDate(str2)-Δms == TimeDate(1900)
@test TimeDate(str2)+(-Δms) == TimeDate(1900)
@test TimeDate(1900)+Δms == TimeDate(str2)
@test TimeDate(1900)-(-Δms) == TimeDate(str2)
@test TimeDate(str2)-Nanosecond(Δms) == TimeDate(1900)

Δtd = td3 - TimeDate(1900)
@test canonical(Δtd) == Nanosecond(3729608940042968750) == mapfoldl(Nanosecond,+,canonical(Δtd))
@test dt2 - dt1 == td2 - td1
let
    td1 = TimeDate(2020,2,28,23,59,59,999,999,990)
    td2 = TimeDate(2020,3,1,0,0,0,0,0,10)
    @test td2 - td1 == Day(1)+Nanosecond(20)
    @test td1 - td2 == Day(-1)+Nanosecond(-20) == Nanosecond(-86400000000020)
end


text:
using TestImages, ImageCore

img = testimage("lena_color_256")

cv = channelview(img)
r=collect(reinterpret(UInt8, cv))
@show typeof(r), size(r) # this is what I have

# rc = r/255.0

# im2 = colorview(RGB, normedview(r))
# im2 = reinterpret(RGB{N0f8}, r)
im2 = colorview(RGB{N0f8}, r)
display(im2)
nothing


text:
#=
Force calculations according to Pammer OConner 2013
where force magnitude is estimated based on the degree of bending
in the same section of whisker in every frame selected to be high
SNR (Point P)

=#

#=
x, y - whisker coordinates
theta_f - whisker angle (radians)
curv - whisker curvature
ii - index of contact
i_p - index of high SNR point - we should use the second DLC point for this (as long as it isn't past pole)
E - Elastic Modulus of Whisker
I_p - Moment of inertia of whisker at point p I_p = 1.0
=#
function calc_force(x::Array{Float64,1},y::Array{Float64,1},theta_f::Float64,curv::Float64,ii::Int,i_p::Int,curv_0=0.0; I_p = 1.0, E=1.0)

    #curv_0 intrinsic curvature of the whisker
    delta_kappa = curv - curv_0 #Change in curvature

    #x_f, y_f follicle x and y
    x_f = x[end]
    y_f = y[end]

    #x_c, y_c contact x and y
    x_c = x[ii]
    y_c = y[ii]

    #x_p, y_p point_p x and y
    x_p = x[i_p]
    y_p = y[i_p]

    theta_contact = 0.0
    try
        #We will first look outward from the point of contact. This will fail if there is not >30 pixels of length
        #beyond that point
        c_i = get_contact_ind(x,y,ii,30.0,false)
        theta_contact = contact_angle(x[ii],x[c_i],y[ii],y[c_i])
    catch
        #if this fails, we will try making a vector from
        c_i = get_contact_ind(x,y,ii,30.0,true)
        theta_contact = contact_angle(x[c_i],x[ii],y[c_i],y[ii])
    end

    (theta_0,r_0) = calc_r_theta(x_f,x_c,y_f,y_c)
    (theta_p,r_p) = calc_r_theta(x_p,x_c,y_p,y_c)

    (F, M_0) = force_moment(delta_kappa,r_p,r_0,theta_p,theta_0,theta_contact,I_p = I_p, E = E)

    (F_ax, F_lat) = decompose_force(F,theta_f,theta_contact)

    (M_0,F_ax,F_lat,F,theta_contact)
end

#=
Calculate Force and Moment
delta_kappa - change in curvature from intrinsic
r_p - distance from point of contact to point p (pixels)
r_0 - distance from follicle to point of contact (pixels)
theta_p - angle from point_p to point of contact (radians)
theta_0 - angle from follicle to point of contact (radians)
theta_contact - angle that whisker is pointing at point of contact (radians)
=#
function force_moment(delta_kappa::Real,r_p::Real,r_0::Real,theta_p::Real,
    theta_0::Real,theta_contact::Real;I_p = 1.0, E=1.0)

    F = abs(delta_kappa * E * I_p / (r_p * cos(theta_p - theta_contact)))

    M_0 = r_0 * F * cos(theta_0 - theta_contact)

    (F, M_0)
end

#=
F - Total Contact Force
theta_f - whisker angle (radians)
theta_contact - angle of contact (radians)
=#
function decompose_force(F::Real,theta_f::Real,theta_contact::Real)
    F_ax = F * sin(theta_f - theta_contact)
    F_lat = F * cos(theta_f - theta_contact)
    (F_ax,F_lat)
end

function calc_r_theta(x1::Real,x2::Real,y1::Real,y2::Real)

    theta_p = atan((y2 - y1),(x2 - x1))
    r_p = sqrt((x2 - x1)^2 + (y2 - y1)^2)
    (theta_p,r_p)
end

function contact_angle(x1::Real,x2::Real,y1::Real,y2::Real)
    atan(y2-y1,x2-x1)
end

#=
Finds
=#
function get_contact_ind(x::Array{T,1},y::Array{T,1},ii,s_thres = 10.0,forward_dir=true) where T
    s = 0.0

    x_0 = x[ii]
    y_0 = y[ii]

    i = ii
    while(s < s_thres)

        s = sqrt((x[i] - x_0)^2 + (y[i] - y_0)^2)

        if forward_dir
            i = i + 1
        else
            i = i - 1
        end
    end
    i
end

#=
Positive sign = retraction
negative sign = protraction
=#
function get_theta_n(c,v::Array{Float64,1},f_t::Array{Float64,1},t_c::Array{Float64,1},t=trues(length(f_t)))

    c_sign = zeros(Float64,length(v))
    t_n = zeros(Float64,length(v))

    i=1
    while (i < length(c_sign))

        if c[i]
            j = findnext(.!c,i)

            rr = i:(j-1)
            rr = rr[t[i:(j-1)]]

            c_sign[rr] .= sign(mean(v[rr]))

            for ii in rr
                if c_sign[ii] > 0.0 #retraction
                    t_n[ii] = t_c[ii] + pi/2
                else  #protraction
                    t_n[ii] = t_c[ii] - pi/2
                end
            end

            i = j
        else
           i = i + 1
        end
    end

    (c_sign, t_n)
end

#=

=#
function calculate_all_forces(xx,yy,p,c,aa::Array{Float64,1},
    curv::Array{Float64,1},tracked=trues(length(c)); i_p_loc=50.0,E_in=1.0,moment_of_inertia_p=1.0)

    F_ax=zeros(Float64,length(c))
    F_lat=zeros(Float64,length(c))
    M=zeros(Float64,length(c))
    F_t=zeros(Float64,length(c))
    theta_c = zeros(Float64,length(c))
    F_calc=falses(length(c))

    for i=1:length(c)
        if ((c[i])&(length(xx[i])>1))&(tracked[i])

            (new_x, new_y) = interpolate_whisker(xx[i],yy[i])

            #ii - index of contact
            ii=calc_p_dist(new_x,new_y,p[i,1],p[i,2])[2]

            #i_p - index of high SNR point
            i_p=culm_dist(new_x,new_y,i_p_loc)

            try
                (M[i],F_ax[i],F_lat[i],F_t[i],theta_c[i])=calc_force(new_x,new_y,aa[i],curv[i],ii,i_p,E=E_in,I_p=moment_of_inertia_p)
                F_calc[i] = true
                !isfinite(F_t[i]) && (F_t[i] = 0.0; F_calc[i] = false)
                !isfinite(F_lat[i]) && (F_lat[i] = 0.0; F_calc[i] = false)
                !isfinite(F_ax[i]) && (F_ax[i] = 0.0; F_calc[i] = false)
                !isfinite(M[i]) && (M[i] = 0.0; F_calc[i] = false)
            catch
            end
        end
    end

    (F_ax,F_lat,M,F_t,theta_c,F_calc)
end


text:
ENV["PYTHON"] = ""
using Pkg
Pkg.build("PyCall")

using Conda
pipcmd = joinpath(Conda.PYTHONDIR,"pip")
run(`$pipcmd install pyneid`)

#=
env = Conda.ROOTENV
Conda.pip_interop(true, env)
Conda.pip("install", "pyneid")
=#


text:
using QuantumOptics

@testset "QuantumOptics" begin

    Random.seed!(0)
    b = SpinBasis(1//2)
    psi0 = spinup(b)
    psi1 = spindown(b)

    @test sprint((io, x) -> dirac(io, x, header=true), psi0) == "Ket(dim=2)\n  basis: Spin(1/2)\n|ψ⟩ = |0⟩\n"
    @test sprint((io, x) -> dirac(io, x, header=true), psi1) == "Ket(dim=2)\n  basis: Spin(1/2)\n|ψ⟩ = |1⟩\n"
    @test sprint((io, x) -> dirac(io, x, header=true), psi0 ⊗ psi1) == "Ket(dim=4)\n  basis: [Spin(1/2) ⊗ Spin(1/2)]\n|ψ⟩ = |10⟩\n"
    @test sprint((io, x) -> dirac(io, x, header=true), psi1 ⊗ psi0) == "Ket(dim=4)\n  basis: [Spin(1/2) ⊗ Spin(1/2)]\n|ψ⟩ = |01⟩\n"
    @test sprint((io, x) -> dirac(io, x, header=false), psi0) == "|ψ⟩ = |0⟩\n"
    @test sprint((io, x) -> dirac(io, x, header=false), psi1) == "|ψ⟩ = |1⟩\n"
    @test sprint((io, x) -> dirac(io, x, header=false), psi0 ⊗ psi1) == "|ψ⟩ = |10⟩\n"
    @test sprint((io, x) -> dirac(io, x, header=false), psi1 ⊗ psi0) == "|ψ⟩ = |01⟩\n"

    @show psi = randstate(b)
    @show psi ⊗ dagger(psi)
    @test sprint((io, x) -> dirac(io, x, header=true), psi) == "Ket(dim=2)\n  basis: Spin(1/2)\n|ψ⟩ = (0.65825+0.727547im)|0⟩+(0.131519+0.141719im)|1⟩\n"
    @test sprint((io, x) -> dirac(io, x, header=true), psi ⊗ dagger(psi)) == "DenseOperator(dim=2x2)\n  basis: Spin(1/2)\nρ = 0.962618|0⟩⟨0|+(0.18968+0.00239967im)|0⟩⟨1|+(0.18968-0.00239967im)|1⟩⟨0|+0.0373817|1⟩⟨1|\n"
    @test sprint((io, x) -> dirac(io, x, header=true), sparse( psi ⊗ dagger(psi)) ) == "SparseOperator(dim=2x2)\n  basis: Spin(1/2)\nρ = 0.962618|0⟩⟨0|+(0.18968+0.00239967im)|0⟩⟨1|+(0.18968-0.00239967im)|1⟩⟨0|+0.0373817|1⟩⟨1|\n"
    @test sprint((io, x) -> dirac(io, x, header=false), psi) == "|ψ⟩ = (0.65825+0.727547im)|0⟩+(0.131519+0.141719im)|1⟩\n"
    @test sprint((io, x) -> dirac(io, x, header=false), psi ⊗ dagger(psi)) == "ρ = 0.962618|0⟩⟨0|+(0.18968+0.00239967im)|0⟩⟨1|+(0.18968-0.00239967im)|1⟩⟨0|+0.0373817|1⟩⟨1|\n"
    @test sprint((io, x) -> dirac(io, x, header=false), sparse( psi ⊗ dagger(psi)) ) == "ρ = 0.962618|0⟩⟨0|+(0.18968+0.00239967im)|0⟩⟨1|+(0.18968-0.00239967im)|1⟩⟨0|+0.0373817|1⟩⟨1|\n"

end # "QuantumOptics"


text:
using Test
using LogicCircuits
using ProbabilisticCircuits
using Random: MersenneTwister
using CUDA: functional
using DataFrames
using Tables

include("../helper/gpu.jl")

function histogram_matches_likelihood(samples::Matrix{Bool}, worlds, loglikelihoods)
    hist = Dict{BitVector,Int}()
    for i = 1:size(samples,1)
        sample = BitVector(samples[i,:])
        hist[sample] = get(hist, sample, 0) + 1
    end
    for i = 1:size(worlds,1)
        exact_prob = exp(loglikelihoods[i])
        ex = BitVector(example(worlds,i))
        estim_prob = get(hist, ex, 0) / size(samples,1)
        @test exact_prob ≈ estim_prob atol=1e-2;
    end

end

@testset "Unconditional Sampling Test" begin

    rng = MersenneTwister(42)

    pc = zoo_psdd("little_4var.psdd");
    s, p = sample(pc;rng)
    @test pc(s...) ≈ p

    worlds = generate_data_all(num_variables(pc));

    loglikelihoods = EVI(pc, worlds)

    Nsamples = 50_000

    samples, _ = sample(pc, Nsamples; rng)
    histogram_matches_likelihood(samples, worlds, loglikelihoods)

    if CUDA.functional()
        samples, _ = sample(pc, Nsamples; rng, gpu = true)
        samples_cpu = to_cpu(samples)
        histogram_matches_likelihood(samples_cpu, worlds, loglikelihoods)
    end

end

@testset "Conditional Sampling Test" begin

    rng = MersenneTwister(42)
    num_samples = 10

    pc = zoo_psdd("little_4var.psdd");
    data_all = generate_data_all(num_variables(pc));

    # sampling given complete data should return same data with its log likelihood

    loglikelihoods = MAR(pc, data_all)
    sample_states, sample_prs = sample(pc, num_samples, data_all; rng)

    for i in 1:num_samples
        @test sample_states[i,:,:] == Tables.matrix(data_all)
        @test sample_prs[i,:] ≈ loglikelihoods atol=1e-6
    end

    # same states on CPU and GPU
    cpu_gpu_agree(data_all) do d
        sample(pc, num_samples, d)[1]
    end

    # same probabilities on CPU and GPU
    cpu_gpu_agree_approx(data_all) do d
        sample(pc, num_samples, d)[2]
    end

    # sampling given partial data invariants

    data_marg = DataFrame([false false false false;
                      false true true false;
                      false false true true;
                      false false false missing;
                      missing true false missing;
                      missing missing missing missing;
                      false missing missing missing], :auto)

    _, map_pr = MAP(pc, data_marg)

    sample_states, sample_prs = sample(pc, num_samples, data_marg; rng)

    for i in 1:num_samples

        # samples keep the partial evidence values
        pairs = collect(zip(sample_states[i,:,:], Tables.matrix(data_marg)))
        @test all(pairs) do (f,m)
            ismissing(m) || f == m
        end

        # probability does not exceed MAP probability
        @test all(sample_prs[i,:] .<= map_pr .+ 1e-6)
    end

end


text:
using JSON

module Bokehjs

typealias RealVect Union{AbstractVector{Int}, AbstractVector{Float64}}
typealias RealMat Union{AbstractMatrix{Int}, AbstractMatrix{Float64}}
typealias RealArray Union{RealMat, RealVect}
# would be nice to parameterize, but more important to constrain dims to 1 or 2
# typealias RealArray{N} Union{AbstractArray{Int, N}, AbstractArray{Float64, N}}

# like nothing except omitted from json rather than being null
type Omit
    v::AbstractString
    Omit() = new("__omitted from json__")
end
const omit = Omit()

# in case we want to restrict value types in future:
typealias BkAny Any # Union{Dict, Array, AbstractString, Number, Bool, Void, UUID}
typealias NullDict Union{Void, Dict{AbstractString, BkAny}}
typealias OmitDict Union{Omit, Dict{Symbol, BkAny}}
typealias NullString Union{Void, AbstractString}
typealias OmitString Union{Omit, AbstractString}
typealias NullSymbol Union{Void, Symbol}
typealias OmitSymbol Union{Omit, Symbol}
typealias NullFloat Union{Float64, Void}
typealias NullInt Union{Int, Void}

uuid4 = Base.Random.uuid4
UUID = Base.Random.UUID

abstract PlotObject

abstract BkRange <: PlotObject

typealias NullBkRange Union{Void, BkRange}

abstract Renderer <: PlotObject

abstract Axis <: PlotObject

type TypeID
    plotob::Union{PlotObject, Void}
end

function TypeID()
    TypeID(nothing)
end

type Plot <: PlotObject
    uuid::UUID
    title::AbstractString
    tools::Vector{BkAny}
    plot_height::Int
    plot_width::Int
    x_range::TypeID
    y_range::TypeID
    # could be Vector{TypeID}?:
    renderers::Vector{BkAny}
    above::Vector{TypeID}
    below::Vector{TypeID}
    left::Vector{TypeID}
    right::Vector{TypeID}
    data_sources::Vector{BkAny}
end

type ColumnDataSource <: PlotObject
    uuid::UUID
    column_names::Vector{Symbol}
    selected::Vector{Any}
    discrete_ranges::Dict{Symbol, BkAny}
    cont_ranges::Dict{Symbol, BkAny}
    data::Dict{Symbol, Vector}
end

function ColumnDataSource(data::Dict{Symbol, Vector})
    ColumnDataSource(uuid4(),
                     collect(keys(data)),
                     BkAny[],
                     Dict{Symbol, BkAny}(),
                     Dict{Symbol, BkAny}(),
                     data)
end

type DataRange1d <: BkRange
    uuid::UUID
    sources::Vector{BkAny}
end

function DataRange1d(cdss::Vector{ColumnDataSource}, columns::Vector{AbstractString})
    source(cds) = Dict{AbstractString, BkAny}("columns" => columns,
                                      "source" => TypeID(cds))
    sources = map(source, cdss)
    DataRange1d(uuid4(), sources)
end

type TickFormatter <: PlotObject
    uuid::UUID
    _type_name::Symbol
    format::OmitDict
    function TickFormatter(name::Symbol)
        @assert name in (:BasicTickFormatter, :DatetimeTickFormatter,
                         :LogTickFormatter)
        # format only seems to occur for DatetimeTickFormatter and even then is empty
        format = name == :DatetimeTickFormatter ? Dict{Symbol, BkAny}() : omit
        new(uuid4(), name, format)
    end
end

type Ticker <: PlotObject
    uuid::UUID
    _type_name::Symbol
    num_minor_ticks::Int64
    Ticker(name::Symbol) = new(uuid4(), name, 5)
end

type LinearAxis <: Axis
    uuid::UUID
    dimension::Int
    bounds::AbstractString
    location::AbstractString
    formatter::TypeID
    ticker::TypeID
    plot::TypeID
end

function LinearAxis(dimension::Int, tf::TickFormatter, t::Ticker, plot::Plot)
    LinearAxis(uuid4(),
               dimension,
               "auto",
               "min",
               TypeID(tf),
               TypeID(t),
               TypeID(plot))
end

type Grid <: Renderer
    uuid::UUID
    dimension::Int
    plot::TypeID
    ticker::TypeID
end

function Grid(dimension::Int, plot::Plot, ticker::Ticker)
    Grid(uuid4(), dimension, TypeID(plot), TypeID(ticker))
end

type Legend <: Renderer
    uuid::UUID
    plot::TypeID
    legends::Vector{Tuple}
    orientation::OmitSymbol
end

function Legend(plot::Plot, legends::Vector{Tuple})
    Legend(plot, legends, nothing)
end

function Legend(plot::Plot, legends::Vector{Tuple}, orientation::NullSymbol)
    orientation = orientation == nothing ? omit : orientation
    @assert orientation in (omit, :top_left, :top_center, :top_right,
                            :right_center, :bottom_right, :bottom_center,
                            :bottom_left, :left_center, :center)
    Legend(uuid4(),
           TypeID(plot),
           [(l, [TypeID(r)]) for (l, r) in legends],
           orientation)
end

type Glyph <: PlotObject
    uuid::UUID
    _type_name::Symbol
    line_color::OmitDict
    line_width::OmitDict
    line_alpha::OmitDict
    fill_color::OmitDict
    fill_alpha::OmitDict
    size::OmitDict
    line_dash::Union{Omit, Vector{Int64}}
    x::Dict{Symbol, Symbol}
    y::Dict{Symbol, Symbol}
end

type GlyphRenderer <: Renderer
    uuid::UUID
    data_source::TypeID
    nonselection_glyph::TypeID
    selection_glyph::TypeID
    glyph::TypeID
    name::Union{Void, AbstractString}
    server_data_source::NullDict
end

typealias NullGlyph Union{Void, Glyph}

function GlyphRenderer(coldata::ColumnDataSource, nonsel_g::NullGlyph,
                       sel_g::NullGlyph, glyph::Glyph)
    GlyphRenderer(uuid4(),
                  TypeID(coldata),
                  TypeID(nonsel_g),
                  TypeID(sel_g),
                  TypeID(glyph),
                  nothing,
                  nothing
                 )
end

type Metatool <: PlotObject
    uuid::UUID
    _type_name::AbstractString
    plot::TypeID
    dimensions::Union{Vector{AbstractString}, Void, Omit}
end

function Metatool(typename::AbstractString, plot::Plot, dimensions)
    plot = TypeID(plot)
    Metatool(uuid4(), typename, plot, dimensions)
end

function Metatool(typename::AbstractString, plot::Plot)
    Metatool(typename, plot, omit)
end

# TODO: these are duplicated as HEIGHT and WIDTH in Bokeh.jl
_DEFAULT_HEIGHT = 600
_DEFAULT_WIDTH = 600

function Plot()
    Plot(uuid4(),
         "",
         Void[],
         _DEFAULT_HEIGHT,
         _DEFAULT_WIDTH,
         TypeID(),
         TypeID(),
         Void[],

         Void[],
         Void[],
         Void[],
         Void[],

         Void[]
        )
end

function Plot(plot::Plot,
              xrange::BkRange,
              yrange::BkRange,
              renderers::Array{PlotObject,1},
              axes,#::Dict{Symbol, Array{PlotObject,1}},
              tools::Array{PlotObject,1},
              title::AbstractString="Bokeh Plot",
              height::Int=_DEFAULT_HEIGHT,
              width::Int=_DEFAULT_WIDTH)
    data_sources = BkAny[]# [TypeID(coldata)]

    Plot(plot.uuid,
         title,
         map(TypeID, tools),
         height,
         width,
         TypeID(xrange),
         TypeID(yrange),
         map(TypeID, renderers),
         map(TypeID, axes[:above]),
         map(TypeID, axes[:below]),
         map(TypeID, axes[:left]),
         map(TypeID, axes[:right]),
         data_sources
        )
end

type PlotContext <: PlotObject
    uuid::UUID
    children::Vector{TypeID}
end

function PlotContext(plot::Plot)
    PlotContext(uuid4(),[TypeID(plot)])
end

end  # module

# extract useful types from Bokehjs module
typealias RealVect Bokehjs.RealVect
typealias RealMat Bokehjs.RealMat
typealias RealArray Bokehjs.RealArray
typealias BkAny Bokehjs.BkAny
typealias omit Bokehjs.omit
typealias Glyph Bokehjs.Glyph
typealias NullString Bokehjs.NullString
typealias NullSymbol Bokehjs.NullSymbol
typealias NullFloat Bokehjs.NullFloat
typealias NullInt Bokehjs.NullInt

JSON.lower(uuid::Bokehjs.UUID) = string(uuid)

function JSON.lower(tid::Bokehjs.TypeID)
    tid.plotob == nothing && return nothing
    attrs = fieldnames(tid.plotob)
    obtype = in(:_type_name, attrs) ? tid.plotob._type_name : typeof(tid.plotob)
    Dict{AbstractString, BkAny}("type" => obtype, "id" => tid.plotob.uuid)
end

JSON.lower{T<:Bokehjs.PlotObject}(::Type{T}) = string(T.name.name)


text:
using HorizonSideRobots

function corner(r::Robot)
    s=0::Int
    n=0::Int
    o=0::Int
    w=0::Int
    while !isborder(r,Sud)
        move!(r,Sud)
        s+=1
    end
    while !isborder(r,West)
        move!(r,West)
        w+=1
    end
    paintover(r)
    while !isborder(r,Sud)
        move!(r,Sud)
    end
    while !isborder(r,West)
        move!(r,West)
    end
    for o in 1:w
        move!(r,Ost)
    end
    for n in 1:s
        move!(r,Nord)
    end
end

function paintover(r)
    for side in [Nord,Ost,Sud,West]
        while (isborder(r,side)==0)
            move!(r,side)
            putmarker!(r)
        end
    putmarker!(r)
    end
end

corner(r)
show(r)

max_stars_repo_name: Cheburkova/mirea-progs

max_issues_repo_name: Cheburkova/mirea-progs

max_forks_repo_name: Cheburkova/mirea-progs


text:
using DirectConvolution
using Test
using LinearAlgebra

@testset "DirectConvolution" begin

    include("utils.jl")

    include("linearFilter.jl")

    include("directConvolution.jl")

    include("SG_Filter.jl")

    include("udwt.jl")

end;
nothing

max_stars_repo_name: vincent-picaud/DirectConvolution

max_issues_repo_name: vincent-picaud/DirectConvolution

max_forks_repo_name: vincent-picaud/DirectConvolution


text:
# Patternshop is a sort of Photoshop for exploring Turing patterns.
# The name is temporary until I think of something I like better.

using TuringPatterns
using HTTP, JSON

write2d(stream, frame) = save(FileIO.DataFormat{:PNG}, stream, Gray.(frame))

# Parameter sweeping
swept(v, x, y) = v
swept(v::Dict, x, y) = haskey(v, "xsweep") ? v["xsweep"][x] : v["ysweep"][y]
# todo: can we use `pairs` in 1.0?
sweep!(d::Dict, x, y) = for (key, val) in d
    d[key] = swept(val, x, y)
end
sweep!(d::Array, x, y) = for (key, val) in enumerate(d)
    d[key] = swept(val, x, y)
end

function handler(req, res)
    params = req |> HTTP.uri |> HTTP.query |> HTTP.unescape |> JSON.parse

    println("Handling request:")
    @show params

    ix, iy = params["ix"], params["iy"]
    sweep!(params, ix, iy)

    width, height, iters, dims, seed = params["width"], params["height"], params["iters"], params["dims"], params["seed"]
    sz = (width, height)
    patterns = []

    for spec in params["patterns"]
        sweep!(spec, ix, iy)
        if spec["type"] == "simple"
            push!(patterns, SimplePattern(Params(spec["r"], spec["R"], spec["effectsize"]), sz))
        else # composite
            sweep!(spec["couplings"], ix, iy)
            subpatterns = SimplePattern{Int(dims)}[]
            for subspec in spec["patterns"]
                sweep!(subspec, ix, iy)
                push!(subpatterns, SimplePattern(Params(subspec["r"], subspec["R"], subspec["effectsize"]), sz))
            end
            len = length(subpatterns)
            couplings = convert(Matrix{Float64}, reshape(spec["couplings"], (len, len)))
            push!(patterns, CompositePattern(subpatterns, couplings))
        end
    end

    srand(seed)

    frame = simulate(
        rand(sz),
        collect(promote(patterns...)),
        BoxBlur(sz),
        iters
    )

    buf = FIFOBuffer()
    write2d(buf, frame)
    close(buf)

    HTTP.Response(
      headers=HTTP.Headers(
        "Access-Control-Allow-Origin" => "*",
        "Content-Type" => "image/png",
      ),
      body=buf
    )
end

function serve()
    server = HTTP.Server(handler, DevNull, ratelimit=100//1)
    HTTP.serve(server, IPv4(127,0,0,1), 8002)
end

serve()


text:
JSON.lower(q::Query) = non_nothing_dict(q)
query_type(q::Query) = q.queryType
JSON.lower(ls::LimitSpec) = non_nothing_dict(ls)
Tables.istable(::QueryResult) = true

include("timeseries.jl")
include("topn.jl")
include("groupby.jl")
include("scan.jl")
include("search.jl")
include("timeboundary.jl")
include("segment_metadata.jl")
include("datasource_metadata.jl")
include("sql.jl")

"""
    execute(client, query; pretty=false)

Executes the native Druid `query` on the `client::Client`. Returns the query
results as a String. Throws exception if query execution fails.
"""
execute(client::Client, query::Query; pretty=false) = execute_native_query(client, query; pretty)

function execute_native_query(client::Client, query::Query; pretty=false)
    url = joinpath(client.url, "druid/v2")
    if pretty
        url = merge(url, query="pretty")
    end
    do_query(url, query)
end

function do_query(url, post_data)
    local response
    try
        response = request("POST", url, ["Content-Type" => "application/json"], json(post_data))
    catch err
        if isa(err, IOError)
            error("I/O Error during query execution")
        elseif isa(err, StatusError)
            local err_message
            try
                err_message = JSON.parse(String(err.response.body))
            catch
                error("Druid error during query execution, error message unavailable\n", "Error status: $(err.status)")
            end
            error("Druid error during query execution\n", err_message, "\nError status: $(err.status)")
        end
    end
    String(response.body)
end

function Base.show(io::IO, q::Query)
    print(io, query_type(q), '(')
    firstprint = true
    for fname in propertynames(q)
        val = getproperty(q, fname)
        if !(val === nothing)
            firstprint || print(io, ", ")
            if firstprint
                firstprint = false
            end
            print(io, fname, '=', repr(val, context=io))
        end
    end
    print(io, ')')
end


text:
const REQUIRED_ARGS_INI = ["resistance_file",
                           "radius",
                           "project_name"]

const REQUIRED_ARGS_DICT = ["radius",
                            "project_name"]

const TRUELIST = ["true", "True", "TRUE"]

const SINGLE = ["single", "Single", "SINGLE"]

const SOLVERS = ["cg+amg", "cholmod"]

const SUPPORTED_ARGS = ["resistance_file", "resistance_is_conductance",
    "source_file", "project_name", "parallelize", "parallel_batch_size",
    "block_size", "radius", "buffer", "source_threshold",
    "source_from_resistance", "r_cutoff", "precision",
    "connect_four_neighbors_only", "solver", "calc_flow_potential",
    "calc_normalized_current", "correct_artifacts", "write_raw_currmap",
    "write_as_tif", "mask_nodata", "suppress_cs_messages", "conditional",
    "n_conditions", "compare_to_future", "condition1_file", "condition2_file",
    "condition1_future_file", "condition2_future_file", "comparison1",
    "comparison2", "condition1_lower", "condition2_lower", "condition1_upper",
    "condition2_upper", "reclassify_resistance", "reclass_table",
    "write_reclassified_resistance", "allow_different_projections"]


text:
module CellValuesMocksTests

using Test
using CellwiseValues
using StaticArrays
using TensorValues

using ..CellValuesMocks

l = 10

for a in [1.2, SVector(1.3,2.0), VectorValue(1.3,2.0)]

  v = TestIterCellValue(a,l)

  w = [a for i in 1:l]
  test_iter_cell_value(v,w)

  v = TestIndexCellValue(a,l)
  test_index_cell_value(v,w)

end

a = 1.2
v = TestIndexCellValue(a,l)

s = string(v)
z = """
1 -> 1.2
2 -> 1.2
3 -> 1.2
4 -> 1.2
5 -> 1.2
6 -> 1.2
7 -> 1.2
8 -> 1.2
9 -> 1.2
10 -> 1.2
"""
@test s == z

end # module


text:
module VariableTemplates

export varsize, Vars, Grad, @vars

using StaticArrays

"""
    varsize(S)

The number of elements specified by the template type `S`.
"""
varsize(::Type{T}) where {T<:Real} = 1
varsize(::Type{Tuple{}}) = 0
varsize(::Type{NamedTuple{(),Tuple{}}}) = 0
varsize(::Type{SVector{N,T}}) where {N,T<:Real} = N

include("var_names.jl")

# TODO: should be possible to get rid of @generated
@generated function varsize(::Type{S}) where {S}
  types = fieldtypes(S)
  isempty(types) ? 0 : sum(varsize, types)
end

function process_vars!(syms, typs, expr)
  if expr isa LineNumberNode
     return
  elseif expr isa Expr && expr.head == :block
    for arg in expr.args
      process_vars!(syms, typs, arg)
    end
    return
  elseif expr.head == :(::)
    push!(syms, expr.args[1])
    push!(typs, expr.args[2])
    return
  else
    error("Invalid expression")
  end
end

macro vars(args...)
  syms = Any[]
  typs = Any[]
  for arg in args
    process_vars!(syms, typs, arg)
  end
  :(NamedTuple{$(tuple(syms...)), Tuple{$(esc.(typs)...)}})
end

struct GetVarError <: Exception
  sym::Symbol
end
struct SetVarError <: Exception
  sym::Symbol
end

"""
    Vars{S,A,offset}(array::A)

Defines property overloading for `array` using the type `S` as a template. `offset` is used to shift the starting element of the array.
"""
struct Vars{S,A,offset}
  array::A
end
Vars{S}(array) where {S} = Vars{S,typeof(array),0}(array)

Base.eltype(v::Vars) = eltype(getfield(v,:array))
Base.propertynames(::Vars{S}) where {S} = fieldnames(S)
Base.similar(v::Vars{S,A,offset}) where {S,A,offset} = Vars{S,A,offset}(similar(getfield(v,:array)))

@generated function Base.getproperty(v::Vars{S,A,offset}, sym::Symbol) where {S,A,offset}
  expr = quote
    Base.@_inline_meta
    array = getfield(v, :array)
  end
  for k in fieldnames(S)
    T = fieldtype(S,k)
    if T <: Real
      retexpr = :($T(array[$(offset+1)]))
      offset += 1
    elseif T <: SHermitianCompact
      LT = StaticArrays.lowertriangletype(T)
      N = length(LT)
      retexpr = :($T($LT($([:(array[$(offset + i)]) for i = 1:N]...))))
      offset += N
    elseif T <: StaticArray
      N = length(T)
      retexpr = :($T($([:(array[$(offset + i)]) for i = 1:N]...)))
      offset += N
    else
      retexpr = :(Vars{$T,A,$offset}(array))
      offset += varsize(T)
    end
    push!(expr.args, :(if sym == $(QuoteNode(k))
      return @inbounds $retexpr
    end))
  end
  push!(expr.args, :(throw(GetVarError(sym))))
  expr
end

@generated function Base.setproperty!(v::Vars{S,A,offset}, sym::Symbol, val) where {S,A,offset}
  expr = quote
    Base.@_inline_meta
    array = getfield(v, :array)
  end
  for k in fieldnames(S)
    T = fieldtype(S,k)
    if T <: Real
      retexpr = :(array[$(offset+1)] = val)
      offset += 1
    elseif T <: SHermitianCompact
      LT = StaticArrays.lowertriangletype(T)
      N = length(LT)
      retexpr = :(array[$(offset + 1):$(offset + N)] .= $T(val).lowertriangle)
      offset += N
    elseif T <: StaticArray
      N = length(T)
      retexpr = :(array[$(offset + 1):$(offset + N)] .= val[:])
      offset += N
    else
      offset += varsize(T)
      continue
    end
    push!(expr.args, :(if sym == $(QuoteNode(k))
      return @inbounds $retexpr
    end))
  end
  push!(expr.args, :(throw(SetVarError(sym))))
  expr
end

"""
    Grad{S,A,offset}(array::A)

Defines property overloading along slices of the second dimension of `array` using the type `S` as a template. `offset` is used to shift the starting element of the array.
"""
struct Grad{S,A,offset}
  array::A
end
Grad{S}(array) where {S} = Grad{S,typeof(array),0}(array)

Base.eltype(g::Grad) = eltype(getfield(g,:array))
Base.propertynames(::Grad{S}) where {S} = fieldnames(S)
Base.similar(g::Grad{S,A,offset}) where {S,A,offset} = Grad{S,A,offset}(similar(getfield(g,:array)))

@generated function Base.getproperty(v::Grad{S,A,offset}, sym::Symbol) where {S,A,offset}
  M = size(A,1)
  expr = quote
    Base.@_inline_meta
    array = getfield(v, :array)
  end
  for k in fieldnames(S)
    T = fieldtype(S,k)
    if T <: Real
      retexpr = :(SVector{$M,$T}($([:(array[$i,$(offset+1)]) for i = 1:M]...)))
      offset += 1
    elseif T <: StaticArray
      N = length(T)
      retexpr = :(SMatrix{$M,$N,$(eltype(T))}($([:(array[$i,$(offset + j)]) for i = 1:M, j = 1:N]...)))
      offset += N
    else
      retexpr = :(Grad{$T,A,$offset}(array))
      offset += varsize(T)
    end
    push!(expr.args, :(if sym == $(QuoteNode(k))
      return @inbounds $retexpr
    end))
  end
  push!(expr.args, :(throw(GetVarError(sym))))
  expr
end

@generated function Base.setproperty!(v::Grad{S,A,offset}, sym::Symbol, val) where {S,A,offset}
  M = size(A,1)
  expr = quote
    Base.@_inline_meta
    array = getfield(v, :array)
  end
  for k in fieldnames(S)
    T = fieldtype(S,k)
    if T <: Real
      retexpr = :(array[:, $(offset+1)] .= val)
      offset += 1
    elseif T <: StaticArray
      N = length(T)
      retexpr = :(array[:, $(offset + 1):$(offset + N)] .= val)
      offset += N
    else
      offset += varsize(T)
      continue
    end
    push!(expr.args, :(if sym == $(QuoteNode(k))
      return @inbounds $retexpr
    end))
  end
  push!(expr.args, :(throw(SetVarError(sym))))
  expr
end

end # module

max_stars_repo_name: akshaysridhar/CLIMA

max_issues_repo_name: akshaysridhar/CLIMA

max_forks_repo_name: akshaysridhar/CLIMA


text:
"""
Return the features of `e` ranked by decreasing
median absolute deviation. Trim to the
first `top` features if a value is provided.
"""
function most_variable_features(e::Experiment; top::Int64 = 0)
    e_mad_ind = e |>
                 getdata |>
                 eachcol |>
                 x -> mad.(x, normalize = true) |>
                 sortperm |>
                 reverse

    # Get symbols from indices
    e_mad_sym = names(getdata(e))[e_mad_ind]

    # Truncate if needed
    if (top > 0) && (length(e_mad_sym) > top)
        e_mad_sym = e_mad_sym[1:top]
    end

    return(e_mad_sym)
end

"""
Return (all or if provided the `top`) features
varying the most in `e` (largest absolute log
fold change), when comparing entries matching
filters `f1` and `f2`. Columns for which the
fold change is negative come last.
"""
function characteristic_features(e::Experiment,
                                 f1::AbstractFilter,
                                 f2::AbstractFilter;
                                 top::Int64 = 0)
    f1_col = filter_entries(e,f1)
    f2_col = filter_entries(e,f2)

    lfc_ind = e.data[:, e.selected_features] |>
         eachcol |>
         y -> map(x -> mean(x[f1_col]) / mean(x[f2_col]), y) |>
         y -> map(x -> x <= 0 ? 0 : abs(log2(x)), y) |>
         sortperm |>
         reverse

    # Get symbols from indices
    sym = names(getdata(e))[lfc_ind]

    # Truncate if needed
    if (top > 0) && (length(sym) > top)
        sym = sym[1:top]
    end

    return(sym)
end

"""
Return (all or if provided the `top`) features
in `e` associated the most with `ref` (absolute
Pearson correlation).
"""
function most_correlated(e::Experiment,
                         ref::AbstractVector;
                         top::Int64 = 0)
    @assert all( [x <: Number for x in eltype.(eachcol(getdata(e)))] )
    mostcor_ind = e |> getdata |>
                   x -> cor(ref, Array(x)) |>
                   x -> abs.(x) |>
                   x -> sortperm([x...]) |>
                   reverse

    # Get symbols from indices
    mostcor = names(getdata(e))[mostcor_ind]

    # Truncate if needed
    if (top > 0) && (length(mostcor) > top)
        mostcor = mostcor[1:top]
    end

    return(mostcor)
end

function most_correlated(e::Experiment,
                         ref::Symbol;
                         top::Int64 = 0)
    ref_vector = e.data[e.selected_entries,ref]
    most_correlated(e,ref_vector,top = top)
end


text:
@testset "Interface Definition" begin

    @test_throws ErrorException Foo()

    a = Foo(fieldA=1, fieldB="A")

    @test a.fieldA == 1
    @test a.fieldB == "A"
    @test a.fieldC === missing
    @test a.fieldD === missing

    b = Foo(fieldA=1, fieldB="A", fieldC="B", fieldD="C")

    @test b.fieldA == 1
    @test b.fieldB == "A"
    @test b.fieldC == "B"
    @test b.fieldD == "C"

    @test Foo(JSON.parse(JSON.json(a))) == a
    @test Foo(JSON.parse(JSON.json(b))) == b

    c = Foo2(fieldA=nothing, fieldB=[1,2])

    @test c.fieldA === nothing
    @test c.fieldB == [1,2]
    @test Foo2(JSON.parse(JSON.json(c))) == c

    d = Foo2(fieldA=3, fieldB=[1,2])
    @test d.fieldA === 3
    @test d.fieldB == [1,2]
    @test Foo2(JSON.parse(JSON.json(d))) == d
end


text:
# %%
import Pkg;
Pkg.activate(".")
# %%
path = pwd() * "/functions"
push!(LOAD_PATH,path)
# %%
using Revise
using JLD2, FileIO
using CompStat
using CustomTypes
# %%
JLD2.@load "output/calibrations.jld2"
# %%
using Plots
pyplot()
# %%
plot_cmpstat_bw(cal_y, "y")
# %%
plot_cmpstat_bw(cal_m, "m", x1 = [0.9, 1.3])
# %%
plot_cmpstat_bw(cal_c, "c")
# %%
plot_cmpstat_bw(cal_ssξ, "ssxi")
# %%
plot_cmpstat_bw(cal_u, "u", x0=[0.9, 6.95], x1=[1.33, 6.95])
# %%
plot_cmpstat_colors(cal_y, "y")
plot_cmpstat_colors(cal_m, "m", x1 = [0.9, 1.3])
plot_cmpstat_colors(cal_c, "c")
plot_cmpstat_colors(cal_ssξ, "ssxi")
plot_cmpstat_colors(cal_u, "u", x0=[0.9, 6.95], x1=[1.33, 6.95])
# %%
plot_transitions(cal_y.F*0.5, cal_y, "y")
# %%


text:
# ALL TESTS THAT CONCERN MODFICIATION

# include
using LatticePhysics
using PyPlot
using Base.Test

# begin the Modification testblock
fermi_surface_testset = @testset "Fermi surfaces" begin

################################################################################
#
#   METHODS FOR CONSTRUCTION OF BANDSTRUCTURES (ALONG PATHS)
#
################################################################################

################################################################################
#
#   1) CALCULATION OF FERMI SURFACE
#
#   2) PLOTTING OF FERMI SURFACE
#       - plotting from points
#       - plotting from unitcell
#
################################################################################

# begin the testset
@testset "2D" begin

    # begin the testset only for calculations
    @testset "honeycomb (calculations)" begin

        # get a unitcell
        unitcell = getUnitcellHoneycomb()

        # test the type of the Fermi surface calculations (bare)
        @test typeof(getFermiSurface(unitcell, 20)) == Array{Float64,2}
        @test size(getFermiSurface(unitcell, 20)) == (20,2)

        # test if modifying the fermi energy is possible
        @test typeof(getFermiSurface(unitcell, 20, fermi_energy=1.0)) == Array{Float64,2}

        # test more options
        @test typeof(getFermiSurface(unitcell, 20, enforce_hermitian=true)) == Array{Float64,2}
        @test typeof(getFermiSurface(unitcell, 20, refold_to_first_BZ=false)) == Array{Float64,2}

    # end the testset here
    end;

    # begin the testset for plotting
    @testset "honeycomb (plotting)" begin

        # get a unitcell
        unitcell = getUnitcellHoneycomb()

        # get the fermi surface
        fermi_surface = getFermiSurface(unitcell, 20)

        # test if modifying the fermi energy is possible
        @test_nowarn close(plotFermiSurface(fermi_surface, showPlot=false))
        @test_nowarn close(plotFermiSurface(fermi_surface, plot_title="myplot", showPlot=false))
        @test_nowarn close(plotFermiSurface(fermi_surface, plot_color="r", showPlot=false))
        @test_nowarn close(plotFermiSurface(fermi_surface, figsize=(10,8), showPlot=false))
        @test_nowarn close(plotFermiSurface(fermi_surface, save_filename="$(FOLDER_SPECTRA)testplot.pdf", showPlot=false))

        # test calculation at the same time
        @test_nowarn close(plotFermiSurface(unitcell, 20, showPlot=false))

    # end the testset here
    end;

    # begin the testset
    @testset "square (calculations)" begin

        # get a unitcell
        unitcell = getUnitcellSquare()

        # test the type of the Fermi surface calculations (bare)
        @test typeof(getFermiSurface(unitcell, 20)) == Array{Float64,2}
        @test size(getFermiSurface(unitcell, 20)) == (20,2)

    # end the testset here
    end;

    # begin the testset for plotting
    @testset "square (plotting)" begin

        # get a unitcell
        unitcell = getUnitcellSquare()

        # get the fermi surface
        fermi_surface = getFermiSurface(unitcell, 20)

        # test if modifying the fermi energy is possible
        @test_nowarn close(plotFermiSurface(fermi_surface, showPlot=false))
        @test_nowarn close(plotFermiSurface(fermi_surface, plot_title="myplot", showPlot=false))
        @test_nowarn close(plotFermiSurface(fermi_surface, plot_color="r", showPlot=false))
        @test_nowarn close(plotFermiSurface(fermi_surface, figsize=(10,8), showPlot=false))
        @test_nowarn close(plotFermiSurface(fermi_surface, save_filename="$(FOLDER_SPECTRA)testplot.pdf", showPlot=false))

        # test calculation at the same time
        @test_nowarn close(plotFermiSurface(unitcell, 20, showPlot=false))

    # end the testset here
    end;

# end the testset here
end;

# begin the testset
@testset "3D" begin

    # begin the testset only for calculations
    @testset "diamond (calculations)" begin

        # get a unitcell
        unitcell = getUnitcellDiamond()

        # test the type of the Fermi surface calculations (bare)
        @test typeof(getFermiSurface(unitcell, 20)) == Array{Float64,2}
        @test size(getFermiSurface(unitcell, 20)) == (20,3)

        # test if modifying the fermi energy is possible
        @test typeof(getFermiSurface(unitcell, 20, fermi_energy=1.0)) == Array{Float64,2}

        # test more options
        @test typeof(getFermiSurface(unitcell, 20, enforce_hermitian=true)) == Array{Float64,2}
        @test typeof(getFermiSurface(unitcell, 20, refold_to_first_BZ=false)) == Array{Float64,2}

    # end the testset here
    end;

    # begin the testset for plotting
    @testset "diamond (plotting)" begin

        # get a unitcell
        unitcell = getUnitcellDiamond()

        # get the fermi surface
        fermi_surface = getFermiSurface(unitcell, 20)

        # test if modifying the fermi energy is possible
        @test_nowarn close(plotFermiSurface(fermi_surface, showPlot=false))
        @test_nowarn close(plotFermiSurface(fermi_surface, plot_title="myplot", showPlot=false))
        @test_nowarn close(plotFermiSurface(fermi_surface, plot_color="r", showPlot=false))
        @test_nowarn close(plotFermiSurface(fermi_surface, figsize=(10,8), showPlot=false))
        @test_nowarn close(plotFermiSurface(fermi_surface, save_filename="$(FOLDER_SPECTRA)testplot.pdf", showPlot=false))

        # test calculation at the same time
        @test_nowarn close(plotFermiSurface(unitcell, 20, showPlot=false))

    # end the testset here
    end;

# end the testset here
end;

# end the testset here
end;