module WilsonFermion_module
    using LinearAlgebra
    #using ..Fermion:FermionFields

    import ..Actions:FermiActionParam,FermiActionParam_Wilson,
                FermiActionParam_WilsonClover,FermiActionParam_Staggered
    import ..Gaugefields:GaugeFields,GaugeFields_1d,SU3GaugeFields,SU2GaugeFields,SU3GaugeFields_1d,SU2GaugeFields_1d,
                staggered_phase,SUn,SU2,SU3,SUNGaugeFields,SUNGaugeFields_1d,SU

    #import ..Fermion:FermionFields,
    #            clear!
    #using ..Fermion
    import ..AbstractFermion:FermionFields,
        Wx!,Wdagx!,clear!,substitute_fermion!,Dx!,fermion_shift!,fermion_shiftB!,add!,set_wing_fermi!,WdagWx!,apply_periodicity,Ddagx!

    """
    Struct for WilsonFermion
    """
    struct WilsonFermion <: FermionFields
        NC::Int64
        NX::Int64
        NY::Int64
        NZ::Int64
        NT::Int64
        f::Array{ComplexF64,6}

        γ::Array{ComplexF64,3}
        rplusγ::Array{ComplexF64,3}
        rminusγ::Array{ComplexF64,3}
        hop::Float64 #Hopping parameter
        r::Float64 #Wilson term
        hopp::Array{ComplexF64,1}
        hopm::Array{ComplexF64,1}
        eps::Float64
        Dirac_operator::String
        MaxCGstep::Int64
        BoundaryCondition::Array{Int8,1}
    end

    function Base.length(x::WilsonFermion)
        return x.NC*x.NX*x.NY*x.NZ*x.NT*4
    end

    function Base.size(x::WilsonFermion)
        return (x.NC,x.NX,x.NY,x.NZ,x.NT,4)
    end

    function Base.iterate(x::WilsonFermion, i::Int = 1)
        i == length(x.f)+1 && return nothing
        #println("i = $i ",x.f[i])
        return (x.f[i], i+1)
    end

    function WilsonFermion(NC,NX,NY,NZ,NT,fparam::FermiActionParam,BoundaryCondition)
        r = fparam.r
        hop = fparam.hop
        eps = fparam.eps
        MaxCGstep = fparam.MaxCGstep
        return WilsonFermion(NC,NX,NY,NZ,NT,r,hop,eps,MaxCGstep,BoundaryCondition)
    end

    function WilsonFermion(NC,NX,NY,NZ,NT,r,hop,eps,MaxCGstep,BoundaryCondition)#r,hop,eps,MaxCGstep)
        γ,rplusγ,rminusγ = mk_gamma(r)
        hopp = zeros(ComplexF64,4)
        hopm = zeros(ComplexF64,4)
        hopp .= hop
        hopm .= hop
        Dirac_operator = "Wilson"
        return WilsonFermion(NC,NX,NY,NZ,NT,zeros(ComplexF64,NC,NX+2,NY+2,NZ+2,NT+2,4),
            γ,rplusγ,rminusγ,hop,r,hopp,hopm,eps,Dirac_operator,MaxCGstep,BoundaryCondition)
    end

    function WilsonFermion(NC,NX,NY,NZ,NT,γ,rplusγ,rminusγ,hop,r,hopp,hopm,eps,fermion,MaxCGstep,BoundaryCondition)
        return WilsonFermion(NC,NX,NY,NZ,NT,zeros(ComplexF64,NC,NX+2,NY+2,NZ+2,NT+2,4),
                    γ,rplusγ,rminusγ,hop,r,hopp,hopm,eps,fermion,MaxCGstep,BoundaryCondition)
    end

    function Base.setindex!(x::WilsonFermion,v,i1,i2,i3,i4,i5,i6)
        x.f[i1,i2 + 1,i3 + 1,i4 + 1,i5 + 1,i6] = v
    end

    function Base.setindex!(x::WilsonFermion,v,ii)
        #ii = (((((ialpha -1)*NT+it-1)*NZ+iz-1)*NY+iy-1)*NX+ix-1)*NC+ic
        ic = (ii - 1) % x.NC + 1
        iii = (ii - ic) ÷ x.NC
        ix = iii % x.NX + 1
        iii = (iii - ix + 1) ÷ x.NX
        iy = iii % x.NY + 1
        iii = (iii - iy + 1) ÷ x.NY
        iz = iii % x.NZ + 1
        iii = (iii - iz + 1) ÷ x.NZ
        it = iii % x.NT + 1
        iii = (iii - it + 1) ÷ x.NT
        ialpha = iii + 1
        x[ic,ix,iy,iz,it,ialpha] = v
    end

    function Base.getindex(x::WilsonFermion,i1,i2,i3,i4,i5,i6)
        return x.f[i1,i2 .+ 1,i3 .+ 1,i4 .+ 1,i5 .+ 1,i6]
    end

    function Base.getindex(x::WilsonFermion,ii)
        #ii = (((((ialpha -1)*NT+it-1)*NZ+iz-1)*NY+iy-1)*NX+ix-1)*NC+ic
        ic = (ii - 1) % x.NC + 1
        iii = (ii - ic) ÷ x.NC
        ix = iii % x.NX + 1
        iii = (iii - ix + 1) ÷ x.NX
        iy = iii % x.NY + 1
        iii = (iii - iy + 1) ÷ x.NY
        iz = iii % x.NZ + 1
        iii = (iii - iz + 1) ÷ x.NZ
        it = iii % x.NT + 1
        iii = (iii - it + 1) ÷ x.NT
        ialpha = iii + 1
        return x[ic,ix,iy,iz,it,ialpha]
    end

    function Base.:*(a::WilsonFermion,b::WilsonFermion)
        c = 0.0im
        for α=1:4
            for it=1:a.NT
                for iz=1:a.NZ
                    for iy=1:a.NY
                        for ix=1:a.NX
                            @simd for ic=1:a.NC
                                c+= conj(a[ic,ix,iy,iz,it,α])*b[ic,ix,iy,iz,it,α]
                            end
                        end
                    end
                end
            end
        end
        return c
    end

    function Base.:*(a::T,b::WilsonFermion) where T <: Number
        c = similar(b)
        for α=1:4
            for it=1:b.NT
                for iz=1:b.NZ
                    for iy=1:b.NY
                        for ix=1:b.NX
                            @simd for ic=1:b.NC
                                c[ic,ix,iy,iz,it,α] = a*b[ic,ix,iy,iz,it,α]
                            end
                        end
                    end
                end
            end
        end
        return c
    end

    function LinearAlgebra.axpy!(a::T,X::WilsonFermion,Y::WilsonFermion) where T <: Number #Y = a*X+Y
        for α=1:4
            for it=1:X.NT
                for iz=1:X.NZ
                    for iy=1:X.NY
                        for ix=1:X.NX
                            @simd for ic=1:X.NC
                                Y[ic,ix,iy,iz,it,α] += a*X[ic,ix,iy,iz,it,α]
                            end
                        end
                    end
                end
            end
        end
        return X
    end

    function LinearAlgebra.rmul!(a::WilsonFermion,b::T) where T <: Number
        for α=1:4
            for it=1:a.NT
                for iz=1:a.NZ
                    for iy=1:a.NY
                        for ix=1:a.NX
                            @simd for ic=1:a.NC
                                a[ic,ix,iy,iz,it,α] = b*a[ic,ix,iy,iz,it,α]
                            end
                        end
                    end
                end
            end
        end
        return a
    end

    function Base.similar(x::WilsonFermion)
        return WilsonFermion(x.NC,x.NX,x.NY,x.NZ,x.NT,
                    x.γ,x.rplusγ,x.rminusγ,x.hop,x.r,x.hopp,x.hopm,x.eps,x.Dirac_operator,x.MaxCGstep,x.BoundaryCondition)
    end

    function substitute_fermion!(H,j,x::WilsonFermion)
        i = 0
        for ialpha = 1:4
            for it=1:x.NT
                for iz=1:x.NZ
                    for iy=1:x.NY
                        for ix=1:x.NX
                            @simd for ic=1:x.NC
                                i += 1
                                H[i,j] = x[ic,ix,iy,iz,it,ialpha]
                            end
                        end
                    end
                end
            end
        end
    end

    function Wx!(xout::WilsonFermion,U::Array{G,1},
        x::WilsonFermion,temps::Array{T,1}) where  {T <: FermionFields,G <: GaugeFields}
        temp = temps[4]
        temp1 = temps[1]
        temp2 = temps[2]

        clear!(temp)
        set_wing_fermi!(x)
        for ν=1:4
            fermion_shift!(temp1,U,ν,x)

            #... Dirac multiplication
            mul!(temp1,view(x.rminusγ,:,:,ν),temp1)

            #
            fermion_shift!(temp2,U,-ν,x)
            mul!(temp2,view(x.rplusγ,:,:,ν),temp2)

            add!(temp,x.hopp[ν],temp1,x.hopm[ν],temp2)

        end

        clear!(xout)
        add!(xout,1,x,-1,temp)

        #display(xout)
        #    exit()
        return
    end

    function Wx!(xout::WilsonFermion,U::Array{G,1},
        x::WilsonFermion,temps::Array{T,1},fparam::FermiActionParam_Wilson) where  {T <: FermionFields,G <: GaugeFields}
        Wx!(xout,U,x,temps)
        return
    end

    function Wx!(xout::WilsonFermion,U::Array{G,1},
        x::WilsonFermion,temps::Array{T,1},fparam::FermiActionParam_WilsonClover) where  {T <: FermionFields,G <: GaugeFields}
        Wx!(xout,U,x,temps,fparam.CloverFμν)
        return
    end

    function Wx!(xout::WilsonFermion,U::Array{G,1},
        x::WilsonFermion,temps::Array{T,1},CloverFμν::AbstractArray) where  {T <: FermionFields,G <: GaugeFields}
        temp = temps[4]
        temp1 = temps[1]
        temp2 = temps[2]

        clear!(temp)
        set_wing_fermi!(x)
        for ν=1:4
            fermion_shift!(temp1,U,ν,x)

            #... Dirac multiplication
            mul!(temp1,view(x.rminusγ,:,:,ν),temp1)

            #
            fermion_shift!(temp2,U,-ν,x)
            mul!(temp2,view(x.rplusγ,:,:,ν),temp2)

            add!(temp,x.hopp[ν],temp1,x.hopm[ν],temp2)

        end

        clear!(xout)
        add!(xout,1,x,-1,temp)

        cloverterm!(xout,CloverFμν,x)
        #println( "xout ",xout*xout)

        #display(xout)
        #    exit()
        return
    end

    function Wdagx!(xout::WilsonFermion,U::Array{G,1},
        x::WilsonFermion,temps::Array{T,1},CloverFμν::AbstractArray) where {T <: FermionFields,G <: GaugeFields}
        temp = temps[4]
        temp1 = temps[1]
        temp2 = temps[2]

        clear!(temp)
        x5 = temps[3]

        mul_γ5x!(x5,x)
        set_wing_fermi!(x5)

        for ν=1:4
            fermion_shift!(temp1,U,ν,x5)

            #... Dirac multiplication
            mul!(temp1,view(x.rminusγ,:,:,ν),temp1)

            #
            fermion_shift!(temp2,U,-ν,x5)

            mul!(temp2,view(x.rplusγ,:,:,ν),temp2)

            add!(temp,x.hopp[ν],temp1,x.hopm[ν],temp2)
        end
        clear!(temp1)
        add!(temp1,1,x5,-1,temp)

        cloverterm!(temp1,CloverFμν,x5)

        mul_γ5x!(xout,temp1)
        return
    end

    function Wdagx!(xout::WilsonFermion,U::Array{G,1},
        x::WilsonFermion,temps::Array{T,1},fparam::FermiActionParam_WilsonClover) where {T <: FermionFields,G <: GaugeFields}
        Wdagx!(xout,U,x,temps,fparam.CloverFμν)
        return

    end

    function Wdagx!(xout::WilsonFermion,U::Array{G,1},
        x::WilsonFermion,temps::Array{T,1}) where  {T <: FermionFields,G <: GaugeFields}
        temp = temps[4]
        temp1 = temps[1]
        temp2 = temps[2]

        clear!(temp)
        set_wing_fermi!(x)
        for ν=1:4
            fermion_shift!(temp1,U,ν,x)

            #... Dirac multiplication
            #mul!(temp1,view(x.rminusγ,:,:,ν),temp1)
            mul!(temp1,view(x.rplusγ,:,:,ν),temp1)

            #
            fermion_shift!(temp2,U,-ν,x)
            #mul!(temp2,view(x.rplusγ,:,:,ν),temp2)
            mul!(temp2,view(x.rminusγ,:,:,ν),temp2)

            add!(temp,x.hopp[ν],temp1,x.hopm[ν],temp2)

        end

        clear!(xout)
        add!(xout,1,x,-1,temp)

        #display(xout)
        #    exit()
        return
    end

    function Wdagx!(xout::WilsonFermion,U::Array{G,1},
        x::WilsonFermion,temps::Array{T,1},fparam::FermiActionParam_Wilson) where {T <: FermionFields,G <: GaugeFields}
        Wdagx!(xout,U,x,temps)
        return
    end

    function Dx!(xout::WilsonFermion,U::Array{G,1},
        x::WilsonFermion,temps::Array{T,1}) where  {T <: FermionFields,G <: GaugeFields}
        temp = temps[4]
        temp1 = temps[1]
        temp2 = temps[2]

        clear!(temp)
        set_wing_fermi!(x)
        for ν=1:4
            fermion_shift!(temp1,U,ν,x)

            #... Dirac multiplication
            mul!(temp1,view(x.rminusγ,:,:,ν),temp1)

            #
            fermion_shift!(temp2,U,-ν,x)
            mul!(temp2,view(x.rplusγ,:,:,ν),temp2)

            add!(temp,0.5,temp1,0.5,temp2)

        end

        clear!(xout)
        add!(xout,1/(2*x.hop),x,-1,temp)

        #display(xout)
        #    exit()
        return
    end

    function Ddagx!(xout::WilsonFermion,U::Array{G,1},
        x::WilsonFermion,temps::Array{T,1}) where  {T <: FermionFields,G <: GaugeFields}
        temp = temps[4]
        temp1 = temps[1]
        temp2 = temps[2]

        clear!(temp)
        set_wing_fermi!(x)
        for ν=1:4
            fermion_shift!(temp1,U,ν,x)

            #... Dirac multiplication
            #mul!(temp1,view(x.rminusγ,:,:,ν),temp1)
            mul!(temp1,view(x.rplusγ,:,:,ν),temp1)

            #
            fermion_shift!(temp2,U,-ν,x)
            #mul!(temp2,view(x.rplusγ,:,:,ν),temp2)
            mul!(temp2,view(x.rminusγ,:,:,ν),temp2)

            add!(temp,0.5,temp1,0.5,temp2)

        end

        clear!(xout)
        add!(xout,1/(2*x.hop),x,-1,temp)

        #display(xout)
        #    exit()
        return
    end

    function LinearAlgebra.mul!(xout::WilsonFermion,A::AbstractMatrix,x::WilsonFermion)
        NX = x.NX
        NY = x.NY
        NZ = x.NZ
        NT = x.NT
        NC = x.NC

        #n6 = size(x.f)[6]
        #f = zeros(ComplexF64,4)
        #e = zeros(ComplexF64,4)

        for ic=1:NC
            for it=1:NT
                for iz=1:NZ
                    for iy=1:NY
                        @simd for ix=1:NX
                                e1 = x[ic,ix,iy,iz,it,1]
                                e2 = x[ic,ix,iy,iz,it,2]
                                e3 = x[ic,ix,iy,iz,it,3]
                                e4 = x[ic,ix,iy,iz,it,4]

                                xout[ic,ix,iy,iz,it,1] = A[1,1]*e1+A[1,2]*e2+A[1,3]*e3+A[1,4]*e4
                                xout[ic,ix,iy,iz,it,2] = A[2,1]*e1+A[2,2]*e2+A[2,3]*e3+A[2,4]*e4
                                xout[ic,ix,iy,iz,it,3] = A[3,1]*e1+A[3,2]*e2+A[3,3]*e3+A[3,4]*e4
                                xout[ic,ix,iy,iz,it,4] = A[4,1]*e1+A[4,2]*e2+A[4,3]*e3+A[4,4]*e4

                                #=
                            for k1=1:4
                                e[k1] = x[ic,ix,iy,iz,it,k1]
                            end

                            for k1=1:4
                                f[k1] = 0
                                for k2=1:4
                                    f[k1] += A[k1,k2]*e[k2]
                                end
                                xout[ic,ix,iy,iz,it,k1] = f[k1]
                            end
                            =#

                        end
                    end
                end
            end
        end
    end

    function LinearAlgebra.mul!(xout::WilsonFermion,x::WilsonFermion,A::AbstractMatrix)
        NX = x.NX
        NY = x.NY
        NZ = x.NZ
        NT = x.NT
        NC = x.NC

        #n6 = size(x.f)[6]
        #f = zeros(ComplexF64,n6)
        #e = zeros(ComplexF64,n6)

        for ic=1:NC
            for it=1:NT
                for iz=1:NZ
                    for iy=1:NY
                        @simd for ix=1:NX
                            e1 = x[ic,ix,iy,iz,it,1]
                            e2 = x[ic,ix,iy,iz,it,2]
                            e3 = x[ic,ix,iy,iz,it,3]
                            e4 = x[ic,ix,iy,iz,it,4]

                            xout[ic,ix,iy,iz,it,1] = A[1,1]*e1+A[2,1]*e2+A[3,1]*e3+A[4,1]*e4
                            xout[ic,ix,iy,iz,it,2] = A[1,2]*e1+A[2,2]*e2+A[3,2]*e3+A[4,2]*e4
                            xout[ic,ix,iy,iz,it,3] = A[1,3]*e1+A[2,3]*e2+A[3,3]*e3+A[4,3]*e4
                            xout[ic,ix,iy,iz,it,4] = A[1,4]*e1+A[2,4]*e2+A[3,4]*e3+A[4,4]*e4

                            #=
                            for k1=1:n6
                                e[k1] = x[ic,ix,iy,iz,it,k1]
                            end

                            for k1=1:n6
                                f[k1] = 0
                                for k2=1:n6
                                    f[k1] += e[k2]*A[k2,k1]
                                end
                                xout[ic,ix,iy,iz,it,k1] = f[k1]
                            end
                            =#

                        end
                    end
                end
            end
        end
    end

    function fermion_shift!(b::WilsonFermion,u::Array{T,1},μ::Int,a::WilsonFermion) where T <: SU3GaugeFields
        if μ == 0
            substitute!(b,a)
            return
        end

        NX = a.NX
        NY = a.NY
        NZ = a.NZ
        NT = a.NT
        NC = 3#a.NC

        #NTrange = get_looprange(NT)
        #println(NTrange)
        if μ > 0
            #idel = zeros(Int64,4)
            #idel[μ] = 1

            n6 = size(a.f)[6]
            for ialpha=1:4
                #for it=NTrange
                for it=1:NT
                    it1 = it + ifelse(μ ==4,1,0) #idel[4]
                    for iz=1:NZ
                        iz1 = iz + ifelse(μ ==3,1,0) #idel[3]
                        for iy=1:NY
                            iy1 = iy + ifelse(μ ==2,1,0) #idel[2]
                            @simd for ix=1:NX
                                ix1 = ix + ifelse(μ ==1,1,0) #idel[1]

                                b[1,ix,iy,iz,it,ialpha] = u[μ][1,1,ix,iy,iz,it]*a[1,ix1,iy1,iz1,it1,ialpha] +
                                                            u[μ][1,2,ix,iy,iz,it]*a[2,ix1,iy1,iz1,it1,ialpha] +
                                                            u[μ][1,3,ix,iy,iz,it]*a[3,ix1,iy1,iz1,it1,ialpha]

                                b[2,ix,iy,iz,it,ialpha] = u[μ][2,1,ix,iy,iz,it]*a[1,ix1,iy1,iz1,it1,ialpha] +
                                                            u[μ][2,2,ix,iy,iz,it]*a[2,ix1,iy1,iz1,it1,ialpha] +
                                                            u[μ][2,3,ix,iy,iz,it]*a[3,ix1,iy1,iz1,it1,ialpha]

                                b[3,ix,iy,iz,it,ialpha] = u[μ][3,1,ix,iy,iz,it]*a[1,ix1,iy1,iz1,it1,ialpha] +
                                                            u[μ][3,2,ix,iy,iz,it]*a[2,ix1,iy1,iz1,it1,ialpha] +
                                                            u[μ][3,3,ix,iy,iz,it]*a[3,ix1,iy1,iz1,it1,ialpha]

                                #=
                                for k1=1:3
                                    b[k1,ix,iy,iz,it,ialpha] = 0
                                    for k2=1:3
                                        b[k1,ix,iy,iz,it,ialpha] += u[μ][k1,k2,ix,iy,iz,it]*a[k2,ix1,iy1,iz1,it1,ialpha]
                                    end
                                end
                                =#
                            end
                        end
                    end
                end
            end

        elseif μ < 0
            #idel = zeros(Int64,4)
            #idel[-μ] = 1
            #n6 = size(b.f)[6]
            for ialpha =1:4
                for it=1:NT
                    it1 = it - ifelse(-μ ==4,1,0) #idel[4]
                    for iz=1:NZ
                        iz1 = iz - ifelse(-μ ==3,1,0) #idel[3]
                        for iy=1:NY
                            iy1 = iy - ifelse(-μ ==2,1,0)  #idel[2]
                            @simd for ix=1:NX
                                ix1 = ix - ifelse(-μ ==1,1,0) #idel[1]

                                b[1,ix,iy,iz,it,ialpha] = conj(u[-μ][1,1,ix1,iy1,iz1,it1])*a[1,ix1,iy1,iz1,it1,ialpha] +
                                                            conj(u[-μ][2,1,ix1,iy1,iz1,it1])*a[2,ix1,iy1,iz1,it1,ialpha] +
                                                            conj(u[-μ][3,1,ix1,iy1,iz1,it1])*a[3,ix1,iy1,iz1,it1,ialpha]

                                b[2,ix,iy,iz,it,ialpha] = conj(u[-μ][1,2,ix1,iy1,iz1,it1])*a[1,ix1,iy1,iz1,it1,ialpha] +
                                                            conj(u[-μ][2,2,ix1,iy1,iz1,it1])*a[2,ix1,iy1,iz1,it1,ialpha] +
                                                            conj(u[-μ][3,2,ix1,iy1,iz1,it1])*a[3,ix1,iy1,iz1,it1,ialpha]

                                b[3,ix,iy,iz,it,ialpha] = conj(u[-μ][1,3,ix1,iy1,iz1,it1])*a[1,ix1,iy1,iz1,it1,ialpha] +
                                                            conj(u[-μ][2,3,ix1,iy1,iz1,it1])*a[2,ix1,iy1,iz1,it1,ialpha] +
                                                            conj(u[-μ][3,3,ix1,iy1,iz1,it1])*a[3,ix1,iy1,iz1,it1,ialpha]
                                #=
                                for k1=1:3
                                    b[k1,ix,iy,iz,it,ialpha] = 0
                                    for k2=1:3
                                        b[k1,ix,iy,iz,it,ialpha] += conj(u[-μ][k2,k1,ix1,iy1,iz1,it1])*a[k2,ix1,iy1,iz1,it1,ialpha]
                                    end
                                end
                                =#
                            end
                        end
                    end
                end
            end
        end

    end

    function fermion_shift_gamma!(b::WilsonFermion,u::Array{T,1},μ::Int,a::WilsonFermion) where T <: SU3GaugeFields
        if μ == 0
            substitute!(b,a)
            return
        end

        NX = a.NX
        NY = a.NY
        NZ = a.NZ
        NT = a.NT
        NC = 3#a.NC

        #NTrange = get_looprange(NT)
        #println(NTrange)
        if μ > 0
            #idel = zeros(Int64,4)
            #idel[μ] = 1

            n6 = size(a.f)[6]
            for ialpha=1:4
                #for it=NTrange
                for it=1:NT
                    it1 = it + ifelse(μ ==4,1,0) #idel[4]
                    for iz=1:NZ
                        iz1 = iz + ifelse(μ ==3,1,0) #idel[3]
                        for iy=1:NY
                            iy1 = iy + ifelse(μ ==2,1,0) #idel[2]
                            @simd for ix=1:NX
                                ix1 = ix + ifelse(μ ==1,1,0) #idel[1]

                                b[1,ix,iy,iz,it,ialpha] = u[μ][1,1,ix,iy,iz,it]*a[1,ix1,iy1,iz1,it1,ialpha] +
                                                            u[μ][1,2,ix,iy,iz,it]*a[2,ix1,iy1,iz1,it1,ialpha] +
                                                            u[μ][1,3,ix,iy,iz,it]*a[3,ix1,iy1,iz1,it1,ialpha]

                                b[2,ix,iy,iz,it,ialpha] = u[μ][2,1,ix,iy,iz,it]*a[1,ix1,iy1,iz1,it1,ialpha] +
                                                            u[μ][2,2,ix,iy,iz,it]*a[2,ix1,iy1,iz1,it1,ialpha] +
                                                            u[μ][2,3,ix,iy,iz,it]*a[3,ix1,iy1,iz1,it1,ialpha]

                                b[3,ix,iy,iz,it,ialpha] = u[μ][3,1,ix,iy,iz,it]*a[1,ix1,iy1,iz1,it1,ialpha] +
                                                            u[μ][3,2,ix,iy,iz,it]*a[2,ix1,iy1,iz1,it1,ialpha] +
                                                            u[μ][3,3,ix,iy,iz,it]*a[3,ix1,iy1,iz1,it1,ialpha]

                                #=
                                for k1=1:3
                                    b[k1,ix,iy,iz,it,ialpha] = 0
                                    for k2=1:3
                                        b[k1,ix,iy,iz,it,ialpha] += u[μ][k1,k2,ix,iy,iz,it]*a[k2,ix1,iy1,iz1,it1,ialpha]
                                    end
                                end
                                =#
                            end
                        end
                    end
                end
            end

        elseif μ < 0
            #idel = zeros(Int64,4)
            #idel[-μ] = 1
            #n6 = size(b.f)[6]
            for ialpha =1:4
                for it=1:NT
                    it1 = it - ifelse(-μ ==4,1,0) #idel[4]
                    for iz=1:NZ
                        iz1 = iz - ifelse(-μ ==3,1,0) #idel[3]
                        for iy=1:NY
                            iy1 = iy - ifelse(-μ ==2,1,0)  #idel[2]
                            @simd for ix=1:NX
                                ix1 = ix - ifelse(-μ ==1,1,0) #idel[1]

                                b[1,ix,iy,iz,it,ialpha] = conj(u[-μ][1,1,ix1,iy1,iz1,it1])*a[1,ix1,iy1,iz1,it1,ialpha] +
                                                            conj(u[-μ][2,1,ix1,iy1,iz1,it1])*a[2,ix1,iy1,iz1,it1,ialpha] +
                                                            conj(u[-μ][3,1,ix1,iy1,iz1,it1])*a[3,ix1,iy1,iz1,it1,ialpha]

                                b[2,ix,iy,iz,it,ialpha] = conj(u[-μ][1,2,ix1,iy1,iz1,it1])*a[1,ix1,iy1,iz1,it1,ialpha] +
                                                            conj(u[-μ][2,2,ix1,iy1,iz1,it1])*a[2,ix1,iy1,iz1,it1,ialpha] +
                                                            conj(u[-μ][3,2,ix1,iy1,iz1,it1])*a[3,ix1,iy1,iz1,it1,ialpha]

                                b[3,ix,iy,iz,it,ialpha] = conj(u[-μ][1,3,ix1,iy1,iz1,it1])*a[1,ix1,iy1,iz1,it1,ialpha] +
                                                            conj(u[-μ][2,3,ix1,iy1,iz1,it1])*a[2,ix1,iy1,iz1,it1,ialpha] +
                                                            conj(u[-μ][3,3,ix1,iy1,iz1,it1])*a[3,ix1,iy1,iz1,it1,ialpha]
                                #=
                                for k1=1:3
                                    b[k1,ix,iy,iz,it,ialpha] = 0
                                    for k2=1:3
                                        b[k1,ix,iy,iz,it,ialpha] += conj(u[-μ][k2,k1,ix1,iy1,iz1,it1])*a[k2,ix1,iy1,iz1,it1,ialpha]
                                    end
                                end
                                =#
                            end
                        end
                    end
                end
            end
        end

    end

    function fermion_shift!(b::WilsonFermion,u::Array{T,1},μ::Int,a::WilsonFermion) where T <: SU2GaugeFields
        if μ == 0
            substitute!(b,a)
            return
        end

        NX = a.NX
        NY = a.NY
        NZ = a.NZ
        NT = a.NT
        NC = a.NC

        #NTrange = get_looprange(NT)
        #println(NTrange)
        if μ > 0
            #idel = zeros(Int64,4)
            #idel[μ] = 1

            n6 = size(a.f)[6]
            for ialpha=1:4
                #for it=NTrange
                for it=1:NT
                    it1 = it + ifelse(μ ==4,1,0) #idel[4]
                    for iz=1:NZ
                        iz1 = iz + ifelse(μ ==3,1,0) #idel[3]
                        for iy=1:NY
                            iy1 = iy + ifelse(μ ==2,1,0) #idel[2]
                            for ix=1:NX
                                ix1 = ix + ifelse(μ ==1,1,0) #idel[1]

                                b[1,ix,iy,iz,it,ialpha] = u[μ][1,1,ix,iy,iz,it]*a[1,ix1,iy1,iz1,it1,ialpha] +
                                                            u[μ][1,2,ix,iy,iz,it]*a[2,ix1,iy1,iz1,it1,ialpha]

                                b[2,ix,iy,iz,it,ialpha] = u[μ][2,1,ix,iy,iz,it]*a[1,ix1,iy1,iz1,it1,ialpha] +
                                                            u[μ][2,2,ix,iy,iz,it]*a[2,ix1,iy1,iz1,it1,ialpha]

                                #=
                                for k1=1:3
                                    b[k1,ix,iy,iz,it,ialpha] = 0
                                    for k2=1:3
                                        b[k1,ix,iy,iz,it,ialpha] += u[μ][k1,k2,ix,iy,iz,it]*a[k2,ix1,iy1,iz1,it1,ialpha]
                                    end
                                end
                                =#
                            end
                        end
                    end
                end
            end

        elseif μ < 0
            #idel = zeros(Int64,4)
            #idel[-μ] = 1
            #n6 = size(b.f)[6]
            for ialpha =1:4
                for it=1:NT
                    it1 = it - ifelse(-μ ==4,1,0) #idel[4]
                    for iz=1:NZ
                        iz1 = iz - ifelse(-μ ==3,1,0) #idel[3]
                        for iy=1:NY
                            iy1 = iy - ifelse(-μ ==2,1,0)  #idel[2]
                            for ix=1:NX
                                ix1 = ix - ifelse(-μ ==1,1,0) #idel[1]

                                b[1,ix,iy,iz,it,ialpha] = conj(u[-μ][1,1,ix1,iy1,iz1,it1])*a[1,ix1,iy1,iz1,it1,ialpha] +
                                                            conj(u[-μ][2,1,ix1,iy1,iz1,it1])*a[2,ix1,iy1,iz1,it1,ialpha]

                                b[2,ix,iy,iz,it,ialpha] = conj(u[-μ][1,2,ix1,iy1,iz1,it1])*a[1,ix1,iy1,iz1,it1,ialpha] +
                                                            conj(u[-μ][2,2,ix1,iy1,iz1,it1])*a[2,ix1,iy1,iz1,it1,ialpha]

                            end
                        end
                    end
                end
            end
        end

    end

    function fermion_shiftB!(b::WilsonFermion,u::Array{GaugeFields{SU{NC}},1},μ,a::WilsonFermion) where NC
        if μ == 0
            substitute!(b,a)
            return
        end

        NX = a.NX
        NY = a.NY
        NZ = a.NZ
        NT = a.NT
        #NC = a.NC

        if μ > 0
            error("""
            Sorry this case if not yet ready
            mu = $mu
            """)

        elseif μ < 0
            #idel = zeros(Int64,4)
            #idel[-μ] = 1
            #n6 = size(b.f)[6]
            for ialpha =1:4
                for it=1:NT
                    it1 = it + ifelse(-μ == 4,1,0) #idel[4]
                    for iz=1:NZ
                        iz1 = iz +ifelse(-μ == 3,1,0)  #idel[3]
                        for iy=1:NY
                            iy1 = iy + ifelse(-μ == 2,1,0) #idel[2]
                            for ix=1:NX
                                ix1 = ix + ifelse(-μ == 1,1,0)  #idel[1]

                                for k1=1:NC
                                    b[k1,ix,iy,iz,it,ialpha] = 0
                                    for k2=1:NC
                                        b[k1,ix,iy,iz,it,ialpha] += conj(a[k2,ix1,iy1,iz1,it1,ialpha])*conj(u[-μ][k1,k2,ix,iy,iz,it])
                                    end
                                end

                            end
                        end
                    end
                end
            end
        end

    end

    function fermion_shiftB!(b::WilsonFermion,evensite,u::Array{GaugeFields{SU{NC}},1},μ,a::WilsonFermion)  where NC
        if μ == 0
            substitute!(b,a)
            return
        end

        ibush = ifelse(evensite,0,1)

        NX = a.NX
        NY = a.NY
        NZ = a.NZ
        NT = a.NT
        #NC = a.NC

        if μ > 0
            error("""
            Sorry this case if not yet ready
            mu = $mu
            """)

        elseif μ < 0
            #idel = zeros(Int64,4)
            #idel[-μ] = 1
            #n6 = size(b.f)[6]
            for ialpha =1:4
                for it=1:NT
                    it1 = it + ifelse(-μ == 4,1,0) #idel[4]
                    for iz=1:NZ
                        iz1 = iz +ifelse(-μ == 3,1,0)  #idel[3]
                        for iy=1:NY
                            iy1 = iy + ifelse(-μ == 2,1,0) #idel[2]
                            xran =1+(1+ibush+iy+iz+it)%2:2:NX
                            for ix in xran
                            #for ix=1:NX
                                ix1 = ix + ifelse(-μ == 1,1,0)  #idel[1]

                                for k1=1:NC
                                    b[k1,ix,iy,iz,it,ialpha] = 0
                                    for k2=1:NC
                                        b[k1,ix,iy,iz,it,ialpha] += conj(a[k2,ix1,iy1,iz1,it1,ialpha])*conj(u[-μ][k1,k2,ix,iy,iz,it])
                                    end
                                end

                            end
                        end
                    end
                end
            end
        end

    end

    function fermion_shiftB!(b::WilsonFermion,u::Array{SU3GaugeFields,1},μ,a::WilsonFermion)
        if μ == 0
            substitute!(b,a)
            return
        end

        NX = a.NX
        NY = a.NY
        NZ = a.NZ
        NT = a.NT
        NC = a.NC

        if μ > 0
            error("""
            Sorry this case if not yet ready
            mu = $mu
            """)

        elseif μ < 0
            #idel = zeros(Int64,4)
            #idel[-μ] = 1
            #n6 = size(b.f)[6]
            for ialpha =1:4
                for it=1:NT
                    it1 = it + ifelse(-μ == 4,1,0) #idel[4]
                    for iz=1:NZ
                        iz1 = iz +ifelse(-μ == 3,1,0)  #idel[3]
                        for iy=1:NY
                            iy1 = iy + ifelse(-μ == 2,1,0) #idel[2]
                            for ix=1:NX
                                ix1 = ix + ifelse(-μ == 1,1,0)  #idel[1]

                                b[1,ix,iy,iz,it,ialpha] = conj(a[1,ix1,iy1,iz1,it1,ialpha])*conj(u[-μ][1,1,ix,iy,iz,it]) +
                                                            conj(a[2,ix1,iy1,iz1,it1,ialpha])*conj(u[-μ][1,2,ix,iy,iz,it]) +
                                                            conj(a[3,ix1,iy1,iz1,it1,ialpha])* conj(u[-μ][1,3,ix,iy,iz,it])

                                b[2,ix,iy,iz,it,ialpha] = conj(a[1,ix1,iy1,iz1,it1,ialpha])*conj(u[-μ][2,1,ix,iy,iz,it])+
                                                            conj(a[2,ix1,iy1,iz1,it1,ialpha])*conj(u[-μ][2,2,ix,iy,iz,it]) +
                                                            conj(a[3,ix1,iy1,iz1,it1,ialpha])*conj(u[-μ][2,3,ix,iy,iz,it])

                                b[3,ix,iy,iz,it,ialpha] = conj(a[1,ix1,iy1,iz1,it1,ialpha])*conj(u[-μ][3,1,ix,iy,iz,it]) +
                                                            conj(a[2,ix1,iy1,iz1,it1,ialpha])*conj(u[-μ][3,2,ix,iy,iz,it]) +
                                                            conj(a[3,ix1,iy1,iz1,it1,ialpha])*conj(u[-μ][3,3,ix,iy,iz,it])

                                #=
                                for k1=1:NC
                                    b[k1,ix,iy,iz,it,ialpha] = 0
                                    for k2=1:NC
                                        b[k1,ix,iy,iz,it,ialpha] += conj(a[k2,ix1,iy1,iz1,it1,ialpha])*conj(u[-μ][k1,k2,ix,iy,iz,it])
                                    end
                                end
                                =#

                            end
                        end
                    end
                end
            end
        end

    end

    function fermion_shiftB!(b::WilsonFermion,evensite,u::Array{SU3GaugeFields,1},μ,a::WilsonFermion)
        if μ == 0
            substitute!(b,a)
            return
        end

        ibush = ifelse(evensite,0,1)

        NX = a.NX
        NY = a.NY
        NZ = a.NZ
        NT = a.NT
        NC = a.NC

        if μ > 0
            error("""
            Sorry this case if not yet ready
            mu = $mu
            """)

        elseif μ < 0
            #idel = zeros(Int64,4)
            #idel[-μ] = 1
            #n6 = size(b.f)[6]
            for ialpha =1:4
                for it=1:NT
                    it1 = it + ifelse(-μ == 4,1,0) #idel[4]
                    for iz=1:NZ
                        iz1 = iz +ifelse(-μ == 3,1,0)  #idel[3]
                        for iy=1:NY
                            iy1 = iy + ifelse(-μ == 2,1,0) #idel[2]
                            xran =1+(1+ibush+iy+iz+it)%2:2:NX
                            for ix in xran
                            #for ix=1:NX
                                ix1 = ix + ifelse(-μ == 1,1,0)  #idel[1]

                                b[1,ix,iy,iz,it,ialpha] = conj(a[1,ix1,iy1,iz1,it1,ialpha])*conj(u[-μ][1,1,ix,iy,iz,it]) +
                                                            conj(a[2,ix1,iy1,iz1,it1,ialpha])*conj(u[-μ][1,2,ix,iy,iz,it]) +
                                                            conj(a[3,ix1,iy1,iz1,it1,ialpha])* conj(u[-μ][1,3,ix,iy,iz,it])

                                b[2,ix,iy,iz,it,ialpha] = conj(a[1,ix1,iy1,iz1,it1,ialpha])*conj(u[-μ][2,1,ix,iy,iz,it])+
                                                            conj(a[2,ix1,iy1,iz1,it1,ialpha])*conj(u[-μ][2,2,ix,iy,iz,it]) +
                                                            conj(a[3,ix1,iy1,iz1,it1,ialpha])*conj(u[-μ][2,3,ix,iy,iz,it])

                                b[3,ix,iy,iz,it,ialpha] = conj(a[1,ix1,iy1,iz1,it1,ialpha])*conj(u[-μ][3,1,ix,iy,iz,it]) +
                                                            conj(a[2,ix1,iy1,iz1,it1,ialpha])*conj(u[-μ][3,2,ix,iy,iz,it]) +
                                                            conj(a[3,ix1,iy1,iz1,it1,ialpha])*conj(u[-μ][3,3,ix,iy,iz,it])

                                #=
                                for k1=1:NC
                                    b[k1,ix,iy,iz,it,ialpha] = 0
                                    for k2=1:NC
                                        b[k1,ix,iy,iz,it,ialpha] += conj(a[k2,ix1,iy1,iz1,it1,ialpha])*conj(u[-μ][k1,k2,ix,iy,iz,it])
                                    end
                                end
                                =#

                            end
                        end
                    end
                end
            end
        end

    end

    function fermion_shiftB!(b::WilsonFermion,u::Array{SU2GaugeFields,1},μ,a::WilsonFermion)
        if μ == 0
            substitute!(b,a)
            return
        end

        NX = a.NX
        NY = a.NY
        NZ = a.NZ
        NT = a.NT
        NC = a.NC

        if μ > 0
            error("""
            Sorry this case if not yet ready
            mu = $mu
            """)

        elseif μ < 0
            #idel = zeros(Int64,4)
            #idel[-μ] = 1
            #n6 = size(b.f)[6]
            for ialpha =1:4
                for it=1:NT
                    it1 = it + ifelse(-μ == 4,1,0) #idel[4]
                    for iz=1:NZ
                        iz1 = iz +ifelse(-μ == 3,1,0)  #idel[3]
                        for iy=1:NY
                            iy1 = iy + ifelse(-μ == 2,1,0) #idel[2]
                            for ix=1:NX
                                ix1 = ix + ifelse(-μ == 1,1,0)  #idel[1]

                                b[1,ix,iy,iz,it,ialpha] = conj(a[1,ix1,iy1,iz1,it1,ialpha])*conj(u[-μ][1,1,ix,iy,iz,it]) +
                                                            conj(a[2,ix1,iy1,iz1,it1,ialpha])*conj(u[-μ][1,2,ix,iy,iz,it])

                                b[2,ix,iy,iz,it,ialpha] = conj(a[1,ix1,iy1,iz1,it1,ialpha])*conj(u[-μ][2,1,ix,iy,iz,it])+
                                                            conj(a[2,ix1,iy1,iz1,it1,ialpha])*conj(u[-μ][2,2,ix,iy,iz,it])

                            end
                        end
                    end
                end
            end
        end

    end

    function fermion_shiftB!(b::WilsonFermion,evensite,u::Array{SU2GaugeFields,1},μ,a::WilsonFermion)
        if μ == 0
            substitute!(b,a)
            return
        end

        ibush = ifelse(evensite,0,1)

        NX = a.NX
        NY = a.NY
        NZ = a.NZ
        NT = a.NT
        NC = a.NC

        if μ > 0
            error("""
            Sorry this case if not yet ready
            mu = $mu
            """)

        elseif μ < 0
            #idel = zeros(Int64,4)
            #idel[-μ] = 1
            #n6 = size(b.f)[6]
            for ialpha =1:4
                for it=1:NT
                    it1 = it + ifelse(-μ == 4,1,0) #idel[4]
                    for iz=1:NZ
                        iz1 = iz +ifelse(-μ == 3,1,0)  #idel[3]
                        for iy=1:NY
                            iy1 = iy + ifelse(-μ == 2,1,0) #idel[2]
                            xran =1+(1+ibush+iy+iz+it)%2:2:NX
                            for ix in xran
                                ix1 = ix + ifelse(-μ == 1,1,0)  #idel[1]

                                b[1,ix,iy,iz,it,ialpha] = conj(a[1,ix1,iy1,iz1,it1,ialpha])*conj(u[-μ][1,1,ix,iy,iz,it]) +
                                                            conj(a[2,ix1,iy1,iz1,it1,ialpha])*conj(u[-μ][1,2,ix,iy,iz,it])

                                b[2,ix,iy,iz,it,ialpha] = conj(a[1,ix1,iy1,iz1,it1,ialpha])*conj(u[-μ][2,1,ix,iy,iz,it])+
                                                            conj(a[2,ix1,iy1,iz1,it1,ialpha])*conj(u[-μ][2,2,ix,iy,iz,it])

                            end
                        end
                    end
                end
            end
        end

    end

    """
    c--------------------------------------------------------------------------c
    c     y = gamma_5 * x
    c     here
    c                  ( -1       )
    c        GAMMA5 =  (   -1     )
    c                  (     +1   )
    c                  (       +1 )
    c--------------------------------------------------------------------------c
        """
        function mul_γ5x!(y::WilsonFermion,x::WilsonFermion)
            NX = x.NX
            NY = x.NY
            NZ = x.NZ
            NT = x.NT
            NC = x.NC
            for ic=1:NC
                for it=1:NT
                    for iz=1:NZ
                        for iy=1:NY
                            @simd for ix=1:NX
                                y[ic,ix,iy,iz,it,1] = -1*x[ic,ix,iy,iz,it,1]
                                y[ic,ix,iy,iz,it,2] = -1*x[ic,ix,iy,iz,it,2]
                                y[ic,ix,iy,iz,it,3] = x[ic,ix,iy,iz,it,3]
                                y[ic,ix,iy,iz,it,4] = x[ic,ix,iy,iz,it,4]
                            end
                        end
                    end
                end
            end
        end

        function mul_1plusγ5x!(y::WilsonFermion,x::WilsonFermion) #(1+gamma_5)/2
            NX = x.NX
            NY = x.NY
            NZ = x.NZ
            NT = x.NT
            NC = x.NC
            for ic=1:NC
                for it=1:NT
                    for iz=1:NZ
                        for iy=1:NY
                            @simd for ix=1:NX
                                y[ic,ix,iy,iz,it,1] = 0#-1*x[ic,ix,iy,iz,it,1]
                                y[ic,ix,iy,iz,it,2] = 0#-1*x[ic,ix,iy,iz,it,2]
                                y[ic,ix,iy,iz,it,3] = x[ic,ix,iy,iz,it,3]
                                y[ic,ix,iy,iz,it,4] = x[ic,ix,iy,iz,it,4]
                            end
                        end
                    end
                end
            end
        end

        function mul_1plusγ5x_add!(y::WilsonFermion,x::WilsonFermion,factor) #x = x +(1+gamma_5)/2
            NX = x.NX
            NY = x.NY
            NZ = x.NZ
            NT = x.NT
            NC = x.NC
            for ic=1:NC
                for it=1:NT
                    for iz=1:NZ
                        for iy=1:NY
                            @simd for ix=1:NX
                                y[ic,ix,iy,iz,it,3] += factor*x[ic,ix,iy,iz,it,3]
                                y[ic,ix,iy,iz,it,4] += factor*x[ic,ix,iy,iz,it,4]
                            end
                        end
                    end
                end
            end
        end

        function mul_1minusγ5x!(y::WilsonFermion,x::WilsonFermion) #(1-gamma_5)/2
            NX = x.NX
            NY = x.NY
            NZ = x.NZ
            NT = x.NT
            NC = x.NC
            for ic=1:NC
                for it=1:NT
                    for iz=1:NZ
                        for iy=1:NY
                            @simd for ix=1:NX
                                y[ic,ix,iy,iz,it,1] = x[ic,ix,iy,iz,it,1]
                                y[ic,ix,iy,iz,it,2] = x[ic,ix,iy,iz,it,2]
                                y[ic,ix,iy,iz,it,3] = 0#x[ic,ix,iy,iz,it,3]
                                y[ic,ix,iy,iz,it,4] = 0#x[ic,ix,iy,iz,it,4]
                            end
                        end
                    end
                end
            end
        end

        function mul_1minusγ5x_add!(y::WilsonFermion,x::WilsonFermion,factor) #+(1-gamma_5)/2
            NX = x.NX
            NY = x.NY
            NZ = x.NZ
            NT = x.NT
            NC = x.NC
            for ic=1:NC
                for it=1:NT
                    for iz=1:NZ
                        for iy=1:NY
                            @simd for ix=1:NX
                                y[ic,ix,iy,iz,it,1] += factor*x[ic,ix,iy,iz,it,1]
                                y[ic,ix,iy,iz,it,2] += factor*x[ic,ix,iy,iz,it,2]
                            end
                        end
                    end
                end
            end
        end

    """
    mk_gamma()
    c----------------------------------------------------------------------c
c     Make gamma matrix
c----------------------------------------------------------------------c
C     THE CONVENTION OF THE GAMMA MATRIX HERE
C     ( EUCLIDEAN CHIRAL REPRESENTATION )
C
C               (       -i )              (       -1 )
C     GAMMA1 =  (     -i   )     GAMMA2 = (     +1   )
C               (   +i     )              (   +1     )
C               ( +i       )              ( -1       )
C
C               (     -i   )              (     -1   )
C     GAMMA3 =  (       +i )     GAMMA4 = (       -1 )
C               ( +i       )              ( -1       )
C               (   -i     )              (   -1     )
C
C               ( -1       )
C     GAMMA5 =  (   -1     )
C               (     +1   )
C               (       +1 )
C
C     ( GAMMA_MU, GAMMA_NU ) = 2*DEL_MU,NU   FOR MU,NU=1,2,3,4
c----------------------------------------------------------------------c
    """
    function mk_gamma(r)
        g0 = zeros(ComplexF64,4,4)
        g1 = zero(g0)
        g2 = zero(g1)
        g3 = zero(g1)
        g4 = zero(g1)
        g5 = zero(g1)
        gamma = zeros(ComplexF64,4,4,5)
        rpg = zero(gamma)
        rmg = zero(gamma)

        g0[1,1]=1.0; g0[1,2]=0.0; g0[1,3]=0.0; g0[1,4]=0.0
        g0[2,1]=0.0; g0[2,2]=1.0; g0[2,3]=0.0; g0[2,4]=0.0
        g0[3,1]=0.0; g0[3,2]=0.0; g0[3,3]=1.0; g0[3,4]=0.0
        g0[4,1]=0.0; g0[4,2]=0.0; g0[4,3]=0.0; g0[4,4]=1.0

        g1[1,1]=0.0; g1[1,2]=0.0; g1[1,3]=0.0; g1[1,4]=-im
        g1[2,1]=0.0; g1[2,2]=0.0; g1[2,3]=-im;  g1[2,4]=0.0
        g1[3,1]=0.0; g1[3,2]=+im;  g1[3,3]=0.0; g1[3,4]=0.0
        g1[4,1]=+im;  g1[4,2]=0.0; g1[4,3]=0.0; g1[4,4]=0.0

        g2[1,1]=0.0; g2[1,2]=0.0; g2[1,3]=0.0; g2[1,4]=-1.0
        g2[2,1]=0.0; g2[2,2]=0.0; g2[2,3]=1.0; g2[2,4]=0.0
        g2[3,1]=0.0; g2[3,2]=1.0; g2[3,3]=0.0; g2[3,4]=0.0
        g2[4,1]=-1.0;g2[4,2]=0.0; g2[4,3]=0.0; g2[4,4]=0.0

        g3[1,1]=0.0; g3[1,2]=0.0; g3[1,3]=-im;  g3[1,4]=0.0
        g3[2,1]=0.0; g3[2,2]=0.0; g3[2,3]=0.0; g3[2,4]=+im
        g3[3,1]=+im;  g3[3,2]=0.0; g3[3,3]=0.0; g3[3,4]=0.0
        g3[4,1]=0.0; g3[4,2]=-im;  g3[4,3]=0.0; g3[4,4]=0.0

        g4[1,1]=0.0; g4[1,2]=0.0; g4[1,3]=-1.0;g4[1,4]=0.0
        g4[2,1]=0.0; g4[2,2]=0.0; g4[2,3]=0.0; g4[2,4]=-1.0
        g4[3,1]=-1.0;g4[3,2]=0.0; g4[3,3]=0.0; g4[3,4]=0.0
        g4[4,1]=0.0; g4[4,2]=-1.0;g4[4,3]=0.0; g4[4,4]=0.0

        g5[1,1]=-1.0;g5[1,2]=0.0; g5[1,3]=0.0; g5[1,4]=0.0
        g5[2,1]=0.0; g5[2,2]=-1.0;g5[2,3]=0.0; g5[2,4]=0.0
        g5[3,1]=0.0; g5[3,2]=0.0; g5[3,3]=1.0; g5[3,4]=0.0
        g5[4,1]=0.0; g5[4,2]=0.0; g5[4,3]=0.0; g5[4,4]=1.0

        gamma[:,:,1] = g1[:,:]
        gamma[:,:,2] = g2[:,:]
        gamma[:,:,3] = g3[:,:]
        gamma[:,:,4] = g4[:,:]
        gamma[:,:,5] = g5[:,:]

        for mu=1:4
            for j=1:4
                for i=1:4
                    rpg[i,j,mu] = r*g0[i,j] + gamma[i,j,mu]
                    rmg[i,j,mu] = r*g0[i,j] - gamma[i,j,mu]
                end
            end
        end

        return gamma,rpg,rmg

    end

    function cloverterm!(vec,CloverFμν,x)
        NT = x.NT
        NZ = x.NZ
        NY = x.NY
        NX = x.NX
        NC = x.NC

        i  =0
        for it=1:NT
            for iz=1:NZ
                for iy=1:NY
                    for ix=1:NX
                        i += 1
                        for k1=1:NC
                            for k2=1:NC

                                c1 = x[k2,ix,iy,iz,it,1]
                                c2 = x[k2,ix,iy,iz,it,2]
                                c3 = x[k2,ix,iy,iz,it,3]
                                c4 = x[k2,ix,iy,iz,it,4]

                                vec[k1,ix,iy,iz,it,1] += CloverFμν[k1,k2,i,1]*(-   c1) +
                                                            + CloverFμν[k1,k2,i,2]*(-im*c2) +
                                                            + CloverFμν[k1,k2,i,3]*(-   c2) +
                                                            + CloverFμν[k1,k2,i,4]*(-   c2) +
                                                            + CloverFμν[k1,k2,i,5]*( im*c2) +
                                                            + CloverFμν[k1,k2,i,6]*(-   c1)

                                vec[k1,ix,iy,iz,it,2] += CloverFμν[k1,k2,i,1]*(   c2) +
                                                            + CloverFμν[k1,k2,i,2]*(im*c1) +
                                                            + CloverFμν[k1,k2,i,3]*(-   c1) +
                                                            + CloverFμν[k1,k2,i,4]*(-   c1) +
                                                            + CloverFμν[k1,k2,i,5]*(-im*c1) +
                                                            + CloverFμν[k1,k2,i,6]*(   c2)

                                vec[k1,ix,iy,iz,it,3] += CloverFμν[k1,k2,i,1]*(   -c3) +
                                                            + CloverFμν[k1,k2,i,2]*(-im*c4) +
                                                            + CloverFμν[k1,k2,i,3]*(   c4) +
                                                            + CloverFμν[k1,k2,i,4]*(-   c4) +
                                                            + CloverFμν[k1,k2,i,5]*(-im*c4) +
                                                            + CloverFμν[k1,k2,i,6]*(   c3)

                                vec[k1,ix,iy,iz,it,4] += CloverFμν[k1,k2,i,1]*(   c4) +
                                                            + CloverFμν[k1,k2,i,2]*(im*c3) +
                                                            + CloverFμν[k1,k2,i,3]*(   c3) +
                                                            + CloverFμν[k1,k2,i,4]*(-   c3) +
                                                            + CloverFμν[k1,k2,i,5]*(im*c3) +
                                                            + CloverFμν[k1,k2,i,6]*( -  c4)

                            end
                        end
                    end
                end

            end
        end

        #println("vec = ",vec*vec)

    end

end


module UNSflow

using Dierckx
export Spline1D, derivative, evaluate

using ForwardDiff
export derivative

using PyPlot
export plot, scatter, figure, xlabel, ylabel, xlim, ylim, xticks, yticks, subplot, axes, legend, markers, axis

using NLsolve
export nlsolve, not_in_place

export camber_calc, update_boundpos, update_kinem, update_indbound,
update_downwash, update_a0anda1, place_tev, place_lev, update_a2toan,
mutual_ind, trapz, update_a2a3adot, update_bv, ind_vel, view_vorts,
wakeroll, lautat, lautat_wakeroll, ldvm, calc_forces, design_solve,
lesp_design_max, transfer_cm, theodorsen, anim_flow,
drone_trajectory_problem, simple_LLT, find_tstep, ldvm_E, update_adot,
ldvm_more, ldvm_E_more, calc_forces_more, calc_forces_E_more, interp,
QScorrect_lautat, QScorrect_ldvm, QSLLT_lautat, QSLLT_ldvm,
lautat_wakeroll_more, LLT_ldvm, update_externalvel

export KinemPar, KinemParwFlap, KinemPar2DOF, KinemPar2DFree,
MotionDef, KinemDef, KinemDefwFlap, EldUpDef, EldUptstartDef,
ConstDef, TwoDOFPar, TwoDFreePar, TwoDSurf, TwoDSurfwFlap,
TwoDSurf_2DOF, TwoDFreeSurf, TwoDVort, TwoDFlowField, KelvinCondition,
KelvinKutta, EldRampReturnDef, EldUpIntDef, EldUpInttstartDef, SinDef,
CosDef, TheoDef, TheoDefwFlap, TwoDFlowData, DelVortDef, patch,
ThreeDSurfSimple, KinemDef3D, ThreeDVort, ThreeDFieldSimple,
KelvinConditionLLTldvm, KelvinKuttaLLTldvm

include("types.jl")
include("calculations.jl")
include("solvers.jl")
include("postprocess.jl")

end


module LED

import ..ioctl

export led_matrix, RGB565, led_clear, @colorant_str

using Colors, ColorTypes, FixedPointNumbers, Dates, Mmap

function _led_fb_device()
    try
        for devname in readdir("/sys/class/graphics")
            sysfname = joinpath("/sys/class/graphics",devname,"name")
            if startswith(devname, "fb") && isfile(sysfname)
                if startswith(readline(sysfname),"RPi-Sense FB")
                    return joinpath("/dev",devname)
                end
            end
        end
    catch e
        @warn("LED: Sense Hat not found.")
        return nothing
    end
end

const LED_FB_DEVICE_PATH = _led_fb_device()
const LED_FB_DEVICE = Ref{IOStream}()

function __init__()
    if LED_FB_DEVICE_PATH != nothing
        LED_FB_DEVICE[] = open(LED_FB_DEVICE_PATH,"w+")
    end
end

const U5 = Normed{UInt8,5}
const U6 = Normed{UInt8,6}

struct RGB565 <: AbstractRGB{Normed{UInt8, 8}}
    data::UInt16
end

function RGB565(r::U5, g::U6, b::U5)
    RGB565( (UInt16(reinterpret(r)) << 11) |
            (UInt16(reinterpret(g)) << 5) |
            (UInt16(reinterpret(b))) )
end

RGB565(r::Real, g::Real, b::Real) =
    RGB565(convert(U5,r), convert(U6,g), convert(U5,b))
RGB565(c::Union{Color,Colorant}) = RGB565(red(c), green(c), blue(c))

ColorTypes.red(c::RGB565) = U5(c.data >> 11, Val{true})
ColorTypes.green(c::RGB565) = U6((c.data >> 5) & 0x3f, Val{true})
ColorTypes.blue(c::RGB565) = U5(c.data & 0x1f, Val{true})

ColorTypes.ccolor(::Type{RGB565}, ::Type{Csrc}) where {Csrc <: Colorant} = RGB565
ColorTypes.base_color_type(::Type{RGB565}) = RGB565

"""
    led_matrix()

Returns an 8x8 matrix of `RGB565` elements that is memory-mapped to the Sense HAT LED Matrix.

While it is possible to invoke the function multiple times (each returning different
arrays), it is generally preferable to assign it once into a `const` variable so as to
minimise the number of open file handlers.

# Example

```
using SenseHat
using ColorTypes

const LED = led_matrix()

LED[:] = SenseHat.JULIA_LOGO
sleep(3)
LED[:] = RGB(0,0,0)
```

"""
led_matrix() = Mmap.mmap(LED_FB_DEVICE[], Array{RGB565,2}, (8,8); grow=false)

const FBIO_GET_GAMMA = 61696
const FBIO_SET_GAMMA = 61697
const FBIO_RESET_GAMMA = 61698

@enum(GammaScale, GAMMA_DEFAULT, GAMMA_LOW, GAMMA_USER)

"""
    getgamma()

Return the current gamma correction table for the Sense HAT LED matrix, as a vector of 32
`UFixed{UInt8,5}` numbers.

"""
function getgamma()
    buffer = zeros(UInt8, 32)
    ioctl(LED_FB_DEVICE[], FBIO_GET_GAMMA, buffer)
    return buffer
end

"""
    setgamma(v::GammaScale)
    setgamma(table::Vector)

Set the gamma correction for the Sense HAT LED matrix. This can either be one of the
predefined settings (`GAMMA_DEFAULT`, `GAMMA_LOW`, `GAMMA_USER`) or a table of values
between 0 and 1 (the actual values will be converted to `UFixed{UInt8, 5}`).

"""
function setgamma end

function setgamma(v::GammaScale)
    ioctl(LED_FB_DEVICE[], FBIO_RESET_GAMMA, Cint(v))
    return nothing
end

function setgamma(table::Vector{U5})
    length(table) == 32 || error("Table must contain 32 values.")
    ioctl(LED_FB_DEVICE[], FBIO_SET_GAMMA, table)
    return nothing
end
setgamma(table::AbstractVector) = setgamma(convert(Vector{U5}, table))

function led_display(X)
    Base.depwarn("`led_display(X)` has been deprecated, use `led_matrix()[:] = X` instead.", :led_display)
    led_matrix()[:] = X
end

"""
    led_clear()

Sets the SenseHat LED matrix to all black.
"""
function led_clear()
    led_matrix()[:] .= RGB565(0,0,0)
    return
end

end # module


using ClimateMachine
const clima_dir = dirname(dirname(pathof(ClimateMachine)));

if parse(Bool, get(ENV, "CLIMATEMACHINE_PLOT_EDMF_COMPARISON", "false"))
    plot_dir = joinpath(clima_dir, "output", "bomex_edmf", "pycles_comparison")
else
    plot_dir = nothing
end

include(joinpath(@__DIR__, "compute_mse.jl"))

#! format: off
best_mse = Dict()
best_mse[:Bomex] = Dict()
best_mse[:Bomex]["ρ"] = 3.4943021267397123e-02
best_mse[:Bomex]["ρu[1]"] = 3.0714039084256679e+03
best_mse[:Bomex]["ρu[2]"] = 1.3375796498101822e-03
best_mse[:Bomex]["moisture.ρq_tot"] = 4.8463531712319707e-02
best_mse[:Bomex]["turbconv.environment.ρatke"] = 6.6626422991098286e+02
best_mse[:Bomex]["turbconv.environment.ρaθ_liq_cv"] = 8.5667099987715503e+01
best_mse[:Bomex]["turbconv.environment.ρaq_tot_cv"] = 1.6439116552012851e+02
best_mse[:Bomex]["turbconv.updraft[1].ρa"] = 7.9577348413012515e+01
best_mse[:Bomex]["turbconv.updraft[1].ρaw"] = 8.4352188057391225e-02
best_mse[:Bomex]["turbconv.updraft[1].ρaθ_liq"] = 9.0101464252959325e+00
best_mse[:Bomex]["turbconv.updraft[1].ρaq_tot"] = 1.0768120864370509e+01
#! format: on

sufficient_mse(computed_mse, best_mse) = computed_mse <= best_mse + eps()

function test_mse(computed_mse, best_mse, key)
    mse_not_regressed = sufficient_mse(computed_mse[key], best_mse[key])
    @test mse_not_regressed
    mse_not_regressed || @show key
end

computed_mse = Dict(
    k => compute_mse(
        solver_config.dg.grid,
        solver_config.dg.balance_law,
        time_data,
        dons_arr,
        data_files[k],
        k,
        best_mse[k],
        plot_dir,
    ) for k in keys(data_files)
)

@testset "BOMEX EDMF Solution Quality Assurance (QA) tests" begin
    #! format: off
    test_mse(computed_mse[:Bomex], best_mse[:Bomex], "ρ")
    test_mse(computed_mse[:Bomex], best_mse[:Bomex], "ρu[1]")
    test_mse(computed_mse[:Bomex], best_mse[:Bomex], "moisture.ρq_tot")
    test_mse(computed_mse[:Bomex], best_mse[:Bomex], "turbconv.updraft[1].ρa")
    test_mse(computed_mse[:Bomex], best_mse[:Bomex], "turbconv.updraft[1].ρaw")
    test_mse(computed_mse[:Bomex], best_mse[:Bomex], "turbconv.updraft[1].ρaθ_liq")
    test_mse(computed_mse[:Bomex], best_mse[:Bomex], "turbconv.updraft[1].ρaq_tot")
    test_mse(computed_mse[:Bomex], best_mse[:Bomex], "turbconv.environment.ρatke")
    test_mse(computed_mse[:Bomex], best_mse[:Bomex], "turbconv.environment.ρaθ_liq_cv")
    test_mse(computed_mse[:Bomex], best_mse[:Bomex], "turbconv.environment.ρaq_tot_cv")
    #! format: on
end


using Random
using Printf
using StatsBase
using StaticArrays
using Distributions
using POMDPs
using AutoViz
using AutomotiveDrivingModels
using AutomotivePOMDPs
using POMDPSimulators
using POMDPPolicies
using DeepQLearning
using Flux
using RLInterface
using POMDPModelTools
using BeliefUpdaters
using TensorBoardLogger
using ArgParse
using BSON
using CSV
using LinearAlgebra
using ProgressMeter
include("environment.jl")
include("generative_mdp.jl")
include("masking.jl")
include("cooperative_IDM.jl")
include("belief_updater.jl")
include("belief_mdp.jl")

# BLAS.set_num_threads(8)

s = ArgParseSettings()
@add_arg_table s begin
    "--seed"
        help = "specify the random seed"
        arg_type = Int
        default = 1
    "--training_steps"
        help = "specify the number of training steps"
        arg_type = Int
        default = 10_000
    "--eps_fraction"
        help = "Fraction of the training set to use to decay epsilon from 1.0 to eps_end"
        arg_type = Float64
        default = 0.5
    "--eps_end"
        help = "epsilon value at the end of the exploration phase"
        arg_type = Float64
        default = 0.01
    "--learning_rate"
        help = "learning rate for DQN"
        arg_type = Float64
        default = 1e-4
    "--target_update_freq"
        help = "target update frequency for DQN"
        default = 5000
    "--logdir"
        help = "Directory in which to save the model and log training data"
        arg_type = String
        default = "log"
    "--n_eval"
        help = "Number of episodes for evaluation"
        arg_type = Int64
        default = 1000
    "--max_cars"
        help = "Number of cars on the main road"
        arg_type = Int64
        default = 12
    "--min_cars"
        help = "Number of cars on the main road"
        arg_type = Int64
        default = 0
    "--cooperation"
        help = "whether the ego vehicle observe the cooperation level or not"
        action = :store_true
    "--driver_type"
        help = "the distribution of driver types, choose among cooperative, aggressive or random"
        arg_type = String
        default = "random"
    "--traffic_speed"
        arg_type = String
        default = "mixed"
    "--recurrent"
        help = "whether to use an RNN in DQN"
        action = :store_true
    "--load"
        help = "load a pretrain policy"
        arg_type = Union{Nothing, String}
        default = nothing
    "--collision_cost"
        help = "collision cost in the mdp formulation"
        arg_type = Float64
        default = -1.0
    "--belief"
        help = "train with belief updater"
        action = :store_true
end
parsed_args = parse_args(s)

seed = parsed_args["seed"]
Random.seed!(seed)
rng = MersenneTwister(seed)

mdp = GenerativeMergingMDP(random_n_cars = true,
                            collision_cost = parsed_args["collision_cost"],
                            max_cars = parsed_args["max_cars"],
                            min_cars = parsed_args["min_cars"],
                           driver_type = Symbol(parsed_args["driver_type"]),
                           traffic_speed = Symbol(parsed_args["traffic_speed"]),
                           observe_cooperation=parsed_args["cooperation"])

if parsed_args["belief"]
    for i=2:mdp.max_cars+1
        mdp.driver_models[i] = CooperativeIDM()
    end
    pomdp = FullyObservablePOMDP(mdp)
    up = MergingUpdater(mdp)
    bmdp = GenerativeBeliefMDP{typeof(pomdp), MergingUpdater, MergingBelief,Int64}(pomdp, up)
    s0 = initialstate(bmdp, rng)
    mdp = bmdp
else
    s0 = initialstate(mdp, rng)
end

svec = convert_s(Vector{Float64}, s0, mdp)

input_dims = length(svec)

if parsed_args["recurrent"]
    model = Chain(LSTM(input_dims, 64), Dense(64, 32, relu), Dense(32, n_actions(mdp)))
else
    model = Chain(Dense(input_dims, 64, relu), Dense(64, 32, relu), Dense(32, n_actions(mdp)))
end

if parsed_args["load"] != nothing
    BSON.@load parsed_args["load"] policy
    Flux.loadparams!(model, Flux.params(policy.qnetwork))
end

solver = DeepQLearningSolver(qnetwork = model,
                      max_steps = parsed_args["training_steps"],
                      eps_fraction = parsed_args["eps_fraction"],
                      eps_end = parsed_args["eps_end"],
                      eval_freq = 10_000,
                      save_freq = 10_000,
                      target_update_freq = parsed_args["target_update_freq"],
                      buffer_size = 10_000,
                      batch_size = 32,
                      learning_rate = parsed_args["learning_rate"],
                      train_start = 1000,
                      buffer_size = 400_000,
                      log_freq = 1000,
                      num_ep_eval = 1000,
                      double_q = true,
                      dueling = false,
                      prioritized_replay = true,
                      verbose = true,
                      recurrence = parsed_args["recurrent"],
                      rng = rng,
                      logdir = parsed_args["logdir"])

@time policy = solve(solver, mdp)

env = MDPEnvironment(mdp)
DeepQLearning.evaluation(solver.evaluation_policy,
                policy, env,
                solver.num_ep_eval,
                solver.max_episode_length,
                solver.verbose)

function quick_evaluation(mdp, policy::Policy, rng::AbstractRNG, n_eval=1000)
    avg_r, avg_dr, c_rate, avg_steps, t_out = 0.0, 0.0, 0.0, 0.0, 0.0
    @showprogress for i=1:n_eval
        s0 = initialstate(mdp, rng)
        hr = HistoryRecorder(rng = rng, max_steps=100)
        hist = simulate(hr, mdp, policy, s0)
        avg_r += undiscounted_reward(hist)
        avg_dr += discounted_reward(hist)
        c_rate += undiscounted_reward(hist) <= -1.0
        t_out += n_steps(hist) >= hr.max_steps ? 1.0 : 0.0
        avg_steps += n_steps(hist)
    end
    avg_r /= n_eval
    avg_dr /= n_eval
    c_rate /= n_eval
    avg_steps /= n_eval
    t_out /= n_eval
    return avg_r, avg_dr, c_rate, avg_steps, t_out
end

avg_r, avg_dr, c_rate, avg_steps, t_out = quick_evaluation(mdp, policy, rng, 1000)

println("Collisions ", c_rate*100)
println("Avg steps ", avg_steps)
println("Time outs ", t_out*100)
println("avg disc reward ", avg_dr)
println("avg reward ", avg_r)

BSON.@save joinpath(parsed_args["logdir"], "policy_$(parsed_args["cooperation"]).bson") policy

BSON.@save "results.bson" c_rate avg_steps t_out avg_dr avg_r


precompile = @timed using PowerSimulationsDynamics

using PowerSimulationsDynamics
const PSID = PowerSimulationsDynamics
using Sundials
using PowerSystems
const PSY = PowerSystems
using OrdinaryDiffEq
using Logging

open("precompile_time.txt", "a") do io
    write(io, "| $(ARGS[1]) | $(precompile.time) |\n")
end

sys = System("test/data_tests/WECC_240_dynamic.json"; runchecks = false)
for l in get_components(PSY.PowerLoad, sys)
    PSY.set_model!(l, PSY.LoadModels.ConstantImpedance)
end

# First runs
try
    for m in [ResidualModel, MassMatrixModel]
        Simulation(
            m,
            sys,
            pwd(),
            (0.0, 20.0), #time span
            BranchTrip(1.0, Line, "CORONADO    -1101-PALOVRDE    -1401-i_10");
            console_level = Logging.Error,
        )
    end
catch e
    @error exception = (e, catch_backtrace())
end

try
    sim_ida, time_build_ida, _, _ = @timed Simulation(
        ResidualModel,
        sys,
        pwd(),
        (0.0, 20.0), #time span
        BranchTrip(1.0, Line, "CORONADO    -1101-PALOVRDE    -1401-i_10");
        console_level = Logging.Error,
    )
    status = execute!(sim_ida, IDA(), dtmax = 0.01, enable_progress_bar = false)
    if status == PSID.SIMULATION_FINALIZED
        res_ida = read_results(sim_ida)
        solve_time = res_ida.time_log[:timed_solve_time]
        open("execute_time.txt", "a") do io
            write(io, "| $(ARGS[1])-Build ResidualModel | $(time_build_ida) |\n")
            write(io, "| $(ARGS[1])-Execute ResidualModel | $(solve_time) |\n")
        end
    else
        error("FAILED TO SOLVE")
    end
catch e
    @error exception = (e, catch_backtrace())
    open("execute_time.txt", "a") do io
        write(io, "| $(ARGS[1])- ResidualModel | FAILED TO TEST |\n")
    end
end

try
    sim_rodas, time_build_rodas, _, _ = @timed Simulation(
        MassMatrixModel,
        sys, #system
        pwd(),
        (0.0, 20.0), #time span
        BranchTrip(1.0, Line, "CORONADO    -1101-PALOVRDE    -1401-i_10"),
        #console_level = Logging.Error,
    ) #Type of Fault
    status = execute!(sim_rodas, Rodas4(), rtol = 1e-9, atol = 1e-9)
    if status == PSID.SIMULATION_FINALIZED
        res_rodas = read_results(sim_rodas)
        solve_time = res_rodas.time_log[:timed_solve_time]
        open("execute_time.txt", "a") do io
            write(io, "| $(ARGS[1])-Build MassMatrixModel | $(time_build_rodas) |\n")
            write(io, "| $(ARGS[1])-Execute MassMatrixModel | $(solve_time) |\n")
        end
    else
        error("FAILED TO SOLVE")
    end
catch e
    @error exception = (e, catch_backtrace())
    open("execute_time.txt", "a") do io
        write(io, "| $(ARGS[1])-MassMatrixModel | FAILED TO TEST |\n")
    end
end


# module Battlesnake

# export Boardstate, Snake, Position, Direction, up, down, left, right, move_from

# ============== DIRECTION ==================#
@enum Direction up down left right

# ============== POSITION ==================#
struct Position
    x::Int
    y::Int
end

function get_neighbors(pos::Position)
    return [(dir, move_from(pos, dir)) for dir in [up, down, left, right]]
end

function move_from(pos::Position, dir::Direction)
    if dir == up
        return Position(pos.x, pos.y + 1)
    elseif dir == down
        return Position(pos.x, pos.y - 1)
    elseif dir == right
        return Position(pos.x + 1, pos.y)
    elseif dir == left
        return Position(pos.x - 1, pos.y)
    end
end

# ============== FOOD ==================#
mutable struct Food
    positions::Vector{Position}
end

# ============== SNAKE ==================#

mutable struct Snake
    body::Vector{Position}
    ate_food::Bool
    health::Int
end

function get_head(snake::Snake)
    return snake.body[1]
end

function get_body_no_head(snake::Snake)
    return snake.body[2:length(snake.body)]
end

function possible_actions(snake::Snake)
    head = snake.body[1]
    neck = snake.body[2]
    # coming from right
    if head.x < neck.x
        return [up, down, left]
    end
    # coming from left
    if head.x > neck.x
        return [up, down, right]
    end

    # coming from below
    if head.y > neck.y
        return [up, right, left]
    end

    # coming from above
    if head.y < neck.y
        return [down, right, left]
    end
    return [up, down, left, right]
end

function move_snake!(snake::Snake, food::Food, action::Direction)
    h = get_head(snake)
    new_head = move_from(h, action)
    insert!(snake.body, 1, new_head)
    if !snake.ate_food
        pop!(snake.body)
    end
    snake.ate_food = false
    snake.health -= 1
    food_eaten = -1
    for i in 1:length(food.positions)
        if new_head == food.positions[i]
            snake.ate_food = true
            food_eaten = i
            snake.health = 100
        end
    end
    if food_eaten > -1
        deleteat!(food.positions, food_eaten)
    end
end

function check_bounds(pos::Position, height::Int, width::Int)
    return pos.x >= 1 && pos.x < width + 1 && pos.y >= 1 && pos.y < height + 1
end

function check_dead(snakes::Vector{Snake}, height::Int, width::Int)
    dead = Set()
    for (p, snek) in enumerate(snakes)
        # starve
        if snek.health == 0
            push!(dead, p)
        end

        for (o, opponent) in enumerate(snakes)
            if o == p
                continue
            end

            if get_head(snek) in opponent.body
                # head2head
                if get_head(snek) == get_head(opponent)
                    # RAM
                    if length(snek.body) > length(opponent.body)
                        push!(dead, o)
                    # draw
                    elseif length(snek.body) == length(opponent.body)
                        push!(dead, o)
                        push!(dead, p)
                    else
                        # lost (head2head) collision with opponent
                        push!(dead, p)
                    end
                else
                    # ran into opponent
                    push!(dead, p)
                end
            end

            # collision with self
            if get_head(snek) in get_body_no_head(snek)
                push!(dead, p)
            end
            if !check_bounds(get_head(snek), height, width)
                push!(dead, p)
            end

        end
    end
    return dead
end

# ============== BOARDSTATE ==================#
mutable struct Boardstate
    snakes::Vector{Snake}
    food::Food
    height::Int
    width::Int
end

function is_occupied_by_snake(pos::Position, state::Boardstate)
    for snake in state.snakes
        if pos in snake.body
            return true
        end
    end
    return false
end

function print_board(state::Boardstate)
    board = fill(0, (state.height, state.width))
    for (i,snek) in enumerate(state.snakes)
        for p in snek.body
            board[p.y, p.x] = i
        end
    end
    for f in state.food.positions
        board[f.y, f.x] = 5
    end
    reverse!(board, dims = 1)
    for r in 1:size(board,1)
        println(board[r, :])
    end
end

function next_player(current_player, num_players)
    return (current_player % num_players) + 1
end

# end # module end


mutable struct QueueRecord{E}
    frames::Vector{Frame{E}}
    timestep::Float64
    nframes::Int # number of active Frames
end
function QueueRecord(::Type{E}, capacity::Int, timestep::Float64, frame_capacity::Int=100) where {E}
    frames = Array{Frame{E}}(undef, capacity)
    for i in 1 : length(frames)
        frames[i] = Frame(E, frame_capacity)
    end
    QueueRecord{E}(frames, timestep, 0)
end

Base.show(io::IO, rec::QueueRecord) = print(io, "QueueRecord(nframes=", rec.nframes, ")")

capacity(rec::QueueRecord) = length(rec.frames)
nframes(rec::QueueRecord) = rec.nframes
function nstates(rec::QueueRecord)
    retval = 0
    for frame_index in 1 : nframes(rec)
        retval += length(rec.frames[frame_index])
    end
    return retval
end

function Base.deepcopy(rec::QueueRecord)
    retval = QueueRecord(capacity(rec), rec.timestep, capacity(rec.frames[1]))
    for i in 1 : rec.nframes
        copyto!(retval.frames[i], rec.frames[i])
    end
    retval
end

pastframe_inbounds(rec::QueueRecord, pastframe::Int) = 1 ≤ 1-pastframe ≤ rec.nframes

"""
Indexed by pastframe, so pastframe == 0 is the current scene, -1 is the previous frame, etc.
"""
Base.getindex(rec::QueueRecord, pastframe::Int) = rec.frames[1 - pastframe]

get_time(rec::QueueRecord, pastframe::Int) = -get_elapsed_time(rec, pastframe)
get_timestep(rec::QueueRecord) = rec.timestep
get_elapsed_time(rec::QueueRecord, pastframe::Int) = -pastframe*rec.timestep
function get_elapsed_time(
    rec::QueueRecord,
    pastframe_farthest_back::Int,
    pastframe_most_recent::Int,
    )

    (pastframe_most_recent - pastframe_farthest_back)*rec.timestep
end

function Base.empty!(rec::QueueRecord)
    rec.nframes = 0
    return rec
end

function push_back_records!(rec::QueueRecord)
    for i in min(rec.nframes+1, capacity(rec)) : -1 : 2
        copyto!(rec.frames[i], rec.frames[i-1])
    end
    return rec
end
function Base.insert!(rec::QueueRecord{E}, frame::Frame{E}, pastframe::Int=0) where {E}
    copyto!(rec[pastframe], frame)
    return rec
end
function Base.get!(frame::Frame{E}, rec::QueueRecord{E}, pastframe::Int=0) where {E}
    copyto!(frame, rec[pastframe])
    frame
end
function update!(rec::QueueRecord{E}, frame::Frame{E}) where {E}
    push_back_records!(rec)
    insert!(rec, frame, 0)
    rec.nframes = min(rec.nframes+1, capacity(rec))
    return rec
end

function allocate_frame(rec::QueueRecord{E}) where {E}
    max_n_objects = maximum(length(rec[j]) for j in 0 : 1-length(rec))
    return Frame(E, max_n_objects)
end

const EntityQueueRecord{S,D,I} = QueueRecord{Entity{S,D,I}}

Base.length(record::QueueRecord) = nframes(record)
function Base.iterate(record::QueueRecord, state::Int64=(-nframes(record)+1))
    return state<=0 ? (record[state], state+1) : nothing
end


##
# sparse sampling solver environments

abstract type AbstractVSSOptions end

# accessors
maxdepth(opt::AbstractVSSOptions) = opt.maxdepth
width(opt::AbstractVSSOptions) = opt.width
action_width(opt::AbstractVSSOptions) = opt.action_width

@with_kw struct VOSSOptions{V<:VOOActionGenerator} <: AbstractVSSOptions
    maxdepth::Int
    width::Int
    action_width::Int
    action_width_decay::Float64
    last_action_null::Bool
    voo::V
end

@with_kw struct VOWSSOptions{V<:VOOActionGenerator} <: AbstractVSSOptions
    maxdepth::Int
    width::Int
    action_width::Int
    action_width_decay::Float64
    last_action_null::Bool
    voo::V
end

struct VOOSparseSamplingSolver
    opt::AbstractVSSOptions
    rng::AbstractRNG
end

VOOSparseSamplingSolver(opt::AbstractVSSOptions; rng=Random.GLOBAL_RNG) = VOOSparseSamplingSolver(opt, rng)
VOOSparseSamplingSolver(;rng=Random.GLOBAL_RNG, kwargs...) = VOOSparseSamplingSolver(VOWSSOptions(kwargs...), rng)

function POMDPs.solve(s::VOOSparseSamplingSolver, m::Union{MDP,POMDP})
    if s.opt isa VOWSSOptions
        @assert !(Nothing <: statetype(m)) "POWSS does not support problems where the state can be nothing (this can be fixed by adding Some in appropriate places)."
    end
    VSSPlanner(m, s.opt, s.rng)
end

struct VSSPlanner{M<:Union{MDP,POMDP},O<:AbstractVSSOptions,R<:AbstractRNG}
    m::M
    opt::O
    rng::R
end


using Random, Plots
using Zygote, ForwardDiff
using OrdinaryDiffEq, DiffEqSensitivity
using LinearAlgebra
using Statistics
using ProgressBars, Printf
using Flux
using Flux.Optimise: update!
using Flux.Losses: mae, mse
using BSON: @save, @load
using DelimitedFiles
using YAML

ENV["GKSwstype"] = "100"

cd(dirname(@__DIR__))
conf = YAML.load_file("./config.yaml")

expr_name = conf["expr_name"]
fig_path = string("./results/", expr_name, "/figs")
ckpt_path = string("./results/", expr_name, "/checkpoint")
config_path = "./results/$expr_name/config.yaml"

is_restart = Bool(conf["is_restart"])
ns = Int64(conf["ns"])
nr = Int64(conf["nr"])
lb = Float64(conf["lb"])
n_epoch = Int64(conf["n_epoch"])
n_plot = Int64(conf["n_plot"])
grad_max = Float64(conf["grad_max"])
maxiters = Int64(conf["maxiters"])

lr_max = Float64(conf["lr_max"])
lr_min = Float64(conf["lr_min"])
lr_decay = Float64(conf["lr_decay"])
lr_decay_step = Int64(conf["lr_decay_step"])
w_decay = Float64(conf["w_decay"])

llb = lb;
global p_cutoff = -1.0

const l_exp = 1:14
n_exp = length(l_exp)

l_train = []
l_val = []
for i = 1:n_exp
    j = l_exp[i]
    if !(j in [2, 6, 9, 12])
        push!(l_train, i)
    else
        push!(l_val, i)
    end
end

opt = Flux.Optimiser(
    ExpDecay(lr_max, lr_decay, length(l_train) * lr_decay_step, lr_min),
    ADAMW(0.005, (0.9, 0.999), w_decay),
);

if !is_restart
    if ispath(fig_path)
        rm(fig_path, recursive = true)
    end
    if ispath(ckpt_path)
        rm(ckpt_path, recursive = true)
    end
end

if ispath("./results") == false
    mkdir("./results")
end

if ispath("./results/$expr_name") == false
    mkdir("./results/$expr_name")
end

if ispath(fig_path) == false
    mkdir(fig_path)
    mkdir(string(fig_path, "/conditions"))
end

if ispath(ckpt_path) == false
    mkdir(ckpt_path)
end

cp("./config.yaml", config_path; force=true)


printstyled("RandSeis\n", color=:light_green)

fs_range = exp10.(range(-6, stop=4, length=50))
fc_range = exp10.(range(-4, stop=2, length=20))

printstyled("  getbandcode\n", color=:light_green)
for fs in fs_range
  for fc in fc_range
    getbandcode(fs, fc=fc)
  end
end

printstyled("  rand_misc\n", color=:light_green)
for i = 1:100
  D = rand_misc(1000)
end

# Similarly for the yp2 codes
printstyled("  iccodes_and_units\n", color=:light_green)
for i = 1:1000
  cha, u = RandSeis.iccodes_and_units('s', true)
  @test isa(cha, String)
  @test isa(u, String)

  cha, u = RandSeis.iccodes_and_units('s', false)
  @test isa(cha, String)
  @test isa(u, String)
end

# check that rand_t produces sane gaps
printstyled("  rand_t\n", color=:light_green)
fs = 100.0
nx = 100

printstyled("    controlled gaps\n", color=:light_green)
t = RandSeis.rand_t(fs, nx, 10, 1)
@test size(t, 1) == 12
t = RandSeis.rand_t(fs, nx, 0, 1)
@test size(t, 1) == 2
t = RandSeis.rand_t(100.0, 1000, 4, 200)
@test size(t, 1) == 6
@test t[:,1] == [1, 200, 400, 600, 800, 1000]

printstyled("    gap < Δ/2 + 1\n", color=:light_green)
for i in 1:1000
  fs = rand(RandSeis.fs_vals)
  nx = round(Int64, rand(1200:7200)*fs)
  t = RandSeis.rand_t(fs, nx, 0, 1)
  δt = div(round(Int64, sμ/fs), 2) + 1
  gaps = t[2:end-1, 2]

  if length(gaps) > 0
    @test minimum(gaps) ≥ δt
  end
  @test minimum(diff(t[:,1])) > 0
end

printstyled("  rand_resp\n", color=:light_green)
R = RandSeis.rand_resp(1.0, 8)
@test length(R.z) == length(R.p) == 8

printstyled("  namestrip\n", color=:light_green)
str = String(0x00:0xff)
S = randSeisData(3)
S.name[2] = str

for key in keys(bad_chars)
  test_str = namestrip(str, key)
  @test length(test_str) == 256 - (32 + length(bad_chars[key]))
end
redirect_stdout(out) do
  test_str = namestrip(str, "Nonexistent List")
end
namestrip!(S)
@test length(S.name[2]) == 210

printstyled("  repop_id!\n", color=:light_green)
S = randSeisData()
S.id[end] = deepcopy(S.id[1])
id = deepcopy(S.id)
RandSeis.repop_id!(S)
@test id != S.id

printstyled("  randSeis*\n", color=:light_green)
for i = 1:10
  randSeisChannel()
  randSeisData()
  randSeisHdr()
  randSeisSrc()
  randSeisEvent()
end

printstyled("    keywords\n", color=:light_green)
printstyled("      a0\n", color=:light_green)
try
  randSeisData(a0=true)
catch err
  println("a0 = true threw error ", err)
end

printstyled("      c\n", color=:light_green)
for i = 1:10
  S = randSeisData(10, c=1.0)
  @test maximum(S.fs) == 0.0
end

printstyled("      fc\n", color=:light_green)
C = randSeisChannel(fc = 2.0)
resp_a0!(C.resp)
r1 = fctoresp(2.0f0, 1.0f0)
r2 = fctoresp(2.0f0, Float32(1.0/sqrt(2)))
resp_a0!(r1)
resp_a0!(r2)
r = C.resp
for f in (:a0, :f0, :z, :p)
  @test isapprox(getfield(r, f), getfield(r1, f), rtol=eps(Float32)) || isapprox(getfield(r, f), getfield(r2, f), rtol=eps(Float32))
end

printstyled("      s\n", color=:light_green)
for i = 1:10
  S = randSeisData(10, s=1.0)
  @test maximum(S.fs) > 0.0
end


# Create a model using Agents.jl

function initialize_model(param_file::AbstractString)
	# Load params
	d = YAML.load_file(param_file)

	r = Researcher(1, d["risk_taking_average"], d["generality_average"])
	d["time"] = 1

	model = ABM(typeof(r), properties=d, scheduler=Schedulers.randomly)

	researchers = start_population(d["population_size"], d["risk_taking_average"], d["generality_average"])
	for res in researchers
		add_agent!(res, model)
	end

	start_experienced_population(model, true)
	model.properties["max_grants"] = d["max_grants"] # restore this param
	model.properties["max_grants_org"] = deepcopy(d["max_grants"]) # restore this param

	return model
end

function initialize_model(;population_size, max_grants, years, risk_taking_average, generality_average, min_prof_exp, average_students, max_exp, mutation_rate, K, max_years, risk_range, A, B, C, save_gens, replicates)

	d = Dict(
		"risk_taking_average" => risk_taking_average,
		"generality_average" => generality_average,
		"population_size" => population_size,
		"max_grants" => max_grants,
		"years" => years,
		"min_prof_exp" => min_prof_exp,
		"average_students" => average_students,
		"max_exp" => max_exp,
		"mutation_rate" => mutation_rate,
		"K" => K,
		"max_years" => max_years,
		"risk_range" => risk_range,
		"A" => A,
		"B" => B,
		"C" => C,
		"save_gens" => save_gens,
		"replicates" => replicates
	)

	r = Researcher(1, d["risk_taking_average"], d["generality_average"])
	d["time"] = 1

	model = ABM(typeof(r), properties=d, scheduler=Schedulers.randomly)

	researchers = start_population(d["population_size"], d["risk_taking_average"], d["generality_average"])
	for res in researchers
		add_agent!(res, model)
	end

	start_experienced_population(model, true)
	model.properties["max_grants"] = d["max_grants"] # restore this param
	model.properties["max_grants_org"] = deepcopy(d["max_grants"]) # restore this param

	return model
end

function model_step!(model::ABM)
	retire!(model)
	exclude_unproductive!(model)
	get_grant!(model)
	# gain experience, do one year of research, publish finished work, receive citations
	update_researcher!(model)
	train_students!(model)
	# update time
	model.properties["time"] += 1
	# update the available grants per year
	model.properties["max_grants"] = deepcopy(model.properties["max_grants_org"])
end

"""
Create a homogeneous population of researchers.

parameters:
population_size: number of researchers.
number_of_grants: how many years of grants exists.
max_years: maximum number of years a researcher can receive a grant.
"""
function start_population(population_size, risk_taking::Float64=0.5, generality::Float64=0.5)
	researchers = [Researcher(i, risk_taking, generality) for i in 1:population_size]
	return researchers
end

"""
Create a heterogenous population of experienced researchers.
"""
function start_experienced_population(model, randomage=true)
	experience = model.properties["min_prof_exp"]
	population_size = model.properties["population_size"]

	if !randomage
		for year in 1:experience
			get_grant!(model)
			for resid in allids(model)
				update_researcher!(model[resid], K)  # gain experience, do one year of research, publish finished work, receive citations
			end
		end
	else
		age_groups = Int(round(population_size/experience))
		allind = collect(allids(model))
		ress = sample(allind, age_groups, replace=false, ordered=true)
		for ind in reverse(ress)
			splice!(allind, findfirst(a-> a==ind, allind)[1])
		end

		for year in 1:experience
			get_grant!(model)
			for resid in ress
				update_researcher!(model[resid], model)  # gain experience, do one year of research, publish finished work, receive citations
			end
			# remove the ress indices from allind and create new ress indices
			if age_groups >= length(allind)
				ress2 = allind
				ress = vcat(ress, allind)
				continue
			else
				ress2 = sample(allind, age_groups, replace=false, ordered=true)
				ress = vcat(ress, ress2)
			end
			for ind in reverse(ress2)
				splice!(allind, findfirst(a-> a==ind, allind))
			end
		end
	end
end


#! /usr/bin/julia

# Rosetta Code, Textonyms

const tcode = (Regex=>Char)[r"A|B|C|Ä|Å|Á|Â|Ç" => '2',
                            r"D|E|F|È|Ê|É" => '3',
                            r"G|H|I|Í" => '4',
                            r"J|K|L" => '5',
                            r"M|N|O|Ó|Ö|Ô|Ñ" => '6',
                            r"P|Q|R|S" => '7',
                            r"T|U|V|Û|Ü" => '8',
                            r"W|X|Y|Z" => '9']

function tpad(str::IOStream)
    tnym = (String=>Array{String,1})[]
    for w in eachline(str)
        w = chomp(w)
        t = uppercase(w)
        for (k,v) in tcode
            t = replace(t, k, v)
        end
        t = replace(t, r"\D", '1')
        tnym[t] = [get(tnym, t, String[]), w]
    end
    return tnym
end

dname = "/usr/share/dict/american-english"
DF = open(dname, "r")
tnym = tpad(DF)
close(DF)

println("The character to digit mapping is done according to")
println("these regular expressions (following uppercase conversion):")
for k in sort(collect(keys(tcode)), by=x->tcode[x])
    println("    ", tcode[k], " -> ", k)
end
println("Unmatched non-digit characters are mapped to 1")

println()
print("There are ", sum(map(x->length(x), values(tnym))))
println(" words in ", dname)
println("  which can be represented by the digit key mapping.")
print("They require ", length(keys(tnym)))
println(" digit combinations to represent them.")
print(sum(map(x->length(x)>1, values(tnym))))
println(" digit combinations represent Textonyms.")

println()
println("The degeneracies of telephone key encodings are:")
println("  Words Encoded   Number of codes")
dgen = zeros(maximum(map(x->length(x), values(tnym))))
for v in values(tnym)
    dgen[length(v)] += 1
end
for (i, d) in enumerate(dgen)
    println(@sprintf "%10d  %15d" i d)
end

println()
dgen = length(dgen) - 2
println("Codes mapping to ", dgen, " or more words:")
for (k, v) in tnym
    dgen <= length(v) || continue
    println(@sprintf "%7s (%2d) %s" k length(v) join(v, ", "))
end


@testset "StockTradingEnv" begin

    env = StockTradingEnvWithTurbulence()

    RLBase.test_interfaces!(env)
    RLBase.test_runnable!(env)
end