text:
"""
    MattssonAlmquistCarpenter2014Extended()

Coefficients of the extended SBP operators given in
- Mattsson, Almquist, Carpenter (2014)
  Optimal diagonal-norm SBP operators.
  Journal of Computational Physics 264, pp. 91-111.
"""
struct MattssonAlmquistCarpenter2014Extended <: SourceOfCoefficients end

function Base.show(io::IO, source::MattssonAlmquistCarpenter2014Extended)
    if get(io, :compact, false)
        summary(io, source)
    else
        print(io,
            "Mattsson, Almquist, Carpenter (2014) \n",
            "  Optimal diagonal-norm SBP operators. \n",
            "  Journal of Computational Physics 264, pp. 91-111.")
    end
end

function first_derivative_coefficients(source::MattssonAlmquistCarpenter2014Extended,
                                       order::Int, T=Float64, mode=FastMode())
    if order == 2
        left_boundary = (
            # d1
            DerivativeCoefficientRow{T,1,3}(SVector(T(-6//5),
                                                    T(7//5),
                                                    T(-1//5) )),
            # d2
            DerivativeCoefficientRow{T,1,3}(SVector(T(-1//2),
                                                    T(0),
                                                    T(1//2) )),
            # d3
            DerivativeCoefficientRow{T,1,4}(SVector(T(1//11),
                                                    T(-7//11),
                                                    T(0),
                                                    T(6//11) )),
        )
        right_boundary = .- left_boundary
        upper_coef = SVector(T(1//2))
        central_coef = zero(T)
        lower_coef = -upper_coef
        left_weights = SVector( T(5//12),
                                T(7//6),
                                T(11//12) )
        right_weights = left_weights
        left_boundary_derivatives = Tuple{}()
        right_boundary_derivatives = left_boundary_derivatives

        DerivativeCoefficients(left_boundary, right_boundary,
                                left_boundary_derivatives, right_boundary_derivatives,
                                lower_coef, central_coef, upper_coef,
                                left_weights, right_weights, mode, 1, order, source)
    elseif order == 4
        left_boundary = (
            # d1
            DerivativeCoefficientRow{T,1,6}(SVector(T(-800//511),
                                                    T(1744244//869211),
                                                    T(-312391//1738422),
                                                    T(-100109//289737),
                                                    T(54107//1738422),
                                                    T(6575//124173) )),
            # d2
            DerivativeCoefficientRow{T,1,6}(SVector(T(-1744244//3752595),
                                                    T(0),
                                                    T(238831//750519),
                                                    T(290357//1501038),
                                                    T(-625//35739),
                                                    T(-220357//7505190) )),
            # d3
            DerivativeCoefficientRow{T,1,6}(SVector(T(312391//3511620),
                                                    T(-238831//351162),
                                                    T(0),
                                                    T(17335//25083),
                                                    T(-74827//702324),
                                                    T(3859//585270) )),
            # d4
            DerivativeCoefficientRow{T,1,6}(SVector(T(100109//1109430),
                                                    T(-290357//1331316),
                                                    T(-17335//47547),
                                                    T(0),
                                                    T(342799//665658),
                                                    T(-149959//6656580) )),
            # d5
            DerivativeCoefficientRow{T,1,7}(SVector(T(-54107//5025510),
                                                    T(625//23931),
                                                    T(74827//1005102),
                                                    T(-342799//502551),
                                                    T(0),
                                                    T(1715156//2512755),
                                                    T(-240//2659) )),
            # d6
            DerivativeCoefficientRow{T,1,8}(SVector(T(-6575//392877),
                                                    T(220357//5500278),
                                                    T(-3859//916713),
                                                    T(149959//5500278),
                                                    T(-1715156//2750139),
                                                    T(0),
                                                    T(9600//14551),
                                                    T(-1200//14551) )),
        )
        right_boundary = .- left_boundary
        upper_coef = SVector(T(2//3), T(-1//12))
        central_coef = zero(T)
        lower_coef = -upper_coef
        left_weights = SVector( T(511//1600),
                                T(3971//2880),
                                T(929//1440),
                                T(587//480),
                                T(2659//2880),
                                T(14551//14400) )
        right_weights = left_weights
        left_boundary_derivatives = Tuple{}()
        right_boundary_derivatives = left_boundary_derivatives

        DerivativeCoefficients(left_boundary, right_boundary,
                                left_boundary_derivatives, right_boundary_derivatives,
                                lower_coef, central_coef, upper_coef,
                                left_weights, right_weights, mode, 1, order, source)
    elseif order == 6
        left_boundary = (
            # d1
            DerivativeCoefficientRow{T,1,8}(SVector(T(-12700800//7497391),
                                                    T(84512094//37486955),
                                                    T(-12926907//149947820),
                                                    T(-1378838//2044743),
                                                    T(-2874//40093),
                                                    T(2857680//7497391),
                                                    T(-39279943//449843460),
                                                    T(-825552//37486955) )),
            # d2
            DerivativeCoefficientRow{T,1,8}(SVector(T(-84512094//193589725),
                                                    T(0),
                                                    T(56626691//464615340),
                                                    T(97579603//232307670),
                                                    T(3642717//77435890),
                                                    T(-246914711//1161538350),
                                                    T(23380423//464615340),
                                                    T(54432//5531135) )),
            # d3
            DerivativeCoefficientRow{T,1,8}(SVector(T(205189//2106340),
                                                    T(-56626691//79619652),
                                                    T(0),
                                                    T(969659//1153908),
                                                    T(-30244385//79619652),
                                                    T(9737351//44233140),
                                                    T(-508309//6634971),
                                                    T(780937//79619652) )),
            # d4
            DerivativeCoefficientRow{T,1,8}(SVector(T(15167218//136593975),
                                                    T(-97579603//273187950),
                                                    T(-22302157//182125300),
                                                    T(0),
                                                    T(8120941//54637590),
                                                    T(13871884//45531325),
                                                    T(-42123859//546375900),
                                                    T(-2075677//273187950) )),
            # d5
            DerivativeCoefficientRow{T,1,8}(SVector(T(537438//10627085),
                                                    T(-3642717//21254170),
                                                    T(6048877//25505004),
                                                    T(-8120941//12752502),
                                                    T(0),
                                                    T(29205107//63762510),
                                                    T(15263287//127525020),
                                                    T(-243381//4250834) )),
            # d6
            DerivativeCoefficientRow{T,1,9}(SVector(T(-45360//513973),
                                                    T(246914711//971408970),
                                                    T(-9737351//215868660),
                                                    T(-13871884//32380299),
                                                    T(-29205107//194281794),
                                                    T(0),
                                                    T(321139631//647605980),
                                                    T(-49594423//971408970),
                                                    T(6720//513973) )),
            # d7
            DerivativeCoefficientRow{T,1,10}(SVector(T(39279943//1409459100),
                                                    T(-23380423//281891820),
                                                    T(508309//23490985),
                                                    T(42123859//281891820),
                                                    T(-15263287//281891820),
                                                    T(-321139631//469819700),
                                                    T(0),
                                                    T(30841499//40270260),
                                                    T(-108864//671171),
                                                    T(12096//671171) )),
            # d8
            DerivativeCoefficientRow{T,1,11}(SVector(T(825552//128159995),
                                                    T(-381024//25631999),
                                                    T(-780937//307583988),
                                                    T(2075677//153791994),
                                                    T(1216905//51263998),
                                                    T(49594423//768959970),
                                                    T(-215890493//307583988),
                                                    T(0),
                                                    T(19051200//25631999),
                                                    T(-3810240//25631999),
                                                    T(423360//25631999) )),
        )
        right_boundary = .- left_boundary
        upper_coef = SVector(T(3//4), T(-3//20), T(1//60))
        central_coef = zero(T)
        lower_coef = -upper_coef
        left_weights = SVector( T(7497391//25401600),
                                T(1106227//725760),
                                T(105317//403200),
                                T(260179//145152),
                                T(303631//725760),
                                T(513973//403200),
                                T(671171//725760),
                                T(25631999//25401600) )
        right_weights = left_weights
        left_boundary_derivatives = Tuple{}()
        right_boundary_derivatives = left_boundary_derivatives

        DerivativeCoefficients(left_boundary, right_boundary,
                                left_boundary_derivatives, right_boundary_derivatives,
                                lower_coef, central_coef, upper_coef,
                                left_weights, right_weights, mode, 1, order, source)
    else
        throw(ArgumentError("Order $order not implemented/derived."))
    end
end


text:
"""
    default_weights(g)

Return the binary undirected adjacency matrix of `g`.
"""
function default_weights(g::G) where {G<:AbstractGraph}
    m = spzeros(nv(g), nv(g))
    for e in edges(g)
        m[src(e), dst(e)] = 1
    end
    return m
end

ext: jl


text:
import Base.in

"""
    in(z,C)
Check if the point `z` is on the circle `C`.
May be invoked as `z ∈ C`.
"""
function in(z::Number, C::Circle)::Bool
    d = abs(z - center(C))
    return abs(d - radius(C)) <= get_tolerance()
end

import Base.issubset

function issubset(C::Circle, D::Circle)
    z = center(C)
    w = center(D)
    d = abs(z - w)  # distance between centers
    r = radius(C)
    s = radius(D)

    return d + r <= s + get_tolerance()
end

function (==)(C::Circle, D::Circle)
    dx = C.x - D.x
    dy = C.y - D.y
    dr = C.r - D.r

    return dx^2 + dy^2 + dr^2 <= get_tolerance()
end

(==)(::Circle, ::Line) = false
(==)(::Line, ::Circle) = false


text:
# Class containing the objects for the windowed version of the integral operator
# for the slow varying quantity along the principal ray direction

type FastMslow
    # type to encapsulate the fast application of M = I + omega^2G*spadiagm(nu)
    GFFT :: Array{Complex128,2}
    nu :: Array{Float64,1}
    x  :: Array{Float64,1}
    y  :: Array{Float64,1}
    # number of points in the extended domain
    ne :: Int64
    me :: Int64
    # number of points in the original domain
    n  :: Int64
    m  :: Int64
    # frequency
    omega :: Float64
    # direction of the plane wave
    e ::Array{Float64,1}
end

import Base.*
import Base.A_mul_B!
import Base.eltype
import Base.size

function size(M::FastMslow, dim::Int64)
  @assert dim<3 && dim>0
    return M.n*M.m
end

function *(M::FastMslow, b::Array{Complex128,1})
    # function to overload the applyication of
    # M using a Toeplitz reduction via a FFT

    # Allocate the space for the extended B
    BExt = zeros(Complex128,M.ne, M.me);
    # Apply spadiagm(nu) and ented by zeros
    BExt[1:M.n,1:M.m]= reshape((exp(1im*M.omega*(M.e[1]*M.x + M.e[2]*M.y)).*M.nu).*b,M.n,M.m) ;

    # Fourier Transform
    BFft = fft(BExt)
    # Component-wise multiplication
    BFft = M.GFFT.*BFft
    # Inverse Fourier Transform
    BExt = ifft(BFft)

    # multiplication by omega^2
    B = M.omega^2*(BExt[M.n:2*M.n-1, M.m:2*M.m-1]);

    return (b + (exp(-1im*M.omega*(M.e[1]*M.x + M.e[2]*M.y)).*(B[:])))
end

function A_mul_B!(Y,
                  M::FastMslow,
                  V)
    # in place matrix matrix multiplication
    @assert(size(Y) == size(V))
    # print(size(V))
    for ii = 1:size(V,2)
        Y[:,ii] = M*V[:,ii]
    end
end

# # TO DO: add more parameters for the window in here
# function buildFastConvolutionSlow(x::Array{Float64,1},y::Array{Float64,1},
#                                   h::Float64,k,nu::Function, e::Array{Float64,1};
#                                   quadRule::String = "trapezoidal",
#                                   window::String = "normal")

#     if quadRule == "trapezoidal"
#         if window == "normal"

#             (ppw,D) = referenceValsTrapRule();
#             D0      = D[round(Int,k*h)];
#             (n,m) = length(x), length(y)
#             Ge    = buildGConv(x,y,h,n,m,D0,k);
#             GFFT  = fft(Ge);
#             X = repmat(x, 1, m)[:]
#             Y = repmat(y', n,1)[:]
#         end
#     return FastM(GFFT,nu(X,Y),2*n-1,2*m-1,n, m, k);

#   elseif quadRule == "Greengard_Vico"

#       Lp = 4*(abs(x[end] - x[1]) + h)
#       L  =   (abs(x[end] - x[1]) + h)*1.5
#       (n,m) = length(x), length(y)
#       X = repmat(x, 1, m)[:]
#       Y = repmat(y', n,1)[:]

#       # this is depending if n is odd or not
#       if mod(n,2) == 0
#         kx = (-(2*n):1:(2*n-1));
#         ky = (-(2*m):1:(2*m-1));

#         KX = (2*pi/Lp)*repmat(kx, 1, 4*m);
#         KY = (2*pi/Lp)*repmat(ky', 4*n,1);

#         S = sqrt(KX.^2 + KY.^2);

#         GFFT = Gtruncated2D(L, k, S)
#         return FastM(GFFT, nu(X,Y), 4*n, 4*m,
#                      n, m, k , quadRule="Greengard_Vico");
#       else
#         # kx = (-2*(n-1):1:2*(n-1) )/4;
#         # ky = (-2*(m-1):1:2*(m-1) )/4;

#         # KX = (2*pi/Lp)*repmat(kx, 1, 4*m-3);
#         # KY = (2*pi/Lp)*repmat(ky', 4*n-3,1);

#         # S = sqrt(KX.^2 + KY.^2);

#         # GFFT = Gtruncated2D(L, k, S)

#         # return FastM(GFFT,nu(X,Y),4*n-3,4*m-3,
#         #              n,m, k,quadRule = "Greengard_Vico");

#         kx = (-2*n:1:2*n-1);
#         ky = (-2*m:1:2*m-1);

#         KX = (2*pi/Lp)*repmat( kx, 1,4*m);
#         KY = (2*pi/Lp)*repmat(ky',4*n,  1);

#         S = sqrt(KX.^2 + KY.^2);

#         GFFT = Gtruncated2D(L, k, S)

#         return FastM(GFFT,nu(X,Y),4*n,4*m,
#                      n,m, k,quadRule = "Greengard_Vico");

#     end
#   end
# end

function buildGConvWindowed(x,y,h::Float64,n::Int64,m::Int64,D0,k::Float64)
    # function to build the convolution vector for the
    # fast application of the convolution, in this case we add a directional
    # filter to eliminate the oscillations

    # this is built for odd n and odd m.

    if mod(n,2) == 1
      # build extended domain
      xe = collect((x[1]-(n-1)/2*h):h:(x[end]+(n-1)/2*h));
      ye = collect((y[1]-(m-1)/2*h):h:(y[end]+(m-1)/2*h));

      Xe = repmat(xe, 1, 2*m-1);
      Ye = repmat(ye', 2*n-1,1);
      # to avoid evaluating at the singularity
      indMiddle = m

    else

      println("so far only works for n odd")
      # to be done
      # # build extended domain
      # xe = collect((x[1]-n/2*h):h:(x[end]+n/2*h));
      # ye = collect((y[1]-m/2*h):h:(y[end]+m/2*h));

      # Xe = repmat(xe, 1, 2*m-1);
      # Ye = repmat(ye', 2*n-1,1);
      # # to avoid evaluating at the singularity
      # indMiddle = m

    end

    R = sqrt(Xe.^2 + Ye.^2);

    # we modify R to remove the zero (so we don't )
    R[indMiddle,indMiddle] = 1;
    # sampling the Green's function
    Ge = sampleGkernelpar(k,R,h)
    #Ge = pmap( x->1im/4*hankelh1(0,k*x)*h^2, R)
    # modiyfin the diagonal with the quadrature
    # modification
    Ge[indMiddle,indMiddle] = 1im/4*D0*h^2;

    return Ge

end

max_stars_repo_name: Forgotten/Lippmann-Schingwer

max_issues_repo_name: Forgotten/Lippmann-Schingwer

max_forks_repo_name: Forgotten/Lippmann-Schingwer


text:
using Pseudea
using Base.Test

# write your own tests here
@test 1 == 1

println("Ten random atomic samples from CZTS...")
for i in 1:10
    println(randCZTS()) #
end

println("A 5x5x5 CZTS random lattice...")
show(randCZTSlattice(5))

println("A 4x4x4 CZTS modulo lattice...")
show(moduloCZTSlattice(4))


text:
using Markdown # hide

# ## Solving single queries and `KnnResult`
# By Eric S. Téllez
#

ENV["GKSwstype"] = "nul" # hide

using SimilaritySearch
#
# This example shows how to perform single queries instead of solving a batch of them.
# This is particularly useful for some applications, and we also show how they are solved,
# which could be used to avoid some memory allocations.
const dim = 8

db = MatrixDatabase(randn(Float32, dim, 10^4))
dist = SqL2Distance()
G = SearchGraph(; dist, db)
index!(G)
nothing # hide

# Suppose you want to compute some $k$ nearest neighbors, for this we use the structure
# `KnnResult` which is a priority queue of maximum size `k`.
res = KnnResult(3) # allocates memory for 10 nearest neighbors
for v in rand(db, 10)
    reuse!(res)  # reuses the res object
    @time search(G, v, res)
    @show minimum(res), maximum(res), argmin(res), argmax(res)
    @show res.id, res.dist
end

# ## `KnnResult`
# This structure is the container for the result and it is also used to specify the number
# of elements to retrieve. As mentioned before, it is a priority queue

reuse!(res)
push!(res, 1, 10)
push!(res, 2, 9)
push!(res, 3, 8)
push!(res, 4, 7)
push!(res, 6, 5)
@show res

# it also supports removals
@show :popfirst! => popfirst!(res)
push!(res, 7, 0.1)
@show :push! => res
@show :pop! => pop!(res)
res
# It can be iterated

for (id, d) in res
    @show id => d
end

#- Using BeutifulMakie style
# ## Dependencies
using Pkg # hide
Pkg.status(["SimilaritySearch"]) # hide

max_stars_repo_name: sadit/SimilaritySearchDemos

max_issues_repo_name: sadit/SimilaritySearchDemos

max_forks_repo_name: sadit/SimilaritySearchDemos


text:
# ------------------------------------------------------------------
# Licensed under the ISC License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    DomainView(domain, locations)

Return a view of `domain` at `locations`.

### Notes

This type implements the `AbstractDomain` interface.
"""
struct DomainView{T,N,
                  D<:AbstractDomain{T,N},
                  I<:AbstractVector{Int}} <: AbstractDomain{T,N}
  domain::D
  locations::I
end

npoints(dv::DomainView) = length(dv.locations)

coordinates!(buff::AbstractVector, dv::DomainView, location::Int) =
  coordinates!(buff, dv.domain, dv.locations[location])

# ------------
# IO methods
# ------------
function Base.show(io::IO, dv::DomainView{T,N,D,I}) where {T,N,
                                                           D<:AbstractDomain{T,N},
                                                           I<:AbstractVector{Int}}
  npts = npoints(dv)
  print(io, "$npts DomainView{$T,$N}")
end

function Base.show(io::IO, ::MIME"text/plain", dv::DomainView)
  println(io, dv)
  Base.print_array(io, coordinates(dv))
end


text:
@testset "Construction" begin

    # basic constructor
    xs = [1,2,3]
    f(x) = x^2
    q = fit(Newton, xs, f)
    @test all(f.(xs) .≈ q.(xs))

    # test tableau
    @test all(q.tableau .≈   [1 3 1; 0 4 5; 0 0 9])

    # test one, zero, variable
    x = variable()
    @test one(q)(x) == one(x)
    @test iszero(zero(q))
    @test variable(q)(x) == x

end

@testset "Conversion" begin

    xs = [1,2,3]
    f(x) = x^2
    p = fit(Newton, xs, f)
    q = convert(Polynomial, p)
    @test coeffs(q)  == [0,0,1]

end

@testset "Evaluation" begin

    xs = sort(rand(5))
    ys = sin.(xs)
    p = fit(Newton, xs, sin)
    @test p.(xs) ≈ sin.(xs)
    us = rand(5)
    @test maximum(abs.([p(x) - sin(x) for x in us])) <= 1e-1

end

@testset "Arithmetic" begin

    p = fit(Newton, [1,2,3], x -> x^2)
    q = fit(Newton, [2,3,4,5], x -> x^3)

    us = rand(5)

    @test all(p.(us) + q.(us) .≈ (p + q).(us))
    @test p.(us) .* q.(us) ≈ (p * q).(us)

end

@testset "Elementwise operations" begin

    p = fit(Newton, [1,2,3,4], x -> x^2)
    p + 1
    -p
    2p

end

@testset "roots" begin

    q = fit(Newton, [2,3,4,5], x -> (x-1)*(x-1/2)*(x+1/2))
    @test all(roots(q) .≈ [-1/2, 1/2, 1])

end


text:
# Generate random channels -

# Copyright Robin Harper 2016-2020

using Distributions
using Random
using LinearAlgebra
using SparseArrays

# Note this is for reproducibility --- for multiple runs or spreadsheets you will need some
# other way to set.
#srand()
Random.seed!

_d=Normal()

#    Based on 1707.06926 (Rudnicki et al) and some others - see workbook for more complete description

"""
    returns the fidelity of the channel (in PauliLiouville basis)
"""
function fidelity(channel)
    diag = tr(channel)-1
    d = sqrt(size(channel)[1])
    p = diag/(d^2-1)
    return ((d-1)*p+1)/d
    #(1+1/3*(channel[2,2]+channel[3,3]+channel[4,4]))/2
end

"""
    returns the unitarity of the channel (PauliLiouville basis)
"""
function unitarity(noise)
    d = sqrt(size(noise)[1])
    return (1/(d^2-1))*tr(noise[2:end,2:end]'*noise[2:end,2:end])
end

"""
    A given fidelity has a minimum unitarity that it can have
    (so perfect fidelity = 1 = unitarity)
    This shows as a percentage how much the unitarity of the channel
    lies between the minimum and the maximum (always 1)
"""
function unitarityPercent(ch)
    fid = fidelity(ch)
    uni = unitarity(ch)
    d = sqrt(size(ch)[1])
    r = 1-fid
    minU = (1-(d*r)/(d-1))^2
    return (uni-minU)/(1-minU)
end

"""
    Simple function that takes a normal distribution and returns a random variable
    that is bounded (non-inclusively) by min and max.
"""
function genTruncated(min,max,average,sigma,normalDistribution)
    x=0.0
    while (x < min) || (x > max)
        x=rand(normalDistribution,1)[1]*sigma+average
    end
    return x
end

"""
    Given the eigenvalues λ return a
    vector of random τ, constrained by the λ
"""
function getτ(λ1,λ2,λ3)
    Maxτ1 = 1-abs(λ1)
    Maxτ2 = 1-abs(λ2)
    Maxτ3 = 1-abs(λ3)
    τ1 = (Maxτ1)*(2*rand()-1)
    τ2 = (Maxτ2)*(2*rand()-1)
    τ3 = (Maxτ3)*(2*rand()-1)
    return [τ1 τ2 τ3]
end

"""
    q(l) and Zη(t,l) See equations 20 and 21 of 1707.06926 (Rudnicki et al)
"""
function q(l)
 return (1+l[1]+l[2]+l[3])*(1 + l[1]-l[2]-l[3])*(1-l[1]+l[2]-l[3])*(1-l[1]-l[2]+l[3])
end

"""
    q(l) and Zη(t,l) See equations 20 and 21 of 1707.06926 (Rudnicki et al)
"""
function Zη(t,l)
    norm(t)^4-2*norm(t)^2-2*sum([l[i]^2*(2*(t[i]^2)-norm(t)^2) for i=1:3]) + q(l)
end

"""
    Main function to generate the 'internal' matrix of the channel i.e. the one
    in the form where the unital part consists of diagonal numbers, and the we have
    a non-unital part. This can then be acted on by unitaries either side (rotations)
    to generate an arbitarty CPTP channel
"""
function genSigma(min,max,average,sigma)
    λ1=genTruncated(min,max,average,sigma,_d)
    λ2=genTruncated(min,max,average,sigma,_d)
    λ3=genTruncated(min,max,average,sigma,_d)
    while 1+λ3 < abs(λ1+λ2) || 1-λ3 < abs(λ2-λ1)
        λ3=genTruncated(min,max,average,sigma,_d)
    end
    τ = getτ(λ1,λ2,λ3)
    while norm(τ)^2 > 1-sum([x^2 for x in [λ1 λ2 λ3]])+2*λ1*λ2*λ3 || Zη(τ,[λ1 λ2 λ3]) < 0
        τ = getτ(λ1,λ2,λ3)
    end
    return [1 0 0 0;τ[1] λ1 0 0;τ[2] 0 λ2 0;τ[3] 0 0 λ3]
end

"""
    genChannel(r1,sm,sv,r2)
    Rotate the sigma channel randomly to generate an arbitrary channel
    r1 represents a rotation. The elements of a vector of three normally distributed elements are multiplied by this and used as an XZX rotation. For high fidelity channels this will be small.
    sm, represents the diagonal channel. This is the mean. High fidelity channels will have this close to 1.
    sv, represents the variance in the diagonal channel.
    r2, like r1.

    For example randomFidelityNoise is defined as   genChannel(0.06,0.998,0.04,0.06)
    whereas the noisier randomMeasureNoise is       genChannel(0.05,0.98,0.15,0.05)
"""
function genChannel(r1,sm,sv,r2)
    createRotation(rand(_d,3)*r1) * genSigma(0.9,1,sm,sv) * createRotation(rand(_d,3)*r2)'
end

"""
    given a vector of three angles, generate a random rotation channel
    by using the vector to control a Z*X*Z rotation
"""
function createRotation(angles)
    rotation = Matrix{Float64}(I,4,4)
    rotation[2:4,2:4] =
        [cos(angles[1]) sin(angles[1]) 0;-sin(angles[1]) cos(angles[1]) 0;0 0 1] *
        [1 0 0;0 cos(angles[2]) sin(angles[2]);0 -sin(angles[2]) cos(angles[2])] *
        [cos(angles[3]) sin(angles[3]) 0;-sin(angles[3]) cos(angles[3]) 0;0 0 1]
    return rotation
end

"""
    Helper function to generate typical, high fidelity channel.
"""
function randomFidelityNoise()
    return genChannel(0.06,0.998,0.04,0.06)
end

"""
    Helper function to generate typical, lower fidelity channel.
    I use it to simulate a state preparation channel.
"""
function randomPrepNoise()
    return genChannel(0.05,0.985,0.15,0.05)
end

"""
    Helper function to generate typical, lower fidelity channel.
    I typically use it to simulate measurement error.
"""
function randomMeasureNoise()
    return genChannel(0.05,0.98,0.15,0.05)
end

"""
    function genChannelMap(f)
    Simple helper function - designed to help generate a map of a specific fidelity (supplied parameter).
    Just generates random channels until you get one of the correct fidelity \$\\pm 0.0001\$.
    In general though you would want to rewrite this to give genChannel
    parameters that will result in random channels that often bracket the desired fidelity
    May take a long time if the parameters are far from the types of channels generated
    by randomFidelityNoise.
"""
function genChannelMap(f)
    while true
        candidate = randomFidelityNoise()
        if abs(f-fidelity(candidate)) < 0.0001
            return candidate
        end
    end
end

pI=[1 0;0 1]
pZ=[1 0;0 -1]

# This can be a bit slow, because of the makeSuper on large qubits.
# I have done this better somewhere.
"""
    function add2QPauliNoise(originalMap,cbit,tbit;pauliToAdd = pZ)

    Adds some noise between two qubits in the noise map.
    You pass in an  noise map (super operator basis)
    the control bit and the target bit and optionally the axis of rotation for the noise.
    will return a noise map with a small amount of random noise added between these qubits.
    The noise itself is calculated as \$e^(i*\\frac{\\text{rand()}}{20}\\text{pauliToAdd})\$
"""
function add2QPauliNoise(originalMap,cbit,tbit;pauliToAdd = pZ)
    nqubitsInMap = round(Int,log2(sqrt(size(originalMap,1))))
    #print("Map is $nqubitsInMap qubits\n")
    ok = false
    ## So we are getting irritating (what I think are) rounding errors here, best way to avoid.
    zz = []
    zeroSelect = SparseArrays.sparse([1; 0]*[1 0])
    oneSelect  = SparseArrays.sparse([0;1]*[0 1])
    spI = SparseArrays.sparse(pI)
    thetaa = rand()/10
    thetab = rand()/10
    if cbit < tbit
            startC=cbit
            endC=tbit
             l1 = SparseArrays.sparse(exp(im*thetaa/2*pI))
             l2 = SparseArrays.sparse(exp(im*thetab/2*pauliToAdd))
             for i=cbit+1:tbit-1
                   l1=kron(spI,l1)
                   l2=kron(spI,l2)
             end
            cbitbit = kron(zeroSelect,l1)+kron(oneSelect,l2)
    else
            startC=tbit
            endC=cbit
            l1 = exp(im*thetaa/2*pI)
            l2 = exp(im*thetab/2*pauliToAdd)
            for i=tbit+1:cbit-1
                 l1=kron(l1,spI)
                 l2=kron(l2,spI)
            end
            cbitbit = kron(l1,zeroSelect)+kron(l2,oneSelect)
    end
    for i = 1:(min(cbit,tbit)-1)
            cbitbit = pI ⊗ cbitbit
    end
    for i = max(cbit,tbit)+1:nqubitsInMap
            cbitbit = cbitbit ⊗ pI
    end
    zz = makeSuper(Array(cbitbit))
    withZ = zz*originalMap
    return withZ
end


text: promote_graph(graph::AbstractMatrix, nf::CuArray) = cu(graph)


text:
module hill_decay_gentrap
using FinEtools
using FinEtoolsHeatDiff
using FinEtoolsHeatDiff.AlgoHeatDiffModule
import LinearAlgebra: cholesky, mul!
using UnicodePlots
# using BenchmarkTools

function hill_decay_t3()
	thermal_conductivity =  [i==j ? 0.2 : zero(FFlt) for i=1:2, j=1:2]; # conductivity matrix
	Width = 60.0
	Height = 40.0
	N = 50
	specific_heat = 1.0
	theta = 1.0; # generalized trapezoidal method parameter
	dt = 2.0; # time step
	tend = 20*dt; # length of the time interval
	T0(x, y) = 500.0/(x^2+y^2+5);# Initial distribution of temperature

	tolerance =Width/N/100;

    fens,fes =T3block(Width, Height, N, N)
    for i in 1:count(fens)
    	fens.xyz[i, 1] -= Width / 2
    	fens.xyz[i, 2] -= Height / 2
    end
    l1  = selectnode(fens; box=[Width/2 Width/2 Height/2 Height/2],  inflate = tolerance)

    geom = NodalField(fens.xyz)
    Temp = NodalField(zeros(size(fens.xyz,1),1))

    for i in 1:count(fens)
    	Temp.values[i] = T0(fens.xyz[i, :]...)
    end

    applyebc!(Temp)
    numberdofs!(Temp)

    cornerdof = Temp.dofnums[l1]

    material = MatHeatDiff(thermal_conductivity, specific_heat)
    femm = FEMMHeatDiff(IntegDomain(fes, TriRule(3)), material)

    K = conductivity(femm, geom, Temp)
    C = capacity(femm, geom, Temp)

    A = cholesky((1/dt*C+theta*K))
    B = (1/dt*C-(1-theta)*K)

    Tn = gathersysvec(Temp)
    Tn1 = deepcopy(Tn)

# Tn = rand(length(Tn))
# dT = rand(length(Tn))

#     @btime (mul!($dT, $K, $Tn); $dT .*= -1)
#     @btime ($dT .= $K * $Tn)

    tstart = time()
    t = 0.0
    ts = Float64[t]
    @show Tn1[cornerdof, 1][1]
    Corner_T = Float64[Tn1[cornerdof, 1][1]];
    while true # Time stepping
    	t = t+dt; # Advance the current time
    	Tn .= Tn1 # current temperature system vector

    	# Solve for the new temperature  vector
    	rhs = B*Tn; # complete right-hand side (loads)
    	Tn1 .= A\rhs; # compute next temperature
    	push!(ts, t)
    	push!(Corner_T, Tn1[cornerdof, 1][1])
    	if (t == tend) # Have we reached the end?  If so jump out.
	    	break;
	    end
	    if (t+dt > tend) # Adjust the last time step so that we exactly reach tend
	    	dt = tend-t;
	    end
	end
	println("Time = $(time() - tstart)")

	# @show Corner_T
	@show minimum(Corner_T), maximum(Corner_T)
	plt = lineplot(vec(ts), vec(Corner_T), canvas = DotCanvas, title = "Transient temperature at the corner", name = "T", xlabel = "Time", ylabel = "T")
	display(plt)
    # File =  "a.vtk"
    # MeshExportModule.vtkexportmesh (File, fes.conn, [geom.values Temp.values], MeshExportModule.T3; scalars=Temp.values, scalars_name ="Temperature")

    true
end

end # module hill_decay_gentrap

using .hill_decay_gentrap
hill_decay_gentrap.hill_decay_t3()

max_stars_repo_name: PetrKryslUCSD/FinEtoolsHeatDiff

max_issues_repo_name: PetrKryslUCSD/FinEtoolsHeatDiff

max_forks_repo_name: PetrKryslUCSD/FinEtoolsHeatDiff


text:
# indicator of the free cone

export IndFree

"""
**Indicator of the free cone**

    IndFree()

Returns the indicator function of the whole space, or "free cone", *i.e.*,
a function which is identically zero.
"""
struct IndFree <: ProximableFunction end

is_separable(f::IndFree) = true
is_convex(f::IndFree) = true
is_affine(f::IndFree) = true
is_cone(f::IndFree) = true
is_smooth(f::IndFree) = true
is_quadratic(f::IndFree) = true

const Zero = IndFree

function (f::IndFree)(x::AbstractArray{T}) where {R, T <: RealOrComplex{R}}
  return zero(R)
end

function prox!(y::AbstractArray{T}, f::IndFree, x::AbstractArray{T}, gamma::R=one(R)) where {R, T <: RealOrComplex{R}}
  y .= x
  return zero(R)
end

prox!(y::AbstractArray{T}, f::IndFree, x::AbstractArray{T}, gamma::AbstractArray{T}) where {R, T <: RealOrComplex{R}} = prox!(y, f, x)

function gradient!(y::AbstractArray{T}, f::IndFree, x::AbstractArray{T}) where {R, T <: RealOrComplex{R}}
  y .= zero(T)
  return zero(R)
end

fun_name(f::IndFree) = "indicator of the free cone"
fun_dom(f::IndFree) = "AbstractArray{Real}, AbstractArray{Complex}"
fun_expr(f::IndFree) = "x ↦ 0"
fun_params(f::IndFree) = "none"

function prox_naive(f::IndFree, x::AbstractArray, gamma=1.0)
  return x, 0.0
end


text:
module MPS

using LinearAlgebra
using TensorOperations
using LinearMaps
using KrylovKit

export MPSTensor, bond_dim, bond_dim_L, bond_dim_R, phys_dim, mps_tensor, mps_tensor_shape,
       rand_MPSTensor_unitary, rand_MPSTensor,
       MPS_TM, TM_dense, eyeTM, TM_dense_op_nn,
       applyTM_op_nn_l, applyTM_l, applyTM_r,
       workvec_applyTM_l, workvec_applyTM_l!, res_applyTM_l, res_applyTM_l!, applyTM_l!,
       workvec_applyTM_r, workvec_applyTM_r!, res_applyTM_r, res_applyTM_r!, applyTM_r!,
       tm_eigs_sparse, tm_eigs_dense, tm_eigs, tm_dominant_eigs,
        gauge_transform, canonicalize_left,
       MPOTensor, IdentityMPOTensor, mpo_tensor_shape, rand_MPOTensor,
       mul_MPO, MPO_open, MPS_MPO_TM, TM_dense_MPO, applyTM_MPO_l, applyTM_MPO_r,
       worklen_applyTM_MPO_l, workvec_applyTM_MPO_l, workvec_applyTM_MPO_l!,
       res_applyTM_MPO_l, res_applyTM_MPO_l!, applyTM_MPO_l!,
       workvec_applyTM_MPO_r, workvec_applyTM_MPO_r!, applyTM_MPO_r!, res_applyTM_MPO_r, res_applyTM_MPO_r!,
       TM_convert

"""
We will use the convention that the middle index is the physical one,
while the first and third are the left and right virtual indices.
In pictures, for an `MPSTensor` `A`:

 1--A--3
    |
    2

"""
MPSTensor{T} = Array{T,3}

bond_dim_L(A::MPSTensor) = size(A,1)
bond_dim_R(A::MPSTensor) = size(A,3)
bond_dim(A::MPSTensor) = bond_dim_L(A)
phys_dim(A::MPSTensor) = size(A,2)
mps_tensor_shape(d::Integer, D::Integer) = (D,d,D)
mps_tensor_shape(d::Integer, D1::Integer, D2::Integer) = (D1,d,D2)

"""
MPO tensor:

     2
     |
 1---o---3
     |
     4
"""
MPOTensor{T} = Array{T,4}

struct IdentityMPOTensor
end

mpo_tensor_shape(d::Integer, M::Integer) = (M,d,M,d)
mpo_tensor_shape(d::Integer, M1::Integer, M2::Integer) = (M1,d,M2,d)

function randunitary(::Type{T}, N::Integer)::Matrix{T} where {T}
    if T <: Complex
        rT = real(T)
        A = complex.(randn(rT, N,N), randn(rT, N,N)) / √2
    else
        A = randn(T, N,N)
    end
    Q, R = qr(A)
    r = diag(R)
    L = diagm(0 => r ./ abs.(r))
    Q*L
end

#Note: An MPS generated this way already has r proportional to I and largest tm eigenvalue = 1
function rand_MPSTensor_unitary(::Type{T}, d::Integer, D1::Integer, D2::Integer)::MPSTensor{T} where {T}
    U = randunitary(T, max(D1,d*D2))
    reshape(U[1:D1,1:d*D2], mps_tensor_shape(d,D1,D2))
end
rand_MPSTensor_unitary(T, d::Integer, D::Integer) = rand_MPSTensor_unitary(T, d, D, D)

function rand_MPSTensor(::Type{T}, d::Integer, D1::Integer, D2::Integer)::MPSTensor{T}  where {T}
    shp = mps_tensor_shape(d,D1,D2)
    if T <: Complex
        rT = real(T)
        A = complex.(randn(rT, shp), randn(rT, shp)) / √2
    else
        A = randn(T, shp)
    end
    A
end
rand_MPSTensor(T, d::Integer, D::Integer) = rand_MPSTensor(T, d, D, D)

function rand_MPOTensor(::Type{T}, d::Integer, M1::Integer, M2::Integer)::MPOTensor{T}  where {T}
    shp = mpo_tensor_shape(d,M1,M2)
    if T <: Complex
        rT = real(T)
        O = complex.(randn(rT, shp), randn(rT, shp)) / √2
    else
        O = randn(T, shp)
    end
    O
end
rand_MPOTensor(T, d::Integer, M::Integer) = rand_MPOTensor(T,d,M,M)

"""
Dense transfer matrix with MPO:

 1----A-----4
      |
 2----o-----5
      |
 3--conj(B)-6
"""
MPS_MPO_TM{T} = Array{T,6}

"""
Dense transfer matrix without MPO:

 1----A-----3
      |
 2--conj(B)-4

"""
MPS_TM{T} = Array{T,4}

#Note: We will generally use MPS_MPO_TM for function arguments, providing functions to convert where possible.

LinearAlgebra.tr(TM::MPS_MPO_TM{T}) where {T} = scalar(@tensor res[] := TM[t,m,b,t,m,b])

#Turn the MPS TM into an MPS TM with (size 1) MPO indices
TM_convert(TM::MPS_TM{T}) where {T} = reshape(TM, (size(TM,1),1,size(TM,2), size(TM,3),1,size(TM,4)))::MPS_MPO_TM{T}

function TM_convert(TM::MPS_MPO_TM{T})::MPS_TM{T} where {T}
    size(TM,2) == size(TM,5) == 1 || error("MPO bond dimensions not equal to 1!")
    dropdims(TM, dims=(2,5))
end

function TM_dense(A::MPSTensor, B::MPSTensor)::MPS_MPO_TM
    GC.enable(false)
    @tensor TM[lA,lB,rA,rB] := A[lA, s, rA] * conj(B[lB, s, rB]) #NOTE: This will generally allocate to do permutations!
    GC.enable(true)
    GC.gc(false)
    TM_convert(TM)
end

eyeTM(T::Type, D::Integer)::MPS_MPO_TM = (e = Matrix{T}(I,D,D); reshape(kron(e,e), (D,1,D, D,1,D)))

function TM_dense_op_nn(A1::MPSTensor, A2::MPSTensor, B1::MPSTensor, B2::MPSTensor, op::Array{T,4})::MPS_MPO_TM where {T}
    @tensor TM[lA1,lB1,rA2,rB2] := ((A1[lA1, p1k, iA] * A2[iA, p2k, rA2]) * op[p1b, p2b, p1k, p2k]) * (conj(B1[lB1, p1b, iB]) * conj(B2[iB, p2b, rB2])) #NOTE: Allocates intermediate arrays
    TM_convert(TM)
end

function applyTM_op_nn_l(A1::MPSTensor, A2::MPSTensor, B1::MPSTensor, B2::MPSTensor, op::Array{T,4}, TM2::MPS_MPO_TM{T}) where {T}
    TM2 = TM_convert(TM2)
    @tensor TM21[k1,b1,k4,b4] := ((((TM2[k1,b1,k2,b2] * A1[k2,s1,k3]) * op[t1,t2,s1,s2]) * conj(B1[b2,t1,b3])) * A2[k3,s2,k4]) * conj(B2[b3,t2,b4])
    TM_convert(TM21)
end

#This is a workaround for StridedArray not including reshaped arrays formed from a FastContiguousSubArray in in Julia <= 0.6.
#Since TensorOperations uses the StridedArray type, it misses these cases. Here, we use pointer manipulation to generate
#views of sections of a vector with an arbitrary "size".
function unsafe_reshaped_subvec(w::Vector{T}, ind1::Integer, sz::NTuple) where {T}
    len = prod(sz)
    @assert length(w) >= ind1+len-1 "Vector is not big enough to hold subarray by $(ind1+len-1 - length(w)) elements!"
    unsafe_wrap(Array, pointer(w, ind1), sz)
end

function res_applyTM_l(A1::MPSTensor{T}, B1::MPSTensor{T}, TM2::MPS_MPO_TM{T}) where {T}
    TM2 = TM_convert(TM2)
    zeros(T, (size(TM2,1), 1, size(TM2,2), bond_dim_R(A1), 1, bond_dim_R(B1)))
end

function res_applyTM_l!(res::MPS_MPO_TM{T}, A1::MPSTensor{T}, B1::MPSTensor{T}, TM2::MPS_MPO_TM{T}) where {T}
    TM2 = TM_convert(TM2)
    sz = (size(TM2,1), 1, size(TM2,2), bond_dim_R(A1), 1, bond_dim_R(B1))
    size(res) == sz ? res : zeros(T, sz)
end

workvec_applyTM_l(A1::MPSTensor{T}, B1::MPSTensor{T}) where {T} = Vector{T}(undef, bond_dim(A1)^2 * bond_dim(B1)^2 * phys_dim(A1) * 2)

function worklen_applyTM_l(A1::MPSTensor{T}, B1::MPSTensor{T}, TM2::MPS_MPO_TM{T}) where {T}
    TM2 = TM_convert(TM2)
    len1 = size(TM2,1) * size(TM2,2) * size(TM2,3) * size(B1,2) * size(B1,3)
    len2 = prod(size(TM2)[1:2]) * size(B1)[3] * size(A1,3)
    max(2len1, len1 + len2)
end

function workvec_applyTM_l(A1::MPSTensor{T}, B1::MPSTensor{T}, TM2::MPS_MPO_TM{T}) where {T}
    Vector{T}(undef, worklen_applyTM_l(A1, B1, TM2))
end

function workvec_applyTM_l!(work::Vector{T}, A1::MPSTensor{T}, B1::MPSTensor{T}, TM2::MPS_MPO_TM{T}) where {T}
    len = worklen_applyTM_l(A1, B1, TM2)
    work = length(work) < len ? resize!(work, len) : work
end

function applyTM_l!(TM21::MPS_MPO_TM{T}, A1::MPSTensor{T}, B1::MPSTensor{T}, TM2::MPS_MPO_TM{T}, work::Vector{T}) where {T}
    TM2 = TM_convert(TM2)
    TM21c = TM_convert(TM21)

    wsz1 = (size(TM2)[1:3]..., size(B1)[2:3]...)
    wlen1 = prod(wsz1)
    TM2B1 = unsafe_reshaped_subvec(work, 1, wsz1)
    @tensor TM2B1[lA,lB, mA,s,rB] = TM2[lA,lB, mA,mB] * conj(B1[mB,s,rB])

    wsz2 = (wsz1[1:2]..., wsz1[5], wsz1[3:4]...)
    wlen2 = prod(wsz2)
    TM2B1p = unsafe_reshaped_subvec(work, wlen1+1, wsz2) #take work vector section after TM2B1 storage
    @tensor TM2B1p[lA,lB, rB,mA,s] = TM2B1[lA,lB, mA,s,rB]
    TM2B1p_copy = unsafe_reshaped_subvec(work, 1, wsz2)
    copyto!(TM2B1p_copy, TM2B1p)
    TM2B1p = TM2B1p_copy

    TM2sz = size(TM2)
    wsz3 = (TM2sz[1:2]..., size(B1)[3], size(A1,3))
    TM21p = unsafe_reshaped_subvec(work, wlen2+1, wsz3)
    @tensor TM21p[lA,lB, rB,rA] = TM2B1p[lA,lB, rB,mA,s] * A1[mA,s,rA]
    @tensor TM21c[lA,lB, rA,rB] = TM21p[lA,lB, rB,rA]

    TM21
end

function applyTM_l(A1::MPSTensor{T}, B1::MPSTensor{T}, TM2::MPS_MPO_TM{T}) where {T}
    TM2 = TM_convert(TM2)
    GC.enable(false)
    @tensor TM21[lA,lB, rA,rB] := (TM2[lA,lB, mA,mB] * conj(B1[mB,s,rB])) * A1[mA,s,rA] #NOTE: Intermediates, final permutation.
    GC.enable(true)
    GC.gc(false)
    TM_convert(TM21)
end

function applyTM_l(A::MPSTensor, B::MPSTensor, x::Matrix)
    GC.enable(false)
    @tensor res[a, d] := conj(A[b, s, a]) * (x[b, c] * B[c, s, d]) #NOTE: This requires at least one intermediate array
    GC.enable(true)
    GC.gc(false)
    res
end
applyTM_l(A::MPSTensor, B::MPSTensor, x::Vector) = applyTM_l(A, B, reshape(x, (bond_dim(A), bond_dim(B))))

function applyTM_r(A::MPSTensor, B::MPSTensor, x::Matrix)
    @tensor res[a, d] := (A[a, s, b] * x[b, c]) * conj(B[d, s, c]) #NOTE: This requires at least one intermediate array
end
applyTM_r(A::MPSTensor, B::MPSTensor, x::Vector) = applyTM_r(A, B, reshape(x, (bond_dim(A), bond_dim(B))))

function applyTM_r(A1::MPSTensor{T}, B1::MPSTensor{T}, TM2::MPS_MPO_TM{T}) where {T}
    TM2 = TM_convert(TM2)
    GC.enable(false)
    @tensor TM12[lA,lB, rA,rB] := conj(B1[lB,s,mB]) * (A1[lA,s,mA] * TM2[mA,mB, rA,rB]) #NOTE: Intermediates, final permutation.
    GC.enable(true)
    GC.gc(false)
    TM_convert(TM12)
end

function applyTM_r!(TM12::MPS_MPO_TM{T}, A1::MPSTensor{T}, B1::MPSTensor{T}, TM2::MPS_MPO_TM{T}, work::Vector{T}) where {T}
    TM2 = TM_convert(TM2)
    TM12c = TM_convert(TM12)

    wsz1 = (size(A1)[1:2]..., size(TM2)[2:4]...)
    wlen1 = prod(wsz1)
    A1TM2 = unsafe_reshaped_subvec(work, 1, wsz1)
    @tensor A1TM2[lA,s,mB, rA,rB] = A1[lA,s,mA] * TM2[mA,mB, rA,rB]

    wsz2 = (wsz1[2:3]..., wsz1[1], wsz1[4:5]...)
    wlen2 = prod(wsz2)
    A1TM2p = unsafe_reshaped_subvec(work, wlen1+1, wsz2) #take work vector section after TM2B1 storage
    @tensor A1TM2p[s,mB,lA, rA,rB] = A1TM2[lA,s,mB, rA,rB]
    A1TM2p_copy = unsafe_reshaped_subvec(work, 1, wsz2)
    copyto!(A1TM2p_copy, A1TM2p)
    A1TM2p = A1TM2p_copy

    TM12sz = size(TM12c)
    wsz3 = (TM12sz[2], TM12sz[1], TM12sz[3:4]...)
    TM12p = unsafe_reshaped_subvec(work, wlen2+1, wsz3)
    @tensor TM12p[lB,lA, rA,rB] = conj(B1[lB,s,mB]) * A1TM2p[s,mB,lA, rA,rB]

    @tensor TM12c[lA,lB, rA,rB] = TM12p[lB,lA, rA,rB]

    TM12
end

function res_applyTM_r(A1::MPSTensor{T}, B1::MPSTensor{T}, TM2::MPS_MPO_TM{T}) where {T}
    TM2 = TM_convert(TM2)
    zeros(T, (bond_dim_L(A1), 1, bond_dim_L(B1), size(TM2,3), 1, size(TM2,4)))
end

function res_applyTM_r!(res::MPS_MPO_TM{T}, A1::MPSTensor{T}, B1::MPSTensor{T}, TM2::MPS_MPO_TM{T}) where {T}
    TM2 = TM_convert(TM2)
    sz = (bond_dim_L(A1), 1, bond_dim_L(B1), size(TM2,3), 1, size(TM2,4))
    size(res) == sz ? res : zeros(T, sz)
end

workvec_applyTM_r(A1::MPSTensor{T}, B1::MPSTensor{T}) where {T} = Vector{T}(bond_dim(A1)^2 * bond_dim(B1)^2 * phys_dim(A1) * 2)

function worklen_applyTM_r(A1::MPSTensor{T}, B1::MPSTensor{T}, TM2::MPS_MPO_TM{T}) where {T}
    TM2 = TM_convert(TM2)
    len1 = prod((size(A1)[1:2]..., size(TM2)[2:4]...))
    len2 = prod((size(A1,1) * size(B1,1), size(TM2)[3:4]...))
    max(2len1, len1 + len2)
end

function workvec_applyTM_r(A1::MPSTensor{T}, B1::MPSTensor{T}, TM2::MPS_MPO_TM{T}) where {T}
    Vector{T}(undef, worklen_applyTM_r(A1, B1, TM2))
end

function workvec_applyTM_r!(work::Vector{T}, A1::MPSTensor{T}, B1::MPSTensor{T}, TM2::MPS_MPO_TM{T}) where {T}
    len = worklen_applyTM_r(A1, B1, TM2)
    work = length(work) < len ? resize!(work, len) : work
end

function tm_eigs_sparse(A::MPSTensor{T}, B::MPSTensor{T}, dirn::Symbol, nev::Integer=1;
                    x0::Matrix{T}=ones(T,(bond_dim(A), bond_dim(B)))) where {T}
    @assert dirn in (:L, :R)

    D = bond_dim(A)
    DB = bond_dim(B)

    x = zeros(T, (D,DB))
    if dirn == :L
        f = (v::AbstractVector) -> vec(applyTM_l(A, B, copyto!(x, v)))
    else
        f = (v::AbstractVector) -> vec(applyTM_r(A, B, copyto!(x, v)))
    end

    fmap = LinearMap{T}(f, D*DB)

    ev, eV, info = eigsolve(f, vec(x0), nev, :LM)
    info.converged < nev && @warn "$(nev - info.converged) eigenvectors did not converge after $(info.numiter) iterations."

    dirn == :L && conj!(ev)

    #Reshape eigenvectors to (D,DB) matrices
    eVm = Matrix[reshape(eV[j], (D,DB)) for j in 1:length(eV)]

    ev, eVm
end

function tm_eigs_dense(A::MPSTensor{T}, B::MPSTensor{T}) where {T}
    TM = TM_dense(A, B)

    DA = bond_dim(A)
    DB = bond_dim(B)

    TM = reshape(TM, (DA*DB,DA*DB))
    ev, eVr = eigen(TM)
    eVl = inv(eVr)'

    eVmr = Matrix[reshape(eVr[:,j], (DA,DB)) for j in 1:size(eVr, 2)]
    eVml = Matrix[reshape(eVl[:,j], (DA,DB)) for j in 1:size(eVl, 2)]

    ev, eVml, eVmr
end

function tm_eigs(A::MPSTensor{T}, B::MPSTensor{T}, min_nev::Integer; D_dense_max::Integer=8) where {T}
    if bond_dim(A) * bond_dim(B) > D_dense_max^2
        evl, eVl = tm_eigs_sparse(A, B, :L, min_nev)
        evr, eVr = tm_eigs_sparse(A, B, :R, min_nev)
        #rely on tm_eigs returning eigenvalues in order of descending magnitude
    else
        evr, eVl, eVr = tm_eigs_dense(A, B)
        evl = evr
    end

    evl, evr, eVl, eVr
end

function tm_dominant_eigs(A::MPSTensor{T}, B::MPSTensor{T}; D_dense_max::Integer=8) where {T}
    evl, evr, eVl, eVr = tm_eigs(A, B, 1, D_dense_max=D_dense_max)

    indmaxl = argmax(abs.(evl))
    indmaxr = argmax(abs.(evr))

    tol = sqrt(eps(real(T))) #If one of the checks hits this upper bound, something went horribly wrong!

    if abs(evl[indmaxl] - evr[indmaxr]) > tol
        @warn "Largest magnitude eigenvalues do not match!" evl[indmaxl] evr[indmaxr]
    end

    dominant_ev = evl[indmaxl]

    if imag(dominant_ev) > tol
        @warn "Largest eigenvalue is not real!" dominant_ev
    end

    l = eVl[indmaxl]
    rmul!(l, 1.0/norm(l)) #tm_eigs_dense does not normalize the left eigenvectors
    r = eVr[indmaxr]

    #We need this to be 1
    n = vec(l) ⋅ vec(r)
    abs_n = abs(n)

    abs_n < tol && @warn "l and r are orthogonal!" abs_n

    phase_n = abs_n/n
    sfac = 1.0/sqrt(abs_n)

    #We can always find Hermitian l and r. Fix the phase on r using its trace, which must be nonzero.
    r_tr = tr(r)
    phase_r = abs(r_tr)/r_tr

    #Fix the phase on l using the phase of n
    phase_l = phase_r * conj(phase_n)

    rmul!(l, phase_l * sfac)
    rmul!(r, phase_r * sfac)

    #force exact Hermiticity
    lh = (l + l')/2
    rh = (r + r')/2

    lh = T <: Real ? real(lh) : lh
    rh = T <: Real ? real(rh) : rh

    l = convert(Matrix{T}, lh)
    r = convert(Matrix{T}, rh)

    norm(lh - l) > tol && @warn "l was not made Hermitian"
    norm(rh - r) > tol && @warn "r was not made Hermitian"

    dominant_ev, l, r
end

function gauge_transform(A::MPSTensor, g::Matrix, gi::Matrix)
    GC.enable(false)
    @tensor Anew[a,s,d] := g[a,b] * A[b,s,c] * gi[c,d] #NOTE: Requires a temporary array
    GC.enable(true)
    GC.gc(false)
    Anew
end

function canonicalize_left(l::Matrix{T}, r::Matrix{T}) where {T}
    tol = sqrt(eps(real(T)))

    evl, Ul = eigen(Hermitian(l))
    norm(Ul * Ul' - I) > tol && @warn "Nonunintary eigenvectors."

    sevl = Diagonal(sqrt.(complex.(evl)))
    g = sevl * Ul'
    gi = Ul * inv(sevl)

    r = g * Hermitian(r) * g'
    r = Hermitian((r + r')/2)

    evr, Ur = eigen(r)
    norm(Ur * Ur' - I) > tol && @warn "Nonunintary eigenvectors."

    gi = gi * Ur
    g = Ur' * g

    #It would be nice to use special matrix types here, but currently our tm functions can't handle them
    rnew = Diagonal(convert(Vector{T}, evr))
    lnew = I

    lnew, rnew, g, gi
end

MPO_open{T} = Vector{MPOTensor{T}}

#Multiplies MPO tensors top to bottom: A*B connects the bottom index of A with the top index of B.
function mul_MPO(A::MPOTensor, B::MPOTensor)
    M1l,d1t,M1r,d1b = size(A)
    M2l,d2t,M2r,d2b = size(B)
    @tensor AB[m1l,m2l,pt,m1r,m2r,pb] := A[m1l,pt,m1r,p] * B[m2l,p,m2r,pb]
    reshape(AB, (M1l*M2l, d1t, M1r*M2r, d2b))
end

function mul_MPO(M1::NTuple{N,MPOTensor{T}}, M2::NTuple{N,MPOTensor{T}}) where {N,T}
    ntuple(j -> mul_MPO(M1[j], M2[j]), N)
end

function mul_MPO(M1::Vector{MPOTensor{T}}, M2::Vector{MPOTensor{T}}) where {T}
    MPOTensor{T}[mul_MPO(M1[j], M2[j]) for j in 1:length(M1)]
end

function TM_dense_MPO(A::MPSTensor, B::MPSTensor, o::MPOTensor)::MPS_MPO_TM
    GC.enable(false)
    @tensor TM[k1,m1,b1, k2,m2,b2] := (A[k1, s1, k2] * o[m1,s1,m2,t1]) * conj(B[b1, t1, b2])
    GC.enable(true)
    GC.gc(false)
    TM
end

function applyTM_MPO_l(A::MPSTensor, B::MPSTensor, o::MPOTensor, TM2::MPS_MPO_TM)::MPS_MPO_TM
    GC.enable(false)
    @tensor TM21[k1,m1,b1, k3,m3,b3] := ((TM2[k1,m1,b1, k2,m2,b2] * conj(B[b2, t2, b3])) * o[m2,s2,m3,t2]) * A[k2, s2, k3]
    GC.enable(true)
    GC.gc(false)
    TM21
end

function applyTM_MPO_r(A::MPSTensor, B::MPSTensor, o::MPOTensor, TM2::MPS_MPO_TM)::MPS_MPO_TM
    GC.enable(false)
    @tensor TM12[k1,m1,b1, k3,m3,b3] :=  conj(B[b1, t1, b2]) * (o[m1,s1,m2,t1] * (A[k1, s1, k2] * TM2[k2,m2,b2, k3,m3,b3]))
    GC.enable(true)
    GC.gc(false)
    TM12
end

function worklen_applyTM_MPO_l(A1::MPSTensor{T}, B1::MPSTensor{T}, o::MPOTensor{T}, TM2::MPS_MPO_TM{T}) where {T}
    len = 0

    wsz1 = (size(TM2)[1:5]..., size(B1)[2:3]...)
    wlen1 = prod(wsz1)
    wsz2 = (wsz1[1:4]..., wsz1[7], wsz1[5:6]...)
    wlen2 = prod(wsz2)

    len = max(len, wlen2 + wlen1)

    opsz = size(o)
    wsz3 = (opsz[1], opsz[4], opsz[2:3]...)
    wlen3 = prod(wsz3)

    wsz4 = (wsz2[1:5]..., wsz3[3:4]...)
    wlen4 = prod(wsz4)

    len = max(len, wlen1 + wlen3 + wlen4)

    wsz5 = (wsz4[1:3]..., wsz4[5], wsz4[7], wsz4[4], wsz4[6])
    wlen5 = prod(wsz5)

    wsz6 = (wsz5[1:5]..., size(A1)[3])
    wlen6 = prod(wsz6)

    if wlen5 <= wlen1+wlen3
        ind5 = 1
    else
        ind5 = wlen1+wlen3+wlen4+1
    end

    len = max(len, ind5-1 + wlen5)

    if ind5 == 1 || wlen6 >= ind5
        len = max(len, ind5-1 + wlen5 + wlen6)
    else
        len = max(len, wlen6)
    end

    len
end

function workvec_applyTM_MPO_l(A1::MPSTensor{T}, B1::MPSTensor{T}, o::MPOTensor{T}, TM2::MPS_MPO_TM{T}) where {T}
    Vector{T}(undef, worklen_applyTM_MPO_l(A1, B1, o, TM2))
end

function workvec_applyTM_MPO_l!(work::Vector{T}, A1::MPSTensor{T}, B1::MPSTensor{T}, o::MPOTensor{T}, TM2::MPS_MPO_TM{T}) where {T}
    len = worklen_applyTM_MPO_l(A1, B1, o, TM2)
    work = length(work) < len ? resize!(work, len) : work
end

function res_applyTM_MPO_l(A1::MPSTensor{T}, B1::MPSTensor{T}, o::MPOTensor{T}, TM2::MPS_MPO_TM{T})::MPS_MPO_TM{T} where {T}
    Array{T,6}(undef, (size(TM2)[1:3]..., bond_dim_R(A1),size(o,3),bond_dim_R(B1)))
end

function res_applyTM_MPO_l!(res::MPS_MPO_TM{T}, A1::MPSTensor{T}, B1::MPSTensor{T}, o::MPOTensor{T}, TM2::MPS_MPO_TM{T})::MPS_MPO_TM{T} where {T}
    sz = (size(TM2)[1:3]..., bond_dim_R(A1), size(o,3), bond_dim_R(B1))
    size(res) == sz ? res : zeros(T, sz)
end

#NOTE: This works, but is pretty horrible. We are doing manual memory management because the garbage collector is super slow!
function applyTM_MPO_l!(TM21::MPS_MPO_TM{T}, A1::MPSTensor{T}, B1::MPSTensor{T}, o::MPOTensor{T}, TM2::MPS_MPO_TM{T}, work::Vector{T})::MPS_MPO_TM{T} where {T}
    wsz1 = (size(TM2)[1:5]..., size(B1)[2:3]...)
    wlen1 = prod(wsz1)
    wsz2 = (wsz1[1:4]..., wsz1[7], wsz1[5:6]...)
    wlen2 = prod(wsz2)

    TM2B1 = unsafe_reshaped_subvec(work, wlen2+1, wsz1)
    @tensor TM2B1[k1,m1,b1, k2,m2,t2,b3] = TM2[k1,m1,b1, k2,m2,b2] * conj(B1[b2, t2, b3])

    TM2B1p = unsafe_reshaped_subvec(work, 1, wsz2)
    @tensor TM2B1p[k1,m1,b1, k2,b3, m2,t2] = TM2B1[k1,m1,b1, k2,m2,t2,b3]

    opsz = size(o)
    wsz3 = (opsz[1], opsz[4], opsz[2:3]...)
    wlen3 = prod(wsz3)

    wsz4 = (wsz2[1:5]..., wsz3[3:4]...)
    wlen4 = prod(wsz4)

    op = unsafe_reshaped_subvec(work, wlen1+1, wsz3)
    @tensor op[m2,t2,s2,m3] = o[m2,s2,m3,t2]

    TM2B1pop = unsafe_reshaped_subvec(work, wlen1+wlen3+1, wsz4)
    @tensor TM2B1pop[k1,m1,b1, k2,b3, s2,m3] = TM2B1p[k1,m1,b1, k2,b3, m2,t2] * op[m2,t2,s2,m3]

    wsz5 = (wsz4[1:3]..., wsz4[5], wsz4[7], wsz4[4], wsz4[6])
    wlen5 = prod(wsz5)

    wsz6 = (wsz5[1:5]..., size(A1)[3])
    wlen6 = prod(wsz6)

    if wlen5 <= wlen1+wlen3
        ind5 = 1
    else
        ind5 = wlen1+wlen3+wlen4+1
    end
    TM2B1pop_p = unsafe_reshaped_subvec(work, ind5, wsz5)
    @tensor TM2B1pop_p[k1,m1,b1, b3, m3, k2,s2] = TM2B1pop[k1,m1,b1, k2,b3, s2,m3]

    if ind5 == 1 || wlen6 >= ind5
        TM21p = unsafe_reshaped_subvec(work, wlen5+ind5, wsz6)
    else
        TM21p = unsafe_reshaped_subvec(work, 1, wsz6)
    end
    @tensor TM21p[k1,m1,b1, b3,m3,k3] = TM2B1pop_p[k1,m1,b1, b3,m3, k2,s2] * A1[k2,s2, k3]

    @tensor TM21[k1,m1,b1, k3,m3,b3] = TM21p[k1,m1,b1, b3,m3,k3]
    TM21
end

#Convenience functions that fall back on MPO-less versions in case o is the Identity.
workvec_applyTM_MPO_l(A1::MPSTensor{T}, B1::MPSTensor{T}, o::IdentityMPOTensor, TM2::MPS_MPO_TM{T})  where {T} = workvec_applyTM_l(A1, B1, TM2)
workvec_applyTM_MPO_l!(work::Vector{T}, A1::MPSTensor{T}, B1::MPSTensor{T}, o::IdentityMPOTensor, TM2::MPS_MPO_TM{T}) where {T} = workvec_applyTM_l!(work, A1, B1, TM2)
function applyTM_MPO_l!(TM21::MPS_MPO_TM{T}, A1::MPSTensor{T}, B1::MPSTensor{T}, o::IdentityMPOTensor, TM2::MPS_MPO_TM{T}, work::Vector{T})::MPS_MPO_TM{T} where {T}
    applyTM_l!(TM21, A1, B1, TM2, work)
end

function workvec_alloc_applyTM_MPO_r(A::MPSTensor{T}, B::MPSTensor{T}, o::MPOTensor{T}, TM2::MPS_MPO_TM{T}) where {T}
    ATM2_sz = (size(A)[1:2]..., size(TM2)[2:end]...)
    ATM2_len = prod(ATM2_sz)

    ATM2p_sz = (ATM2_sz[2:3]..., ATM2_sz[1], ATM2_sz[4:end]...)
    ATM2p_len = prod(ATM2p_sz)

    op_sz = (size(o)[1], size(o)[4], size(o)[2], size(o)[3])
    op_len = prod(op_sz)

    oATM2_sz = (op_sz[1:2]..., ATM2p_sz[3:end]...)
    oATM2_len = prod(oATM2_sz)

    oATM2p_sz = (oATM2_sz[2], oATM2_sz[4], oATM2_sz[1], oATM2_sz[3], oATM2_sz[5:end]...)
    oATM2p_len = prod(oATM2p_sz)

    TM21p_sz = (size(B)[1], oATM2p_sz[3:end]...)
    TM21p_len = prod(TM21p_sz)

    ATM2_ind = ATM2p_len+1
    ATM2p_ind = 1
    op_ind = ATM2p_len+1
    oATM2_ind = ATM2p_len+op_len+1
    if ATM2p_len+op_len >= oATM2_len
        oATM2p_ind = 1
    else
        oATM2p_ind = ATM2p_len+op_len+oATM2_len+1
    end
    if oATM2p_ind >= TM21p_len
        TM21p_ind = 1
    else
        TM21p_ind = oATM2p_ind + oATM2p_len
    end

    (ATM2_sz, ATM2p_sz, op_sz, oATM2_sz, oATM2p_sz, TM21p_sz), (ATM2_ind, ATM2p_ind, op_ind, oATM2_ind, oATM2p_ind, TM21p_ind)
end

function workvec_applyTM_MPO_r!(work::Vector{T}, A::MPSTensor{T}, B::MPSTensor{T}, o::MPOTensor{T}, TM2::MPS_MPO_TM{T}) where {T}
    szs, inds = workvec_alloc_applyTM_MPO_r(A, B, o, TM2)

    len = maximum(inds[j] + prod(szs[j]) for j in 1:6)

    work = length(work) < len ? resize!(work, len) : work
end

workvec_applyTM_MPO_r(A::MPSTensor{T}, B::MPSTensor{T}, o::MPOTensor{T}, TM2::MPS_MPO_TM{T}) where {T} =
    workvec_applyTM_MPO_r!(Vector{T}(), A, B, o, TM2)

function res_applyTM_MPO_r(A::MPSTensor{T}, B::MPSTensor{T}, o::MPOTensor{T}, TM2::MPS_MPO_TM{T})::MPS_MPO_TM{T} where {T}
    Array{T,6}(undef, (bond_dim_L(A),size(o,1),bond_dim_L(B),size(TM2)[4:6]...))
end

function res_applyTM_MPO_r!(res::MPS_MPO_TM{T}, A1::MPSTensor{T}, B1::MPSTensor{T}, o::MPOTensor{T}, TM2::MPS_MPO_TM{T})::MPS_MPO_TM{T} where {T}
    sz = ((bond_dim_L(A),size(o,1),bond_dim_L(B),size(TM2)[4:6]...))
    size(res) == sz ? res : zeros(T, sz)
end

function applyTM_MPO_r!(TM21::MPS_MPO_TM{T}, A::MPSTensor{T}, B::MPSTensor{T}, o::MPOTensor{T}, TM2::MPS_MPO_TM{T}, work::Vector{T})::MPS_MPO_TM{T} where {T}
    szs, inds = workvec_alloc_applyTM_MPO_r(A, B, o, TM2)
    ATM2_sz, ATM2p_sz, op_sz, oATM2_sz, oATM2p_sz, TM21p_sz = szs
    ATM2_ind, ATM2p_ind, op_ind, oATM2_ind, oATM2p_ind, TM21p_ind = inds

    ATM2 = unsafe_reshaped_subvec(work, ATM2_ind, ATM2_sz)
    @tensor ATM2[k1,s1, m2,b2, k3,m3,b3] = A[k1, s1, k2] * TM2[k2,m2,b2, k3,m3,b3]

    ATM2p = unsafe_reshaped_subvec(work, ATM2p_ind, ATM2p_sz)
    @tensor ATM2p[s1,m2, k1, b2, k3,m3,b3] = ATM2[k1,s1, m2,b2, k3,m3,b3]

    op = unsafe_reshaped_subvec(work, op_ind, op_sz)
    @tensor op[m1,t1,s1,m2] = o[m1,s1,m2,t1]

    oATM2 = unsafe_reshaped_subvec(work, oATM2_ind, oATM2_sz)
    @tensor oATM2[m1,t1, k1, b2, k3,m3,b3] = op[m1,t1,s1,m2] * ATM2p[s1,m2, k1, b2, k3,m3,b3]

    oATM2p = unsafe_reshaped_subvec(work, oATM2p_ind, oATM2p_sz)
    @tensor oATM2p[t1,b2, m1, k1, k3,m3,b3] = oATM2[m1,t1, k1, b2, k3,m3,b3]

    TM21p = unsafe_reshaped_subvec(work, TM21p_ind, TM21p_sz)
    @tensor TM21p[b1, m1, k1, k3,m3,b3] = conj(B[b1, t1, b2]) * oATM2p[t1,b2, m1, k1, k3,m3,b3]

    @tensor TM21[k1,m1,b1, k3,m3,b3] = TM21p[b1, m1, k1, k3,m3,b3]
end

workvec_applyTM_MPO_r(A1::MPSTensor{T}, B1::MPSTensor{T}, o::IdentityMPOTensor, TM2::MPS_MPO_TM{T}) where {T} = workvec_applyTM_r(A1, B1, TM2)
workvec_applyTM_MPO_r!(work::Vector{T}, A1::MPSTensor{T}, B1::MPSTensor{T}, o::IdentityMPOTensor, TM2::MPS_MPO_TM{T}) where {T} = workvec_applyTM_r!(work, A1, B1, TM2)
function applyTM_MPO_r!(TM21::MPS_MPO_TM{T}, A1::MPSTensor{T}, B1::MPSTensor{T}, o::IdentityMPOTensor, TM2::MPS_MPO_TM{T}, work::Vector{T})::MPS_MPO_TM{T} where {T}
    applyTM_r!(TM21, A1, B1, TM2, work)
end

end

max_stars_repo_name: FuTen/puMPS

max_issues_repo_name: FuTen/puMPS

max_forks_repo_name: FuTen/puMPS


text:
using Random
using StatsBase: sample!
using OffsetArrays
#implements https://arxiv.org/pdf/2006.07060.pdf
#thousands of times faster than https://github.com/manhtuando97/KDD-20-Hypergraph/blob/master/Code/Generator/hyper_preferential_attachment.py

function hyper_pa(degree_distribution, edgesize_distribution, max_edgesize::Integer, nodes::I;
    rng::AbstractRNG=Random.GLOBAL_RNG) where I <: Integer

    edges = [[a,a+I(1)] for a in I(1):I(2):I(max_edgesize-1)]
    edges_by_size = [Vector{I}[] for _ in 1:max_edgesize]
    edges_by_size[2] = copy(edges)

    cum_sum_source = Vector{Int}(undef, max_edgesize)

    binomial_table = [binomial(a,b) for a in 1:max_edgesize, b in 1:max_edgesize-1]

    for n in last(last(edges))+1:nodes
        for _ in 1:rand(degree_distribution)
            new_edgesize = rand(edgesize_distribution)
            new_edge = Vector{I}(undef, new_edgesize)
            new_edge[1] = n

            if new_edgesize > 1
                #binom = 1
                acc = 0
                cum_sum = OffsetVector(view(cum_sum_source, 1:max_edgesize-new_edgesize+2), (new_edgesize-1):max_edgesize) #13%
                for source_edgesize in eachindex(cum_sum)
                    binom = binomial_table[source_edgesize,new_edgesize-1]
                    acc += length(edges_by_size[source_edgesize])*binom
                    cum_sum[source_edgesize] = acc
                    #binom = binom*(source_edgesize+1)÷(source_edgesize-new_edgesize+2)
                end

                total = last(cum_sum)
                if total == 0
                    sample!(rng, 1:(n-1), (@view new_edge[2:end]), replace=false) # sample! mutates its 3rd argument only.
                else
                    key = rand(rng, 1:total)
                    source_edgesize = firstindex(cum_sum)
                    while #=@inbounds=# cum_sum[source_edgesize] < key
                        source_edgesize += 1
                    end

                    #Huzzah! we have a source edgesize!

                    source_edge = rand(rng, edges_by_size[source_edgesize])

                    #Huzzah! and a specific source edge

                    # This sampling takes ~38% of runtime. ~20% of that (~7.5% total) can be
                    # alleviated by passing a workspace vector through. StatsBase doesn't
                    # support that, though.
                    sample!(rng, source_edge, (@view new_edge[2:end]), replace=false) # sample! mutates its 3rd argument only.

                    #Huzzah! and an actual edge to use
                end
            end

            push!(edges, new_edge)
            @assert length(new_edge) == new_edgesize
            push!(edges_by_size[new_edgesize], new_edge)
        end
    end

    edges
end

max_stars_repo_name: LilithHafner/FBD

max_issues_repo_name: LilithHafner/FBD

max_forks_repo_name: LilithHafner/FBD


text:
module SimpleNavierStokes
export LidDrivenCavity, ShowStreamlines

using LinearAlgebra, SparseArrays, IterativeSolvers, Plots

include("functions.jl")
include("lid_driven_cavity.jl")
include("plot_utils.jl")

end # module


text:
"""
    $(SIGNATURES)

Columns are symbol, description, value.

Flow:
1. Constructor
2. Add rows, either as string (e.g. `\\midrule`) or as triple of values
3. Write table
"""
mutable struct ParameterTable
    tb :: Table
end

"""
    $(SIGNATURES)

Constructor.
"""
function ParameterTable()
    pt = ParameterTable(Table("lll"));
    add_row!(pt, "Parameter", "Description", "Value");
    add_row!(pt, "\\midrule");
    return pt
end

# Add row from string
# Mainly for headers, lines
function add_row!(p :: ParameterTable, rowStr :: T1) where T1 <: AbstractString
    add_row!(p.tb, rowStr)
    return nothing
end

# Add row from parameter values
function add_row!(p :: ParameterTable, symStr :: T1, descrStr :: T1, valueStr :: T1) where
    T1 <: AbstractString

    add_row!(p.tb, symStr * " & " * descrStr * " & " * valueStr);
    return nothing
end

function write_table(p :: ParameterTable, filePath :: T1) where T1 <: AbstractString
    @assert nrows(p.tb) > 0  "Table is empty"
    write_table(p.tb, filePath);
    return nothing
end

# ------------


text:
function loadGAP(instance::Symbol)
    file_name = joinpath(data_path, string(instance) * ".zip")
    if !isfile(file_name)
        println("File $(string(instance)) not found!")
        return nothing
    end

    name = splitext(basename(file_name))[1]
    file = ZipFile.Reader(file_name)
    values = parse.(Int64, split(read(file.files[1], String)))

    return load(values, name)
end

function loadGAP(file_name::String)
    if !isfile(file_name)
        println("File $file_name not found!")
        return nothing
    end

    name = splitext(basename(file_name))[1]
    values = parse.(Int64, split(read(file_name, String)))

    return load(values, name)
end

function load(values::Array{Int64}, name::String)
    m = values[1]
    n = values[2]

    counter = 3

    costs = zeros(Int64, n, m)
    for j in 1:m
        costs[:, j] = values[counter:counter + n - 1]
        counter += n
    end

    consumptions = zeros(Int64, n, m)
    for j in 1:m
        consumptions[:, j] = values[counter:counter + n - 1]
        counter += n
    end

    capacities = values[counter:counter + m - 1]

    agents = [ GAPAgent(j, capacities[j], costs[:, j], consumptions[:, j]) for j in 1:m ]
    jobs = [ GAPJob(i, costs[i, :], consumptions[i, :]) for i in 1:n ]

    return GAPData(name, agents, jobs, typemin(Int64), typemax(Int64))
end


text:
function learn!(env::E, qpolicy::Q, mem::M, num_eps, γ;
                maxn=200, opt=ADAM(0.00001), update_freq=3000, chkpt_freq=3000,
                train_freq=100, chkpt_filename="model_checkpoint.bson",
                cb_ep = () -> (), cb_step = () -> (),
                show_progress=true) where {E<:Reinforce.AbstractEnvironment,
                                           Q<:QPolicy,
                                           M<:ReplayMemoryBuffer}

    # Build an epsilon greedy policy for the learning
    π = ϵGreedyPolicy(LinearSequence(0.01,0.01,num_eps,iter=0), qpolicy)

    # Track number of sucessful attempts
    num_successes = 0

    # Params to optimize
    p = get_params(qpolicy)

    losses = Float64[]

    # Track the number of training steps completed so far
    numStepsCompleted = 1
    show_progress && (progress = Progress(num_eps, 3))
    for i ∈ 1:num_eps
        ep = Episode(env, π; maxn = maxn)

        for (s, a, r, s′) ∈ ep
            # Save the step into the replay buffer
            addexp!(mem, s(), a, r, s′(), finished(env, s′))

            # Fill the buffer before training or lowering ϵ
            if length(mem) < mem.batch_size
                # cb_step()
                continue
            end

            # Training - every train_freq steps
            if numStepsCompleted % train_freq == 0
                td_errs = batchTrain!(π, mem, qpolicy, γ, p, opt)
                cb_step()
            end

            # If needed, update the target Network
            update_freq > 0 && numStepsCompleted % update_freq == 0 && update_target(qpolicy)
            # update_freq > 0 && numStepsCompleted % update_freq == 0 && push!(losses, Flux.huber_loss(td_errs))

            # If desired, save the network
            chkpt_freq > 0 && numStepsCompleted % chkpt_freq == 0 && save_policy(qpolicy)

            # Record if this episode was successful
            finished(env, s′) && (num_successes += 1)

            numStepsCompleted += 1

            # Run the step callback
            # cb_step()

        end # end of episode

        # Change β
        update_β!(mem)

        # Run the episode callback
        cb_ep()

        # Update the progress
        show_progress && next!(progress)
    end
    chkpt_freq > 0 && save_policy(qpolicy, chkpt_filename)
    num_successes, losses
end

function batchTrain!(π, mem, qpolicy, γ, params, opt)
    # Decrease ϵ
    update_ϵ!(π)

    # Sample a batch from the replay buffer
    (s_batch, a_batch, r_batch, s′_batch, done_batch, ids, weights) = sample(mem)

    # Get the target values based on the next states
    target = get_target(qpolicy, γ, r_batch, done_batch, s′_batch)

    # May need this to help set the data type for loss - not sure
    loss = 0.0

    # Track the td errs
    td_errs = similar(target)

    # Track gradients while calculating the loss
    gs = Flux.gradient(params) do
        currentQ_SA = get_QValues(qpolicy, s_batch)[a_batch]
        td_errs = target .- currentQ_SA
        loss = Flux.huber_loss(td_errs.*weights)
    end

    # Train the network(s)
    Flux.Optimise.update!(opt, params, gs)

    # Update Priorities for selected memory elements
    update_priorities!(mem, ids, td_errs)

    # return the td_errs so we can save them for later
    return td_errs
end


text:
using StatsBase

typealias  Graph Dict{Int, Set{Int}}

function addnode!(g::Graph, u::Int)
    if !haskey(g, u)
        g[u] = valtype(g)()
    end
    nothing
end

function delnode!(g::Graph, u::Int)
    if haskey(g, u)
        for v in g[u]
            delete!(g[v], u)
        end
        delete!(g, u)
    end
    nothing
end

function addedge!(g::Graph, u::Int, v::Int)
    addnode!(g, u)
    addnode!(g, v)
    push!(g[u], v)
    push!(g[v], u)
    nothing
end

function deledge!(g::Graph, u::Int, v::Int)
    if haskey(g, u) && haskey(g, v)
        delete!(g[u], v)
        delete!(g[v], u)
    end
    nothing
end

function ne(g::Graph)
    m = 0
    for u in keys(g)
        m += length(g[u])
    end
    return div(m, 2)
end

neighbors(g::Graph, u::Int) = g[u]

nodes(g::Graph) = keys(g)

hasnode(g::Graph, u::Int) = haskey(g, u)

hasedge(g::Graph, u::Int, v::Int) = hasnode(g, u) && in(v, g[u])

"""Type to record neighbor labels and their counts."""
type NeighComm
    neigh_pos::Vector{Int}
    neigh_cnt::Vector{Int}
    neigh_last::Int
end

type NodeStatus
    neigh_pos::Vector{Int}
    neigh_cnt::Vector{Bool}
    neigh_last::Int
end

function range_shuffle!(r, a::AbstractVector)
    @inbounds for i=r:-1:2
        j = StatsBase.randi(i)
        a[i],a[j] = a[j],a[i]
    end
end

"""Return the most frequency label."""
function pre_vote!(g, m, c::NeighComm, u, node_status)
    @inbounds for i=1:c.neigh_last-1
        c.neigh_cnt[c.neigh_pos[i]] = -1
    end
    c.neigh_last = 1
    c.neigh_pos[1] = m[u]
    c.neigh_cnt[c.neigh_pos[1]] = 0
    c.neigh_last = 2
    max_cnt = 0
    for neigh in neighbors(g,u)
        #if in(neigh, all_active_nodes)
        if node_status.neigh_cnt[neigh]
            neigh_comm = m[neigh]
            if c.neigh_cnt[neigh_comm] < 0
                c.neigh_cnt[neigh_comm] = 0
                c.neigh_pos[c.neigh_last] = neigh_comm
                c.neigh_last += 1
            end
            c.neigh_cnt[neigh_comm] += 1
            if c.neigh_cnt[neigh_comm] > max_cnt
                max_cnt = c.neigh_cnt[neigh_comm]
            end
        end
    end
    # ties breaking randomly
    range_shuffle!(c.neigh_last-1, c.neigh_pos)
    for lbl in c.neigh_pos
        if c.neigh_cnt[lbl] == max_cnt
            return lbl
        end
    end
end

function pre_update!(g, m, c::NeighComm, active_nodes, random_order, node_status)
    while !isempty(active_nodes)
        #random_order = collect(active_nodes)
        i = 0
        for u in active_nodes
            i += 1
            random_order[i] = u
        end
        #shuffle!(random_order)
        range_shuffle!(i, random_order)
        for j=1:i
            u = random_order[j]
            old_comm = m[u]
            new_comm = pre_vote!(g, m, c, u, node_status)
            if new_comm != old_comm
                for v in neighbors(g,u)
                    #if in(v, all_active_nodes)
                    if node_status.neigh_cnt[v]
                        push!(active_nodes, v)
                        m[u] = new_comm
                    end
                end
            else
                delete!(active_nodes, u)
            end
        end
    end
end

"""Return the most frequency label."""
function vote!(g, m, c::NeighComm, u)
    @inbounds for i=1:c.neigh_last-1
        c.neigh_cnt[c.neigh_pos[i]] = -1
    end
    c.neigh_last = 1
    c.neigh_pos[1] = m[u]
    c.neigh_cnt[c.neigh_pos[1]] = 0
    c.neigh_last = 2
    max_cnt = 0
    for neigh in neighbors(g,u)
        neigh_comm = m[neigh]
        if c.neigh_cnt[neigh_comm] < 0
            c.neigh_cnt[neigh_comm] = 0
            c.neigh_pos[c.neigh_last] = neigh_comm
            c.neigh_last += 1
        end
        c.neigh_cnt[neigh_comm] += 1
        if c.neigh_cnt[neigh_comm] > max_cnt
            max_cnt = c.neigh_cnt[neigh_comm]
        end
    end
    # ties breaking randomly
    range_shuffle!(c.neigh_last-1, c.neigh_pos)
    for lbl in c.neigh_pos
        if c.neigh_cnt[lbl] == max_cnt
            return lbl
        end
    end
end

function update!(g, m, c::NeighComm, active_nodes, random_order)
    num_active_nodes = 0
    while !isempty(active_nodes)
        if length(active_nodes) > num_active_nodes
            num_active_nodes = length(active_nodes)
        end
        #random_order = collect(active_nodes)
        i = 0
        for u in active_nodes
            i += 1
            random_order[i] = u
        end
        #shuffle!(random_order)
        range_shuffle!(i, random_order)
        for j=1:i
            u = random_order[j]
            old_comm = m[u]
            new_comm = vote!(g, m, c, u)
            if new_comm != old_comm
                for v in neighbors(g,u)
                    push!(active_nodes, v)
                    m[u] = new_comm
                end
            else
                delete!(active_nodes, u)
            end
        end
    end
    num_active_nodes
end

function lpa_addnode!(g, m, u)
    if !hasnode(g, u)
        addnode!(g, u)
        m[u] = u
    end
end

function lpa_deletenode!(g, m, c, u, active_nodes, random_order, node_status)
    if hasnode(g, u)
        empty!(active_nodes)
        u_comm = m[u]
        delnode!(g, u)
        #delete!(m, u)
        m[u] = 0
        node_status.neigh_last = 0
        for v in nodes(g)
            if m[v] == u_comm
                push!(active_nodes, v)
                node_status.neigh_last += 1
                node_status.neigh_pos[node_status.neigh_last] = v
                node_status.neigh_cnt[v] = true
                m[v] = v
            end
        end
        pre_update!(g, m, c, active_nodes, random_order, node_status)
        for i=1:node_status.neigh_last
            push!(active_nodes, node_status.neigh_pos[i])
            node_status.neigh_cnt[node_status.neigh_pos[i]] = false
        end
        return update!(g, m, c, active_nodes, random_order)
    end
    return 0
end

function lpa_addedge!(g, m, c, u, v, active_nodes, random_order, node_status)
    lpa_addnode!(g, m, u)
    lpa_addnode!(g, m, v)
    if !hasedge(g, u, v)
        u_comm = m[u]
        v_comm = m[v]
        if u_comm != v_comm
            ku_int = 0
            ku_ext = 0
            kv_int = 0
            kv_ext = 0
            for neigh in neighbors(g,u)
                if m[neigh] == u_comm
                    ku_int += 1
                else
                    ku_ext += 1
                end
            end
            for neigh in neighbors(g,v)
                if m[neigh] == v_comm
                    kv_int += 1
                else
                    kv_ext += 1
                end
            end
            if ku_ext+1 >= ku_int || kv_ext+1 >= kv_int
                #empty!(active_nodes)
                #empty!(active_nodes)
                node_status.neigh_last = 0
                for i in nodes(g)
                    i_comm = m[i]
                    if i_comm == u_comm || i_comm == v_comm
                        push!(active_nodes, i)
                        #push!(active_nodes, i)
                        node_status.neigh_last += 1
                        node_status.neigh_pos[node_status.neigh_last] = i
                        node_status.neigh_cnt[i] = true
                        m[i] = i
                    end
                end
                addedge!(g, u, v)
                pre_update!(g, m, c, active_nodes, random_order, node_status)
                #empty!(pre_active_nodes)
                for i=1:node_status.neigh_last
                    push!(active_nodes, node_status.neigh_pos[i])
                    node_status.neigh_cnt[node_status.neigh_pos[i]] = false
                end
                return update!(g, m, c, active_nodes, random_order)
            else
                addedge!(g, u, v)
            end
        else
            addedge!(g, u, v)
        end
    end
    return 0
end

function lpa_deleteedge!(g, m, c, u, v, active_nodes, random_order, node_status)
    if hasedge(g, u, v)
        u_comm = m[u]
        v_comm = m[v]
        if u_comm == v_comm
            empty!(active_nodes)
            node_status.neigh_last = 0
            for i in nodes(g)
                if m[i] == u_comm
                    #push!(pre_active_nodes, i)
                    push!(active_nodes, i)
                    node_status.neigh_last += 1
                    node_status.neigh_pos[node_status.neigh_last] = i
                    node_status.neigh_cnt[i] = true
                    m[i] = i
                end
            end
            deledge!(g, u, v)
            pre_update!(g, m, c, active_nodes, random_order, node_status)
            for i=1:node_status.neigh_last
                push!(active_nodes, node_status.neigh_pos[i])
                node_status.neigh_cnt[node_status.neigh_pos[i]] = false
            end
            return update!(g, m, c, active_nodes, random_order)
        else
            deledge!(g, u, v)
        end
    end
    return 0
end

using ParserCombinator: Parsers.GML

function _gml_read_one_graph(gs)
    nodes = [x[:id] for x in gs[:node]]
    g = Graph{Int,Int}()
    sds = [(Int(x[:source]), Int(x[:target])) for x in gs[:edge]]
    for (s,d) in (sds)
        addedge!(g, s, d)
    end
    return g
end

function loadgml(gname::AbstractString)
    p = GML.parse_dict(readall(gname))
    for gs in p[:graph]
        return _gml_read_one_graph(gs)
    end
    error("Graph $gname not found")
end

function savegml(gname::AbstractString, g, c=Vector{Integer}())
    io = open(gname, "w")
    println(io, "graph")
    println(io, "[")
    println(io, "directed 0")
    i = 0
    nodemap = Dict{Int,Int}()
    for u in keys(g)
        nodemap[u] = i
        println(io,"\tnode")
        println(io,"\t[")
        println(io,"\t\tid $i")
        println(io,"\t\tlabel $u")
        length(c) > 0 && println(io,"\t\tvalue $(c[u])")
        println(io,"\t]")
        i += 1
    end
    for u in keys(g)
        for v in keys(g[u])
            if u < v
                println(io,"\tedge")
                println(io,"\t[")
                println(io,"\t\tsource $(nodemap[u])")
                println(io,"\t\ttarget $(nodemap[v])")
                println(io,"\t]")
            end
        end
    end
    println(io, "]")
    close(io)
    return 1
end

function graphdiff(g1::Graph, g2::Graph)
    new_edges = Set{Pair{Int,Int}}()
    for u in nodes(g1)
        for v in neighbors(g1,u)
            if u<=v
                push!(new_edges, Pair(u,v))
            end
        end
    end
    for u in nodes(g2)
        for v in neighbors(g2,u)
            if u<=v
                delete!(new_edges, Pair(u,v))
            end
        end
    end
    new_edges
end

function graphdiff(graph1, graph2)
    edges_g1 = readdlm(graph1, Int)
    edges_g2 = readdlm(graph2, Int)
    new_edges = Set{Pair{Int,Int}}()
    for i=1:size(edges_g1, 1)
        u = edges_g1[i,1]
        v = edges_g1[i,2]
        if u>v
            u,v = v,u
        end
        push!(new_edges, Pair(u,v))
    end
    for i=1:size(edges_g2, 1)
        u = edges_g2[i,1]
        v = edges_g2[i,2]
        if u>v
            u,v = v,u
        end
        delete!(new_edges, Pair(u,v))
    end
    new_edges
end

function modularity(g, c)
    labels = unique(c)
    length(labels) == 1 && return 0.0
    no_of_edges = ne(g)
    m = 0
    e = Dict{Int,Int}()
    a = Dict{Int,Int}()
    m = no_of_edges
    for u in nodes(g)
        for v in neighbors(g,u)
            if u<v
                c1 = c[u]
                c2 = c[v]
                if c1 == c2
                    e[c1] = get(e,c1,0) + 2
                end
                a[c1] = get(a,c1,0) + 1
                a[c2] = get(a,c2,0) + 1
            end
        end
    end
    modularity_value = 0.0
    if m > 0
        for i in labels
            if i > 0
                tmp = haskey(a,i) ? a[i]/2/m : 0
                modularity_value += haskey(e,i) ? e[i]/2/m : 0
                modularity_value -= tmp*tmp
            end
        end
    end
    modularity_value
end


text:
function oldChooseMove(b::Chess.Board; variant = "standard")
    @debug("choosing move", b)
    move = search_in_opening_books(b; variant = variant)
    if !isnothing(move)
        @debug("returning bookmove")
        return move
    end
    @debug("no move found in book")
    myColor = sidetomove(b)
    possible = moves(b)

    if !isempty(possible)
        freeRealEstate = attacked_but_undefended(b, -myColor)
        anyFree = !isempty(freeRealEstate)
        @debug("freeRealEstate: ")
        pprint(b, highlight = freeRealEstate)
        takeable = attacked_with_pieces(b, -myColor)
        anyTakeable = !isempty(takeable)
        @debug("takeable squares/pieces")
        pprint(b, highlight = takeable)
        takingFree = Move[]
        takingAny = Move[]
        @debug("$(length(possible)) moves to evaluate")
        mCounter = 0
        bestPawnEval = 0
        bestPawnEvalFor = nothing
        @debug("starting loop")
        for m in possible
            @debug("start of iteration")
            undo = domove!(b, m)
            @debug("did move, undo information saved")
            if ischeckmate(b)
                @debug("found checkmate")
                undomove!(b, undo)
                return m
            end
            @debug("checked for amte")
            if any(mov -> move_is_mate(b, mov), moves(b))
                @debug("skipping move to not get mated in 1")
                undomove!(b, undo)
                continue
            end
            @debug("checked for incoming mate")
            if isstalemate(b)
                @debug("skipping move to avoid stalemating")
                undomove!(b, undo)
                continue
            end
            pE = pawnEval(b, myColor)
            @debug("pawnEval done")
            if pE > bestPawnEval
                bestPawnEval = pE
                bestPawnEvalFor = m
            end
            undomove!(b, undo)
            if anyFree && to(m) in freeRealEstate
                @debug("found move that takes unprotected")
                @show takingFree
                @debug(typeof(m))
                push!(takingFree, m)
                @debug(takingFree)
                #return m
            elseif anyTakeable && to(m) in takeable
                takingValue = value(pieceon(b, from(m)))
                takenValue = value(pieceon(b, to(m)))
                @debug "takingValue = $takingValue"
                @debug "takenValue = $takenValue"
                if takingValue <= takenValue
                    @debug("found move that takes higher valued but protected")
                    push!(takingAny, m)
                end
            end
            mCounter += 1
            @debug("\r", mCounter, "/", length(possible), " evaluated")
        end
        @show takingAny
        @show takingFree
        if !isempty(takingFree)
            @debug("returning random take of unprotected")
            return rand(takingFree)
        end
        if !isempty(takingAny)
            @debug("returning random move that takes higher valued")
            return rand(takingAny)
        end
        if !isnothing(bestPawnEvalFor)
            return bestPawnEvalFor
        end
        @debug("returning completely random move")
        return rand(possible)
        if !isempty(takingAny)
            @debug("returning random move that takes higher valued")
            return rand(takingAny)
        end
        @debug("returning completely random move, no mate no free")
        return rand(possible)
    end
    move
end

function attacked_but_undefended(board, color)
    # copied from Chess.jl tutorial at https://romstad.github.io/Chess.jl/stable/manual/#Square-Sets-1

    attacker = -color  # The opposite color

    # Find all attacked squares
    attacked = SS_EMPTY  # The empty square set
    for s ∈ pieces(board, attacker)
        attacked = attacked ∪ attacksfrom(board, s)
    end

    # Find all defended squares
    defended = SS_EMPTY
    for s ∈ pieces(board, color)
        defended = defended ∪ attacksfrom(board, s)
    end

    # Return all attacked, but undefended squares containing pieces of
    # the desired color:
    attacked ∩ -defended ∩ pieces(board, color)
end

function attacked_but_undefended(board, color, attacked)
    defended = SS_EMPTY
    for s ∈ pieces(board, color)
        defended = defended ∪ attacksfrom(board, s)
    end
    attacked ∩ -defended
end

function attacked_with_pieces(board::Board, color = sidetomove(board))
    attacker = -color
    attacked = SS_EMPTY
    for s ∈ pieces(board, attacker)
        attacked = attacked ∪ attacksfrom(board, s)
    end

    attacked ∩ pieces(board, color)
end

function move_is_mate(board, move)
    #from CHess.jl Tutorial
    # Do the move
    u = domove!(board, move)

    # Check if the resulting board is checkmate
    result = ischeckmate(board)

    # Undo the move
    undomove!(board, u)

    # Return result
    result
end

function value(p::Chess.Piece)
    p == PIECE_TYPE_NONE && return -1
    return ptype(p).val
end

function search_in_opening_books(b; variant = "standard")
    if variant == "atomic"
        move = pickbookmove(
            b,
            bookfile = "/home/dave/LichessGames/atomicOpening.obk",
        )
    else
        move = pickbookmove(
            b,
            bookfile = "/home/dave/LichessGames/lichess_elite.obk",
        )
        if isnothing(move)
            @debug("no move found in elite db, looking in own games")
            move = pickbookmove(
                b,
                bookfile = "/home/dave/LichessGames/myGames.obk",
            )
        end
    end
    return move
end

function pawnEval(b::Board, c::PieceColor, m::Move)
    undo = nothing
    homerank = c == WHITE ? RANK_1 : RANK_8
    undo = domove!(b, m)
    myPawns = pawns(b, c)
    isempty(myPawns) && return 0
    undomove!(b, undo)
    sum(square -> distance(homerank, rank(square))^2, myPawns)
end

function pawnEval(b::Board, c::PieceColor)
    homerank = c == WHITE ? RANK_1 : RANK_8
    myPawns = pawns(b, c)
    isempty(myPawns) && return 0
    sum(square -> distance(homerank, rank(square))^2, myPawns)
end

colorMult = Dict{Chess.PieceColor,Int}(((Chess.WHITE, 1), (Chess.BLACK, -1)))

pVals = [100, 300, 300, 500, 900, 1000, -1000]

function pieceValue(p::Piece)
    pVals[ptype(p).val%8]
end

function simpleEval(
    g::Union{Game,SimpleGame},
    d::Integer = 1;
    matescore::Integer = 10000,
)
    isdraw(g) && return 0
    simpleEval(board(g), d; matescore = matescore)
end

function simpleEval(b::Board, d::Integer = 1; matescore::Integer = 10000)
    score = 0
    if ischeckmate(b)
        return -max(d, 1) * matescore
    end
    if isdraw(b)
        return 0
    end
    c = sidetomove(b)
    multiplier = colorMult[c]
    for i = 1:14
        p = Piece(i)
        !isok(p) && continue
        myPiece = multiplier * colorMult[pcolor(p)]
        score += pieceValue(p) * myPiece * count_ones(pieces(b, p).val)
    end
    if count_ones(knights(b, c).val) > 1
        score += 50
    end

    if count_ones(bishops(b, c).val) > 1
        score += 50
    end

    if count_ones(knights(b, -c).val) > 1
        score -= 50
    end

    if count_ones(bishops(b, -c).val) > 1
        score -= 50
    end

    attacked = attacked_with_pieces(b, sidetomove(b))
    for s in attacked
        score += pieceValue(pieceon(b, s)) ÷ 5
    end
    undefended = attacked_but_undefended(b, sidetomove(b), attacked)
    for s in undefended
        score += pieceValue(pieceon(b, s)) ÷ 2
    end

    attacked = attacked_with_pieces(b, -sidetomove(b))
    for s in attacked
        score -= pieceValue(pieceon(b, s)) ÷ 5
    end
    undefended = attacked_but_undefended(b, -sidetomove(b), attacked)
    for s in undefended
        score -= pieceValue(pieceon(b, s)) ÷ 2
    end
    #score += attacked_with_pieces(b, sidetomove(b))÷3
    score += pawnEval(b, c)
    score
end

reallySimpleEval(g::Union{Game,SimpleGame}, d = 1; matescore = 10000)::Integer =
    isdraw(g) ? 0 : reallySimpleEval(board(g), d; matescore = matescore)

function reallySimpleEval(b::Board, d::Integer = 1; matescore::Integer = 10000)
    score = 0
    if ischeckmate(b)
        return -max(d, 1) * matescore
    end
    if isdraw(b)
        return 0
    end
    any(m -> LichessBot.move_is_mate(b, m), moves(b)) &&
        return max(d, 1) * matescore
    c = sidetomove(b)
    multiplier = colorMult[c]
    for i = 1:14
        p = Piece(i)
        !isok(p) && continue
        myPiece = multiplier * colorMult[pcolor(p)]
        score += pieceValue(p) * myPiece * count_ones(pieces(b, p).val)
    end
    d >= 3 && @info score
    attacked = attacked_with_pieces(b, sidetomove(b))
    for s in attacked
        score += pieceValue(pieceon(b, s)) ÷ 40
    end
    undefended = attacked_but_undefended(b, sidetomove(b), attacked)
    for s in undefended
        score += pieceValue(pieceon(b, s)) ÷ 2
    end
    #
    attacked = attacked_with_pieces(b, -sidetomove(b))
    for s in attacked
        score -= pieceValue(pieceon(b, s)) ÷ 50
    end
    undefended = attacked_but_undefended(b, -sidetomove(b), attacked)
    for s in undefended
        score -= pieceValue(pieceon(b, s)) ÷ 3
    end
    score += pawnEval(b, c)
    score -= pawnEval(b, -c)
    score += length(moves(b)) ÷ 2
    u = donullmove!(b)
    score -= length(moves(b)) ÷ 2
    undomove!(b, u)
    return score
end

checkVal = [90, 450, 5000]

function threeCheckPst(g, d::Integer, table = psttable; matescore = 10000)
    s = pstAndMate(g, d, table; matescore = matescore)
    c = sidetomove(board(g))
    bs = boards(g)
    checkScore = 0
    checkCounts =
        Dict{Chess.PieceColor,Int}(((Chess.WHITE, 0), (Chess.BLACK, 0)))
    for b in bs
        if ischeck(b)
            checkCounts[sidetomove(b)] += 1
            checkScore +=
                colorMult[sidetomove(b)] *
                4 *
                exp(2.4 * checkCounts[sidetomove(b)])
        end
    end
    s + colorMult[c] * checkScore
end

function isrepetitiondraw(g::Union{Game,SimpleGame})
    curBoard = board(g).key
    count(b -> b.key == curBoard, boards(g)) >= 3
end