####################################################################
# Predictive online PDPS for optical flow with known velocity field
####################################################################

__precompile__()

module AlgorithmFBDual

using Printf

using AlgTools.Util
import AlgTools.Iterate
using ImageTools.Gradient

using ..OpticalFlow: Image,
                     ImageSize,
                     flow!

#########################
# Iterate initialisation
#########################

function init_rest(x::Image)
    imdim=size(x)

    y = zeros(2, imdim...)
    Δx = copy(x)
    Δy = copy(y)

    return x, y, Δx, Δy
end

function init_iterates(xinit::Image)
    return init_rest(copy(xinit))
end

function init_iterates(dim::ImageSize)
    return init_rest(zeros(dim...))
end

############
# Algorithm
############

function solve( :: Type{DisplacementT};
               dim :: ImageSize,
               iterate = AlgTools.simple_iterate,
               params::NamedTuple) where DisplacementT

    ################################                                        
    # Extract and set up parameters
    ################################                    

    α, ρ = params.α, params.ρ
    τ₀, τ̃₀ =  params.τ₀, params.τ̃₀

    R_K² = ∇₂_norm₂₂_est²    
    τ = τ₀/R_K²

    ######################
    # Initialise iterates
    ######################

    x, y, Δx, Δy = init_iterates(dim)

    ####################
    # Run the algorithm
    ####################

    v = iterate(params) do verbose :: Function,
                           b :: Image,
                           v_known :: DisplacementT,
                           🚫unused_b_next :: Image

        ##################
        # Prediction step
        ##################

        # Δx is a temporary storage variable of correct dimensions
        flow!(@view(y[1,:, :]), Δx, v_known)
        flow!(@view(y[2,:, :]), Δx, v_known)
        
        ∇₂ᵀ!(Δx, y)
        @. x = b - Δx
        ∇₂!(Δy, x)
        @. y = (y - τ*Δy)/(1 + τ*ρ/α)
        proj_norm₂₁ball!(y, α)

        ################################
        # Give function value if needed
        ################################
        v = verbose() do            
            ∇₂ᵀ!(Δx, y)
            @. x = b - Δx
            ∇₂!(Δy, x)
            value = norm₂²(b-x)/2 + α*γnorm₂₁(Δy, ρ)
            value, x, [NaN, NaN], nothing
        end

        v
    end
    
    @warn "Using old x value. Better approach unimplemented as this algorithm is not good."
    # ∇₂ᵀ!(Δx, y)
    # @. x = b - Δx

    return x, y, v
end

end # Module


