using LatticeGPU
using TOML
using TimerOutputs
using ArgParse
using CUDA

function load_fields()
    global U = vector_field(SU3{Float64}, lp)
    fill!(U,one(SU3{Float64}));

    global U_CPU = Array(U)
    global psi = scalar_field(Spinor{4,SU3fund{Float64}},lp)
    global psi_CPU = Array(psi)
    global psit_CPU = Array(psi)

    global pp_density = Array{Float64}(undef,lp.bsz,lp.rsz);
    global ap_density = Array{Complex{Float64}}(undef,lp.bsz,lp.rsz);

    global pp_corr_t0 = fill(zero(Float64),(lp.iL[4],params["Frontflow"]["N_noise"]));
    global ap_corr_t0 = fill(zero(ComplexF64),(lp.iL[4],params["Frontflow"]["N_noise"]));

    global pp_corr_t = fill(zero(Float64),(lp.iL[4],params["Frontflow"]["N_noise"],params["Frontflow"]["nsteps"]+1));
    global ap_corr_t = fill(zero(ComplexF64),(lp.iL[4],params["Frontflow"]["N_noise"],params["Frontflow"]["nsteps"]+1));

    global pp_corr_tfl = fill(zero(Float64),(lp.iL[4],params["Backflow"]["N_noise"],length(flow_times)));
    global ap_corr_tfl = fill(zero(ComplexF64),(lp.iL[4],params["Backflow"]["N_noise"],length(flow_times)));

    global Eoft = Array{Complex{Float64}}(undef,2+ffnsteps,lp.iL[4]);
    global Eofpla = Array{Complex{Float64}}(undef,lp.iL[4],lp.npls);

    global phat_t = Array{Float64}(undef,params["Frontflow"]["N_noise"],length(flow_times));
    global Quark_cond = Array{Complex{Float64}}(undef,lp.iL[4],params["Backflow"]["N_noise"],length(flow_times));
    global Quark_cond_cfl = Array{Complex{Float64}}(undef,lp.iL[4],params["Backflow"]["N_noise"],length(flow_times));

    global Quark_cond2 = Array{Complex{Float64}}(undef,lp.iL[4],params["Backflow"]["N_noise"],length(flow_times));
    global Quark_cond2_cfl = Array{Complex{Float64}}(undef,lp.iL[4],params["Backflow"]["N_noise"],length(flow_times));
    global ChiDchi = Array{Complex{Float64}}(undef,lp.iL[4],params["Backflow"]["N_noise"],length(flow_times));
    global ChiDchi_cfl = Array{Complex{Float64}}(undef,lp.iL[4],params["Backflow"]["N_noise"],length(flow_times));

    return nothing
end

"""
function two_pt()

Stores the two point function in the variables 'pp_corr_t0', 'ap_corr_t0', 'pp_corr_t' and 'ap_corr_t'
"""
function two_pt() # Will be Frontflow

    println(log_file,"\nMeasuring 2pt...")
    flush(log_file)

    Eoft_plaq(Eofpla,U, gp, lp, ymws)
    Eoft[1,:] .= sum(Eofpla,dims = 2)

    for noi in 1:params["Frontflow"]["N_noise"]

        niter = propagator!(psi,U, dpar, dws, lp,params["Fermion"]["maxiter"], params["Fermion"]["tolerance"],params["Frontflow"]["tsource"])

        println(log_file,"CG converged after ",niter," iterations with residue ",CUDA.mapreduce(x -> norm2(x), +, dws.sr))
        flush(log_file)

        pp_density .= Array(norm2.(psi))
        ap_density .= Array(dot.(psi,dmul.(Gamma{4},psi)))

        pp_corr_t0[:,noi] .= zero(Float64)
        ap_corr_t0[:,noi] .= zero(ComplexF64)
        for t in 1:lp.iL[4]
            for i in 1:lp.iL[1] for j in 1:lp.iL[2] for k in 1:lp.iL[3]
                b,r = point_index(CartesianIndex{lp.ndim}((i,j,k,t)),lp)
                pp_corr_t0[t,noi] += pp_density[b,r]
                ap_corr_t0[t,noi] += ap_density[b,r]
            end end end
        end

        _,epslist = flw_adapt(U, psi, int, params["Frontflow"]["t_zero"], gp, dpar, lp, ymws, dws)

        for tstep in 1:params["Frontflow"]["nsteps"]

            if noi ==1
                Eoft_plaq(Eofpla,U, gp, lp, ymws)
                Eoft[tstep+1,:] .= sum(Eofpla,dims = 2)
            end

            pp_density .= Array(norm2.(psi))
            ap_density .= Array(dot.(psi,dmul.(Gamma{4},psi)))
            pp_corr_t[:,noi,tstep] .= zero(Float64)
            ap_corr_t[:,noi,tstep] .= zero(ComplexF64)
            for t in 1:lp.iL[4]
                for i in 1:lp.iL[1] for j in 1:lp.iL[2] for k in 1:lp.iL[3]
                    b,r = point_index(CartesianIndex{lp.ndim}((i,j,k,t)),lp)
                    pp_corr_t[t,noi,tstep] += pp_density[b,r]
                    ap_corr_t[t,noi,tstep] += ap_density[b,r]
                end end end
            end

            flw(U, psi, int, 1, params["Frontflow"]["epsilon"], gp, dpar, lp, ymws, dws)
        end

        pp_density .= Array(norm2.(psi))
        ap_density .= Array(dot.(psi,dmul.(Gamma{4},psi)))
        pp_corr_t[:,noi,end] .= zero(Float64)
        ap_corr_t[:,noi,end] .= zero(ComplexF64)
        for t in 1:lp.iL[4]
            for i in 1:lp.iL[1] for j in 1:lp.iL[2] for k in 1:lp.iL[3]
                b,r = point_index(CartesianIndex{lp.ndim}((i,j,k,t)),lp)
                pp_corr_t[t,noi,end] += pp_density[b,r]
                ap_corr_t[t,noi,end] += ap_density[b,r]
            end end end
        end

        @timeit "CPU to GPU" copyto!(U,U_CPU)
    end

    pp_corr_t0 .= pp_corr_t0./prod(lp.iL[1:3])
    ap_corr_t0 .= ap_corr_t0./prod(lp.iL[1:3])
    pp_corr_t .= pp_corr_t./prod(lp.iL[1:3])
    ap_corr_t .= ap_corr_t./prod(lp.iL[1:3])

    println(log_file,"Finished measuring 2pt")
    flush(log_file)

    return nothing
end

"""
function one_pt()

Stores the one point function in the variable 'Quark_cond' and 'Quark_cond_cfl'. The final result for the quark condensate is 'Quark_cond' + C_fl*'Quark_cond_cfl'.
The two point function with source in possitive flow time is also stored in 'pp_corr_tfl' and 'ap_corr_tfl'
"""
function one_pt() # Will be backflow

    @timeit "CPU to GPU" copyto!(U,U_CPU)

    for fl in 1:length(flow_times)
        println(log_file,"\nMeasuring 1pt at t = "*string(flow_times[fl])*"...")
        flush(log_file)

        for noi in 1:params["Backflow"]["N_noise"]

            ## The change in the field before randomize should be moved to the package
            psi .= 0.0*psi
            pfrandomize!(psi,lp,params["Backflow"]["tsource"])
            @timeit "GPU to CPU" psit_CPU .= Array(psi)

            backflow(psi, U, flow_times[fl], , gp, dpar, lp, ymws, dws)

            @timeit "GPU to CPU" psi_CPU .= Array(psi)

            @timeit "Inversion" begin
                dws.sp .= dmul.(Gamma{5},psi)
                g5Dw!(psi,U,dws.sp,dpar,dws,lp)

                niter = CG!(psi,U,DwdagDw!,dpar,lp,dws,params["Fermion"]["maxiter"], params["Fermion"]["tolerance"])
            end

            println(log_file,"CG converged after ",niter," iterations with residue ",CUDA.mapreduce(x -> norm2(x), +, dws.sr))
            flush(log_file)

            @timeit "CPU to GPU" copyto!(dws.st,psi_CPU)

            Quark_cond[:,noi,fl] .= zero(ComplexF64)
            ap_density .= Array(dot(dws.st[b,r],psi[b,r]))
            for t in 1:lp.iL[4]
                for i in 1:lp.iL[1] for j in 1:lp.iL[2] for k in 1:lp.iL[3]
                    b,r = point_index(CartesianIndex{lp.ndim}((i,j,k,t)),lp)
                    CUDA.@allowscalar Quark_cond[t,noi,fl] += ap_density[b,r]
                end end end
            end

            Quark_cond_cfl[:,noi,fl] .= zero(ComplexF64)
            pp_density .= Array(norm2.(dws.st))
            for t in 1:lp.iL[4]
                for i in 1:lp.iL[1] for j in 1:lp.iL[2] for k in 1:lp.iL[3]
                    b,r = point_index(CartesianIndex{lp.ndim}((i,j,k,t)),lp)
                    CUDA.@allowscalar Quark_cond_cfl[t,noi,fl] += -pp_density[b,r]
                end end end
            end

            pp_density .= Array(norm2.(psi))
            ap_density .= Array(dot.(psi,dmul.(Gamma{4},psi)))
            pp_corr_tfl[:,noi,fl] .= zero(Float64)
            ap_corr_tfl[:,noi,fl] .= zero(ComplexF64)
            for t in 1:lp.iL[4]
                for i in 1:lp.iL[1] for j in 1:lp.iL[2] for k in 1:lp.iL[3]
                    b,r = point_index(CartesianIndex{lp.ndim}((i,j,k,t)),lp)
                    pp_corr_tfl[t,noi,fl] += pp_density[b,r]
                    ap_corr_tfl[t,noi,fl] += ap_density[b,r]
                end end end
            end

            @timeit "CPU to GPU" copyto!(U,U_CPU)

            @timeit "ChiDchi" begin

                flw_adapt(U, psi, int, flow_times[fl], gp, dpar, lp, ymws, dws)
                Dw!(dws.sp,U,psi,dpar,dws,lp)

                @timeit "CPU to GPU" copyto!(dws.st,psit_CPU)

                ChiDchi[:,noi,fl] .= zero(Float64)
                ap_density .= Array(dot(dws.st[b,r],dws.sp[b,r]))
                for t in 1:lp.iL[4]
                    for i in 1:lp.iL[1] for j in 1:lp.iL[2] for k in 1:lp.iL[3]
                        b,r = point_index(CartesianIndex{lp.ndim}((i,j,k,t)),lp)
                        CUDA.@allowscalar ChiDchi[t,noi,fl] += ap_density[b,r]
                    end end end
                end
                Dw!(dws.sp,U,dws.st,dpar,dws,lp)
                ap_density .= Array(dot(dws.sp[b,r],psi[b,r]))
                for t in 1:lp.iL[4]
                    for i in 1:lp.iL[1] for j in 1:lp.iL[2] for k in 1:lp.iL[3]
                        b,r = point_index(CartesianIndex{lp.ndim}((i,j,k,t)),lp)
                        CUDA.@allowscalar ChiDchi[t,noi,fl] += -ap_density[b,r]
                    end end end
                end

                Quark_cond2[:,noi,fl] .= zero(ComplexF64)
                ap_density .= Array(dot(dws.st[b,r],psi[b,r]))
                for t in 1:lp.iL[4]
                    for i in 1:lp.iL[1] for j in 1:lp.iL[2] for k in 1:lp.iL[3]
                        b,r = point_index(CartesianIndex{lp.ndim}((i,j,k,t)),lp)
                        CUDA.@allowscalar Quark_cond2[t,noi,fl] += ap_density[b,r]
                    end end end
                end

                @timeit "CPU to GPU" copyto!(U,U_CPU)
                @timeit "CPU to GPU" copyto!(psi,psi_CPU)

                flw_adapt(U, psi, int, flow_times[fl], gp, dpar, lp, ymws, dws)
                Dw!(dws.sp,U,psi,dpar,dws,lp)

                ChiDchi_cfl[:,noi,fl] .= zero(ComplexF64)
                ap_density .= Array(dot(dws.st[b,r],dws.sp[b,r]))
                for t in 1:lp.iL[4]
                    for i in 1:lp.iL[1] for j in 1:lp.iL[2] for k in 1:lp.iL[3]
                        b,r = point_index(CartesianIndex{lp.ndim}((i,j,k,t)),lp)
                        CUDA.@allowscalar ChiDchi_cfl[t,noi,fl] += ap_density[b,r]
                    end end end
                end
                Dw!(dws.sp,U,dws.st,dpar,dws,lp)
                ap_density .= Array(dot(dws.sp[b,r],psi[b,r]))
                for t in 1:lp.iL[4]
                    for i in 1:lp.iL[1] for j in 1:lp.iL[2] for k in 1:lp.iL[3]
                        b,r = point_index(CartesianIndex{lp.ndim}((i,j,k,t)),lp)
                        CUDA.@allowscalar ChiDchi_cfl[t,noi,fl] += -ap_density[b,r]
                    end end end
                end

                Quark_cond2_cfl[:,noi,fl] .= zero(ComplexF64)
                ap_density .= Array(dot(dws.st[b,r],psi[b,r]))
                for t in 1:lp.iL[4]
                    for i in 1:lp.iL[1] for j in 1:lp.iL[2] for k in 1:lp.iL[3]
                        b,r = point_index(CartesianIndex{lp.ndim}((i,j,k,t)),lp)
                        CUDA.@allowscalar Quark_cond2_cfl[t,noi,fl] += ap_density[b,r]
                    end end end
                end

            end


            @timeit "CPU to GPU" copyto!(U,U_CPU)

        end
    end

    Quark_cond .= Quark_cond./prod(lp.iL[1:3])
    Quark_cond_cfl .= Quark_cond_cfl./prod(lp.iL[1:3])

    pp_corr_tfl .= pp_corr_tfl./prod(lp.iL[1:3])
    ap_corr_tfl .= ap_corr_tfl./prod(lp.iL[1:3])

    ChiDchi .= ChiDchi./prod(lp.iL[1:3])
    Quark_cond2 .= Quark_cond2./prod(lp.iL[1:3])

    ChiDchi_cfl .= ChiDchi_cfl./prod(lp.iL[1:3])
    Quark_cond2_cfl .= Quark_cond2_cfl./prod(lp.iL[1:3])

    println(log_file,"Finished measuring 1pt")

    return nothing
end

"""
function two_pt_rnd()

Stores the random field evolved in flow time in the variable 'phat_t'
"""
function two_pt_rnd()

    copyto!(U,U_CPU)

    println(log_file,"Measuring PhatPt correlator...")
    flush(log_file)
    for noi in 1:params["Frontflow"]["N_noise"]

        pfrandomize!(psi,lp)

        delta_t = flow_times .- [0.0,flow_times[1:end-1]...]
        neps = int.eps_ini
        for fl in 1:length(flow_times)
            _,epslist = flw_adapt(U, psi, int, delta_t[fl], neps, gp, dpar, lp, ymws, dws)
            neps = epslist[end]

            phat_t[noi,fl] = sum(norm2.(psi))
        end

        copyto!(U,U_CPU)
    end

    phat_t .= phat_t./V4

    println(log_file,"Finished measuring PhatPt")
    flush(log_file)

    return nothing
end
