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 pphat_t = Array{Float64}(undef,lp.iL[4],params["Frontflow"]["N_noise"],params["Frontflow"]["nsteps"] + 1);
    global pptilde_t = Array{Complex{Float64}}(undef,lp.iL[4],params["Frontflow"]["N_noise"],params["Frontflow"]["nsteps"] + 1);

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

    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()

Computes the correlation functions involving Frontflow of fermionic fields.

Stores the two point function in the variables 'pp_corr_t0', 'ap_corr_t0', 'pp_corr_t' and 'ap_corr_t'
"""
function Frontflow_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)
        @timeit "Volume sum" for b in 1:lp.bsz for r in 1:lp.rsz
            t = point_time((b,r),lp)
            pp_corr_t0[t,noi] += pp_density[b,r]
            ap_corr_t0[t,noi] += ap_density[b,r]
        end end

        if params["Frontflow"]["t_zero"]>0.0
            _,epslist = flw_adapt(U, psi, int, params["Frontflow"]["t_zero"], gp, dpar, lp, ymws, dws)
            println(log_file,"Flowed correlator to t = ",params["Frontflow"]["t_zero"]," with last stepsize ",epslist[end-1])
            flush(log_file)
        end

        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_corr_t[:,noi,tstep] .= zero(Float64)
            ap_corr_t[:,noi,tstep] .= zero(ComplexF64)
            @timeit "Volume sum performance test" begin
                pp_corr_t[:,noi,tstep] .= scalar_contraction(psi)
                ap_corr_t[:,noi,tstep] .= gammazero_contraction(psi)
            end

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

        if noi ==1
            Eoft_plaq(Eofpla,U, gp, lp, ymws)
            Eoft[end,:] .= sum(Eofpla,dims = 2)
        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)
        @timeit "Volume sum" for b in 1:lp.bsz for r in 1:lp.rsz
            t = point_time((b,r),lp)
            pp_corr_t[t,noi,end] += pp_density[b,r]
            ap_corr_t[t,noi,end] += ap_density[b,r]
        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 Backflow_pt()

Computes the correlation functions involving Backflow of fermionic fields.

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'. The correlation function 'ChiDchi' is also updated, as well as a second computation of the condensate.
"""
function Backflow_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], params["Backflow"]["Nsaves"] , 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,psi))
            Quark_cond_cfl[:,noi,fl] .= zero(ComplexF64)
            pp_density .= Array(norm2.(dws.st))
            @timeit "Volume sum" for b in 1:lp.bsz for r in 1:lp.rsz
                t = point_time((b,r),lp)
                Quark_cond[t,noi,fl] += ap_density[b,r]
                Quark_cond_cfl[t,noi,fl] += -pp_density[b,r]
            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)
            @timeit "Volume sum" for b in 1:lp.bsz for r in 1:lp.rsz
                t = point_time((b,r),lp)
                pp_corr_tfl[t,noi,fl] += pp_density[b,r]
                ap_corr_tfl[t,noi,fl] += ap_density[b,r]
            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,dws.sp))
                @timeit "Volume sum" for b in 1:lp.bsz for r in 1:lp.rsz
                    t = point_time((b,r),lp)
                    ChiDchi[t,noi,fl] += ap_density[b,r]
                end end
                Dw!(dws.sp,U,dws.st,dpar,dws,lp)
                ap_density .= Array(dot.(dws.sp,psi))
                @timeit "Volume sum" for b in 1:lp.bsz for r in 1:lp.rsz
                    t = point_time((b,r),lp)
                    ChiDchi[t,noi,fl] += -ap_density[b,r]
                end end

                Quark_cond2[:,noi,fl] .= zero(ComplexF64)
                ap_density .= Array(dot.(dws.st,psi))
                @timeit "Volume sum" for b in 1:lp.bsz for r in 1:lp.rsz
                    t = point_time((b,r),lp)
                    Quark_cond2[t,noi,fl] += ap_density[b,r]
                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,dws.sp))
                @timeit "Volume sum" for b in 1:lp.bsz for r in 1:lp.rsz
                    t = point_time((b,r),lp)
                    ChiDchi_cfl[t,noi,fl] += ap_density[b,r]
                end end
                Dw!(dws.sp,U,dws.st,dpar,dws,lp)
                ap_density .= Array(dot.(dws.sp,psi))
                @timeit "Volume sum" for b in 1:lp.bsz for r in 1:lp.rsz
                    t = point_time((b,r),lp)
                    ChiDchi_cfl[t,noi,fl] += -ap_density[b,r]
                end end

                Quark_cond2_cfl[:,noi,fl] .= zero(ComplexF64)
                ap_density .= Array(dot.(dws.st,psi))
                @timeit "Volume sum" for b in 1:lp.bsz for r in 1:lp.rsz
                    t = point_time((b,r),lp)
                    Quark_cond2_cfl[t,noi,fl] += ap_density[b,r]
                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_lagrange()

Computes the two point functions containing Lagrange multiplier fields using front flow. The convention is such that the quark condentate is 'pptilde_t' + c_fl*'pphat_t'

Stores the two point functions in the variables 'pphat_t' and 'pptilde_t'.
"""
function Two_pt_lagrange() # Will be 2pt lagrange mult

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

    println(log_file,"Measuring Lagrange multiplier correlators...")
    flush(log_file)
    for noi in 1:params["Frontflow"]["N_noise"]

        pfrandomize!(psi,lp)

        @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)

        _,epslist = flw_adapt(U, psi, int, params["Frontflow"]["t_zero"], gp, dpar, lp, ymws, dws)
        @timeit "GPU to CPU" psit_CPU .= Array(psi)
        @timeit "CPU to GPU" copyto!(U,U_CPU)
        @timeit "CPU to GPU" copyto!(psi,psi_CPU)
        if params["Frontflow"]["t_zero"]>0.0
            _,epslist = flw_adapt(U, psi, int, params["Frontflow"]["t_zero"], gp, dpar, lp, ymws, dws)
            println(log_file,"Flowed correlators to t = ",params["Frontflow"]["t_zero"]," with last stepsize ",epslist[end-1])
            flush(log_file)
        end

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

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

            pphat_t[:,noi,fl] .= zero(Float64)
            pp_density .= Array(norm2.(psi))
            pptilde_t[:,noi,fl] .= zero(Float64)
            ap_density .= Array(dot.(dws.st,psi))
            @timeit "Volume sum" for b in 1:lp.bsz for r in 1:lp.rsz
                t = point_time((b,r),lp)
                pphat_t[t,noi,fl] += -pp_density[b,r]
                pptilde_t[t,noi,fl] += ap_density[b,r]
            end end
            ymws.U1 .= U
            flw(U, psi, int, 1, params["Frontflow"]["epsilon"], gp, dpar, lp, ymws, dws)
            U .= ymws.U1
            flw(U, dws.st, int, 1, params["Frontflow"]["epsilon"], gp, dpar, lp, ymws, dws)

        end
        pphat_t[:,noi,end] .= zero(Float64)
        pp_density .= Array(norm2.(psi))
        pptilde_t[:,noi,end] .= zero(Complex{Float64})
        ap_density .= Array(dot.(dws.st,psi))
        @timeit "Volume sum" for b in 1:lp.bsz for r in 1:lp.rsz
            t = point_time((b,r),lp)
            pphat_t[t,noi,end] += -pp_density[b,r]
            pptilde_t[t,noi,end] += ap_density[b,r]
        end end

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

    end

    pphat_t .= pphat_t./prod(lp.iL[1:3])
    pptilde_t .= pptilde_t./prod(lp.iL[1:3])


    println(log_file,"Finished measuring Lagrange multiplier correlators")
    flush(log_file)

    return nothing
end



function krnl_scalar_contraction!(rm, psi, lp)
    @inbounds begin
        b = Int64(CUDA.threadIdx().x)
        r = Int64(CUDA.blockIdx().x)
        I = point_coord((b,r), lp)
        rm[I] = norm2(psi[b,r])
    end
    return nothing
end
function scalar_contraction(psi)
    CUDA.@sync begin
        CUDA.@cuda threads=lp.bsz blocks=lp.rsz krnl_scalar_contraction!(ymws.rm, psi, lp)
    end
    return reshape(Array(CUDA.reduce(+, ymws.rm; dims=(1,2,3))),lp.iL[end])
end

function krnl_gammazero_contraction!(cm, psi, lp)
    @inbounds begin
        b = Int64(CUDA.threadIdx().x)
        r = Int64(CUDA.blockIdx().x)
        I = point_coord((b,r), lp)
        cm[I] = dot(psi[b,r],dmul(Gamma{4},psi[b,r]))
    end
    return nothing
end
function gammazero_contraction(psi)
    CUDA.@sync begin
        CUDA.@cuda threads=lp.bsz blocks=lp.rsz krnl_gammazero_contraction!(ymws.cm, psi, lp)
    end
    return reshape(Array(CUDA.reduce(+, ymws.cm; dims=(1,2,3))),lp.iL[end])
end
