using LinearAlgebra

@testset "U1Tensor" begin
    u1chs, u1ds, io = U1Charges(-2:2), fill(2,5), InOut(-1,1)
    @testset "Initialization" begin
        a = DASTensor{Float64,2}(U1(), (u1chs, u1chs), (u1ds, u1ds), io)

        @test symmetry(a) == U1()
        @test chargetype(a) == chargetype(U1())
        @test chargestype(a) == chargestype(U1())
        @test eltype(a) == Float64
        @test isempty(tensor(a))

        #Initialize with zeros
        b0 = initwithzero!(deepcopy(a))
        @test isinvariant(b0)
        @tensor bs = b0[1,1]

        @test charge(b0) == zero(U1Charge)
        @test bs == zero(Float64)

        #Initialize with rand
        br = initwithrand!(deepcopy(a))
        bd = todense(br)
        @tensor begin
            sc = br[1,1]
            scd = bd[1,1]
        end

        @test sc isa Float64
        @test scd isa Float64
        @test sc ≈ scd
    end

    @testset "TensorOperations" begin
        f = DASTensor{Complex{Float64},3}(U1(), ntuple(x -> u1chs, 3),
                 ntuple(x -> u1ds, 3), InOut(-1,-1,1))
        fr = initwithrand!(deepcopy(f))
        #permutation
        perms = (p for p in Iterators.product(ntuple(i -> 1:3, 3)...)
            if sort(collect(p)) == 1:3)
        @test all(toarray(tensorcopy(fr, p)) ≈ tensorcopy(toarray(fr), p)
            for p in perms)
        ##inner product direct
        @tensor ref = fr[1,2,3] * fr'[1,2,3]
        ##inner product direct with permutation
        @test ref ≈ @tensor fr'[2,1,3] * fr[2,1,3]
        ##outer product then trace
        @test ref ≈ @tensor begin
            fr2[1,2,3,4,5,6] := fr[1,2,3] * fr'[4,5,6]
            fr2[1,2,3,1,2,3]
        end
        ##partial contraction then partial trace
        @test ref ≈ @tensor begin
            fr3[1,2,3,4] := fr[1,2,-1] * fr'[3,4,-1]
            fr3[1,2,1,2]
        end
    end

    @testset "KrylovKit" begin
        #compare operations from krylovkit
        f = DASTensor{Complex{Float64},3}(U1(), ntuple(x -> u1chs, 3),
                 ntuple(x -> u1ds, 3), InOut(-1,-1,1))
        fr = initwithrand!(deepcopy(f))

        ##fill!
        fr0 = fill!(deepcopy(fr), 0)
        @test all(iszero(v) for v in values(fr0))

        ##mul! -> w = a .* v
        fr2 = deepcopy(fr)
        mul!(fr2, fr, 2)
        @test fr2 ≈ fr * 2
        mul!(fr2, fr, 1/2)
        @test fr2 ≈ fr/2
        mul!(fr2, fr, 1im)
        @test fr2 ≈ fr * 1im

        ##rmul! w = a .* w
        fr2 = deepcopy(fr)
        rmul!(fr2, 2)
        @test fr2 ≈ fr * 2
        rmul!(fr2, 0)
        @test all(iszero(v) for v in values(fr2))

        ##axpy! a,v,w -> w = a * v + w
        fr2 = deepcopy(fr)
        axpy!(1,fr, fr2)
        @test fr2 ≈ 2 * fr
        axpy!(0,fr, fr2)
        @test fr2 ≈ 2 * fr

        ##axpby! a,v,b,w -> w = a * v + b * w
        fr2 = deepcopy(fr)
        axpby!(1,fr,1,fr2)
        @test fr2 ≈ 2 * fr
        initwithrand!(fr2)
        axpby!(1,fr,0,fr2)
        @test fr2 ≈ fr
        initwithrand!(fr2)
        @tensor frsum[1,2,3] := fr[1,2,3] + fr2[1,2,3]
        axpby!(1,fr,1,fr2)
        @test fr2 ≈ frsum

        ##dot, norm
        f = DASTensor{Complex{Float64},3}(U1(), ntuple(x -> u1chs, 3),
                 ntuple(x -> u1ds, 3), InOut(-1,-1,1))
        fr = initwithrand!(deepcopy(f))
        @tensor ref = fr[1,2,3] * fr'[1,2,3]
        @test ref ≈ dot(fr,fr)
        @test ref ≈ norm(fr)^2
        nfr = fr / norm(fr)
        @test norm(nfr) ≈ 1
        @tensor ref = nfr[1,2,3] * conj(nfr)[1,2,3]
    end

    @testset "TensorFactorization" begin
        ##w/o splitting and fusing
        ar = DASTensor{Float64,2}(U1(), (u1chs, u1chs), (u1ds, u1ds), io)
        initwithrand!(ar)
        ua, sa, vda = tensorsvd(ar)
        @tensor ar2[1,2] := ua[1,-1] * sa[-1,-2] * vda[-2,2]
        @test ar2 ≈ ar
        @test sum(diag(sa)) ≈ @tensor sa[1,1]
        @tensor sa2[1,2] := ua'[-1,1] * ar[-1,-2] * vda'[2,-2]
        @test sa2 ≈ sa

        ##with splitting and fusing
        f = DASTensor{Complex{Float64},3}(U1(), ntuple(x -> u1chs, 3),
                 ntuple(x -> u1ds, 3), InOut(-1,-1,1))
        fr = initwithrand!(deepcopy(f))
        uf, sf, vdf = tensorsvd(fr,((1,2),3))
        @tensor fr2[1,2,3] := uf[1,2,-1] * sf[-1,-2] * vdf[-2,3]
        @test fr2 ≈ fr

        ## with charge
        ar = DASTensor{Float64,2}(U1(), (u1chs, u1chs), (u1ds, u1ds), io, U1Charge(1))
        initwithrand!(ar)
        @test charge(ar) == U1Charge(1)
        ua, sa, vda = tensorsvd(ar)
        @tensor ar2[1,2] := ua[1,-1] * sa[-1,-2] * vda[-2,2]
        @test ar2 ≈ ar
        @test sum(diag(sa)) ≈ @tensor sa[1,1]
        @tensor sa2[1,2] := ua'[-1,1] * ar[-1,-2] * vda'[2,-2]
        @test sa2 ≈ sa

        ## with truncation
        ar = DASTensor{Float64,2}(U1(), (U1Charges(0:0), U1Charges(0:0)), ([200], [200]), io)
        initwithrand!(ar)
        ua, sa, vda = tensorsvd(ar, svdtrunc = svdtrunc_maxerror(0.001))
        @test all(sizes(sa,1) .<= 200)

        @tensor ar2[1,2] := ua[1,-1] * sa[-1,-2] * vda[-2,2]
        ar2 /= norm(ar2)
        ar /= norm(ar)
        @test dot(ar2,ar) > 1-0.001

        ua, sa, vda = tensorsvd(ar, svdtrunc = svdtrunc_maxcumerror(0.001))
        @test all(sizes(sa,1) .<= 200)

        @tensor ar2[1,2] := ua[1,-1] * sa[-1,-2] * vda[-2,2]
        ar2 /= norm(ar2)
        ar /= norm(ar)
        @test dot(ar2,ar) > 1-0.001

        ##tensorqr
        ar = DASTensor{Float64,2}(U1(), (u1chs, u1chs), (u1ds, u1ds), io)
        initwithrand!(ar)
        q,r = tensorqr(ar,((1,),(2,)))
        @tensor id[1,2] := q[1,-1] * q'[2,-1]
        @test toarray(id) ≈ I
        @tensor id[1,2] := q'[1,-1] * q[2,-1]
        @test toarray(id) ≈ I
        @test ar ≈ @tensor ar2[1,2] := q[1,-1] * r[-1,2]

        fr = DASTensor{Float64,3}(U1(), (u1chs, u1chs, u1chs), (u1ds, u1ds, u1ds), InOut(1,1,-1))
        fr = initwithrand!(fr)
        q,r = tensorqr(fr, ((1,2),(3,)))
        @tensor id[1,2] := q'[-1,-2,1] * q[-1,-2,2]
        @test toarray(id) ≈ I
        @test fr ≈ @tensor fr2[1,2,3] := q[1,2,-1] * r[-1,3]

        r,q = tensorrq(fr, ((1,),(2,3)))
        @tensor id[1,2] := q'[1,-1,-2] * q[2,-1,-2]
        @test toarray(id) ≈ I
        @test fr ≈ @tensor fr2[1,2,3] := r[1,-1] * q[-1,2,3]
    end

    @testset "pinv" begin
        ar = DASTensor{Float64,2}(
            U1(),
            (U1Charges(-1:1), U1Charges(-1:1)),
            ([10,10,10], [10,10,10]),
            InOut(1,-1))
        initwithrand!(ar)
        ari = pinv(ar)
        @tensor id[1,2] := ar[1,-1] * ari[-1,2]
        @test all(v ≈ I for v in values(id))

        ar = DASTensor{Float64,2}(
            U1(),
            (U1Charges(-1:1), U1Charges(-1:1)),
            (fill(10,3),fill(4,3)),
            InOut(1,-1))
        initwithrand!(ar)
        ari = pinv(ar)
        @tensor id[1,2] := ari[1,-1] * ar[-1,2]
        @test all(v ≈ I for v in values(id))
    end

    @testset "Reshaping" begin
        ##without permutation
        f = DASTensor{Complex{Float64},6}(U1(), ntuple(x -> u1chs, 6),
                 ntuple(x -> u1ds, 6), InOut(-1,-1,-1,1,1,1))
        fr = initwithrand!(deepcopy(f))
        frf, rs2 = fuselegs(fr, ((1,2),(3,4),(5,6)))
        frff, rs3 = fuselegs(frf, ((1,2,3),))
        @test abs(norm(frf)  - norm(fr)) < 1e-10
        @test abs(norm(frff) - norm(fr)) < 1e-10

        frffs = initwithzero!(deepcopy(frf))
        frffss = initwithzero!(deepcopy(fr))
        splitlegs!(frffs,  frff,  ((1,1,1),(1,1,2),(1,1,3)), rs3...)
        splitlegs!(frffss, frffs, ((1,1,1),(1,1,2),(2,2,1),(2,2,2),(3,3,1),(3,3,3)), rs2...)
        @test frffss ≈ fr

        ##with permutation
        f = DASTensor{Complex{Float64},4}(U1(), ntuple(x -> u1chs, 4),
                 ntuple(x -> u1ds, 4), InOut(-1,-1,1,1))
        fr = initwithrand!(deepcopy(f))
        frf, rs = fuselegs(fr, ((1,3),(2,4)))
        @test abs(norm(frf)  - norm(fr)) < 1e-10

        frfs = initwithzero!(deepcopy(fr))
        splitlegs!(frfs,  frf,  ((1,1,1),(2,2,1),(1,1,2),(2,2,2)), rs...)
        @test frfs ≈ fr
    end
end

@testset "Z2" begin
    z2chs, z2ds, io = Z2Charges(), fill(2,2), InOut(-1,1)
    @testset "Initialization" begin
        a = DASTensor{Float64,2}(Z2(), (z2chs, z2chs), (z2ds, z2ds), io)
        @test symmetry(a) == Z2()
        @test chargetype(a) == chargetype(Z2())
        @test chargestype(a) == chargestype(Z2())
        @test eltype(a) == Float64
        @test isempty(tensor(a))

        #Initialize with zeros
        b0 = initwithzero!(deepcopy(a))
        @tensor bs = b0[1,1]
        @test isinvariant(b0)

        @test charge(b0) == zero(Z2Charge)
        @test bs == zero(Float64)

        #Initialize with rand
        br = initwithrand!(deepcopy(a))
        bd = todense(br)
        @tensor begin
            sc = br[1,1]
            scd = bd[1,1]
        end

        @test sc isa Float64
        @test scd isa Float64
        @test sc ≈ scd
    end

    @testset "TensorOperations" begin
        f = DASTensor{Complex{Float64},3}(Z2(), ntuple(x -> z2chs, 3),
                 ntuple(x -> z2ds, 3), InOut(-1,-1,1))
        fr = initwithrand!(deepcopy(f))
        #permutation
        perms = (p for p in Iterators.product(ntuple(i -> 1:3, 3)...)
            if sort(collect(p)) == 1:3)
        @test all(toarray(tensorcopy(fr, p)) ≈ tensorcopy(toarray(fr), p)
            for p in perms)
        ##inner product direct
        @tensor ref = fr[1,2,3] * fr'[1,2,3]
        ##inner product direct with permutation
        @test ref ≈ @tensor fr'[2,1,3] * fr[2,1,3]
        ##outer product then trace
        @test ref ≈ @tensor begin
            fr2[1,2,3,4,5,6] := fr[1,2,3] * fr'[4,5,6]
            fr2[1,2,3,1,2,3]
        end
        ##partial contraction then partial trace
        @test ref ≈ @tensor begin
            fr3[1,2,3,4] := fr[1,2,-1] * fr'[3,4,-1]
            fr3[1,2,1,2]
        end
    end

    @testset "KrylovKit" begin
        #compare operations from krylovkit
        f = DASTensor{Complex{Float64},3}(Z2(), ntuple(x -> z2chs, 3),
                 ntuple(x -> z2ds, 3), InOut(-1,-1,1))
        fr = initwithrand!(deepcopy(f))

        ##fill!
        fr0 = fill!(deepcopy(fr), 0)
        @test all(iszero(v) for v in values(fr0))

        ##mul! -> w = a .* v
        fr2 = deepcopy(fr)
        mul!(fr2, fr, 2)
        @test fr2 ≈ fr * 2
        mul!(fr2, fr, 1/2)
        @test fr2 ≈ fr/2
        mul!(fr2, fr, 1im)
        @test fr2 ≈ fr * 1im

        ##rmul! w = a .* w
        fr2 = deepcopy(fr)
        rmul!(fr2, 2)
        @test fr2 ≈ fr * 2
        rmul!(fr2, 0)
        @test all(iszero(v) for v in values(fr2))

        ##axpy! a,v,w -> w = a * v + w
        fr2 = deepcopy(fr)
        axpy!(1,fr, fr2)
        @test fr2 ≈ 2 * fr
        axpy!(0,fr, fr2)
        @test fr2 ≈ 2 * fr

        ##axpby! a,v,b,w -> w = a * v + b * w
        fr2 = deepcopy(fr)
        axpby!(1,fr,1,fr2)
        @test fr2 ≈ 2 * fr
        initwithrand!(fr2)
        axpby!(1,fr,0,fr2)
        @test fr2 ≈ fr
        initwithrand!(fr2)
        @tensor frsum[1,2,3] := fr[1,2,3] + fr2[1,2,3]
        axpby!(1,fr,1,fr2)
        @test fr2 ≈ frsum

        ##dot, norm
        f = DASTensor{Complex{Float64},3}(Z2(), ntuple(x -> z2chs, 3),
                 ntuple(x -> z2ds, 3), InOut(-1,-1,1))
        fr = initwithrand!(deepcopy(f))
        @tensor ref = fr[1,2,3] * fr'[1,2,3]
        @test ref ≈ dot(fr,fr)
        @test ref ≈ norm(fr)^2
        nfr = fr / norm(fr)
        @test norm(nfr) ≈ 1
    end

    @testset "TensorFactorization" begin
        ##w/o splitting and fusing
        ar = DASTensor{Float64,2}(Z2(), (z2chs, z2chs), (z2ds, z2ds), io)
        initwithrand!(ar)
        ua, sa, vda = tensorsvd(ar)
        @tensor ar2[1,2] := ua[1,-1] * sa[-1,-2] * vda[-2,2]
        @test ar2 ≈ ar
        @test sum(diag(sa)) ≈ @tensor sa[1,1]
        @tensor sa2[1,2] := ua'[-1,1] * ar[-1,-2] * vda'[2,-2]
        @test sa2 ≈ sa

        ##with splitting and fusing
        f = DASTensor{Complex{Float64},3}(Z2(), ntuple(x -> z2chs, 3),
                 ntuple(x -> z2ds, 3), InOut(-1,-1,1))
        fr = initwithrand!(deepcopy(f))
        uf, sf, vdf = tensorsvd(fr,((1,2),3))
        @tensor fr2[1,2,3] := uf[1,2,-1] * sf[-1,-2] * vdf[-2,3]
        @test fr2 ≈ fr

        ## with charge
        ar = DASTensor{Float64,2}(Z2(), (z2chs, z2chs), (z2ds, z2ds), io, Z2Charge(1))
        initwithrand!(ar)
        @test charge(ar) == Z2Charge(1)
        ua, sa, vda = tensorsvd(ar)
        @tensor ar2[1,2] := ua[1,-1] * sa[-1,-2] * vda[-2,2]
        @test ar2 ≈ ar
        @test sum(diag(sa)) ≈ @tensor sa[1,1]
        @tensor sa2[1,2] := ua'[-1,1] * ar[-1,-2] * vda'[2,-2]
        @test sa2 ≈ sa

        ## with truncation
        ar = DASTensor{Float64,2}(Z2(), (Z2Charges(), Z2Charges()), ([200,200], [200,200]), io)
        initwithrand!(ar)
        ua, sa, vda = tensorsvd(ar, svdtrunc = svdtrunc_maxerror(0.001))
        @test all(sizes(sa,1) .<= 200)

        @tensor ar2[1,2] := ua[1,-1] * sa[-1,-2] * vda[-2,2]
        ar2 /= norm(ar2)
        ar /= norm(ar)
        @test dot(ar2,ar) > 1-0.001

        ua, sa, vda = tensorsvd(ar, svdtrunc = svdtrunc_maxcumerror(0.001))
        @test all(sizes(sa,1) .<= 200)

        @tensor ar2[1,2] := ua[1,-1] * sa[-1,-2] * vda[-2,2]
        ar2 /= norm(ar2)
        ar /= norm(ar)
        @test dot(ar2,ar) > 1-0.001

        ##tensorqr
        ar = DASTensor{ComplexF64,2}(ZN{2}(),
            ntuple(i -> ZNCharges{2}(),2),
            ntuple(i -> [10,10], 2),
            InOut(ntuple(i -> -1, 2)...))
        initwithrand!(ar)
        q,r = tensorqr(ar,((1,),(2,)))
        @tensor id[1,2] := q[1,-1] * q'[2,-1]
        @test toarray(id) ≈ I
        @tensor id[1,2] := q'[1,-1] * q[2,-1]
        @test toarray(id) ≈ I
        @test ar ≈ @tensor ar2[1,2] := q[1,-1] * r[-1,2]

        fr = DASTensor{ComplexF64,3}(ZN{2}(),
            ntuple(i -> ZNCharges{2}(),3),
            ntuple(i -> [10,10], 3),
            InOut(1,1,-1))
        fr = initwithrand!(fr)
        q,r = tensorqr(fr, ((1,2),(3,)))
        @tensor id[1,2] := q'[-1,-2,1] * q[-1,-2,2]
        @test toarray(id) ≈ I
        @test fr ≈ @tensor fr2[1,2,3] := q[1,2,-1] * r[-1,3]

        r,q = tensorrq(fr, ((1,),(2,3)))
        @tensor id[1,2] := q'[1,-1,-2] * q[2,-1,-2]
        @test toarray(id) ≈ I
        @test fr ≈ @tensor fr2[1,2,3] := r[1,-1] * q[-1,2,3]
    end

    @testset "pinv" begin
        ar = DASTensor{ComplexF64,2}(ZN{2}(),
            ntuple(i -> ZNCharges{2}(),2),
            ([4,4], [4,4]),
            InOut(ntuple(i -> -1, 2)...))
        initwithrand!(ar)
        ari = pinv(ar)
        @tensor id[1,2] := ar[1,-1] * ari[-1,2]
        @test all(v ≈ I for v in values(id))

        ar = DASTensor{ComplexF64,2}(ZN{2}(),
            ntuple(i -> ZNCharges{2}(),2),
            ([10,10], [4,4]),
            InOut(ntuple(i -> -1, 2)...))
        initwithrand!(ar)
        ari = pinv(ar)
        @tensor id[1,2] := ari[1,-1] * ar[-1,2]
        @test all(v ≈ I for v in values(id))
    end

    @testset "Reshaping" begin
        ##without permutation
        f = DASTensor{Complex{Float64},6}(Z2(), ntuple(x -> z2chs, 6),
                 ntuple(x -> z2ds, 6), InOut(-1,-1,-1,1,1,1))
        fr = initwithrand!(deepcopy(f))
        frf, rs2 = fuselegs(fr, ((1,2),(3,4),(5,6)))
        frff, rs3 = fuselegs(frf, ((1,2,3),))
        @test abs(norm(frf)  - norm(fr)) < 1e-10
        @test abs(norm(frff) - norm(fr)) < 1e-10

        frffs = initwithzero!(deepcopy(frf))
        frffss = initwithzero!(deepcopy(fr))
        splitlegs!(frffs,  frff,  ((1,1,1),(1,1,2),(1,1,3)), rs3...)
        splitlegs!(frffss, frffs, ((1,1,1),(1,1,2),(2,2,1),(2,2,2),(3,3,1),(3,3,3)), rs2...)
        @test frffss ≈ fr

        ##with permutation
        f = DASTensor{Complex{Float64},4}(Z2(), ntuple(x -> z2chs, 4),
                 ntuple(x -> z2ds, 4), InOut(-1,-1,1,1))
        fr = initwithrand!(deepcopy(f))
        frf, rs = fuselegs(fr, ((1,3),(2,4)))
        @test abs(norm(frf)  - norm(fr)) < 1e-10

        frfs = initwithzero!(deepcopy(fr))
        splitlegs!(frfs,  frf,  ((1,1,1),(2,2,1),(1,1,2),(2,2,2)), rs...)
        @test frfs ≈ fr
    end
end

@testset "NDAS" begin
    ndaschs, ndasds, io = NDASCharges(Z2Charges(), U1Charges(-2:2)), fill(2,10), InOut(-1,1)
    @testset "Initialization" begin
        a = DASTensor{Float64,2}(NDAS(Z2(),U1()), (ndaschs, ndaschs), (ndasds, ndasds), io)
        @test symmetry(a) == NDAS(Z2(),U1())
        @test chargetype(a) == chargetype(a)
        @test chargestype(a) == chargestype(a)
        @test eltype(a) == Float64
        @test isempty(tensor(a))

        #Initialize with zeros
        b0 = initwithzero!(deepcopy(a))
        @test isinvariant(b0)
        @tensor bs = b0[1,1]

        @test charge(b0) == zero(chargetype(a))
        @test bs == zero(Float64)

        #Initialize with rand
        br = initwithrand!(deepcopy(a))
        bd = todense(br)
        @tensor begin
            sc = br[1,1]
            scd = bd[1,1]
        end

        @test sc isa Float64
        @test scd isa Float64
        @test sc ≈ scd
    end

    @testset "TensorOperations" begin
        f = DASTensor{Complex{Float64},3}(NDAS(Z2(),U1()), ntuple(x -> ndaschs, 3),
                 ntuple(x -> ndasds, 3), InOut(-1,-1,1))
        fr = initwithrand!(deepcopy(f))
        #permutation
        perms = (p for p in Iterators.product(ntuple(i -> 1:3, 3)...)
            if sort(collect(p)) == 1:3)
        @test all(toarray(tensorcopy(fr, p)) ≈ tensorcopy(toarray(fr), p)
            for p in perms)
        ##inner product direct
        @tensor ref = fr[1,2,3] * fr'[1,2,3]
        ##inner product direct with permutation
        @test ref ≈ @tensor fr'[2,1,3] * fr[2,1,3]
        ##outer product then trace
        @test ref ≈ @tensor begin
            fr2[1,2,3,4,5,6] := fr[1,2,3] * fr'[4,5,6]
            fr2[1,2,3,1,2,3]
        end
        ##partial contraction then partial trace
        @test ref ≈ @tensor begin
            fr3[1,2,3,4] := fr[1,2,-1] * fr'[3,4,-1]
            fr3[1,2,1,2]
        end
    end

    @testset "KrylovKit" begin
        #compare operations from krylovkit
        f = DASTensor{Complex{Float64},3}(NDAS(Z2(),U1()), ntuple(x -> ndaschs, 3),
                 ntuple(x -> ndasds, 3), InOut(-1,-1,1))
        fr = initwithrand!(deepcopy(f))

        ##fill!
        fr0 = fill!(deepcopy(fr), 0)
        @test all(iszero(v) for v in values(fr0))

        ##mul! -> w = a .* v
        fr2 = deepcopy(fr)
        mul!(fr2, fr, 2)
        @test fr2 ≈ fr * 2
        mul!(fr2, fr, 1/2)
        @test fr2 ≈ fr/2
        mul!(fr2, fr, 1im)
        @test fr2 ≈ fr * 1im

        ##rmul! w = a .* w
        fr2 = deepcopy(fr)
        rmul!(fr2, 2)
        @test fr2 ≈ fr * 2
        rmul!(fr2, 0)
        @test all(iszero(v) for v in values(fr2))

        ##axpy! a,v,w -> w = a * v + w
        fr2 = deepcopy(fr)
        axpy!(1,fr, fr2)
        @test fr2 ≈ 2 * fr
        axpy!(0,fr, fr2)
        @test fr2 ≈ 2 * fr

        ##axpby! a,v,b,w -> w = a * v + b * w
        fr2 = deepcopy(fr)
        axpby!(1,fr,1,fr2)
        @test fr2 ≈ 2 * fr
        initwithrand!(fr2)
        axpby!(1,fr,0,fr2)
        @test fr2 ≈ fr
        initwithrand!(fr2)
        @tensor frsum[1,2,3] := fr[1,2,3] + fr2[1,2,3]
        axpby!(1,fr,1,fr2)
        @test fr2 ≈ frsum

        ##dot, norm
        f = DASTensor{Complex{Float64},3}(NDAS(Z2(),U1()), ntuple(x -> ndaschs, 3),
                 ntuple(x -> ndasds, 3), InOut(-1,-1,1))
        fr = initwithrand!(deepcopy(f))
        @tensor ref = fr[1,2,3] * fr'[1,2,3]
        @test ref ≈ dot(fr,fr)
        @test ref ≈ norm(fr)^2
        nfr = fr / norm(fr)
        @test norm(nfr) ≈ 1
    end

    @testset "TensorFactorization" begin
        ##w/o splitting and fusing
        ar = DASTensor{Float64,2}(NDAS(Z2(),U1()), (ndaschs, ndaschs), (ndasds, ndasds), io)
        initwithrand!(ar)
        ua, sa, vda = tensorsvd(ar)
        @tensor ar2[1,2] := ua[1,-1] * sa[-1,-2] * vda[-2,2]
        @test ar2 ≈ ar
        @test sum(diag(sa)) ≈ @tensor sa[1,1]
        @tensor sa2[1,2] := ua'[-1,1] * ar[-1,-2] * vda'[2,-2]
        @test sa2 ≈ sa

        ##with splitting and fusing
        f = DASTensor{Complex{Float64},3}(NDAS(Z2(),U1()), ntuple(x -> ndaschs, 3),
                 ntuple(x -> ndasds, 3), InOut(-1,-1,1))
        fr = initwithrand!(deepcopy(f))
        uf, sf, vdf = tensorsvd(fr,((1,2),3))
        @tensor fr2[1,2,3] := uf[1,2,-1] * sf[-1,-2] * vdf[-2,3]
        @test fr2 ≈ fr

        ## with charge
        ar = DASTensor{Float64,2}(
            NDAS(Z2(),U1()),
            (ndaschs, ndaschs),
            (ndasds, ndasds),
            io,
            NDASCharge(Z2Charge(1), U1Charge(1)))
        initwithrand!(ar)
        @test charge(ar) == NDASCharge(Z2Charge(1), U1Charge(1))
        ua, sa, vda = tensorsvd(ar)
        @tensor ar2[1,2] := ua[1,-1] * sa[-1,-2] * vda[-2,2]
        @test ar2 ≈ ar
        @test sum(diag(sa)) ≈ @tensor sa[1,1]
        @tensor sa2[1,2] := ua'[-1,1] * ar[-1,-2] * vda'[2,-2]
        @test sa2 ≈ sa

        ## with truncation
        ar = DASTensor{Float64,2}(NDAS(Z2(),U1()),
            (NDASCharges(Z2Charges(), U1Charges(0:0)),
             NDASCharges(Z2Charges(), U1Charges(0:0))),
             ([200,200], [200,200]), io)
        initwithrand!(ar)
        ua, sa, vda = tensorsvd(ar, svdtrunc = svdtrunc_maxerror(0.001))
        @test all(sizes(sa,1) .<= 200)

        @tensor ar2[1,2] := ua[1,-1] * sa[-1,-2] * vda[-2,2]
        ar2 /= norm(ar2)
        ar /= norm(ar)
        @test dot(ar2,ar) > 1-0.001

        ua, sa, vda = tensorsvd(ar, svdtrunc = svdtrunc_maxcumerror(0.001))
        @test all(sizes(sa,1) .<= 200)

        @tensor ar2[1,2] := ua[1,-1] * sa[-1,-2] * vda[-2,2]
        ar2 /= norm(ar2)
        ar /= norm(ar)
        @test dot(ar2,ar) > 1-0.001

        ##tensorqr
        ar = DASTensor{Float64,2}(NDAS(Z2(),U1()),
            (NDASCharges(Z2Charges(), U1Charges(0:0)),
             NDASCharges(Z2Charges(), U1Charges(0:0))),
             ([200,200], [200,200]), io)
        initwithrand!(ar)
        q,r = tensorqr(ar,((1,),(2,)))
        @tensor id[1,2] := q[1,-1] * q'[2,-1]
        @test toarray(id) ≈ I
        @tensor id[1,2] := q'[1,-1] * q[2,-1]
        @test toarray(id) ≈ I
        @test ar ≈ @tensor ar2[1,2] := q[1,-1] * r[-1,2]

        f = DASTensor{Complex{Float64},3}(NDAS(Z2(),U1()), ntuple(x -> ndaschs, 3),
                 ntuple(x -> ndasds, 3), InOut(-1,-1,1))
        fr = initwithrand!(deepcopy(f))
        fr = initwithrand!(fr)
        q,r = tensorqr(fr, ((1,2),(3,)))
        @tensor id[1,2] := q'[-1,-2,1] * q[-1,-2,2]
        @test toarray(id) ≈ I
        @test fr ≈ @tensor fr2[1,2,3] := q[1,2,-1] * r[-1,3]

        r,q = tensorrq(fr, ((1,),(2,3)))
        @tensor id[1,2] := q'[1,-1,-2] * q[2,-1,-2]
        @test toarray(id) ≈ I
        @test fr ≈ @tensor fr2[1,2,3] := r[1,-1] * q[-1,2,3]
    end

    @testset "pinv" begin
        ar = DASTensor{Float64,2}(NDAS(Z2(),U1()),
            (NDASCharges(Z2Charges(), U1Charges(0:0)),
             NDASCharges(Z2Charges(), U1Charges(0:0))),
             ([200,200], [200,200]), InOut(1,-1))
        initwithrand!(ar)
        ari = pinv(ar)
        @tensor id[1,2] := ar[1,-1] * ari[-1,2]
        @test all(v ≈ I for v in values(id))

        ar = DASTensor{Float64,2}(NDAS(Z2(),U1()),
            (NDASCharges(Z2Charges(), U1Charges(0:0)),
             NDASCharges(Z2Charges(), U1Charges(0:0))),
             ([200,200], [20,20]), InOut(1,-1))
        initwithrand!(ar)
        ari = pinv(ar)
        @tensor id[1,2] := ari[1,-1] * ar[-1,2]
        @test all(v ≈ I for v in values(id))
    end

    @testset "Reshaping" begin
        ##without permutation
        f = DASTensor{Complex{Float64},6}(NDAS(Z2(),U1()), ntuple(x -> ndaschs, 6),
                 ntuple(x -> ndasds, 6), InOut(-1,-1,-1,1,1,1))
        fr = initwithrand!(deepcopy(f))
        frf, rs2 = fuselegs(fr, ((1,2),(3,4),(5,6)))
        frff, rs3 = fuselegs(frf, ((1,2,3),))
        @test abs(norm(frf)  - norm(fr)) < 1e-10
        @test abs(norm(frff) - norm(fr)) < 1e-10

        frffs = initwithzero!(deepcopy(frf))
        frffss = initwithzero!(deepcopy(fr))
        splitlegs!(frffs,  frff,  ((1,1,1),(1,1,2),(1,1,3)), rs3...)
        splitlegs!(frffss, frffs, ((1,1,1),(1,1,2),(2,2,1),(2,2,2),(3,3,1),(3,3,3)), rs2...)
        @test frffss ≈ fr

        ##with permutation
        f = DASTensor{Complex{Float64},4}(NDAS(Z2(),U1()), ntuple(x -> ndaschs, 4),
                 ntuple(x -> ndasds, 4), InOut(-1,-1,1,1))
        fr = initwithrand!(deepcopy(f))
        frf, rs = fuselegs(fr, ((1,3),(2,4)))
        @test abs(norm(frf)  - norm(fr)) < 1e-10

        frfs = initwithzero!(deepcopy(fr))
        splitlegs!(frfs,  frf,  ((1,1,1),(2,2,1),(1,1,2),(2,2,2)), rs...)
        @test frfs ≈ fr
    end
end

@testset "DTensor" begin
    @testset "Initialization" begin
        #Initialize with zeros
        a = DTensor{Complex{Float64}}((8,8))
        b0 = initwithzero!(deepcopy(a))
        @tensor bs = b0[1,1]

        @test bs == zero(Complex{Float64})

        #Initialize with rand
        br = initwithrand!(deepcopy(a))
        @tensor sc = br[1,1]
        @test sc isa Complex{Float64}
    end

    @testset "TensorOperations" begin
        f = DTensor{Complex{Float64},3}((4,4,4))
        fr = initwithrand!(deepcopy(f))
        #permutation
        perms = (p for p in Iterators.product(ntuple(i -> 1:3, 3)...)
            if sort(collect(p)) == 1:3)
        @test all(toarray(tensorcopy(fr, p)) ≈ tensorcopy(toarray(fr), p)
            for p in perms)
        ##inner product direct
        @tensor ref = fr[1,2,3] * fr'[1,2,3]
        ##inner product direct with permutation
        @test ref ≈ @tensor fr'[2,1,3] * fr[2,1,3]
        ##outer product then trace
        @test ref ≈ @tensor begin
            fr2[1,2,3,4,5,6] := fr[1,2,3] * fr'[4,5,6]
            fr2[1,2,3,1,2,3]
        end
        ##partial contraction then partial trace
        @test ref ≈ @tensor begin
            fr3[1,2,3,4] := fr[1,2,-1] * fr'[3,4,-1]
            fr3[1,2,1,2]
        end
    end

    @testset "KrylovKit" begin
        #compare operations from krylovkit
        f = DTensor{Complex{Float64},3}((8,8,8))
        fr = initwithrand!(deepcopy(f))

        ##fill!
        fr0 = fill!(deepcopy(fr), 0)
        @test iszero(fr0.array)

        ##mul! -> w = a .* v
        fr2 = deepcopy(fr)
        mul!(fr2, fr, 2)
        @test fr2 ≈ fr * 2
        mul!(fr2, fr, 1/2)
        @test fr2 ≈ fr/2
        mul!(fr2, fr, 1im)
        @test fr2 ≈ fr * 1im

        ##rmul! w = a .* w
        fr2 = deepcopy(fr)
        rmul!(fr2, 2)
        @test fr2 ≈ fr * 2
        rmul!(fr2, 0)
        @test iszero(fr2.array)

        ##axpy! a,v,w -> w = a * v + w
        fr2 = deepcopy(fr)
        axpy!(1,fr, fr2)
        @test fr2 ≈ 2 * fr
        axpy!(0,fr, fr2)
        @test fr2 ≈ 2 * fr

        ##axpby! a,v,b,w -> w = a * v + b * w
        fr2 = deepcopy(fr)
        axpby!(1,fr,1,fr2)
        @test fr2 ≈ 2 * fr
        initwithrand!(fr2)
        axpby!(1,fr,0,fr2)
        @test fr2 ≈ fr
        initwithrand!(fr2)
        @tensor frsum[1,2,3] := fr[1,2,3] + fr2[1,2,3]
        axpby!(1,fr,1,fr2)
        @test fr2 ≈ frsum

        ##dot, norm
        f = DTensor{Complex{Float64},3}((8,8,8))
        fr = initwithrand!(deepcopy(f))
        @tensor ref = fr[1,2,3] * fr'[1,2,3]
        @test ref ≈ dot(fr,fr)
        @test ref ≈ norm(fr)^2
        nfr = fr / norm(fr)
        @test norm(nfr) ≈ 1
    end

    @testset "TensorFactorization" begin
        ##w/o splitting and fusing
        ar = DTensor{Float64,2}((20,20))
        initwithrand!(ar)
        ua, sa, vda = tensorsvd(ar)
        @tensor ar2[1,2] := ua[1,-1] * sa[-1,-2] * vda[-2,2]
        @test ar2 ≈ ar
        @test sum(diag(sa)) ≈ @tensor sa[1,1]
        @tensor sa2[1,2] := ua'[-1,1] * ar[-1,-2] * vda'[2,-2]
        @test sa2 ≈ sa

        ##with splitting and fusing
        f = DTensor{Complex{Float64},3}((5,5,5))
        fr = initwithrand!(deepcopy(f))
        uf, sf, vdf = tensorsvd(fr,((1,2),3))
        @tensor fr2[1,2,3] := uf[1,2,-1] * sf[-1,-2] * vdf[-2,3]
        @test fr2 ≈ fr

        ## with truncation
        ar = DTensor{Float64,2}((200,200))
        initwithrand!(ar)
        ua, sa, vda = tensorsvd(ar, svdtrunc = svdtrunc_maxerror(0.001))
        @test all(size(sa) .<= 200)

        @tensor ar2[1,2] := ua[1,-1] * sa[-1,-2] * vda[-2,2]
        ar2 /= norm(ar2)
        ar /= norm(ar)
        @test dot(ar2,ar) > 1-0.001

        ua, sa, vda = tensorsvd(ar, svdtrunc = svdtrunc_maxcumerror(0.001))
        @test all(size(sa) .<= 200)

        @tensor ar2[1,2] := ua[1,-1] * sa[-1,-2] * vda[-2,2]
        ar2 /= norm(ar2)
        ar /= norm(ar)
        @test dot(ar2,ar) > 1-0.001

        ##tensorqr
        ar = DTensor{Float64,2}((200,200))
        initwithrand!(ar)
        q,r = tensorqr(ar,((1,),(2,)))
        @tensor id[1,2] := q[1,-1] * q'[2,-1]
        @test toarray(id) ≈ I
        @tensor id[1,2] := q'[1,-1] * q[2,-1]
        @test toarray(id) ≈ I
        @test ar ≈ @tensor ar2[1,2] := q[1,-1] * r[-1,2]

        fr = DTensor{Complex{Float64},3}((5,5,5))
        fr = initwithrand!(fr)
        q,r = tensorqr(fr, ((1,2),(3,)))
        @tensor id[1,2] := q'[-1,-2,1] * q[-1,-2,2]
        @test toarray(id) ≈ I
        @test fr ≈ @tensor fr2[1,2,3] := q[1,2,-1] * r[-1,3]

        r,q = tensorrq(fr, ((1,),(2,3)))
        @tensor id[1,2] := q'[1,-1,-2] * q[2,-1,-2]
        @test toarray(id) ≈ I
        @test fr ≈ @tensor fr2[1,2,3] := r[1,-1] * q[-1,2,3]
    end

    @testset "pinv" begin
        ar =  DTensor{ComplexF64,2}((200,200))
        initwithrand!(ar)
        ari = pinv(ar)
        @tensor id[1,2] := ar[1,-1] * ari[-1,2]
        @test toarray(id) ≈ I

        ar =  DTensor{ComplexF64,2}((200,20))
        initwithrand!(ar)
        ari = pinv(ar)
        @tensor id[1,2] := ari[1,-1] * ar[-1,2]
        @test toarray(id) ≈ I
    end

    @testset "Reshaping" begin
        ##without permutation
        f = DTensor{Complex{Float64},6}(ntuple(i -> 2, 6))
        fr = initwithrand!(deepcopy(f))
        frf, rs2 = fuselegs(fr, ((1,2),(3,4),(5,6)))
        frff, rs3 = fuselegs(frf, ((1,2,3),))
        @test abs(norm(frf)  - norm(fr)) < 1e-10
        @test abs(norm(frff) - norm(fr)) < 1e-10

        frffs = initwithzero!(deepcopy(frf))
        frffss = initwithzero!(deepcopy(fr))
        splitlegs!(frffs,  frff,  ((1,1,1),(1,1,2),(1,1,3)), rs3...)
        splitlegs!(frffss, frffs, ((1,1,1),(1,1,2),(2,2,1),(2,2,2),(3,3,1),(3,3,3)), rs2...)
        @test frffss ≈ fr

        ##with permutation
        f = DTensor{Complex{Float64},4}(ntuple(i -> 2, 4))
        fr = initwithrand!(deepcopy(f))
        frf, rs = fuselegs(fr, ((1,3),(2,4)))
        @test abs(norm(frf)  - norm(fr)) < 1e-10

        frfs = initwithzero!(deepcopy(fr))
        splitlegs!(frfs,  frf,  ((1,1,1),(2,2,1),(1,1,2),(2,2,2)), rs...)
        @test frfs ≈ fr
    end
end


function levelset_normal(levelset, p::V, invjac) where {V<:AbstractVector}
    g = vec(gradient(levelset, p))
    n = invjac .* g
    return n / norm(n)
end

function normalize_normals!(normals)
    dim,npts = size(normals)
    for i = 1:npts
        n = normals[:,i]
        normals[:,i] .= n/norm(n)
    end
end

function levelset_normal(levelset, points::M, invjac) where {M<:AbstractMatrix}
    npts = size(points)[2]
    if npts == 0
        return []
    else
        g = hcat([gradient(levelset, points[:, i])' for i = 1:npts]...)
        normals = diagm(invjac) * g
        normalize_normals!(normals)
        return normals
    end
end

function levelset_coefficients(distancefunc,mesh)
    nodalcoordinates = nodal_coordinates(mesh)
    return distancefunc(nodalcoordinates)
end

function reference_cell_area()
    return 4.0
end

function reference_face_normals()
    n1 = [0.,-1.]
    n2 = [1.,0.]
    n3 = [0.,1.]
    n4 = [-1.,0.]
    return [n1,n2,n3,n4]
end

function reference_bottom_face_midpoint()
    [0.0, -1.0]
end

function reference_right_face_midpoint()
    [1.0, 0.0]
end

function reference_top_face_midpoint()
    [0.0, 1.0]
end

function reference_left_face_midpoint()
    [-1.0, 0.0]
end

function reference_face(faceid)
    if faceid == 1
        return (2, -1.0)
    elseif faceid == 2
        return (1, +1.0)
    elseif faceid == 3
        return (2, +1.0)
    elseif faceid == 4
        return (1, -1.0)
    else
        error("Expected faceid ∈ {1,2,3,4}, got faceid = $faceid")
    end
end

function reference_face_midpoints()
    [
        reference_bottom_face_midpoint(),
        reference_right_face_midpoint(),
        reference_top_face_midpoint(),
        reference_left_face_midpoint(),
    ]
end

function extend_to_face(points, faceid)
    dir, coordval = reference_face(faceid)
    @assert dir == 1 || dir == 2
    flipdir = dir == 1 ? 2 : 1
    return extend([coordval], flipdir, points)
end

function opposite_face(faceid)
    if faceid == 1
        return 3
    elseif faceid == 2
        return 4
    elseif faceid == 3
        return 1
    elseif faceid == 4
        return 2
    else
        error("Expected faceid ∈ {1,2,3,4} got faceid = $faceid")
    end
end

function tangents(normals)
    rot = [
        0.0 -1.0
        1.0 0.0
    ]
    return rot * normals
end

function scale_area(normals,invjac)
    t = tangents(normals)
    den = sqrt.((t .^ 2)' * (invjac .^ 2))
    return 1.0 ./ den
end

function scale_area(cellmap::CellMap, normals)
    invjac = inverse_jacobian(cellmap)
    return scale_area(normals,invjac)
end

function cell_sign_to_row(s)
    (s == -1 || s == +1) || error("Use ±1 to index into rows (i.e. phase), got index = $s")
    row = s == +1 ? 1 : 2
    return row
end

function cell_couple_sign_to_row(s1,s2)
    (s1 == -1 || s1 == +1) || error("Expected sign ∈ {-1,1}, got sign $s1")
    (s2 == -1 || s2 == +1) || error("Expected sign ∈ {-1,1}, got sign $s2")
    if s1 == +1 && s2 == +1
        return 1
    elseif s1 == +1 && s2 == -1
        return 2
    elseif s1 == -1 && s2 == +1
        return 3
    elseif s1 == -1 && s2 == -1
        return 4
    end
end

function dimension(basis::TensorProductBasis{dim}) where {dim}
    return dim
end

function number_of_basis_functions(basis::TensorProductBasis{dim,T,NF}) where {dim,T,NF}

    return NF
end

function interpolation_points(basis::TensorProductBasis{D,T,N}) where {D,T,N}
    return basis.points
end

function number_of_symmetric_degrees_of_freedom(dim)
    if dim == 2
        return 3
    else
        error("Current support for dim = 2, got dim = $dim")
    end
end


using Flux.Optimise
using Flux.Optimise: runall
using Flux.Tracker
using Test
@testset "Optimise" begin
  w = randn(10, 10)
  @testset for opt in [ADAMW(), ADAGrad(0.1), AdaMax(), ADADelta(0.9), AMSGrad(),
                       NADAM(), Descent(0.1), ADAM(), Nesterov(), RMSProp(),
                       Momentum()]
    w′ = param(randn(10, 10))
    loss(x) = Flux.mse(w*x, w′*x)
    for t = 1: 10^5
      θ = Params([w′])
      θ̄ = gradient(() -> loss(rand(10)), θ)
      Optimise.update!(opt, θ, θ̄)
    end
    @test Flux.mse(w, w′) < 0.01
  end
end

@testset "Optimiser" begin
  w = randn(10, 10)
  @testset for Opt in [InvDecay, WeightDecay, ExpDecay]
    w′ = param(randn(10, 10))
    loss(x) = Flux.mse(w*x, w′*x)
    opt = Optimiser(Opt(), ADAM(0.001))
    for t = 1:10^5
      l = loss(rand(10))
      back!(l)
      delta = Optimise.apply!(opt, w′.data, w′.grad)
      w′.data .-= delta
    end
    @test Flux.mse(w, w′) < 0.01
  end
end

@testset "Training Loop" begin
  i = 0
  l = param(1)

  Flux.train!(() -> (sleep(0.1); i += 1; l),
              (),
              Iterators.repeated((), 100),
              Descent(),
              cb = Flux.throttle(() -> (i > 3 && Flux.stop()), 1))

  @test 3 < i < 50

  # Test multiple callbacks
  x = 0
  fs = [() -> (), () -> x = 1]
  cbs = runall(fs)
  cbs()
  @test x == 1
end

@testset "ExpDecay" begin
    w = randn(10, 10)
    o = ExpDecay(0.1, 0.1, 1000, 1e-4)
    w1 = param(randn(10,10))
    loss(x) = Flux.mse(w*x, w1*x)
    flag = 1
    decay_steps = []
    for t = 1:10^5
      l = loss(rand(10))
      back!(l)
      prev_eta = o.eta
      prev_grad = collect(w1.grad)
      delta = Optimise.apply!(o, w1.data, w1.grad)
      w1.data .-= delta
      new_eta = o.eta
      if new_eta != prev_eta
        push!(decay_steps, t)
      end
      array = fill(o.eta, size(prev_grad))
      if array .* prev_grad != delta
        flag = 0
      end
    end
    @test flag == 1
    # Test to check if decay happens at decay steps. Eta reaches clip value eventually.
    ground_truth = []
    for i in 1:11
      push!(ground_truth, 1000*i)  # Expected decay steps for this example.
    end
    @test decay_steps == ground_truth
    @test o.eta == o.clip
end


module TwoBodyUnitTests

using Test
using UnitfulAstrodynamics
using UnitfulAstrodynamics.TwoBody.Systems: Earth

@testset "Transforms" begin

    rᵢ = [0.0, 11681.0, 0.0] * u"km"
    vᵢ = [5.134, 4.226, 2.787] * u"km/s"
    orbit = CartesianOrbit(rᵢ, vᵢ, Earth, 0.0u"s") |> KeplerianOrbit

    @test semimajor_axis(orbit) == 24509.265399338536 * u"km"
    @test eccentricity(orbit) == 0.723452708202361
    @test inclination(orbit) == 151.50460766373865 * u"°"
    @test true_anomoly(orbit) == 89.99652573907436  * u"°"

    @test orbit ≈ CartesianOrbit(orbit)

    e      =  0.3
    a      =  semimajor_axis(orbit)
    i      =  10.      * u"°"
    Ω      =  0.       * u"°"
    ω      =  10.      * u"°"
    ν      =  0.       * u"°"
    orbit  =  KeplerianOrbit(e, a, i, Ω, ω, ν, Earth, 0.0u"s")

    @test isapprox(orbit, CartesianOrbit(orbit); atol=1e-6)

end

@testset "Kepler" begin

    rᵢ = [0.0, 11681.0, 0.0]u"km"
    vᵢ = [5.134, 4.226, 2.787]u"km/s"
    orbit = CartesianOrbit(rᵢ, vᵢ, Earth, 0.0u"s")

    @test kepler(orbit) ≈ orbit

end

@testset "Lambert" begin

    rᵢ = [0.0, 11681.0, 0.0]u"km"
    vᵢ = [5.134, 4.226, 2.787]u"km/s"
    initial = CartesianOrbit(rᵢ, vᵢ, Earth, 0.0u"s")

    Δt = 1000u"s"
    final = kepler(initial, Δt; tol=1e-12)

    v₁, v₂ = lambert(position_vector(initial), position_vector(final), mass_parameter(Earth), Δt, :short; tol=1e-6, max_iter=1000)

    @test isapprox(ustrip.(u"km/s", v₁), ustrip.(u"km/s", velocity_vector(initial)); atol=1e-6)
    @test isapprox(ustrip.(u"km/s", v₂), ustrip.(u"km/s", velocity_vector(final)); atol=1e-6)

end

end


function compute_L(file, species)
    w, (x,y,z), px, py, pz = file["weight/"*species,
                                  "grid/"*species,
                                  "px/"*species,
                                  "py/"*species,
                                  "pz/"*species]

    r = SVector{3}.(x,y,z)
    p = VectorVariable(px, py, pz)

    L = w .* r .× p
end

compute_L(dir::Symbol, file, species) = compute_L(Val(dir), file, species)

function compute_L(::Val{:x}, file, species)
    w, r, py, pz = file["weight/"*species,
                        "grid/"*species,
                        "py/"*species,
                        "pz/"*species]

    y = r[2]
    z = r[3]
    Lx = @. w * (y * pz - z * py)
end


using Revise
using MendelIHT
using SnpArrays
using Random
using GLM
using DelimitedFiles
using Distributions
using LinearAlgebra
using CSV
using DataFrames
using StatsBase
using TraitSimulation
BLAS.set_num_threads(1)

"""
    βi ~ N(0, 0.1) chosen randomly across genome

k = Number of causal SNPs
p = Total number of SNPs
traits = Number of traits (phenotypes)
overlap = number of causal SNPs shared in each trait
"""
function simulate_random_beta(k::Int, p::Int, traits::Int; overlap::Int=0, βσ=1.0)
    d = Normal(0, βσ)
    true_b = zeros(p, traits)
    if overlap == 0
        causal_snps = sample(1:(traits * p), k, replace=false)
        true_b[causal_snps] = rand(d, k)
    else
        shared_snps = sample(1:p, overlap, replace=false)
        weight_vector = aweights(1 / (traits * (p - overlap)) * ones(traits * p))
        for i in 1:traits
            weight_vector[i*shared_snps] .= 0.0 # avoid sampling from shared snps
        end
        @assert sum(weight_vector) ≈ 1.0 "sum(weight_vector) = $(sum(weight_vector)) != 1.0"
        # simulate β for shared predictors
        for i in 1:traits
            true_b[shared_snps, i] = rand(d, overlap)
        end
        # simulate β for none shared predictors
        nonshared_snps = sample(1:(traits * p), weight_vector, k - (traits * overlap), replace=false)
        true_b[nonshared_snps] = rand(d, k - (traits * overlap))
    end

    return true_b
end

"""
    βi ~ Uniform([-0.5, -0.45, ..., -0.05, 0.05, ..., 0.5]) chosen uniformly across genome

k = Number of causal SNPs
p = Total number of SNPs
traits = Number of traits (phenotypes)
overlap = number of pleiotropic SNPs (affects each trait with effect size βi / r)
"""
function simulate_fixed_beta(k::Int, p::Int, traits::Int; overlap::Int=0)
    true_b = zeros(p, traits)
    effect_sizes = collect(0.05:0.05:0.5)
    # num_causal_snps = (k - (traits - 1) * overlap) :: Int
    # @assert num_causal_snps > 0 "number of causal SNPs should be positive but was $num_causal_snps"
    # idx_causal_snps = sample(1:p, num_causal_snps, replace=false)
    # @assert length(idx_causal_snps) == num_causal_snps "length(idx_causal_snps) = $(length(idx_causal_snps)) != num_causal_snps"
    # shuffle!(idx_causal_snps)
    k_indep = k - 2overlap # pleiotropic SNPs affect 2 phenotypes
    num_causal_snps = k_indep + overlap
    @assert num_causal_snps > 0 "number of causal SNPs should be positive but was $num_causal_snps"
    idx_causal_snps = sample(1:p, num_causal_snps, replace=false)
    @assert length(idx_causal_snps) == num_causal_snps "length(idx_causal_snps) = $(length(idx_causal_snps)) != num_causal_snps"
    shuffle!(idx_causal_snps)

    # pleiotropic SNPs affect 2 phenotypes
    for i in 1:overlap
        j = idx_causal_snps[i]
        rs = sample(1:traits, 2, replace=false)
        for r in rs
            true_b[j, r] = rand(-1:2:1) * effect_sizes[rand(1:10)]
        end
    end
    # non pleiotropic SNPs affect only 1 phenotype
    for i in (overlap+1):length(idx_causal_snps)
        idx = idx_causal_snps[i]
        true_b[idx, rand(1:traits)] = rand(-1:2:1) * effect_sizes[rand(1:10)]
    end

    @assert count(!iszero, true_b) == k "count(!iszero, true_b) = $(count(!iszero, true_b)) != k = $k"

    return true_b
end

"""
    every causal SNP affects every phenotype, chosen fixed distance with each other across genome
    Their total effect size sums up to 0.5

k = Number of causal SNPs (each SNP affects every trait)
p = Total number of SNPs
traits = Number of traits (phenotypes)
"""
function simulate_pleiotropic_beta(k::Int, p::Int, traits::Int)
    true_b = zeros(p, traits)
    effect_size = 0.5 / traits
    idx_causal_snps = collect(div(p, k):div(p, k):p)
    shuffle!(idx_causal_snps)

    # pleiotropic SNPs affect every phenotype
    for i in idx_causal_snps
        for r in 1:traits
            true_b[i, r] = effect_size
        end
    end

    @assert count(!iszero, true_b) == k*traits "count(!iszero, true_b) should be equal to $(k*traits)"

    return true_b
end

"""
# Arguments
xla = simulated genotype matrix (converted to a SnpLinAlg)
k = number of causal SNPs
r = number of traits
Φ = estimated GRM (using GEMMA)

# Optional arguments
seed = random seed for reproducibility
σ2 = contribution of GRM
σe = random environmental effect
βoverlap = number of causal SNPs shared in all traits
"""
function simulate_multivariate_polygenic(
    plinkname::String, n::Int, p::Int, k::Int, r::Int;
    seed::Int=2021, σg=0.6, σe=0.4, βoverlap=2
    )
    # set seed
    Random.seed!(seed)

    # simulate `.bed` file with no missing data
    x = simulate_random_snparray("sim$seed/" * plinkname * ".bed", n, p)
    xla = SnpLinAlg{Float64}(x, model=ADDITIVE_MODEL, impute=true, center=true, scale=true)
    make_bim_fam(xla, "sim$seed/" * plinkname)

    # intercept is the only nongenetic covariate
    Z = ones(n, 1)
    intercepts = zeros(r)' # each trait have 0 intercept

    # simulate β
    # B = simulate_pleiotropic_beta(k, p, r)
    B = simulate_fixed_beta(k, p, r, overlap=βoverlap)
    writedlm("sim$(seed)/trueb.txt", B)

    # between trait covariance matrix
    Σ = random_covariance_matrix(r)

    # between sample covariance is identity + GRM (2x because OpenMendel always uses half the GRM)
    run(`./gemma -bfile sim$(seed)/$plinkname -gk 1 -o $plinkname`)
    run(`mv ./output/$(plinkname).cXX.txt sim$(seed)`)
    Φ = readdlm("sim$(seed)/" * plinkname * ".cXX.txt")
    V = σg * Φ + σe * I

    # simulate y using TraitSimulations.jl (https://github.com/OpenMendel/TraitSimulation.jl/blob/master/src/modelframework.jl#L137)
    vc = @vc Σ ⊗ V
    μ = zeros(n, r)
    μ_null = zeros(n, r)
    LinearAlgebra.mul!(μ_null, Z, intercepts)
    mul!(μ, xla, B)
    BLAS.axpby!(1.0, μ_null, 1.0, μ)
    VCM_model = VCMTrait(Z, intercepts, xla, B, vc, μ)
    Y = Matrix(Transpose(simulate(VCM_model)))

    # simulate using Distributions.jl
    # μ = z * intercepts + xla * B
    # Y = rand(MatrixNormal(μ', Σ, V))

    return xla, Matrix(Z'), B, Σ, Y
end

function simulate_multivariate_sparse(
    plinkname::String, n::Int, p::Int, k::Int, r::Int;
    seed::Int=2021, σg=0.6, σe=0.4, βoverlap=2
    )
    # set seed
    Random.seed!(seed)

    # simulate `.bed` file with no missing data
    x = simulate_random_snparray("sim$seed/" * plinkname * ".bed", n, p)
    xla = SnpLinAlg{Float64}(x, model=ADDITIVE_MODEL, impute=true, center=true, scale=true)
    make_bim_fam(xla, "sim$seed/" * plinkname)

    # intercept is the only nongenetic covariate
    Z = ones(n, 1)
    intercepts = zeros(r)' # each trait have 0 intercept

    # simulate β
    # B = simulate_pleiotropic_beta(k, p, r)
    B = simulate_fixed_beta(k, p, r, overlap=βoverlap)
    writedlm("sim$(seed)/trueb.txt", B)

    # between trait covariance matrix (but GEMMA GRM still needed)
    run(`./gemma -bfile sim$(seed)/$plinkname -gk 1 -o $plinkname`)
    run(`mv ./output/$(plinkname).cXX.txt sim$(seed)`)
    Σ = random_covariance_matrix(r)
    writedlm("sim$(seed)/trueCovariance.txt", Σ)

    # simulate multivariate normal phenotype for each sample
    μ = xla * B + Z*intercepts

    # simulate response
    Y = zeros(n, r)
    for i in 1:n
        μi = @view(μ[i, :])
        Y[i, :] = rand(MvNormal(μi, Σ))
    end

    return xla, Matrix(Z'), B, Σ, Matrix(Y')
end

"""
simulate under IHT's model
"""
function simulate_NFBC1966_sparse(
    plinkname::String, k::Int, r::Int;
    seed::Int=2021, σg=0.6, σe=0.4, βoverlap=2
    )
    # set seed
    Random.seed!(seed)

    # simulate `.bed` file with no missing data
    x = SnpArray(plinkname * ".bed")
    xla = SnpLinAlg{Float64}(x, model=ADDITIVE_MODEL, impute=true, center=true, scale=true)
    n, p = size(xla)

    # intercept is the only nongenetic covariate
    Z = ones(n, 1)
    intercepts = zeros(r)' # each trait have 0 intercept

    # simulate β
    # B = simulate_pleiotropic_beta(k, p, r)
    B = simulate_fixed_beta(k, p, r, overlap=βoverlap)
    writedlm("sim$(seed)/trueb.txt", B)

    # between trait covariance matrix
    Σ = random_covariance_matrix(r)
    writedlm("sim$(seed)/trueCovariance.txt", Σ)

    # simulate multivariate normal phenotype for each sample
    μ = xla * B + Z*intercepts

    # simulate response
    Y = zeros(n, r)
    for i in 1:n
        μi = @view(μ[i, :])
        Y[i, :] = rand(MvNormal(μi, Σ))
    end

    return xla, Matrix(Z'), B, Σ, Matrix(Y')
end

"""
Trait covariance matrix is σg * Φ + σe * I where Φ is the GRM.
"""
function simulate_NFBC1966_polygenic(
    plinkname::String, k::Int, r::Int;
    seed::Int=2021, σg=0.1, σe=0.9, βoverlap=2,
    )
    # set seed
    Random.seed!(seed)

    # simulate `.bed` file with no missing data
    x = SnpArray(plinkname * ".bed")
    xla = SnpLinAlg{Float64}(x, model=ADDITIVE_MODEL, impute=true, center=true, scale=true)
    n, p = size(x)

    # intercept is the only nongenetic covariate
    Z = ones(n, 1)
    intercepts = zeros(r)' # each trait have 0 intercept

    # simulate β
    # B = simulate_pleiotropic_beta(k, p, r)
    B = simulate_fixed_beta(k, p, r, overlap=βoverlap)
    writedlm("sim$(seed)/trueb.txt", B)

    # between trait covariance matrix
    Σ = random_covariance_matrix(r)
    writedlm("sim$(seed)/true_cov.txt", Σ)

    # between sample covariance is identity + GRM
    Φ = readdlm(plinkname * ".cXX.txt")
    V = σg * Φ + σe * I

    # simulate y using TraitSimulations.jl (https://github.com/OpenMendel/TraitSimulation.jl/blob/master/src/modelframework.jl#L137)
    vc = @vc Σ ⊗ V
    μ = zeros(n, r)
    μ_null = zeros(n, r)
    LinearAlgebra.mul!(μ_null, Z, intercepts)
    mul!(μ, xla, B)
    BLAS.axpby!(1.0, μ_null, 1.0, μ)
    VCM_model = VCMTrait(Z, intercepts, xla, B, vc, μ)
    Y = Matrix(Transpose(simulate(VCM_model)))

    # simulate using Distributions.jl
    # μ = z * intercepts + xla * B
    # Y = rand(MatrixNormal(μ', Σ, V))

    return xla, Matrix(Z'), B, Σ, Y
end

function make_bim_fam(x::AbstractMatrix, name::String)
    n, p = size(x)

    #create .bim file structure: https://www.cog-genomics.org/plink2/formats#bim
    open(name * ".bim", "w") do f
        for i in 1:p
            write(f, "1\tsnp$i\t0\t1\t1\t2\n")
        end
    end

    open(name * ".fam", "w") do f
        for i in 1:n
            write(f, "$i\t1\t0\t0\t1\t1.0\n")
        end
    end
end

function make_GEMMA_fam_file(x::AbstractMatrix, y::AbstractVecOrMat, name::String)
    ly = size(y, 1)
    n, p = size(x)

    # put 1st phenotypes in 6th column, 2nd phenotype in 7th column ... etc
    traits = size(y, 1)
    open(name * ".fam", "w") do f
        for i in 1:n
            write(f, "$i\t1\t0\t0\t1")
            for j in 1:traits
                write(f, "\t$(y[j, i])")
            end
            write(f, "\n")
        end
    end
end

function make_MVPLINK_fam_and_phen_file(x::AbstractMatrix, y::AbstractVecOrMat, name::String)
    ly = size(y, 1)
    n, p = size(x)

    # put a random phenotype in fam file
    traits = size(y, 1)
    open(name * ".fam", "w") do f
        for i in 1:n
            println(f, "$i\t1\t0\t0\t1\t1")
        end
    end

    # save phenotypes in separate `.phen` file
    open(name * ".phen", "w") do io
        print(io, "FID\tIID")
        for j in 1:traits
            print(io, "\tT$j")
        end
        print(io, "\n")
        for i in 1:n
            print(io, "$i\t1")
            for j in 1:traits
                write(io, "\t$(y[j, i])")
            end
            print(io, "\n")
        end
    end
end

"""
Computes power and false positive rates
- p: total number of SNPs
- pleiotropic_snps: Indices (or ID) of the true causal SNPs that affect >1 phenotype
- independent_snps: Indices (or ID) of the true causal SNPs that affect exactly 1 phenotype
- signif_snps: Indices (or ID) of SNPs that are significant after testing

returns: pleiotropic SNP's power, independent SNP's power, number of false positives, and false positive rate
"""
function power_and_fpr(p::Int, pleiotropic_snps, independent_snps, signif_snps)
    pleiotropic_power = length(signif_snps ∩ pleiotropic_snps) / length(pleiotropic_snps)
    independent_power = length(signif_snps ∩ independent_snps) / length(independent_snps)
    correct_snps = pleiotropic_snps ∪ independent_snps
    FP = length(signif_snps) - length(signif_snps ∩ correct_snps) # number of false positives
    TN = p - length(signif_snps) # number of true negatives
    FPR = FP / (FP + TN)
    return pleiotropic_power, independent_power, FP, FPR
end

# https://github.com/OpenMendel/MendelPlots.jl/blob/master/src/gwasplots.jl#L108
function genomic_inflation(pvalues::AbstractVecOrMat)
    λ = median(quantile.(Chisq(1), 1 .- pvalues) ./ quantile(Chisq(1), 0.5))
end

"""
- filename: gemma's output file name
- pleiotropic_snpid: ID of causal SNPs that affect > 1 phenotypes
- independent_snpid: ID of causal SNPs that affect 1 phenotypes

returns: power, number of false positives, and false positive rate
"""
function process_gemma_result(filename, pleiotropic_snpid, independent_snpid)
    # read GEMMA result
    gemma_df = CSV.read(filename, DataFrame)
    snps = size(gemma_df, 1)

    # pvalues
    pval_wald = gemma_df[!, :p_wald]
#    pval_lrt = gemma_df[!, :p_lrt]
#     pval_score = gemma_df[!, :p_score]

    # estimated beta
#     estim_β1 = gemma_df[!, :beta_1]
#     estim_β2 = gemma_df[!, :beta_2]

#     # estimated covariance matrix
#     estim_σ11 = gemma_df[!, :Vbeta_1_1]
#     estim_σ12 = gemma_df[!, :Vbeta_1_2]
#     estim_σ22 = gemma_df[!, :Vbeta_2_2];

    # check how many real SNPs were recovered
    signif_snps = findall(x -> x ≤ 0.05 / snps, pval_wald) # gemma's selected snps
    signif_snpid = gemma_df[signif_snps, :rs]

    # return power, false positives, false positive rate, and genomic inflation
    pleiotropic_power, independent_power, FP, FPR = power_and_fpr(snps,
        pleiotropic_snpid, independent_snpid, signif_snpid)
    λ = genomic_inflation(pval_wald)
    return pleiotropic_power, independent_power, FP, FPR, λ
end

"""
- filename: mvPLINK's output file name
- pleiotropic_snps: indices for causal SNPs that affect > 1 phenotypes
- independent_snps: indices for causal SNPs that affect 1 phenotypes

returns: power, number of false positives, and false positive rate
"""
function process_mvPLINK(filename, pleiotropic_snps, independent_snps)
    # read mvPLINK result
    mvplink_df = CSV.read(filename, DataFrame, delim=' ', ignorerepeated=true)
    snps = size(mvplink_df, 1)

    # get pvalues, possibly accounting for "NA"s
    if eltype(mvplink_df[!, :P]) == Float64
        pval = mvplink_df[!, :P]
    else
        mvplink_df[findall(x -> x == "NA", mvplink_df[!, :P]), :P] .= "1.0"
        pval = parse.(Float64, mvplink_df[!, :P])
    end

    # SNPs passing threshold
    signif_snps = findall(x -> x ≤ 0.05 / snps, pval)

    # compute power, false positives, and false positive rate
    pleiotropic_power, independent_power, FP, FPR = power_and_fpr(snps,
        pleiotropic_snps, independent_snps, signif_snps)
    λ = genomic_inflation(pval)
    return pleiotropic_power, independent_power, FP, FPR, λ
end

function one_NFBC_simulation(
    k::Int, r::Int;
    seed::Int=2021, σg=0.1, σe=0.9, βoverlap=2,
    path=5:5:50, init_beta=false, model=:polygenic, debias=100
    )
    isdir("sim$seed") ? (return nothing) : mkdir("sim$seed")
    plinkname = "NFBC.qc.imputeBy0.chr.1"

    # simulate data
    Random.seed!(seed)
    if model == :polygenic
        xla, Z, B, Σ, Y = simulate_NFBC1966_polygenic("../" * plinkname, k, r,
            seed=seed, σg=σg, σe=σe, βoverlap=βoverlap)
    elseif model == :sparse
        xla, Z, B, Σ, Y = simulate_NFBC1966_sparse("../" * plinkname, k, r,
            seed=seed, σg=σg, σe=σe, βoverlap=βoverlap)
    else
        error("model misspecified!")
    end

    correct_snps = unique([x[1] for x in findall(!iszero, B)])
    pleiotropic_snps, independent_snps = Int[], Int[]
    for snp in correct_snps
        count(x -> abs(x) > 0, @view(B[snp, :])) > 1 ?
            push!(pleiotropic_snps, snp) : push!(independent_snps, snp)
    end
    snpdata = SnpData("../" * plinkname)
    pleiotropic_snp_rsid = snpdata.snp_info[pleiotropic_snps, :snpid]
    independent_snp_rsid = snpdata.snp_info[independent_snps, :snpid]

    # run GEMMA (GRM is precomputed already)
    cd("sim$seed")
    run(`cp ../../$(plinkname).bed .`)
    run(`cp ../../$(plinkname).bim .`)
    run(`cp ../../$(plinkname).cXX.txt .`)
    make_GEMMA_fam_file(xla, Y, plinkname)
    pheno_columns = [string(ri) for ri in 1:r]
    gemma_time = @elapsed begin
        run(`../../gemma -bfile $plinkname -k $(plinkname).cXX.txt -notsnp -lmm 1 -n $pheno_columns -o gemma.sim$seed`)
    end
    gemma_pleiotropic_power, gemma_independent_power, gemma_FP, gemma_FPR, gemma_λ =
        process_gemma_result("output/gemma.sim$seed.assoc.txt", pleiotropic_snp_rsid, independent_snp_rsid)
    println("GEMMA time = $gemma_time, pleiotropic power = $gemma_pleiotropic_power, independent power = $gemma_independent_power, FP = $gemma_FP, FPR = $gemma_FPR, gemma_λ=$gemma_λ")
    mv("output/gemma.sim$seed.assoc.txt", "gemma.sim$seed.assoc.txt")
    mv("output/gemma.sim$seed.log.txt", "gemma.sim$seed.log.txt")

    # run multivariate IHT
    mIHT_time = @elapsed begin
        mses = cross_validate(plinkname, MvNormal, path=path, phenotypes=collect(1:r).+5;
            init_beta=init_beta, debias=debias)
        k_rough_guess = path[argmin(mses)]
        dense_path = (k_rough_guess - 4):(k_rough_guess + 4)
        mses_new = cross_validate(plinkname, MvNormal, path=dense_path, phenotypes=collect(1:r).+5;
            init_beta=init_beta, debias=debias, cv_summaryfile="miht.cviht.summary.txt")
        iht_result = iht(plinkname, dense_path[argmin(mses_new)], MvNormal, phenotypes=collect(1:r).+5;
            init_beta=init_beta, debias=debias, summaryfile="miht.summary.txt")
    end
    detected_snps = Int[]
    for i in 1:r # save each beta separately
        β = iht_result.beta[i, :]
        detected_snps = detected_snps ∪ findall(!iszero, β)
        writedlm("multivariate_iht_beta$i.txt", β)
    end
    mIHT_pleiotropic_power, mIHT_independent_power, mIHT_FP, mIHT_FPR = power_and_fpr(size(B, 1), pleiotropic_snps, independent_snps, detected_snps)
    println("multivariate IHT time = $mIHT_time, pleiotropic power = $mIHT_pleiotropic_power, independent power = $mIHT_independent_power, FP = $mIHT_FP, FPR = $mIHT_FPR")

    # run multiple univariate IHT
    detected_snps = Int[]
    uIHT_time = @elapsed begin
        for trait in 1:r
            mses = cross_validate(plinkname, Normal, path=path, phenotypes=trait+5;
                init_beta=init_beta, debias=debias)
            k_rough_guess = path[argmin(mses)]
            dense_path = (k_rough_guess == 5) ? (0:5) : ((k_rough_guess - 4):(k_rough_guess + 4))
            mses_new = cross_validate(plinkname, Normal, path=dense_path, phenotypes=trait+5;
                init_beta=init_beta, debias=debias, cv_summaryfile="uiht.cviht.summary$trait.txt")
            best_k = dense_path[argmin(mses_new)]
            if best_k > 0
                iht_result = iht(plinkname, best_k, Normal, phenotypes=trait+5;
                    init_beta=init_beta, debias=debias, summaryfile="uiht.summary$trait.txt")
                β = iht_result.beta
            else
                β = zeros(size(B, 2))
            end

            # save results
            detected_snps = detected_snps ∪ findall(!iszero, β)
            writedlm("univariate_iht_beta$trait.txt", β)
        end
    end
    uIHT_pleiotropic_power, uIHT_independent_power, uIHT_FP, uIHT_FPR = power_and_fpr(size(B, 1), pleiotropic_snps, independent_snps, detected_snps)
    println("univariate IHT time = $uIHT_time, pleiotropic power = $uIHT_pleiotropic_power, independent power = $uIHT_independent_power, FP = $uIHT_FP, FPR = $uIHT_FPR")

    # run MVPLINK
    phenofile = plinkname * ".phen"
    make_MVPLINK_fam_and_phen_file(xla, Y, plinkname)
    mvplink_time = @elapsed run(`../../plink.multivariate --bfile $plinkname --noweb --mult-pheno $phenofile --mqfam`)
    mvPLINK_pleitropic_power, mvPLINK_independent_power, mvPLINK_FP, mvPLINK_FPR, mvPLINK_λ =
        process_mvPLINK("plink.mqfam.total", pleiotropic_snps, independent_snps)
    println("mvPLINK time = $mvplink_time, pleiotropic power = $mvPLINK_pleitropic_power, independent power = $mvPLINK_independent_power, FP = $mvPLINK_FP, FPR = $mvPLINK_FPR, mvPLINK_λ=$mvPLINK_λ")

    # clean up
    # rm("plink.hh", force=true)
    rm("$(plinkname).fam", force=true)
    rm("$(plinkname).bed", force=true)
    rm("$(plinkname).bim", force=true)
    rm("$(plinkname).cXX.txt", force=true)

    # save summary stats
    n, p = size(xla)
    open("summary.txt", "w") do io
        println(io, "Simulation $seed summary")
        println(io, "n = $n, p = $p, k = $k, r = $r, βoverlap=$βoverlap")
        println(io, "debias=$debias, init_beta=$init_beta")
        model == :polygenic ? println(io, "model = $model, σg=$σg, σe=$σe") : println(io, "model = $model")
        println(io, "")
        println(io, "mIHT time = $mIHT_time seconds, pleiotropic power = $mIHT_pleiotropic_power, independent power = $mIHT_independent_power, FP = $mIHT_FP, FPR = $mIHT_FPR, λ = NaN")
        println(io, "uIHT time = $uIHT_time seconds, pleiotropic power = $uIHT_pleiotropic_power, independent power = $uIHT_independent_power, FP = $uIHT_FP, FPR = $uIHT_FPR, λ = NaN")
        println(io, "mvPLINK time = $mvplink_time seconds, pleiotropic power = $mvPLINK_pleitropic_power, independent power = $mvPLINK_independent_power, FP = $mvPLINK_FP, FPR = $mvPLINK_FPR, λ = $mvPLINK_λ")
        println(io, "GEMMA time = $gemma_time seconds, pleiotropic power = $gemma_pleiotropic_power, independent power = $gemma_independent_power, FP = $gemma_FP, FPR = $gemma_FPR, λ = $gemma_λ")
    end
    cd("../")

    return nothing
end

#
# polygenic/independent beta = sign * Uniform{0.05, …, 0.25}
# max condition number = 10
# q = 5, iterates ≥5 times, init_beta=true, debias=false)
# sim set 1 are for k = 10, r = 2, βoverlap = 3 each affects 2 traits, path = 5:5:50 (then search around best k)
# sim set 2 are for k = 20, r = 3, βoverlap = 5 each affects 2 traits, path = 5:5:50 (then search around best k)
# sim set 3 are for k = 30, r = 5, βoverlap = 7 each affects 2 traits, path = 5:5:50 (then search around best k)
# sim set 4 are for k = 10, r = 10, βoverlap = 3 each affects 2 traits, path = 5:5:50 (then search around best k)
# sim set 5 are for k = 20, r = 50, βoverlap = 5 each affects 2 traits, path = 5:5:50 (then search around best k)
# sim set 6 are for k = 30, r = 100, βoverlap = 7 each affects 2 traits, path = 5:5:50 (then search around best k)
#

function run_simulation(set::Int, model::Symbol)
    σg = 0.1
    σe = 0.9
    init_beta = true
    debias = 100
    path = 5:5:50
    βoverlap = [3, 5, 7, 3, 5, 7]
    k = [10, 20, 30, 10, 20, 30]
    r = [2, 3, 5, 10, 50, 100]

    println("Simulation model = $model, set $set has k = $(k[set]), r = $(r[set]), βoverlap = $(βoverlap[set])")

    cur_dir = pwd() * "/set$set"
    isdir(cur_dir) || mkdir(cur_dir)
    k_cur = k[set]
    r_cur = r[set]
    βoverlap_cur = βoverlap[set]

    for seed in 1:100
        try
            cd(cur_dir)
            one_NFBC_simulation(k_cur, r_cur, seed = seed, path = path, βoverlap=βoverlap_cur,
                σg=σg, σe=σe, init_beta=init_beta, model=model, debias=debias)
        catch e
            bt = catch_backtrace()
            msg = sprint(showerror, e, bt)
            println("set $set sim $seed threw an error!")
            println(msg)
            continue
        end
    end
end

set = parse(Int, ARGS[1])
model = :polygenic
run_simulation(set, model)


__precompile__()

module BayesianSumProductNetworks

    import Base.eltype
    import Base.length
    import Base.rand
    import Base.copy

    # Required packages
    using Reexport
    @reexport using Distributions
    using AxisArrays
    using SpecialFunctions
    using StatsFuns
    import StatsFuns.normlogpdf

    # Required base packages.
    using Statistics
    using Random
    using SparseArrays

    # Exports
    export NormalInverseGamma
    export DirichletSufficientStats, NormalInvGammaSufficientStats
    export GammaSufficienStats, BetaSufficientStats
    export AbstractSufficientStats
    export postparams, update!
    export RegionGraphNode, PartitionGraphNode
    export FactorizedDistributionGraphNode, AbstractRegionGraphNode
    export FactorizedMixtureGraphNode, InfiniteSumNode
    export haschildren, topoligicalorder, id
    export scope, randomscopes!, nscope, hasscope, setscope!, updatescope!
    export haschildren, children
    export logpdf!, logmllh!, logpdf

    ######################
    ## Type Definitions ##
    ######################

    abstract type AbstractSufficientStats end

    """
        Gamma prior sufficient stats.
    """
    mutable struct GammaSufficienStats <: AbstractSufficientStats
        n::Int
        sn::Float64
    end

    GammaSufficienStats() = GammaSufficienStats(0, 0.0)

    function update!(s::GammaSufficienStats, x::AbstractVector)
        if isempty(x)
            s.n = 0
            s.sn = 0.0
        else
            @assert all(x .>= 0) "Not all positive, $x"
            s.n = length(x)
            s.sn = sum(x)
        end
        return s
    end

    @inline postparams(d::Gamma, T::GammaSufficienStats) = (d.α+T.sn, (1/d.θ)+T.n)

    """
        Beta prior sufficient stats.
    """
    mutable struct BetaSufficientStats <: AbstractSufficientStats
        n::Int
        count::Int
    end

    function update!(s::BetaSufficientStats, x::T) where {T<:AbstractVector{<:Real}}
        if isempty(x)
            s.n = 0
            s.count = 0
        else
            @assert all((y) -> isone(y) | iszero(y), x)
            s.n = length(x)
            s.count = sum(x)
        end
        return s
    end
    @inline BetaSufficientStats() = BetaSufficientStats(0,0)
    @inline postparams(d::Beta, T::BetaSufficientStats) = (T.count+d.α,(T.n-T.count)+d.β)

    """
        Dirichlet prior sufficient stats.
    """
    struct DirichletSufficientStats <: AbstractSufficientStats
        counts::Vector{Int}
    end
    function update!(s::DirichletSufficientStats, x::AbstractVector)
        if isempty(x)
            fill!(s.counts, 0)
        else
            s.counts[:] = map(k -> sum(x .== k), 1:length(s.counts))
        end
        return s
    end
    @inline DirichletSufficientStats(n::Int) = DirichletSufficientStats(zeros(Int, n))
    @inline postparams(d::Dirichlet, T::DirichletSufficientStats) = T.counts.+d.alpha

    ## Gaussian with Normal Inverse Gamma prior ##
    struct NormalInverseGamma <: Distribution{Multivariate,Continuous}
        μ::Float64
        ν::Float64
        a::Float64
        b::Float64
    end

    length(d::NormalInverseGamma) = 2

    function logpdf(d::NormalInverseGamma, θ...)
        (μ, σ²) = θ
        l = _invgammalogpdf(d.a, d.b, σ²)
        l += normlogpdf(d.μ, sqrt(σ²) / d.ν, μ)
        return l
    end
    @inline _invgammalogpdf(a::Real, b::Real, x::Real) = a*log(b)-lgamma(a)-(a+1)*log(x)-b/x

    """
        Normal-Inverse-Gamma prior sufficient stats.
    """
    mutable struct NormalInvGammaSufficientStats <: AbstractSufficientStats
        n::Int
        s::Float64
        sdiff::Float64
    end

    NormalInvGammaSufficientStats() = NormalInvGammaSufficientStats(0,0.0,0.0)

    function update!(s::NormalInvGammaSufficientStats, x::AbstractVector)
        if isempty(x)
            s.n = 0
            s.s = 0
            s.sdiff = 0
        else
            s.n = length(x)
            s.s = sum(x)
            s.sdiff = sum((x .- mean(x)).^2)
        end
        s
    end

    function rand(d::NormalInverseGamma)
        σ² = rand(InverseGamma(d.a, d.b))
        μ = rand(Normal(d.μ, sqrt(σ² / d.ν )))
        return [μ, σ²]
    end

    @inline params(d::NormalInverseGamma) = (d.μ, d.ν, d.a, d.b)
    function postparams(d::NormalInverseGamma, T::NormalInvGammaSufficientStats)

        if T.n == 0
            return (d.μ, d.ν, d.a, d.b)
        end

        μ = (d.ν*d.μ + T.s) / (d.ν + T.n)
        ν = d.ν + T.n
        a = d.a + (T.n / 2)
        b = d.b
        b += T.sdiff/2 + ((T.n * d.ν) / (d.ν + T.n)) * (((T.s / T.n) - d.μ)^2 / 2)

        return (μ, ν, a, b)
    end

    # include implementations
    include("mathutil.jl")
    include("distUtils.jl")
    include("templateStructure.jl")
    include("regiongraph.jl")
    include("nonparametrics.jl")
    include("ancestralsampler.jl")
    include("initialization.jl")
    include("structureinference.jl")

end # module


using LinearAlgebra

_normalize(P, normalized) = normalized ? P : normalize([
    P[1, 1] P[1, 2]
    P[2, 1] P[2, 2]
], 1)

_papb(P) = P[1, 1] + P[1, 2], P[1, 1] + P[2, 1]

_D_max(P, D) = let (pa, pb) = _papb(P)
    D < 0 ?
    max(-pa * pb, -(1 - pa) * (1 - pb)) :
    min(pa * (1 - pb), (1 - pa) * pb)
end

"""
Coefficient of Linkage Disequilibrium

`P` is a matrix for the (not necessary normalized) frequency of haplotypes for two loci A and B.
"""
function D(P; normalized = false)
    P = _normalize(P, normalized)
    pa, pb = _papb(P)
    P[1, 1] - pa * pb
end

"""
A normalized measure of LD.
"""
function D_prime(P; normalized = false)
    P = _normalize(P, normalized)
    vD = D(P; normalized = true)
    vD / _D_max(P, D)
end

"""
The correlation coefficient
"""
function r²(P; normalized = false)
    P = _normalize(P, normalized)
    pa, pb = _papb(P)
    vD = D(P; normalized = true)
    vD ^ 2 / (pa * (1 - pa) * pb * (1 - pb))
end

"""
`snp1` and `snp2` is n-element boolean arrays indicating alleles at two SNPs within n haplotypes.
"""
function r²(snp1::AbstractVector{Bool}, snp2::AbstractVector{Bool})
    P = zeros(2, 2)
    n = length(snp1)
    for i in 1:n
        P[2 - snp1[i], 2 - snp2[i]] += 1
    end
    # P includes numbers of SNP pairs that has following values
    # [(1,1) (1,0)
    #  (0,1) (0,0)]
    r²(P)
end

"""
`snps` is a nxm boolean matrix indicating alleles at m SNPs within n haplotypes.

Return:

`R²`: A matrix of `r²` values for each pair of SNPs.
"""
function r²(snps::Matrix{Bool})
    n, m = size(snps)
    R² = zeros(m, m)
    for i in 1:m-1
        R²[i, i] = 1.0
        for j in i + 1:m
            R²[i, j] = R²[j, i] = r²(snps[:, i], snps[:, j])
        end
    end
    R²
end


#=
Decomposition of Return
    rₜ=√hₜ zₜ, zₜ~N(0,1) so rₜ|Fₜ₋₁~N(0,hₜ). hₜ=Var(rₜ|Fₜ₋₁)

GARCH(p,q)(Bollerslev 1986)

    hₜ₊₁=ω+Σβⱼhₜ₋ⱼ+Σαᵢ(rₜ₋ᵢ²-hₜ₋ᵢ)

    αᵢ,βⱼ,ω≧0

EGARCH(p,q)(Nelson 1991)

    loghₜ₊₁=ω+Σαᵢg(zₜ)+Σβⱼloghₜ₋ⱼ

    where g(z)=θ₁z+θ₂(|z|-E|z|) (Leverage function)

REGARCH(1,1,K)(Hansen and Zhuo 2015)

    GARCH equation

        loghₜ₊₁=ω + τ(zₜ) + βloghₜ + γ′uₜ

        where τ(z)=τ₁z+τ₂(z²-1)

    Measurement Equation

        log xₖ,ₜ = ξₖ + ϕₖ log hₜ + δ₍ₖ₎(zₜ) + uₖ,ₜ for k=1,...,K (we assume ϕₖ=1)

        where δ₍ₖ₎(z)=δₖ,₁z + δₖ,₂(z²-1)

=#

module GARCHModels

    using Optim
    using Optim: TwiceDifferentiable, optimize, minimizer
    using FiniteDiff, ForwardDiff
    using LinearAlgebra, Statistics, Distributions
    using DataFrames

    export REGARCH, GARCH, EGARCH

    mutable struct Volatilityestimation
        parameter
        robustcov
        conditionalvariance
        LL
    end

    subscript(i::Integer) = i<0 ? error("$i is negative") : join('₀'+d for d in reverse(digits(i)))

    function GARCH_coefnames(p,q)
        names = Array{String, 1}()
        push!(names,"ω")
        for k in 1:p
            push!(names, "β"*subscript(k))
        end
        for k in 1:q
            push!(names, "α"*subscript(k))
        end
        push!(names,"h0")
        return names
    end
    function GARCH_startingvals(p,q,rt)
        x0 = zeros(p+q+2)
        x0[1]=0.01
        x0[2:p+1] = 0.7/p
        x0[p+2:p+1+q] .= 0.2/q
        x0[end] = (var(rt))
        return x0
    end
    function GARCH_constraints(p,q)
        lower = -ones(p+q+2)
        upper = ones(p+q+2)
        lower[end] = 0
        return lower, upper
    end
    function GARCH_loglik(params::AbstractVector{T},p,q,rt,output) where T
    # Return data must be a Vector, not a Matrix.
        lower, upper = GARCH_constraints(p,q)
        all(lower.<params.<upper) || return -Inf
        n = size(rt)[1]
        presample = max(p,q)
        names = Array{String, 1}()

        ω = params[1]
        β = params[2:p+1]
        α = params[p+2:p+q+1]
        h0 = params[end]
        ht = zeros(T,n)
        zt = zeros(T,n)
        ht[1:presample] .= h0
        zt[1:presample] = rt[1:presample]./sqrt.(ht[1:presample])
        for t in presample+1:n
            mht = ω
            for i in 1:p
                mht += β[i]*ht[t-i]#log(xt[t,k]) - ξ[k] - log(ht[t]) - δ₁[k]*zt[t] - δ₂[k]*(zt[t]^2 - 1)
            end
            for j in 1:q
                mht += α[j]*(rt[t-j]^2 - ht[t-j])
            end
            ht[t] = mht
            zt[t] = rt[t]./sqrt.(ht[t])
        end
        LL = -log.(ht) .- zt.^2
        if output=="LL"
            return sum(LL)
        elseif output=="ht"
            return ht
        elseif output=="LLs"
            return LL
        end
    end
    function GARCH_scores(params,p,q,rt)
        f = x -> GARCH_loglik(x,p,q,rt,"LLs")
        S = ForwardDiff.jacobian(f, params)
    end
    function GARCH_InformationMatrix(params,p,q,rt)
        g = x -> GARCH_loglik(x,p,q,rt,"LL")
        H = ForwardDiff.hessian(g, params)
        J = -H
    end

    function GARCH_fit(p,q,rt,xt)
        obj = params -> -GARCH_loglik(params,p,q,rt,"LL")
        params = GARCH_startingvals(p,q,rt)
        res = optimize(obj, params, autodiff=:forward,iterations =10000,show_trace=true)
        params .= Optim.minimizer(res)
    end

    function GARCH(p,q,rt)
        params = GARCH_fit(p,q,rt)
        I = GARCH_InformationMatrix(params,p,q,rt)
        S = GARCH_scores(params,p,q,rt)
        ht = GARCH_loglik(params,p,q,rt,"ht")
        J = S'*S
        robustcov = I\J/I
        LL = GARCH_loglik(params,p,q,rt,"LL")
        coefnames = GARCH_coefnames(K)
        z = params./sqrt.(diag(robustcov))
        pvalue = (1 .- cdf.(Normal(),abs.(z)))*2
        parameters = DataFrame([coefnames,params,sqrt.(diag(robustcov)),z,pvalue])
        DataFrames.rename!(parameters,["Parameter","Estimate","Std","z","P(>|z|)"])
        res = Volatilityestimation(parameters,robustcov,ht,LL)
    end

    function EGARCH_coefnames(p,q)
        names = Array{String, 1}()
        push!(names,"ω")
        for k in 1:p
            push!(names, "β"*subscript(k))
        end
        for k in 1:q
            push!(names, "α"*subscript(k))
        end
        for i in 1:2
            push!(θ, "θ"*subscript(k))
        end
        push!(names,"logh0")
        return names
    end

    function EGARCH_startingvals(p,q,rt)
        x0 = zeros(p+q+4)
        x0[1]=0
        x0[2:p+1] = 0.7/p
        x0[p+2:p+1+q] .= 0.2/q
        x0[p+q+2:p+q+3] .= 0
        x0[end] = (var(rt))
        return x0
    end

    function EGARCH_constraints(p,q)
        lower = -ones(p+q+4)
        upper = ones(p+q+4)
        lower[end] = -10
        return lower, upper
    end

    function EGARCH_loglik(params::AbstractVector{T},p,q,rt,output) where T
    # Return data must be a Vector, not a Matrix.
        lower, upper = EGARCH_constraints(p,q)
        all(lower.<params.<upper) || return -Inf

        n = size(rt)[1]
        presample = max(p,q)
        names = Array{String, 1}()

        ω = params[1]
        β = params[2:p+1]
        α = params[p+2:p+q+1]
        θ = params[p+q+2:p+q+3]
        logh0 = params[end]
        ht = zeros(T,n)
        zt = zeros(T,n)
        ht[1:presample] .= exp(logh0)
        zt[1:presample] = rt[1:presample]./sqrt.(ht[1:presample])
        for t in presample+1:n
            mht = ω
            for i in 1:p
                mht += β[i]*log(ht[t-i])#log(xt[t,k]) - ξ[k] - log(ht[t]) - δ₁[k]*zt[t] - δ₂[k]*(zt[t]^2 - 1)
            end
            for j in 1:q
                mht += α[j]*(abs(zt[t-j]) - sqrt(2/pi))
            end
            ht[t] = mht
            zt[t] = rt[t]./sqrt.(ht[t])
        end
        LL = -log.(ht) .- zt.^2
        if output=="LL"
            return sum(LL)
        elseif output=="ht"
            return ht
        elseif output=="LLs"
            return LL
        end
    end
    function EGARCH_scores(params,p,q,rt)
        f = x -> EGARCH_loglik(x,p,q,rt,"LLs")
        S = ForwardDiff.jacobian(f, params)
    end
    function EGARCH_InformationMatrix(params,p,q,rt)
        g = x -> EGARCH_loglik(x,p,q,rt,"LL")
        H = ForwardDiff.hessian(g, params)
        J = -H
    end

    function EGARCH_fit(p,q,rt,xt)
        obj = params -> -EGARCH_loglik(params,p,q,rt,"LL")
        params = EGARCH_startingvals(p,q,rt)
        res = optimize(obj, params, autodiff=:forward,iterations =10000,show_trace=true)
        params .= Optim.minimizer(res)
    end

    function EGARCH(p,q,rt)
        params = EGARCH_fit(p,q,rt)
        I = EGARCH_InformationMatrix(params,p,q,rt)
        S = EGARCH_scores(params,p,q,rt)
        ht = EGARCH_loglik(params,p,q,rt,"ht")
        J = S'*S
        robustcov = I\J/I
        LL = EGARCH_loglik(params,p,q,rt,"LL")
        coefnames = EGARCH_coefnames(K)
        z = params./sqrt.(diag(robustcov))
        pvalue = (1 .- cdf.(Normal(),abs.(z)))*2
        parameters = DataFrame([coefnames,params,sqrt.(diag(robustcov)),z,pvalue])
        DataFrames.rename!(parameters,["Parameter","Estimate","Std","z","P(>|z|)"])
        res = Volatilityestimation(parameters,robustcov,ht,LL)
    end

    function REGARCH_coefnames(p,q)
        names = Array{String, 1}()
        push!(names,"ω")
        push!(names, "β")
        push!(names, "τ₁")
        push!(names, "τ₂")
        for k in 1:K
            push!(names, "γ"*subscript(k))
        end
        for k in 1:K
            push!(names, "ξ"*subscript(k))
        end
        for k in 1:K
            push!(names,"δ₁"*subscript(k))
        end
        for k in 1:K
            push!(names,"δ₂"*subscript(k))
        end
        push!(names,"logh0")
        return names
    end
    function REGARCH_startingvals(K,rt)
        x0 = zeros(5+4*K)
        x0[1]=0.01
        x0[2] = 0.8
        x0[5:4+K] .= 0.3/K
        x0[5+K:4+2*K] .= -0.3
        x0[end] = log(var(rt))
        return x0
    end
    function REGARCH_constraints(K)
        lower = -ones(4*K+5)
        upper = ones(4*K+5)
        upper[5+2*K:4+3*K] .= 2
        lower[end] = -50
        return lower, upper
    end
    function REGARCH_loglik(params::AbstractVector{T},K,rt,xt,output) where T
    # Return data must be a Vector, not a Matrix.
        size(rt)[1]==size(xt)[1] || return error("Check the length of the data")
        lower, upper = REGARCH_constraints(K)
        all(lower.<params.<upper) || return -Inf
        n = size(rt)[1]
        presample = 1
        nparam = 5+5*K
        ω = params[1]
        β = params[2]
        τ₁ = params[3]
        τ₂ = params[4]
        γ = params[5:4+K]
        ξ = params[5+K:4+2*K]
        δ₁ = params[5+2*K:4+3*K]
        δ₂ = params[5+3*K:4+4*K]
        logh0 = params[end]
        ht = zeros(T,n)
        ut = zeros(T,n,K)
        zt = zeros(T,n)
        ht[1:presample] .= exp(logh0)
        ut[1:presample] .= 0
        zt[1:presample] = rt[1:presample]./sqrt.(ht[1:presample])
        for t in presample+1:n
            mlht = ω + β*log(ht[t-1]) .+ τ₁*zt[t-1] .+ τ₂*(zt[t-1]^2-1) .+ sum(γ.*(ut[t-1,:]))
            ht[t] = exp(mlht)
            zt[t] = rt[t]./sqrt.(ht[t])
            for k in 1:K
                ut[t,k] = log(xt[t,k]) - ξ[k] - log(ht[t]) - δ₁[k]*zt[t] - δ₂[k]*(zt[t]^2 - 1)
            end
        end
        Σ = ut'*ut/n
        LL = -log.(ht) .- zt.^2 .- log(det(Σ))
        if output=="LL"
            return sum(LL)
        elseif output=="ht"
            return ht
        elseif output=="LLs"
            return LL
        end
    end
    function REGARCH_scores(params,K,rt,xt)
        f = x -> REGARCH_loglik(x,K,rt,xt,"LLs")
        S = ForwardDiff.jacobian(f, params)
    end
    function REGARCH_InformationMatrix(params,K,rt,xt)
        g = x -> REGARCH_loglik(x,K,rt,xt,"LL")
        H = ForwardDiff.hessian(g, params)
        J = -H
    end

    function REGARCH_fit(K,rt,xt)
        obj = params -> -REGARCH_loglik(params,K,rt,xt,"LL")
        params = REGARCH_startingvals(K,rt)
        res = optimize(obj, params, autodiff=:forward,iterations =10000,show_trace=true)
        params .= Optim.minimizer(res)
    end

    function REGARCH(K,rt,xt)
        params = REGARCH_fit(K,rt,xt)
        I = REGARCH_InformationMatrix(params,K,rt,xt)
        S = REGARCH_scores(params,K,rt,xt)
        ht = REGARCH_loglik(params,K,rt,xt,"ht")
        J = S'*S
        robustcov = I\J/I
        LL = REGARCH_loglik(params,K,rt,xt,"LL")
        coefnames = REGARCH_coefnames(K)
        z = params./sqrt.(diag(robustcov))
        pvalue = (1 .- cdf.(Normal(),abs.(z)))*2
        parameters = DataFrame([coefnames,params,sqrt.(diag(robustcov)),z,pvalue])
        DataFrames.rename!(parameters,["Parameter","Estimate","Std","z","P(>|z|)"])
        res = Volatilityestimation(parameters,robustcov,ht,LL)
    end

end


function inferExact(nodepot, edgepot, es::EdgeStruct)

  nnodes = size(nodepot,1)
  nedges = size(edgepot,3)

  nodebel = zeros(size(nodepot))
  edgebel = zeros(size(edgepot))
  y = ones(nnodes)
  Z = 0
  i = 1
  while true

      pot = configPotential(y,nodepot,edgepot,es)

      for n = 1:nnodes
          nodebel[n,y[n]] = nodebel[n,y[n]]+pot;
      end

      for e = 1:nedges
          n1 = es.edgeends[e,1]
          n2 = es.edgeends[e,2]
          edgebel[y[n1],y[n2],e] = edgebel[y[n1],y[n2],e]+pot
      end

      Z += pot

      for i = 1:nnodes
          y[i] += 1
          if y[i] <= es.nstates[i]
              break
          else
              y[i] = 1;
          end
      end

      if  i == nnodes && y[end] == 1
          break;
      end
  end

  nodebel = nodebel./Z;
  edgebel = edgebel./Z;
  logZ = log(Z);

  return nodebel, edgebel, logZ, Z
end


using PkgSRA
using Plots, Random, Distributions, Interpolations
pyplot()
Random.seed!(11)
using BSON: @save
include("../scripts/paper_settings.jl");

#####
##### Define the controlled dynamical system
#####
# Load example problem
include(EXAMPLE_FOLDERNAME*"example_falling_ball.jl")
U_starts = [6.0]
    U_widths = 0.2;
    amplitude = 75.0;
U_func_kick(t, u) = U_func_time(t, u,
                                U_widths, U_starts,
                                F_dim=2,
                                amplitude=amplitude)

# Define the spatial forcing function

U_func_wall(X) = U_func_spring(X, k=1e3, r=1.0);

# Solve the system
prob = ODEProblem((du, u, p, t)->
        ball_system(du, u, p, t,
             U_func_time=U_func_kick,
             U_func_space=U_func_wall),
         u1, tspan, p1)
dt = ts[2] - ts[1]

dyn_with_ctr = solve(prob, Tsit5(), saveat=ts, adaptive=false, dt = dt);
# dyn_with_ctr = solve_ball_system(U_func_space=U_func_wall,
#                                  U_func_time=U_func_kick)
dat = Array(dyn_with_ctr)
dat_with_const = vcat(dat, ones(1,size(dat,2)))
numerical_grad = numerical_derivative(dat, ts)

# Get just the control signals
dyn_control_wall = zeros(size(dat))
for i in 1:size(dat, 2)
    dyn_control_wall[:,i] = U_func_wall(dat[:, i])
end

dyn_control_kick = zeros(size(dat))
for (i, t) in enumerate(ts)
    dyn_control_kick[:,i] = U_func_kick(t, dat[:, i])
end

U_true = dyn_control_kick .+ dyn_control_wall

#####
##### Fit a naive linear model (with constant)
#####
naive_model = dmdc(dat_with_const, numerical_grad)[1]

f(u,p,t) = vcat(naive_model * u, [1])
prob = ODEProblem(f, vcat(u1,[1]), tspan, p1)
dat_naive = Array(solve(prob, Tsit5(), saveat=ts))

# plot(dat_naive[1,:]); plot!(dat[1,:])

#####
##### Build SRA object and analyze
#####
# Initialize
this_model = sra_stateful_object(ts, tspan, dat, u1, numerical_grad)
prams = this_model.parameters
prams.sindyc_ensemble_parameters[:selection_criterion] =
    sindy_cross_validate;
# Several changes are required because there are only 2 variables here
# prams.variable_names = ["x", "v"];
prams.sindy_terms_list = Iterators.product(1:2, 1:2)
prams.sindy_library["cross_terms"] = [2, 3] # Also include cubic terms

fit_first_model(this_model, 1);

# U_true = vcat(dyn_control_kick, dyn_control_wall)
ind = 1:500
i = 2
    plot(residual_unctr[i,ind], ribbon=sample_noise)
    plot!(U_true[:,ind]', label="True")
    title!("Residual and true control")

## Subsample
accepted_ind = subsample_using_residual(
            residual_unctr, sample_noise,
            min_length=1, shorten_length=2)
num_pts = 400; start_ind = 101
subsample_ind = accepted_ind[start_ind:num_pts+start_ind-1]

U_sub = U_true[:,subsample_ind]
plot(U_sub')
    # plot!(residual[:,subsample_ind]', ribbon=sample_trajectory_noise)
    title!("Control signals in the subsampled dataset")

chain_ctr, best_sindy_ctr = calc_distribution_of_models(
    dat[:,subsample_ind],
    numerical_grad[:,subsample_ind],
    sindy_library,
    val_list = Iterators.product(1:3,1:3), # DEBUG
    # val_list = combinations(1:3, 2),
    # val_list = calc_permutations(3,2),
    chain_opt = (iterations=300, train_ind=1:num_pts)
)

# print_equations(best_sindy_ctr)
# plot(chain_ctr)

# Calculate updated control signal
(residual_ctr, _, noise_ctr, _) =
        calc_distribution_of_residuals(
        dat, numerical_grad, chain_ctr, 1:length(ts), best_sindy_ctr)
ctr_final = process_residual(residual_ctr, noise_ctr)

plot(ctr_final', label="Learned", lw=3, color=:black)
    plot!(U_true', label="True")

# Get the actually controlled model
# val_list = calc_permutations(3,2)
# val_list = combinations(1:3, 2)
val_list = Iterators.product(1:31:3) # DEBUG
(final_sindy_model,best_criterion,all_criteria,_, _) =
     sindyc_ensemble(dat, numerical_grad, sindy_library, val_list,
                     U=ctr_final, ts=ts,
                     selection_criterion=my_aicc,
                     sparsification_mode="num_terms",
                     selection_dist=Normal(0,1),
                     use_clustering_minimization=false)

scatter(sum.(val_list), all_criteria)

prob = ODEProblem(final_sindy_model, dat[:,1], tspan)
# ts2 = range(tspan[1], tspan[end], length=10000)
dt = ts[2] - ts[1]
sol = solve(prob, AB3(), dt=dt, saveat=ts)
dat_final = Array(sol)

plot(sol)
#####
##### Save
#####
this_dat_name = DAT_FOLDERNAME*"dat_ball_"

# Raw data
fname = this_dat_name*"raw.bson";
@save fname dat numerical_grad dyn_control_kick dyn_control_wall

# Naive model
fname = this_dat_name*"naive_model.bson";
@save fname dat_naive naive_model

# Controlled model
fname = this_dat_name*"controlled_model.bson";
@save fname chain_ctr ctr_final dat_final


function layer_norm(sz::Integer, activation=identity; affine::Bool=true,
                    epsilon::Real=1.0f-5)
    if affine
        return Chain(WrappedFunction(x -> normalise(x, identity; dims=1, epsilon=epsilon)),
                     Scale(sz, activation))
    else
        return WrappedFunction(x -> normalise(x, activation; dims=1, epsilon=epsilon))
    end
end


export graph_ps, graph_ps_degopt

"""
     (graph,crefs)=graph_ps(a; input=:A,
                          B_base=:B, C_base=:C, P_base=:P)

Generates the graph for the Paterson–Stockmeyer procedure with monomial basis
coefficieents. More precisely, it corresponds to evaluation of the polynomial

    p(s) = a[1] + a[2]*s + ... + a[n]*s^(n-1).

The code follows the description in the second of the papers referenced below.

**References**

1. M. Paterson and L. Stockmeyer. "On the number of nonscalar multiplications
   necessary to evaluate polynomials". SIAM Journal on Scientific Computing,
   2(1):60-66, 1973. DOI: [10.1137/0202007](https://doi.org/10.1137/0202007)

2. M. Fasi. "Optimality of the Paterson–Stockmeyer method for evaluating matrix
   polynomials and rational matrix functions". Linear Algebra and its
   Applications, 574, 2019.
   DOI: [10.1016/j.laa.2019.04.001](https://doi.org/10.1016/j.laa.2019.04.001)
"""
function graph_ps(a; input = :A, B_base = :B, C_base = :C, P_base = :P)

    # Initial setup
    n = length(a)
    cref = Vector{Tuple{Symbol,Int}}(undef, n)
    graph = Compgraph(eltype(a))

    k = n - 1 # Polynomial degree

    # Degenerate cases k = 0 and k = 1
    if k < 2
        if k == 0
            add_lincomb!(graph, :P0, a[1], :I, 0.0, :I)
            cref[1] = (:P0, 1)
        elseif k == 1
            add_lincomb!(graph, :P0, a[2], :A, a[1], :I)
            cref[1] = (:P0, 1)
            cref[2] = (:P0, 2)
        end
        add_output!(graph, :P0)
        return (graph, cref)
    else
        s = ceil(Int, sqrt(k))
        v = floor(Int, k / s)
    end

    # Create monomial basis
    prevkey = input
    # up to s-1 are needed for sub-polynomials, and s needed for outer monomials
    for i = 2:s
        key = Symbol("$input$i")
        add_mult!(graph, key, input, prevkey)
        prevkey = key
    end
    Amax = Symbol("$input$s")

    # Create intermediate coefficient-polynomials. (3) and following equation in
    # Fasi 2019
    for i = 0:(v-1)
        # The number of coefficients in the sub-coefficient-polynomial
        # Note, degree is one lower than the index.
        lower_idx = i * s + 1 #plus 1 for 1-indexing
        upper_idx = lower_idx + s - 1 # s number of coeffients, so degree s-1
        idx = lower_idx:upper_idx
        c = view(a, idx)
        ccref = view(cref, idx)

        # B-keys are cumulative sums of the coefficient-polynomials.
        Bkey = Symbol("$(B_base)_$(i)_1")
        add_lincomb!(graph, Bkey, c[1], :I, c[2], input)
        ccref[1] = (Bkey, 1)
        ccref[2] = (Bkey, 2)
        for k = 3:s
            #As above, degree is one lower than the index
            powkey = Symbol("$input$(k-1)")
            prevBkey = Bkey
            Bkey = Symbol("$(B_base)_$(i)_$(k-1)")
            add_lincomb!(graph, Bkey, c[k], powkey, 1.0, prevBkey)
            ccref[k] = (Bkey, 1)
        end
    end

    # Handle highest powers separately. i=v. Part of (3), the following
    # equation, and (4) in Fasi 2019.
    ks = rem(k, s) # Degree of remaining highest order term. We compute B_v_ks
    if ks == 0 # Only constant term
        Pkey = Symbol("$(P_base)$(v-1)")
        Bkey = Symbol("$(B_base)_$(v-1)_$(s-1)")
        # P_{v-1} = a[n]*A^s + B_[v-1}
        add_lincomb!(graph, Pkey, a[n], Amax, 1.0, Bkey)
        cref[n] = (Pkey, 1)
    else # ks > 0. At least constant and linear term.
        # Repeat above procedue with possibly lower-degree polynomial
        idx = (n-(ks+1)+1):n
        c = view(a, idx)
        ccref = view(cref, idx)

        # P_{v-1} = A^s*B_v + B_{v-1}, i.e., have to create B_v

        # B-keys are cumulative sums of the coefficient-polynomials.
        Bkey = Symbol("$(B_base)_$(v)_1")
        add_lincomb!(graph, Bkey, c[1], :I, c[2], input)
        ccref[1] = (Bkey, 1)
        ccref[2] = (Bkey, 2)
        for k = 3:ks+1
            # As above, degree is one lower than the index
            powkey = Symbol("$input$(k-1)")
            prevBkey = Bkey
            Bkey = Symbol("$(B_base)_$(v)_$(k-1)")
            add_lincomb!(graph, Bkey, c[k], powkey, 1.0, prevBkey)
            ccref[k] = (Bkey, 1)
        end
        # C_{v-1} = A^s * B_v
        Ckey = Symbol("$(C_base)$(v-1)")
        Bkey = Symbol("$(B_base)_$(v)_$(ks)")
        add_mult!(graph, Ckey, Bkey, Amax)
        # P_{v-1} = C_{v-1}+ B_{v-1}
        Pkey = Symbol("$(P_base)$(v-1)")
        Bkey = Symbol("$(B_base)_$(v-1)_$(s-1)")
        add_lincomb!(graph, Pkey, 1.0, Ckey, 1.0, Bkey)
    end

    prevPkey = Pkey

    # Evaluate outer-polynomial with Horner-type scheme. (4) in Fasi 2019.
    # Pout = :Nothing;
    Pout = Pkey
    for i in reverse(0:(v-2))
        # C_i = A^s * P_{i+1}
        Ckey = Symbol("$(C_base)$i")
        add_mult!(graph, Ckey, prevPkey, Amax)

        # P_i = C_i + B_i
        Bkey = Symbol("$(B_base)_$(i)_$(s-1)")
        Pkey = Symbol("$(P_base)$i")
        add_lincomb!(graph, Pkey, 1.0, Ckey, 1.0, Bkey)

        prevPkey = Pkey
        Pout = Pkey
    end

    add_output!(graph, Pout)
    return (graph, cref)
end

"""
     (graph,crefs)=graph_ps_degopt(a; input=:A)

Generates the same polynomial as [`graph_ps`](@ref), i.e., the graph for the
Paterson–Stockmeyer procedure with monomial basis coefficieents. However, it
does so by wrapping a call to [`graph_degopt`](@ref), resulting in more degrees
of freedom in `crefs`.
"""
function graph_ps_degopt(a; input = :A)

    # Initial setup
    n = length(a)
    T = eltype(a)
    graph = Compgraph(T)
    k = n - 1 # Polynomial degree

    if k <= 2 # Degenerate cases k = 0, k = 1, and k = 2
        if k == 0
            error("Does not implement degree-zero polynomial.")
        else
            # PS is not more efficient than monomial evaluation for low degrees,
            # but gives complications in code below.
            x = Vector{Tuple{Vector{T},Vector{T}}}(undef, k - 1)
            for i = 1:(k-1)
                x[i] = (
                    vcat(zero(T), one(T), zeros(i - 1)),
                    vcat(zeros(i), one(T)),
                )
            end
            return graph_degopt(x, a)
        end
    end

    s = ceil(Int, sqrt(k))
    v = floor(Int, k / s)
    ks = rem(k, s) # Degree of remaining highest order term.
    # Treat special if only constant. Save one multiplication
    if ks == 0
        i_adj = 1
    else
        i_adj = 0
    end

    x = Vector{Tuple{Vector{T},Vector{T}}}(undef, s + v - 1 - i_adj)
    # Create monomial basis (The first s-1 elements in the x-vector  <=>
    # The first s+1 slots in an x-component)
    for i = 1:(s-1)
        x[i] =
            (vcat(zeros(T, i), one(T)), vcat(zero(T), one(T), zeros(T, i - 1)))
    end

    # Evaluate outer-polynomial with Horner-type scheme, and at the same time
    # create the intermediate coefficient-polynomials
    if ks != 0 # If highest order term non-constant we create the polynomial
        idx = (n-(ks+1)+1):n
        diff = s - (ks + 1)
        c = view(a, idx)
        # A^s*P_{v-1}
        x[s] = (vcat(zeros(T, s), one(T)), vcat(c, zeros(T, diff), zero(T)))
    end

    for i in reverse(0:v-2)
        lower_idx = (i + 1) * s + 1
        upper_idx = lower_idx + s - 1
        idx = lower_idx:upper_idx
        c = view(a, idx)
        # Coeffs for A^s
        xl = vcat(zeros(T, s), one(T), zeros(T, v - 1 - i - i_adj))
        # If highest order term is constant we incorporate it here
        if (i == v - 2) && (ks == 0)
            xr = vcat(c, a[n]) # Coeffs for P_{v-i} + A^s*a[n]
        else
            # Coeffs for P_{v-i} + A^s*P_{v+1-i}
            xr = vcat(c, zero(T), zeros(T, v - 2 - i - i_adj), one(T))
        end
        x[(s-i_adj)+v-1-i] = (xl, xr) # A^s*( P_{v-i} + A^s*P_{v+1-i} )
    end

    z = vcat(a[1:s], zero(T), zeros(T, v - 1 - i_adj), one(T))
    return graph_degopt(x, z, input = input)
end


using AbstractAlgebra
using Logging

global_logger(Logging.ConsoleLogger(Logging.Info))

struct RewriteSystem
    tails  # vector of tails
    rules  # vector of rewriters
end

struct Rewriter
    idx::Integer
    monom
    terminal
end

function addrule!(RS, i, m)
    push!(RS.rules, Rewriter(i, m, isone(m)))
end

function addtail!(RS, t)
    push!(RS.tails, (leading_monomial(t), tail(t)))
    addrule!(RS, length(RS.tails), one(t))
end

function findrewriter(monom, RS)
    for rule in reverse(RS.rules)
        t, d = divides(monom, rule.monom*first(RS.tails[rule.idx]))
        if t
            @info "found rule" rule
            addrule!(RS, rule.idx, rule.monom*d)
            return d, rule.monom*last(RS.tails[rule.idx])
        end
    end
    return monom, zero(monom)
end

function rewrite(mult, f, RS)
    newmonoms = []
    @warn "rewriting" mult f
    for m in monomials(f)
        mm = mult*m
        mult1, rewriter = findrewriter(mm, RS)
        @info m mm mult1 rewriter
        if !iszero(rewriter)
            append!(newmonoms, .- rewrite(mult1, rewriter, RS))
        else
            push!(newmonoms, mm)
        end
    end
    newmonoms
end

function rewrite(f, RS)
    sum(rewrite(one(f), f, RS))
end

"""

"""
function interreduce(gens::Vector{Poly}) where {Poly}
    sort!(gens, by=total_degree ∘ leading_monomial)

    RS = RewriteSystem([], [])

    # rewriters[1] = x1 + x2 + x3

    @warn "interreducing" gens

    newgens = similar(gens)
    newgens[1] = map_coefficients(c -> c // leading_coefficient(gens[1]), gens[1])

    addtail!(RS, newgens[1])

    i = 2

    while i <= length(gens)
        newgens[i] = rewrite(gens[i], RS)
        newgens[i] = map_coefficients(c -> c // leading_coefficient(newgens[i]), newgens[i])
        addtail!(RS, newgens[i])
        @error "End of $i" RS
        i += 1
    end

    newgens
end