using CornerSpaceRenorm, QuantumOptics
using LightGraphs, GraphPlot
using LinearAlgebra, IterativeSolvers, Random
using Test, InteractiveUtils

macro no_error(ex)
    quote
        try
            $(esc(ex))
            true
        catch
            false
        end
    end
end

@testset "CornerSpaceRenorm" begin
    @testset "Operators" begin
        lb = FockBasis(1)
        lH = randoperator(lb)
        tH = randoperator(lb)

        for pbc in [true, false]
            L = NdLattice((3,2,1);periodic=pbc)
            H1 = hamiltonian(L,lH,(1.,1.,1.),tH)
            H2 = hamiltonian(L,lH,(1.,tH))
            H3 = hamiltonian(L,lH,1.,tH)
            @test maximum(abs2.(H1.data .- H2.data)) ≈ 0. atol=eps(Float64)
            @test maximum(abs2.(H1.data .- H3.data)) ≈ 0. atol=eps(Float64)
        end
    end;

    @testset "Types" begin
        @testset "NdLattice" begin
            L = NdLattice((2,2))

            # Test method extensions
            # TO DO: extend tests to instances of all subtypes of Lattice
            @test eltype(L) == eltype(L.L)
            @test vertices(L) == vertices(L.L)
            @test edges(L) == edges(L.L)
            @test nv(L) == nv(L.L)
            @test ne(L) == ne(L.L)
            @test all([has_edge(L,v1,v2) == has_edge(L.L,v1,v2) for v1 in vertices(L), v2 in vertices(L)])
            @test all([has_edge(L,e) == has_edge(L.L,e) for e in edges(L)])
            @test all([has_vertex(L,v) == has_vertex(L.L,v) for v in collect(vertices(L)) ∪ collect(-10:-1)])
            @test all([inneighbors(L,v) == inneighbors(L.L,v) for v in vertices(L)])
            @test all([outneighbors(L,v) == outneighbors(L.L,v) for v in vertices(L)])
            @test is_directed(L) == false
            @test all(is_directed.(subtypes(Lattice)) .== false)

            function compare_lattices(L1,L2)
                @test typeof(L1) == typeof(L2)
                @test all([getfield(L1,f) == getfield(L2,f) for f in fieldnames(typeof(L1))])
            end

            # Test lattice merging
            L1 = NdLattice((2,3))
            L2 = NdLattice((2,1))
            L3_12 = union(L1,L2,2)
            L3_21 = union(L2,L1,2)
            compare_lattices(L3_12,L3_21)
            @test_throws DimensionMismatch union(L1,L2,1)

            L1 = NdLattice((3,2))
            L2 = NdLattice((1,2))
            L3_12 = union(L1,L2,1)
            L3_21 = union(L2,L1,1)
            @test typeof(L3_12) == typeof(L3_21)
            # TO DO: test boundaries
            @test LightGraphs.Experimental.could_have_isomorph(L3_12.L, L3_21.L)
            @test L3_12.shape[1] == L1.shape[1] + L2.shape[1]
            compare_lattices(L3_12, union(L1,L2,1))
            @test_throws DimensionMismatch union(L1,L2,2)

            Lpbc = NdLattice((2,2);periodic=true)
            Lobc = NdLattice((2,2);periodic=false)
            @test_throws AssertionError union(Lpbc,Lobc,1)
            @test_throws AssertionError union(Lpbc,Lobc,2)

            L = NdLattice((2,);periodic=true)
            LuL = union(L,L,1)
            L2 = NdLattice((4,);periodic=true)
            compare_lattices(LuL,L2)

            L = NdLattice((2,2);periodic=true)
            LuL = union(L,L,1)
            L2 = NdLattice((4,2);periodic=true)
            @test typeof(LuL) == typeof(L2)
            @test LightGraphs.Experimental.could_have_isomorph(LuL.L, L2.L)

            L = NdLattice((2,2);periodic=true)
            LuL = union(L,L,2)
            L2 = NdLattice((2,4);periodic=true)
            @test typeof(LuL) == typeof(L2)
            compare_lattices(LuL,L2)

            L = NdLattice((1,);periodic=true)
            for i in 1:3
                L = union(L,L,1)
            end
            compare_lattices(L, NdLattice((8,); periodic=true))

            # Test pbc_from_obc and obc_from_pbc
            Lo = NdLattice((5,5);periodic=false)
            Lp = NdLattice((5,5);periodic=true)
            compare_lattices(pbc_from_obc(Lo), Lp)
            compare_lattices(Lo, obc_from_pbc(Lp))

            Lo = NdLattice((1,1);periodic=false)
            Lo = union(Lo,Lo,1)
            Lo = union(Lo,Lo,1)
            Lo = union(Lo,Lo,1)
            Lo = union(Lo,Lo,2)
            Lo = union(Lo,Lo,2)
            Lo = union(Lo,Lo,2)
            Lp = NdLattice((8,8);periodic=true)
            compare_lattices(pbc_from_obc(Lo), Lp)
            compare_lattices(Lo, obc_from_pbc(Lp))

            # Dummy tests on plotting
            L = NdLattice((2,2);periodic=true)
            @test @no_error gplot(L)
            @test @no_error gplot(L, Float64.(collect(1:nv(L))), Float64.(collect(1:nv(L))))
            x = gplot(L)
            y = gplot(L.L)
            @test typeof(x) == typeof(y)
            x = gplot(L, Float64.(collect(1:nv(L))), Float64.(collect(1:nv(L))))
            y = gplot(L.L, Float64.(collect(1:nv(L))), Float64.(collect(1:nv(L))))
            @test all([typeof(getfield(x,f)) == typeof(getfield(y,f)) for f in fieldnames(typeof(x))])

            # Test pretty-printing
            L = NdLattice((2,2);periodic=true)
            @test @no_error repr(L)
            @test @no_error repr(MIME("text/plain"),L)
            @test @no_error repr(MIME("application/prs.juno.inline"),L)
            L = NdLattice((2,2);periodic=false)
            @test @no_error repr(L)
            @test @no_error repr(MIME("text/plain"),L)
            @test @no_error repr(MIME("application/prs.juno.inline"),L)
        end;

        @testset "NdSystem" begin
            # Local basis
            b_spin = SpinBasis(1//2)

            # Build some operators in the local basis
            sm = sigmam(b_spin)
            sp = sigmap(b_spin)
            sz = sigmaz(b_spin)
            sx = sigmax(b_spin)
            sy = sigmay(b_spin)

            # Build a set of local observables in the local basis
            lobs = Dict("sigmax"=>sx,"sigmay"=>sy,"sigmaz"=>sz)

            g = 1.
            gamma = 1.
            V = 2/2.

            # Test assertions in constructors
            L = NdLattice((4,))
            H = hamiltonian(L, g/2 * sx, V/4, sz)
            J = dissipators(L, [sqrt(2gamma) * sm])
            # Incompatible lattice
            L_err = NdLattice((5,))
            @test_throws AssertionError NdSystem(L_err, H, (V/4,), sz, J, lobs)
            # Number of tunnelling rates incompatible with number of dimensions
            @test_throws AssertionError NdSystem(L, H, (V/4,V/4), sz, J, lobs)
            # Incompatible operators
            @test_throws ErrorException NdSystem(L, H, (V/4,), destroy(FockBasis(1)), J, lobs)

            # Test constructor with missing observables
            s = NdSystem(L, H, (V/4,), sz, J)
            @test eltype(s.observables) == Dict{String,typeof(H)}

            # Broken: cannot construct a system with 1 site.

            function compare_systems(s1::AbstractSystem,s2::AbstractSystem)
                for f in fieldnames(typeof(s1))
                    if f == :lattice
                        L1 = getfield(s1,f)
                        L2 = getfield(s2,f)
                        for ff in fieldnames(typeof(L1))
                            @test getfield(L1,ff) == getfield(L2,ff)
                        end
                    else
                        @test getfield(s1,f) == getfield(s2,f)
                    end
                end
            end
            # Simpler constructor for homogeneous coupling
            L = NdLattice((2,1,1,1))
            H = hamiltonian(L, g/2 * sx, V/4, sz)
            J = dissipators(L, [sqrt(2gamma) * sm])
            s1 = NdSystem(L, H, V/4, sz, J, lobs)
            s2 = NdSystem(L, H, Tuple([V/4 for i in 1:4]), sz, J, lobs)
            compare_systems(s1,s2)
            s1 = NdSystem(L, H, V/4, sz, J)
            s2 = NdSystem(L, H, Tuple([V/4 for i in 1:4]), sz, J)
            compare_systems(s1,s2)

            # Test new constructor
            L = NdLattice((2,3); periodic=true)
            H = hamiltonian(L, g/2 * sx, V/4/2., sz)
            J = dissipators(L, [sqrt(gamma) * sm])
            s1 = NdSystem(L, H, (V/4,V/4), (sz,sz), J, lobs)
            s2 = NdSystem(L, H, V/4, sz, J, lobs)
            compare_systems(s1,s2)
            s3 = NdSystem(L, H, (V/4, V/4), sz, J, lobs)
            compare_systems(s1,s3)
            s4 = NdSystem(L, H, V/4, (sz,sz), J, lobs)
            compare_systems(s1,s4)

            # Test multiple coupling
            @test @no_error s1 = NdSystem(L, H, V/4, ([sx,sz],[sx,sz]), J, lobs)
            @test @no_error s2 = NdSystem(L, H, (V/4,V/4), ([sx,sz],[sx,sz]), J, lobs)
            compare_systems(s1,s2)

            # Test construction and merging
            # 1D
            for pbc in [true, false]
                L = NdLattice((4,); periodic=pbc)
                H = hamiltonian(L, g/2 * sx, V/4, sz)
                J = dissipators(L, [sqrt(2gamma) * sm])
                s = NdSystem(L, H, (V/4,), sz, J, lobs)
                ρ = steadystate.master(s)[2][end]
                sUs = merge(s,s,1,ρ,ρ,256)

                L2 = NdLattice((8,); periodic=pbc)
                H2 = hamiltonian(L2, g/2 * sx, V/4, sz)
                J2 = dissipators(L2, [sqrt(2gamma) * sm])
                s2 = NdSystem(L2, H2, (V/4,), sz, J2, lobs)

                @test maximum(abs2.(eigvals(sUs.H.data) .- eigvals(Matrix(s2.H.data)))) ≈ 0. atol=eps(Float64)
                @test maximum([maximum(abs2.(eigvals(sUs.J[i].data'sUs.J[i].data) .- eigvals(Matrix(s2.J[i].data's2.J[i].data)))) for i in 1:length(sUs.J)]) ≈ 0. atol=eps(Float64)

                d = 1
                @test maximum([maximum(abs2.(eigvals(sUs.Htint[d][i].data) .- eigvals(Matrix(s2.Htint[d][i].data)))) for i in 1:length(sUs.Htint[d])]) ≈ 0. atol=eps(Float64)
                @test maximum([maximum(abs2.(eigvals(sUs.Htext[d][i].data) .- eigvals(Matrix(s2.Htext[d][i].data)))) for i in 1:length(sUs.Htext[d])]) ≈ 0. atol=eps(Float64)
            end

            # 2D
            for pbc in [true, false]
                L = NdLattice((2,2); periodic=pbc)
                H = hamiltonian(L, g/2 * sx, V/4, sz)
                J = dissipators(L, [sqrt(2gamma) * sm])
                s = NdSystem(L, H, (V/4, V/4), sz, J, lobs)
                ρ = steadystate.master(s)[2][end]
                sUs = merge(s,s,2,ρ,ρ,256)

                L2 = NdLattice((2,4); periodic=pbc)
                H2 = hamiltonian(L2, g/2 * sx, V/4, sz)
                J2 = dissipators(L2, [sqrt(2gamma) * sm])
                s2 = NdSystem(L2, H2, (V/4, V/4), sz, J2, lobs)

                @test maximum(abs2.(eigvals(sUs.H.data) .- eigvals(Matrix(s2.H.data)))) ≈ 0. atol=eps(Float64)
                @test maximum([maximum(abs2.(eigvals(sUs.J[i].data'sUs.J[i].data) .- eigvals(Matrix(s2.J[i].data's2.J[i].data)))) for i in 1:length(sUs.J)]) ≈ 0. atol=eps(Float64)

                for d in 1:2
                    @test maximum([maximum(abs2.(eigvals(sUs.Htint[d][i].data) .- eigvals(Matrix(s2.Htint[d][i].data)))) for i in 1:length(sUs.Htint[d])]) ≈ 0. atol=eps(Float64)
                    @test maximum([maximum(abs2.(eigvals(sUs.Htext[d][i].data) .- eigvals(Matrix(s2.Htext[d][i].data)))) for i in 1:length(sUs.Htext[d])]) ≈ 0. atol=eps(Float64)
                end

                # Test pretty-printing
                @test @no_error repr(s)
                @test @no_error repr(MIME("text/plain"),s)
                @test @no_error repr(MIME("application/prs.juno.inline"),s)
                @test @no_error repr(sUs)
                @test @no_error repr(MIME("text/plain"),sUs)
                @test @no_error repr(MIME("application/prs.juno.inline"),sUs)
            end
            # TO DO: Test construction and merging of generic lattices

            # Dummy tests on plotting
            L = NdLattice((4,))
            H = hamiltonian(L, g/2 * sx, V/4, sz)
            J = dissipators(L, [sqrt(2gamma) * sm])
            s = NdSystem(L, H, (V/4,), sz, J)
            @test @no_error gplot(s);
            @test @no_error gplot(s, Float64.(collect(1:nv(L))), Float64.(collect(1:nv(L))));
            @test @no_error plot_system(s);
            @test @no_error plot_system(s, Float64.(collect(1:nv(L))), Float64.(collect(1:nv(L))));
            x = gplot(s);
            y = gplot(L.L);
            @test typeof(x) == typeof(y)
            x = plot_system(s);
            @test typeof(x) == typeof(y)
            x = gplot(s, Float64.(collect(1:nv(L))), Float64.(collect(1:nv(L))));
            y = gplot(L.L, Float64.(collect(1:nv(L))), Float64.(collect(1:nv(L))));
            @test all([typeof(getfield(x,f)) == typeof(getfield(y,f)) for f in fieldnames(typeof(x))])
            x = plot_system(s, Float64.(collect(1:nv(L))), Float64.(collect(1:nv(L))));
            @test all([typeof(getfield(x,f)) == typeof(getfield(y,f)) for f in fieldnames(typeof(x))])
        end;

        @testset "CornerBasis" begin
            # Test constructors
            b1 = SpinBasis(1//2)
            b2 = CompositeBasis([b1 for i in 1:10])
            bC = CornerBasis(b1, b2, 10)
            @test bC == CornerBasis([10],10,[[2], fill(2,10)])
            bC2 = CornerBasis(bC, bC, 10)
            @test bC2 == CornerBasis([10],10,vcat(bC.local_shapes, bC.local_shapes))

            # Test Base.== overload
            bC = CornerBasis(10)
            @test bC == deepcopy(bC)
            @test bC !== SpinBasis(1//2)
            @test bC !== FockBasis(1)
            @test bC !== GenericBasis(1)
            bC2 = CornerBasis([11],bC.M,bC.local_shapes)
            @test bC !== bC2
            bC2 = CornerBasis(bC.shape,bC.M+1,bC.local_shapes)
            @test bC !== bC2
            bC2 = CornerBasis(bC.shape,bC.M,vcat(bC.local_shapes, [[1]]))
            @test bC !== bC2

            # Test pretty-printing
            @test @no_error repr(bC)
            @test @no_error repr(MIME("text/plain"),bC)
            @test @no_error repr(MIME("application/prs.juno.inline"),bC)
        end;
    end;

    @testset "Corner algorithm" begin
        @testset "Core" begin
            # Test orthonormalisation
            N = 10
            U = diagm([i=>ones(ComplexF64,N-i) for i in 0:N]...)
            @test maximum(abs2.(CornerSpaceRenorm.cgs(U) .- Diagonal(ones(ComplexF64,N)))) ≈ 0. atol=eps(Float64)
            @test maximum(abs2.(CornerSpaceRenorm.mgs(U) .- Diagonal(ones(ComplexF64,N)))) ≈ 0. atol=eps(Float64)
            V = copy(U); CornerSpaceRenorm.cgs!(V)
            @test maximum(abs2.(V .- Diagonal(ones(ComplexF64,N)))) ≈ 0. atol=eps(Float64)
            V = copy(U); CornerSpaceRenorm.mgs!(V)
            @test maximum(abs2.(V .- Diagonal(ones(ComplexF64,N)))) ≈ 0. atol=eps(Float64)

            # Test maxk
            x = randperm(1000)
            idcs, y = CornerSpaceRenorm.maxk(x,100)
            # Test that indices match sorted values
            @test all([x[idcs[i]] == y[i] for i in 1:100])
            # Test that k sorted values are the largest
            @test all([val == y[i] for (i,val) in enumerate(1000:-1:901)])

            # Test max_prod_pairs
            x1 = 1:50; x2 = 1:50
            idcs, pairs = CornerSpaceRenorm.max_prod_pairs(x1,x2,10)
            @test all([x1[idx[1]] * x2[idx[2]] for idx in idcs] .== pairs)
            @test all(sort([x1[i] * x2[j] for i in 1:50, j in 1:50][:],rev=true)[1:10] .== pairs)

            # Test hermitianize
            b = GenericBasis(50)
            Op = randoperator(b)
            OpH = hermitianize(Op)
            @test ishermitian(OpH.data)
            @test maximum(abs2.(OpH.data .- (Op + dagger(Op)).data/2.)) ≈ 0. atol=eps(Float64)
            Op = randoperator(b);
            OpH = copy(Op);
            hermitianize!(OpH)
            @test ishermitian(OpH.data)
            @test maximum(abs2.(OpH.data .- (Op + dagger(Op)).data/2.)) ≈ 0. atol=eps(Float64)
        end;

        @testset "Merging" begin
            @testset "NdSystem" begin
                # Local basis
                b_spin = SpinBasis(1//2)

                # Build some operators in the local basis
                sm = sigmam(b_spin)
                sp = sigmap(b_spin)
                sz = sigmaz(b_spin)
                sx = sigmax(b_spin)
                sy = sigmay(b_spin)

                # Build a set of local observables in the local basis
                lobs = Dict("sigmax"=>sx,"sigmay"=>sy,"sigmaz"=>sz)

                g = 1.
                gamma = 1.
                V = 2/2.

                L1 = NdLattice((2,1))
                H1 = hamiltonian(L1, g/2 * sx, V/4, sz)
                J1 = dissipators(L1, [sqrt(2gamma) * sm])
                s1 = NdSystem(L1, H1, (V/4,V/4), sz, J1, lobs)
                ρ1 = steadystate.master(s1.H,s1.J)[2][end]
                L2 = NdLattice((1,2))
                H2 = hamiltonian(L2, g/2 * sx, V/4, sz)
                J2 = dissipators(L2, [sqrt(2gamma) * sm])
                s2 = NdSystem(L2, H2, (V/4,V/4), sz, J2, lobs)
                ρ2 = steadystate.master(s2.H,s2.J)[2][end]

                @test_throws DimensionMismatch merge(s1, s2, 1, ρ1, ρ2, 2)
                @test_throws DimensionMismatch merge(s1, s2, 2, ρ1, ρ2, 2)

                L1_pbc = NdLattice((2,1);periodic=true)
                H1_pbc = hamiltonian(L1_pbc, g/2 * sx, V/4, sz)
                J1_pbc = dissipators(L1_pbc, [sqrt(2gamma) * sm])
                s1_pbc = NdSystem(L1_pbc, H1_pbc, (V/4,V/4), sz, J1_pbc, lobs)
                ρ1_pbc = steadystate.master(s1_pbc.H,s1_pbc.J)[2][end]

                @test_throws AssertionError merge(s1,s1_pbc,1,ρ1,ρ1_pbc,2)

                # Full space corner vs exact system
                # 1D
                for pbc in [true, false]
                    L1 = NdLattice((2,); periodic=pbc)
                    H1 = hamiltonian(L1, g/2 * sx, V/4, sz)
                    J1 = dissipators(L1, [sqrt(2gamma) * sm])
                    s1 = NdSystem(L1, H1, (V/4,), sz, J1, lobs)
                    ρ1 = steadystate.master(s1.H,s1.J)[2][end]
                    s1 = merge(s1,s1,1,ρ1,ρ1,16)
                    ρ1 = steadystate.master(s1.H,s1.J;tol=1e-6)[2][end]

                    L2 = NdLattice((4,); periodic=pbc)
                    H2 = hamiltonian(L2, g/2 * sx, V/4, sz)
                    J2 = dissipators(L2, [sqrt(2gamma) * sm])
                    s2 = NdSystem(L2, H2, (V/4,), sz, J2, lobs)
                    ρ2 = steadystate.master(s2.H,s2.J;tol=1e-6)[2][end]

                    @test maximum(abs2.(eigvals(Matrix(s1.H.data)) .- eigvals(Matrix(s2.H.data)))) ≈ 0. atol=eps(Float64)
                    @test maximum([maximum(abs2.(eigvals(Matrix(s1.J[i].data's1.J[i].data)) .- eigvals(Matrix(s2.J[i].data's2.J[i].data)))) for i in 1:length(s1.J)]) ≈ 0. atol=eps(Float64)
                    @test maximum(abs2.(eigvals(ρ1.data) .- eigvals(ρ2.data))) ≈ 0. atol=1e-5
                end

                # 2D
                for pbc in [true, false]
                    L1 = NdLattice((2,1); periodic=pbc)
                    H1 = hamiltonian(L1, g/2 * sx, V/4, sz)
                    J1 = dissipators(L1, [sqrt(2gamma) * sm])
                    s1 = NdSystem(L1, H1, (V/4,V/4), sz, J1, lobs)
                    ρ1 = steadystate.master(s1.H,s1.J)[2][end]
                    s1 = merge(s1,s1,1,ρ1,ρ1,16)
                    ρ1 = steadystate.master(s1.H,s1.J;tol=1e-5)[2][end]

                    L2 = NdLattice((4,1); periodic=pbc)
                    H2 = hamiltonian(L2, g/2 * sx, V/4, sz)
                    J2 = dissipators(L2, [sqrt(2gamma) * sm])
                    s2 = NdSystem(L2, H2, (V/4,V/4), sz, J2, lobs)
                    ρ2 = steadystate.master(s2.H,s2.J;tol=1e-5)[2][end]

                    @test maximum(abs2.(eigvals(Matrix(s1.H.data)) .- eigvals(Matrix(s2.H.data)))) ≈ 0. atol=eps(Float64)
                    @test maximum([maximum(abs2.(eigvals(Matrix(s1.J[i].data's1.J[i].data)) .- eigvals(Matrix(s2.J[i].data's2.J[i].data)))) for i in 1:length(s1.J)]) ≈ 0. atol=eps(Float64)
                    @test maximum(abs2.(eigvals(ρ1.data) .- eigvals(ρ2.data))) ≈ 0. atol=1e-5
                end

                # Check returning the product density matrix when merging
                L1 = NdLattice((2,))
                H1 = hamiltonian(L1, g/2 * sx, V/4, sz)
                J1 = dissipators(L1, [sqrt(2gamma) * sm])
                s1 = NdSystem(L1, H1, (V/4,), sz, J1, lobs)
                ρ1 = steadystate.master(s1.H,s1.J)[2][end]
                s1, ρ02 = merge(s1,s1,1,ρ1,ρ1,16;return_dm=true)
                @test tr(ρ02) ≈ 1.0 atol=1e-9
                @test opnorm(ρ02.data .- ρ02.data) ≈ 0.0 atol=1e-9
                @test all(eigvals(hermitianize(ρ02).data) .> 0.)
                tss = steadystate.master(s1.H,s1.J;tol=1e-5)[1][end]
                tss_init = steadystate.master(s1.H,s1.J;rho0=ρ02,tol=1e-5)[1][end]
                @test tss_init < tss

                # Merging different lattices
                Ls = NdLattice((3,))
                Hs = hamiltonian(Ls, g/2 * sx, V/4, sz)
                Js = dissipators(Ls, [sqrt(2gamma) * sm])
                ss = NdSystem(Ls, Hs, V/4, sz, Js, lobs)
                ρs = steadystate.master(ss)[2][end]

                t = 1.
                Δ = 1.
                κ = 1.
                F = 5e-1 * κ

                lb = FockBasis(1)
                Lb = NdLattice((3,),periodic=false)
                Hb = hamiltonian(Lb, -Δ * number(lb) + F * (create(lb) + destroy(lb)), t/2., destroy(lb))
                Jb = dissipators(Lb,[sqrt(κ) * destroy(lb)])
                sb = NdSystem(Lb,Hb,(t/2.,),destroy(lb),Jb)
                ρb = steadystate.master(sb)[2][end]

                @test_logs (:info,"The two input systems do not have equal tunnelling rates along all directions.") s_sb = merge(ss,sb,1,ρs,ρb,10)
            end;
        end;
    end;

    @testset "steadystate" begin
        # Local basis
        b_spin = SpinBasis(1//2)

        # Build some operators in the local basis
        sm = sigmam(b_spin)
        sp = sigmap(b_spin)
        sz = sigmaz(b_spin)
        sx = sigmax(b_spin)
        sy = sigmay(b_spin)

        # Build a set of local observables in the local basis
        lobs = Dict("sigmax"=>sx,"sigmay"=>sy,"sigmaz"=>sz)

        g = 1.
        gamma = 1.
        V = 2/2.

        L1 = NdLattice((6,))
        H1 = hamiltonian(L1, g/2 * sx, V/4, sz)
        J1 = dissipators(L1, [sqrt(2gamma) * sm])
        s1 = NdSystem(L1, H1, (V/4,), sz, J1, lobs)
        ρ1 = steadystate.master(s1.H,s1.J; tol=1e-5)[2][end]
        ρ2 = steadystate_bicg(s1)
        @test fidelity(ρ1, ρ2) ≈ 1. atol=1e-5
        ρ3, log = steadystate_bicg(s1; log=true)
        @test typeof(log) == ConvergenceHistory{true,Nothing}
        ρ4 = CornerSpaceRenorm.steadystate_bicg_LtL(s1)
        @test fidelity(ρ1, ρ4) ≈ 1. atol=1e-5
        ρ5, log = steadystate_bicg!(ρ3, s1; log=true)
        @test log.iters == 0
        ρ, log = steadystate_bicg(s1; log=true)
        @test typeof(log) == ConvergenceHistory{true,Nothing}
        ρ, log = CornerSpaceRenorm.steadystate_bicg_LtL(s1; log=true)
        @test typeof(log) == ConvergenceHistory{true,Nothing}
    end;

end;


using Variography
using Meshes
using GeoStatsBase
using GeoStatsImages
using Unitful
using Distances
using LinearAlgebra
using CSV, DelimitedFiles
using Plots; gr(size=(600,400))
using ReferenceTests, ImageIO
using Test, Random
import CoDa: Composition

# workaround for GR warnings
ENV["GKSwstype"] = "100"

# environment settings
isCI = "CI" ∈ keys(ENV)
islinux = Sys.islinux()
visualtests = !isCI || (isCI && islinux)
datadir = joinpath(@__DIR__,"data")

# list of tests
testfiles = [
  "empirical.jl",
  "partition.jl",
  "varioplane.jl",
  "theoretical.jl",
  "nesting.jl",
  "sampling.jl",
  "pairwise.jl",
  "fitting.jl"
]

@testset "Variography.jl" begin
  for testfile in testfiles
    include(testfile)
  end
end


export AbstractSequence, Sequence, sequence, generator

abstract type AbstractSequence{T,D} end

struct Generator{D}
    generator::D
end

(G::Generator)(i) = G.generator(i)
Base.getindex(G::Generator, i) = G(i)

mutable struct Sequence{T,D} <: AbstractSequence{T,D}
    sequence::Vector{T}
    generator::Generator{D}
end

Sequence(generator::D) where D<:Base.Callable = Sequence([generator(1)],Generator(generator))

sequence(S::AbstractSequence) = S.sequence
generator(S::AbstractSequence) = S.generator

Base.eltype(::AbstractSequence{T}) where T = T
Base.size(::AbstractSequence) = (Inf,)
Base.length(::AbstractSequence) = Inf

Base.getindex(S::AbstractSequence,i::Integer) = (generate!(S,i); sequence(S)[i])
Base.getindex(S::AbstractSequence,ir::AbstractRange) = (generate!(S,maximum(ir)); sequence(S)[ir])


using Test, DataLoaders

include("dataset.jl")

@testset "DataLoder" begin
  batchsize = 16
  dataset = getdataset((10,10), 100)
  loader  = DataLoader(dataset, batchsize, ntasks = 8, shuffle=true, droplast=false)

  @test length(loader) == round(Int, 100 / 16, RoundDown)

  (x, y), index = iterate(loader)

  @test size(x) == (10, 10, batchsize)
  @test size(y) == (batchsize,)
  @test index   == batchsize + 1
end


module State

using PyCall
import ..@pyfunc

#qulacs = PyNULL()

state = PyNULL()
const state_functions = [
    :inner_product,
    :tensor_product,
    :permutate_qubit,
    :drop_qubit,
    :partial_trace,
]

#=
for func in state_functions
    @eval begin
        function $(func)(args...; kwargs...)
            qulacs.state.$(func)(args...; kwargs...)
        end
        export $(func)
    end
end
=#

for func in state_functions
    @eval begin
        @pyfunc state $(func)
        export $(func)
    end
end

function __init__()
    #copy!(qulacs, pyimport("qulacs"))
    copy!(state, pyimport("qulacs.state"))
end

end # module State


# Feynman/diagrams.jl
#
# MIT License
#
# Copyright (c) 2019 Brandon Gomes
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import LightGraphs: AbstractEdge
import MetaGraphs: MetaGraph, MetaDiGraph

include("utils/graph_wrapper.jl")

export Topology, Diagram, topologies

"""
"""
struct Topology{I} <: GraphWrapper{I}
    graph::MetaGraph{I}
end

"""
"""
graph(t::Topology) = t.graph

"""
"""
function (::Type{Topology})(mdg::MetaDiGraph)
    Topology(MetaGraph(mdg))
end

"""
"""
function topologies(loops::Integer, external::Integer)
    error("Not Implemented")
end

"""
"""
function topologies(loops::Integer, incoming::Integer, outgoing::Integer)
    topologies(loops, incoming + outgoing)
end

"""
"""
function topologies(loops::Integer, inout::Pair{<: Integer, <: Integer})
    topologies(loops, first(inout), last(inout))
end

"""
"""
struct Diagram{I} <: AbstractGraph{I}
    graph::MetaDiGraph{I}
end

"""
"""
graph(t::Diagram) = t.graph


const OneDimensionalEnsembleModel = HydrostaticFreeSurfaceModel{TS, E, A, S, <:OneDimensionalEnsembleGrid, T, V, B, R, F, P, U, C, Φ, K, AF} where {TS, E, A, S, T, V, B, R, F, P, U, C, Φ, K, AF}

ensemble_size(model::OneDimensionalEnsembleModel) = model.grid.Nx
batch_size(model::OneDimensionalEnsembleModel) = model.grid.Ny

"""
    set!(model::OneDimensionalEnsembleModel,
         observations::OneDimensionalTimeSeriesBatch, time_index)

Set columns of each field in `model` to the model profile columns in `observations`,
where every field column in `model` that corresponds to the ith `OneDimensionalTimeSeries` object in `observations`
is set to the field column in `observations[i]` at time index `time_indices[i]`.
"""
function set!(model::OneDimensionalEnsembleModel,
              observations::OneDimensionalTimeSeriesBatch, time_indices::Vector)

    ensemble(x) = column_ensemble_interior(observations, x, time_indices, model.grid.Nx)

    set!(model, b = ensemble(:b),
                u = ensemble(:u),
                v = ensemble(:v),
                e = ensemble(:e)
        )
end

set!(model::OneDimensionalEnsembleModel, observations::OneDimensionalTimeSeriesBatch, time_index) = set!(model, observations, [time_index for i in observations])

"""
    initialize_forward_run!(model, observations::OneDimensionalTimeSeriesBatch, params, time_indices::Vector)

Set columns of each field in `model` to the corresponding profile columns in `observations`,
where every field column in `model` that corresponds to the ith `OneDimensionalTimeSeries` object in `observations`
is set to the field column in `observations[i]` at `time_indices[i]`.
"""
function initialize_forward_run!(model::OneDimensionalEnsembleModel, observations::OneDimensionalTimeSeriesBatch, params, time_indices::Vector)
    set!(model, params)
    set!(model, observations, time_indices)
    model.clock.time = 0.0
    model.clock.iteration = 0
    return nothing
end


using BenchmarkTools
using JLD
using Test
using MIToS
using MIToS.MSA
using MIToS.Information
using MIToS.PDB
using MIToS.SIFTS
using MIToS.Pfam
using MIToS.Utils

# --------------------------------------------------------------------------- #
# MSA module benchmarks
const msa = BenchmarkGroup()

## Set up
### Files
const msafile_long_sth_gz = "../../data/PF00089.stockholm.gz"
const msafile_long_sth    = "../../data/PF00089.sth"
const msafile_long_fas_gz = "../../data/PF00089.fasta.gz"
const msafile_long_fas    = "../../data/PF00089.fasta"
const msafile_wide_sth_gz = "../../data/PF16957.stockholm.gz"
const msafile_wide_sth    = "../../data/PF16957.sth"
const msafile_wide_fas_gz = "../../data/PF16957.fasta.gz"
const msafile_wide_fas    = "../../data/PF16957.fasta"
### MSAs
const msa_long    = read(msafile_long_sth   , Stockholm)
const msa_wide    = read(msafile_wide_sth   , Stockholm)

#### Parse benchmarks

##### input
msa["input"] = BenchmarkGroup()

for (file,gzipped,shape,format) in [(msafile_long_sth_gz, "gzipped", "long", MIToS.MSA.Stockholm),
                                    (msafile_long_sth,  "ungzipped", "long", MIToS.MSA.Stockholm),
                                    (msafile_long_fas_gz, "gzipped", "long", MIToS.MSA.FASTA),
                                    (msafile_long_fas,  "ungzipped", "long", MIToS.MSA.FASTA),
                                    (msafile_wide_sth_gz, "gzipped", "wide", MIToS.MSA.Stockholm),
                                    (msafile_wide_sth,  "ungzipped", "wide", MIToS.MSA.Stockholm),
                                    (msafile_wide_fas_gz, "gzipped", "wide", MIToS.MSA.FASTA),
                                    (msafile_wide_fas,  "ungzipped", "wide", MIToS.MSA.FASTA)]

    # Default parser
    msa["input"][string(shape,"_",gzipped,"_",format)] = @benchmarkable read($file, $format)::MIToS.MSA.AnnotatedMultipleSequenceAlignment
    # With mapping
    msa["input"][string(shape,"_",gzipped,"_",format,"_mapping")] = @benchmarkable read($file, $format, generatemapping=true, useidcoordinates=true)::MIToS.MSA.AnnotatedMultipleSequenceAlignment
end

##### output
msa["output"] = BenchmarkGroup()

for (file,gzipped,shape,format) in [(msafile_long_sth_gz, "gzipped", "long", MIToS.MSA.Stockholm),
                                    (msafile_long_sth,  "ungzipped", "long", MIToS.MSA.Stockholm),
                                    (msafile_long_fas_gz, "gzipped", "long", MIToS.MSA.FASTA),
                                    (msafile_long_fas,  "ungzipped", "long", MIToS.MSA.FASTA),
                                    (msafile_wide_sth_gz, "gzipped", "wide", MIToS.MSA.Stockholm),
                                    (msafile_wide_sth,  "ungzipped", "wide", MIToS.MSA.Stockholm),
                                    (msafile_wide_fas_gz, "gzipped", "wide", MIToS.MSA.FASTA),
                                    (msafile_wide_fas,  "ungzipped", "wide", MIToS.MSA.FASTA)]
    outfile = string("./tmp/",split(file,"/")[end])
    aln = read(file, format)
    msa["output"][string(shape,"_",format,"_",gzipped)] = @benchmarkable write($outfile, $aln, $format)
end

##### Identity
msa["identity"] = BenchmarkGroup()

for (aln,label) in ((msa_long,"long"), (msa_wide,"wide"))
    for t in (Float16,Float32)
        msa["identity"][string("matrix_",t,"_",label)] = @benchmarkable percentidentity($aln, $t)
    end
    if label == "wide"
        msa["identity"][string("matrix_Float64_wide")] = @benchmarkable percentidentity($aln, Float64)
    end
    msa["identity"][string("mean_",label)] = @benchmarkable meanpercentidentity($aln)
end

##### Clustering
msa["hobohmI"] = BenchmarkGroup()

for (aln,label) in ((msa_long,"long"), (msa_wide,"wide"))
    for pid in 10:10:90
        msa["hobohmI"][string(pid,"_",label)] = @benchmarkable hobohmI($aln, $pid)
    end
end

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

# --------------------------------------------------------------------------- #
# Information module benchmarks
const information = BenchmarkGroup()

##### estimateincolumns
information["estimateincolumns"] = BenchmarkGroup()

for (aln,label) in ((msa_long,"long"), (msa_wide,"wide"))
    residues = getresidues(aln)
    information["estimateincolumns"][string("Entropy_Probability_",label)] = @benchmarkable estimateincolumns($residues, Int, ResidueProbability{Float64,1,false}, Entropy{Float64}())
    information["estimateincolumns"][string("Entropy_Count_",label)] = @benchmarkable estimateincolumns($residues, ResidueCount{Int,2,false}, Entropy{Float64}())
    information["estimateincolumns"][string("MI_Probability_",label)] = @benchmarkable estimateincolumns($residues, Int, ResidueProbability{Float64,2,false}, MutualInformation{Float64}())
    information["estimateincolumns"][string("MI_Count_",label)] = @benchmarkable estimateincolumns($residues, ResidueCount{Int,2,false}, MutualInformation{Float64}())
    information["estimateincolumns"][string("Entropy_Probability_Gapped_",label)] = @benchmarkable estimateincolumns($residues, Int, ResidueProbability{Float64,1,true}, Entropy{Float64}())
    information["estimateincolumns"][string("Entropy_Count_Gapped_",label)] = @benchmarkable estimateincolumns($residues, ResidueCount{Int,2,true}, Entropy{Float64}())
    information["estimateincolumns"][string("MI_Probability_Gapped_",label)] = @benchmarkable estimateincolumns($residues, Int, ResidueProbability{Float64,2,true}, MutualInformation{Float64}())
    information["estimateincolumns"][string("MI_Count_Gapped_",label)] = @benchmarkable estimateincolumns($residues, ResidueCount{Int,2,true}, MutualInformation{Float64}())
end

##### high level
information["highlevel"] = BenchmarkGroup()

for (aln,label) in ((msa_long,"long"), (msa_wide,"wide"))
    # default
    information["highlevel"][string("Buslje09_",label)] = @benchmarkable buslje09($aln)
    information["highlevel"][string("ZBLMIp_",label)] = @benchmarkable BLMI($aln)
end

##### low level
information["lowlevel"] = BenchmarkGroup()

const residues = getresidues(msa_long)
const column_i = residues[:,10]
const column_j = residues[:,20]
const column_k = residues[:,30]
const Pij = probabilities(column_i, column_j)
const Gij = ContingencyTable{Float64,  2, UngappedAlphabet}(UngappedAlphabet())
const nseq_msa_long = nsequences(msa_long)
const clusters_long = hobohmI(msa_long, 62)

information["lowlevel"]["count_col"] = @benchmarkable count($column_i)
information["lowlevel"]["count_col_col"] = @benchmarkable count($column_i, $column_j)
information["lowlevel"]["count_col_col_col"] = @benchmarkable count($column_i, $column_j, $column_k)
information["lowlevel"]["count_col_clusters"] = @benchmarkable count($column_i, weight=$clusters_long)
information["lowlevel"]["count_col_col_clusters"] = @benchmarkable count($column_i, $column_j, weight=$clusters_long)
information["lowlevel"]["count_col_col_col_clusters"] = @benchmarkable count($column_i, $column_j, $column_k, weight=$clusters_long)
information["lowlevel"]["probabilities_col"] = @benchmarkable probabilities($column_i)
information["lowlevel"]["probabilities_col_col"] = @benchmarkable probabilities($column_i, $column_j)
information["lowlevel"]["probabilities_col_col_col"] = @benchmarkable probabilities($column_i, $column_j, $column_k)

information["lowlevel"]["blosum_pseudofrequencies"] = @benchmarkable blosum_pseudofrequencies!($Gij, $Pij)
information["lowlevel"]["probabilities_blosum"] = @benchmarkable probabilities($column_i, $column_j, pseudofrequencies=BLOSUM_Pseudofrequencies(nsequences(msa), 8.512))

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

# --------------------------------------------------------------------------- #
# PDB module benchmarks
const mitos_pdb = BenchmarkGroup()

## Set up
### Files
const pdb_pdb_gz = "../../data/2XWB.pdb.gz"
const pdb_xml_gz = "../../data/2XWB.xml.gz"
const pdb_pdb    = "../../data/2XWB.pdb"
const pdb_xml    = "../../data/2XWB.xml"
### Residues
const pdb_residues = read(pdb_xml_gz, PDBML)

#### Parse benchmarks

##### input
mitos_pdb["input"] = BenchmarkGroup()

for (file,gzipped,label,format) in [(pdb_pdb_gz, "gzipped", "pdb", PDBFile),
                                    (pdb_xml_gz, "gzipped", "xml", PDBML),
                                    (pdb_pdb,  "ungzipped", "pdb", PDBFile),
                                    (pdb_xml,  "ungzipped", "xml", PDBML)]
    # Default parser
    mitos_pdb["input"][string(label,"_",format,"_",gzipped)] = @benchmarkable read($file, $format)
end

##### output
mitos_pdb["output"] = BenchmarkGroup()

for (file,gzipped,label,format) in [(pdb_pdb_gz, "gzipped", "pdb", PDBFile),
                                    (pdb_xml_gz, "gzipped", "xml", PDBML),
                                    (pdb_pdb,  "ungzipped", "pdb", PDBFile),
                                    (pdb_xml,  "ungzipped", "xml", PDBML)]
    outfile = string("./tmp/",split(file,"/")[end])
    mitos_pdb["output"][string(label,"_",format,"_",gzipped)] = @benchmarkable write($outfile, $pdb_residues, $format)
end

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

# --------------------------------------------------------------------------- #
# Pfam module (pipeline) benchmarks
const pfam = BenchmarkGroup()

# Set up
const aln = read(msafile_long_sth_gz, Stockholm, generatemapping=true, useidcoordinates=true)
const col2res = msacolumn2pdbresidue(aln, "CFAB_HUMAN/481-752", "2XWB", "F", "PF00089","../../data/2xwb.xml")
const pdb = read("../../data/2XWB.xml.gz", PDBML)
const resdict = @residuesdict pdb model "1" chain "F" group "ATOM" residue All
const cmap = msacontacts(aln, resdict, col2res)
const ZMIp, MIp = buslje09(aln)

pfam["read_pfam_gzipped"] = @benchmarkable read($msafile_long_sth_gz, Stockholm, generatemapping=true, useidcoordinates=true)
pfam["getseq2pdb"] = @benchmarkable getseq2pdb($aln)
pfam["msacolumn2pdbresidue_sifts"] = @benchmarkable msacolumn2pdbresidue($aln, "CFAB_HUMAN/481-752", "2XWB", "F", "PF00089","../../data/2xwb.xml")
pfam["msacolumn2pdbresidue_sifts_gzipped"] = @benchmarkable msacolumn2pdbresidue($aln, "CFAB_HUMAN/481-752", "2XWB", "F", "PF00089","../../data/2xwb.xml.gz")
pfam["read_PDBML_gzipped"] = @benchmarkable read("../../data/2XWB.xml.gz", PDBML)
pfam["residue_list_to_dict"] = @benchmarkable residuesdict($pdb,"1","F","ATOM",All)
pfam["msaresidues"] = @benchmarkable msaresidues($aln, $resdict, $col2res)
pfam["hasresidues"] = @benchmarkable hasresidues($aln, $col2res)
pfam["contact_map"] = @benchmarkable msacontacts($aln, $resdict, $col2res)
pfam["buslje09"] = @benchmarkable buslje09($aln)
pfam["AUC"] = @benchmarkable AUC($ZMIp, $cmap)

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

function SetUp!(;module_msa::Bool=true,module_information::Bool=true,module_pdb::Bool=true,module_pfam::Bool=true)
    if module_msa
        tune!(msa)
        JLD.save("msa.jld", "msa", params(msa))
    end
    if module_information
        tune!(information)
        JLD.save("information.jld", "information", params(information))
    end
    if module_pdb
        tune!(mitos_pdb)
        JLD.save("mitos_pdb.jld", "mitos_pdb", params(mitos_pdb))
    end
    if module_pfam
        tune!(pfam)
        JLD.save("pfam.jld", "pfam", params(pfam))
    end
end

function Run!(;module_msa::Bool=true,module_information::Bool=true,module_pdb::Bool=true,module_pfam::Bool=true)
    loadparams!(msa, JLD.load("msa.jld", "msa"), :evals, :samples);
    loadparams!(information, JLD.load("information.jld", "information"), :evals, :samples);
    loadparams!(mitos_pdb, JLD.load("mitos_pdb.jld", "mitos_pdb"), :evals, :samples);
    loadparams!(pfam, JLD.load("pfam.jld", "pfam"), :evals, :samples);
    bench = Dict{ASCIIString,BenchmarkGroup}()
    if module_msa
        bench["msa"] = run(msa)
    end
    if module_information
        bench["information"] = run(information)
    end
    if module_pdb
        bench["mitos_pdb"] = run(mitos_pdb)
    end
    if module_pfam
        bench["pfam"] = run(pfam)
    end
    bench
end

# @elapsed include("Benchmark.jl") # 486 s
# @elapsed SetUp!() # 5907 s
# @elapsed result = Run!() # 5642 s
# JLD.save("result_Benchmark.jld", "result", result);


struct QuerySynthesisOCCOracle <: Oracle
    classifier::SVDD.OCClassifier
end

function QuerySynthesisOCCOracle(classifier_type, init_strategy, data_file::String, solver; classifier_params=Dict{Symbol, Any}(), adjust_K=true)
    data, labels = load_data(data_file)
    return QuerySynthesisOCCOracle(classifier_type, init_strategy, data, labels, solver, classifier_params=classifier_params, adjust_K=adjust_K)
end

function QuerySynthesisOCCOracle(classifier_type, init_strategy, data, labels, solver; classifier_params=Dict{Symbol, Any}(), adjust_K=true)
    pools = convert_labels_to_learning(labels)
    oracle = SVDD.instantiate(classifier_type, data, pools, classifier_params)
    SVDD.initialize!(oracle, init_strategy)
    SVDD.set_adjust_K!(oracle, true)
    SVDD.fit!(oracle, solver)
    return QuerySynthesisOCCOracle(oracle)
end

function QuerySynthesisOCCOracle(data, labels, params::Dict{Symbol, Any})
    # inline function declaration
    check_oracle_parameter(params, key) = key in keys(params) || throw(ArgumentError("Parameter '$key' missing for QuerySynthesisOCCOracle."))
    check_oracle_parameter(params, :classifier_type)
    check_oracle_parameter(params, :init_strategy)
    check_oracle_parameter(params, :solver)
    additional_params = Dict(k => params[k] for k in [:classifier_params, :adjust_K] if k in keys(params))
    return QuerySynthesisOCCOracle(params[:classifier_type], params[:init_strategy], data, labels, params[:solver], additional_params...)
end

function ask_oracle(oracle::QuerySynthesisOCCOracle, query_objects::Array{T, 2})::Vector{Symbol} where T <: Real
    return SVDD.classify.(SVDD.predict(oracle.classifier, query_objects))
end


# Example showing the use of waveform triggering
using AcqSynth, PlotlyJS
const acq = AcqSynth

# As in the first example, we setup the board
boardnum = 0
acq.init_board(boardnum)

# Let's change to single channel mode
chan = acq.AIN0
acq.set_channels(boardnum,chan)
# We can verify that this worked properly
acq.get_channels(boardnum) == chan

# Check the frequency, in single channel mode, this will be twice the normal
acq.get_frequency(boardnum)

# Configure the board for waveform triggering
hysteresis = 32
threshold = 2^11-220 # middle value for a 12-bit board
ttype = acq.WAVEFORM_TRIGGER
# This function configures the waveform parameters
acq.set_waveform_trigger_params(boardnum,threshold,hysteresis)
# This function sets the trigger type
acq.set_trigger(boardnum,ttype,acq.RISING_EDGE,chan)

# Setup acquisition, the card will wait until a trigger is detected
numblocks = 2 # let's grab 2 blocks
acq.setup_acquire(boardnum,numblocks)

# Grab the data from the card
data = acq.get_samples_12(boardnum,numblocks) # transfer data as 12-bit samples

# Plot the stuff
t = (1:length(data))./4E9 # time vector, we are sampling at 4GSPS
plot(scatter(x=t[1:2000],y=data[1:2000],name="AIN0"))


# Find All Occurrences of a Pattern in a String

function find_pattern_occurrences(text, pattern)
    k = length(pattern)
    matching_indices = Int64[]
    for ix in 1:(length(text)-k+1)
        substring = text[ix:ix+k-1]
        if substring == pattern
            push!(matching_indices, ix - 1)
        end
    end
    matching_indices
end

function main()
    pattern = readline()
    genome = readline()
    result = find_pattern_occurrences(genome, pattern)
    println(join(result, " "))
end

main()


using Combinatorics
using DataFrames
using DataStructures
using Dates
using LightGraphs
using LinearAlgebra
using Random
using StaticArrays
using Statistics
using TeaFiles
using Test
using TimeDag
using TimeZones

using TimeDag: Node
using TimeDag: IntersectAlignment, LeftAlignment, UnionAlignment
using TimeDag: duplicate, evaluate_until!, start_at

#! format: off

@testset "TimeDag.jl" begin
    # Perform these tests first, because they do naughty things to the global identity map
    # which would e.g. break nodes defined in common.jl.
    @testset "identity map" begin include("identity_map.jl") end

    include("common.jl")

    @testset "block" begin include("block.jl") end
    @testset "constants" begin include("constants.jl") end
    @testset "maybe" begin include("maybe.jl") end
    @testset "output_type" begin include("output_type.jl") end

    @testset "ops" begin
        @testset "align" begin include("ops/align.jl") end
        @testset "array" begin include("ops/array.jl") end
        @testset "common" begin include("ops/common.jl") end
        @testset "conditional" begin include("ops/conditional.jl") end
        @testset "history" begin include("ops/history.jl") end
        @testset "lagging" begin include("ops/lagging.jl") end
        @testset "random" begin include("ops/random.jl") end
        @testset "simple" begin include("ops/simple.jl") end
        @testset "sources" begin include("ops/sources.jl") end
        @testset "window" begin include("ops/window.jl") end
    end
end

#! format: on


# Convenience type to allow dispatch on solution objects that were created by Trixi
#
# This is a union of a Trixi-specific DiffEqBase.ODESolution and of Trixi's own
# TimeIntegratorSolution.
#
# Note: This is an experimental feature and may be changed in future releases without notice.
const TrixiODESolution = Union{ODESolution{T, N, uType, uType2, DType, tType, rateType, P} where
    {T, N, uType, uType2, DType, tType, rateType, P<:ODEProblem{uType_, tType_, isinplace, P_, F_} where
     {uType_, tType_, isinplace, P_<:AbstractSemidiscretization, F_<:ODEFunction{true, typeof(rhs!)}}}, TimeIntegratorSolution}

"""
    PlotData2D

Holds all relevant data for creating 2D plots of multiple solution variables and to visualize the
mesh.

!!! warning "Experimental implementation"
    This is an experimental feature and may change in future releases.
"""
struct PlotData2D{Coordinates, Data, VariableNames, Vertices}
  x::Coordinates
  y::Coordinates
  data::Data
  variable_names::VariableNames
  mesh_vertices_x::Vertices
  mesh_vertices_y::Vertices
end

"""
    PlotData2D(u, semi;
               solution_variables=cons2prim,
               grid_lines=true, max_supported_level=11, nvisnodes=nothing,
               slice_axis=:z, slice_axis_intercept=0)

Create a new `PlotData2D` object that can be used for visualizing 2D/3D DGSEM solution data array
`u` with `Plots.jl`. All relevant geometrical information is extracted from the semidiscretization
`semi`. By default, the conservative variables from the solution are used for plotting. This can be
changed by passing an appropriate conversion function to `solution_variables`.

If `grid_lines` is `true`, also extract grid vertices for visualizing the mesh. The output
resolution is indirectly set via `max_supported_level`: all data is interpolated to
`2^max_supported_level` uniformly distributed points in each spatial direction, also setting the
maximum allowed refinement level in the solution. `nvisnodes` specifies the number of visualization
nodes to be used. If it is `nothing`, twice the number of solution DG nodes are used for
visualization, and if set to `0`, exactly the number of nodes in the DG elements are used.

When visualizing data from a three-dimensional simulation, a 2D slice is extracted for plotting.
`slice_axis` specifies the axis orthogonal to the slice plane and may be `:x`, `:y`, or `:z`. The
point on the slice axis where it intersects with the slice plane is given in `slice_axis_intercept`.
Both of these values are ignored when visualizing 2D data.

!!! warning "Experimental implementation"
    This is an experimental feature and may change in future releases.

# Examples
```julia
julia> using Trixi, Plots

julia> trixi_include(default_example())
[...]

julia> pd = PlotData2D(sol)
PlotData2D(...)

julia> plot(pd) # To plot all available variables

julia> plot(pd["scalar"]) # To plot only a single variable

julia> plot!(getmesh(pd)) # To add grid lines to the plot
```
"""
function PlotData2D(u, semi;
                    solution_variables=cons2prim,
                    grid_lines=true, max_supported_level=11, nvisnodes=nothing,
                    slice_axis=:z, slice_axis_intercept=0)
  mesh, equations, solver, _ = mesh_equations_solver_cache(semi)
  @assert ndims(mesh) in (2, 3) "unsupported number of dimensions $ndims (must be 2 or 3)"

  # Extract mesh info
  center_level_0 = mesh.tree.center_level_0
  length_level_0 = mesh.tree.length_level_0
  leaf_cell_ids = leaf_cells(mesh.tree)
  coordinates = mesh.tree.coordinates[:, leaf_cell_ids]
  levels = mesh.tree.levels[leaf_cell_ids]

  unstructured_data = get_unstructured_data(u, semi, solution_variables)
  x, y, data, mesh_vertices_x, mesh_vertices_y = get_data_2d(center_level_0, length_level_0,
                                                             leaf_cell_ids, coordinates, levels,
                                                             ndims(mesh), unstructured_data,
                                                             nnodes(solver), grid_lines,
                                                             max_supported_level, nvisnodes,
                                                             slice_axis, slice_axis_intercept)
  variable_names = SVector(varnames(solution_variables, equations))

  return PlotData2D(x, y, data, variable_names, mesh_vertices_x, mesh_vertices_y)
end

"""
    PlotData2D(u_ode::AbstractVector, semi; kwargs...)

Create a `PlotData2D` object from a one-dimensional ODE solution `u_ode` and the semidiscretization
`semi`.

!!! warning "Experimental implementation"
    This is an experimental feature and may change in future releases.
"""
PlotData2D(u_ode::AbstractVector, semi; kwargs...) = PlotData2D(wrap_array(u_ode, semi), semi; kwargs...)

"""
    PlotData2D(sol::Union{DiffEqBase.ODESolution,TimeIntegratorSolution}; kwargs...)

Create a `PlotData2D` object from a solution object created by either `OrdinaryDiffEq.solve!` (which
returns a `DiffEqBase.ODESolution`) or Trixi's own `solve!` (which returns a
`TimeIntegratorSolution`).

!!! warning "Experimental implementation"
    This is an experimental feature and may change in future releases.
"""
PlotData2D(sol::TrixiODESolution; kwargs...) = PlotData2D(sol.u[end], sol.prob.p; kwargs...)

# Auxiliary data structure for visualizing a single variable
#
# Note: This is an experimental feature and may be changed in future releases without notice.
struct PlotDataSeries2D{PD<:PlotData2D}
  plot_data::PD
  variable_id::Int
end

# Show only a truncated output for convenience (the full data does not make sense)
function Base.show(io::IO, pds::PlotDataSeries2D)
  print(io, "PlotDataSeries2D{", typeof(pds.plot_data), "}(<plot_data::PlotData2D>, ",
        pds.variable_id, ")")
end

"""
    Base.getindex(pd::PlotData2D, variable_name)

Extract a single variable `variable_name` from `pd` for plotting with `Plots.plot`.

!!! warning "Experimental implementation"
    This is an experimental feature and may change in future releases.
"""
function Base.getindex(pd::PlotData2D, variable_name)
  variable_id = findfirst(isequal(variable_name), pd.variable_names)

  if isnothing(variable_id)
    throw(KeyError(variable_name))
  end

  return PlotDataSeries2D(pd, variable_id)
end

# Define additional methods for convenience
Base.firstindex(pd::PlotData2D) = first(pd.variable_names)
Base.lastindex(pd::PlotData2D) = last(pd.variable_names)
Base.length(pd::PlotData2D) = length(pd.variable_names)
Base.size(pd::PlotData2D) = (length(pd),)
Base.keys(pd::PlotData2D) = tuple(pd.variable_names...)
Base.eltype(pd::PlotData2D) = Pair{String, PlotDataSeries2D}
function Base.iterate(pd::PlotData2D, state=1)
  if state > length(pd)
    return nothing
  else
    return (pd.variable_names[state] => pd[pd.variable_names[state]], state + 1)
  end
end

# Show only a truncated output for convenience (the full data does not make sense)
function Base.show(io::IO, pd::PlotData2D)
  print(io, "PlotData2D{",
            typeof(pd.x), ",",
            typeof(pd.data), ",",
            typeof(pd.variable_names), ",",
            typeof(pd.mesh_vertices_x),
            "}(<x>, <y>, <data>, <variable_names>, <mesh_vertices_x>, <mesh_vertices_y>)")
end

# Auxiliary data structure for visualizing the mesh
#
# Note: This is an experimental feature and may be changed in future releases without notice.
struct PlotMesh2D{PD<:PlotData2D}
  plot_data::PD
end

# Show only a truncated output for convenience (the full data does not make sense)
function Base.show(io::IO, pm::PlotMesh2D)
  print(io, "PlotMesh2D{", typeof(pm.plot_data), "}(<plot_data::PlotData2D>)")
end

"""
    getmesh(pd::PlotData2D)

Extract grid lines from `pd` for plotting with `Plots.plot`.

!!! warning "Experimental implementation"
    This is an experimental feature and may change in future releases.
"""
getmesh(pd::PlotData2D) = PlotMesh2D(pd)

# Visualize a single variable in a 2D plot (default: heatmap)
#
# Note: This is an experimental feature and may be changed in future releases without notice.
@recipe function f(pds::PlotDataSeries2D)
  @unpack plot_data, variable_id = pds
  @unpack x, y, data, variable_names = plot_data

  # Set geometric properties
  xlims --> (x[begin], x[end])
  ylims --> (y[begin], y[end])
  aspect_ratio --> :equal

  # Set annotation properties
  legend -->  :none
  title --> variable_names[variable_id]
  colorbar --> :true

  # Set series properties
  seriestype --> :heatmap

  # Return data for plotting
  x, y, data[variable_id]
end

# Visualize the mesh in a 2D plot
#
# Note: This is an experimental feature and may be changed in future releases without notice.
@recipe function f(pm::PlotMesh2D)
  @unpack plot_data = pm
  @unpack x, y, mesh_vertices_x, mesh_vertices_y = plot_data

  # Set geometric and annotation properties
  xlims --> (x[begin], x[end])
  ylims --> (y[begin], y[end])
  aspect_ratio --> :equal
  legend -->  :none

  # Set series properties
  seriestype := :path
  linecolor := :black
  linewidth := 1

  # Return data for plotting
  mesh_vertices_x, mesh_vertices_y
end

# Plot all available variables at once for convenience
#
# Note: This is an experimental feature and may be changed in future releases without notice.
@recipe function f(pd::PlotData2D)
  # Create layout that is as square as possible, with a preference for more columns than rows if not
  cols = ceil(Int, sqrt(length(pd)))
  rows = ceil(Int, length(pd)/cols)
  layout := (rows, cols)

  # Plot all existing variables
  for (i, (variable_name, series)) in enumerate(pd)
    @series begin
      subplot := i
      series
    end
  end

  # Fill remaining subplots with empty plot
  for i in (length(pd)+1):(rows*cols)
    @series begin
      subplot := i
      axis := false
      ticks := false
      legend := false
      [], []
    end
  end
end

# Create a PlotData2D plot directly from a TrixiODESolution for convenience
#
# Note: This is an experimental feature and may be changed in future releases without notice.
#
# Note: If you change the defaults values here, you need to also change them in the PlotData2D
#       constructor.
@recipe function f(sol::TrixiODESolution;
                   solution_variables=cons2prim,
                   grid_lines=true, max_supported_level=11, nvisnodes=nothing, slice_axis=:z,
                   slice_axis_intercept=0)
  return PlotData2D(sol;
                    solution_variables=solution_variables,
                    grid_lines=grid_lines, max_supported_level=max_supported_level,
                    nvisnodes=nvisnodes, slice_axis=slice_axis,
                    slice_axis_intercept=slice_axis_intercept)
end


module TestLazyArrays

include("preamble.jl")
using LazyArrays: Vcat, Hcat

@testset "Vcat" begin
    @testset for xs in @expressions begin
            Vcat(1:5, 2:3)
            Vcat(zeros(1, 2), ones(2, 2))
        end
        xs = @eval $xs
        @test collect(Map(identity), xs) == collect(xs)[:]
    end

end

@testset "Hcat" begin
    @testset for xs in @expressions begin
            Hcat(1:2, 3:4)
            Hcat(zeros(2, 1), ones(2, 2))
        end
        xs = @eval $xs
        @test collect(Map(identity), xs) == collect(xs)[:]
    end

end

end  # module


import Pkg
Pkg.activate(".")
using Test
using PandaModels; const _PdM = PandaModels
import PowerModels; const _PM = PowerModels

_PM.silence()

pdm_path = joinpath(dirname(pathof(PandaModels)), "..")
data_path = joinpath(pdm_path, "test", "data")
cast_ts = joinpath(data_path, "cigre_with_timeseries.json")

##
pm2 = _PdM.load_pm_from_json(cast_ts)
param = _PdM.extract_params!(pm2)

##
pm = _PdM.load_pm_from_json(cast_ts)
_PdM.active_powermodels_silence!(pm)
pm = _PdM.check_powermodels_data!(pm)
# params = _PdM.extract_params!(pm)
# pm = delete!(pm, "user_defined_params")
model = _PdM.get_model(pm["pm_model"])
solver = _PdM.get_solver(pm)
mn = _PdM.set_pq_values_from_timeseries(pm)
# pm_mn = _PM.instantiate_model(mn, model, _PM.build_mn_opf, multinetwork=true)

result = _PdM._run_v_stab_ts(
    mn,
    model,
    solver,
    setting = Dict("output" => Dict("branch_flows" => true)),
    ext = _PdM.extract_params!(pm),
)

# using InfrastructureModels; const _IM = InfrastructureModels
#
# for (n, network) in _IM.nws(mn)
#         variable_bus_voltage(pm, nw=n)
#         variable_gen_power(pm, nw=n)
#         variable_branch_power(pm, nw=n)
#         variable_dcline_power(pm, nw=n)
#
#         constraint_model_voltage(pm, nw=n)
#
#         for i in ids(pm, :ref_buses, nw=n)
#             constraint_theta_ref(pm, i, nw=n)
#         end
#
#         for i in ids(pm, :bus, nw=n)
#             constraint_power_balance(pm, i, nw=n)
#         end
#
#         for i in ids(pm, :branch, nw=n)
#             constraint_ohms_yt_from(pm, i, nw=n)
#             constraint_ohms_yt_to(pm, i, nw=n)
#
#             constraint_voltage_angle_difference(pm, i, nw=n)
#
#             constraint_thermal_limit_from(pm, i, nw=n)
#             constraint_thermal_limit_to(pm, i, nw=n)
#         end
#
#         for i in ids(pm, :dcline, nw=n)
#             constraint_dcline_power_losses(pm, i, nw=n)
#         end
#     end

# _PdM.active_powermodels_silence!(pm)
# pm = _PdM.check_powermodels_data!(pm)
# model = _PdM.get_model(pm["pm_model"])
# solver = _PdM.get_solver(pm)
# result = _PdM._run_q_flex(
#     pm,
#     model,
#     solver,
#     setting = Dict("output" => Dict("branch_flows" => true)),
#     ext = _PdM.extract_params!(pm),
# )
# mn = set_pq_values_from_timeseries(pm)

# result = _run_mn_vd(
#     pm,
#     model,
#     solver,
#     setting = Dict("output" => Dict("branch_flows" => true)),
#     ext = extract_params!(pm),
# )

# result = run_pandamodels_q_flex(case_q_flex)
# import InfrastructureModels
# using JuMP
# pm = _PdM.load_pm_from_json(case_q_flex)
# # active_powermodels_silence!(pm)
# pm =  _PdM.check_powermodels_data!(pm)
# model =  _PdM.get_model(pm["pm_model"])
# solver =  _PdM.get_solver(pm)
# #
# # # ext = _PdM.extract_params!(pm)
# #
# # result = _run_q_flex(
# # pm,
# # model,
# # solver,
# # setting = Dict("output" => Dict("branch_flows" => true)),
# # ext = _PdM.extract_params!(pm),
# # )
#
#
# # pm = load_pm_from_json(json_path)
# # active_powermodels_silence!(pm)
# # pm = check_powermodels_data!(pm)
# # model = get_model(pm["pm_model"])
# # solver = get_solver(pm)
# ##  Load System Data
# # pm = load_pm_from_json(json_path)
# # solver = get_solver(pm["pm_solver"], pm["pm_nl_solver"], pm["pm_mip_solver"],
# # pm["pm_log_level"], pm["pm_time_limit"], pm["pm_nl_time_limit"], pm["pm_mip_time_limit"])
#
# # if haskey(pm, "user_defined_params")
#     # params = Dict{Symbol,Dict{String,Any}}()
#     # for key in keys(pm["user_defined_params"])
#     #     params[Symbol(key)] = pm["user_defined_params"][key]
#     # end
#     # data = delete!(pm, "user_defined_params")
# # end
# #
# if haskey(pm, "user_defined_params")
#     user_defined_params = pm["user_defined_params"]
#     data = delete!(pm, "user_defined_params")
# else
#     data = pm
# end
#
# # additional modification
# for (i, branch) in pm["branch"]
#     pm["branch"][i]["angmin"] = -6.28
#     pm["branch"][i]["angmax"] = 6.28
# end
#
# # use build_ref to filter out inactive components
# ref = PowerModels.build_ref(data)[:it][:pm][:nw][0]
#
# ###############################################################################
# # 1. Building the Optimal Power Flow Model
# ###############################################################################
# # Initialize a JuMP Optimization Model
# #-------------------------------------
# model = Model(solver)
# #nlp_optimizer = optimizer_with_attributes(Ipopt.Optimizer)
# #model = Model(nlp_optimizer)
# # Add Optimization and State Variables
# # ------------------------------------
#
# # Add voltage angles va for each bus
# @variable(model, va[i in keys(ref[:bus])])
# # note: [i in keys(ref[:bus])] adds one `va` variable for each bus in the network
#
# # Add voltage angles vm for each bus
# @variable(model, ref[:bus][i]["vmin"] <= vm[i in keys(ref[:bus])] <= ref[:bus][i]["vmax"], start=1.0)
# # note: this vairable also includes the voltage magnitude limits and a starting value
#
# # Add active power generation variable pg for each generator (including limits)
# @variable(model, ref[:gen][i]["pmin"] <= pg[i in keys(ref[:gen])] <= ref[:gen][i]["pmax"])
# # Add reactive power generation variable qg for each generator (including limits)
# @variable(model, ref[:gen][i]["qmin"] <= qg[i in keys(ref[:gen])] <= ref[:gen][i]["qmax"])
#
# # Add power flow variables p to represent the active power flow for each branch
# @variable(model, -ref[:branch][l]["rate_a"] <= p[(l,i,j) in ref[:arcs]] <= ref[:branch][l]["rate_a"])
# # Add power flow variables q to represent the reactive power flow for each branch
# @variable(model, -ref[:branch][l]["rate_a"] <= q[(l,i,j) in ref[:arcs]] <= ref[:branch][l]["rate_a"])
# # note: ref[:arcs] includes both the from (i,j) and the to (j,i) sides of a branch
#
# # Add Objective Function
# # ----------------------
#
# @objective(model, Min, sum((q[(content["element_index"],
#                             content["f_bus"],
#                             content["t_bus"])] - content["value"])^2
#                         for (i, content) in user_defined_params["setpoint_q"]))
#
# # Add Constraints
# # ---------------
#
# # Fix the voltage angle to zero at the reference bus ## why?
# for (i,bus) in ref[:ref_buses]
#     @constraint(model, va[i] == 0)
# end
#
# # Nodal power balance constraints
# for (i,bus) in ref[:bus]
#     # Build a list of the loads and shunt elements connected to the bus i
#     bus_loads = [ref[:load][l] for l in ref[:bus_loads][i]]
#     bus_shunts = [ref[:shunt][s] for s in ref[:bus_shunts][i]]
#
#     # Active power balance at node i
#     @constraint(model,
#         sum(p[a] for a in ref[:bus_arcs][i]) +                  # sum of active power flow on lines from bus i +
#         sum(p_dc[a_dc] for a_dc in ref[:bus_arcs_dc][i]) ==     # sum of active power flow on HVDC lines from bus i =
#         sum(pg[g] for g in ref[:bus_gens][i]) -                 # sum of active power generation at bus i -
#         sum(load["pd"] for load in bus_loads) -                 # sum of active load consumption at bus i -
#         sum(shunt["gs"] for shunt in bus_shunts)*vm[i]^2        # sum of active shunt element injections at bus i
#     )
#
#     # Reactive power balance at node i
#     @constraint(model,
#         sum(q[a] for a in ref[:bus_arcs][i]) +                  # sum of reactive power flow on lines from bus i +
#         sum(q_dc[a_dc] for a_dc in ref[:bus_arcs_dc][i]) ==     # sum of reactive power flow on HVDC lines from bus i =
#         sum(qg[g] for g in ref[:bus_gens][i]) -                 # sum of reactive power generation at bus i -
#         sum(load["qd"] for load in bus_loads) +                 # sum of reactive load consumption at bus i -
#         sum(shunt["bs"] for shunt in bus_shunts)*vm[i]^2        # sum of reactive shunt element injections at bus i
#     )
# end
#
# # Branch power flow physics and limit constraints
# for (i,branch) in ref[:branch]
#     # Build the from variable id of the i-th branch, which is a tuple given by (branch id, from bus, to bus)
#     f_idx = (i, branch["f_bus"], branch["t_bus"])
#     # Build the to variable id of the i-th branch, which is a tuple given by (branch id, to bus, from bus)
#     t_idx = (i, branch["t_bus"], branch["f_bus"])
#     # note: it is necessary to distinguish between the from and to sides of a branch due to power losses
#
#     p_fr = p[f_idx]                     # p_fr is a reference to the optimization variable p[f_idx]
#     q_fr = q[f_idx]                     # q_fr is a reference to the optimization variable q[f_idx]
#     p_to = p[t_idx]                     # p_to is a reference to the optimization variable p[t_idx]
#     q_to = q[t_idx]                     # q_to is a reference to the optimization variable q[t_idx]
#     # note: adding constraints to p_fr is equivalent to adding constraints to p[f_idx], and so on
#
#     vm_fr = vm[branch["f_bus"]]         # vm_fr is a reference to the optimization variable vm on the from side of the branch
#     vm_to = vm[branch["t_bus"]]         # vm_to is a reference to the optimization variable vm on the to side of the branch
#     va_fr = va[branch["f_bus"]]         # va_fr is a reference to the optimization variable va on the from side of the branch
#     va_to = va[branch["t_bus"]]         # va_fr is a reference to the optimization variable va on the to side of the branch
#
#     # Compute the branch parameters and transformer ratios from the data
#     g, b = PowerModels.calc_branch_y(branch)
#     tr, ti = PowerModels.calc_branch_t(branch)
#     g_fr = branch["g_fr"]
#     b_fr = branch["b_fr"]
#     g_to = branch["g_to"]
#     b_to = branch["b_to"]
#     tm = branch["tap"]^2
#     # note: tap is assumed to be 1.0 on non-transformer branches
#
#
#     # AC Power Flow Constraints
#
#     # From side of the branch flow
#     @NLconstraint(model, p_fr ==  (g+g_fr)/tm*vm_fr^2 + (-g*tr+b*ti)/tm*(vm_fr*vm_to*cos(va_fr-va_to)) + (-b*tr-g*ti)/tm*(vm_fr*vm_to*sin(va_fr-va_to)) )
#     @NLconstraint(model, q_fr == -(b+b_fr)/tm*vm_fr^2 - (-b*tr-g*ti)/tm*(vm_fr*vm_to*cos(va_fr-va_to)) + (-g*tr+b*ti)/tm*(vm_fr*vm_to*sin(va_fr-va_to)) )
#
#     # To side of the branch flow
#     @NLconstraint(model, p_to ==  (g+g_to)*vm_to^2 + (-g*tr-b*ti)/tm*(vm_to*vm_fr*cos(va_to-va_fr)) + (-b*tr+g*ti)/tm*(vm_to*vm_fr*sin(va_to-va_fr)) )
#     @NLconstraint(model, q_to == -(b+b_to)*vm_to^2 - (-b*tr+g*ti)/tm*(vm_to*vm_fr*cos(va_fr-va_to)) + (-g*tr-b*ti)/tm*(vm_to*vm_fr*sin(va_to-va_fr)) )
#
#     # Voltage angle difference limit
#     @constraint(model, va_fr - va_to <= branch["angmax"])
#     @constraint(model, va_fr - va_to >= branch["angmin"])
#
#     # Apparent power limit, from side and to side
#     @constraint(model, p_fr^2 + q_fr^2 <= branch["rate_a"]^2)
#     @constraint(model, p_to^2 + q_to^2 <= branch["rate_a"]^2)
# end
#
#
# ###############################################################################
# # 3. Solve the Optimal Power Flow Model and Review the Results
# ###############################################################################
#
# # Solve the optimization problem
# optimize!(model)
# ###############################################################################
# # 4. Create Result Dictionary such that the PowerModels Results can be used by pandapower
# ###############################################################################
# solution =  Dict{String,Any}()
# push!(solution, "baseMVA" => data["baseMVA"])
# push!(solution, "per_unit" => data["per_unit"])
# push!(solution, "gen" => data["gen"])
# push!(solution, "bus" => data["bus"])
# push!(solution, "branch" => data["branch"])
# push!(solution, "multinetwork" => false)
# push!(solution, "multiinfrastructrue" => false)
#
# for (i, gen) in solution["gen"]
#     index = gen["index"]
#     gen["qg"] = value(model[:qg][index])
#     gen["pg"] = value(model[:pg][index])
# end
#
# for (i, bus) in solution["bus"]
#     index = bus["index"]
#     bus["vm"] = value(model[:vm][index])
#     bus["va"] = value(model[:va][index])
# end
#
# for (i, branch) in solution["branch"]
#     index = branch["index"]
#     push!(branch, "qf" => value(model[:q][(index, branch["f_bus"], branch["t_bus"])]))
#     push!(branch, "qt" => value(model[:q][(index, branch["t_bus"], branch["f_bus"])]))
#     push!(branch, "pf" => value(model[:p][(index, branch["f_bus"], branch["t_bus"])]))
#     push!(branch, "pt" => value(model[:p][(index, branch["t_bus"], branch["f_bus"])]))
# end
# println("The solver termination status is")
#
# result = Dict{String,Any}(
#         "optimizer" => JuMP.solver_name(model),
#         "termination_status" => JuMP.termination_status(model),   # "LOCALLY_SOLVED",
#         "primal_status" => JuMP.primal_status(model),
#         "dual_status" => JuMP.dual_status(model),
#         "objective" => InfrastructureModels._guard_objective_value(model),
#         "objective_lb" => InfrastructureModels._guard_objective_bound(model),
#         "solve_time" => solve_time,
#         "solution" => solution)
#
# return result


struct LinearPolicyWithFluxBasis{TP, TB} <: AbstractPolicy where {TP, TB}
    π::TP
    ϕ::TB
end

function params(π::LinearPolicyWithFluxBasis)
    return Flux.params(π.ϕ, params(π.π)...)
end

function paramsvec(π::LinearPolicyWithFluxBasis)
    xs = Zygote.Buffer([])
    m = π.ϕ
    fmap(m) do x
        x isa AbstractArray && push!(xs, x)
        return x
    end
    ps = params(π.π)
    for p in ps
        push!(xs, p)
    end
    return vcat(vec.(copy(xs))...)
end

function paramsfromvec!(π::LinearPolicyWithFluxBasis, θ)
    i = 0
    m = π.ϕ
    fmap(m) do x
        x isa AbstractArray || return x
        x .= reshape(@view θ[i.+(1:length(x))], size(x))
        i += length(x)
        return x
    end
    length(xs) == i || @warn "Expected $(i) params, got $(length(xs))"
    paramsfromvec!(π.π, @view θ[i+1:end])
    return π
end

function rrule(::typeof(paramsfromvec!), π::LinearPolicyWithFluxBasis, θ)
    i = 0
    m = π.ϕ
    m̄ = fmap(m) do x
        x isa AbstractArray || return x
        x .= reshape(@view θ[i.+(1:length(x))], size(x))
        i += length(x)
        return x
    end
    length(xs) == i || @warn "Expected $(i) params, got $(length(xs))"
    paramsfromvec!(π.π, @view θ[i+1:end])

    function paramsfromvec_linearwithflux!(ȳ)
        xs = Zygote.Buffer([])
        dm = ȳ.ϕ
        fmap(dm) do x
            x isa AbstractArray && push!(xs, x)
            return x
        end
        dpi = ȳ.π
        for d in dpi
            push!(xs, p)
        end
        dθ = vcat(vec.(copy(xs))...)
        return NoTangent(), ȳ, dθ
    end
    return π, paramsfromvec_linearwithflux!
end

# function _restructure(m, xs)
#     i = 0
#     m̄ = fmap(m) do x
#         x isa AbstractArray || return x
#         x = reshape(xs[i.+(1:length(x))], size(x))
#         i += length(x)
#         return x
#     end
#     length(xs) == i || @warn "Expected $(i) params, got $(length(xs))"
#     return m̄
# end

# @adjoint function _restructure(m, xs)
#     m̄, numel = _restructure(m, xs), length(xs)
#     function _restructure_pullback(dm)
#         xs′ = destructure(dm)[1]
#         numel == length(xs′) || @warn "Expected $(numel) params, got $(length(xs′))"
#         return (nothing, xs′)
#     end
#     return m̄, _restructure_pullback
# end

# """
#     destructure(m)
# Flatten a model's parameters into a single weight vector.
#     julia> m = Chain(Dense(10, 5, σ), Dense(5, 2), softmax)
#     Chain(Dense(10, 5, σ), Dense(5, 2), softmax)
#     julia> θ, re = destructure(m);
#     julia> θ
#     67-element Vector{Float32}:
#     -0.1407104
#     ...
# The second return value `re` allows you to reconstruct the original network after making
# modifications to the weight vector (for example, with a hypernetwork).
#     julia> re(θ .* 2)
#     Chain(Dense(10, 5, σ), Dense(5, 2), softmax)
# """
# function destructure(m)
#     xs = Zygote.Buffer([])
#     fmap(m) do x
#         x isa AbstractArray && push!(xs, x)
#         return x
#     end
#     return vcat(vec.(copy(xs))...), p -> _restructure(m, p)
# end

function (π::LinearPolicyWithFluxBasis)(s)
    feats = π.ϕ(s)
    return π.π(feats)
end

function logpdf(π::LinearPolicyWithFluxBasis, s, a)
    feats = π.ϕ(s)
    return logpdf(π.π, feats, a)
end

function grad_logpdf(π::LinearPolicyWithFluxBasis, s, a)
    ps = params(π)
    logp, gp = Zygote.pullback(()->logpdf(π, s, a), ps)
    tmp = gp(1)
    ψ = [tmp[p] for p in ps]
    return logp, ψ
end

function sample_with_trace(π::LinearPolicyWithFluxBasis, s)
    ps = params(π)
    action = []
    function f!(res, π)
        d = π(s)

        Zygote.ignore() do
            a = rand(d)
            push!(res, a)
        end
        return logpdf(d, res[1])
    end
    logp, back = Zygote.pullback(()->f!(action, π), ps)
    tmp = back(1)
    ψ = [tmp[p] for p in ps]
    return action[1], logp, ψ
end


const ValBool = Union{Val{false}, Val{true}}

# One-dimensional distributed FFT plan.
struct PencilPlan1D{
        Ti <: Number,  # input type
        To <: Number,  # output type
        Pi <: Pencil,
        Po <: Pencil,
        Tr <: AbstractTransform,
        FFTPlanF <: Transforms.Plan,
        FFTPlanB <: Transforms.Plan,
    }
    # Each pencil pair describes the decomposition of input and output FFT
    # data. The two pencils will be different for transforms that do not
    # preserve the size of the data (e.g. real-to-complex transforms).
    # Otherwise, they will be typically identical.
    pencil_in  :: Pi  # pencil before transform
    pencil_out :: Po  # pencil after transform
    transform  :: Tr  # transform type

    fft_plan   :: FFTPlanF  # forward FFTW plan
    bfft_plan  :: FFTPlanB  # backward FFTW plan (unnormalised)

    scale_factor :: Int  # scale factor for backward transform

    function PencilPlan1D{Ti}(p_i, p_o, tr, fw, bw, scale) where {Ti}
        To = eltype_output(tr, Ti)
        new{Ti, To, typeof(p_i), typeof(p_o), typeof(tr), typeof(fw), typeof(bw)}(
            p_i, p_o, tr, fw, bw, scale,
        )
    end
end

Transforms.eltype_input(::PencilPlan1D{Ti}) where {Ti} = Ti
Transforms.eltype_output(::PencilPlan1D{Ti,To}) where {Ti,To} = To
Transforms.is_inplace(p::PencilPlan1D) = is_inplace(p.transform)

"""
    PencilFFTPlan{T,N} <: AbstractFFTs.Plan{T}

Plan for N-dimensional FFT-based transform on MPI-distributed data, where input
data has type `T`.

---

    PencilFFTPlan(p::Pencil, transforms; kwargs...)

Create a `PencilFFTPlan` for distributed arrays following a given
[`Pencil`](https://jipolanco.github.io/PencilArrays.jl/dev/Pencils/#PencilArrays.Pencils.Pencil)
configuration.
See variant below for details on the specification of `transforms` and on
possible keyword arguments.

---

    PencilFFTPlan(
        A::PencilArray, transforms;
        fftw_flags = FFTW.ESTIMATE,
        fftw_timelimit = FFTW.NO_TIMELIMIT,
        permute_dims = Val(true),
        transpose_method = Transpositions.PointToPoint(),
        timer = TimerOutput(),
    )

Create plan for `N`-dimensional transform on MPI-distributed `PencilArray`s.

# Extended help

This creates a `PencilFFTPlan` for arrays sharing the same properties as `A`
(dimensions, MPI decomposition, memory layout, ...), which describe data on an
`N`-dimensional domain.

## Transforms

The transforms to be applied along each dimension are specified by the
`transforms` argument. Possible transforms are defined as subtypes of
[`Transforms.AbstractTransform`](@ref), and are listed in [Transform
types](@ref). This argument may be either:

- a tuple of `N` transforms to be applied along each dimension. For instance,
  `transforms = (Transforms.R2R(FFTW.REDFT01), Transforms.RFFT(), Transforms.FFT())`;

- a single transform to be applied along all dimensions. The input is
  automatically expanded into `N` equivalent transforms. For instance, for a
  three-dimensional array, `transforms = Transforms.RFFT()` specifies a 3D
  real-to-complex transform, and is equivalent to passing `(Transforms.RFFT(),
  Transforms.FFT(), Transforms.FFT())`.

Note that forward transforms are applied from left to right. In the last
example, this means that a real-to-complex transform (`RFFT`) is first performed along
the first dimension. This is followed by complex-to-complex transforms (`FFT`)
along the second and third dimensions.

## Input data layout

The input `PencilArray` must satisfy the following constraints:

- array dimensions must *not* be permuted. This is the default when constructing
  `PencilArray`s.

- for an `M`-dimensional domain decomposition (with `M < N`), the input array
  must be decomposed along the *last `M` dimensions*. For example, for a 2D
  decomposition of 3D data, the decomposed dimensions must be `(2, 3)`. In
  particular, the first array dimension must *not* be distributed among
  different MPI processes.

  In the PencilArrays package, the decomposed dimensions are specified
  at the moment of constructing a [`Pencil`](https://jipolanco.github.io/PencilArrays.jl/dev/Pencils/#PencilArrays.Pencils.Pencil).

- the element type must be compatible with the specified transform. For
  instance, real-to-complex transforms (`Transforms.RFFT`) require the input to
  be real floating point values. Other transforms, such as `Transforms.R2R`,
  accept both real and complex data.

## Keyword arguments

- The keyword arguments `fftw_flags` and `fftw_timelimit` are passed to the
  `FFTW` plan creation functions (see [`AbstractFFTs`
  docs](https://juliamath.github.io/AbstractFFTs.jl/stable/api/#AbstractFFTs.plan_fft)).

- `permute_dims` determines whether the indices of the output data should be
  reversed. For instance, if the input data has global dimensions
  `(Nx, Ny, Nz)`, then the output of a complex-to-complex FFT would have
  dimensions `(Nz, Ny, Nx)`. This enables FFTs to always be performed along
  the first (i.e. fastest) array dimension, which could lead to performance
  gains. This option is enabled by default. For type inference reasons, it must
  be a value type (`Val(true)` or `Val(false)`).

- `transpose_method` allows to select between implementations of the global
  data transpositions. See
  [PencilArrays docs](https://jipolanco.github.io/PencilArrays.jl/dev/Transpositions/#PencilArrays.Transpositions.Transposition)
  docs for details.

- `timer` should be a `TimerOutput` object.
  See [Measuring performance](@ref PencilFFTs.measuring_performance) for details.

---

    PencilFFTPlan(
        dims_global::Dims{N}, transforms, proc_dims::Dims{M}, comm::MPI.Comm,
        [real_type = Float64]; extra_dims = (), kws...
    )

Create plan for N-dimensional transform.

# Extended help

Instead of taking a `PencilArray` or a `Pencil`, this constructor requires the
global dimensions of the input data, passed via the `size_global` argument.

The data is distributed over the MPI processes in the `comm` communicator.
The distribution is performed over `M` dimensions (with `M < N`) according to
the values in `proc_dims`, which specifies the number of MPI processes to put
along each dimension.

`PencilArray`s that may be transformed with the returned plan can be created
using [`allocate_input`](@ref).

## Optional arguments

- The floating point precision can be selected by setting `real_type` parameter,
  which is `Float64` by default.

- `extra_dims` may be used to specify the sizes of one or more extra dimensions
  that should not be transformed. These dimensions will be added to the rightmost
  (i.e. slowest) indices of the arrays. See **Extra dimensions** below for usage
  hints.

- see the other constructor for more keyword arguments.

## Extra dimensions

One possible application of `extra_dims` is for describing the components of a
vector or tensor field. However, this means that different `PencilFFTPlan`s
would need to be created for each kind of field (scalar, vector, ...).
To avoid the creation of multiple plans, a possibly better alternative is to
create tuples (or arrays) of `PencilArray`s using [`allocate_input`](@ref) and
[`allocate_output`](@ref).

Another more legitimate usage of `extra_dims` is to specify one or more
Cartesian dimensions that should not be transformed nor split among MPI
processes.

## Example

Suppose we want to perform a 3D FFT of real data. The data is to be
decomposed along two dimensions, over 8 MPI processes:

```julia
size_global = (64, 32, 128)  # size of real input data

# Perform real-to-complex transform along the first dimension, then
# complex-to-complex transforms along the other dimensions.
transforms = (Transforms.RFFT(), Transforms.FFT(), Transforms.FFT())
# transforms = Transforms.RFFT()  # this is equivalent to the above line

proc_dims = (4, 2)  # 2D decomposition
comm = MPI.COMM_WORLD

plan = PencilFFTPlan(size_global, transforms, proc_dims, comm)
```

"""
struct PencilFFTPlan{
        T,   # element type of input data
        N,   # dimension of arrays (= Nt + Ne)
        I,   # in-place (Bool)
        Nt,  # number of transformed dimensions
        Nd,  # number of decomposed dimensions
        Ne,  # number of extra dimensions
        G <: GlobalFFTParams,
        P <: NTuple{Nt, PencilPlan1D},
        TransposeMethod <: AbstractTransposeMethod,
    } <: AbstractFFTs.Plan{T}

    global_params :: G
    topology      :: MPITopology{Nd}

    extra_dims :: Dims{Ne}

    # One-dimensional plans, including data decomposition configurations.
    plans :: P

    # Scale factor to be applied after backwards transforms.
    scale_factor :: Float64

    # `method` parameter passed to `transpose!`
    transpose_method :: TransposeMethod

    # Temporary data buffers.
    ibuf :: Vector{UInt8}
    obuf :: Vector{UInt8}

    # Runtime timing.
    # Should be used along with the @timeit_debug macro, to be able to turn it
    # off if desired.
    timer :: TimerOutput

    function PencilFFTPlan(
            A::PencilArray, transforms::AbstractTransformList;
            fftw_flags = FFTW.ESTIMATE,
            fftw_timelimit = FFTW.NO_TIMELIMIT,
            permute_dims::ValBool = Val(true),
            transpose_method::AbstractTransposeMethod =
                Transpositions.PointToPoint(),
            timer::TimerOutput = TimerOutput(),
            ibuf = UInt8[], obuf = UInt8[],  # temporary data buffers
        )
        T = eltype(A)
        dims_global = size_global(pencil(A), LogicalOrder())
        g = GlobalFFTParams(dims_global, transforms, real(T))
        check_input_array(A, g)
        inplace = is_inplace(g)
        fftw_kw = (; flags = fftw_flags, timelimit = fftw_timelimit)

        # Options for creation of 1D plans.
        plans = _create_plans(
            A, g;
            permute_dims = permute_dims,
            ibuf = ibuf,
            timer = timer,
            fftw_kw = fftw_kw,
        )

        scale = prod(p -> float(p.scale_factor), plans)

        # If the plan is in-place, the buffers won't be needed anymore, so we
        # free the memory.
        if inplace
            resize!.((ibuf, obuf), 0)
        end

        edims = extra_dims(A)
        Nt = length(transforms)
        Ne = length(edims)
        N = Nt + Ne
        G = typeof(g)
        P = typeof(plans)
        TM = typeof(transpose_method)
        t = topology(A)
        Nd = ndims(t)

        new{T, N, inplace, Nt, Nd, Ne, G, P, TM}(
            g, t, edims, plans, scale, transpose_method, ibuf, obuf, timer)
    end
end

function PencilFFTPlan(
        pen::Pencil{Nt}, transforms::AbstractTransformList{Nt}, ::Type{Tr} = Float64;
        extra_dims::Dims = (), timer = TimerOutput(), ibuf = UInt8[], kws...,
    ) where {Nt, Tr <: FFTReal}
    T = _input_data_type(Tr, transforms...)
    A = _temporary_pencil_array(T, pen, ibuf, extra_dims)
    PencilFFTPlan(A, transforms; timer = timer, ibuf = ibuf, kws...)
end

function PencilFFTPlan(
        dims_global::Dims{Nt}, transforms::AbstractTransformList{Nt},
        proc_dims::Dims, comm::MPI.Comm, ::Type{Tr} = Float64;
        timer = TimerOutput(), kws...,
    ) where {Nt, Tr}
    t = MPITopology(comm, proc_dims)
    pen = _make_input_pencil(dims_global, t, timer)
    PencilFFTPlan(pen, transforms, Tr; timer = timer, kws...)
end

function PencilFFTPlan(A, transform::AbstractTransform, args...; kws...)
    N = _ndims_transformable(A)
    transforms = expand_dims(transform, Val(N))
    PencilFFTPlan(A, transforms, args...; kws...)
end

@inline _ndims_transformable(dims::Dims) = length(dims)
@inline _ndims_transformable(p::Pencil) = ndims(p)
@inline _ndims_transformable(A::PencilArray) = _ndims_transformable(pencil(A))

"""
    Transforms.is_inplace(p::PencilFFTPlan)

Returns `true` if the given plan operates in-place on the input data, `false`
otherwise.
"""
Transforms.is_inplace(p::PencilFFTPlan{T,N,I}) where {T,N,I} = I :: Bool

"""
    Transforms.eltype_input(p::PencilFFTPlan)

Returns the element type of the input data.
"""
Transforms.eltype_input(p::PencilFFTPlan) = eltype_input(first(p.plans))

"""
    Transforms.eltype_output(p::PencilFFTPlan)

Returns the element type of the output data.
"""
Transforms.eltype_output(p::PencilFFTPlan) = eltype_output(last(p.plans))

pencil_input(p::PencilFFTPlan) = first(p.plans).pencil_in
pencil_output(p::PencilFFTPlan) = last(p.plans).pencil_out

function check_input_array(A::PencilArray, g::GlobalFFTParams)
    # TODO relax condition to ndims(A) >= N and transform the first N
    # dimensions (and forget about extra_dims)
    N = ndims(g)
    if ndims(pencil(A)) != N
        throw(ArgumentError(
            "number of transforms ($N) must be equal to number " *
            "of transformable dimensions in array (`ndims(pencil(A))`)"
        ))
    end

    if permutation(A) != NoPermutation()
        throw(ArgumentError("dimensions of input array must be unpermuted"))
    end

    decomp = decomposition(pencil(A))  # decomposed dimensions, e.g. (2, 3)
    M = length(decomp)
    decomp_expected = input_decomposition(N, Val(M))
    if decomp != decomp_expected
        throw(ArgumentError(
            "decomposed dimensions of input data must be $decomp_expected" *
            " (got $decomp)"
        ))
    end

    T = eltype(A)
    T_expected = input_data_type(g)
    if T_expected !== T
        throw(ArgumentError("wrong input datatype $T, expected $T_expected\n$g"))
    end

    nothing
end

input_decomposition(N, ::Val{M}) where {M} = ntuple(d -> N - M + d, Val(M))

function _create_plans(A::PencilArray, g::GlobalFFTParams; kws...)
    Tin = input_data_type(g)
    transforms = g.transforms
    _create_plans(Tin, g, A, nothing, transforms...; kws...)
end

# Create 1D plans recursively.
function _create_plans(
        ::Type{Ti}, g::GlobalFFTParams{T,N} where T,
        Ai::PencilArray, plan_prev, transform_fw::AbstractTransform,
        transforms_next::Vararg{AbstractTransform,Ntr};
        timer, ibuf, fftw_kw, permute_dims,
    ) where {Ti, N, Ntr}
    dim = Val(N - Ntr)  # current dimension index
    n = N - Ntr
    si = g.size_global_in
    so = g.size_global_out

    Pi = pencil(Ai)

    # Output transform along dimension `n`.
    Po = let dims = ntuple(j -> j ≤ n ? so[j] : si[j], Val(N))
        if dims === size_global(Pi)
            Pi  # in this case Pi and Po are the same
        else
            Pencil(Pi, size_global=dims, timer=timer)
        end
    end

    To = eltype_output(transform_fw, Ti)

    # Note that Ai and Ao may share memory, but that's ok here.
    Ao = _temporary_pencil_array(To, Po, ibuf, extra_dims(Ai))
    plan_n = _make_1d_fft_plan(dim, Ti, Ai, Ao, transform_fw; fftw_kw = fftw_kw)

    # These are both `nothing` when there's no transforms left
    Pi_next = _make_intermediate_pencil(
        g, topology(Pi), Val(n + 1), plan_n, timer, permute_dims)
    Ai_next = _temporary_pencil_array(To, Pi_next, ibuf, extra_dims(Ai))

    (
        plan_n,
        _create_plans(
            To, g, Ai_next, plan_n, transforms_next...;
            timer = timer, ibuf = ibuf, fftw_kw = fftw_kw, permute_dims = permute_dims,
        )...,
    )
end

# No transforms left!
_create_plans(::Type, ::GlobalFFTParams, ::Nothing, plan_prev; kws...) = ()

function _make_input_pencil(dims_global, topology, timer)
    # This is the case of the first pencil pair.
    # Generate initial pencils for the first dimension.
    # - Decompose along dimensions "far" from the first one.
    #   Example: if N = 5 and M = 2, then decomp_dims = (4, 5).
    # - No permutation is applied for input data: arrays are accessed in the
    #   natural order (i1, i2, ..., iN).
    N = length(dims_global)
    M = ndims(topology)
    decomp_dims = input_decomposition(N, Val(M))
    perm = NoPermutation()
    Pencil(topology, dims_global, decomp_dims; permute=perm, timer=timer)
end

function _make_intermediate_pencil(
        g::GlobalFFTParams{T,N} where T,
        topology::MPITopology{M}, dim::Val{n},
        plan_prev::PencilPlan1D, timer,
        permute_dims::ValBool,
    ) where {N, M, n}
    @assert n ≥ 2  # this is an intermediate pencil
    n > N && return nothing
    Po_prev = plan_prev.pencil_out

    # (i) Determine permutation of pencil data.
    perm = _make_permutation_in(permute_dims, dim, Val(N))

    # (ii) Determine decomposed dimensions from the previous
    # decomposition `n - 1`.
    # If `n` was decomposed previously, shift its associated value
    # in `decomp_prev` to the left.
    # Example: if n = 3 and decomp_prev = (1, 3), then decomp = (1, 2).
    decomp_prev = decomposition(Po_prev)
    decomp = ntuple(Val(M)) do i
        p = decomp_prev[i]
        p == n ? p - 1 : p
    end

    # Note that if `n` was not decomposed previously, then the
    # decomposed dimensions stay the same.
    @assert n ∈ decomp_prev || decomp === decomp_prev

    # If everything is done correctly, there should be no repeated
    # decomposition dimensions.
    @assert allunique(decomp)

    # Create new pencil sharing some information with Po_prev.
    # (Including dimensions, MPI topology and data buffers.)
    Pencil(Po_prev, decomp_dims=decomp, permute=perm, timer=timer)
end

# No permutations
_make_permutation_in(permute_dims::Val{false}, etc...) = NoPermutation()

function _make_permutation_in(::Val{true}, dim::Val{n}, ::Val{N}) where {n, N}
    @assert n ≥ 2
    # Here the data is permuted so that the n-th logical dimension is the first
    # (fastest) dimension in the arrays.
    # The chosen permutation is equivalent to (n, (1:n-1)..., (n+1:N)...).
    t = ntuple(i -> (i == 1) ? n : (i ≤ n) ? (i - 1) : i, Val(N))
    @assert isperm(t)
    @assert t == (n, (1:n-1)..., (n+1:N)...)
    Permutation(t)
end

# Case n = N
function _make_permutation_in(::Val{true}, dim::Val{N}, ::Val{N}) where {N}
    # This is the last transform, and I want the index order to be
    # exactly reversed (easier to work with than the alternative above).
    Permutation(ntuple(i -> N - i + 1, Val(N)))  # (N, N-1, ..., 2, 1)
end

function _make_1d_fft_plan(
        dim::Val{n}, ::Type{Ti}, A_fw::PencilArray, A_bw::PencilArray,
        transform_fw::AbstractTransform; fftw_kw) where {n, Ti}
    Pi = pencil(A_fw)
    Po = pencil(A_bw)
    perm = permutation(Pi)

    dims = if PencilArrays.isidentity(perm)
        n  # no index permutation
    else
        # Find index of n-th dimension in the permuted array.
        # If we permuted data to have the n-th dimension as the fastest
        # (leftmost) index, then the result of `findfirst` should be 1.
        findfirst(==(n), Tuple(perm)) :: Int
    end

    d = size(parent(A_fw), 1)  # input length along transformed dimension
    transform_bw = binv(transform_fw, d)

    # Scale factor to be applied after backward transform.
    # The passed array must have the dimensions of the backward transform output
    # (i.e. the forward transform input)
    scale_bw = scale_factor(transform_bw, parent(A_bw), dims)

    # Generate forward and backward FFTW transforms.
    plan_fw = plan(transform_fw, parent(A_fw), dims; fftw_kw...)
    plan_bw = plan(transform_bw, parent(A_bw), dims; fftw_kw...)

    PencilPlan1D{Ti}(Pi, Po, transform_fw, plan_fw, plan_bw, scale_bw)
end

function Base.show(io::IO, p::PencilFFTPlan)
    show(io, p.global_params)
    edims = extra_dims(p)
    isempty(edims) || println(io, "Extra dimensions: $edims")
    println(io)
    show(io, p.topology)
    nothing
end

"""
    get_comm(p::PencilFFTPlan)

Get MPI communicator associated to a `PencilFFTPlan`.
"""
get_comm(p::PencilFFTPlan) = get_comm(topology(p))

"""
    scale_factor(p::PencilFFTPlan)

Get scale factor associated to a `PencilFFTPlan`.
"""
scale_factor(p::PencilFFTPlan) = p.scale_factor

"""
    timer(p::PencilFFTPlan)

Get `TimerOutput` attached to a `PencilFFTPlan`.

See [Measuring performance](@ref PencilFFTs.measuring_performance) for details.
"""
timer(p::PencilFFTPlan) = p.timer

# For consistency with AbstractFFTs, this gives the global dimensions of the input.
Base.size(p::PencilFFTPlan) = size_global(pencil_input(p), LogicalOrder())

topology(p::PencilFFTPlan) = p.topology
extra_dims(p::PencilFFTPlan) = p.extra_dims


module LennardJones

export generate_cubic_lattice, generate_lattice, LATTICE, State, VerletList,
       update_serial!, serial, simd, naive, cuda

const FORCES = Val{true}
const NOFORCES = Val{false}

include("state.jl")       # Allows particle models to be saved as States
include("models.jl")      # Allows particle model generation and printing
include("verlet_list.jl") # Allows the usage of Verlet lists
include("serial.jl")      # Calculation of energy and forces with serial methods
include("simd.jl")        # Calculation of energy and forces with SIMD methods
include("cuda.jl")        # Calculation of energy and forces with CUDA methods

end # module