text:
using Inflation
using Printf
using Test

@testset "convert_01tuples" begin
	# Test 1
	str  = "0xx"
	char = 'x'
	exp_output = ["000", "001", "010", "011"]
	act_output = convert_01tuples(str, char)
	@test exp_output == act_output

	# Test 2
	str  = "01x"
	char = 'x'
	exp_output = ["010", "011"]
	act_output = convert_01tuples(str, char)
	@test exp_output == act_output

	# Test 3
	str  = "y1x0"
	char = 'y'
	exp_output = ["01x0", "11x0"]
	act_output = convert_01tuples(str, char)
	@test exp_output == act_output
end;

@testset "rows_and_cols" begin
	# Test 1
	infl_strings = ["yyx", "yxy", "xyy"]
	(a, b) = rows_and_cols(infl_strings)
	@test (a, b) == (12, 8)
end;

@testset "get_M_matrix" begin
	# Test 1
	infl_strings = ["yyx", "yxy", "xyy"]
	exp_output   = [1 1 0 0 0 0 0 0;
					0 0 1 1 0 0 0 0;
					0 0 0 0 1 1 0 0;
					0 0 0 0 0 0 1 1;
					1 0 1 0 0 0 0 0;
					0 1 0 1 0 0 0 0;
					0 0 0 0 1 0 1 0;
					0 0 0 0 0 1 0 1;
					1 0 0 0 1 0 0 0;
					0 1 0 0 0 1 0 0;
					0 0 1 0 0 0 1 0;
					0 0 0 1 0 0 0 1]
	act_output = get_M_matrix(infl_strings)
	@test exp_output == act_output
end;

@testset "pfunc_series" begin
	# Testing pfunc1
	@test pfunc1(0, 0, 0) == 0.5
	@test pfunc1(0, 0, 1) == 0.0
	@test pfunc1(0, 1, 0) == 0.0
	@test pfunc1(0, 1, 1) == 0.0
	@test pfunc1(1, 0, 0) == 0.0
	@test pfunc1(1, 0, 1) == 0.0
	@test pfunc1(1, 1, 0) == 0.0
	@test pfunc1(1, 1, 1) == 0.5

	@test pfunc1("000") == 0.5
	@test pfunc1("001") == 0.0

	# Testing pfunc2
	@test pfunc2(0, 0, 0) == 0.0
	@test pfunc2(0, 0, 1) == 1.0/3.0
	@test pfunc2(0, 1, 0) == 1.0/3.0
	@test pfunc2(0, 1, 1) == 0.0
	@test pfunc2(1, 0, 0) == 1.0/3.0
	@test pfunc2(1, 0, 1) == 0.0
	@test pfunc2(1, 1, 0) == 0.0
	@test pfunc2(1, 1, 1) == 0.0

	@test pfunc2("000") == 0.0
	@test pfunc2("001") == 1.0/3.0

	# Testing pfunct
	pfunct = get_pfunct(0.5)
	@test pfunct("000") == 1.0/4.0
	@test pfunct("100") == 1.0/12.0
end;

@testset "integrate_pfunc" begin
	# Test 1 : pfunc1(0, 0, 0) + pfunc1(0, 0, 1) == 0.5 + 0.0 == 0.5
	@test integrate_pfunc(pfunc1, "00x", 'x') == 0.5

	# Test 2 : pfunc1(0, 1, 1) + pfunc1(0, 0, 1) == 0.0 + 0.0 == 0.0
	@test integrate_pfunc(pfunc1, "0x1", 'x') == 0.0

	# Test 3 :   pfunc1(0, 1, 1) + pfunc1(0, 0, 1)
	#        : + pfunc1(0, 0, 0) + pfunc1(0, 1, 0)
	#        : = 0.0 + 0.0 + 0.5 + 0.0 == 0.5
	@test integrate_pfunc(pfunc1, "0xx", 'x') == 0.5
end;

@testset "spiral_inflation_test" begin
	# The spiral inflation for the triangle causal model has a very
	# large M-matrix and b-array, so we can test a few terms

	infl_strings = ["yxyxyx", "yxxyxy", "xyyxxy", "xyxyyx", "xyxyxy"]

	# creating the M-matrix and b-array to be used in the Convex.jl optimization
	M = get_M_matrix(infl_strings)
	b = get_b_array(infl_strings, pfunc2, spiral)

	############################################################################
	# P_{A1,B1,C1}(0, 0, 0) term?
	############################################################################
	# What is the index for this term?
	# index :   8 * [(position in infl_strings) - 1] + [binary('000') + 1]
	# index : = 8*(0) + 1 = 1

	# What does the row in the M-matrix look like?
	# infl_string = '0x0x0x'
	# length = 2^6 = 64
	# '000000' => 1      '000001' => 2      '000100' => 5      '000101' => 6
	# '010000' => 17     '010001' => 18     '010100' => 21     '010101' => 22
	Mvec000 = zeros(Int64, 64)
	Mvec000[1]  = 1  ;  Mvec000[2]  = 1  ;  Mvec000[5]  = 1  ;  Mvec000[6]  = 1
	Mvec000[17] = 1  ;  Mvec000[18] = 1  ;  Mvec000[21] = 1  ;  Mvec000[22] = 1

	# What does the b-array element look like?
	# P_{A1,B1,C1}(0,0,0) = P_{A,B,C}(0,0,0) = 0.0
	b000 = 0.0

	# check if correct
	@test M[1,:] == Mvec000
	@test b[1] == b000

	############################################################################
	# P_{A2,B2,C1}(1, 0, 0) term?
	############################################################################
	# What is the index for this term?
	# index :   8 * [(position in infl_strings) - 1] + [binary('100') + 1]
	# index : = 8*(3) + 5 = 29

	# What does the row in the M-matrix look like?
	# infl_string = 'x1x00x'
	# length = 2^6 = 64
	# '010000' => 17     '010001' => 18     '011000' => 25     '011001' => 26
	# '110000' => 49     '110001' => 50     '111000' => 57     '111001' => 58
	Mvec100 = zeros(Int64, 64)
	Mvec100[17] = 1  ;  Mvec100[18] = 1  ;  Mvec100[25] = 1  ;  Mvec100[26] = 1
	Mvec100[49] = 1  ;  Mvec100[50] = 1  ;  Mvec100[57] = 1  ;  Mvec100[58] = 1

	# What does the b-array element look like?
	# P_{A1,B1,C1}(1,0,0) = P_{A,C}(1,0) P_{B}(0) = [0 + 1/3][0 + 1/3 + 1/3 + 0]
	#                     = 2/9 = 0.2222...
	b100 = 2.0/9.0

	@test M[29,:] == Mvec100
	@test b[29] == b100
end;


text:
# "Simple" matrix-vector multiplication, efficient for both dense and
# sparse matrices.
# Assumes A is efficiently indexed in column-major order.
function matvecmul!(dest::AbstractVector, A::AbstractMatrix, x::AbstractVector)
    fill!(dest, 0)
    for (j, vx) in couple(index(A, ?, :), x)
        for (vA, idest) in couple(stored(A, :, j), index(dest))
            dest[idest] += vA*vx
        end
    end
    dest
end

# Computing A'*x, assuming column-major order.
function matvecmul!(dest::AbstractVector, AT::MatrixTranspose, x::AbstractVector)
    A = parent(AT)
    fill!(dest, 0)
    for (j, idest) in couple(index(A, ?, :), index(dest))
        val = dest[idest]
        for (vA, vx) in couple(stored(A, :, j), x)
            val += vA*vx
        end
        dest[idest] = val
    end
    dest
end

# A more challenging variant: Inf/NaN-aware sparse
# matrix multiplication with a hypothetical fillval
# Trying the algorithm in https://github.com/JuliaLang/julia/pull/15579#issuecomment-200604174 for CSC
function matvecmul!(dest::AbstractVector, A::AbstractSparseMatrix, x::AbstractVector)
    inds(A, 2) == inds(x, 1) || throw(DimensionMismatch("blah blah"))
    inds(A, 1) == inds(dest, 1) || throw(DimensionMismatch("blah blah"))
    if isnan(A.fillval)
        # Acol holds the previous stored column for each row. Don't
        # assume 1-based. This does assume the indexes are integers.
        Acol = fill(first(inds(A, 2))-1, inds(A, 1))
        k = 0
        for (col, vx) in couple(index(A, ?, :), x)
            for (vA, idest, colold) in couple(stored(A, :, col), index(dest), Acol)
                inc = vA*vx
                dest[idest] += ifelse(colold==col+1, inc, oftype(inc, NaN))
            end
        end
    elseif isinf(A.fillval)
        fill!(dest, 0)
        Acol = fill(first(inds(A, 2))-1, inds(A, 1)) # prev stored col
        # Cumulative count of number of zeros of x
        xz0 = cumsum(y->y==0, x)     # should have the same inds as x
        xz = icat(first(xz0), xz0)   # starts one earlier
        for (col, jx) in couple(index(A, ?, :), index(x))
            vx, z = v[jx], xz[jx]
            for (vA, idest, colold) in couple(stored(A, :, col), index(dest), Acol)
                inc = vA*vx
                dest[idest] += ifelse(z == xz[colold], inc, oftype(inc, NaN))
            end
        end
    else
        for (col, vx) in couple(index(A, ?, :), x)
            for (vA, idest) in couple(stored(A, :, col), index(dest))
                dest[idest] += vA*vx
            end
        end
    end
    dest
end

# Simple matrix-matrix multiplication A'*B. Efficient if output is
# dense; inputs can be whatever.
function matmatmul!(dest::AbstractMatrix, AT::MatrixTranspose, B::AbstractMatrix)
    A = parent(AT)
    inds(A, 1) == inds(B, 1) || throw(DimensionMismatch("blah blah"))
    fill!(dest, 0)
    for (jB, jdest) in couple(index(B, ?, :), index(dest, ?, :))
        for (iA, idest) in couple(index(A, ?, :), index(dest, :, jdest))
            val = dest[idest,jdest]
            for (vA, vB) in couple(stored(A, :, iA), stored(B, :, jB))
                val += vA*vB
            end
            dest[idest,jdest] = val
        end
    end
    dest
end

# Cholesky decomposition
# This works for arrays that have numeric indexes, but would fail for
# an array indexed by A[:cat, :dog]
function chol!(A::AbstractMatrix{T}, ::Type{Val{:L}}) where T
    ind = inds(A, 1)
    inds(A, 2) == ind || throw(DimensionMismatch("blah blah"))
    @inbounds begin
        for k in ind
            Akkm = A[k,k]
            for (Arow, Acol) in zip(view(A, k, first:k-1), view(A, first:k-1, k))
                Akkm -= Arow*Acol'
            end
            Akk = chol!(Akkm, Val{:L})
            A[k,k] = Akk
            AkkInv = inv(Akk)
            Ak = view(A, :, k)
            for j in ind[first:findfirst(ind, k)-1]  # would be nice to have a notation for this
                Aj = view(A, :, j)
                c = Aj[k]'*AkkInv'
                for (i, aij) in zip(index(Ak, k+1:end), slice(Aj, k+1:end))
                    if j == first(ind)
                        Ak[i] = Ak[i]*AkkInv'
                    end
                    if j < k
                        Ak[i] -= aij*c
                    end
                end
            end
        end
     end
    return LowerTriangular(A)
end

ext: jl


text:
using ITensors,
      Test

@testset "IndexSet" begin
  idim = 2
  jdim = 3
  kdim = 4
  ldim = 5
  i = Index(idim,"i")
  j = Index(jdim,"j")
  k = Index(kdim,"k")
  l = Index(ldim,"l")
  @testset "show" begin
    indset = IndexSet(i,j,k)
    @test length(sprint(show, indset)) > 1
  end
  @testset "Basic constructors" begin
    I = IndexSet(i,j,k)
    @test IndexSet(I) === I
    @test l ∈ IndexSet(I, l)
    @test l ∈ IndexSet(l, I)
    @test l ∈ IndexSet( (I, IndexSet(l)) )
    #TODO: what should size(::IndexSet) do?
    #@test size(I) == (3,)
  end
  @testset "Convert to Index" begin
    @test Index(IndexSet(i)) === i
    @test_throws ErrorException Index(IndexSet(i, j))
  end
  @testset "Index dimensions" begin
    I = IndexSet(i,j,k)
    @test dim(I) == idim*jdim*kdim
    @test dims(I) == (idim,jdim,kdim)
    @test dim(I,1) == idim
    @test dim(I,2) == jdim
    @test dim(I,3) == kdim

    @test maxDim(I) == max(idim,jdim,kdim)
  end
  @testset "Set operations" begin
    I1 = IndexSet(i,j,k)
    I2 = IndexSet(k,l)
    I3 = IndexSet(j,l)
    @test hassameinds(I1,(k,j,i))
    @test uniqueindex(I1,(I2,I3)) == i
    @test isnothing(uniqueindex(I1,IndexSet(k, j, i)))
    @test uniqueinds(I1,I2) == IndexSet(i,j)
    @test setdiff(I1,I2) == IndexSet(i,j)
    @test hassameinds(uniqueinds(I1,I2),(j,i))
    @test commoninds(I1,I2) == IndexSet(k)
    @test commonindex(I1,I2) == k
    @test isnothing(commonindex(I1,IndexSet(l)))
    @test commoninds(I1,(j,l)) == IndexSet(j)
    @test commonindex(I1,(j,l)) == j
    @test commoninds(I1,(j,k)) == IndexSet(j,k)
    @test hassameinds(commoninds(I1,(j,k,l)),(j,k))
    @test findinds(I1,"i") == IndexSet(i)
    @test findindex(I1,"j") == j
    @test isnothing(findindex(I1,"l"))
    @test indexposition(I1,i) == 1
    @test indexposition(I1,j) == 2
    @test indexposition(I1,k) == 3
    @test isnothing(indexposition(I1,Index(2)))
  end
  @testset "commoninds index ordering" begin
    I = IndexSet(i,k,j)
    J = IndexSet(j,l,i)
    # Test that commoninds respects the ordering
    # of the indices in the first IndexSet
    @test commoninds(I,J) == IndexSet(i,j)
    @test commoninds(J,I) == IndexSet(j,i)
  end
  @testset "adjoint" begin
    I = IndexSet(i,k,j)
    @test adjoint(I) == IndexSet(i', k', j')
  end
  @testset "mapprime" begin
    I = IndexSet(i',k'',j)
    @test mapprime(I,1,5) == IndexSet(i^5,k'',j)
    @test mapprime(I,2,0) == IndexSet(i',k,j)

    J = IndexSet(i,j,k')
    @test mapprime(J,0,2) == IndexSet(i'',j'',k')

    mapprime!(J,1,5)
    @test J == IndexSet(i,j,k^5)
  end
  @testset "compute_strides" begin
    I = IndexSet(i, j)
    @test ITensors.compute_strides(I) == [1, idim]
  end
  @testset "setprime" begin
    I = IndexSet(i, j)
    J = setprime(I, 2, i)
    @test i'' ∈ J
  end
  @testset "prime" begin
    I = IndexSet(i, j)
    J = prime(I, j)
    @test i ∈ J
    @test j' ∈ J
  end
  @testset "noprime" begin
    I = IndexSet(i'', j')
    J = noprime(I)
    @test i ∈ J
    @test j ∈ J
  end
  @testset "swapprime" begin
    I = IndexSet(i,j)
    @test swapprime(I,0,1) == IndexSet(i',j')
    @test swapprime(I,0,4) == IndexSet(i^4,j^4)
    I = IndexSet(i,j'')
    @test swapprime(I,2,0) == IndexSet(i'',j)
    I = IndexSet(i,j'',k,l)
    @test swapprime(I,2,0) == IndexSet(i'',j,k'',l'')
    I = IndexSet(i,k'',j'')
    @test swapprime(I,2,1) == IndexSet(i,k',j')
    # In-place version:
    I = IndexSet(i,k'',j''')
    swapprime!(I,2,0)
    @test I == IndexSet(i'',k,j''')
    # With tags specified:
    I = IndexSet(i,k,j)
    @test swapprime(I,0,1,"i") == IndexSet(i',k,j)
    @test swapprime(I,0,1,"j") == IndexSet(i,k,j')

    I = IndexSet(i,i',j)
    @test swapprime(I,0,1,"i") == IndexSet(i',i,j)
    @test swapprime(I,0,1,"j") == IndexSet(i,i',j')
  end
  @testset "swaptags" begin
    i1 = Index(2,"Site,A")
    i2 = Index(2,"Site,B")
    is = IndexSet(i1,i2)
    sis = swaptags(is,"Site","Link")
    for j in sis
      @test !hastags(j,"Site")
      @test hastags(j,"Link")
    end
  end
end


text:
ENV["GKSwstype"] = "100" #src

# # Density Estimation
#
# Understanding the shape of your data is important for any data analysis.
#
# - Do you have continuous data? Categorical data? Both?
# - Are variables correlated?
# - What insights can you gain from univariate and bivariate data summaries?
#
# There are many ways to estimate and visualize the distribution of data.  In this module,
# we'll cover: TODO
#
# ## Empirical Density Estimation
#
# ### Empirical Cumulative Density Function (ECDF)
#
# The **`StatsBase`** package provides the `ecdf` function, which returns a function that
# can calculate the ECDF evaluated at a point, defined as
#
# $$f_{\text{ECDF}}(x) = \frac{\text{# values less than or equal to x}}{\text{total # values}}.$$
#
# In other words, the number returned by the ECDF is the probability that a random sample
# from the dataset is less than or equal to the argument.

using StatsBase, Plots

y = 1:10

f = ecdf(y)

plot(f, 0, 11, label="", xlab="x", ylab = "P(data <= x)")

# ### Univariate Data
#
# First we will look at how to estimate/plot univariate distributions of data.
# As we saw in the intro module, we can plot histograms of continuous datafrom the
# **`StatsBase`** package.

using StatsBase, Plots

y = randn(10_000)

h = fit(Histogram, y; closed = :left)

plot(h)

# Alternatively, we can create the plot with the `histogram` function from ``*Plots*``.

histogram(y)

# #### Kernel Density Estimation (KDE)
#
# Kernel Density Estimation (KDE) is a method of estimating the probability density function
# by averaging together a kernel function applied to neighboring points.  A KDE with kernel
# $K$ and bandwidth parameter $h>0$ is a function of data points $(x_1,\ldots,x_n)$ that looks like
#
# $$\hat f(x) = \frac{1}{nh} \sum_{i=1}^n K\left(\frac{x-x_i}{h}\right).$$
#
# Here are some examples of kernel functions:
#
# ![](https://user-images.githubusercontent.com/8075494/30523575-acd48de2-9bb1-11e7-8f0f-3ce2ab09c713.png)
#
# With the **`KernelDensity`** package, let's plot some KDEs with a variety of bandwidths
# using the default (Normal) kernel.  Note that the bandwidth choice has a very big effect
# on the shape of the estimated density.

using KernelDensity, StatsPlots

k = kde(y)

p = plot(k, label = "Bandwidth: Auto", w=2)

for bw in .1:.2:.9
    plot!(p, kde(y; bandwidth = bw), label = "Bandwidth: $bw", w=2)
end

histogram!(p, y, alpha = .1, normed=true)

# #### Averaged Shifted Histograms (ASH)
#
# The Averaged Shifted Histogram (ASH) estimator is a bit of a misnomer; It is essentially
# kernel density estimation performed on a fine-partition (many small bins) histogram.
#
# ASH has an advantage over KDEs in terms of performance since the density is calculated
# over the bins instead of all observations.  Similarly, ASH has an advantage over histograms
# in that smoothing the density over many small bins gives a more fine-grained view of the
# distribution.

using AverageShiftedHistograms

a = ash(y)

plot(a)

# The line is a bit noisy, but we can increase the amount of smoothing with parameter `m`

ash!(a; m = 15)

plot(a)

# But make sure you don't oversmooth!

ash!(a; m = 50)
plot(a)

# Let's compare all three of our density estimates.

plot(ash!(a, m = 15), hist=false)
histogram!(y, normed=true, alpha=.2, label = "Histogram")
plot!(k, label = "KDE")

# #### Categorical Data

using OnlineStats

y = rand([1,2,2,3,3,3,4,4,5], 10^5)

o = fit!(CountMap(Int), y)

plot(o)

# ### Bivariate Relationships
#
# So far, we've only looked at univariate distributions.  However, we can also use
# Histograms, KDEs, and ASH estimators on bivariate data.

x = randn(10^4)
y = x .+ randn(10^4)

h = fit(Histogram, (x,y); closed = :left)
k = kde((x,y))
a = ash(x, y)

plot(
    plot(h; title="Histogram"),
    plot(k; title="KDE"),
    plot(a; title = "ASH");
    layout=3
)

# ## Parametric Density Estimation
#
# Suppose you already know some things about a dataset's distribution or you want to see
# how well the data fits a specific distribution.  Here we'll use the **`Distributions`**
# package to fit data to distributions using maximum likelihood estimation (MLE).  We
# also need the **`StatsPlots`** package for plotting the distributions.

using Distributions

# In **`Distributions`**, each distribution is its own type, and there are a variety of
# methods that work for any distribution

for dist in [Normal(0,1), Gamma(5, 1), Poisson(5)]
    println("The probability density function for $dist evaluated at 0 is $(pdf(dist, 0))")
end

# The `fit` function can be used to pick a reasonable method of fitting the distribution
# (typically MLE).  MLE can be explicitly chosen using the `fit_mle` function.

y = 20 .+ randn(1000)

histogram(y, normed=true, label = "Histogram", xlim = extrema(y))

plot!(fit(Normal, y), label = "Normal Fit")
plot!(fit(Gamma, y), label = "Gamma Fit")

max_stars_repo_name: fercarozzi/JuliaAcademyMaterials

max_issues_repo_name: fercarozzi/JuliaAcademyMaterials

max_forks_repo_name: fercarozzi/JuliaAcademyMaterials


text:
using LinearSolvers
using Test

@testset "LinearSolvers.jl" begin
    # Write your own tests here.
end


text:
# doubt anyone is ever gonna edit this but you never know.

get_redirect_uri()      = "http://localhost:8688"
get_auth_server_port()  = 8688


text:
using LazyBandedMatrices, BlockBandedMatrices, BandedMatrices, LazyArrays, BlockArrays,
            ArrayLayouts, MatrixFactorizations, Random, Test
import LinearAlgebra
import LinearAlgebra: qr, rmul!, lmul!
import LazyArrays: Applied, resizedata!, FillLayout, MulStyle, arguments, colsupport, rowsupport, LazyLayout, ApplyStyle,
                    PaddedLayout, paddeddata, call, ApplyLayout, LazyArrayStyle, simplifiable
import LazyBandedMatrices: VcatBandedMatrix, BroadcastBlockBandedLayout, BroadcastBandedLayout,
                    ApplyBandedLayout, ApplyBlockBandedLayout, ApplyBandedBlockBandedLayout, BlockKron, LazyBandedLayout, BroadcastBandedBlockBandedLayout
import BandedMatrices: BandedStyle, _BandedMatrix, AbstractBandedMatrix, BandedRows, BandedColumns
import ArrayLayouts: StridedLayout, OnesLayout
import BlockArrays: blockcolsupport, blockrowsupport

Random.seed!(0)

include("test_tridiag.jl")
include("test_bidiag.jl")
include("test_special.jl")

struct PseudoBandedMatrix{T} <: AbstractMatrix{T}
    data::Array{T}
    l::Int
    u::Int
end

Base.size(A::PseudoBandedMatrix) = size(A.data)
function Base.getindex(A::PseudoBandedMatrix, j::Int, k::Int)
    l, u = bandwidths(A)
    if -l ≤ k-j ≤ u
        A.data[j, k]
    else
        zero(eltype(A.data))
    end
end
function Base.setindex!(A::PseudoBandedMatrix, v, j::Int, k::Int)
    l, u = bandwidths(A)
    if -l ≤ k-j ≤ u
        A.data[j, k] = v
    else
        error("out of band.")
    end
end

struct PseudoBandedLayout <: AbstractBandedLayout end
Base.BroadcastStyle(::Type{<:PseudoBandedMatrix}) = BandedStyle()
BandedMatrices.MemoryLayout(::Type{<:PseudoBandedMatrix}) = PseudoBandedLayout()
BandedMatrices.isbanded(::PseudoBandedMatrix) = true
BandedMatrices.bandwidths(A::PseudoBandedMatrix) = (A.l , A.u)
BandedMatrices.inbands_getindex(A::PseudoBandedMatrix, j::Int, k::Int) = A.data[j, k]
BandedMatrices.inbands_setindex!(A::PseudoBandedMatrix, v, j::Int, k::Int) = setindex!(A.data, v, j, k)
LinearAlgebra.fill!(A::PseudoBandedMatrix, v) = fill!(A.data,v)
ArrayLayouts.lmul!(β::Number, A::PseudoBandedMatrix) = (lmul!(β, A.data); A)
LinearAlgebra.lmul!(β::Number, A::PseudoBandedMatrix) = (lmul!(β, A.data); A)

struct MyLazyArray{T,N} <: AbstractArray{T,N}
    data::Array{T,N}
end

Base.size(A::MyLazyArray) = size(A.data)
Base.getindex(A::MyLazyArray, j::Int...) = A.data[j...]
LazyArrays.MemoryLayout(::Type{<:MyLazyArray}) = LazyLayout()
LinearAlgebra.factorize(A::MyLazyArray) = factorize(A.data)

@testset "LazyBlock" begin
    @test Block(5) in BroadcastVector(Block, [1,3,5])
    @test Base.broadcasted(LazyArrayStyle{1}(), Block, 1:5) ≡ Block.(1:5)
    @test Base.broadcasted(LazyArrayStyle{1}(), Int, Block.(1:5)) ≡ 1:5
    @test Base.broadcasted(LazyArrayStyle{0}(), Int, Block(1)) ≡ 1
end

@testset "LazyBlockArray Triangle Recurrences" begin
    N = 1000
    n = mortar(BroadcastArray(Fill,Base.OneTo(N),Base.OneTo(N)))
    k = mortar(BroadcastArray(Base.OneTo,Base.OneTo(N)))

    @test view(n, Block(5)) ≡ Fill(5,5)
    @test view(k,Block(5)) ≡ Base.OneTo(5)
    a = b = c = 0.0
    # for some reason the following was causing major slowdown. I think it
    # went pass a limit to Base.Broadcast.flatten which caused `bc.f` to have a strange type.
    # bc = Base.Broadcast.instantiate(Base.broadcasted(/, Base.broadcasted(*, k, Base.broadcasted(-, Base.broadcasted(-, k, n), a)), Base.broadcasted(+, 2k, b+c-1)))

    bc = Base.Broadcast.instantiate(Base.broadcasted((k,n,a,b,c) -> k * (k-n-a) / (2k+(b+c-1)), k, n, a, b, c))
    @test axes(n,1) ≡ axes(k,1) ≡ axes(bc)[1] ≡ blockedrange(Base.OneTo(N))
    u = (k .* (k .- n .- a) ./ (2k .+ (b+c-1)))
    @test u == (Vector(k) .* (Vector(k) .- Vector(n) .- a) ./ (2Vector(k) .+ (b+c-1)))
    @test copyto!(u, bc) == (k .* (k .- n .- a) ./ (2k .+ (b+c-1)))
    @test @allocated(copyto!(u, bc)) ≤ 1000
    # not clear why allocatinos so high: all allocations are coming from checking
    # axes

    u = PseudoBlockArray{Float64}(undef, collect(1:N))
    @test copyto!(u, bc) == (k .* (k .- n .- a) ./ (2k .+ (b+c-1)))
    @test @allocated(copyto!(u, bc)) ≤ 1000
end

@testset "Padded" begin
    @testset "Banded padded" begin
        A = _BandedMatrix((1:10)', 10, -1,1)
        x = Vcat(1:3, Zeros(10-3))
        @test MemoryLayout(x) isa PaddedLayout
        @test A*x isa Vcat{Float64,1,<:Tuple{<:Vector,<:Zeros}}
        @test length((A*x).args[1]) == length(x.args[1]) + bandwidth(A,1) == 2
        @test A*x == A*Vector(x)

        A = _BandedMatrix(randn(3,10), 10, 1,1)
        x = Vcat(randn(10), Zeros(0))
        @test A*x isa Vcat{Float64,1,<:Tuple{<:Vector,<:Zeros}}
        @test length((A*x).args[1]) == 10
        @test A*x ≈ A*Vector(x)

        A = Vcat(Zeros(1,10), brand(9,10,0,2))
        @test bandwidths(A) == (1,1)
        @test BandedMatrix(A) == Array(A) == A

        A = Hcat(Zeros(5,2), brand(5,5,1,1))
        @test bandwidths(A) == (-1,3)
        @test BandedMatrix(A) == Array(A) == A
    end
    @testset "BlockBanded and padded" begin
        A = BlockBandedMatrix{Float64}(undef, 1:4, 1:4, (1,0)); A.data .= randn.();
        D = mortar(Diagonal([randn(k,k) for k=1:4]))
        c = Vcat(randn(3), Zeros(7))
        b = PseudoBlockVector(c, (axes(A,2),))
        @test MemoryLayout(b) isa PaddedLayout
        @test MemoryLayout(A*b) isa PaddedLayout
        @test MemoryLayout(A*c) isa PaddedLayout
        @test A*b ≈ A*c ≈ Matrix(A)*Vector(b)
        @test D*b ≈ D*c ≈ Matrix(D)*Vector(b)

        @test b[Block.(2:3)] isa PseudoBlockVector{Float64,<:ApplyArray}
        @test MemoryLayout(b[Block.(2:3)]) isa PaddedLayout
        @test b[Block.(2:3)] == b[2:6]
    end

    @testset "BroadcastBanded * Padded" begin
        A = BroadcastArray(*, randn(5), brand(5,5,1,2))
        @test axes(A) == (Base.OneTo(5), Base.OneTo(5))
        B = BroadcastArray(*, randn(5,5), brand(5,5,1,2))
        b = Vcat(randn(2), Zeros(3))
        @test A*b ≈ Matrix(A)b
        @test B*b ≈ Matrix(B)b
    end

    @testset "Apply * Banded" begin
        B = brand(5,5,2,1)
        A = ApplyArray(*, B, B)
        @test A * Vcat([1,2], Zeros(3)) ≈ B*B*[1,2,0,0,0]
    end

    @testset "block padded" begin
        c = PseudoBlockVector(Vcat(1, Zeros(5)), 1:3)
        @test paddeddata(c) == [1]
        @test paddeddata(c) isa PseudoBlockVector
        @test blockcolsupport(c) == Block.(1:1)
        C = PseudoBlockArray(Vcat(randn(2,3), Zeros(4)), 1:3, [1,2])
        @test blockcolsupport(C) == Block.(1:2)
        @test blockrowsupport(C) == Block.(1:2)

        @test C[Block.(1:2),1:3] == C[Block.(1:2),Block.(1:2)] == C[1:3,Block.(1:2)] == C[1:3,1:3]

        H = PseudoBlockArray(Hcat(1, Zeros(1,5)), [1], 1:3)
        @test paddeddata(H) == Ones(1,1)
    end

    @testset "Banded Perturbed" begin
        n = 1000
        D = Diagonal(1:n)
        P = ApplyArray(hvcat, 2, randn(3,3), Zeros(3,n-3), Zeros(n-3,3), Zeros(n-3,n-3))
        @test isbanded(P)
        @test bandwidths(P) == (2,2)

        B = BroadcastArray(+, D, P)
        @test MemoryLayout(B) isa BroadcastBandedLayout
        @test bandwidths(B) == (2,2)

        B = BroadcastArray(+, P, D)
        @test MemoryLayout(B) isa BroadcastBandedLayout
        @test bandwidths(B) == (2,2)

        C = ApplyArray(hvcat, 2, 1, 2, 3, 4)
        @test bandwidths(C) == (1,1)
    end
end

@testset "MulMatrix" begin
    @testset "MulBanded" begin
        A = brand(6,5,0,1)
        B = brand(5,5,1,0)

        M = ApplyArray(*, A)
        @test BandedMatrix(M) == copyto!(similar(A), M) == A

        M = ApplyArray(*,A,B)
        @test isbanded(M) && isbanded(Applied(M))
        @test bandwidths(M) == bandwidths(Applied(M))
        @test BandedMatrix(M) == A*B == copyto!(BandedMatrix(M), M)
        @test MemoryLayout(M) isa ApplyBandedLayout{typeof(*)}
        @test arguments(M) == (A,B)
        @test call(M) == *
        @test colsupport(M,1) == colsupport(Applied(M),1) == 1:2
        @test rowsupport(M,1) == rowsupport(Applied(M),1) == 1:2

        @test Base.BroadcastStyle(typeof(M)) isa BandedStyle
        @test M .+ A isa BandedMatrix
        @test M .+ A == M .+ Matrix(A) == Matrix(A) .+ M

        V = view(M,1:4,1:4)
        @test bandwidths(V) == (1,1)
        @test MemoryLayout(V) == MemoryLayout(M)
        @test M[1:4,1:4] isa BandedMatrix
        @test colsupport(V,1) == 1:2
        @test rowsupport(V,1) == 1:2

        @test MemoryLayout(view(M, [1,3], [2,3])) isa ApplyLayout{typeof(*)}

        A = brand(5,5,0,1)
        B = brand(6,5,1,0)
        @test_throws DimensionMismatch ApplyArray(*,A,B)

        A = brand(6,5,0,1)
        B = brand(5,5,1,0)
        C = brand(5,6,2,2)
        M = applied(*,A,B,C)
        @test @inferred(eltype(M)) == Float64
        @test bandwidths(M) == (3,3)
        @test M[1,1] ≈ (A*B*C)[1,1]

        M = @inferred(ApplyArray(*,A,B,C))
        @test @inferred(eltype(M)) == Float64
        @test bandwidths(M) == (3,3)
        @test BandedMatrix(M) ≈ A*B*C ≈ copyto!(BandedMatrix(M), M)

        M = ApplyArray(*, A, Zeros(5))
        @test colsupport(M,1) == colsupport(Applied(M),1)
        @test_skip colsupport(M,1) == 1:0

        @testset "inv" begin
            A = brand(6,5,0,1)
            B = brand(5,5,1,0)
            C = randn(6,2)
            M = ApplyArray(*,A,B)
            @test M \ C ≈ Matrix(M) \ C
        end

        @testset "Sym/Herm" begin
            A = brand(5,5,0,1)
            B = brand(5,5,1,0)
            M = ApplyArray(*,A,B)
            C = ApplyArray(*,A,im*B)
            @test MemoryLayout(Symmetric(M)) isa SymmetricLayout{LazyBandedLayout}
            @test MemoryLayout(Symmetric(C)) isa SymmetricLayout{LazyBandedLayout}
            @test MemoryLayout(Hermitian(M)) isa SymmetricLayout{LazyBandedLayout}
            @test MemoryLayout(Hermitian(C)) isa HermitianLayout{LazyBandedLayout}
        end
    end
    @testset "MulBlockBanded" begin
        A = BlockBandedMatrix{Float64}(undef, 1:4, 1:4, (1,0)); A.data .= randn.();
        B = BlockBandedMatrix{Float64}(undef, 1:4, 1:4, (1,1)); B.data .= randn.();
        M = ApplyMatrix(*, A, B)
        @test blockbandwidths(M) == (2,1)
        @test MemoryLayout(M) isa ApplyBlockBandedLayout{typeof(*)}
        @test Base.BroadcastStyle(typeof(M)) isa LazyArrayStyle{2}
        @test BlockBandedMatrix(M) ≈ A*B
        @test arguments(M) == (A,B)
        V = view(M, Block.(1:2), Block.(1:2))
        @test MemoryLayout(V) isa ApplyBlockBandedLayout{typeof(*)}
        @test arguments(V) == (A[Block.(1:2),Block.(1:2)], B[Block.(1:2),Block.(1:2)])
        @test M[Block.(1:2), Block.(1:2)] isa BlockBandedMatrix
    end
    @testset "MulBandedBlockBanded" begin
        A = BandedBlockBandedMatrix{Float64}(undef, 1:4, 1:4, (1,0), (1,0)); A.data .= randn.();
        B = BandedBlockBandedMatrix{Float64}(undef, 1:4, 1:4, (1,1), (1,1)); B.data .= randn.();
        M = ApplyMatrix(*, A, B)
        @test blockbandwidths(M) == (2,1)
        @test subblockbandwidths(M) == (2,1)
        @test MemoryLayout(M) isa ApplyBandedBlockBandedLayout{typeof(*)}
        @test Base.BroadcastStyle(typeof(M)) isa LazyArrayStyle{2}
        @test BandedBlockBandedMatrix(M) ≈ BlockBandedMatrix(M) ≈ A*B
        @test arguments(M) == (A,B)
        V = view(M, Block.(1:2), Block.(1:2))
        @test MemoryLayout(V) isa ApplyBandedBlockBandedLayout{typeof(*)}
        @test arguments(V) == (A[Block.(1:2),Block.(1:2)], B[Block.(1:2),Block.(1:2)])
        @test M[Block.(1:2), Block.(1:2)] isa BandedBlockBandedMatrix
        V = view(M, 1:3, 1:3)
        @test MemoryLayout(V) isa ApplyLayout{typeof(*)}
        @test arguments(V) == (A[1:3,1:3], B[1:3,1:3])
        @test M[1:3, 1:3] ≈ (A*B)[1:3,1:3]

        @test M[Block(2)[1:2],Block(2)[1:2]] isa BandedMatrix
        @test M[Block(2)[1:2],Block(2)] isa BandedMatrix
        @test M[Block(2),Block(2)[1:2]] isa BandedMatrix
        @test M[Block.(1:2), Block.(2:3)] isa BandedBlockBandedMatrix
        @test M[Block(2),Block.(2:3)] isa PseudoBlockArray
        @test M[Block.(2:3),Block(2)] isa PseudoBlockArray
        @test M[Block.(2:3),Block(2)[1:2]] isa PseudoBlockArray
        @test M[Block(2)[1:2],Block.(2:3)] isa PseudoBlockArray
    end
    @testset "Pseudo Mul" begin
        A = PseudoBandedMatrix(rand(5, 4), 1, 2)
        B = PseudoBandedMatrix(rand(4, 4), 2, 3)
        C = PseudoBandedMatrix(zeros(5, 4), 3, 4)
        D = zeros(5, 4)

        @test (C .= applied(*, A, B)) ≈ (D .= applied(*, A, B)) ≈ A*B
    end
    @testset "MulStyle" begin
        A = brand(5,5,0,1)
        B = brand(5,5,1,0)
        C = BroadcastMatrix(*, A, 2)
        M = ApplyArray(*,A,B)
        @test M^2 isa BandedMatrix
        @test M*C isa ApplyMatrix{Float64,typeof(*)}
        @test C*M isa ApplyMatrix{Float64,typeof(*)}
    end
    @testset "Apply*Broadcast" begin
        A = randn(5,5)
        B = randn(5,5)
        C = brand(5,5,1,1)
        D = brand(5,5,1,1)
        @test ApplyArray(*, A, B) * BroadcastArray(*, A, B) ≈ (A*B) * (A .* B)
        @test ApplyArray(*, C, D) * BroadcastArray(*, A, B) ≈ (C*D) * (A .* B)
        @test ApplyArray(*, A, B) * BroadcastArray(*, C, D) ≈ (A*B) * (C .* D)
        @test BroadcastArray(*, A, B) * ApplyArray(*, A, B) ≈ (A .* B) * (A*B)
        @test BroadcastArray(*, C, D) * ApplyArray(*, A, B) ≈ (C .* D) * (A*B)
        @test BroadcastArray(*, A, B) * ApplyArray(*, C, D) ≈ (A .* B) * (C*D)

        @test ApplyArray(*, A, B) \ BroadcastArray(*, A, B) ≈ (A*B) \ (A .* B)
        @test BroadcastArray(*, A, B) \ ApplyArray(*, A, B) ≈ (A .* B) \ (A * B)
        @test BroadcastArray(*, A, B) \ BroadcastArray(*, A, B) ≈ (A .* B) \ (A .* B)
        @test BroadcastArray(*, C, D) \ BroadcastArray(*, C, D) ≈ (C .* D) \ (C .* D)
        @test ApplyArray(*, C, D) \ BroadcastArray(*, C, D) ≈ (C * D) \ (C .* D)
    end

    @testset "flatten" begin
        A = brand(5,5,1,1)
        B = brand(5,5,1,1)
        C = brand(5,5,1,1)
        @test LazyArrays.flatten(ApplyArray(*, A, ApplyArray(*, B, C))) ≈ A * B *C
    end

    @testset "Eye simplifiable" begin
        A = Eye(5)
        B = brand(5,5,1,1)
        C = brand(5,5,1,1)
        D = Diagonal(Fill(2,5))
        @test simplifiable(*, A, BroadcastArray(*, B, C)) == Val(true)
        @test simplifiable(*, BroadcastArray(*, B, C), A) == Val(true)

        @test D * BroadcastArray(*, B, C) ≈ D * (B .* C)
        @test BroadcastArray(*, B, C) * D ≈ (B .* C) * D
    end
end

@testset "InvMatrix" begin
    D = brand(5,5,0,0)
    L = brand(5,5,2,0)
    U = brand(5,5,0,2)
    B = brand(5,5,1,2)

    @test bandwidths(ApplyMatrix(inv,D)) == (0,0)
    @test bandwidths(ApplyMatrix(inv,L)) == (4,0)
    @test bandwidths(ApplyMatrix(inv,U)) == (0,4)
    @test bandwidths(ApplyMatrix(inv,B)) == (4,4)

    @test colsupport(ApplyMatrix(inv,D) ,3) == 3:3
    @test colsupport(ApplyMatrix(inv,L), 3) == 3:5
    @test colsupport(ApplyMatrix(inv,U), 3) == 1:3
    @test colsupport(ApplyMatrix(inv,B), 3) == 1:5

    @test rowsupport(ApplyMatrix(inv,D) ,3) == 3:3
    @test rowsupport(ApplyMatrix(inv,L), 3) == 1:3
    @test rowsupport(ApplyMatrix(inv,U), 3) == 3:5
    @test rowsupport(ApplyMatrix(inv,B), 3) == 1:5

    @test bandwidths(ApplyMatrix(\,D,B)) == (1,2)
    @test bandwidths(ApplyMatrix(\,L,B)) == (4,2)
    @test bandwidths(ApplyMatrix(\,U,B)) == (1,4)
    @test bandwidths(ApplyMatrix(\,B,B)) == (4,4)

    @test colsupport(ApplyMatrix(\,D,B), 3) == 1:4
    @test colsupport(ApplyMatrix(\,L,B), 4) == 2:5
    @test colsupport(ApplyMatrix(\,U,B), 3) == 1:4
    @test colsupport(ApplyMatrix(\,B,B), 3) == 1:5

    @test rowsupport(ApplyMatrix(\,D,B), 3) == 2:5
    @test rowsupport(ApplyMatrix(\,L,B), 2) == 1:4
    @test rowsupport(ApplyMatrix(\,U,B), 3) == 2:5
    @test rowsupport(ApplyMatrix(\,B,B), 3) == 1:5
end

@testset "Cat" begin
    A = brand(6,5,2,1)
    H = Hcat(A,A)
    @test H[1,1] == applied(hcat,A,A)[1,1] == A[1,1]
    @test isbanded(H)
    @test bandwidths(H) == (2,6)
    @test BandedMatrix(H) == BandedMatrix(H,(2,6)) == hcat(A,A) == hcat(A,Matrix(A)) ==
            hcat(Matrix(A),A) == hcat(Matrix(A),Matrix(A))
    @test hcat(A,A) isa BandedMatrix
    @test hcat(A,Matrix(A)) isa Matrix
    @test hcat(Matrix(A),A) isa Matrix
    @test isone.(H) isa BandedMatrix
    @test bandwidths(isone.(H)) == (2,6)
    @test @inferred(colsupport(H,1)) == 1:3
    @test Base.replace_in_print_matrix(H,4,1,"0") == "⋅"

    H = Hcat(A,A,A)
    @test isbanded(H)
    @test bandwidths(H) == (2,11)
    @test BandedMatrix(H) == hcat(A,A,A) == hcat(A,Matrix(A),A) == hcat(Matrix(A),A,A) ==
            hcat(Matrix(A),Matrix(A),A) == hcat(Matrix(A),Matrix(A),Matrix(A))
    @test hcat(A,A,A) isa BandedMatrix
    @test isone.(H) isa BandedMatrix
    @test bandwidths(isone.(H)) == (2,11)

    V = Vcat(A,A)
    @test V isa VcatBandedMatrix
    @test isbanded(V)
    @test bandwidths(V) == (8,1)
    @test BandedMatrix(V) == vcat(A,A) == vcat(A,Matrix(A)) == vcat(Matrix(A),A) == vcat(Matrix(A),Matrix(A))
    @test vcat(A,A) isa BandedMatrix
    @test vcat(A,Matrix(A)) isa Matrix
    @test vcat(Matrix(A),A) isa Matrix
    @test isone.(V) isa BandedMatrix
    @test bandwidths(isone.(V)) == (8,1)
    @test Base.replace_in_print_matrix(V,1,3,"0") == "⋅"

    V = Vcat(A,A,A)
    @test bandwidths(V) == (14,1)
    @test BandedMatrix(V) == vcat(A,A,A) == vcat(A,Matrix(A),A) == vcat(Matrix(A),A,A) ==
            vcat(Matrix(A),Matrix(A),A) == vcat(Matrix(A),Matrix(A),Matrix(A))
    @test vcat(A,A,A) isa BandedMatrix
    @test isone.(V) isa BandedMatrix
    @test bandwidths(isone.(V)) == (14,1)
end

@testset "BroadcastMatrix" begin
    @testset "BroadcastBanded" begin
        A = BroadcastMatrix(*, brand(5,5,1,2), brand(5,5,2,1))
        @test eltype(A) == Float64
        @test bandwidths(A) == (1,1)
        @test colsupport(A, 1) == 1:2
        @test rowsupport(A, 1) == 1:2
        @test A == broadcast(*, A.args...) == BandedMatrix(A)
        @test MemoryLayout(A) isa BroadcastBandedLayout{typeof(*)}

        @test MemoryLayout(A') isa BroadcastBandedLayout{typeof(*)}
        @test bandwidths(A') == (1,1)
        @test colsupport(A',1) == rowsupport(A', 1) == 1:2
        @test A' == BroadcastArray(A') == Array(A)' == BandedMatrix(A')

        V = view(A, 2:3, 3:5)
        @test MemoryLayout(V) isa BroadcastBandedLayout{typeof(*)}
        @test bandwidths(V) == (1,0)
        @test colsupport(V,1) == 1:2
        @test V == BroadcastArray(V) == Array(A)[2:3,3:5]
        @test bandwidths(view(A,2:4,3:5)) == (2,0)

        V = view(A, 2:3, 3:5)'
        @test MemoryLayout(V) isa BroadcastBandedLayout{typeof(*)}
        @test bandwidths(V) == (0,1)
        @test colsupport(V,1) == 1:1
        @test V == BroadcastArray(V) == Array(A)[2:3,3:5]'

        B = BroadcastMatrix(+, brand(5,5,1,2), 2)
        @test B == broadcast(+, B.args...)

        C = BroadcastMatrix(+, brand(5,5,1,2), brand(5,5,3,1))
        @test bandwidths(C) == (3,2)
        @test MemoryLayout(C) == BroadcastBandedLayout{typeof(+)}()
        @test isbanded(C) == true
        @test BandedMatrix(C) == C == copyto!(BandedMatrix(C), C)

        D = BroadcastMatrix(*, 2, brand(5,5,1,2))
        @test bandwidths(D) == (1,2)
        @test MemoryLayout(D) == BroadcastBandedLayout{typeof(*)}()
        @test isbanded(D) == true
        @test BandedMatrix(D) == D == copyto!(BandedMatrix(D), D) == 2*D.args[2]

        @testset "band" begin
            @test A[band(0)] == Matrix(A)[band(0)]
            @test B[band(0)] == Matrix(B)[band(0)]
            @test C[band(0)] == Matrix(C)[band(0)]
            @test D[band(0)] == Matrix(D)[band(0)]
        end

        @testset "non-simple" begin
            A = BroadcastMatrix(sin,brand(5,5,1,2))
            @test bandwidths(A) == (1,2)
            @test BandedMatrix(A) == Matrix(A) == A
        end

        @testset "Complex" begin
            C = BroadcastMatrix(*, 2, im*brand(5,5,2,1))
            @test MemoryLayout(C') isa ConjLayout{BroadcastBandedLayout{typeof(*)}}
        end
    end
    @testset "BroadcastBlockBanded" begin
        A = BlockBandedMatrix(randn(6,6),1:3,1:3,(1,1))
        B = BroadcastMatrix(*, 2, A)
        @test blockbandwidths(B) == (1,1)
        @test MemoryLayout(B) == BroadcastBlockBandedLayout{typeof(*)}()
        @test BandedBlockBandedMatrix(B) == B == copyto!(BandedBlockBandedMatrix(B), B) == 2*B.args[2]
        @test MemoryLayout(B') isa LazyBandedMatrices.LazyBlockBandedLayout
        @test BlockBandedMatrix(B') == B'

        C = BroadcastMatrix(*, 2, im*A)
        @test MemoryLayout(C') isa LazyBandedMatrices.LazyBlockBandedLayout
        @test MemoryLayout(transpose(C)) isa LazyBandedMatrices.LazyBlockBandedLayout

        E = BroadcastMatrix(*, A, 2)
        @test MemoryLayout(E) == BroadcastBlockBandedLayout{typeof(*)}()

        D = Diagonal(PseudoBlockArray(randn(6),1:3))
        @test MemoryLayout(BroadcastMatrix(*, A, D)) isa BroadcastBlockBandedLayout{typeof(*)}
        @test MemoryLayout(BroadcastMatrix(*, D, A)) isa BroadcastBlockBandedLayout{typeof(*)}

        F = BroadcastMatrix(*, A, A)
        @test MemoryLayout(F) == BroadcastBlockBandedLayout{typeof(*)}()
    end
    @testset "BroadcastBandedBlockBanded" begin
        A = BandedBlockBandedMatrix(randn(6,6),1:3,1:3,(1,1),(1,1))

        B = BroadcastMatrix(*, 2, A)
        @test blockbandwidths(B) == (1,1)
        @test subblockbandwidths(B) == (1,1)
        @test MemoryLayout(B) == BroadcastBandedBlockBandedLayout{typeof(*)}()
        @test BandedBlockBandedMatrix(B) == B == copyto!(BandedBlockBandedMatrix(B), B) == 2*B.args[2]
        @test MemoryLayout(B') isa LazyBandedMatrices.LazyBandedBlockBandedLayout
        @test BandedBlockBandedMatrix(B') == B'

        C = BroadcastMatrix(*, 2, im*A)
        @test MemoryLayout(C') isa LazyBandedMatrices.LazyBandedBlockBandedLayout
        @test MemoryLayout(transpose(C)) isa LazyBandedMatrices.LazyBandedBlockBandedLayout

        E = BroadcastMatrix(*, A, 2)
        @test MemoryLayout(E) == BroadcastBandedBlockBandedLayout{typeof(*)}()

        D = Diagonal(PseudoBlockArray(randn(6),1:3))
        @test MemoryLayout(BroadcastMatrix(*, A, D)) isa BroadcastBandedBlockBandedLayout{typeof(*)}
        @test MemoryLayout(BroadcastMatrix(*, D, A)) isa BroadcastBandedBlockBandedLayout{typeof(*)}

        F = BroadcastMatrix(*, Ones(axes(A,1)), A)
        @test blockbandwidths(F) == (1,1)
        @test subblockbandwidths(F) == (1,1)
        @test F == A
    end
end

@testset "Cache" begin
    A = _BandedMatrix(Fill(1,3,10_000), 10_000, 1, 1)
    C = cache(A);
    @test C.data isa BandedMatrix{Int,Matrix{Int},Base.OneTo{Int}}
    @test colsupport(C,1) == rowsupport(C,1) == 1:2
    @test bandwidths(C) == bandwidths(A) == (1,1)
    @test isbanded(C)
    resizedata!(C,1,1);
    @test C[1:10,1:10] == A[1:10,1:10]
    @test C[1:10,1:10] isa BandedMatrix
end

@testset "NaN Bug" begin
    C = BandedMatrix{Float64}(undef, (1,2), (0,2)); C.data .= NaN;
    A = brand(1,1,0,1)
    B = brand(1,2,0,2)
    C .= applied(*, A,B)
    @test C == A*B

    C.data .= NaN
    C .= @~ 1.0 * A*B + 0.0 * C
    @test C == A*B
end

@testset "Applied" begin
    A = brand(5,5,1,2)
    @test applied(*,Symmetric(A),A) isa Applied{MulStyle}
    B = apply(*,A,A,A)
    @test B isa BandedMatrix
    @test all(B .=== (A*A)*A)
    @test bandwidths(B) == (3,4)
end

@testset "QL tests" begin
    for T in (Float64,ComplexF64,Float32,ComplexF32)
        A=brand(T,10,10,3,2)
        Q,L=ql(A)
        @test ql(A).factors ≈ ql!(Matrix(A)).factors
        @test ql(A).τ ≈ ql!(Matrix(A)).τ
        @test Matrix(Q)*Matrix(L) ≈ A
        b=rand(T,10)
        @test mul!(similar(b),Q,mul!(similar(b),Q',b)) ≈ b
        for j=1:size(A,2)
            @test Q' * A[:,j] ≈ L[:,j]
        end

        A=brand(T,14,10,3,2)
        Q,L=ql(A)
        @test ql(A).factors ≈ ql!(Matrix(A)).factors
        @test ql(A).τ ≈ ql!(Matrix(A)).τ
        @test_broken Matrix(Q)*Matrix(L) ≈ A

        for k=1:size(A,1),j=1:size(A,2)
            @test Q[k,j] ≈ Matrix(Q)[k,j]
        end

        A=brand(T,10,14,3,2)
        Q,L=ql(A)
        @test ql(A).factors ≈ ql!(Matrix(A)).factors
        @test ql(A).τ ≈ ql!(Matrix(A)).τ
        @test Matrix(Q)*Matrix(L) ≈ A

        for k=1:size(Q,1),j=1:size(Q,2)
            @test Q[k,j] ≈ Matrix(Q)[k,j]
        end

        A=brand(T,10,14,3,6)
        Q,L=ql(A)
        @test ql(A).factors ≈ ql!(Matrix(A)).factors
        @test ql(A).τ ≈ ql!(Matrix(A)).τ
        @test Matrix(Q)*Matrix(L) ≈ A

        for k=1:size(Q,1),j=1:size(Q,2)
            @test Q[k,j] ≈ Matrix(Q)[k,j]
        end

        A=brand(T,100,100,3,4)
        @test ql(A).factors ≈ ql!(Matrix(A)).factors
        @test ql(A).τ ≈ ql!(Matrix(A)).τ
        b=rand(T,100)
        @test ql(A)\b ≈ Matrix(A)\b
        b=rand(T,100,2)
        @test ql(A)\b ≈ Matrix(A)\b
        @test_throws DimensionMismatch ql(A) \ randn(3)
        @test_throws DimensionMismatch ql(A).Q'randn(3)

        A=brand(T,102,100,3,4)
        @test ql(A).factors ≈ ql!(Matrix(A)).factors
        @test ql(A).τ ≈ ql!(Matrix(A)).τ
        b=rand(T,102)
        @test_broken ql(A)\b ≈ Matrix(A)\b
        b=rand(T,102,2)
        @test_broken ql(A)\b ≈ Matrix(A)\b
        @test_throws DimensionMismatch ql(A) \ randn(3)
        @test_throws DimensionMismatch ql(A).Q'randn(3)

        A=brand(T,100,102,3,4)
        @test ql(A).factors ≈ ql!(Matrix(A)).factors
        @test ql(A).τ ≈ ql!(Matrix(A)).τ
        b=rand(T,100)
        @test_broken ql(A)\b ≈ Matrix(A)\b

        A = LinearAlgebra.Tridiagonal(randn(T,99), randn(T,100), randn(T,99))
        @test ql(A).factors ≈ ql!(Matrix(A)).factors
        @test ql(A).τ ≈ ql!(Matrix(A)).τ
        b=rand(T,100)
        @test ql(A)\b ≈ Matrix(A)\b
        b=rand(T,100,2)
        @test ql(A)\b ≈ Matrix(A)\b
        @test_throws DimensionMismatch ql(A) \ randn(3)
        @test_throws DimensionMismatch ql(A).Q'randn(3)
    end

    @testset "lmul!/rmul!" begin
        for T in (Float32, Float64, ComplexF32, ComplexF64)
            A = brand(T,100,100,3,4)
            Q,R = qr(A)
            x = randn(T,100)
            b = randn(T,100,2)
            @test lmul!(Q, copy(x)) ≈ Matrix(Q)*x
            @test lmul!(Q, copy(b)) ≈ Matrix(Q)*b
            @test lmul!(Q', copy(x)) ≈ Matrix(Q)'*x
            @test lmul!(Q', copy(b)) ≈ Matrix(Q)'*b
            c = randn(T,2,100)
            @test rmul!(copy(c), Q) ≈ c*Matrix(Q)
            @test rmul!(copy(c), Q') ≈ c*Matrix(Q')

            A = brand(T,100,100,3,4)
            Q,L = ql(A)
            x = randn(T,100)
            b = randn(T,100,2)
            @test lmul!(Q, copy(x)) ≈ Matrix(Q)*x
            @test lmul!(Q, copy(b)) ≈ Matrix(Q)*b
            @test lmul!(Q', copy(x)) ≈ Matrix(Q)'*x
            @test lmul!(Q', copy(b)) ≈ Matrix(Q)'*b
            c = randn(T,2,100)
            @test rmul!(copy(c), Q) ≈ c*Matrix(Q)
            @test rmul!(copy(c), Q') ≈ c*Matrix(Q')
        end
    end

    @testset "Mixed types" begin
        A=brand(10,10,3,2)
        b=rand(ComplexF64,10)
        Q,L=ql(A)
        @test L\(Q'*b) ≈ ql(A)\b ≈ Matrix(A)\b
        @test Q*L ≈ A

        A=brand(ComplexF64,10,10,3,2)
        b=rand(10)
        Q,L=ql(A)
        @test Q*L ≈ A
        @test L\(Q'*b) ≈ ql(A)\b ≈ Matrix(A)\b

        A = BandedMatrix{Int}(undef, (2,1), (4,4))
        A.data .= 1:length(A.data)
        Q, L = ql(A)
        @test_broken Q*L ≈ A
    end
end

struct FiniteDifference{T} <: AbstractBandedMatrix{T}
    n::Int
end

FiniteDifference(n) = FiniteDifference{Float64}(n)

Base.getindex(F::FiniteDifference{T}, k::Int, j::Int) where T =
    if k == j
        -2*one(T)*F.n^2
    elseif abs(k-j) == 1
        one(T)*F.n^2
    else
        zero(T)
    end

BandedMatrices.bandwidths(F::FiniteDifference) = (1,1)
Base.size(F::FiniteDifference) = (F.n,F.n)

@testset "Misc" begin
    @testset "Block banded Kron" begin
        n = 10
        h = 1/n
        D² = BandedMatrix(0 => Fill(-2,n), 1 => Fill(1,n-1), -1 => Fill(1,n-1))/h^2

        @time D_xx = BandedBlockBandedMatrix(BlockKron(D², Eye(n)))
        @time D_yy = BandedBlockBandedMatrix(BlockKron(Eye(n),D²))
        @test D_xx == blockkron(D², Eye(n))
        @time Δ = D_xx + D_yy

        @test Δ isa BandedBlockBandedMatrix
        @test blockbandwidths(Δ) == subblockbandwidths(Δ) == (1,1)
        @test Δ == blockkron(Matrix(D²), Matrix(I,n,n)) + blockkron(Matrix(I,n,n), Matrix(D²))

        n = 10
        D² = FiniteDifference(n)
        D̃_xx = BlockKron(D², Eye(n))
        @test blockbandwidths(D̃_xx) == (1,1)
        @test subblockbandwidths(D̃_xx) == (0,0)

        V = view(D̃_xx, Block(1,1))
        @test bandwidths(V) == (0,0)

        @test BandedBlockBandedMatrix(D̃_xx) ≈ D_xx

        D̃_yy = BlockKron(Eye(n), D²)
        @test blockbandwidths(D̃_yy) == (0,0)
        @test subblockbandwidths(D̃_yy) == (1,1)

        V = view(D̃_yy, Block(1,1))
        @test bandwidths(V) == (1,1)

        @test BandedBlockBandedMatrix(D̃_yy) ≈ D_yy
    end

    @testset "Diagonal interface" begin
        n = 10
        h = 1/n
        D² = BandedMatrix(0 => Fill(-2,n), 1 => Fill(1,n-1), -1 => Fill(1,n-1))/h^2
        D_xx = BandedBlockBandedMatrix(BlockKron(D², Eye(n)))

        D = Diagonal(randn(n^2))
        @test D_xx + D isa BandedBlockBandedMatrix
        @test blockbandwidths(D_xx + D) == blockbandwidths(D_xx)
        @test subblockbandwidths(D_xx + D) == subblockbandwidths(D_xx)
        @test D_xx + D == Matrix(D_xx) + D

        @test D_xx - D isa BandedBlockBandedMatrix
        @test blockbandwidths(D_xx - D) == blockbandwidths(D_xx)
        @test subblockbandwidths(D_xx - D) == subblockbandwidths(D_xx)
        @test D_xx - D == Matrix(D_xx) - D

        @test D_xx*D == Matrix(D_xx)*D
        @test D_xx*D isa BandedBlockBandedMatrix
        @test blockbandwidths(D_xx*D) == blockbandwidths(D_xx)
        @test subblockbandwidths(D_xx*D) == subblockbandwidths(D_xx)

        @test D*D_xx == D*Matrix(D_xx)
        @test D*D_xx isa BandedBlockBandedMatrix
        @test blockbandwidths(D*D_xx) == blockbandwidths(D_xx)
        @test subblockbandwidths(D*D_xx) == subblockbandwidths(D_xx)
    end

    @testset "Banded Vcat" begin
        A = Vcat(Zeros(1,10), brand(9,10,1,1))
        @test isbanded(A)
        @test bandwidths(A) == (2,0)
        @test MemoryLayout(A) isa ApplyBandedLayout{typeof(vcat)}
        @test BandedMatrix(A) == Array(A) == A
        @test A*A isa MulMatrix
        @test A*A ≈ BandedMatrix(A)*A ≈ A*BandedMatrix(A) ≈ BandedMatrix(A*A)
        @test A[1:5,1:5] isa BandedMatrix
    end

    @testset "Banded Hcat" begin
        A = Hcat(Zeros(10), brand(10,9,1,1))
        @test isbanded(A)
        @test bandwidths(A) == (0,2)
        @test MemoryLayout(A) isa ApplyBandedLayout{typeof(hcat)}
        @test BandedMatrix(A) == Array(A) == A
        @test A*A isa MulMatrix
        @test A*A ≈ BandedMatrix(A)*A ≈ A*BandedMatrix(A) ≈ BandedMatrix(A*A)
        @test A[1:5,1:5] isa BandedMatrix

        A = Hcat(Zeros(10,3), brand(10,9,1,1))
        @test isbanded(A)
        @test bandwidths(A) == (-2,4)
        @test MemoryLayout(A) isa ApplyBandedLayout{typeof(hcat)}
        @test BandedMatrix(A) == Array(A) == A
        @test A[1:5,1:5] isa BandedMatrix
    end

    @testset "resize" begin
        A = brand(4,5,1,1)
        @test LazyBandedMatrices.resize(A,6,5)[1:4,1:5] == A
        @test LazyBandedMatrices.resize(view(A,2:3,2:5),5,5) isa BandedMatrix
        @test LazyBandedMatrices.resize(view(A,2:3,2:5),5,5)[1:2,1:4] == A[2:3,2:5]
    end

    @testset "Lazy banded * Padded" begin
        A = _BandedMatrix(Vcat(BroadcastArray(exp, 1:5)', Ones(1,5)), 5, 1, 0)
        @test MemoryLayout(A) isa BandedColumns{LazyLayout}
        x = Vcat([1,2], Zeros(3))
        @test A*x isa Vcat
        @test A*A*x isa Vcat

        B = PaddedArray(randn(3,4),5,5)
        @test MemoryLayout(A*B) isa PaddedLayout
        @test A*B ≈ Matrix(A)Matrix(B)
        C = BandedMatrix(1 => randn(4))
        @test C*B ≈ Matrix(C)Matrix(B)
        D = BandedMatrix(-1 => randn(4))
        @test D*B ≈ Matrix(D)Matrix(B)

        B.args[2][end,:] .= 0
        C = PaddedArray(randn(3,4),5,5)
        @test muladd!(1.0, A, B, 2.0, deepcopy(C)) ≈ A*B + 2C
    end

    @testset "Lazy banded" begin
        A = _BandedMatrix(Ones{Int}(1,10),10,0,0)'
        B = _BandedMatrix((-2:-2:-20)', 10,-1,1)
        C = Diagonal( BroadcastVector(/, 2, (1:2:20)))
        C̃ = _BandedMatrix(BroadcastArray(/, 2, (1:2:20)'), 10, -1, 1)
        D = MyLazyArray(randn(10,10))
        M = ApplyArray(*,A,A)
        M̃ = ApplyArray(*,randn(10,10),randn(10,10))
        @test MemoryLayout(A) isa BandedRows{OnesLayout}
        @test MemoryLayout(B) isa BandedColumns{UnknownLayout}
        @test MemoryLayout(C) isa DiagonalLayout{LazyLayout}
        @test MemoryLayout(C̃) isa BandedColumns{LazyLayout}
        BC = BroadcastArray(*, B, permutedims(MyLazyArray(Array(C.diag))))
        @test MemoryLayout(BC) isa BroadcastBandedLayout
        @test A*BC isa MulMatrix
        @test BC*B isa MulMatrix
        @test BC*BC isa MulMatrix
        @test C*C̃ isa MulMatrix
        @test C̃*C isa MulMatrix
        @test C̃*D isa MulMatrix
        @test D*C̃ isa MulMatrix
        @test C̃*M isa MulMatrix
        @test M*C̃ isa MulMatrix
        @test C̃*M̃ isa MulMatrix
        @test M̃*C̃ isa MulMatrix
    end

    @testset "banded-block-banded Kron" begin
        n = 4
        h = 1/n
        D² = BandedMatrix(0 => Fill(-2,n), 1 => Fill(1,n-1), -1 => Fill(1,n-1))/h^2

        D_xx = BandedBlockBandedMatrix(BlockKron(D², Eye(n)))
        D_yy = BandedBlockBandedMatrix(BlockKron(Eye(n),D²))
        Δ = BroadcastArray(+, D_xx, D_yy)
        @test MemoryLayout(Δ) isa BroadcastBandedBlockBandedLayout
        @test blockbandwidths(Δ) == (1,1)
        @test subblockbandwidths(Δ) == (1,1)
        @test Δ[Block.(1:4),Block.(2:4)] == Δ[:,5:end]

        @testset "irradic indexing" begin
            B = cache(Δ);
            resizedata!(B,16,1);
            # we do blockwise
            @test B.data[1:16,1:4] == Δ[1:16,1:4]
            resizedata!(B,1,16);
            @test B.datasize == (16,16)
            @test B.data == Δ
        end

        @testset "Cache-block indexing" begin
            B = cache(Δ);
            @test B[Block(1,1)] == B[Block(1),Block(1)] == Δ[Block(1,1)]
            @test B[Block.(1:3),Block.(1:2)] == Δ[Block.(1:3),Block.(1:2)]
            @test B[:, Block(1)] == Δ[:, Block(1)]
            @test B[Block(1), :] == Δ[Block(1), :]
            @test B[Block(1), [1,2,3]] == Δ[Block(1), [1,2,3]]
            @test B[[1,2,3], Block(1)] == Δ[[1,2,3], Block(1)]
        end

        @testset "resizedata!" begin
            B = cache(Δ);
            resizedata!(B,5,5);
            @test B.data[1:5,1:5] == Δ[1:5,1:5]
            @test B == Δ
        end
    end

    @testset "Padded Block" begin
        b = PseudoBlockArray(cache(Zeros(55)),1:10);
        b[10] = 5;
        @test MemoryLayout(b) isa PaddedLayout{DenseColumnMajor}
        @test paddeddata(b) isa PseudoBlockVector
        @test paddeddata(b) == [zeros(9); 5]
    end

    # @testset "Padded columns" begin
    #     B = brand(8,8,1,2)
    #     v = view(B,:,4)
    #     w = view(B,3,:)
    #     @test MemoryLayout(v) isa PaddedLayout
    #     @test_broken MemoryLayout(w) isa PaddedLayout
    #     @test paddeddata(v) isa Vcat
    #     paddeddata(v) == B[:,4]
    # end

    @testset "Banded rot" begin
        A = brand(5,5,1,2)
        R = ApplyArray(rot180, A)
        @test MemoryLayout(R) isa BandedColumns{StridedLayout}
        @test bandwidths(R) == (2,1)
        @test BandedMatrix(R) == R == rot180(Matrix(A)) == rot180(A)

        A = brand(5,4,1,2)
        R = ApplyArray(rot180, A)
        @test MemoryLayout(R) isa BandedColumns{StridedLayout}
        @test bandwidths(R) == (3,0)
        @test BandedMatrix(R) == R == rot180(Matrix(A)) == rot180(A)

        A = brand(5,6,1,-1)
        R = ApplyArray(rot180, A)
        @test MemoryLayout(R) isa BandedColumns{StridedLayout}
        @test bandwidths(R) == (-2,2)
        @test BandedMatrix(R) == R == rot180(Matrix(A)) == rot180(A)

        A = brand(6,5,-1,1)
        R = ApplyArray(rot180, A)
        @test MemoryLayout(R) isa BandedColumns{StridedLayout}
        @test bandwidths(R) == (2,-2)
        @test BandedMatrix(R) == R == rot180(Matrix(A)) == rot180(A)

        B = brand(5,4,1,1)
        R = ApplyArray(rot180, ApplyArray(*, A, B))
        MemoryLayout(R) isa ApplyBandedLayout{typeof(*)}
        @test bandwidths(R) == (4,-2)
        @test R == rot180(A*B)
    end

    @testset "invlayout * structured banded (#21)" begin
        A = randn(5,5)
        B = BroadcastArray(*, brand(5,5,1,1), 2)
        @test A * B ≈ A * Matrix(B)
        @test A \ B ≈ A \ Matrix(B)
    end

    @testset "Triangular bandwidths" begin
        B = brand(5,5,1,2)
        @test bandwidths(ApplyArray(\, Diagonal(randn(5)), B)) == (1,2)
        @test bandwidths(ApplyArray(\, UpperTriangular(randn(5,5)), B)) == (1,4)
        @test bandwidths(ApplyArray(\, LowerTriangular(randn(5,5)), B)) == (4,2)
        @test bandwidths(ApplyArray(\, randn(5,5), B)) == (4,4)
    end

    @testset "Lazy block" begin
        b = PseudoBlockVector(randn(5),[2,3])
        c = BroadcastVector(exp,1:5)
        @test c .* b isa BroadcastVector
        @test b .* c isa BroadcastVector
        @test (c .* b)[Block(1)] == c[1:2] .* b[Block(1)]
    end

    @testset "concat" begin
        @test MemoryLayout(Vcat(1,1)) isa ApplyLayout{typeof(vcat)}
        @test MemoryLayout(Vcat(1,Zeros(5),1)) isa ApplyLayout{typeof(vcat)}

        @test bandwidths(Hcat(1,randn(1,5))) == (0,5)
        @test bandwidths(Vcat(1,randn(5,1))) == (5,0)

        V = Vcat(brand(5,5,1,1), brand(4,5,0,1))
        @test arguments(view(V,:,1:3)) == (V.args[1][:,1:3], V.args[2][:,1:3])
        H = ApplyArray(hvcat, 2, 1, Hcat(1, Zeros(1,10)), Vcat(1, Zeros(10)), Diagonal(1:11))
        @test bandwidths(H) == (1,1)
        H = ApplyArray(hvcat, 2, 1, Hcat(0, Zeros(1,10)), Vcat(0, Zeros(10)), Diagonal(1:11))
        @test bandwidths(H) == (0,0)
        H = ApplyArray(hvcat, (2,2), 1, Hcat(1, Zeros(1,10)), Vcat(1, Zeros(10)), Diagonal(1:11))
        @test_broken bandwidths(H) == (1,1)

        @test bandwidths(Vcat(Diagonal(1:3), Zeros(3,3))) == (0,0)
        @test bandwidths(Hcat(1, Zeros(1,3))) == (0,0)
        c = cache(Zeros(5));
        @test bandwidths(c) == (-1,0)
    end

    @testset "zeros mul" begin
        A = _BandedMatrix(BroadcastVector(exp,1:10)', 10, -1,1)
        @test ArrayLayouts.mul(Zeros(5,10),A) ≡ Zeros(5,10)
        @test ArrayLayouts.mul(A,Zeros(10,5)) ≡ Zeros(10,5)
    end
end

@testset "QR" begin
    A = brand(100_000,100_000,1,1)
    F = qr(A)
    b = Vcat([1,2,3],Zeros(size(A,1)-3))
    @test F.Q'b == apply(*,F.Q',b)
end

include("test_blockkron.jl")
include("test_blockconcat.jl")


text:
using Oceananigans

using Oceananigans.TurbulenceClosures.CATKEVerticalDiffusivities:
    CATKEVerticalDiffusivity,
    MixingLength

using OceanTurbulenceParameterEstimation.EnsembleSimulations: ensemble_column_model_simulation
using OceanTurbulenceParameterEstimation.Observations: FieldTimeSeriesCollector
using JLD2
using GLMakie
using Printf
using Statistics

name = "weak_wind_strong_cooling"
xz_filepath = name * "_xz_slice.jld2"
yz_filepath = name * "_yz_slice.jld2"
xy_filepath = name * "_xy_slice.jld2"
statistics_filepath = name * "_instantaneous_statistics.jld2"

#####
##### Metadata
#####

xy_file = jldopen(xy_filepath)

Nx = xy_file["grid/Nx"]
Ny = xy_file["grid/Ny"]
Nz = xy_file["grid/Nz"]
Hx = xy_file["grid/Hx"]
Hy = xy_file["grid/Hy"]
Hz = xy_file["grid/Hz"]
Lx = xy_file["grid/Lx"]
Ly = xy_file["grid/Ly"]
Lz = xy_file["grid/Lz"]

iterations = parse.(Int, keys(xy_file["timeseries/t"]))
times = [xy_file["timeseries/t/$i"] for i in iterations]
Nt = length(times)

Qᵀ = xy_file["parameters/temperature_flux"]
Qᵘ = xy_file["parameters/momentum_flux"]
f  = xy_file["parameters/coriolis_parameter"]
α  = xy_file["parameters/thermal_expansion_coefficient"]
f  = xy_file["parameters/coriolis_parameter"]
dTdz = xy_file["parameters/dθdz_deep"]

close(xy_file)

H = Hz
chop(a::AbstractArray{<:Any, 1}) = a[1+H:end-H]
chop(a::AbstractArray{<:Any, 2}) = a[1+H:end-H, 1+H:end-H]

function extract_slices(filepath; dims, name="w")
    file = jldopen(filepath)
    slices = [chop(dropdims(file["timeseries/$name/$i"]; dims)) for i in iterations]
    close(file)
    return slices
end

#####
##### Construct 1D model
#####

# Data
u_avg_series = extract_slices(statistics_filepath, dims=(1, 2), name="u")
v_avg_series = extract_slices(statistics_filepath, dims=(1, 2), name="v")
T_avg_series = extract_slices(statistics_filepath, dims=(1, 2), name="T")

# column_ensemble_size = ColumnEnsembleSize(Nz=Nz, ensemble=(1, 1))
# column_ensemble_halo_size = ColumnEnsembleSize(Nz=0, Hz=Hz)

one_dimensional_grid = RectilinearGrid(size=Nz, topology=(Flat, Flat, Bounded), z=(-Lz, 0))

boundary_conditions = (
    u = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵘ)),
    T = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵀ), bottom = GradientBoundaryCondition(dTdz))
)

coriolis = FPlane(; f)
equation_of_state = LinearEquationOfState(; α)
buoyancy = SeawaterBuoyancy(; equation_of_state, constant_salinity=true)
tracers = (:T, :e)

mixing_length = MixingLength(Cᴬu   = 0.0,
                             Cᴬc   = 1.0,
                             Cᴬe   = 0.0,
                             Cᴷu⁻  = 0.072,
                             Cᴷc⁻  = 0.088,
                             Cᴷe⁻  = 0.50,
                             Cᴷuʳ  = 1.0,
                             Cᴷcʳ  = 0.069,
                             Cᴷeʳ  = 0.5,
                             CᴷRiʷ = 0.05,
                             CᴷRiᶜ = 0.1)

surface_tke_flux = SurfaceTKEFlux(CᵂwΔ=4.99, Cᵂu★=2.5)

closure = CATKEVerticalDiffusivity(; Cᴰ=2.44, surface_tke_flux, mixing_length)

model = HydrostaticFreeSurfaceModel(; tracers, buoyancy, coriolis, closure,
                                    boundary_conditions, grid = one_dimensional_grid)

# Initial condition
model.tracers.T .= reshape(T_avg_series[1], 1, 1, Nz)

times = round.(times)
simulation = Simulation(model; Δt=2.0, stop_time=times[end])

collector = FieldTimeSeriesCollector(merge(model.velocities, model.tracers), times)
simulation.callbacks[:collector] = Callback(collector, SpecifiedTimes(times...))

progress(sim) = @info "Sim time: " * prettytime(sim) * ", wall time: " * prettytime(sim.run_wall_time)
simulation.callbacks[:progress] = Callback(progress, IterationInterval(100))

@info "Running 1D simulation..."
start = time_ns()
run!(simulation)
elapsed = 1e-9 * (time_ns() - start)
@info "    ... done (" * prettytime(elapsed) * ")."

#####
##### 3D visualization
#####

grid = RectilinearGrid(size=(Nx, Ny, Nz), halo=(Hx, Hy, Hz), x=(0, Lx), y=(0, Ly), z=(-Lz, 0))

w_yz_series = extract_slices(yz_filepath, dims=1, name="w")
w_xz_series = extract_slices(xz_filepath, dims=2, name="w")
w_xy_series = extract_slices(xy_filepath, dims=3, name="w")

T_yz_series = extract_slices(yz_filepath, dims=1, name="T")
T_xz_series = extract_slices(xz_filepath, dims=2, name="T")
T_xy_series = extract_slices(xy_filepath, dims=3, name="T")

x, y, z = nodes((Center, Center, Center), grid)

x_xz = repeat(x, 1, Nz)
y_xz = 0.995 * Ly * ones(Nx, Nz)
z_xz = repeat(reshape(z, 1, Nz), Nx, 1)

x_yz = 0.995 * Lx * ones(Ny, Nz)
y_yz = repeat(y, 1, Nz)
z_yz = repeat(reshape(z, 1, Nz), grid.Ny, 1)

# Slight displacements to "stitch" the cube together
x_xy = x
y_xy = y
z_xy = - 0.001 * Lz * ones(grid.Nx, grid.Ny)

#####
##### Animate!
#####

fig = Figure(resolution=(1800, 600))

n = Observable(1)

###
### Vertical velocity
###

w_yz = @lift w_yz_series[$n]
w_xz = @lift w_xz_series[$n]
w_xy = @lift w_xy_series[$n]

w_max = maximum(maximum.(abs, w_xz_series[end]))

colormap = :balance
colorrange = (-w_max/2, w_max/2)

ax_w = Axis3(fig[1:7, 1:4], aspect=:data, xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", azimuth=0.85, elevation=0.70)
ax_T = Axis3(fig[1:7, 5:8], aspect=:data, xlabel="x (m)", ylabel="y (m)", zlabel="z (m)", azimuth=0.86, elevation=0.38)

sfc_w = surface!(ax_w, x_xz, y_xz, z_xz; color=w_xz, colormap, colorrange)
        surface!(ax_w, x_yz, y_yz, z_yz; color=w_yz, colormap, colorrange)
        surface!(ax_w, x_xy, y_xy, z_xy; color=w_xy, colormap, colorrange)

###
### Temperature
###

T_yz = @lift T_yz_series[$n]
T_xz = @lift T_xz_series[$n]
T_xy = @lift T_xy_series[$n]

colormap = :oslo

# Compute maximum temperature for colorrange with a moving average
ΔT = 0.02 # Absolute range
Δn = 5 # Half-width of moving average
Tf_max = maximum(T_xy_series[end])
T0_max = maximum(T_xy_series[1])

colorrange = @lift begin
    if $n < Δn + 1
        T_max = T0_max
    elseif $n < Nt - Δn
        T_max = mean([maximum(T_xy_series[nn]) for nn in $n-Δn:$n+Δn])
    else
        T_max = Tf_max
    end

    T_min = T_max - ΔT

    (T_min, T_max)
end

sfc_T = surface!(ax_T, x_xz, y_xz, z_xz; color=T_xz, colormap, colorrange)
        surface!(ax_T, x_yz, y_yz, z_yz; color=T_yz, colormap, colorrange)
        surface!(ax_T, x_xy, y_xy, z_xy; color=T_xy, colormap, colorrange)

cp = fig[8, 1:4] = Colorbar(fig, sfc_w, vertical=false, flipaxis=false, label="Vertical velocity (m s⁻¹)")
cp = fig[8, 5:8] = Colorbar(fig, sfc_T, vertical=false, flipaxis=false, label="Temperature (ᵒC)")

ax_T_avg = Axis(fig[2:6, 9:10], xlabel = "Average temperature (ᵒC)", ylabel = "z (m)")
ax_u_avg = Axis(fig[2:6, 11:12], xlabel = "Average velocity \n components (m s⁻¹)", ylabel = "z (m)")

u_model = @lift interior(collector.field_time_serieses.u[$n])[1, 1, :]
v_model = @lift interior(collector.field_time_serieses.v[$n])[1, 1, :]
T_model = @lift interior(collector.field_time_serieses.T[$n])[1, 1, :]

u_LES = @lift u_avg_series[$n]
v_LES = @lift v_avg_series[$n]
T_LES = @lift T_avg_series[$n]

T_model_lines = lines!(ax_T_avg, T_model, z, label="Calibrated CATKE model")
T_LES_lines = lines!(ax_T_avg, T_LES, z, label="Large eddy simulation data")

model_color = T_model_lines.attributes.color.val
LES_color = T_LES_lines.attributes.color.val

lines!(ax_u_avg, u_model, z, color=model_color, label="u")
lines!(ax_u_avg, v_model, z, color=model_color, linestyle=:dash, label="v")

lines!(ax_u_avg, u_LES, z, color=LES_color, label="u")
lines!(ax_u_avg, v_LES, z, color=LES_color, linestyle=:dash, label="v")

xlims!(ax_u_avg, -0.1, 0.1)

Legend(fig[7:8, 9:12], ax_T_avg)
axislegend(ax_u_avg, position=:rb, merge=true)

title = @lift "Ocean surface boundary layer at t = " * prettytime(times[$n])
lab = Label(fig[0, :], title, textsize=18)

record(fig, "$name.mp4", 1:Nt; framerate=16) do nn
    @info "Drawing frame $nn of $Nt..."
    n[] = nn
end


text:
"""
    ComposedTransform(outer, inner)

A subtype of `CoordinateTransform` that lazily combines index `outer` and `inner`. Indexing a
`ComposedTransform` whith `i` is equivalent to `outer[inner[i]]`.
"""
struct ComposedTransform{NI,NO,O,I} <: CoordinateTransform{NI,NO}
    outer::O
    inner::I

    function ComposedTransform(o::CoordinateTransform{O,N}, i::CoordinateTransform{N,I}) where {O,N,I}
        @assert !(i isa ComposedTransform)
        new{O,I,typeof(o),typeof(i)}(o, i)
    end
end

@inline (t::ComposedTransform)(x) = getfield(t, :outer)(getfield(t, :inner)(x))

function Base.show(io::IO, m::MIME"text/plain", @nospecialize(x::ComposedTransform))
    show(io, m, x.outer)
    print(io, " ∘ ")
    show(io, m, x.inner)
end

@inline (t::CoordinateTransform)(::IdentityTransform) = t
@inline (t1::CoordinateTransform)(t2::CoordinateTransform) = ComposedTransform(t2, t1)

@inline function (t1::VecPermute)(t2::StrideTransform{1})
    s = getfield(getfield(t1, :strides), 1)
    StrideTransform((s, s), (static(1), getfield(getfield(t1, :offsets), 1)))
end

@inline function (t2::Permute{I1,I2,N})(t1::StrideTransform{N}) where {N,I1,I2}
    StrideTransform{N}(t2(getfield(x, :strides)), t2(getfield(x, :offsets)))
end

function (v::View{NI,NO,I})(t::StrideTransform{NO,S,O}) where {NI,NO,I,S,O}
    _combined_sub_strides(ArrayInterface.stride_preserving_index(I), t, v)
end

@inline function _combined_sub_strides(::False, t::StrideTransform, v::View{NI,NO,I}) where {NI,NO,I}
    __map_strides(View(_map_strides(
        Static.eachop(ArrayInterface._ndims_index, Static.nstatic(Val(known_length(I))), I),
        getfield(v, :indices),
        getfield(t, :strides),
        getfield(t, :offsets)
        )))
end
@inline function __map_strides(v::View{<:Any,O}) where {O}
    ComposedTransform(ComposedTransform(Offset(static(1)), ReduceTransform{O,typeof(+)}(+)), v)
end
@generated function _map_strides(::NDimsIndex, inds::I, s::S, o::O) where {NDimsIndex,I,S,O}
    t = Expr(:tuple)
    dim = 0
    ndi = known(NDimsIndex)
    for i in 1:length(ndi)
        indsexpr = :(@inbounds(getfield(inds, $i)))
        stride_expr = Expr(:tuple)
        offset_expr = Expr(:tuple)
        for j in 1:ndi[i]
            dim += 1
            push!(stride_expr.args, :(@inbounds(getfield(s, $dim))))
            push!(offset_expr.args, :(@inbounds(getfield(o, $dim))))
        end
        push!(t.args, :(StrideTransform($stride_expr, $offset_expr)($indsexpr)))
    end
    Expr(:block, Expr(:meta, :inline), t)
end
@inline function _combined_sub_strides(::True,
    x::StrideTransform{NO,S,O},
    i::View{NI,NO,I}
) where {NI,NO,S,O,I}
    pdims = ArrayInterface._to_sub_dims(I)
    o = getfield(x, :offsets)
    s = getfield(x, :strides)
    inds = getfield(i, :indices)
    out = StrideTransform(
        eachop(getmul, pdims, map(maybe_static_step, inds), s),
        Static.permute(o, pdims)
    )
    ComposedTransform(Offset(Static.reduce_tup(+, map(*, map(_diff, inds, o), s))), out)
end
getmul(x::Tuple, y::Tuple, ::StaticInt{i}) where {i} = getfield(x, i) * getfield(y, i)
@inline _diff(::Base.Slice, ::Any) = Zero()
@inline _diff(x::AbstractRange, o) = static_first(x) - o
@inline _diff(x::Integer, o) = x - o


text:
include("../src/hypergraph-affinity-functions.jl")
using StatsBase

dataset = "congress-bills"
r = 25
H, classes = read_hypergraph_data(dataset,r)
# 1 = democrat
# 2 = republican
m,n = size(H)
classes = classes .- 1
# 0 = democrat = 2 mod 0
# 1 = republican
class1 = "Republican"
class2 = "Democrat"

Hyp = incidence2elist(H)

## Bootstrap
alpha, B1, B2, H1, H2, tR1, tR2, N = Hypergraph_to_Scores(H,classes,15)

for k =  5:5:15
B = 100
TR1 = tR1[k,1:k]  # True affinity scores
TR0 = tR2[k,1:k]

TA1 = H1[k,1:k]  # True affinity scores
TA0 = H2[k,1:k]

# Bootstraps
R0, R1, MR1, MR0, SR1, SR0, A0, A1, MA1, MA0, SA1, SA0 = Bootstrap_Affinities(H,k,B,classes)

E1 = vec(maximum(R1,dims = 1))
E1m = vec(minimum(R1,dims = 1))
E0 = vec(maximum(R0,dims = 1))
E0m = vec(minimum(R0,dims = 1))

## Plot Ratio Scores
lw = 2.5
ms = 8
gfs = 15
tfs = 15
titlesize = 14
s1 = 400
s2 = 300

sc = 2
color1 = :red
color2 = :blue
title = "$k-person bills"
title = ""
leg = false
xlab = "Affinity type t"
ylab = "Affinity / Baseline "
p = plot()
p = Plots.plot([1; k],[1; 1],fontfamily = "helvetica",linewidth = 2,yscale = :log10,legend = false,grid = false, size = (s1,s2),color = :gray)
title!(p,title,titlefontsize = titlesize)
xaxis!(p,xlab, xticks = k/5:k/5:k,tickfontsize=tfs,guidefontsize = gfs)
yaxis!(p,ylab,tickfontsize=tfs,guidefontsize = gfs)

# [M0-2*SE0 M0+2*SE]
# [E0m E0]
# [E1m E1]
p = plot!(1:k,[TR0 TR0], fillrange=[MR0-sc*SR0 MR0+sc*SR0], fillalpha=0.3, color = color2, linewidth = lw, markersize = ms,
    markershape = :circle, markerstrokecolor = color2, label = class2)
p = plot!(1:k,[TR1 TR1], fillrange=[MR1-sc*SR1 MR1+sc*SR1],fillalpha=0.3,color = color1, linewidth = lw, markersize = ms,
    markershape = :circle,markerstrokecolor = color1, label = class1)
savefig("Figures/Ratios_$(k)_Congress_Bootstrap.pdf")

## Plot just affinity scores
lw = 2.5
ms = 8
gfs = 15
tfs = 15
titlesize = 14
s1 = 400
s2 = 300

color1 = :red
color2 = :blue
title = "$k-person bills"
title = ""
leg = false
xlab = "Affinity type t"
ylab = "Affinity"

p = Plots.plot(fontfamily = "helvetica",linewidth = 2,yscale = :identity,legend = false,grid = false, size = (s1,s2),color = :gray)
title!(p,title,titlefontsize = titlesize)
xaxis!(p,xlab, xticks = k/5:k/5:k,tickfontsize=tfs,guidefontsize = gfs)
yaxis!(p,ylab,tickfontsize=tfs,guidefontsize = gfs)

# [M0-2*SE0 M0+2*SE]
# [E0m E0]
# [E1m E1]
p = plot!(1:k,[TA0 TA0], fillrange=[MA0-sc*SA0 MA0+sc*SA0], fillalpha=0.3, color = color2, linewidth = lw, markersize = ms,
    markershape = :circle, markerstrokecolor = color2, label = class2)
p = plot!(1:k,[TA1 TA1], fillrange=[MA1-sc*SA1 MA1+sc*SA1],fillalpha=0.3,color = color1, linewidth = lw, markersize = ms,
    markershape = :circle,markerstrokecolor = color1, label = class1)
savefig("Figures/Affinities_$(k)_Congress_Bootstrap.pdf")
end

max_stars_repo_name: nveldt/HypergraphHomophily

max_issues_repo_name: nveldt/HypergraphHomophily

max_forks_repo_name: nveldt/HypergraphHomophily


text:
using UUIDs

mutable struct Notebook
    path::String

    "Cells are ordered in a `Notebook`, and this order can be changed by the user. Cells will always have a constant UUID."
    cells::Array{Cell,1}

    uuid::UUID
    combined_funcdefs::Dict{Symbol, SymbolsState}

    # buffer must contain all undisplayed outputs
    pendingupdates::Channel

    executetoken::Channel
end
# We can keep 128 updates pending. After this, any put! calls (i.e. calls that push an update to the notebook) will simply block, which is fine.
# This does mean that the Notebook can't be used if nothing is clearing the update channel.
Notebook(path::String, cells::Array{Cell,1}, uuid) = let
    et = Channel{Nothing}(1)
    put!(et, nothing)
    Notebook(path, cells, uuid, Dict{Symbol, SymbolsState}(), Channel(128), et)
end
Notebook(path::String, cells::Array{Cell,1}) = Notebook(path, cells, uuid1())
Notebook(cells::Array{Cell,1}) = Notebook(tempname() * ".jl", cells)

function selectcell_byuuid(notebook::Notebook, uuid::UUID)::Union{Cell,Nothing}
    cellIndex = findfirst(c->c.uuid == uuid, notebook.cells)
    if cellIndex === nothing
        @warn "Requested non-existing cell with UUID $(uuid)\nTry refreshing the page in your browser."
        return nothing
    end
    notebook.cells[cellIndex]
end

# We use a creative delimiter to avoid accidental use in code
_uuid_delimiter = "# ╔═╡ "
_order_delimited = "# ○ "
_cell_appendix = "\n\n"

emptynotebook(path) = Notebook(path, [Cell("")])
emptynotebook() = emptynotebook(tempname() * ".jl")

function samplenotebook()
    cells = Cell[]

    push!(cells, Cell("100*a + b"))
    push!(cells, Cell("a = 1"))
    push!(cells, Cell("b = let\n\tx = a + a\n\tx*x\nend"))
    push!(cells, Cell("html\"
Hoi!
\n
My name is kiki
\""))
    push!(cells, Cell("""md"# Cześć!
    My name is **baba** and I like \$\\LaTeX\$ _support!_

    \$\$\\begin{align}
    \\varphi &= \\sum_{i=1}^{\\infty} \\frac{\\left(\\sin{x_i}^2 + \\cos{x_i}^2\\right)}{i^2} \\\\
    b &= \\frac{1}{2}\\,\\log \\exp{\\varphi}
    \\end{align}\$\$

    ### The spectacle before us was indeed sublime.
    Apparently we had reached a great height in the atmosphere, for the sky was a dead black, and the stars had ceased to twinkle. By the same illusion which lifts the horizon of the sea to the level of the spectator on a hillside, the sable cloud beneath was dished out, and the car seemed to float in the middle of an immense dark sphere, whose upper half was strewn with silver. Looking down into the dark gulf below, I could see a ruddy light streaming through a rift in the clouds."
    """))

    Notebook(tempname() * ".jl", cells)
end

function save_notebook(io, notebook::Notebook)
    write(io, "### A Pluto.jl notebook ###\n")
    write(io, "# " * PLUTO_VERSION_STR * "\n")

    # TODO: this can be optimised by caching the topological order:
    # maintain cache with ordered UUIDs
    # whenever a run_reactive is done, move the found cells **up** until they are in one group, and order them topologcally within that group. Errable cells go to the bottom.

    # we first move cells to the front if they call an import
    # MergeSort because it is a stable sort: leaves cells in order if they are in the same category
    prelim_order = sort(notebook.cells, alg=MergeSort, by=(c -> !isempty(c.module_usings)))
    # the next call took 2ms for a small-medium sized notebook: (so not too bad)
    celltopology = dependent_cells(notebook, prelim_order)

    cells_ordered = union(celltopology.runnable, keys(celltopology.errable))

    for c in cells_ordered
        write(io, _uuid_delimiter * string(c.uuid) * "\n")
        write(io, c.code)
        write(io, _cell_appendix)
    end

    write(io, _uuid_delimiter * "Cell order:" * "\n")
    for c in notebook.cells
        write(io, _order_delimited * string(c.uuid) * "\n")
    end
end

function save_notebook(notebook::Notebook, path::String)
    open(path, "w") do io
        save_notebook(io, notebook)
    end
end

save_notebook(notebook::Notebook) = save_notebook(notebook, notebook.path)

function load_notebook_nobackup(io, path)
    firstline = String(readline(io))

    if firstline != "### A Pluto.jl notebook ###"
        @error "File is not a Pluto.jl notebook"
    end

    file_VERSION_STR = readline(io)[3:end]
    if file_VERSION_STR != PLUTO_VERSION_STR
        @warn "Loading a notebook saved with Pluto $(file_VERSION_STR). This is Pluto $(PLUTO_VERSION_STR)."
    end

    collected_cells = Dict()

    # ignore first bits of file
    readuntil(io, _uuid_delimiter)

    last_read = ""
    while !eof(io)
        uuid_str = String(readline(io))
        if uuid_str == "Cell order:"
            break
        else
            uuid = UUID(uuid_str)
            code = String(readuntil(io, _uuid_delimiter))
            # Change windows line endings to linux; remove the cell appendix.
            code_normalised = replace(code, "\r\n" => "\n")[1:end - ncodeunits(_cell_appendix)]

            read_cell = Cell(uuid, code_normalised)

            collected_cells[uuid] = read_cell
        end
    end

    ordered_cells = Cell[]
    while !eof(io)
        uuid_str = String(readline(io))
        if startswith(uuid_str, _order_delimited)
            uuid = let
                # Because we support Unicode, this is not just `length(_order_delimited) + 1`.
                uuid_index = ncodeunits(_order_delimited) + 1
                UUID(uuid_str[uuid_index:end])
            end
            push!(ordered_cells, collected_cells[uuid])
        end
    end

    Notebook(path, ordered_cells)
end

function load_notebook_nobackup(path::String)
    local loaded
    open(path, "r") do io
        loaded = load_notebook_nobackup(io, path)
    end
    loaded
end

function load_notebook(path::String)
    backupPath = path
    while isfile(backupPath)
        backupPath = backupPath * ".backup"
    end
    cp(path, backupPath)

    loaded = load_notebook_nobackup(path)

    save_notebook(loaded)

    if only_versions_or_lineorder_differ(path, backupPath)
        rm(backupPath)
    else
        @warn "Old Pluto notebook might not have loaded correctly. Backup saved to: " backupPath
    end

    loaded
end

function move_notebook(notebook::Notebook, newpath::String)
    # Will throw exception and return if anything goes wrong, so at least one file is guaranteed to exist.
    oldpath = notebook.path
    save_notebook(notebook, oldpath)
    save_notebook(notebook, newpath)
    @assert only_versions_differ(oldpath, newpath)
    notebook.path = newpath
    rm(oldpath)
end

"Check if two savefiles are identical, up to their version numbers and a possible line shuffle.

If a notebook has not yet had all of its cells run, we can't deduce the topological cell order."
function only_versions_or_lineorder_differ(pathA::AbstractString, pathB::AbstractString)
    Set(readlines(pathA)[3:end]) == Set(readlines(pathB)[3:end])
end

function only_versions_differ(pathA::AbstractString, pathB::AbstractString)
    readlines(pathA)[3:end] == readlines(pathB)[3:end]
end


text:
using Plots
using Random
Random.seed!(1)

include_ddp()

# Model
include_model("double_integrator")

function f(model::DoubleIntegratorContinuous, x, u, w)
    [x[2]; u[1]]
end

function fd(model::DoubleIntegratorContinuous{Midpoint, FixedTime}, x, u, w, h, t)
	x + h * f(model, x + 0.5 * h * f(model, x, u, w), u, w)
end

model = DoubleIntegratorContinuous{Midpoint, FixedTime}(2, 1, 1)
n = model.n
m = model.m

# Time
T = 11
h = 0.1

# Initial conditions, controls, disturbances
x1 = [1.0; 0.0]
x_ref = [[0.0; 0.0] for t = 1:T]
xT = [0.0; 0.0]
ū = [1.0 * rand(model.m) for t = 1:T-1]
u_ref = [zeros(model.m) for t = 1:T-1]
w = [zeros(model.d) for t = 1:T-1]

# Rollout
x̄ = rollout(model, x1, ū, w, h, T)

# Objective
Q = [(t < T ? 0.0 * Diagonal([1.0; 1.0]) : 0.0 * Diagonal([1.0; 1.0])) for t = 1:T]
q = [-2.0 * Q[t] * x_ref[t] for t = 1:T]
R = [Diagonal(1.0e-5 * ones(model.m)) for t = 1:T-1]
r = [zeros(model.m) for t = 1:T-1]

obj = StageCosts([QuadraticCost(h * Q[t], h * q[t],
	t < T ? h * R[t] : nothing, t < T ? h * r[t] : nothing) for t = 1:T], T)

function g(obj::StageCosts, x, u, t)
	T = obj.T
    if t < T
		Q = obj.cost[t].Q
		q = obj.cost[t].q
	    R = obj.cost[t].R
		r = obj.cost[t].r
        return x' * Q * x + q' * x + u' * R * u + r' * u
    elseif t == T
		Q = obj.cost[T].Q
		q = obj.cost[T].q
        return x' * Q * x + q' * x
    else
        return 0.0
    end
end

# Constraints
ul = [-10.0]
uu = [10.0]
p = [t < T ? 2 * m : n for t = 1:T]
info_t = Dict(:ul => ul, :uu => uu, :inequality => (1:2 * m))
info_T = Dict(:xT => xT)
con_set = [StageConstraint(p[t], t < T ? info_t : info_T) for t = 1:T]

function c!(c, cons::StageConstraints, x, u, t)
	T = cons.T
	p = cons.con[t].p

	if t < T
		ul = cons.con[t].info[:ul]
		uu = cons.con[t].info[:uu]
		c .= [ul - u; u - uu]
		nothing
	else
		c .= x - cons.con[T].info[:xT]
	end
end

prob = problem_data(model, obj, con_set, copy(x̄), copy(ū), w, h, T)

# Solve
@time constrained_ddp_solve!(prob,
	max_iter = 1000, max_al_iter = 10,
	ρ_init = 1.0, ρ_scale = 10.0,
	con_tol = 1.0e-5)

x, u = current_trajectory(prob)
x̄, ū = nominal_trajectory(prob)

objective(obj, x, u)
objective(prob.m_data.obj, x, u)

# Visualize
using Plots
t = range(0, stop = h * (T - 1), length = T)
plot(t, hcat([[0.0; 0.0] for t = 1:T]...)',
    width = 2.0, color = :black, label = "")
plot!(t, hcat([x̄[t][1:2] for t = 1:T]...)', color = :magenta, label = "")
plot(t, hcat(ū..., u[end])', linetype = :steppost)

# free final time
function f(model::DoubleIntegratorContinuous, x, u, w)
    [x[2]; u[1]]
end

function fd(model::DoubleIntegratorContinuous{Midpoint, FixedTime}, x, u, w, h, t)
	if t == 1
		# h = max(0.001, u[2])
		h = u[2] * u[2]
    	return [x + h * f(model, x + 0.5 * h * f(model, x, view(u, 1:1), w), view(u, 1:1), w); u[2]]
	else
		h = x[3] * x[3]
		return [view(x, 1:2) + h * f(model, view(x, 1:2) + 0.5 * h * f(model, view(x, 1:2), u, w), u, w); x[3]]
	end
end

model = DoubleIntegratorContinuous{Midpoint, FixedTime}(2, 1, 1)
n = [t == 1 ? model.n : model.n + 1 for t = 1:T]
m = [t == 1 ? model.m + 1 : model.m for t = 1:T]

ū_alt = [t == 1 ? [ū[t]; sqrt(h)] : ū[t] for t = 1:T-1]

# Rollout
x̄ = rollout(model, x1, ū_alt, w, nothing, T)

# Objective
obj = StageCosts([NonlinearCost() for t = 1:T], T)

function g(obj::StageCosts, x, u, t)
	T = obj.T
    if t < T
		t == 1 ? (h = u[2] * u[2]) : (h = x[3] * x[3])
        return h * (view(x, 1:2)' * Q[t] * view(x, 1:2) + q[t]' * view(x, 1:2) + view(u, 1:1)' * R[t] * view(u, 1:1) + r[t]' * view(u, 1:1)) + 1.0 * h
    elseif t == T
		h = x[3] * x[3]
        return view(x, 1:2)' * Q[t] * view(x, 1:2) + q[t]' * view(x, 1:2) + 1.0 * h
    else
        return 0.0
    end
end

# Constraints
hl = 0.1 * sqrt(h)
hu = 1.0 * sqrt(h)
p = [t == 1 ? 4 : (t == T ? 2 + 2 : 2 + 2) for t = 1:T]
info_1 = Dict(:ul => [ul[1]; hl], :uu => [uu[1]; hu], :inequality => (1:4))
info_t = Dict(:ul => ul, :uu => uu, :hl => [hl], :hu => [hu], :inequality => (1:4))
info_T = Dict(:xT => xT, :hl => [hl], :hu => [hu], :inequality => (3:4))
con_set = [StageConstraint(p[t], t < T ? (t == 1 ? info_1 : info_t) : info_T) for t = 1:T]

function c!(c, cons::StageConstraints, x, u, t)
	T = cons.T
	if t == 1
		c[1:2] .= view(u, 1:2) - cons.con[t].info[:uu]
		c[3:4] .= cons.con[t].info[:ul] - view(u, 1:2)
	end
	if t > 1 && t < T
		c[1:1] .= view(u, 1:1) - cons.con[t].info[:uu]
		c[2:2] .= cons.con[t].info[:ul] - view(u, 1:1)
		c[3:3] .= view(x, 3:3) - cons.con[t].info[:hu]
		c[4:4] .= cons.con[t].info[:hl] - view(x, 3:3)
	end
	if t == T
		c[1:2] .= view(x, 1:2) - cons.con[t].info[:xT]
		c[3:3] .= view(x, 3:3) - cons.con[t].info[:hu]
		c[4:4] .= cons.con[t].info[:hl] - view(x, 3:3)
	end
end

prob = problem_data(model, obj, con_set, copy(x̄), copy(ū_alt), w, h, T, n = n, m = m)

# Solve
@time constrained_ddp_solve!(prob,
    max_iter = 1000, max_al_iter = 10,
	ρ_init = 1.0, ρ_scale = 10.0)

# # Solve
# @time ddp_solve!(prob,
#     max_iter = 100, verbose = true)
x, u = current_trajectory(prob)
x̄, ū = nominal_trajectory(prob)
objective(obj, x, u)
objective(prob.m_data.obj, x, u)
h_opt = u[1][2] * u[1][2]
@show h_opt * (T - 1)
@show h_opt * (T - 1)

# Visualize
using Plots
t = range(0, stop = h_opt * (T - 1), length = T)
plot(t, hcat([xt[1:2] for xt in x]...)[:, 1:T]', color = :black, label = "")
plot(t, hcat([ut[1:1] for ut in u]..., u[end])', color = :black, label = "", linetype = :steppost)

max_issues_repo_name: thowell/DirectMotionPlanning


text:
"""
Utilities for distances between coordinates on the sphere and matching
sets of coordinates.
"""
module Coordinates

import NearestNeighbors: KDTree, knn

"""
    angular_separation(λ_1, ϕ_1, λ_2, ϕ_2)

Angular separatation in *degrees* between two points on the sphere at
longitude λ and latitude ϕ, both given in *degrees*.
"""
function angular_separation(λ_1, ϕ_1, λ_2, ϕ_2)
    Δλ = λ_2 - λ_1
    sin_Δλ = sind(Δλ)
    cos_Δλ = cosd(Δλ)
    sin_ϕ1 = sind(ϕ_1)
    sin_ϕ2 = sind(ϕ_2)
    cos_ϕ1 = cosd(ϕ_1)
    cos_ϕ2 = cosd(ϕ_2)
    return rad2deg(atan2(hypot(cos_ϕ2 * sin_Δλ,
                               cos_ϕ1 * sin_ϕ2 - sin_ϕ1 * cos_ϕ2 * cos_Δλ),
                         sin_ϕ1 * sin_ϕ2 + cos_ϕ1 * cos_ϕ2 * cos_Δλ))
end

"""
convert (lon, lat) -> (x, y, z) unit vector for multiple coordinates
where spherical coordinates are given in *degrees*.
"""
function _spherical_to_cartesian(lon::AbstractVector, lat::AbstractVector)
    length(lon) == length(lat) || error("lengths of `lon` and `lat` must match")

    T = promote_type(eltype(lon), eltype(lat))
    result = Array{T}(3, length(lon))
    for i in eachindex(lon)
        coslat = cosd(lat[i])
        result[1, i] = coslat*cosd(lon[i])  # x
        result[2, i] = coslat*sind(lon[i])  # y
        result[3, i] = sind(lat[i])  # z
    end
    return result
end

"""
    match_coordinates_sky(lon1, lat1, lon2, lat2; nneighbor=1)

For each position in the set of coordinates `(lon1, lat1)`, find the
`nneighbor`th nearest on-sky match in the set of coordinates `(lon2,
lat2)`. Both sets of coordinates are in *degrees*.  Returns the indicies
of each closest match (`Vector{Int}` of length `length(lon1)`) and the
corresponding distances in degrees (`Vector{Float64}`). `nneighbor=2` can be
used to match a catalog to itself.

# Notes

This function correctly accounts for wrap around and behavior near the poles.
Works by transforming spherical coordinates to cartesian and using a KD tree
to efficiently find nearest neighbors.

An alternative approach that *might* be faster would be to use a hierarchical
triangular mesh (HTM) or HEALPix to spatially index coordinates and then only
search nearby indicies for matches.

This function could potentially be incorporated into the SkyCoords
package eventually, but for now it's unclear where it best fits.
"""
function match_coordinates(lon1::AbstractVector, lat1::AbstractVector,
                           lon2::AbstractVector, lat2::AbstractVector;
                           nneighbor::Int=1)
    cartcoords1 = _spherical_to_cartesian(lon1, lat1)
    cartcoords2 = _spherical_to_cartesian(lon2, lat2)
    kdtree = KDTree(cartcoords2; leafsize = 10)
    all_idxs, all_cartdists = knn(kdtree, cartcoords1, nneighbor)

    # extract nneighbor-th index
    idxs = [v[nneighbor] for v in all_idxs]

    # extract nneighbor-th distance and convert from cartesian to spherical
    dists = [2.0 * asind(v[nneighbor] / 2.0) for v in all_cartdists]

    return idxs, dists
end

end  # module


text:
using CairoMakie: Cairo
using CairoMakie
using JLD2

FILE_NAME = "NDE_18sim_windcooling_windheating_18simBFGST0.8nograd_divide1f5_gradient_smallNN_leakyrelu_rate_2e-4_T0.8_1e-4"
DATA_PATH = joinpath(pwd(), "extracted_training_output", "$(FILE_NAME)_extracted.jld2")
OUTPUT_PATH = "C:\\Users\\xinle\\Documents\\OceanParameterizations.jl"

file = jldopen(DATA_PATH, "r")
train_parameters = file["training_info/parameters"]
close(file)

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

function diffusivity(Ri, ν₀, ν₋, ΔRi, Riᶜ, Pr)
    tanh_step(x) = (1 - tanh(x)) / 2
    ν = ν₀ + ν₋ * tanh_step((Ri - Riᶜ) / ΔRi)
    return ν
end

function diffusivity_T(Ri, ν₀, ν₋, ΔRi, Riᶜ, Pr)
    return diffusivity(Ri, ν₀, ν₋, ΔRi, Riᶜ, Pr) / Pr
end

Ris = range(-2, 2, length=1000)

ν = diffusivity.(Ris, ν₀, ν₋, ΔRi, Riᶜ, Pr)
ν_T = diffusivity_T.(Ris, ν₀, ν₋, ΔRi, Riᶜ, Pr)

fig = Figure(resolution=(1500, 750))
ax = fig[1,1] = CairoMakie.Axis(fig, yscale=log10, xlabel="Ri", ylabel="Diffusivity/ m² s⁻¹")

ν_line = CairoMakie.lines!(ax, Ris, ν)
ν_T_line = CairoMakie.lines!(ax, Ris, ν_T)

axislegend(ax, [ν_line, ν_T_line], ["Momentum Diffusivity", "Temperature Diffusivity"])
fig

save(joinpath("C:\\Users\\xinle\\Documents\\OceanParameterizations.jl", "$(FILE_NAME)_mpp.png"), fig, px_per_unit=4)


text:
# Structs and functions related to the various iterators in eccodes

export eachpoint

mutable struct codes_keys_iterator
end

struct EachKey
    ptr::Ptr{codes_keys_iterator}
    msg::Message
end

function Base.iterate(iter::EachKey, state=())
    next = ccall((:codes_keys_iterator_next, eccodes), Cint, (Ptr{codes_keys_iterator},), iter.ptr)
    if next == 0
        destroy(iter)
        return nothing
    else
        element = ccall((:codes_keys_iterator_get_name, eccodes), Cstring,
                        (Ptr{codes_keys_iterator},),
                        iter.ptr)
        elestring = unsafe_string(element)

        # Skip unreadable keys and keys with unsupported types
        native = Tuple{}
        try
            native = getnativetype(iter.msg, elestring)
        catch e
            return iterate(iter, ())
        end
        if native == Tuple{}
            return iterate(iter, ())
        else
            return (elestring, ())
        end
    end
end

Base.IteratorSize(::Type{EachKey}) = Base.SizeUnknown()
Base.eltype(::Type{EachKey}) = String

function destroy(iter::EachKey)
    err = ccall((:codes_keys_iterator_delete, eccodes), Cint, (Ptr{codes_keys_iterator},), iter.ptr)
    errorcheck(err)
end

struct EachValue
    keyiter::EachKey
end

function Base.iterate(iter::EachValue, state=())
    ret = iterate(iter.keyiter)
    if isnothing(ret)
        nothing
    else
        (iter.keyiter.msg[ret[1]], ())
    end
end

Base.IteratorSize(::Type{EachValue}) = Base.SizeUnknown()

function Base.iterate(iter::Message, state=(keys(iter),))
    keyitr = state[1]
    ret = iterate(keyitr)
    if isnothing(ret)
        nothing
    else
        (ret[1] => iter[ret[1]], (keyitr,))
    end
end

Base.IteratorSize(::Type{Message}) = Base.SizeUnknown()
Base.eltype(::Type{Message}) = Pair{String, Any}

"""
    keys(message::Message)

Iterate through each key in the message.

# Example
```Julia
for key in keys(message)
    println(key)
end
```
"""
function Base.keys(handle::Message)
    ptr = ccall((:codes_keys_iterator_new, eccodes), Ptr{codes_keys_iterator},
          (Ptr{codes_handle}, Culong, Cstring), handle.ptr, CODES_KEYS_ITERATOR_ALL_KEYS, C_NULL)
    EachKey(ptr, handle)
end

"""
    values(message::Message)

Iterate through each value in the message.

# Example
```Julia
for v in values(message)
    println(v)
end
```
"""
function Base.values(handle::Message)
    keyiter = keys(handle)
    EachValue(keyiter)
end

mutable struct codes_iterator
end

struct GribIterator
    ptr::Ptr{codes_iterator}
    npoints::Int
end

"""
    eachpoint(message::Message)

Iterate through each point in the message, returning the lon, lat, and value.

# Example
```Julia
for (lon, lat, val) in eachpoint(message)
    println("Value at (\$lat, \$lon) is \$val")
end
```
"""
function eachpoint(handle::Message)
    eref = Ref(Int32(0))
    ptr = ccall((:codes_grib_iterator_new, eccodes),
                Ptr{codes_iterator}, (Ptr{codes_handle}, Culong, Ref{Cint}),
                handle.ptr, 0, eref)
    errorcheck(eref[])
    if ptr == C_NULL
        throw(ErrorException("Could not create GribIterator"))
    end
    GribIterator(ptr, handle["numberOfPoints"])
end

function Base.iterate(iter::GribIterator, state=())
    lon = 0.0
    lat = 0.0
    value = 0.0
    lonref = Ref(lon)
    latref = Ref(lat)
    valref = Ref(value)
    next = ccall((:codes_grib_iterator_next, eccodes), Cint,
                (Ptr{codes_iterator}, Ref{Cdouble}, Ref{Cdouble}, Ref{Cdouble}),
                iter.ptr, latref, lonref, valref)
    if next == 0
        destroy(iter)
        nothing
    else
        rets = (lonref[], latref[], valref[])
        (rets, ())
    end
end

function destroy(iter::GribIterator)
    err = ccall((:codes_grib_iterator_delete, eccodes), Cint, (Ptr{codes_iterator},), iter.ptr)
    errorcheck(err)
end

Base.eltype(::Type{GribIterator}) = Tuple{Float64, Float64, Float64}
Base.length(iter::GribIterator) = iter.npoints


text:
immutable ErrorBarGeometry <: Gadfly.GeometryElement
end

immutable XErrorBarGeometry <: Gadfly.GeometryElement
end

immutable YErrorBarGeometry <: Gadfly.GeometryElement
end

const errorbar = ErrorBarGeometry
const xerrorbar = XErrorBarGeometry
const yerrorbar = YErrorBarGeometry

function element_aesthetics(::ErrorBarGeometry)
    [:x, :y, :xmin, :xmax, :ymin, :ymax]
end

function element_aesthetics(::YErrorBarGeometry)
    [:x, :ymin, :ymax]
end

function element_aesthetics(::XErrorBarGeometry)
    [:y, :xmin, :xmax]
end

# Generate a form for the errorbar geometry.
#
# Args:
#   geom: errorbar geometry
#   theme: the plot's theme
#   aes: aesthetics
#
# Returns:
#   A compose Form.
#
function render(geom::ErrorBarGeometry, theme::Gadfly.Theme, aes::Gadfly.Aesthetics)
    # check for X and Y error bar aesthetics
    if isempty(Gadfly.undefined_aesthetics(aes, element_aesthetics(xerrorbar())...))
        xctx = render(xerrorbar(), theme, aes)
    else
        xctx = nothing
    end
    if isempty(Gadfly.undefined_aesthetics(aes, element_aesthetics(yerrorbar())...))
        yctx = render(yerrorbar(), theme, aes)
    else
        yctx = nothing
    end
    compose(context(order=3), xctx, yctx)
end

function render(geom::YErrorBarGeometry, theme::Gadfly.Theme, aes::Gadfly.Aesthetics)
    Gadfly.assert_aesthetics_defined("Geom.errorbar", aes,
                                     element_aesthetics(geom)...)
    Gadfly.assert_aesthetics_equal_length("Geom.errorbar", aes,
                                          element_aesthetics(geom)...)

    default_aes = Gadfly.Aesthetics()
    default_aes.color = PooledDataArray(RGB{Float32}[theme.default_color])
    aes = inherit(aes, default_aes)
    caplen = theme.errorbar_cap_length/2

    return compose!(
        context(order=3),

        # top cap
        Compose.line([[(x*cx - caplen, ymax), (x*cx + caplen, ymax)]
                      for (x, ymax) in zip(aes.x, aes.ymax)]),

        # error bar
        Compose.line([[(x*cx, ymax), (x*cx, ymin)]
                      for (x, ymin, ymax) in zip(aes.x, aes.ymin, aes.ymax)]),

        # bottom cap
        Compose.line([[(x*cx - caplen, ymin), (x*cx + caplen, ymin)]
                      for (x, ymin) in zip(aes.x, aes.ymin)]),

        stroke([theme.stroke_color(c) for c in aes.color]),
        linewidth(theme.line_width),
        aes.color_key_continuous == true ?
            svgclass("geometry") :
            svgclass([string("geometry ", svg_color_class_from_label(aes.color_label([c])[1]))
                      for c in aes.color]))
end

function render(geom::XErrorBarGeometry, theme::Gadfly.Theme, aes::Gadfly.Aesthetics)
    Gadfly.assert_aesthetics_defined("Geom.errorbar", aes,
                                     element_aesthetics(geom)...)
    Gadfly.assert_aesthetics_equal_length("Geom.errorbar", aes,
                                          element_aesthetics(geom)...)

    colored = aes.color != nothing
    default_aes = Gadfly.Aesthetics()
    default_aes.color = PooledDataArray(RGB{Float32}[theme.default_color])
    aes = inherit(aes, default_aes)
    caplen = theme.errorbar_cap_length/2

    return compose!(
        context(order=3),

        # top cap
        Compose.line([[(xmin, y*cy - caplen), (xmin, y*cy + caplen)]
                      for (xmin, y) in zip(aes.xmin, aes.y)]),

        # error bar
        Compose.line([[(xmin, y*cy), (xmax, y*cy)]
                      for (xmin, xmax, y) in zip(aes.xmin, aes.xmax, aes.y)]),

        # right cap
        Compose.line([[(xmax, y*cy - caplen), (xmax, y*cy + caplen)]
                      for (xmax, y) in zip(aes.xmax, aes.y)]),

        stroke([theme.stroke_color(c) for c in aes.color]),
        linewidth(theme.line_width),
        (aes.color_key_continuous == true || !colored) ?
            svgclass("geometry") :
            svgclass([string("geometry ", svg_color_class_from_label(aes.color_label([c])[1]))
                      for c in aes.color]))
end