using PowerDynamics: AbstractLine, construct_edge
using LinearAlgebra: I
import PowerDynamics: PiModel
using NetworkDynamics: ODEEdge

begin
    @__doc__ struct RLLine <: AbstractLine
            from
            to
            R
            L
            ω0
            RLLine(; from, to, R, L, ω0) = new(from, to, R, L, ω0)
        end
    function PowerDynamics.construct_edge(par::RLLine)
        from = par.from
        to = par.to
        R = par.R
        L = par.L
        ω0 = par.ω0
        @assert R > 0 "resistance (R) should be positive"
        @assert L > 0 "inductance (L) should be positive"
        @assert ω0 > 0 "rated frequency (ω0) should be positive"
        Z = [R -ω0*L; ω0*L R]
        function rhs!(de, e, v_s, v_d, p, t)
            # the different minus signs are due to the PowerDynamics sign convention for currents
            i_left_right = [e[1]; e[2]]
            i_right_left = [e[3]; e[4]]
            v_left = [v_s[1]; v_s[2]]
            v_right = [v_d[1]; v_d[2]]
            v_left_right = (- Z * i_left_right .- v_left .+ v_right) ./ L
            v_right_left = (- Z * i_right_left .+ v_right .- v_left) ./ L
            de .= [v_left_right; v_right_left]
        end
        #TODO check if there's a bug in the macro missing the empty space
        return ODEEdge(f! = rhs!, dim=4, mass_matrix=I, sym=Symbol[:i_r, :i_i, :ir_r, :ir_i])
    end
    symbolsof(::RLLine) = begin
            [:i_r, :i_i, :ir_r, :ir_i]
    end
    dimension(::RLLine) = begin
            4
    end
end

PiModel(l::RLLine) = PiModel(1/(l.R + im *l.ω0*l.L), 0, 0, 1, 1)


@testset "Spine2D" begin
    @polyvar x y
    f = x^2*y + y^2 + 3x^2*y^3 + y^4 + x^4*y^4

    S = spine(f)
    @test length(components_complement(S)) == 5
    @test sprint(show, S) ==
        """Spine2D{DynamicPolynomials.Polynomial{true,Tropical{Float64}}} with:
            $(" ")* 4 vertices
            $(" ")* 4 half-rays
            $(" ")* 1 bounded cells"""
    S = spine(f, minimal_component_size=0.05)
    @test length(components_complement(S)) == 5

    S = spine(f, domain=(-5, 5, -5, 5), minimal_component_size=0.01)
    @test length(components_complement(S)) == 5
end


using QuantumComputing
using Base.Test

# write your own tests here
a = Ket(0,1,0,1,0)
b = Ket(1,0,1,0,1)

n = length(a)
@test n==5
@test a == a
@test a != b

c = Swap(n,1,1)*a
@test c == a
c = Swap(n,1,2)*a
@test c == Ket(1,0,0,1,0)
c = Swap(n,1,n)*a
@test c == a

c = Swap(2,2,2)
for i=0:1, j=0:1
@test c*Ket(i,j)==Ket(i,j)
end

c = Not(5,1)
@test c*Ket(0,0,0,0,0)==Ket(1,0,0,0,0)
c = Not(n,n)
@test c*Ket(1,1,1,1,1)==Ket(1,1,1,1,0)

a = Ket(1,1)
c = CNot(2,1,2)*a
println(c)
@test c == Ket(0,1)

c = GCNot(2,2,1)
for i=0:1, j=0:1
@test c*Ket(i,j) == GGate(CNOT)*Ket(i,j)
end


function KronGraph500NoPerm(SCALE,EdgePerVertex)
    N = 2^SCALE

    M = round(EdgePerVertex * N)

    A = 0.57
    B = 0.19
    C = 0.19
    D = 1-(A+B+C)

    ij = Array{Int}(ones(2,M))
    ab = A+B
    c_norm = C/(1-(A+B))
    a_norm = A/(A+B)

    for ib = 1:SCALE
        ii_bit = rand(1,M) .> ab
        jj_bit = rand(1,M) .> (c_norm * ii_bit + a_norm * !(ii_bit))
        ij = ij + 2^(ib-1) * [ii_bit;jj_bit]
    end
    return reshape(ij[1,:]',size(ij,2)), reshape(ij[2,:]',size(ij,2))
end


# Adapted from:
# SOSDEMO9 --- Matrix SOS decomposition
# Section 3.9 of SOSTOOLS User's Manual

@testset "SOSDEMO9 with $(factory.constructor)" for factory in sdp_factories
    @polyvar x1 x2 x3

    P = [x1^4+x1^2*x2^2+x1^2*x3^2 x1*x2*x3^2-x1^3*x2-x1*x2*(x2^2+2*x3^2);
         x1*x2*x3^2-x1^3*x2-x1*x2*(x2^2+2*x3^2) x1^2*x2^2+x2^2*x3^2+(x2^2+2*x3^2)^2]

    # Test if P(x1,x2,x3) is an SOS matrix
    m = SOSModel(factory)
    # TODO return H so that P = H.'*H
    @SDconstraint m P ⪰ 0

    JuMP.optimize!(m)
    @test JuMP.primal_status(m) == MOI.FeasiblePoint
end


function joInvertibleMatrix(A::AbstractMatrix{F};DDT::DataType=F,RDT::DataType=promote_type(F,DDT)) where {F<:Number}
    (m,n) = size(A)
    m==n || throw(ArgumentError("A must be square"))
    Fact = lu(A);
    (L,U,p,q,R) = Fact.L, Fact.U, Fact.p, Fact.q, Fact.Rs
    P = joPermutation(p,DDT=DDT)
    Q = joPermutation(q,DDT=DDT)
    Ut = U'
    Lt = L'
    forw_div = v->Q'*(U\(L\(P*(R.*v))))
    adj_div = v->conj(R).*(P'*(Lt\(Ut\(Q*v))))

    return joLinearFunction_A(n,n,
                              v->A*v,
                              v->A'*v,
                              forw_div,
                              adj_div,
                              DDT,RDT,name="joInvertibleMatrix",fMVok=true,iMVok=true)
end


using Plots

x = -π:0.01:π
ys=hcat(sin.(x), cos.(x),tan.(x),sinh.(x))
plot(x,ys,layout=(4,1))


import FFNetwork
using Test

@testset "FFNetwork.jl" begin

    println();
    printstyled("Training with MNIST,\n This could take a few minutes and migth require internet!!!"; color = :green);
    println();

    FFNetwork.loadMNIST();
    traindata, testdata = FFNetwork.get_MNIST_train_and_test_data();

    println();
    println("Train set length: 50000");
    println("Mini_batch_size 20");
    println("Eta 3.0");
    println("Test set length: 10000");
    println();

    println("Running 10 epochs ");
    println("Negative is Good, Positive is Bad");
    println();

    println("Using ReLU function as act_function ");
    println();
    FFNetwork.act_function(z) = FFNetwork.ReLU(z);
    FFNetwork.act_function_prime(z) = FFNetwork.ReLU_prime(z);
    netReLU = FFNetwork.Network(28*28,16,10);
    FFNetwork.train!(netReLU, FFNetwork.data.(traindata) , FFNetwork.data.(testdata) , 10 , 20, 3.0; verbose = true);

    println("Using Sigmoid function as act_function ");
    println();
    FFNetwork.act_function(z) = FFNetwork.sigmoid(z);
    FFNetwork.act_function_prime(z) = FFNetwork.sigmoid_prime(z);
    netSigmoid = FFNetwork.Network(28*28,16,10);
    FFNetwork.train!(netSigmoid, FFNetwork.data.(traindata) , FFNetwork.data.(testdata) , 10, 20, 3.0; verbose = true);

    sigmoidacc = FFNetwork.evaluate!(netSigmoid, FFNetwork.get_MNIST_img_data());
    ReLUacc = FFNetwork.evaluate!(netReLU, FFNetwork.get_MNIST_img_data());
    println();

    println("Evaluation over whole MNIST data");
    println("Sigmoid accuracy $sigmoidacc");
    println("ReLU accuracy $ReLUacc");
    println();

    @test 1==1;

end


using Base.Test, DataFrames, MixedModels

contra = readtable(joinpath(dirname(@__FILE__), "data", "Contraception.csv.gz"))
for c in [:district, :use, :urban, :livch, :urbdist]
    contra[c] = pool(contra[c])
end
contra[:age2] = abs2(contra[:age])
gm1 = fit!(glmm(use01 ~ 1 + age + age2 + urban + livch + (1 | urbdist), contra, Bernoulli()));
@test lowerbd(gm1) == push!(fill(-Inf, 7), 0.)
@test isapprox(LaplaceDeviance(gm1), 2361.5457555; atol = 0.0001)
@test isapprox(loglikelihood(gm1), -1180.7729; atol = 0.001)
# There may be multiple optima here.
#@show logdet(gm1), sumabs2(gm1.u[1]), sum(gm1.devresid)
#@test isapprox(logdet(gm1), 75.717; atol = 0.001)
#@test isapprox(sumabs2(gm1.u[1]), 48.475; atol = 0.001)
#@test isapprox(sum(gm1.devresid), 2237.354; atol = 0.001)
#@test isapprox(fixef(gm1), [-1.0592,0.00328263,-0.00447496,0.770853,0.834578,0.913848,0.927232];
#    atol = 0.001)
show(IOBuffer(), gm1)

cbpp = readtable(joinpath(dirname(@__FILE__), "data", "cbpp.csv.gz"))
for c in [:herd, :period]
    cbpp[c] = pool(oftype(Int8[], cbpp[c]))
end
for c in [:incidence, :size]
    cbpp[c] = oftype(Int8[], cbpp[c])
end
cbpp[:prop] = cbpp[:incidence] ./ cbpp[:size]
gm2 = fit!(glmm(prop ~ 1 + period + (1 | herd), cbpp, Binomial(), LogitLink(); wt = cbpp[:size]));

@test isapprox(LaplaceDeviance(gm2), 100.095856; atol = 0.0001)
@test isapprox(sumabs2(gm2.u[1]), 9.72305; atol = 0.0001)
@test isapprox(logdet(gm2), 16.901054; atol = 0.0001)
@test isapprox(sum(gm2.devresid), 73.47143; atol = 0.001)
@test isapprox(loglikelihood(gm2), -92.02628; atol = 0.001)
@test sdest(gm2) == 1
@test varest(gm2) == 1
#VerbAgg = rcopy("lme4::VerbAgg")
#VerbAgg[:r201] = [Int8(x == "N" ? 0 : 1) for x in VerbAgg[:r2]]
#m3 = glmm(r201 ~ 1 + Anger + Gender + btype + situ + (1 | id) + (1 | item), VerbAgg, Bernoulli());
#f = r201 ~ 1 + Anger + Gender + btype + situ + (1 | id)


include("FEM1D.jl")

using LinearAlgebra, SparseArrays

function BPX_half_matrix(L::Integer, l::Integer)
    ψ = Array{Float64, 1}(undef, 2^(L-l))
    χ = Array{Float64, 1}(undef, 2^(L-l))
    ψ .= range(1, 2^(L-l), step=1)/2^(L-l)
    χ .= 1 .- ψ
    n = 2^l
    Id = spdiagm(0 => ones(n))
    S = spdiagm(-1 => ones(n-1))
    B = (kron(Id, ψ) + kron(S, χ))[1:(end-1), 1:(end-1)]
    return B
end

function Schwarz_matrix_3D(A::SparseMatrixCSC{Float64, Int64}, L::Integer)
    B = spdiagm(0 => diag(A).^-1)
    for l=1:(L-1)
        P_x = BPX_half_matrix(L, l)
        P = kron(P_x, P_x, P_x)
        if l == 1
            A_P = sparse(inv(Array(transpose(P)*A*P)))
        else
            A_P = spdiagm(0 => diag(transpose(P)*A*P).^-1)
        end
        B += P*A_P*transpose(P)
    end
    return B
end

function Laplace_3D(L::Integer)
    N = 2^L - 1
    A = spdiagm(0 => -2*ones(N), 1 => ones(N-1), -1 => ones(N-1))
    D = spdiagm(0 => ones(N))
    T = [[D, D, A], [A, D, D], [D, A, D]]
    M = sum([kron(t[1], kron(t[2], t[3])) for t = T])
    return M
end

function CG_solver!(A::SparseMatrixCSC{Float64, Int64}, b::Array{<:Real, 1}, x::Array{<:Real, 1}, tol::Real)
    r::Array{Float64, 1} = b - A*x
    p::Array{Float64, 1} = r .+ 0.0
    p_A::Array{Float64, 1} = r .+ 0.0
    β::Float64 = 1.0
    α::Float64 = 1.0
    γ_0::Float64 = norm(r, Inf)
    error::Float64 = 1.0
    γ::Float64 = γ_0
    N_sweeps::Integer = 0
    while error/γ_0 > tol^2
        γ = dot(r, r)
        p_A .= A*p
        α = γ/dot(p, p_A)
        x .+= α*p
        r .-= α*(p_A)
        β = dot(r, r)/γ
        p .= r + β*p
        error = norm(b - A*x, Inf)
        N_sweeps += 1
    end
    return N_sweeps
end

function PCG_solver!(A::SparseMatrixCSC{Float64, Int64}, B::SparseMatrixCSC{Float64, Int64}, b::Array{<:Real, 1}, x::Array{<:Real, 1}, tol::Real)
    r::Array{Float64, 1} = b - A*x
    p::Array{Float64, 1} = r .+ 0.0
    p_A::Array{Float64, 1} = r .+ 0.0
    z::Array{Float64, 1} = r .+ 0.0
    β::Float64 = 1.0
    α::Float64 = 1.0
    γ_0::Float64 = norm(r, Inf)
    error::Float64 = 1.0
    γ::Float64 = γ_0
    N_sweeps::Integer = 0
    while error/γ_0 > tol^2
        γ = dot(r, z)
        p_A .= A*p
        α = γ/dot(p, p_A)
        x .+= α*p
        r .-= α*(p_A)
        z .= B*r
        β = dot(r, z)/γ
        p .= z + β*p
        error = norm(b - A*x, Inf)
        N_sweeps += 1
    end
    return N_sweeps
end

function test_CG(L::Integer)
    A = Laplace_3D(L)
    N = 2^L - 1
    h = 1/2^L
    x = randn(N^3)
    b = randn(N^3)
    N_sweeps = CG_solver!(A, b, x, 1e-6)
    println(N_sweeps)
end

function test_PCG(L::Integer)
    A = Laplace_3D(L)
    B = Schwarz_matrix_3D(A, L)
    N = 2^L - 1
    h = 1/2^L
    x = randn(N^3)
    b = randn(N^3)
    N_sweeps = PCG_solver!(A, B, b, x, 1e-6)
    println(N_sweeps)
end

function test_sparse(L::Integer)
    A = Laplace_3D(L)
    N = 2^L - 1
    b = randn(N^3)
    F = lu(A)
    x = F \ b
    return nothing
end


using MatrixPolynomials

Jcol = colBlock([1,2,3])
Jrow = rowBlock([1,2,3])
Jinf = infBlock([1,2,3])

J2 = eigBlock(2, [1,2])
J3 = eigBlock(3, [2,2])

J = blkdiag([Jcol, Jrow, Jinf, J2, J3])


"""
module ``FloatSafeInts``
\n
Use `safemax`, `safemin` to get the extremal integer in a connected sequence
of integer values that interconvert exactly with their same sized float type.
\n
Use `isfloatsafe`, `isintsafe` to determine whether a value safely interconverts.
"""
module FloatSafeInts

import Base: trailing_ones, significand_mask

export safemax, safemin, isfloatsafe, isintsafe

safemax(::Type{Float64}) = one(Int64) << (1 + trailing_ones(significand_mask(Float64)))
safemax(::Type{Float32}) = one(Int32) << (1 + trailing_ones(significand_mask(Float32)))
safemax(::Type{Float16}) = one(Int16) << (1 + trailing_ones(significand_mask(Float16)))

safemin(::Type{Float64}) = -safemax(Float64)
safemin(::Type{Float32}) = -safemax(Float32)
safemin(::Type{Float16}) = -safemax(Float16)

safemax(::Type{Int64}) = safemax(Float64)
safemax(::Type{Int32}) = safemax(Float32)
safemax(::Type{Int16}) = safemax(Float16)
safemin(::Type{Int64}) = safemin(Float64)
safemin(::Type{Int32}) = safemin(Float32)
safemin(::Type{Int16}) = safemin(Float16)

typealias StdFloat Union{Float64,Float32,Float16}
typealias StdInt   Union{Int64,Int32,Int16}

isfloatsafe{T<:StdInt}(x::T)   = abs(x) < (safemax(T) + one(T))
isfloatsafe{T<:StdFloat}(x::T) = ((x % 1) == zero(T)) & (abs(x) < (safemax(T) + one(T)))

isintsafe{T<:StdInt}(x::T)     = abs(x) < (safemax(T) + one(T))
isintsafe{T<:StdFloat}(x::T)   = ((x % 1) == zero(T)) & (abs(x) < (safemax(T) + one(T)))

@doc """safemax(FloatType|IntType)
\n
The largest positive integer with sizeof() matching given type such that it
and its predecessor interconvert with sizeof matching float exactly.
""" safemax

@doc """safemin(FloatType|IntType)
\n
The largest negative integer with sizeof() matching given type such that it
and its predecessor interconvert with sizeof matching float exactly.
""" safemin

@doc """isfloatsafe(FloatValue|IntValue)
\n
Tests that the abs(value) and its predecessor interconvert to IntValue|FloatValue exactly.
""" isfloatsafe

@doc """isintsafe(IntValue|FloatValue)
\n
Tests that the abs(value) and its predecessor interconvert to FloatValue|IntValue exactly.
""" isintsafe

end # FloatSafeInts


module GenericFieldTest

using Test

using TowerFields.BigIntField: addmod, submod, mulmod, sqmod, invmod, powmod, frobmod, FQ2, FQ3

const Q = 7

mulroot(x) = -x

@testset "addmod: FQ" begin
    @test addmod(Q, 0, 1) == 1  # additive identity
    @test addmod(Q, 2, 6) == 1  # basic arithmetic with modulus
    @test addmod(Q, 1, -2) == Q - 1  # signed arithmetic
end

@testset "addmod: FQ2" begin
    @test addmod(Q, (0, 0), (1, 1)) == (1, 1)  # additive identity
    @test addmod(Q, (1, 2), (2, 6)) == (3, 1)  # basic arithmetic with modulus
end

@testset "addmod: FQ3" begin
    @test addmod(Q, (0, 0, 0), (1, 1, 1)) == (1, 1, 1)  # additive identity
    @test addmod(Q, (1, 2, 3), (2, 3, 4)) == (3, 5, 0)  # basic arithmetic and modulus
end

@testset "submod: FQ" begin
    @test submod(Q, 1, 0) == 1  # additive identity
    @test submod(Q, 1, 2) == Q - 1  # basic arithmetic with modulus
    @test submod(Q, 1, -2) == 3  # signed arithmetic
end

@testset "submod: FQ2" begin
    @test submod(Q, (1, 1), (0, 0)) == (1, 1)  # additive identity
    @test submod(Q, (2, 3), (3, 2)) == (Q - 1, 1)  # basic arithmetic with modulus
end

@testset "submod: FQ3" begin
    @test submod(Q, (1, 1, 1), (0, 0, 0)) == (1, 1, 1) # additive identity
    @test submod(Q, (2, 3, 4), (4, 3, 2)) == (Q - 2, 0, 2)  # basic arithmetic and modulus
end

@testset "submod: FQ3" begin
    @test submod(Q, (1, 1, 1), (0, 0, 0)) == (1, 1, 1)  # additive identity
    @test submod(Q, (2, 3, 4), (4, 3, 2)) == (Q - 2, 0, 2)  # basic arithmetic and modulus
end

@testset "submod: FQ3" begin
    @test submod(Q, (1, 1, 1), (0, 0, 0)) == (1, 1, 1) # additive identity
    @test submod(Q, (2, 3, 4), (4, 3, 2)) == (Q - 2, 0, 2)  # basic arithmetic and modulus
end

@testset "mulmod: FQ1" begin
    @test mulmod(Q, mulroot, 2, 0) == 0 # additive identity
    @test mulmod(Q, mulroot, 2, 1) == 2 # multiplicative identity
    @test mulmod(Q, mulroot, 2, 4) == 1  # basic arithmetic and modulus
    @test mulmod(Q, mulroot, 1, -1) == Q - 1  # signed arithmetic
end

@testset "mulmod: FQ2" begin
    x = (2, 3)
    @test mulmod(Q, mulroot, x, (0, 0)) == (0, 0) # additive identity
    @test mulmod(Q, mulroot, x, (1, 0)) == x # multiplicative identity
    @test mulmod(Q, mulroot, x, (3, 4)) == (1, 3)  # basic arithmetic and modulus
    @test mulmod(Q, mulroot, x, (-3, -4)) == (6, 4)  # signed arithmetic
    @test mulmod(Q, mulroot, x, 2) == addmod(Q, x, x)  # scalar arithmetic
    @test mulmod(Q, mulroot, 2, x) == addmod(Q, x, x)  # scalar arithmetic
end

@testset "mulmod: FQ3" begin
    x = (2, 3, 4)
    @test mulmod(Q, mulroot, x, (0, 0, 0)) == (0, 0, 0) # additive identity
    @test mulmod(Q, mulroot, x, (1, 0, 0)) == x # multiplicative identity
    @test mulmod(Q, mulroot, x, (3, 4, 5)) == (3, 4, 6)  # basic arithmetic and modulus
    @test mulmod(Q, mulroot, x, (-3, -4, -5)) == (4, 3, 1)  # signed arithmetic
    @test mulmod(Q, mulroot, x, 2) ==  addmod(Q, x, x)  # scalar arithmetic
    @test mulmod(Q, mulroot, 2, x) ==  addmod(Q, x, x)  # scalar arithmetic
end

@testset "sqmod: FQ1" begin
    x = 4
    @test sqmod(Q, mulroot, 0) == 0 # additive identity
    @test sqmod(Q, mulroot, 1) == 1 # multiplicative identity
    @test sqmod(Q, mulroot, x) == mulmod(Q, mulroot, x, x) # basic arithmetic and modulus
    @test sqmod(Q, mulroot, -x) == mulmod(Q, mulroot, -x, -x) # signed arithmetic
end

@testset "sqmod: FQ2" begin
    x = (2, 3)
    @test sqmod(Q, mulroot, (0, 0)) == (0, 0) # additive identity
    @test sqmod(Q, mulroot, (1, 0)) == (1, 0) # multiplicative identity
    @test sqmod(Q, mulroot, x) == mulmod(Q, mulroot, x, x) # basic arithmetic and modulus
    xneg = mulmod(Q, mulroot, -1, x)
    @test sqmod(Q, mulroot, xneg) == mulmod(Q, mulroot, xneg, xneg) # signed arithmetic
end

@testset "sqmod: FQ3" begin
    x = (5, 2, 3)
    @test sqmod(Q, mulroot, (0, 0, 0)) == (0, 0, 0) # additive identity
    @test sqmod(Q, mulroot, (1, 0, 0)) == (1, 0, 0) # multiplicative identity
    @test sqmod(Q, mulroot, x) == mulmod(Q, mulroot, x, x) # basic arithmetic and modulus
    xneg = mulmod(Q, mulroot, -1, x)
    @test sqmod(Q, mulroot, xneg) == mulmod(Q, mulroot, xneg, xneg) # signed arithmetic
end

@testset "invmod: FQ1" begin
    x = 3
    @test invmod(Q, mulroot, 1) == 1 # multiplicative identity
    y = invmod(Q, mulroot, x)
    @test mulmod(Q, mulroot, x, y) == 1
end

@testset "invmod: FQ2" begin
    x = (2, 3)
    @test invmod(Q, mulroot, (1, 0)) == (1, 0) # multiplicative identity
    y = invmod(Q, mulroot, x)
    @test mulmod(Q, mulroot, x, y) == (1, 0)
end

@testset "invmod: FQ3" begin
    x = (2, 3, 4)
    @test invmod(Q, mulroot, (1, 0, 0)) == (1, 0, 0) # multiplicative identity
    y = invmod(Q, mulroot, x)
    @test mulmod(Q, mulroot, x, y) == (1, 0, 0)
end

@testset "powmod: FQ1" begin
    x, e = 3, 3
    @test powmod(Q, mulroot, 2, 0) == 1  # special exponent 0
    @test powmod(Q, mulroot, 2, 1) == 2  # special exponent 1
    @test powmod(Q, mulroot, x, 2) == sqmod(Q, mulroot, x)  # basic arithmetic and modulus
    @test powmod(Q, mulroot, x, Q - 1) == 1  # fermat's theorem
    y = mulmod(Q, mulroot, x, sqmod(Q, mulroot, x))
    @test powmod(Q, mulroot, x, e) == y  # basic arithmetic and modulus
    @test powmod(Q, mulroot, x, -e) == invmod(Q, mulroot, y)  # negative exponent
end

@testset "powmod: FQ2" begin
    x, e = (2, 3), 3
    @test powmod(Q, mulroot, x, 0) == (1, 0)  # special exponent 0
    @test powmod(Q, mulroot, x, 1) == x  # special exponent 1
    @test powmod(Q, mulroot, x, 2) == sqmod(Q, mulroot, x)  # basic arithmetic and modulus
    y = mulmod(Q, mulroot, x, sqmod(Q, mulroot, x))
    @test powmod(Q, mulroot, x, e) == y  # basic arithmetic and modulus
    @test powmod(Q, mulroot, x, -e) == invmod(Q, mulroot, y)  # negative exponent
end

@testset "powmod: FQ3" begin
    x, e = (2, 3, 4), 3
    @test powmod(Q, mulroot, x, 0) == (1, 0, 0)  # special exponent 0
    @test powmod(Q, mulroot, x, 1) == x  # special exponent 1
    @test powmod(Q, mulroot, x, 2) == sqmod(Q, mulroot, x)  # basic arithmetic and modulus
    y = mulmod(Q, mulroot, x, sqmod(Q, mulroot, x))
    @test powmod(Q, mulroot, x, e) == y  # basic arithmetic and modulus
    @test powmod(Q, mulroot, x, -e) == invmod(Q, mulroot, y)  # negative exponent
end

end


using Plots
plotly()

using simulator.fresnelltools
using simulator.materials.spesifics
using simulator.matrixcore

θs = 0:1e-3:π/2
θ = 0.806
λs = 350e-9:1e-9:1800e-9
λ = 1434e-9

d_ags = 10e-9:1e-9:200e-9
d_ag = 50e-9
d_sis = 10e-9:1e-9:200e-9
d_si = 20e-9
d_pds = 10e-9:1e-9:200e-9
d_pd = 5e-9

n_fiber = SiO2_core_Sellmeier
n_ag = Ag
n_silicafilm = SiO2_thinfilm_Ciprian
n_palladium(λ) = Pd_Werner(λ) - 0im
n_air = 1

function structure(λ, θ, d_ag, d_si, d_pd)
    interface_fiber_ag_te, interface_fiber_ag_tm, θag = FresnellBoundrary(n_fiber(λ), n_ag(λ), θ)
    bulk_ag = FresnellSlab(n_ag(λ), 2*π/λ, d_ag, θag)
    interface_ag_si_te, interface_ag_si_tm, θsi = FresnellBoundrary(n_ag(λ), n_silicafilm(λ), θag)
    bulk_si = FresnellSlab(n_silicafilm(λ), 2*π/λ, d_si, θsi)
    interface_si_pd_te, interface_si_pd_tm, θpd = FresnellBoundrary(n_silicafilm(λ), n_palladium(λ), θsi)
    bulk_pd = FresnellSlab(n_palladium(λ), 2*π/λ, d_pd, θpd)
    interface_pd_air_te, interface_pd_air_tm, _ = FresnellBoundrary(n_palladium(λ), n_air, θpd)

    te_system = CascadeScattering([interface_fiber_ag_te, bulk_ag, interface_ag_si_te, bulk_si, interface_si_pd_te, bulk_pd, interface_pd_air_te])
    tm_system = CascadeScattering([interface_fiber_ag_tm, bulk_ag, interface_ag_si_tm, bulk_si, interface_si_pd_tm, bulk_pd, interface_pd_air_tm])

    te_system, tm_system
end

function structure_inside(λ, θ, d_ag, d_si, d_pd)
    bulk_fiber(d) = FresnellSlab(n_fiber(λ), 2*π/λ, d, θ)
    interface_fiber_ag_te, interface_fiber_ag_tm, θ_ag = FresnellBoundrary(n_fiber(λ), n_ag(λ), θ)
    interface_fiber_ag_tm_f(d) = interface_fiber_ag_tm
    interface_fiber_ag_te_f(d) = interface_fiber_ag_te
    bulk_ag(d) = FresnellSlab(n_ag(λ), 2*π/λ, d, θ_ag)
    interface_ag_si_te, interface_ag_si_tm, θ_si = FresnellBoundrary(n_ag(λ), n_silicafilm(λ), θ_ag)
    interface_ag_si_tm_f(d) = interface_ag_si_tm
    interface_ag_si_te_f(d) = interface_ag_si_te
    bulk_si(d) = FresnellSlab(n_silicafilm(λ), 2*π/λ, d, θ_si)

    # interface_si_pd_te, interface_si_pd_tm, θ_pd = FresnellBoundrary(n_silicafilm(λ), n_palladium(λ), θ_si)
    # interface_si_pd_te_f(d) = interface_si_pd_te
    # interface_si_pd_tm_f(d) = interface_si_pd_tm
    # bulk_pd(d) = FresnellSlab(n_palladium(λ), 2*π/λ, d, θ_pd)
    # interface_pd_air_te, interface_pd_air_tm, θ_air = FresnellBoundrary(n_palladium(λ), n_air, θ_air)
    # interface_pd_air_te_f(d) = interface_pd_air_te
    # interface_pd_air_tm_f(d) = interface_pd_air_tm
    # bulk_air(d) = FresnellSlab(n_air, 2*π/λ, d, θ_air)

    interface_si_air_te, interface_si_air_tm, θ_air = FresnellBoundrary(n_silicafilm(λ), n_air, θ_si)
    interface_si_air_te_f(d) = interface_si_air_te
    interface_si_air_tm_f(d) = interface_si_air_tm
    bulk_air(d) = FresnellSlab(n_air, 2*π/λ, d, θ_air)

    interrogated = (Up, Un) -> abs(Up + Un)^2

    u_values_tm, d_values_tm = Interrogator(
        [bulk_fiber, interface_fiber_ag_tm_f, bulk_ag, interface_ag_si_tm_f, bulk_si, interface_si_air_tm_f, bulk_air],
        [10e-9, 0, d_ag, 0, d_si, 0, 10e-9], 1e-10,
        interrogated)
    u_values_te, d_values_te = Interrogator(
        [bulk_fiber, interface_fiber_ag_te_f, bulk_ag, interface_ag_si_te_f, bulk_si, interface_si_air_te_f, bulk_air],
        [10e-9, 0, d_ag, 0, d_si, 0, 10e-9], 1e-10,
        interrogated)

    u_values_te, d_values_te, u_values_tm, d_values_tm
end

R_te(λ, θ, d_ag, d_si, d_pd) = abs((structure(λ, θ, d_ag, d_si, d_pd)[1] * [1, 0])[2])^2
R_tm(λ, θ, d_ag, d_si, d_pd) = abs((structure(λ, θ, d_ag, d_si, d_pd)[2] * [1, 0])[2])^2
R_ratio(λ, θ, d_ag, d_si, d_pd) = R_te(λ, θ, d_ag, d_si, d_pd) / R_tm(λ, θ, d_ag, d_si, d_pd)

interrogated_u_te, interrogated_d_te, interrogated_u_tm, interrogated_d_tm = structure_inside(λ, θ, d_ag, d_si, d_pd)

# plot(λs, R_ratio.(λs, θ, d_ag, d_si, d_pd), label="θ=0.856", title="TE Reflection / TM Reflection at d=53nm")
# plot(θs, (θ) -> (R_ratio.(λ, θ + 1e-9im, d_ag, d_si, d_pd)))
# plot(d_ags, R_ratio.(λ, θ, d_ags, d_si, d_pd))
# plot(d_sis, R_ratio.(λ, θ, d_ag, d_sis, d_pd))
# plot(d_pds, R_ratio.(λ, θ, d_ag, d_si, d_pds))

# plot(λs, R_te.(λs, θ, d_ag, d_si, d_pd), label="te, Pd-structure, Different Pd", title="Reflection at silica-palladium-air", legend=:outertopright)
# plot!(λs, R_tm.(λs, θ, d_ag, d_si, d_pd), label="tm, Pd-structure, Different Pd", xlabel="wavelength", ylabel="reflected power")

# plot(interrogated_d_te .* 1e9, interrogated_u_te, label="te, λ=$λ", title="Sum of fields at plasmone conditions", legend=:topright, xlabel="distance in structure", ylabel="total power")
# plot!(interrogated_d_tm .* 1e9, interrogated_u_tm, label="tm λ=$λ")


#### Frank wolfe for matrix completion
using FirstOrderOptimization
using PyPlot
import Base: axpy!, scale!

## Generate data

# make it deterministic
srand(10)

m, n, r, rtrue = 100, 100, 10, 2
X = randn(m, rtrue) * randn(rtrue, n)
smooth_inf = log(min(m,n))
# smooth_inf and smooth_1 are conjugate
# 1/p + 1/q = 1 => p = 1/(1-1/q) = q / (q - 1)
smooth_1 = smooth_inf / (smooth_inf - 1)
u,s,v = svd(X)
tau = sum(s)
tau_smooth = sum(s.^(smooth_1))^(1/smooth_1)

pobs = .5
obs = sprand(m, n, pobs).>0
iobs = find(obs)
A = IndexingOperator(m, n, iobs)
b = A*X

srand(10)

objective_sketched(z) = .5*sum((z - b).^2)
grad_objective_sketched(z) = LowRankOperator(A, z-b, transpose = Symbol[:T, :N]) # concretely, should be A'*(z-b)

const_nucnorm(z) = tau # we'll always saturate the constraint, don't bother computing it
function min_lin_st_nucnorm_sketched(g, tau)
	u,s,v = svds(Array(g), nsv=1) # for this case, g is a sparse matrix so representing it is O(m)
	return LowRankOperator(-tau*u, v')
end
# I can't think of a pretty way to compute
# <g, tilde_x - x> = <A'(z-b), u tau v' - x> = <z-b, A(u tau v') - A(x)> = <z-b, A(u tau v') - z>
# dot_g_w_tx_minus_x(g, tilde_x, x) = dot(g.factors[2], g.factors[1]*tilde_x - x)
dot_g_w_tx_minus_x(g, Delta, z) = dot(vec(g.factors[2]), g.factors[1]*Delta - z)
# we end up computing A*Delta twice with this scheme
update_var!(z, Delta, a) = (scale!(z, 1-a); axpy!(a, A*Delta, z))

z = zeros(length(b))
ch = ConvergenceHistory("fw_sketched")
X_sketched = frank_wolfe_sketched(
	z,
	objective_sketched, grad_objective_sketched,
	const_nucnorm,
	tau,
	min_lin_st_nucnorm_sketched,
	AsymmetricSketch(m,n,r),
	FrankWolfeParams(100, 1e-2, DecreasingStepSize(2,1)),
	ch,
	dot_g_w_tx_minus_x = dot_g_w_tx_minus_x,
	update_var! = update_var!,
	#LB = 0,
	verbose=true
	)

figure()
# yscale("log")
plot(ch.objective, label="fw")
# plot(ch.objective, label="fw")
xlabel("Iterations")
ylabel("Objective")

## now smooth version

for (nsv, nsv_label) = [(r, "r"), (Int(ceil(smooth_inf)), "log(n)"), (min(m,n), "n")]

	function min_lin_st_nucnorm_sketched(g, tau,
			smooth_inf = smooth_inf,
			nsv = nsv)
		u,s,v = svds(Array(g), nsv=nsv) # for this case, g is a sparse matrix so representing it is O(m)
		s = s.^smooth_inf
		s = s / sum(s.^smooth_1).^(1/smooth_1) * -tau # normalize so the smooth_1 norm of the solution is tau
		return LowRankOperator(u, spdiagm(s), v')
	end
	const_nucnorm(z) = tau_smooth # we'll always saturate the constraint, don't bother computing it

	srand(10)

	ch_smooth = ConvergenceHistory("smooth_fw_sketched")
	z = zeros(length(b))
	X_sketched_smooth = frank_wolfe_sketched(
		z,
		objective_sketched, grad_objective_sketched,
		const_nucnorm,
		tau_smooth,
		min_lin_st_nucnorm_sketched,
		AsymmetricSketch(m,n,r),
		FrankWolfeParams(100, 1e-2, DecreasingStepSize(2,1)),
		ch_smooth,
		dot_g_w_tx_minus_x = dot_g_w_tx_minus_x,
		update_var! = update_var!,
		#LB = 0,
		verbose=true
		)

	plot(ch_smooth.objective, label="smoothed fw top "*nsv_label)
	#@show(vecnorm(Array(X_sketched_smooth) - X))
end
legend()


# Hinge loss function

export HingeLoss

"""
**Hinge loss**

  HingeLoss(y, μ=1.0)

Returns the function
```math
f(x) = μ⋅∑_i \\max\\{0, 1 - y_i ⋅ x_i\\},
```
where `y` is an array and `μ` is a positive parameter.
"""
HingeLoss(y::T, mu::R=1.0) where {T <: AbstractArray, R <: Real} = Postcompose(PrecomposeDiagonal(SumPositive(), -y, 1.0), mu)


using Test
using PolynomialBasis
using ImplicitDomainQuadrature
# using Revise
using CutCell
include("useful_routines.jl")

polyorder = 1
numqp = 2

basis = TensorProductBasis(2, polyorder)
levelset = InterpolatingPolynomial(1, basis)
nf = CutCell.number_of_basis_functions(basis)
mesh = CutCell.Mesh([0.0, 0.0], [3.0, 1.0], [3, 1], nf)
nodalcoordinates = CutCell.nodal_coordinates(mesh)
nodalconnectivity = CutCell.nodal_connectivity(mesh)

normal = [1.0, 0.0]
x0 = [0.5, 0.0]
tol = perturbation = 1e-3
levelsetcoeffs = plane_distance_function(nodalcoordinates, normal, x0)

cellsign = CutCell.cell_sign!(levelset, levelsetcoeffs, nodalconnectivity, tol, perturbation)
cellmap = CutCell.cell_map(mesh,1)

interfacequads = CutCell.InterfaceQuadratures(
    cellsign,
    levelset,
    levelsetcoeffs,
    nodalconnectivity,
    numqp,
    cellmap,
)
@test size(interfacequads.quads) == (2,1)
testcelltoquad = [1, 0, 0]
@test allequal(interfacequads.celltoquad, testcelltoquad)

testnormals = repeat(normal,inner=(1,numqp))
@test allapprox(testnormals,CutCell.interface_normals(interfacequads,1))

@test interfacequads[1,1] ≈ interfacequads[-1,1]


const PWGRID_VERSION = 4

mutable struct GkVectors
    kpts::Array{Float64,2}
    Ngwx::Int
    Ngw::Array{Int,1}
    idx_gk::Array{Array{Int64,1},1}
end

mutable struct GVectors
    Ng::Int
    G::Array{Float64,2}
    G2::Array{Float64}
end

mutable struct PWGrid
    ecutwfc::Float64
    ecutrho::Float64
    Ns::Array{Int64}
    LatVecs::Array{Float64,2}
    RecVecs::Array{Float64,2}
    Ω::Float64
    r::Array{Float64,2}
    gvec::GVectors
    gkvec::GkVectors
end

function PWGrid( ecutwfc::Float64, LatVecs::Array{Float64,2}, kpts_red )
    ecutrho = 4.0*ecutwfc
    #
    println("LatVecs = ", LatVecs)
    RecVecs = 2*pi*inv(LatVecs)'
    println("RecVecs = ", RecVecs)
    Ω = det(LatVecs)
    #
    LatVecsLen = Array{Float64}(undef,3)
    LatVecsLen[1] = norm(LatVecs[1,:])
    LatVecsLen[2] = norm(LatVecs[2,:])
    LatVecsLen[3] = norm(LatVecs[3,:])

    Ns = Array{Int64}(undef,3)
    Ns[1] = 2*round( Int64, sqrt(ecutrho/2)*LatVecsLen[1]/pi ) + 1
    Ns[2] = 2*round( Int64, sqrt(ecutrho/2)*LatVecsLen[2]/pi ) + 1
    Ns[3] = 2*round( Int64, sqrt(ecutrho/2)*LatVecsLen[3]/pi ) + 1

    # Use even sampling numbers
    Ns[1] = Ns[1] % 2 == 1 ? Ns[1] + 1 : Ns[1]
    Ns[2] = Ns[2] % 2 == 1 ? Ns[2] + 1 : Ns[2]
    Ns[3] = Ns[3] % 2 == 1 ? Ns[3] + 1 : Ns[3]

    Npoints = prod(Ns)
    r = init_grid_R( Ns, LatVecs )

    gvec = init_grid_G( Ns, RecVecs )

    #kpts = (kpts_red' * RecVecs)'

    kpts = RecVecs*kpts_red # new following convention of PWDFT.jl

    #kpts = (kpts_red' * RecVecs')'

    #kpts = RecVecs' * kpts_red
    #kpts = RecVecs * kpts_red

    #Nkpts = size(kpts_red)[2]
    #kpts = zeros(Float64, 3, Nkpts)
    #for ik = 1:Nkpts
    #    kpts[:,ik] = RecVecs' * kpts_red[:,ik]
    #end
    #for ik = 1:Nkpts
    #    for i = 1:3
    #        kpts[i,ik] = kpts_red[i,ik] * RecVecs[1,:]
    #    end
    #end
    gkvec = init_gkvec( ecutwfc, gvec.G, kpts )

    return PWGrid( ecutwfc, ecutrho, Ns, LatVecs, RecVecs, Ω, r, gvec, gkvec )
end

function mm_to_nn(mm::Int,S::Int)
    if mm > S/2
        return mm - S
    else
        return mm
    end
end

function init_grid_G( Ns, RecVecs )

    Ng = prod(Ns)

    G  = Array{Float64}(undef,3,Ng)
    G2 = Array{Float64}(undef,Ng)

    ig = 0
    for k in 0:Ns[3]-1
    for j in 0:Ns[2]-1
    for i in 0:Ns[1]-1
        ig = ig + 1
        gi = mm_to_nn( i, Ns[1] )
        gj = mm_to_nn( j, Ns[2] )
        gk = mm_to_nn( k, Ns[3] )
        G[1,ig] = RecVecs[1,1]*gi + RecVecs[1,2]*gj + RecVecs[1,3]*gk
        G[2,ig] = RecVecs[2,1]*gi + RecVecs[2,2]*gj + RecVecs[2,3]*gk
        G[3,ig] = RecVecs[3,1]*gi + RecVecs[3,2]*gj + RecVecs[3,3]*gk
        G2[ig] = G[1,ig]^2 + G[2,ig]^2 + G[3,ig]^2
    end
    end
    end

    return GVectors( Ng, G, G2 )
end

function init_gkvec( ecutwfc, G, kpts )
    Nkpts = size(kpts)[2]
    idx_gk = Array{Array{Int64,1},1}(undef,Nkpts)
    Ng = size(G)[2]
    Gk = zeros( 3, Ng )
    Gk2 = zeros( Ng )
    Ngw = zeros( Int64, Nkpts )
    for ik = 1:Nkpts
        #@printf("\n")
        for ig = 1:Ng
            Gk[:,ig] = G[:,ig] + kpts[:,ik]
            Gk2[ig]  = Gk[1,ig]^2 + Gk[2,ig]^2 + Gk[3,ig]^2
            #@printf("%8d %18.10f %18.10f\n", ig, Gk2[ig], ecutwfc)
        end
        idx_gk[ik] = findall( 0.5*Gk2 .< ecutwfc )
        Ngw[ik] = length(idx_gk[ik])
        #@printf("kpt = %d, Ngw = %d\n", ik, Ngw[ik])
        #@printf("kpts[%d]: %f %f %f\n", ik, kpts[1,ik], kpts[3,ik], kpts[3,ik] )
    end
    Ngwx = maximum(Ngw)
    return GkVectors( kpts, Ngwx, Ngw, idx_gk )
end

function init_grid_R( Ns, LatVecs )
    #
    Npoints = prod(Ns)
    #
    R = Array{Float64}(undef,3,Npoints)
    ip = 0
    for k in 0:Ns[3]-1
    for j in 0:Ns[2]-1
    for i in 0:Ns[1]-1
        ip = ip + 1
        R[1,ip] = LatVecs[1,1]*i/Ns[1] + LatVecs[1,2]*j/Ns[2] + LatVecs[1,3]*k/Ns[3]
        R[2,ip] = LatVecs[2,1]*i/Ns[1] + LatVecs[2,2]*j/Ns[2] + LatVecs[2,3]*k/Ns[3]
        R[3,ip] = LatVecs[3,1]*i/Ns[1] + LatVecs[3,2]*j/Ns[2] + LatVecs[3,3]*k/Ns[3]
    end
    end
    end
    #
    return R
end


# This file generates the plot in the logo
using Ripserer
using Plots
gr()
using Random
Random.seed!(1337)

function circle(n; r=0.7, center, s=0)
    points = NTuple{2,Float64}[]
    x, y = center
    for θ in range(0, 2π; length=n + 1)[circshift(1:n, s)]
        push!(points, (x + r * sin(θ), y + r * cos(θ)))
    end
    return points
end

data = [
    circle(20; center=(-1, 0), s=0)
    circle(20; center=(1, 0), s=0)
    circle(20; center=(0, √3), s=0)
]

res = ripserer(data; reps=true, dim_max=2)

ripserer_logo = plot(;
    aspect_ratio=1,
    xlim=(-2.1, 2.1),
    ylim=(-1.1, 3.1),
    xlabel=:x,
    ylabel=:y,
    ticks=nothing,
    legend=false,
    border=:none,
)
for (i, int) in enumerate(sort(res[2]; by=death, rev=true))
    if i == 1
        col = 2 # red
        α = 1
    elseif i == 2
        col = 4 # purple
        α = 1
    elseif i == 3
        col = 3 # green
        α = 1
    else
        col = 1
        α = 0.2
    end

    plot!(
        ripserer_logo,
        int,
        data;
        alpha=α,
        color=col,
        threshold_strict=death(int),
        linewidth=5,
    )
end
scatter!(ripserer_logo, data; markersize=4, color=1)
display(ripserer_logo)

diagrams_logo = plot(
    res;
    title="",
    ticks=nothing,
    legend=false,
    xlabel="",
    ylabel="",
    markersize=3,
    size=(100, 100),
)


using RecursiveArrayTools, DataFrames, GLM
struct Regression_Sensitivity_Coefficients
    Pearson
    Standard_Regression
    Partial_Correlation
    Spearman
    Standard_Rank_Regression
    Partial_Rank_Correlation
end

function regression_sensitivity(f,p_range,p_fixed,n;coeffs=:rank)

    pearson_coeffs = []
    src_coeffs = []
    partial_coeffs = []
    spearman_coeffs = []
    srr_coeffs = []
    partial_rank_coeffs = []

    for i in 1:length(p_range)
        pcs = []
        srcs = []
        params = []
        for j in 1:length(p_range)
            if i != j
                push!(params,p_fixed[j])
            else
                push!(params,0)
            end
        end
        xis = []
        yis = Array{Float64}[]
        for j in 1:n
            params[i] = (p_range[i][2]-p_range[i][1])*rand() + p_range[i][1]
            x = params
            y = f(params)
            push!(xis,copy(x))
            push!(yis,copy(y))
        end
        yis = VectorOfArray(yis)
        x_mean = mean(xis)
        x_vrs = [sum(xis[i] - x_mean) for o in 1:length(xis)]
        x_lm = [j[i] for j in xis]
        x_rnk = sortperm(x_lm)
        x_rnk_mean = mean(x_rnk)
        x_rnk_vrs = [x_rnk[i] - x_rnk_mean for o in 1:length(x_rnk)]
        x_var = var(x_lm)
        x_rnk_var = var(x_rnk)
        pcs,srcs = pcs_and_srcs(yis,x_lm,x_vrs,x_var)
        pcc = pcc_f(x_lm,yis)
        yis_rnk = zero(yis)
        for k in 1:size(yis)[1]
            for j in 1:size(yis)[2]
                yis_rnk[k,j,:] = sortperm(yis[k,j,:])
            end
        end
        if coeffs == :rank
            prcc = pcc_f(x_rnk,yis_rnk)
            spear_rcc, srrc = pcs_and_srcs(yis_rnk,x_rnk,x_rnk_vrs,x_rnk_var)
        end
        push!(pearson_coeffs,pcs)
        push!(src_coeffs,srcs)
        push!(partial_coeffs,pcc)
        if coeffs == :rank
            push!(partial_rank_coeffs,prcc)
            push!(spearman_coeffs,spear_rcc)
            push!(srr_coeffs,srrc)
        end
    end
    if coeffs == :rank
        regre_coeff = Regression_Sensitivity_Coefficients(pearson_coeffs,src_coeffs,partial_coeffs,spearman_coeffs,partial_rank_coeffs,srr_coeffs)
    else
        regre_coeff = Regression_Sensitivity_Coefficients(pearson_coeffs,src_coeffs,partial_coeffs,nothing,nothing,nothing)
    end
    regre_coeff
end

function pcs_and_srcs(yis,x_lm,x_vrs,x_var)
    pcs = []
    srcs = []
    for k in 1:size(yis)[1]
        pc = []
        src = []
        for j in 2:size(yis)[2]
            pear_coeff_num = sum(x_vrs .* (yis[k,j,:] - mean(yis[k,j,:])))
            pear_coeff_deno = sqrt(sum(x_vrs.^2)) * sqrt(sum((yis[k,j,:] - mean(yis[k,j,:])).^2))
            pear_coeff = pear_coeff_num/pear_coeff_deno
            push!(pc,pear_coeff)
            df = DataFrame(X=x_lm,Y=yis[k,j,:])
            lin_model = lm(@formula(Y ~ X),df)
            push!(src,coef(lin_model)[2].*sqrt.(x_var ./ var(yis[k,j,:])))
        end
        push!(pcs,pc)
        push!(srcs,src)
    end
    pcs,srcs
end

function pcc_f(x_lm,yis)
    df_arr = Array{Float64}[Float64[] for i in 1:length(x_lm)]
    for o in 1:length(x_lm)
        for j in 1:length(x_lm)
            push!(df_arr[o],x_lm[(j+o)%length(x_lm)+1])
        end
    end
    df = DataFrame(df_arr)
    last_ind = length(x_lm)
    formula = "@formula(x$last_ind ~ "
    for j in 1:length(x_lm)-1
        formula *= "x$j +"
    end
    formula = formula[1:end-1] * ")"
    ols_x = lm(eval(Meta.parse(formula)),df)
    x_cap = x_lm .- predict(ols_x)
    y_caps = []
    for k in 1:size(yis)[1]
        y_cps = []
        for j in 2:size(yis)[2]
            df_arr = Array{Float64}[Float64[] for i in 1:length(yis[k,j,:])]
            for o in 1:length(yis[k,j,:])
                for l in 1:length(yis[k,j,:])
                    push!(df_arr[o],yis[k,j,:][(l+o)%length(yis[k,j,:])+1])
                end
            end
            df = DataFrame(df_arr)
            last_ind = length(yis[k,j,:])
            formula = "@formula(x$last_ind ~ "
            for j in 1:length(yis[k,j,:])-1
                formula *= "x$j +"
            end
            formula = formula[1:end-1] * ")"
            ols_y = lm(eval(Meta.parse(formula)),df)
            y_cap = yis[k,j,:] .- predict(ols_y)
            push!(y_cps,y_cap)
        end
        push!(y_caps,y_cps)
    end
    pcc = []
    for k in 1:length(y_caps)
        pcc_ = []
        for j in 1:length(y_caps[k])
            num = sum(x_cap .* y_caps[k][j])
            deno = sqrt(sum(x_cap.^2)*sum(y_caps[k][j].^2))
            push!(pcc_,num/deno)
        end
        push!(pcc,pcc_)
    end
    pcc
end

function regression_sensitivity(prob::DiffEqBase.DEProblem,alg,t,p_range,p_fixed,n;coeffs=:rank)
    f = function (p)
        prob1 = remake(prob;p=p)
        Array(solve(prob1,alg;saveat=t))
    end
    regression_sensitivity(f,p_range,p_fixed,n,coeffs=coeffs)
end


mutable struct ILU0Preconditioner{Tv, Ti} <: AbstractPreconditioner{Tv,Ti}
    A::ExtendableSparseMatrix{Tv,Ti}
    xdiag::Array{Tv,1}
    idiag::Array{Ti,1}
    phash::UInt64
    function ILU0Preconditioner{Tv,Ti}() where {Tv,Ti}
        p=new()
        p.phash=0
        p
    end
end

"""
```
ILU0Preconditioner(;valuetype=Float64,indextype=Int64)
ILU0Preconditioner(matrix)
```

ILU preconditioner with zero fill-in.
"""
ILU0Preconditioner(;valuetype::Type=Float64, indextype::Type=Int64)=ILU0Preconditioner{valuetype,indextype}()

function update!(precon::ILU0Preconditioner{Tv,Ti}) where {Tv,Ti}
    flush!(precon.A)
    cscmatrix=precon.A.cscmatrix
    colptr=cscmatrix.colptr
    rowval=cscmatrix.rowval
    nzval=cscmatrix.nzval
    n=cscmatrix.n

    if precon.phash==0
        n=size(precon.A,1)
        precon.xdiag=Array{Tv,1}(undef,n)
        precon.idiag=Array{Ti,1}(undef,n)
    end

    xdiag=precon.xdiag
    idiag=precon.idiag

    # Find main diagonal index and
    # copy main diagonal values
    if precon.phash != precon.A.phash
        @inbounds for j=1:n
            @inbounds for k=colptr[j]:colptr[j+1]-1
                i=rowval[k]
                if i==j
                    idiag[j]=k
                    break
                end
            end
        end
        precon.phash=precon.A.phash
    end

    @inbounds for j=1:n
        xdiag[j]=one(Tv)/nzval[idiag[j]]
        @inbounds for k=idiag[j]+1:colptr[j+1]-1
            i=rowval[k]
            for l=colptr[i]:colptr[i+1]-1
                if rowval[l]==j
                    xdiag[i]-=nzval[l]*xdiag[j]*nzval[k]
                    break
                end
            end
        end
    end
    precon
end

function  LinearAlgebra.ldiv!(u::AbstractArray{T,1}, precon::ILU0Preconditioner, v::AbstractArray{T,1}) where T
    cscmatrix=precon.A.cscmatrix
    colptr=cscmatrix.colptr
    rowval=cscmatrix.rowval
    n=cscmatrix.n
    nzval=cscmatrix.nzval
    xdiag=precon.xdiag
    idiag=precon.idiag

    @inbounds for j=1:n
        x=zero(T)
        @inbounds for k=colptr[j]:idiag[j]-1
            x+=nzval[k]*u[rowval[k]]
        end
        u[j]=xdiag[j]*(v[j]-x)
    end

    @inbounds for j=n:-1:1
        x=zero(T)
        @inbounds for k=idiag[j]+1:colptr[j+1]-1
            x+=u[rowval[k]]*nzval[k]
        end
        u[j]-=x*xdiag[j]
    end
end

function LinearAlgebra.ldiv!(precon::ILU0Preconditioner, v::AbstractArray{T,1} where T)
    ldiv!(v, precon, v)
end


export RCAMspg

function RCAMspg(b, alt, r, eta, Lprev, Rprev, L::Tx, weight, iter)where {ETx<:Number, Tx<:AbstractMatrix{ETx}}

    m,n = size(b);

     #change for different rows and cols 2021.01.28
    #indi = 1:1:(maximum[m,n]);
    indin = 1:1:n;
    indim = 1:1:m;

    R = zeros(Complex{Float32},size(Rprev));

        Wr = 1.;
        Ql = 1.;

    for k = 1:alt
       time1 = @elapsed R = minFrobNormspg(L,b,1, eta,indin,Lprev,weight, iter, Wr);
       println("Time1 = ", string(time1))
       time2 = @elapsed L = minFrobNormspg(R,b',1, eta,indim,Rprev,weight, iter, Ql);
       println("Time1 = ", string(time2))
    end
    return L,R
end


@inline function metropolis_acc(kT::Real, dU::Real, sθa::Real, sθb::Real, ϵ::Real)
  return ϵ <= (exp(-dU / kT) * sθb / sθa);
end

function kawasaki_acc(kT::Real, dU::Real, ϵ::Real)
  boltz = exp(-dU / (2*kT));
  anti_boltz = exp(dU / (2*kT));
  return ( ϵ <= (boltz / (boltz + anti_boltz)) );
end

abstract type Acceptor end;

mutable struct Metropolis <: Acceptor
  logπ_prev::Float64;
  weight_function::WeightFunction;
end

logπ_chain(chain::EAPChain) = -chain.U / chain.kT + chain.Ω;

function logπ_chain(chain::EAPChain, log_wf::WeightFunction)
  return -chain.U / chain.kT + chain.Ω + log_wf(chain);
end

function Metropolis(chain::EAPChain, wf::WeightFunction)
  return Metropolis(logπ_chain(chain, wf), wf);
end

# alpha is a ratio of trial move probabilities (related to clustering)
function (metro::Metropolis)(chain::EAPChain, ϵ::Real; α::Real = 1.0)
  metro.weight_function(chain);
  logπ_curr = logπ_chain(chain, metro.weight_function) + log(α);
  if (logπ_curr >= metro.logπ_prev) || (ϵ < exp(logπ_curr - metro.logπ_prev))
    metro.logπ_prev = logπ_curr;
    return true;
  else
    return false;
  end
  return false;
end


function [X,J,dXdx]=JacobiCart(ksi,eta,F,Grid)
n1=size(ksi,1);
n2=size(eta,2);
X=zeros(n1,n2,2);
% P1=Grid.Nodes(F.N(1)).P;
% P2=Grid.Nodes(F.N(2)).P;
% P3=Grid.Nodes(F.N(3)).P;
% P4=Grid.Nodes(F.N(4)).P,
% X1(:,:,1)=P1*(1-ksi)*(1-eta);
X(:,:,1)=0.25*(Grid.Nodes(F.N(1)).P(1).*(1-ksi)*(1-eta)+...
  Grid.Nodes(F.N(2)).P(1).*(1-ksi)*(1+eta)+...
  Grid.Nodes(F.N(3)).P(1).*(1+ksi)*(1+eta)+...
  Grid.Nodes(F.N(4)).P(1).*(1+ksi)*(1-eta));
X(:,:,2)=0.25*(Grid.Nodes(F.N(1)).P(2).*(1-ksi)*(1-eta)+...
  Grid.Nodes(F.N(2)).P(2).*(1-ksi)*(1+eta)+...
  Grid.Nodes(F.N(3)).P(2).*(1+ksi)*(1+eta)+...
  Grid.Nodes(F.N(4)).P(2).*(1+ksi)*(1-eta));

if nargout>2
  dXdx=zeros(n1,n2,2,2);
  dXdx(:,:,1,1)=0.25*(-Grid.Nodes(F.N(1)).P(1)-Grid.Nodes(F.N(2)).P(1)+...
    Grid.Nodes(F.N(3)).P(1)+Grid.Nodes(F.N(4)).P(1)+...
    ( Grid.Nodes(F.N(1)).P(1)-Grid.Nodes(F.N(2)).P(1)+...
    Grid.Nodes(F.N(3)).P(1)-Grid.Nodes(F.N(4)).P(1))*(ones(n1,1)*eta));
  dXdx(:,:,1,2)=0.25*(-Grid.Nodes(F.N(1)).P(1)+Grid.Nodes(F.N(2)).P(1)+...
    Grid.Nodes(F.N(3)).P(1)-Grid.Nodes(F.N(4)).P(1)+...
    ( Grid.Nodes(F.N(1)).P(1)-Grid.Nodes(F.N(2)).P(1)+...
    Grid.Nodes(F.N(3)).P(1)-Grid.Nodes(F.N(4)).P(1))*(ksi*ones(1,n2)));
  dXdx(:,:,2,1)=0.25*(-Grid.Nodes(F.N(1)).P(2)-Grid.Nodes(F.N(2)).P(2)+...
    Grid.Nodes(F.N(3)).P(2)+Grid.Nodes(F.N(4)).P(2)+...
    ( Grid.Nodes(F.N(1)).P(2)-Grid.Nodes(F.N(2)).P(2)+...
    Grid.Nodes(F.N(3)).P(2)-Grid.Nodes(F.N(4)).P(2))*(ones(n1,1)*eta));
  dXdx(:,:,2,2)=0.25*(Grid.Nodes(F.N(1)).P(2)-Grid.Nodes(F.N(2)).P(2)+...
    Grid.Nodes(F.N(3)).P(2)-Grid.Nodes(F.N(4)).P(2)+...
    ( Grid.Nodes(F.N(1)).P(2)-Grid.Nodes(F.N(2)).P(2)+...
    Grid.Nodes(F.N(3)).P(2)-Grid.Nodes(F.N(4)).P(2))*(ksi*ones(1,n2)));
  J=dXdx(:,:,1,1).*dXdx(:,:,2,2)-dXdx(:,:,1,2).*dXdx(:,:,2,1);
end

end


function maximumtaylor(f, (a, b), n; absmax = false)
    x = setinterval(a, b)
    maybeabs = ifelse(absmax, abs, identity)
    PP = ArbPolyRing(parent(x), :x)

    # Compute the rest term of the Taylor expansion
    y = f(arb_series(PP([x, parent(x)(1)]), n + 1))

    restterm = (x - midpoint(x))^n*y[n]

    # If the rest term is not finite the result always be not finite
    if !isfinite(restterm)
        return restterm
    end

    # Compute the Taylor polynomial at the midpoint of x
    y = f(arb_series(PP([midpoint(x), parent(x)(1)]), n))

    # Enclose the Taylor polynomial evaluated at (a, b) - midpoint((a, b))
    mid = 0.5*(a + b)
    res = maximumpoly(y.poly, (a - mid, b - mid), absmax = absmax)
    p = res

    res += restterm

    return res
end


""
function constraint_voltage_magnitude_link(pm::_PM.AbstractACRModel, n_1::Int, n_2::Int, i::Int)
    vr_1 = var(pm, n_1, :vr, i)
    vi_1 = var(pm, n_1, :vi, i)

    vr_2 = var(pm, n_2, :vr, i)
    vi_2 = var(pm, n_2, :vi, i)

    JuMP.@constraint(pm.model, vr_1^2 + vi_1^2 == vr_2^2 + vi_2^2)
end

""
function constraint_power_balance_shunt_dispatch(pm::_PM.AbstractACRModel, n::Int, i::Int, bus_arcs, bus_arcs_dc, bus_arcs_sw, bus_gens, bus_storage, bus_shunts_var, bus_pd, bus_qd, bus_gs_const, bus_bs_const)
    vi = var(pm, n, :vi, i)
    vr = var(pm, n, :vr, i)
    p    = get(var(pm, n),    :p, Dict()); _PM._check_var_keys(p, bus_arcs, "active power", "branch")
    q    = get(var(pm, n),    :q, Dict()); _PM._check_var_keys(q, bus_arcs, "reactive power", "branch")
    pg   = get(var(pm, n),   :pg, Dict()); _PM._check_var_keys(pg, bus_gens, "active power", "generator")
    qg   = get(var(pm, n),   :qg, Dict()); _PM._check_var_keys(qg, bus_gens, "reactive power", "generator")
    ps   = get(var(pm, n),   :ps, Dict()); _PM._check_var_keys(ps, bus_storage, "active power", "storage")
    qs   = get(var(pm, n),   :qs, Dict()); _PM._check_var_keys(qs, bus_storage, "reactive power", "storage")
    psw  = get(var(pm, n),  :psw, Dict()); _PM._check_var_keys(psw, bus_arcs_sw, "active power", "switch")
    qsw  = get(var(pm, n),  :qsw, Dict()); _PM._check_var_keys(qsw, bus_arcs_sw, "reactive power", "switch")
    p_dc = get(var(pm, n), :p_dc, Dict()); _PM._check_var_keys(p_dc, bus_arcs_dc, "active power", "dcline")
    q_dc = get(var(pm, n), :q_dc, Dict()); _PM._check_var_keys(q_dc, bus_arcs_dc, "reactive power", "dcline")

    bs = get(var(pm, n), :bs, Dict()); _PM._check_var_keys(bs, bus_shunts_var, "reactive power", "shunt")

    # possibly can save 2x in function eval, but no the dominant runtime at this moment
    #vm_sqr = @variable(pm.model, start=1.0, base_name="$(0)_vm_sqr_$(i)")

    #JuMP.@constraint(pm.model, vm_sqr == vi^2 + vr^2)
    #cstr_p = JuMP.@constraint(pm.model, 0 == - sum(p[a] for a in bus_arcs) + sum(pg[g] for g in bus_gens) - sum(pd for pd in values(bus_pd)) - sum(gs for gs in values(bus_gs_const))*vm_sqr)
    #cstr_q = JuMP.@constraint(pm.model, 0 == - sum(q[a] for a in bus_arcs) + sum(qg[g] for g in bus_gens) - sum(qd for qd in values(bus_qd)) + sum(bs for bs in values(bus_bs_const))*vm_sqr + sum(bs[s]*vm_sqr for s in bus_shunts_var))

    cstr_p = JuMP.@constraint(pm.model, 0 == - sum(p[a] for a in bus_arcs) + sum(pg[g] for g in bus_gens) - sum(pd for pd in values(bus_pd)) - sum(gs for gs in values(bus_gs_const))*(vi^2 + vr^2))
    cstr_q = JuMP.@NLconstraint(pm.model, 0 == - sum(q[a] for a in bus_arcs) + sum(qg[g] for g in bus_gens) - sum(qd for qd in values(bus_qd)) + sum(bs for bs in values(bus_bs_const))*(vi^2 + vr^2) + sum(bs[s]*(vi^2 + vr^2) for s in bus_shunts_var))

    if _IM.report_duals(pm)
        sol(pm, n, :bus, i)[:lam_kcl_r] = cstr_p
        sol(pm, n, :bus, i)[:lam_kcl_i] = cstr_q
    end
end

function constraint_ohms_yt_from_goc(pm::_PM.AbstractACRModel, n::Int, f_bus, t_bus, f_idx, t_idx, g, b, g_fr, b_fr, tr, ti, tm)
    p_fr = var(pm, n, :p, f_idx)
    q_fr = var(pm, n, :q, f_idx)
    vr_fr = var(pm, n, :vr, f_bus)
    vr_to = var(pm, n, :vr, t_bus)
    vi_fr = var(pm, n, :vi, f_bus)
    vi_to = var(pm, n, :vi, t_bus)

    JuMP.@constraint(pm.model, p_fr ==  (g/tm^2+g_fr)*(vr_fr^2 + vi_fr^2) + (-g*tr+b*ti)/tm^2*(vr_fr*vr_to + vi_fr*vi_to) + (-b*tr-g*ti)/tm^2*(vi_fr*vr_to - vr_fr*vi_to) )
    JuMP.@constraint(pm.model, q_fr == -(b/tm^2+b_fr)*(vr_fr^2 + vi_fr^2) - (-b*tr-g*ti)/tm^2*(vr_fr*vr_to + vi_fr*vi_to) + (-g*tr+b*ti)/tm^2*(vi_fr*vr_to - vr_fr*vi_to) )
end

""
function expression_branch_power_ohms_yt_from_goc(pm::_PM.AbstractACRModel, n::Int, f_bus, t_bus, f_idx, t_idx, g, b, g_fr, b_fr, tr, ti, tm)
    vr_fr = var(pm, n, :vr, f_bus)
    vr_to = var(pm, n, :vr, t_bus)
    vi_fr = var(pm, n, :vi, f_bus)
    vi_to = var(pm, n, :vi, t_bus)

    var(pm, n, :p)[f_idx] =  (g/tm^2+g_fr)*(vr_fr^2 + vi_fr^2) + (-g*tr+b*ti)/tm^2*(vr_fr*vr_to + vi_fr*vi_to) + (-b*tr-g*ti)/tm^2*(vi_fr*vr_to - vr_fr*vi_to)
    var(pm, n, :q)[f_idx] = -(b/tm^2+b_fr)*(vr_fr^2 + vi_fr^2) - (-b*tr-g*ti)/tm^2*(vr_fr*vr_to + vi_fr*vi_to) + (-g*tr+b*ti)/tm^2*(vi_fr*vr_to - vr_fr*vi_to)
end

""
function expression_bus_withdrawal(pm::_PM.AbstractACRModel, n::Int, i::Int, bus_storage, bus_pd, bus_qd, bus_gs, bus_bs)
    vr = var(pm, n, :vr, i)
    vi = var(pm, n, :vi, i)
    ps = get(var(pm, n), :ps, Dict()); _PM._check_var_keys(ps, bus_storage, "active power", "storage")
    qs = get(var(pm, n), :qs, Dict()); _PM._check_var_keys(ps, bus_storage, "reactive power", "storage")

    ps_total = 0.0
    qs_total = 0.0
    if length(bus_storage) > 0
        ps_total = sum(ps[s] for s in bus_storage)
        qs_total = sum(qs[s] for s in bus_storage)
    end

    pd_total = 0.0
    qd_total = 0.0
    if length(bus_pd) > 0
        pd_total = sum(pd for pd in values(bus_pd))
    end
        if length(bus_qd) > 0
        qd_total = sum(qd for qd in values(bus_qd))
    end

    gs_total = 0.0
    bs_total = 0.0
    if length(bus_gs) > 0
        gs_total = sum(gs for gs in values(bus_gs))*(vr^2 + vi^2)
    end
    if length(bus_bs) > 0
        bs_total = sum(bs for bs in values(bus_bs))*(vr^2 + vi^2)
    end

    var(pm, n, :bus_wdp)[i] = ps_total + pd_total + gs_total
    var(pm, n, :bus_wdq)[i] = qs_total + qd_total + bs_total
end


using FourierCore, FourierCore.Grid, FourierCore.Domain
using FFTW, LinearAlgebra, BenchmarkTools, Random, JLD2
using KernelAbstractions
# using CUDA
import FourierCore.OverEngineering: PointWiseSpectralOperator

include("ode.jl")

arraytype = Array
Ω = S¹(2π)^2
N = 2^7 # number of gridpoints
grid = FourierGrid(N, Ω, arraytype=arraytype)
(; nodes, wavenumbers) = grid

x = nodes[1]
y = nodes[2]
kˣ = wavenumbers[1]
kʸ = wavenumbers[2]

# construct filter
kxmax = maximum(kˣ)
kymax = maximum(kˣ)
filter = @. (kˣ)^2 + (kʸ)^2 ≤ ((kxmax / 2)^2 + (kymax / 2)^2)

# Operators
∂x = im * kˣ
∂y = im * kʸ

# to convert things to damping timescale of smallest wavelength
δ = min(x[2] - x[1], y[2] - y[1])

Δ = @. ∂x^2 + ∂y^2
Δᵞ = @. δ^2 * Δ - δ^4 * Δ^2 + δ^6 * Δ^3 - δ^8 * Δ^4 # δ^2 * Δ * exp(-2 * δ^2 * Δ)
# Δᵞ = @. δ^2 * Δ * exp(-δ^2 * Δ)

Δ⁻¹ = 1 ./ Δ
bools = (!).(isnan.(Δ⁻¹))
Δ⁻¹ .*= bools # hack in the fact that false * NaN = 0

kmax = maximum(kˣ)
kcutoff = floor(Int, kmax / 2)
s = 1
a = log(1 + 4π / N) / (kmax - kcutoff)^s
k = @. sqrt((kˣ)^2 + (kʸ)^2)
filter_mask = @. k ≤ kcutoff
expfilter = @. exp(-a * (k - kcutoff)^s) * (1 - filter_mask) + filter_mask

# Note that the first order derivative operators only return a smooth derivative
∂ˣ = PointWiseSpectralOperator(filter .* ∂x)
∂ʸ = PointWiseSpectralOperator(filter .* ∂y)
∇² = PointWiseSpectralOperator(Δ)
∇⁻² = PointWiseSpectralOperator(Δ⁻¹)
∇ᵞ = PointWiseSpectralOperator(Δᵞ)
Filter = PointWiseSpectralOperator(expfilter)

# Ics
ψ = @. sin(x) * sin(y) + 0im
ζ = @. -sin(x) * sin(y) + 0im
dζ = @. -sin(x) * sin(y) + 0im
q = @. sin(x) + 0im * y
u = @. sin(x) * cos(y) + 0im
v = @. -cos(x) * sin(y) + 0im

# Tests: Quick Check
norm(-∂ˣ(ψ) - v)
norm(∂ʸ(ψ) - u)

#=
random forcing prep
rng = MersenneTwister(1234)
# Random.seed!(123456789)
Random.seed!(12)
φ̇ = arraytype(zeros(12,12))
rng = MersenneTwister(1234)
randn!(rng, φ̇)
φ̇ .*= sqrt(1 / 12)
=#

function rhs!(dS, S, b)
    # unpack state
    ζ = view(S, :, :, 1)
    dζ = view(dS, :, :, 1)
    # q = view(S, :, :, 2)
    # dq = view(dS, :, :, 2)
    q₁ = view(S, :, :, 2)
    dq₁ = view(dS, :, :, 2)

    q₁ .= Filter(q₁)
    ζ .= Filter(ζ)

    τ = 1e-2
    γ = 1.0
    e = 0.1
    c = 1 / τ * q₁ .* (real.(q₁) .> 0)
    # compute rhs
    ψ .= ∇⁻²(ζ)
    dζ .= ∂ʸ(ψ) .* ∂ˣ(ζ) - ∂ʸ(ζ) .* ∂ˣ(ψ) - 0.03 * ζ + 0.00 .* ∇ᵞ(ζ)
    # should evolve q₁ = q - qₛ where qₛ = γ * y # q = ⟨q⟩ + q'
    # q = q₁ + qₛ
    # dq .= ∂ʸ(ψ) .* ∂ˣ(q) - ∂ʸ(q) .* ∂ˣ(ψ) + 0.01 .* ∇ᵞ(q)
    dq₁ .= ∂ʸ(ψ) .* ∂ˣ(q₁) - ∂ʸ(q₁) .* ∂ˣ(ψ) - γ * ∂ˣ(ψ) + 0.00 .* ∇ᵞ(q₁)
    dq₁ .+= e .- c

    dq₁ .= Filter(dq₁)
    dζ .= Filter(dζ)
    return nothing
end

function stochastic_update!(S, dt, x, y)
    ζ = view(S, :, :, 1)
    a = 1.0
    a1 = a * 0.3 * randn()
    a2 = a * 0.3 * randn()
    a3 = a * 0.3 * randn()
    a4 = a * 0.3 * randn()
    a5 = a * 0.3 * randn()
    a6 = a * 0.3 * randn()
    a7 = a * 0.3 * randn()
    a8 = a * 0.3 * randn()
    a9 = a * 0.3 * randn()
    forcing = @. a1 * cos(2 * x) * cos(2 * y) + a2 * cos(3 * x) * cos(3 * y) + a5 * cos(2 * x) * cos(3 * y) + a6 * cos(3 * x) * cos(2 * y)
    forcing = @. a3 * sin(1 * x) * sin(3 * y) + a4 * sin(3 * x) * sin(1 * y)
    forcing = @. a7 * sin(2 * x) * cos(2 * y) + a8 * cos(2 * x) * sin(2 * y) + a9 * sin(2 * x) * sin(2 * y)
    ζ .+= sqrt(dt) .* forcing
end

# Legit forcing from Shafer Smith! Not that legit of a description though!!!
# c = 0.99
# theta = 2\pi * rand()
# F̂\_n = c F_n−1 + (1 − c\^2)\^(1/2) * exp(i * \theta)
# F\_n = havg(\psi F̂\_n)\^-1 F̂\_n
# havg here denotes the horizontal average across the domain.

# Create State array and unpack
S = arraytype(zeros(ComplexF64, size(ψ)..., 2))
ζ = view(S, :, :, 1)
q = view(S, :, :, 2)

min_δ = min(x[2] - x[1], y[2] - y[1])
cfl = 0.05
wavespeed = 2.0
dt = cfl * min_δ / wavespeed
# rhs!(dζ, ζ, 0)

# odesolver = LSRK144(rhs!, ζ, dt; t0=0)
odesolver = LSRK144(rhs!, S, dt; t0=0)
@. ζ = -sin(1.0 * x) * sin(y) + -sin(2.0 * x) * sin(y) + 0im
@. q = -sin(1.0 * x) + sin(0.0 * y) + 0im

timeseries = []
timeend = 10
iterations = floor(Int, timeend / dt)
for i in 1:iterations
    dostep!(S, odesolver)
    stochastic_update!(S, dt, x, y)
    if i % 40 == 0
        push!(timeseries, Array(real.(copy(S))))
    end
    if i % 100 == 0
        println("done with iteration ", i)
        println("the extrema of vorticity are ", extrema(real.(ζ)))
        println("the extrema of the tracer are ", extrema(real.(q)))
    end
end

##
using GLMakie
fig = Figure()
ax = Axis(fig[1, 1], title="vorticity")
ax2 = Axis(fig[1, 3], title="saturation")

timeslider = Slider(fig[2, 1:2], range=1:1:length(timeseries), startvalue=1)
timeindex = timeslider.value
field1 = @lift(timeseries[$timeindex][:, :, 1])
field2 = @lift(timeseries[$timeindex][:, :, 2])
hm_ζ = heatmap!(ax, field1, colorrange=(-1, 1), colormap=:balance, interpolate=true)
Colorbar(fig[1, 2], hm_ζ, height=Relative(3 / 4), width=25, ticklabelsize=30, labelsize=30, ticksize=25, tickalign=1,)

hm_q = heatmap!(ax2, field2, colorrange=(-5, 0.02), colormap=:solar, interpolate=true)
Colorbar(fig[1, 4], hm_q, height=Relative(3 / 4), width=25, ticklabelsize=30, labelsize=30, ticksize=25, tickalign=1,)
display(fig)


using Random
using LinearAlgebra
using Base.Threads
using IterativeSolvers
using TensorOperations
import Base.isempty
import Base.eltype
import Base.copy

"""
TT constructor of a tensor
``C[\\mu_1,…,\\mu_d] = A_1[1][\\mu_1]*⋯*A_d[d][\\mu_d] \\in \\mathbb{K}^{n_1 \\times ... \\times n_d}``
The following properties are stored
	* ttv_vec: the TT cores A_k as a list of 3-order tensors ``(A_1,...,A_d)`` where ``A_k = A_k[\\mu_k,\\alpha_{k-1},\\alpha_k]``, ``1 \\leq \\alpha_{k-1} \\leq r_{k-1}``, ``1 \\leq \\alpha_{k} \\leq r_{k}``, ``1 \\leq \\mu_k \\leq n_k``
	* ttv_dims: the dimension of the tensor along each mode
	* ttv_rks: the TT ranks ``(r_0,...,r_d)`` where ``r_0=r_d=1``
	* ttv_ot: the orthogonality of the TT where
		* ttv_ot[i] = 1 iff ``A_i`` is right-orthogonal *i.e.* ``\\sum_{\\mu_i} A_i[\\mu_i]^T A_i[\\mu_i] = I_{r_i}``
		* ttv_ot[i] = -1 iff ``A_i`` is left-orthogonal *i.e.* ``\\sum_{\\mu_i} A_i[\\mu_i] A_i[\\mu_i]^T = I_{r_{i-1}}``
		* ttv_ot[i] = 0 if nothing is known
"""
struct ttvector{T<:Number}
	ttv_vec :: Array{Array{T,3},1}
	ttv_dims :: Array{Int64,1}
	ttv_rks :: Array{Int64,1}
	ttv_ot :: Array{Int64,1}
end

Base.eltype(::ttvector{T}) where T<:Number = T
"""
Vidal representation of TT vector
"""
struct tt_vidal{T<:Number}
	#Vidal representation of a higher-order tensor
	#C[μ_1,…,μ_L] = core[μ_1] * Diagonal(Σ_1) * … * Diagonal(Σ_{L-1}) * core[μ_L]
	#Cores are orthogonal and Σ are the higher-order singular values
	core :: Array{Array{T,3},1}
	Σ :: Array{Array{Float64,1},1}
	dims :: Array{Int64,1}
	rks :: Array{Int64,1}
end

"""
TT constructor of a matrix ``M[i_1,..,i_d;j_1,..,j_d] \\in \\mathbb{K}^{n_1 \\cdots n_d \\times n_1 \\cdots n_d}`` where
``MM[i_1,..,i_d;j_1,..,j_d] = A_1[i_1,j_1] ... A_L[i_d,j_d]``
The following properties are stored
	* tto_vec: the TT cores A_k as a list of 4-order tensors ``(A_1,...,A_d)``
	* ttv_dims: the dimension of the tensor along each mode
	* ttv_rks: the TT ranks ``(r_0,...,r_d)`` where ``r_0=r_d=1``
"""
struct ttoperator{T<:Number}
	tto_vec :: Array{Array{T,4},1}
	tto_dims :: Array{Int64,1}
	tto_rks :: Array{Int64,1}
	tto_ot :: Array{Int64,1}
end

Base.eltype(::ttoperator{T}) where T<:Number = T

"""
returns a zero ttvector with dimensions `dims` and ranks `rks`
"""
function zeros_tt(dims,rks;T=Float64,ot=zeros(Int,length(dims)))
	@assert length(dims)+1==length(rks) "Dimensions and ranks are not compatible"
	d = length(dims)
	tt_vec = Array{Array{T,3}}(undef,d)
	for i in 1:d
		tt_vec[i] = zeros(T,dims[i],rks[i],rks[i+1])
	end
	return ttvector{T}(tt_vec,dims,copy(rks),copy(ot))
end

#returns partial isometry Q ∈ R^{n x m}
function rand_orthogonal(n,m;T=Float64)
    N = max(n,m)
    q,r = qr(rand(T,N,N))
    return Matrix(q)[1:n,1:m]
end

"""
Returns a random TTvector with dimensions `dims` and ranks `rks`
"""
function rand_tt(dims,rks;T=Float64)
	@assert length(dims)+1==length(rks) "Dimensions and ranks are not compatible"
	d = length(dims)
	tt_vec = Array{Array{T,3}}(undef,d)
	for i in 1:d
		tt_vec[i] = randn(T,dims[i],rks[i],rks[i+1])
	end
	return ttvector{T}(tt_vec,dims,copy(rks),zeros(Int,d))
end

function Base.copy(x_tt::ttvector{T}) where T<:Number
	y_tt = zeros_tt(x_tt.ttv_dims,x_tt.ttv_rks;T=T,ot=x_tt.ttv_ot)
	@threads for i in 1:length(x_tt.ttv_dims)
		y_tt.ttv_vec[i] = copy(x_tt.ttv_vec[i])
	end
	return y_tt
end

"""
TT decomposition by the Hierarchical SVD algorithm
	* Oseledets, I. V. (2011). Tensor-train decomposition. *SIAM Journal on Scientific Computing*, 33(5), 2295-2317.
	* Schollwöck, U. (2011). The density-matrix renormalization group in the age of matrix product states. *Annals of physics*, 326(1), 96-192.
The *root* of the TT decomposition is at index *i.e.* ``A_i`` for ``i < index`` are left-orthogonal and ``A_i`` for ``i > index`` are right-orthogonal. Singular values lower than tol are discarded.
"""
function ttv_decomp(tensor::Array{T};index=1,tol=1e-12) where T<:Number
	# Decomposes a tensor into its tensor train with core matrices at i=index
	dims = collect(size(tensor)) #dims = [n_1,...,n_d]
	n_max = maximum(dims)
	d = length(dims)
	ttv_vec = Array{Array{T}}(undef,d)
	# ttv_ot[i]= -1 if i < index
	# ttv_ot[i] = 0 if i = index
	# ttv_ot[i] = 1 if i > index
	ttv_ot = -ones(Int64,d)
	ttv_ot[index] = 0
	if index < d
		ttv_ot[index+1:d] = ones(d-index)
	end
	rks = ones(Int, d+1)
	tensor_curr = tensor
	# Calculate ttv_vec[i] for i < index
	for i = 1 : (index - 1)
		# Reshape the currently left tensor
		tensor_curr = reshape(tensor_curr, Int(rks[i] * dims[i]), :)
		# Perform the singular value decomposition
		u, s, v = svd(tensor_curr)
		# Define the i-th rank
		rks[i+1] = length(s[s .>= tol])
		# Initialize ttv_vec[i]
		ttv_vec[i] = zeros(T,dims[i],rks[i],rks[i+1])
		# Fill in the ttv_vec[i]
		for x = 1 : dims[i]
			ttv_vec[i][x, :, :] = u[(rks[i]*(x-1) + 1):(rks[i]*x), :]
		end
		# Update the currently left tensor
		tensor_curr = Diagonal(s[1:rks[i+1]])*v'[1:rks[i+1],:]
	end

	# Calculate ttv_vec[i] for i > index
	if index < d
		for i = d : (-1) : (index + 1)
			# Reshape the currently left tensor
			tensor_curr = reshape(tensor_curr, :, dims[i] * rks[i+1])
			# Perform the singular value decomposition
			u, s, v = svd(tensor_curr)
			# Define the (i-1)-th rank
			rks[i]=length(s[s .>= tol])
			# Initialize ttv_vec[i]
			ttv_vec[i] = zeros(T, dims[i], rks[i], rks[i+1])
			# Fill in the ttv_vec[i]
			i_vec = zeros(Int,rks[i+1])
			for x = 1 : dims[i]
				i_vec = dims[i]*((1:rks[i+1])-ones(Int, rks[i+1])) + x*ones(Int,rks[i+1])
				ttv_vec[i][x, :, :] = v'[1:rks[i],i_vec] #(rks[i+1]*(x-1)+1):(rks[i+1]*x)
			end
			# Update the currently left tensor
			tensor_curr = u[:,1:rks[i]]*Diagonal(s[1:rks[i]])
		end
	end
	# Calculate ttv_vec[i] for i = index
	# Reshape the currently left tensor
	tensor_curr = reshape(tensor_curr, Int(dims[index]*rks[index]),:)
	# Initialize ttv_vec[i]
	ttv_vec[index] = zeros(T, dims[index], rks[index], rks[index+1])
	# Fill in the ttv_vec[i]
	for x = 1 : dims[index]
		ttv_vec[index][x, :, :] =
		tensor_curr[Int(rks[index]*(x-1) + 1):Int(rks[index]*x), 1:rks[index+1]]
	end

	# Define the return value as a ttvector
	return ttvector{T}(ttv_vec, dims, rks, ttv_ot)
end

"""
Returns the tensor corresponding to x_tt
"""
function ttv_to_tensor(x_tt :: ttvector{T}) where T<: Number
	d = length(x_tt.ttv_dims)
	r_max = maximum(x_tt.ttv_rks)
	# Initialize the to be returned tensor
	tensor = zeros(T, x_tt.ttv_dims...)
	# Fill in the tensor for every t=(x_1,...,x_d)
	for t in CartesianIndices(tensor)
		curr = ones(T,r_max)
		a = collect(Tuple(t))
		for i = d:-1:1
			curr[1:x_tt.ttv_rks[i]] = x_tt.ttv_vec[i][a[i],:,:]*curr[1:x_tt.ttv_rks[i+1]]
		end
		tensor[t] = curr[1]
	end
	return tensor
end

#returns the Vidal representation of a TT
function tt_to_vidal(x_tt::ttvector{T};tol=1e-14) where T<:Number
	d = length(x_tt.ttv_dims)
	core = Array{Array{T,3},1}(undef,d)
	Σ = Array{Array{Float64,1},1}(undef,d-1)
	y_tt = orthogonalize(x_tt)
	y_rks = copy(y_tt.ttv_rks)
	for j in 1:d-1
		A = zeros(y_tt.ttv_dims[j],y_rks[j],y_tt.ttv_dims[j+1],y_rks[j+2])
		@tensor A[a,b,c,d] = y_tt.ttv_vec[j][a,z,b]*y_tt.ttv_vec[j+1][z,d,c]
		u,s,v = svd(reshape(A,size(A,1)*size(A,2),:))
		Σ[j] = s[s.>tol]
		y_rks[j+1] = length(Σ[j])
		core[j] = permutedims(reshape(u[:,s.>tol],y_rks[j],y_tt.ttv_dims[j],:),[1 3 2])
		y_tt.ttv_vec[j+1] = reshape(v[:,s.>tol],:,y_rks[j+2],y_tt.ttv_dims[j+1])
	end
	core[d] = y_tt.ttv_vec[d]
	return tt_vidal(core,Σ,y_tt.ttv_dims,y_rks)
end

"""
Transforms a ttoperator into a ttvector
"""
function tto_to_ttv(A::ttoperator{T}) where T<:Number
	d = length(A.tto_dims)
	xtt_vec = Array{Array{T,3},1}(undef,d)
	A_rks = A.tto_rks
	for i in 1:d
		xtt_vec[i] = reshape(A.tto_vec[i],A.tto_dims[i]^2,A_rks[i],A_rks[i+1])
	end
	return ttvector{T}(xtt_vec,A.tto_dims.^2,A.tto_rks,A.tto_ot)
end

"""
Transforms a ttvector (coming from a ttoperator) into a ttoperator
"""
function ttv_to_tto(x::ttvector{T}) where T<:Number
	d = length(x.ttv_dims)
	@assert(isqrt.(x.ttv_dims).^2 == x.ttv_dims, DimensionMismatch)
	Att_vec = Array{Array{T,4},1}(undef,d)
	x_rks = x.ttv_rks
	A_dims = zeros(Int64,d)
	for i in 1:d
		A_dims[i] = isqrt(x.ttv_dims[i])
		Att_vec[i] = reshape(x.ttv_vec[i],A_dims[i],A_dims[i],x_rks[i],x_rks[i+1])
	end
	return ttoperator{T}(Att_vec,A_dims,x.ttv_rks,x.ttv_ot)
end

"""
Returns the TT decomposition of a matrix using the HSVD algorithm
"""
function tto_decomp(tensor::Array{T}; index=1) where T<:Number
	# Decomposes a tensor operator into its tensor train
	# with core matrices at i=index
	# The tensor is given as tensor[x_1,...,x_d,y_1,...,y_d]
	d = Int(ndims(tensor)/2)
	tto_dims = collect(size(tensor))[1:d]
	n_max = maximum(tto_dims)
	dims_sq = tto_dims.^2
	# The tensor is reorder  into tensor[x_1,y_1,...,x_d,y_d],
	# reshaped into tensor[(x_1,y_1),...,(x_d,y_d)]
	# and decomposed into its tensor train with core matrices at i= index
	index_sorted = reshape(Transpose(reshape(1:(2*d),:,2)),1,:)
	ttv = ttv_decomp(reshape(permutedims(tensor,index_sorted),(dims_sq[1:(end-1)]...), :); index=index)
	# Define the array of ranks [r_0=1,r_1,...,r_d]
	rks = ttv.ttv_rks
	r_max = maximum(rks)
	# Initialize tto_vec
	tto_vec = Array{Array{T}}(undef,d)
	# Fill in tto_vec
	for i = 1:d
		# Initialize tto_vec[i]
		tto_vec[i] = zeros(T, tto_dims[i], tto_dims[i], rks[i], rks[i+1])
		# Fill in tto_vec[i]
		tto_vec[i] = reshape(ttv.ttv_vec[i], tto_dims[i], tto_dims[i], :, rks[i+1])
	end
	return ttoperator{T}(tto_vec, tto_dims, rks, ttv.ttv_ot)
end

function tto_to_tensor(tto :: ttoperator{T}) where T<:Number
	d = length(tto.tto_dims)
	# Define the array of ranks [r_0=1,r_1,...,r_d]
	rks = tto.tto_rks
	r_max = maximum(rks)
	# The tensor has dimensions [n_1,...,n_d,n_1,...,n_d]
	dims = zeros(Int, 2*d)
	dims[1:d] = tto.tto_dims
	dims[(d+1):(2*d)] = tto.tto_dims
	tensor = zeros(T,dims...)
	# Fill in the tensor for every t=(x_1,...,x_d,y_1,...,y_d)
	for t in CartesianIndices(tensor)
		curr = ones(T,r_max)
		a = collect(Tuple(t))
		for i = d:-1:1
			curr[1:rks[i]] = tto.tto_vec[i][t[i], t[d + i], :, :]*curr[1:rks[i+1]]
		end
		tensor[t] = curr[1]
	end
	return tensor
end

#TTO representation of the identity matrix
function id_tto(d;n_dim=2,T=Float64)
	dims = n_dim*ones(Int64,d)
	A = Array{Array{T,4},1}(undef,d)
	for j in 1:d
		A[j] = zeros(2,2,1,1)
		A[j][:,:,1,1] = Matrix{T}(I,2,2)
	end
	return ttoperator{T}(A,dims,ones(Int64,d+1),zeros(d))
end


#Gravity related stuff

#Additional source terms
function source_terms(hyd::data2d, dt)

    #self-gravity
    #gxflx, gyflx = selfgravity(hyd.rho, hyd.nx, hyd.ny, r32x, r32y)

    #constant gravity

    #x-dir
    dir = 1

    #+
    gxflx, gyflx = ygravity(hyd.rhop[:,:,dir], hyd.nx, hyd.ny)
    hyd.velxp[:,:,dir] += 0.5*gxflx
    hyd.velyp[:,:,dir] += 0.5*gyflx
    hyd.epsp[:,:,dir] += 0.5*((hyd.velxp[:,:,dir] .* gxflx) .+ (hyd.velyp[:,:,dir] .* gyflx))

    #-
    gxflx, gyflx = ygravity(hyd.rhom[:,:,dir], hyd.nx, hyd.ny)
    hyd.velxm[:,:,dir] += 0.5*gxflx
    hyd.velym[:,:,dir] += 0.5*gyflx
    hyd.epsm[:,:,dir] += 0.5*((hyd.velxm[:,:,dir] .* gxflx) .+ (hyd.velym[:,:,dir] .* gyflx))

    #y-dir
    dir = 2

    #+
    gxflx, gyflx = ygravity(hyd.rhop[:,:,dir], hyd.nx, hyd.ny)
    hyd.velxp[:,:,dir] += 0.5*gxflx
    hyd.velyp[:,:,dir] += 0.5*gyflx
    hyd.epsp[:,:,dir] += 0.5*((hyd.velxp[:,:,dir] .* gxflx) .+ (hyd.velyp[:,:,dir] .* gyflx))

    #-
    gxflx, gyflx = ygravity(hyd.rhom[:,:,dir], hyd.nx, hyd.ny)
    hyd.velxm[:,:,dir] += 0.5*gxflx
    hyd.velym[:,:,dir] += 0.5*gyflx
    hyd.epsm[:,:,dir] += 0.5*((hyd.velxm[:,:,dir] .* gxflx) .+ (hyd.velym[:,:,dir] .* gyflx))

    return hyd
end

#Constant y-dir gravity
function ygravity(rho::AbstractMatrix,
                  nx::Int64,
                  ny::Int64;
                  gx=0.0,
                  gy=-1.0)

    gxflux = gx .* rho
    gyflux = gy .* rho

    return gxflux, gyflux
end

#Calculate x(x^2 + y^2)^-3/2 and y(x^2 + y^2)^-3/2 kernel
#for convolution used in self-gravity
function r32kernel(x, y)

    Nx = length(x)
    Ny = length(y)
    dx = x[2] - x[1]
    dy = y[2] - y[1]

    r32x = zeros(2Ny-1, 2Nx-1)
    r32y = zeros(2Ny-1, 2Nx-1)

    r32q = zeros(Ny, Nx)

    #construct 1 quarter of the kernel
    for j = 0:(Ny-1), i = 0:(Nx-1)
        r32q[j+1, i+1] = ((i*dx)^2.0 + (j*dy)^2.0)^(-1.5)
    end
    r32q[1,1] = 0.0

    #x - axis component
    r32x[1:Ny, Nx:2Nx-1] = flipud(r32q)
    r32x[Ny:2Ny-1, Nx:2Nx-1] = r32q

    #multiply by x
    for j = 1:2Ny-1, i = Nx:2Nx-1
        r32x[j, i] = r32x[j, i]*dx*(i-Nx)
    end

    #flip negative values
    r32x[:, 1:Nx-1] = -1.0*fliplr(r32x[:, Nx+1:end])

    #y - axis component
    r32y[Ny:2Ny-1, 1:Nx] = fliplr(r32q)
    r32y[Ny:2Ny-1, Nx:2Nx-1] = r32q

    #multiply by y
    for j = Ny:2Ny-1, i = 1:2Nx-1
        r32y[j, i] = r32y[j, i]*dy*(j-Ny)
    end

    #flip negative values
    r32y[1:Ny-1, :] = -1.0*flipud(r32y[Ny+1:end, :])

    return r32x, r32y
end

#Self-gravity
function selfgravity(rho::AbstractMatrix,
                     nx::Int64,
                     ny::Int64,
                     r32x::AbstractMatrix,
                     r32y::AbstractMatrix)

    xr = nx:(2nx-1)
    yr = ny:(2ny-1)

    gxflux = conv2(rho, r32x)[xr, yr]
    gyflux = conv2(rho, r32y)[xr, yr]

    gxflux = gxflux .* rho
    gyflux = gyflux .* rho

    return gxflux, gyflux
end


using Jecco
using Jecco.AdS5_3_1
using Plots
gr()
include("../src/AdS5_3_1/GW_2.jl")
#/Users/apple/Dropbox/CollisionNewPotential/3+1 project/runs/spinodal2D_model1_en0.9_run01/#
dirname = "/Users/apple/Documents/Jecco.jl/data/spinodal2D_e_0.817_Lxy_30_adaptive"
field = :energy
dt = 0.1
alg = AB3()
h, h_t, Tk, kx, ky = solve_GW(dirname, field, dt=dt, alg=alg)

#=
tmax = (length(h[:,1,1])-1)*dt
t = 0:dt:tmax
imax = 1:length(t);
j=15;
plot(t[imax], real.(h[imax,j,1]), label="h", lw=3)
plot!(t[imax], real.(h_t[imax,j,1]), label="h_t", lw=3)
plot!(t[imax], real.(Tk[imax,j,1]), label="Tk", lw=3)
=#


triangle(n::Int) = n >= 0 ? [[binomial(i, j) for j in 0:i] for i in 0:n-1] : throw(DomainError(n, "n must be strictly positive"))


using Primes

let x = monom(QQ{Int}[:x], 1), P = ZZ{Int}[:x]

@testset "cyclotomic" begin
    @test deg(cyclotomic(P, 385)) == totient(385)
end

@testset "kronecker ($n, $k, $r)" for (n, k, r) in ((1,1,1), (1,13,1), (12,1,1), (5, 8, -1), (12,18,0), (6, 17, -1), (6, -17, -1), (-6, 17, -1), (-6, -17, 1))
    @test kronecker(n, k) == r
end

@testset "moebius ($n, $r)" for (n, r) in ((1,1), (2,-1),(3,-1),(4,0),(5,-1),(6,1),(8,0),(10,1))
    @test moebius(n) == r
end

@testset "necklace ($n, $r)" for (n, r) in ((1,x), (2,(x^2 - x)/2), (6,(x^6-x^3-x^2+x)/6))
    @test necklace(x, n) == r
    @test necklace(17, n) == r(17)
end
end


# Vicenty's inverse (two points -> ellipsoidal distance and azimuths)
const ANTIPODAL_ERROR =
    ArgumentError("Points too near antipodal for distance algorithm to converge")

function vicentys_inverse{T <: LL}(a::T, b::T)
    d = ellipsoid(T)
    f = 1 - d.b/d.a

    # Reduced latitude (latitude on the auxiliary sphere)
    U₁ = atan(d.b / d.a * tand(a.lat))
    U₂ = atan(d.b / d.a * tand(b.lat))

    sinU₁, cosU₁ = sin(U₁), cos(U₁)
    sinU₂, cosU₂ = sin(U₂), cos(U₂)

    cosU₁cosU₂ = cosU₁ * cosU₂
    sinU₁sinU₂ = sinU₁ * sinU₂
    sinU₁cosU₂ = sinU₁ * cosU₂
    cosU₁sinU₂ = cosU₁ * sinU₂

    ΔL = deg2rad(b.lon - a.lon)

    # Needed in scope after the loop
    oldλ = sinλ = cosλ = sinσ = cosσ = σ = cos²α = cos2σm = cos²2σm = Inf
    itr = 0

    # Initialize λ to ΔL for first iteration
    λ = ΔL

    for itr=1:200
        oldλ = λ

        sinλ, cosλ = sin(λ), cos(λ)

        x = cosU₂ * sinλ
        y = cosU₁sinU₂ - sinU₁cosU₂ * cosλ
        sinσ = sqrt(x*x + y*y)
        cosσ = sinU₁sinU₂ + cosU₁cosU₂ * cosλ
        σ = atan2(sinσ, cosσ)

        sinα = cosU₁cosU₂ * sinλ / sinσ
        cos²α = 1 - sinα*sinα
        cos2σm = cosσ - 2sinU₁sinU₂ / cos²α
        cos²2σm = cos2σm * cos2σm
        C = f / 16 * cos²α * (4 + f*(4-3cos²α))

        λ = ΔL + (1-C) * f * sinα *
            (σ + C * sinσ * (cos2σm + C * cosσ * (-1 + 2cos²2σm)))

        abs(λ - oldλ) < 1e-12 && break
    end

    # Failure to converge happens when abs(λ) > π after the first iteration
    itr == 200 && throw(ANTIPODAL_ERROR)

    a², b² = d.a * d.a, d.b * d.b
    u² = cos²α * (a² - b²) / b²
    A = 1 + u² / 16384 * (4096 + u² * (-768 + u² * (320 - 175u²)))
    B = u² / 1024 * (256 + u² * (-128 + u² * (74 - 47u²)))
    Δσ = B * sinσ * (cos2σm + B / 4 *
        (cosσ * (-1 + 2cos²2σm) -
         B / 6 * cos2σm * (-3 + 4sinσ*sinσ) * (-3 + 4cos²2σm)))

    s = d.b * A * (σ - Δσ) # ellipsoidal distance
    α₁ = atan2(cosU₂ * sinλ,  cosU₁sinU₂ - sinU₁cosU₂ * cosλ) # start azimuth
    α₂ = atan2(cosU₁ * sinλ, -sinU₁cosU₂ + cosU₁sinU₂ * cosλ) # end azimuth

    return s, rad2deg(α₁), rad2deg(α₂)
end


module day10

using OffsetArrays

readInput() = parse.(Int, eachline(joinpath(@__DIR__, "input.txt")))

function part1(input)
    list = sort(input)
    pushfirst!(list, 0)
    push!(list, list[end] + 3)
    nbOnes, nbThrees = 0, 0
    for i = 1:length(list) - 1
        if list[i + 1] - list[i] == 1
            nbOnes += 1
        elseif list[i + 1] - list[i] == 3
            nbThrees += 1
        end
    end
    nbOnes * nbThrees
end

function part2(input)
    list = sort(input)
    push!(list, list[end] + 3)
    nbWays = OffsetVector(zeros(Int, list[end] + 1), -1)
    nbWays[0] = 1
    for i in list
        nbWays[i] = nbWays[i - 1] + (i > 1 ? nbWays[i - 2] : 0) + (i > 2 ? nbWays[i - 3] : 0)
    end
    nbWays[end]
end

end


#####
# ListNode
#####
import LinearAlgebra
Base.@kwdef mutable struct ListNode{V}
    val::V = 0
    next::Union{Nothing,ListNode{V}} = nothing
end

ListNode(x) = ListNode(x, nothing)

function ListNode{V}(x::Vector{V}) where {V}
    head = cur = ListNode(x[1])
    for i in 2:length(x)
        cur.next = ListNode(x[i])
        cur = cur.next
    end
    return head
end

val(x::ListNode) = x.val
val!(x::ListNode, y) = x.val = y
next(x::ListNode) = x.next
next!(x::ListNode, y) = x.next = y

Base.:(==)(x::ListNode, y::ListNode) = x.val == y.val && x.next == y.next

#####
# TreeNode
#####

Base.@kwdef mutable struct TreeNode{V}
    val::V = 0
    left::Union{Nothing,TreeNode} = nothing
    right::Union{Nothing,TreeNode} = nothing
end

TreeNode(x) = TreeNode(; val=x)
TreeNode(x, left) = TreeNode(; val=x, left=left)

function Base.:(==)(t1::TreeNode, t2::TreeNode)
    return t1.val == t2.val && t1.left == t2.left && t1.right == t2.right
end

function Base.convert(::Type{TreeNode{V}}, xs::Vector) where {V}
    @assert length(xs) > 0
    root = _build_tree!(V, xs)
    return root
end

TreeNode{V}(xs::Vector) where {V} = convert(TreeNode{V}, xs)

function _build_tree!(::Type{T}, xs::Vector) where T
    res = TreeNode(xs[1])
    q1 = TreeNode{T}[res]
    q2 = TreeNode{T}[]
    idx, len = 2, length(xs)
    while idx <= len
        while !isempty(q1)
            qf = popfirst!(q1)
            idx > len && break
            if isnothing(xs[idx])
                qf.left = nothing
            else
                qf.left = TreeNode(xs[idx])
                push!(q2, qf.left)
            end
            (idx += 1) > len && break
            if isnothing(xs[idx])
                qf.right = nothing
            else
                qf.right = TreeNode(xs[idx])
                push!(q2, qf.right)
            end
            idx += 1
        end
        q1, q2 = q2, q1
    end
    return res
end

function next_perm!(itr)::Bool
    (isempty(itr) || length(itr) == 1) && return false

    i = findlast(idx -> itr[idx + 1] > itr[idx], 1:lastindex(itr) - 1)
    isnothing(i) && (reverse!(itr); return false)

    j = findlast(idx -> itr[idx] > itr[i], i+1:lastindex(itr)) + i
    itr[i], itr[j] = itr[j], itr[i]
    reverse!(itr, i + 1)
    return true
end

"""
``res = mat ^ p \\pmod m``
"""
function mat_fast_mul(mat::AbstractMatrix{T}, p::T, m::T=typemax(T)) where T
    p == 0 && return (ones(T, size(mat)))
    (m == 1 || m == -1) && return (zeros(T, size(mat)))
    base = mat .% m
    res = copy(base)
    tmp = typeof(mat)(undef, size(mat))
    t = prevpow(2, p)
    p -= t
    while true
        if p >= t
            LinearAlgebra.mul!(tmp, res, base)
            tmp, res = res, tmp
            res .%= m
            p -= t
        end
        t >>>= 1
        t <= 0 && break
        LinearAlgebra.mul!(tmp, res, res)
        tmp, res = res, tmp
        res .%= m
    end
    res
end


Random.seed!(0)

#sig = [1 2; 2 2]
#sig = [1 1; 2 3]
sig = [2 3]
#sig = [2 2]
#sig = [3 2; 2 3]

S = random_S0Graph(sig)

w = random_bounded(S.n)

@time opt1, x1 = dsw(S, w, eps=solver_eps)
@time opt2, x2 = dsw(S, w, use_diag_optimization=false, eps=solver_eps)
@test opt1 ≈ opt2  atol=tol


"""
```
visualize_flow(flow, ColorBased(), RasterConvention())
visualize_flow(flow, ColorBased(), CartesianConvention())
```

Returns an image that matches the dimensions of the input matrix and that depicts the
orientation and magnitude of the optical flow vectors using the HSV color space.

# Details
Hue encodes angle between optical flow vector and the x-axis in the image plane;
saturation encodes the ratio between the individual vector magnitudes and the maximum magnitude
among the whole motion field, and the values always equal one.

# Arguments
The flow parameter needs to be a two-dimensional arrays of length-2 vectors (of type SVector)
which represent the displacement of each pixel. The displacement can be expressed in a
coordinate system based on a RasterConvention() (i.e. rows and columns of a matrix)
or a CartesianConvention().

# Example

Compute the HSV encoded visualization of flow vectors in (row, column) convention.
```julia

using ImageTracking

hsv = visualize_flow(flow, ColorBased(), RasterConvention())

imshow(RGB.(hsv))
```

Compute the HSV encoded visualization of flow vectors in (x, y) convention.
```julia

using ImageTracking

hsv = visualize_flow(flow, ColorBased(), CartesianConvention())

imshow(RGB.(hsv))
```

# References
[1] S. Baker, D. Scharstein, JP Lewis, S. Roth, M.J. Black, and R. Szeliski. A database and
evaluation methodology for optical flow. International Journal of Computer Vision, 92(1):1–31, 2011.
"""
function visualize_flow(flow::Array{SVector{2, Float64}, 2}, method::ColorBased, convention::RasterConvention)
	# Convert from (row,column) to (x,y) convention.
	map!(x-> SVector(last(x),first(x)), flow, flow)

    # Display optical flow as an image, with hue encoding the orientation and
    # saturation encoding the relative magnitude.
    max_norm = maximum(map(norm, flow))
    normalised_flow = map(PolarFromCartesian(), flow / max_norm)
    hsv = zeros(HSV{Float32},size(flow))

    for i in eachindex(flow)
        hsv[i] = HSV((normalised_flow[i].θ + pi) * 180 / pi, normalised_flow[i].r, 1)
    end

    return hsv
end

function visualize_flow(flow::Array{SVector{2, Float64}, 2}, method::ColorBased, convention::CartesianConvention)
    # Display optical flow as an image, with hue encoding the orientation and
    # saturation encoding the relative magnitude.
    max_norm = maximum(map(norm, flow))
    normalised_flow = map(PolarFromCartesian(), flow / max_norm)
    hsv = zeros(HSV{Float32},size(flow))

    for i in eachindex(flow)
        hsv[i] = HSV((normalised_flow[i].θ + pi) * 180 / pi, normalised_flow[i].r, 1)
    end

    return hsv
end

function is_flow_known(p::SVector{2, Float64})
    return !isnan(first(p)) && !isnan(last(p)) && abs(first(p)) < 1e9 && abs(last(p)) < 1e9
end

"""
```
read_flow_file(file_path)
```

Returns a 2-Dimensional array that depicts the optical flow.

# Details
Reads simple .flo file format which is used for optical flow evaluation.
The .flo file should store 2-band float image for horizontal (u) and vertical (v) flow components.
Floats are stored in little endian format.

bytes contain -
  0-3     tag: "PIEH" in ASCII, which in little endian happens to be the float 202021.25
          (just a sanity check that floats are represented correctly)
  4-7     width as an integer
  8-11    height as an integer
  12-end  data (width*height*2*4 bytes total)
          the float values for u and v, interleaved, in row order, i.e.,
          u[row0,col0], v[row0,col0], u[row0,col1], v[row0,col1], ...

# Arguements
The file_path parameter needs to be a string, storing the address of .flo file.

# Example

Read flow file from .flo file and visualize the flow in HSV
```julia

using ImageTracking

file_path = "./Example_flow.flo"
flow = read_flow_file(file_path)
result = visualize_flow(flow, ColorBased(), RasterConvention())
imshow(RGB.(result))
```

# References
[1] http://vision.middlebury.edu/flow/data/
"""
# first four bytes to be same in little endian
# check this when reading a .flo file
FLOAT_TAG = 202021.25
function read_flow_file(file_path::String)

    io = open(file_path, "r")
    seekstart(io)

    tag = read(io, Float32) # reading the initial tag

    if tag != FLOAT_TAG
        print("read_from_file: Problem reading file ", filepath)
        return -1
    end

    width  = read(io, Int32)
    height = read(io, Int32)

    flow = Array{SVector{2, Float64}, 2}(undef,height,width)
    for i in 1:height
        for j in 1:width
            if eof(io) == false
                fx = read(io, Float32)
                fy = read(io, Float32)
                flow[i, j] = SVector{2, Float64}(fx,fy)
            end
        end
    end

    close(io)
    return flow
end

"""
```
evaluate_flow_error(ground_truth_flow, estimated_flow, EndpointError())
```

Returns a 2-Dimensional array that matches the dimensions of the input flow vector and
that depicts the end point error between the estimated flow and the ground truth flow.

# Details
If the estimated flow at a point is (u0, v0) and ground truth flow is (u1, v1), then
error will be sqrt[(u0 - u1)^2 + (v0 - v1)^2] at that point.

# Arguments
The flow parameters needs to be two-dimensional arrays of length-2 vectors (of type SVector)
which represent the displacement of each pixel.

# Example

Compute the end point error between two flows.
```julia

using ImageTracking

result = evaluate_flow_error(ground_truth_flow, estimated_flow, EndpointError())

imshow(result)
```

# References
[1] S. Baker, D. Scharstein, JP Lewis, S. Roth, M.J. Black, and R. Szeliski. A database and
evaluation methodology for optical flow. International Journal of Computer Vision, 92(1):1–31, 2011.
"""
function evaluate_flow_error(ground_truth_flow::Array{SVector{2, Float64}, 2}, estimated_flow::Array{SVector{2, Float64}, 2}, error_type::EndpointError)
    result = similar(Array{Float64, 2}, axes(estimated_flow))
    for i in CartesianIndices(estimated_flow)
        p1 = estimated_flow[i]
        p2 = ground_truth_flow[i]

        if is_flow_known(p1) && is_flow_known(p2)
            δ = p1 - p2
            result[i] = sqrt(sum(δ.^2))
        else
            result[i] = NaN
        end
    end

    return result
end

"""
```
evaluate_flow_error(ground_truth_flow, estimated_flow, AngularError())
```

Returns a 2-Dimensional array that matches the dimensions of the input flow vector and
that depicts the angle error between the estimated flow and the ground truth flow.

# Details
If the estimated flow at a point is (u0, v0) and ground truth flow is (u1, v1), it
calculates the angle between (u0, v0, 1) and (u1, v1, 1) vectors as measure for error.

# Arguments
The flow parameters needs to be two-dimensional arrays of length-2 vectors (of type SVector)
which represent the displacement of each pixel.

# Example

Compute the angle error between two flows.
```julia

using ImageTracking

result = evaluate_flow_error(ground_truth_flow, estimated_flow, AngularError())

imshow(result)
```

# References
[1] S. Baker, D. Scharstein, JP Lewis, S. Roth, M.J. Black, and R. Szeliski. A database and
evaluation methodology for optical flow. International Journal of Computer Vision, 92(1):1–31, 2011.
"""
function evaluate_flow_error(ground_truth_flow::Array{SVector{2, Float64}, 2}, estimated_flow::Array{SVector{2, Float64}, 2}, error_type::AngularError)
    result = similar(Array{Float64, 2}, axes(estimated_flow))
    for i in CartesianIndices(estimated_flow)
        p1 = ground_truth_flow[i]
        p2 = estimated_flow[i]
        if is_flow_known(p2) && is_flow_known(p1)
            p1 = push(p1, 1.0)
            p2 = push(p2, 1.0)
            cosine = dot(p1, p2)/(norm(p1)*norm(p2))
            if cosine > 1
                result[i] = 0
            else
                result[i] = acos(cosine)
            end
            if result[i] <= 1e-5
                result[i] = 0
            end
        else
            result[i] = NaN
        end
    end

    return result
end

"""
```
calculate_statistics(error)
```

Quantifies the flow error in terms of a mean and standard deviation as well as two additional
statistics that measure "robustness".

# Details
RX and AX are the robustness statistics. RX denotes the percentage of pixels that have
an error measure over X. AX denotes the accuracy of the error measure at the Xth percentile.

# Arguements
The error parameters needs to be two-dimensional arrays of length-2 vectors (of type SVector)
which represent the error between two flows.

# Example

Compute the Mean, SD, RX, AX stats of the flow error calculated using endpoint error method.
```julia

using ImageTracking

error = evaluate_flow_error(ground_truth_flow, estimated_flow, EndpointError())
mean, sd, rx, ax = calculate_statistics(error)
```

Compute the mean, SD, RX, AX stats of the flow error calculated using angula error method.
```julia

using ImageTracking

error = evaluate_flow_error(ground_truth_flow, estimated_flow, AngularError())
mean, sd, rx, ax = calculate_statistics(error)
```

# References
[1] S. Baker, D. Scharstein, JP Lewis, S. Roth, M.J. Black, and R. Szeliski. A database and
evaluation methodology for optical flow. International Journal of Computer Vision, 92(1):1–31, 2011.
"""
function calculate_statistics(error::Array{Float64, 2})
    R_thresholds = [0.5, 1.0, 2.0, 3.0, 5.0]
    A_thresholds = [0.5, 0.75, 0.95]

    # mean and standard deviations
    mean_value = mean(skipmissing(error))
    std_value = std(skipmissing(error); mean=mean_value)

    # RX stats
    RX_count = Dict()
    for rx in R_thresholds
        count_above_thresh = count(x->(x > rx), error)
        RX_count[rx] = count_above_thresh / (size(error)[1]*size(error)[2])
    end

    # AX stats
    AX_count = Dict()
    edges, counts = imhist(error)
    max_val = maximum(counts)
    total_pix = length(error)
    pixel_cumulative_count = cumsum(counts)
    for ax in A_thresholds
        cutoff = floor(ax*3.0 + 0.5)
        count_below_thresh = count(x->(x < cutoff), pixel_cumulative_count)
        AX_count[ax] = count_below_thresh / (length(edges)*max_val)
    end

    return mean_value, std_value, RX_count, AX_count
end


export logsumexp!, softmax3!, logsumexp_and_softmax!

"""
    sumprod(f, x, y, u, v) = sum(x .* y .* f.(u, v))
"""
function sumprod(f::Function, x::AbstractArray{T}, y::AbstractArray{T}, u::AbstractArray{T}, v::AbstractArray{T}) where {T<:Real}
    length(x)==length(y)==length(u)==length(v) || throw(DimensionMismatch())
    s = zero(T)
    @inbounds @simd for i in eachindex(x)
        s += x[i] * y[i] * f(u[i], v[i])
    end
    return s
end

"""
    sumprod3(x, y, z)

sumprod3 = sum( x .* y .* z)
"""
function sumprod3(x::AbstractArray{T}, y::AbstractArray{T}, z::AbstractArray{T}) where {T<:Real}
    length(x)==length(y)==length(z) || throw(DimensionMismatch())
    s = zero(T)
    @inbounds @simd for i in eachindex(x)
        s += x[i] * y[i] * z[i]
    end
    return s
end

function sumprod!(red::AbstractArray3, big::AbstractArray4, small::AbstractArray3)
     nz, nψ, nk, nd = size(big)
    (nz, nψ,     nd) == size(small) || throw(DimensionMismatch())
    (nz, nψ, nk,   ) == size(red)   || throw(DimensionMismatch())

    # second set w/ plus equals
    @inbounds for d in 1:nd, k in 1:nk
        if d == 1
            @views red[:,:,k] .=  small[:,:,d] .* big[:,:,k,d]
        else
            @views red[:,:,k] .+= small[:,:,d] .* big[:,:,k,d]
        end
    end
end

function sumprod!(red::AbstractMatrix, big::AbstractArray3, small::AbstractArray3)
     nz, nψ, nd  = size(big)
    (nz, nψ, nd) == size(small) || throw(DimensionMismatch())
    (nz, nψ)     == size(red) || throw(DimensionMismatch())

    fill!(red, 0)
    @inbounds for d in 1:nd
        @views red .+= small[:,:,d] .* big[:,:,d]
    end
end

# TODO would be good to evaluate whether I can do better with tmapreduce
# https://github.com/jw3126/ThreadingTools.jl
# https://discourse.julialang.org/t/how-to-speed-up-this-simple-code-multithreading-simd-inbounds/19681/39
# https://discourse.julialang.org/t/parallelizing-for-loop-in-the-computation-of-a-gradient/9154/7?u=tkoolen
# https://discourse.julialang.org/t/innefficient-paralellization-need-some-help-optimizing-a-simple-dot-product/9723/31
# https://discourse.julialang.org/t/for-loops-acceleration/24011/7

# tmapreduce((q,c,u,v) -> q*c*dψdρ(u,v), +, qm,cm,um,vm; init=zero(eltype(qm))) * βψ

# https://discourse.julialang.org/t/poor-performance-on-cluster-multithreading/12248/39

"""
    logsumexp!(r, x)

Compute `r` = softmax(x) and return `logsumexp(x)`.

Based on code from
https://arxiv.org/pdf/1412.8695.pdf eq 3.8 for p(y)
https://discourse.julialang.org/t/fast-logsumexp/22827/7?u=baggepinnen for stable logsumexp
"""
function logsumexp!(r::AbstractArray{T}, x::AbstractArray{T}) where {T}
    n = length(x)
    length(r) == n || throw(DimensionMismatch())
    isempty(x) && return -T(Inf)
    1 == stride1(r) == stride1(x) || throw(error("Arrays not strided"))

    u = maximum(x)                                       # max value used to re-center
    abs(u) == Inf && return any(isnan, x) ? T(NaN) : u   # check for non-finite values

    s = zero(T)
    for i = 1:n
        tmp = exp(x[i] - u)
        r[i] = tmp
        s += tmp
    end

    invs = inv(s)
    r .*= invs

    return log1p(s-1) + u
end

logsumexp!(x) = logsumexp!(x,x)

add_1_dim(x::AbstractArray) = reshape(x, size(x)..., 1)

"""
    softmax3!(q, lse, tmp, x, maxk)

set `q[i,j,k] = Pr( k | x[i,j,:])` for k in 1:maxk

Sets `lse = ∑_k exp(x[i,j,k])`

Uses temporary array `tmp`
"""
function softmax3!(q::AA, lse::A, tmpmax::A, x::AA, maxk=size(q, ndims(q)) ) where {T<:Real, A<:Array{T}, AA<:AbstractArray{T}}
        ndims(q) == 1+ndims(lse) || throw(DimensionMismatch())
        xsizes = size(x)
        xsizes == size(q) || throw(DimensionMismatch("size(x) = $(size(x)) but size(q) = $(size(q))"))
        nk = last(xsizes)
        for i = OneTo(ndims(lse))
            size(q,i) == size(lse,i) == size(tmpmax,i) || throw(DimensionMismatch("size(x) = $(size(x)),  size(lse) = $(size(lse)), and size(tmpmax) = $(size(tmpmax))"))
        end
        0 < maxk <= nk || throw(DomainError(maxk))
        1 == stride1(q) == stride1(x) || throw(error("Arrays not strided"))

        isempty(x) && throw(error("x empty"))
        check_finite(x)

        maximum!(add_1_dim(tmpmax), x)
        fill!(lse, zero(T))

        xx = reshape(x, :, nk)
        qq = reshape(q, :, nk)
        tmpmaxvec = vec(tmpmax)
        lsevec = vec(lse)

        for k in OneTo(maxk)
            qk = uview(qq, :, k)
            xk = uview(xx, :, k)
            for i in eachindex(lsevec)
                tmp = exp(xk[i] - tmpmaxvec[i])
                lsevec[i] += tmp
                qk[i] = tmp
            end
        end

        for k in maxk+1:nk
            qk = uview(qq, :, k)
            xk = uview(xx, :, k)
            for i in eachindex(lsevec)
                tmp = exp(xk[i] - tmpmaxvec[i])
                lsevec[i] += tmp
            end
        end

        qq[:,OneTo(maxk)] ./= vec(lse)
end

"use in vfit!"
function logsumexp_and_softmax!(lse, q, tmpmax, x, args...)
    softmax3!(q, lse, tmpmax, x, args...)
    lse .-= 1
    lse .= log1p.(lse) .+ tmpmax
end

function softmax3test!(q::AbstractArray3, x::AbstractArray3)
    ni, nj, nk = size(x)
    (ni, nj, nk) == size(q) || throw(DimensionMismatch())
    @inbounds for j = OneTo(nj)
        for i in OneTo(ni)
            @views softmax!(q[i,j,:], x[i,j,:])
        end
    end
end

function logsumexp3test!(x::AbstractArray3{T}) where {T}
    ni, nj, nk = size(x)
    lse = zeros(T,ni,nj)
    tmp = zeros(T,nk)
    # @inbounds
    for j = OneTo(nj)
        for i in OneTo(ni)
            tmp .= view(x, i, j, :)
            lse[i,j] = logsumexp!(tmp)
        end
    end
    return lse
end

@inline function sumsq(x, start::Integer, stop::Integer)
    @assert start >= 1
    @assert stop <= length(x)
    s = zero(eltype(x))
    @inbounds @simd for i in start:stop
        s += x[i]*x[i]
    end
    return s
end


"An `AbstractMappedGrid` represents the lazy application of a map to a grid."
abstract type AbstractMappedGrid{T,N} <: SimpleLazyGrid{T,N} end

const AbstractMappedGrid1d{T<:Number,N} = AbstractMappedGrid{T,N}

for op in (:length, :size, :eachindex, :indextype, :isperiodic)
	@eval $op(g::AbstractMappedGrid) = $op(supergrid(g))
end

covering(g::AbstractMappedGrid) = map_domain(forward_map(g), covering(supergrid(g)))

for op in (:minimum, :maximum)
	# TODO: the map could flip minimum and maximum...
	@eval $op(g::AbstractMappedGrid1d) = forward_map(g).($op(supergrid(g)))
end

resize(g::AbstractMappedGrid, n::Int) = map_grid(forward_map(g), resize(supergrid(g), n))

function rescale(g::AbstractGrid1d, a, b)
	m = mapto(covering(g), a..b)
	map_grid(m, g)
end

(→)(g::AbstractGrid1d, d::AbstractInterval) = rescale(g, infimum(d), supremum(d))

map_grid(grid::AbstractGrid, map) = (@warn "Interchange arguments of map_grid here"; map_grid(map, grid))
map_grid(map, grid::AbstractGrid) = map_grid1(map, grid)
map_grid1(map, grid::AbstractGrid) = map_grid2(map, grid)
map_grid2(map, grid) = MappedGrid(map, grid)

# some simplifications
map_grid1(map, g::AbstractMappedGrid) = map_grid(map∘forward_map(g), supergrid(g))
map_grid2(map::IdentityMap, grid) = grid

# Convenience function, similar to apply_map for Dictionary's
apply_map(grid::AbstractGrid, map::AbstractMap) = map_grid(map, grid)

forward_map(g::AbstractGrid, x...) = forward_map(g)(x...)
inverse_map(g::AbstractGrid, x...) = inverse_map(g)(x...)

unsafe_grid_getindex(g::AbstractMappedGrid, I...) = forward_map(g, supergrid(g, I...))

# Preserve tensor product structure
function rescale(g::ProductGrid, a::SVector{N}, b::SVector{N}) where {N}
	scaled_grids = [ rescale(component(g, i), a[i], b[i]) for i in 1:N]
	ProductGrid(scaled_grids...)
end

# Base.show(io::IO, mime::MIME"text/plain", g::AbstractMappedGrid) =
# 	composite_show(io, mime, g)
# Display.displaystencil(g::AbstractMappedGrid) =
#     DomainSets.map_stencil_broadcast(forward_map(g), supergrid(g))
# Display.object_parentheses(g::AbstractMappedGrid) =
#     Display.object_parentheses(forward_map(g))
# Display.stencil_parentheses(g::AbstractMappedGrid) =
#     Display.stencil_parentheses(forward_map(g))

"""
A MappedGrid consists of a grid and a map. Each grid point of the mapped grid
is the map of the corresponding point of the underlying grid.
"""
struct MappedGrid{T,N,M,G} <: AbstractMappedGrid{T,N}
	map		::	M
	grid	::	G
end

const MappedGrid1d{T<:Number,N,M,G} = MappedGrid{T,N,M,G}

MappedGrid(map, grid::AbstractGrid{T,N}) where {T,N} =
	MappedGrid{T,N}(map, grid)
MappedGrid(map::Map{S}, grid::AbstractGrid{T,N}) where {S,T,N} =
	MappedGrid{codomaintype(map),N}(map, grid)
MappedGrid{T,N}(map, grid) where {T,N} =
	MappedGrid{T,N,typeof(map),typeof(grid)}(map, grid)

MappedGrid(grid::AbstractGrid, map) = MappedGrid(map, grid)

forward_map(g::MappedGrid) = g.map
inverse_map(g::MappedGrid) = inverse(forward_map(g))
inverse_map(g::MappedGrid, x) = inverse(forward_map(g), x)


using MASW, JLD2, Plots
Plots.pyplot()

#===
# example to run masw(). Example and SampleData are based on matlab script written by John Schuh.

Matlab script: Copyright (c) 2018, John Schuh All rights reserved.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

•    John Schuh (2020). MASW Dispersion Curve (https://www.mathworks.com/matlabcentral/fileexchange/65138-masw-dispersion-curve), MATLAB Central File Exchange.
      Retrieved January 14, 2020.
===#

curdir = pwd()
fi = jldopen(curdir*"/../example/SampleData.jld", "r")
U = fi["suma_sg"]
close(fi)

fs=1000; # sampling frequency Hz
min_x = 1; # min offset [m] of reciever spread
max_x = 100; # max offset [m] of reciever spread
d_x = 2; # distance between receivers [m]

x = min_x:d_x:max_x; # spatial sampling vector (receiver position)
t=1/fs:1/fs:length(U[:,1])/fs;  # time vector in seconds

figdir = "./"
figname = "dispersionimage.png"

ct=300:1:1200;
freqlimits=(15.0, 80.0);

(f, ct, Udisp, DispersionCurve)=masw(U, x, t, ct, plotfreqlimits=freqlimits,figdir="./",figname="functiontest.png")


if lowercase(get(ENV, "CI", "false")) == "true"
    include("install_dependencies.jl")
end

using Test
using CausalityToolsBase
using StateSpaceReconstruction
using PerronFrobenius

# Individual constructors (i.e. rectangularinvariantmeasure(pts, ϵ))
include("test_rectangularmeasure_constructor.jl")
#include("test_inducedmeasure.jl")
#include("test_inducedmeasure_constructor.jl")
#include("test_averagemeasure.jl")
#include("test_averagemeasure_constructor.jl")

# Test the estimators (i.e. invariantmeasure(pts, binningscheme))
include("test_invariantmeasure_estimators.jl")

include("test_gridestimator.jl")
include("test_transferoperator_rectangular.jl")

#include("triangulations.jl")
include("test_triangulationestimators.jl")