text:
# *****************************************************************************
# Written by Ritchie Lee, ritchie.lee@sv.cmu.edu
# *****************************************************************************
# Copyright ã 2015, United States Government, as represented by the
# Administrator of the National Aeronautics and Space Administration. All
# rights reserved.  The Reinforcement Learning Encounter Simulator (RLES)
# platform is licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License. You
# may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0. Unless required by applicable
# law or agreed to in writing, software distributed under the License is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the specific language
# governing permissions and limitations under the License.
# _____________________________________________________________________________
# Reinforcement Learning Encounter Simulator (RLES) includes the following
# third party software. The SISLES.jl package is licensed under the MIT Expat
# License: Copyright (c) 2014: Youngjun Kim.
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED
# "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
# NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
# PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
# *****************************************************************************

"""
(Tree-based) Genetic programming for the ACASX problem.
Example usage: config=configure(ACASX_GP,"normal","nvn_dasc"); acasx_gp(;config...)
"""
module ACASX_GP

export configure, acasx_gp

import Compat.ASCIIString
using ExprSearch.GP
using Datasets
using RLESUtils, FileUtils, Configure, LogSystems, Loggers
import RLESTypes.SymbolTable

using GrammarExpts
using ACASXProblem, DerivTreeVis
import Configure.configure

const CONFIGDIR = joinpath(dirname(@__FILE__), "..", "config")
const RESULTDIR = joinpath(dirname(@__FILE__), "..", "..", "..", "results")

configure(::Type{Val{:ACASX_GP}}, configs::AbstractString...) = configure_path(CONFIGDIR, configs...)

#nmacs vs nonnmacs
"""
Example call:
config=configure(ACASX_GP, "nvn_dasc", "normal")
acasx_ge(; config...)
"""
function acasx_gp(;outdir::AbstractString=joinpath(RESULTDIR, "./ACASX_GP"),
                  seed=1,
                  logfileroot::AbstractString="acasx_gp_log",

                  runtype::Symbol=:nmacs_vs_nonnmacs,
                  data::AbstractString="dasc",
                  manuals::AbstractString="",
                  clusterdataname::AbstractString="",

                  pop_size::Int64=100,
                  maxdepth::Int64=10,
                  iterations::Int64=10,
                  tournament_size::Int64=20,
                  top_keep::Float64=0.1,
                  crossover_frac::Float64=0.3,
                  mutate_frac::Float64=0.3,
                  rand_frac::Float64=0.2,
                  default_code::Any=:(eval(false)),

                  vis::Bool=true)
    srand(seed)
    mkpath(outdir)

    problem = ACASXClustering(runtype, data, manuals, clusterdataname)

    logsys = GP.logsystem()
    empty_listeners!(logsys)
    send_to!(STDOUT, logsys, ["verbose1", "current_best_print", "result"])
    logs = TaggedDFLogger()
    send_to!(logs, logsys, ["code", "computeinfo", "current_best", "elapsed_cpu_s", "fitness",
        "fitness5", "parameters", "result"])

    gp_params = GPESParams(pop_size, maxdepth, iterations, tournament_size, top_keep,
        crossover_frac, mutate_frac, rand_frac, default_code, logsys;
        userargs=SymbolTable(:ids=>collect(1:length(problem.Dl))))

    result = exprsearch(gp_params, problem)

    #manually push! extra info to log
    push!(logs, "parameters", ["seed", seed])
    push!(logs, "parameters", ["runtype", runtype])
    push!(logs, "parameters", ["data", data])
    add_folder!(logs, "expression", [ASCIIString, ASCIIString, ASCIIString],
        ["raw", "pretty", "natural"])
    push!(logs, "expression", [string(result.expr), pretty_string(result.tree, FMT_PRETTY),
        pretty_string(result.tree, FMT_NATURAL, true)])

    #save log
    outfile = joinpath(outdir, "$(logfileroot).txt")
    save_log(outfile, logs)

    if vis
        derivtreevis(get_derivtree(result), joinpath(outdir, "$(logfileroot)_derivtreevis"))
    end

    textfile(joinpath(outdir, "summary.txt"), "gp", seed=seed, iteration=iterations,
           fitness=result.fitness, expr=string(result.expr))
    result
end

end #module


text:
## Author: Pearson
## Date: April 2020
## Purpose: Simulation of EKNR tow cities case
## Julia version: 1.3.1

#=
This function governs the equations for solving the steady state.
=#
function Fun_2Cty_SS!(
    res_ss::AbstractArray, # residuals
    guess_ss::AbstractArray, # guess
    exos_ss::NamedTuple, # exogenous variables
    params_ss::NamedTuple # parameters
)
    # Unpack exogenous variables and parameters
    @unpack χ, Aᴰ, L, dₙᵢ, Yˢₛₛ = exos_ss
    @unpack α, βᴸ, βᴷ, δ, ρ, θ, B = params_ss

    # Initialize vectors
    Yₛₛ = similar(guess_ss, 2)
    Kₛₛ = similar(guess_ss, 2)
    wₛₛ = similar(guess_ss, 2)
    rₛₛ = similar(guess_ss, 2)
    bₛₛ = similar(guess_ss, 2)
    pᴰₛₛ = similar(guess_ss, 2)
    πₙᵢₛₛ = similar(guess_ss, 2, 2)

    # Assign guess
    Yₛₛ[:] = guess_ss[:,1]
    Kₛₛ[:] = guess_ss[:,2]

    # Derive other variables in terms of guess
    Yᴰₛₛ = Yₛₛ-Yˢₛₛ
    for n = 1:size(guess_ss,1)
        wₛₛ[n] = βᴸ*Yₛₛ[n]/L[n]
        rₛₛ[n] = βᴷ*Yₛₛ[n]/Kₛₛ[n]
        bₛₛ[n] = wₛₛ[n]^βᴸ*rₛₛ[n]^βᴷ
    end

    for n = 1:size(guess_ss,1)
        pᴰₛₛ[n] = ((bₛₛ[1]*dₙᵢ[n,1]/Aᴰ[1])^-θ+(bₛₛ[2]*dₙᵢ[n,2]/Aᴰ[2])^-θ)^(-1/θ)
        for i = 1:size(guess_ss,1)
            πₙᵢₛₛ[n,i] = (bₛₛ[i]*dₙᵢ[n,i]/pᴰₛₛ[n]/Aᴰ[i])^-θ
        end
    end

    Xᴰₛₛ = πₙᵢₛₛ'\Yᴰₛₛ

    # Equations
    nf = 0
    for n = 1:size(guess_ss,1)
        res_ss[nf+1] = abs(Xᴰₛₛ[n]/pᴰₛₛ[n]/Kₛₛ[n]/(δ/χ[n])^(1/α)-1) # capital accumulation
        nf += 1
        res_ss[nf+1] = abs((1-ρ*(1-δ))*pᴰₛₛ[n]/ρ/χ[n]/(Xᴰₛₛ[n]/pᴰₛₛ[n]/Kₛₛ[n])^α/pᴰₛₛ[n]/((1-α)+
            α*βᴷ*Yₛₛ[n]/Xᴰₛₛ[n])-1) # euler equation
        nf += 1
    end

    return res_ss, Yₛₛ, Kₛₛ
end

#=
This function governs the equations for solving the problem in levels.
=#
function Fun_2Cty_Levels!(
    res_lev::AbstractArray, # residuals
    guess_lev::AbstractArray, # guess
    init_lev::NamedTuple, # initial conditions
    exos_lev::NamedTuple, # exogenous variables
    params_lev::NamedTuple # parameters
)
    # Unpack initial conditions, exogenous variables and parameters
    @unpack K₁ = init_lev
    @unpack χ, Aᴰ, L, dₙᵢ, Kₛₛ, Yˢ = exos_lev
    @unpack α, βᴷ, βᴸ, δ, ρ, θ, T = params_lev

    # Initialize vectors
    Y = similar(guess_lev)
    K = similar(guess_lev)
    w = similar(guess_lev)
    r = similar(guess_lev)
    b = similar(guess_lev)
    pᴰ = similar(guess_lev)
    πₙᵢ = similar(guess_lev, size(dₙᵢ))
    Yᴰ = similar(guess_lev)
    Xᴰ = similar(guess_lev)

    # Assign guess
    Y[:,:] = guess_lev

    # Equations and other variables
    Yᴰ = Y-Yˢ
    K[:,1] = K₁

    for n = 1:size(guess_lev,1)
        w[n,1] = βᴸ*Y[n,1]/L[n,1]
        r[n,1] = βᴷ*Y[n,1]/K[n,1]
        b[n,1] = w[n,1]^βᴸ*r[n,1]^βᴷ
    end

    for n = 1:size(guess_lev,1)
        pᴰ[n,1] = ((b[1,1]*dₙᵢ[n,1,1]/Aᴰ[1,1])^-θ+(b[2,1]*dₙᵢ[n,2,1]/Aᴰ[2,1])^-θ)^(-1/θ)
        for i = 1:size(guess_lev,1)
            πₙᵢ[n,i,1] = (b[i,1]*dₙᵢ[n,i,1]/pᴰ[n,1]/Aᴰ[i,1])^-θ
        end
    end

    Xᴰ[:,1] = (πₙᵢ[:,:,1]')\Yᴰ[:,1]

    for tt = 2:T
        for n = 1:size(guess_lev,1)
            K[n,tt] = χ[n,tt-1]*(Xᴰ[n,tt-1]/pᴰ[n,tt-1])^α*K[n,tt-1]^(1-α)+(1-δ)*K[n,tt-1]
            w[n,tt] = βᴸ*Y[n,tt]/L[n,tt]
            r[n,tt] = βᴷ*Y[n,tt]/K[n,tt]
            b[n,tt] = w[n,tt]^βᴸ*r[n,tt]^βᴷ
        end

        for n = 1:size(guess_lev,1)
            pᴰ[n,tt] = ((b[1,tt]*dₙᵢ[n,1,tt]/Aᴰ[1,tt])^-θ+(b[2,tt]*dₙᵢ[n,2,tt]/Aᴰ[2,tt])^-θ)^(-1/θ)
            for i = 1:size(guess_lev,1)
                πₙᵢ[n,i,tt] = (b[i,tt]*dₙᵢ[n,i,tt]/pᴰ[n,tt]/Aᴰ[i,tt])^-θ
            end
        end

        Xᴰ[:,tt] = (πₙᵢ[:,:,tt]')\Yᴰ[:,tt]

        for n = 1:size(guess_lev,1)
            res_lev[n,tt-1] = abs(
                (pᴰ[n,tt-1]/ρ/χ[n,tt-1]*(Xᴰ[n,tt-1]/pᴰ[n,tt-1]/K[n,tt-1])^(1-α))/
                (pᴰ[n,tt]/χ[n,tt]*(Xᴰ[n,tt]/pᴰ[n,tt]/K[n,tt])^(1-α)*
                (χ[n,tt]*(1-α)*(Xᴰ[n,tt]/pᴰ[n,tt]/K[n,tt])^α+(1-δ))+α*r[n,tt])-1
            ) # euler equations(T-1)
        end
    end

    for n = 1:size(guess_lev,1)
        res_lev[n,T] = abs((χ[n,T]*(Xᴰ[n,T]/pᴰ[n,T])^α*K[n,T]^(1-α)+(1-δ)*K[n,T])/Kₛₛ[n]-1) # terminal condition(1)
    end

    return res_lev, K, Y, Yᴰ, Xᴰ, πₙᵢ
end

function Fun_2Cty_Changes!(
    res_hat::AbstractArray, # residuals
    guess_hat::AbstractArray, # guess
    init_hat::NamedTuple, # initial conditions
    exos_hat::NamedTuple, # exogenous variables
    params_hat::NamedTuple # parameters
)
    # Unpack initial conditions, exogenous variables and parameters
    @unpack Y₁, Yᴰ₁, πₙᵢ₁ = init_hat
    @unpack χ̂, Âᴰ, L̂, d̂ₙᵢ, ϕ̂ = exos_hat
    @unpack α, βᴷ, βᴸ, δ, ρ, θ, T = params_hat

    # Initialize vectors for changes
    K̂ = similar(guess_hat)
    Ŷ = similar(guess_hat)
    ŵ = similar(guess_hat)
    r̂ = similar(guess_hat)
    b̂ = similar(guess_hat)
    p̂ᴰ = similar(guess_hat)
    π̂ₙᵢ = similar(guess_hat, size(d̂ₙᵢ))
    Ŷᴰ = similar(guess_hat)
    X̂ᴰ = similar(guess_hat)

    # Initialize vectors for levels of flow
    Y = similar(guess_hat)
    Yᴰ = similar(guess_hat)
    πₙᵢ = similar(guess_hat, size(d̂ₙᵢ))
    Xᴰ = similar(guess_hat)
    Yˢ = similar(guess_hat)

    # Assign guess and initial conditions
    K̂[:,1] = guess_hat[:,1]
    Ŷ[:,1:T-1] = guess_hat[:,2:T]
    Y[:,1] = Y₁
    Yᴰ[:,1] = Yᴰ₁
    πₙᵢ[:,:,1] = πₙᵢ₁
    Xᴰ[:,1] = (πₙᵢ[:,:,1]')\Yᴰ[:,1]
    Yˢ[:,1] = Y[:,1]-Yᴰ[:,1]
    for n = 1:1:size(guess_hat,1)
        for ts = 2:T
            Yˢ[n,ts] = Yˢ[n,ts-1].*ϕ̂[n,ts-1]
        end
    end

    # Equations and other variables
    for tt = 2:T
        for n = 1:1:size(guess_hat,1)
        # other variables and equations
            ŵ[n,tt-1] = Ŷ[n,tt-1]/L̂[n,tt-1]
            r̂[n,tt-1] = Ŷ[n,tt-1]/K̂[n,tt-1]
            b̂[n,tt-1] = ŵ[n,tt-1]^βᴸ*r̂[n,tt-1]^βᴷ
        end

        for n = 1:1:size(guess_hat,1)
            p̂ᴰ[n,tt-1] = (πₙᵢ[n,1,tt-1]*(b̂[1,tt-1]*d̂ₙᵢ[n,1,tt-1]/Âᴰ[1,tt-1])^-θ+
                πₙᵢ[n,2,tt-1]*(b̂[2,tt-1]*d̂ₙᵢ[n,2,tt-1]/Âᴰ[2,tt-1])^-θ)^(-1/θ)
        end

        for n = 1:1:size(guess_hat,1)
            for i = 1:1:size(guess_hat,1)
                π̂ₙᵢ[n,i,tt-1] = (b̂[i,tt-1]*d̂ₙᵢ[n,i,tt-1]/p̂ᴰ[n,tt-1]/Âᴰ[i,tt-1])^-θ
                πₙᵢ[n,i,tt] = π̂ₙᵢ[n,i,tt-1]*πₙᵢ[n,i,tt-1]
            end
            Y[n,tt] = Ŷ[n,tt-1]*Y[n,tt-1]
            Yᴰ[n,tt] = Y[n,tt]-Yˢ[n,tt]
        end

        Xᴰ[:,tt] = (πₙᵢ[:,:,tt]')\(Yᴰ[:,tt])

        for n = 1:1:size(guess_hat,1)
            X̂ᴰ[n,tt-1] = Xᴰ[n,tt]/Xᴰ[n,tt-1]
            K̂[n,tt] = χ̂[n,tt-1]*(X̂ᴰ[n,tt-1]/p̂ᴰ[n,tt-1]/K̂[n,tt-1])^α*(K̂[n,tt-1]-(1-δ))+(1-δ)
            res_hat[n,tt-1] = K̂[n,tt-1]/(K̂[n,tt-1]-(1-δ))/ρ/(X̂ᴰ[n,tt-1]*((1-α)+(1-δ)*
            (K̂[n,tt-1]*p̂ᴰ[n,tt-1]/X̂ᴰ[n,tt-1])^α/χ̂[n,tt-1]/(K̂[n,tt-1]-(1-δ)))+
            α*βᴷ*Y[n,tt]/Xᴰ[n,tt-1])-1 # euler equations(T-1)
        end
    end

    for n = 1:1:size(guess_hat,1)
        res_hat[n,T] = abs(K̂[n,T]/1-1) # terminal condition(1)
    end

    return res_hat, K̂, Ŷ, Ŷᴰ, X̂ᴰ, π̂ₙᵢ
end

ext: jl


text:
@doc raw"""
```
Frankot()
```
Defines the Frankot integrator which contians the Frankot-Chellappa method of
integration. The Frankot-Chellappa method is a fast, reliable method for
integration surface normals while enforcing integrability using Fourier methods.
# Output
`Frankot()` returns a Frankot integrator which can then be called to run the
Frankot-Chellappa method on a gradient field.
# Details
Frankot-Chellappa method uses Fourier methods to attempt to solve the Poission
equation ``\nabla^2z = \partial_up + \partial_vq``. By taking the Fourier transform
of both sides we get:
```math
−(\omega^2_u + \omega^2_v)\hat{z}(\omega_u, \omega_v) = \imath \omega_u\hat{p}
(\omega_u, \omega_v) + \imath \omega_v\hat{q}(\omega_u, \omega_v)
```
By rearranging the above equation we arrive at an equation for ``\hat{z}``;
```math
\hat{z}(\omega_u, \omega_v) = \frac{\omega_u\hat{p}(\omega_u, \omega_v) +
\omega_v\hat{q}(\omega_u, \omega_v)}{\imath(\omega^2_u + \omega^2_v)}
```
From which the final surface can be found by taking the inverse Fourier transform
of ``\hat{z}``.

Due to the way ``(\omega_u, \omega_v)`` is defined the algorithm works best when
the input dimensions are odd in length. To accommodate this the integrator will
pad the edge of the inputs if they are even before running the algorithm. This
padding will be removed before returning a value hence output size will be
unaffected.
# Parameters
`Frankot` integrator take no parameters.
# Example
The following example demonstrates  the use of the `Frankot` integrator.
```julia
using ShapeFromShading, Makie

# Generate synthetic gradients
p, q = synthetic_gradient(SynthSphere(38), img_size = 151)

# Create a Frankot() integrator
frankot = Frankot()

# Calculate the heightmap from the gradients
Z = frankot(p, q)

# Normalize to maximum of 1 (not necessary but makes displaying easier)
Z = Z./maximum(Z)

# Display using Makie (Note: Makie can often take several minutes first time)
r = 0.0:0.1:4
surface(r, r, Z)
```
# Reference
[1] R. T. Frankot and R. Chellappa, "A method for enforcing integrability in shape from shading algorithms," in IEEE Transactions on Pattern Analysis and Machine Intelligence, vol. 10, no. 4, pp. 439-451, July 1988. [doi: 10.1109/34.3909](https://doi.org/10.1109/34.3909 )
"""
function (scheme::Frankot)(pin::AbstractArray, qin::AbstractArray)
    # Resize grid to be odd dimensions. (Requared to setup transofrm values)
    M, N = size(pin)
    p = zeros(Complex{Float64}, M + 1 - (M%2), N + 1 - (N%2))
    p[1:M, 1:N] = Complex{Float64}.(pin)
    q = zeros(Complex{Float64}, M + 1 - (M%2), N + 1 - (N%2))
    q[1:M, 1:N] = Complex{Float64}.(qin)

    wx, wy = setup_transform_values(N,M)
    fft!(p)
    fft!(q)
    Z = zeros(Complex{Float64}, size(p))
    Z = (-1im .* wx .* p .+ 1im .* wy .* q) ./ (2 .* π .* (wx.^2 .+ wy.^2 .+ eps()))
    ifft!(Z)
    Z = abs.(Z)
    return Z[1:M,1:N]
end

@doc raw"""
```
Path()
```
Creates a `Path()` integrator which utilises the average of two path integrals
along varying paths. Each path integral reconstructs the surface with
accumulating error along the path, hence averaging two different paths can
minimise this error, although the method still suffers if the gradient field is
not integrable at some points.
# Output
`Path()` returns a Path integrator which can then be called to integrate a
gradient field.
# Details
Under the assumption that the surface normals are approximately integrable
everywhere (``\frac{\partial p}{\partial y}\approx\frac{\partial q}{\partial x}``),
then surface can be reconstructed using the path integral defined as:
```math
z(x,y)=\oint_c\left(\frac{\partial z}{\partial x},\frac{\partial z}{\partial y}\right)\cdot dl
```
Which can be broken into two integrals representing the value at each point on
the surface as shown below for a path which integrates along the first column
then along the row.
```math
z(u,v)=\int_0^v\frac{\partial z}{\partial y}(0,y)dy + \int_0^u\frac{\partial z}{\partial x}(x,v)dx
```
The second path used in the algorithm is simply the transpose of the first,
integrating along the first row then down the column represented mathematically as:
```math
z(u,v)=\int_0^u\frac{\partial z}{\partial x}(x,0)dx + \int_0^v\frac{\partial z}{\partial y}(u,y)dy
```
The algorithm can be written, then discretised as shown below:
```math
\begin{gathered}
z(u,v)=\frac{1}{2}\left(\int_0^v\frac{\partial z}{\partial y}(0,y)dy + \int_0^u\frac{\partial z}{\partial x}(x,v)dx + \int_0^u\frac{\partial z}{\partial x}(x,0)dx + \int_0^v\frac{\partial z}{\partial y}(u,y)dy\right)\\
z(u,v)=\frac{1}{2}\left(\sum_{i=0}^vq(0,i) + \sum_{j=0}^up(j,v) + \sum_{j=0}^up(j,0) + \sum_{i=0}^vq(u,i)\right)\\
z(u,v)=\frac{1}{2}\left(\sum_{i=0}^v(q(0,i) + q(u,i)) + \sum_{j=0}^u(p(j,0) + p(j,v))\right)\\
\end{gathered}
```
It is important to note as mentioned above if there are non-integrable points in
the normal field then artefacts can appear in the reconstruction. This is
seen in the example below where the otherwise smooth sphere appears "spiky".
This can be corrected post reconstruction by smoothing but ideally a different
integrator should be used.
# Parameters
`Path` integrator take no parameters.
# Example
The following example demonstrates the use of the `Path` integrator.
```julia
using ShapeFromShading, Makie

# Generate synthetic gradients
p, q = synthetic_gradient(SynthSphere(38), img_size = 151)

# Create a Path() integrator
path = Path()

# Calculate the heightmap from the gradients
Z = path(p, q)

# Normalize to maximum of 1 (not necessary but makes displaying easier)
Z = Z./maximum(Z)

# Display using Makie (Note: Makie can often take several minutes first time)
r = 0.0:0.1:4
surface(r, r, Z)
```
# Reference
[1] D. Forsyth and J. Ponce, Computer vision: a modern approach. Upper Saddle River, N.J: Prentice Hall, 2003, pp. 84-86.
"""
function (scheme::Path)(p::AbstractArray, q::AbstractArray)
    R, C = size(p)
    Z₁ = zeros(Float64, size(p))
    Z₂ = zeros(Float64, size(p))

    #path 1
    for i = 2:R
        Z₁[i,1] = Z₁[i-1,1] - q[i,1]
    end
    for i = 1:R
        for j = 2:C
            Z₁[i,j] = Z₁[i,j-1] + p[i,j]
        end
    end

    #path 2
    for i = 2:C
        Z₂[1,i] = Z₂[1,i-1] + p[1,i]
    end
    for i = 2:R
        for j = 1:C
            Z₂[i,j] = Z₂[i-1,j] - q[i,j]
        end
    end
    Z = (Z₁ .+ Z₂) / 2
    return Z
end

@doc raw"""
```
SplitPath()
```
Creates a `SplitPath()` integrator which utilizes the average of two path integrals
along varying paths averaging the value at each step. Each path integral
reconstructs the surface with accumlating error along the path, hence averaging
two different paths at each step reduces the global error at the cost of local
error, although the method still suffers if the gradient field is not integrable
at some points it does less so the `Path()` from which it extends.
# Output
`SplitPath()` returns a SplitPath integrator which can then be called to integrate
a gradient field.
# Details
Under the assumption that the surface normals are approximately integrable
everywhere (``\frac{\partial p}{\partial y}\approx\frac{\partial q}{\partial x}``),
then surface can be reconstructed using the path integral defined as:
```math
z(x,y)=\oint_c\left(\frac{\partial z}{\partial x},\frac{\partial z}{\partial y}\right)\cdot dl
```
By expanding on this principle and the discreate summation from `Path()` we can
arrive at the discreate expresion for the value at each point, assuming all
values prior to that point have been calculated, as follows:
```math
z_{u,v} = \frac{1}{2}(z_{u-1,v}+p_{u-1,v}+z_{u,v-1}+q_{u,v-1})
```
As with other similar methods (see `Horn()`) care must be taken with regards to
boundaries which can be calculated, to a constant value ``z(0,0)`` which is assumed
to be the zero point, using:
```math
\begin{gathered}
z_{u,0} = z_{u-1,0}+p_{u-1,0}\\
z_{0,v} = z_{0,v-1}+q_{0,v-1}
\end{gathered}
```
It is important to note as mentioned above if there are non-integrable points in
the normal field then artefacts can appear in the reconstruction. These errors
gradually average out but will lead to "streaks" appearing in the reconstruction.
This is seen in the example below where the otherwise smooth sphere appears has
ripple like structures pointing toward to top right corner. This can be corrected
post reconstruction by smoothing but ideally a different integrator should be used.
It is also interesting to note the parallels between this method and the Horn
and Brooks method, with this method being effectively the forward component of
Horn's method. As such this algorithm provided a useful middle ground between
direct integration algorithms and iterative algorithms such as the Horn and Brooks
method.
# Parameters
`SplitPath` integrator take no parameters.
# Example
The following example demonstrates the use of the `SplitPath` integrator.
```julia
using ShapeFromShading, Makie

# Generate synthetic gradients
p, q = synthetic_gradient(SynthSphere(38), img_size = 151)

# Create a Path() integrator
splitPath = SplitPath()

# Calculate the heightmap from the gradients
Z = splitPath(p, q)

# Normalize to maximum of 1 (not necessary but makes displaying easier)
Z = Z./maximum(Z)

# Display using Makie (Note: Makie can often take several minutes first time)
r = 0.0:0.1:4
surface(r, r, Z)
```
# Reference
[1] D. Forsyth and J. Ponce, Computer vision: a modern approach. Upper Saddle River, N.J: Prentice Hall, 2003, pp. 84-86.
[2] B. Horn and M. Brooks, "The variational approach to shape from shading", Computer Vision, Graphics, and Image Processing, vol. 33, no. 2, pp. 174-208, 1986. [doi: 10.1016/0734-189x(86)90114-3](https://doi.org/10.1016/0734-189x(86)90114-3 )
"""
function (scheme::SplitPath)(p::AbstractArray, q::AbstractArray)
    R, C = size(p)
    Z = zeros(Float64, size(p))
    for i = 2:R
        Z[i,1]=Z[i-1,1] + q[i,1]
        Z[1,i]=Z[1,i-1] - p[1,i]
    end
    for i = 2:R
        for j = i:R
            Z[i,j] = ((Z[i,j-1] + p[i,j]) + (Z[i-1,j] - q[i,j])) / 2
            Z[j,i] = ((Z[j,i-1] + p[j,i]) + (Z[j-1,i] - q[j,i])) / 2
        end
    end
    return Z
end

@doc raw"""
```
Horn(ϵ::Real = 1.0, max_iter::Real = 10000)
```
Implements the Horn and Brook's method of integrating surface normals. This
algorithm offers an iterative solution to the Poisson equation describing the
surface providing good reconstructions under most conditions.
# Output
`Horn()` returns a Horn integrator which can then be called to integrate
a gradient field.
# Details
The Horn and Brook's method attempts to solve the Poisson equation ``\nabla^2z = \partial_up + \partial_vq``
by the discretization given below:
```math
z_{u+1,v}+z_{u,v+1}+z_{u-1,v}+z_{u,v-1}-4z_{u,v}=\frac{p_{u+1,v}-p_{u-1,v}}{2}+\frac{q_{u,v+1}-q_{u,v-1}}{2}
```
Which can be rearranged to give the iterative scheme provided by:
```math
z_{u,v}^{k+1}= \frac{z_{u+1,v}^k + z_{u,v+1}^k + z_{u-1,v}^k + z_{u,v-1}^k}{4} - \frac{p_{u+1,v}-p_{u-1,v}}{8} - \frac{q_{u,v+1}-q_{u,v-1}}{8}
```
This scheme will always converge to a solution however the rate of convergence
may depend upon the initial solution. This implementation will initilize with a
zero solution. Neumann boundary conditions are imposed at the edges where the
scheme would otherwise go out of bounds.
# Parameters
The function parameters are described in more detail below.
##  `Max_iter`:
An `Int` which controls the number of iterations the algorithm will run for.
## `ϵ`:
A `Real` representing the distance between pixels. This will Control how tall
the final reconstruction is relative the array grid.
# Example
The following example demonstrates the use of the `Horn` integrator.
```julia
using ShapeFromShading, Makie

# Generate synthetic gradients
p, q = synthetic_gradient(SynthSphere(38), img_size = 151)

# Create a Horn() integrator
horn = Horn(ϵ = 0.03, max_iter = 10000)

# Calculate the heightmap from the gradients
Z = horn(p, q)

# Display using Makie (Note: Makie can often take several minutes first time)
r = 0.0:0.1:4
surface(r, r, Z)
```
# Reference
[1] B. Horn and M. Brooks, "The variational approach to shape from shading", Computer Vision, Graphics, and Image Processing, vol. 33, no. 2, pp. 174-208, 1986. [doi: 10.1016/0734-189x(86)90114-3](https://doi.org/10.1016/0734-189x(86)90114-3 )
"""
function (scheme::Horn)(p::AbstractArray, q::AbstractArray)
    iter = copy(scheme.max_iter)
    ϵ = copy(scheme.ϵ)
    Z = zeros(Float64, size(p))
    Zᵏ⁺¹ = zeros(Float64, size(p))
    h = zeros(Float64, size(p))
    v = zeros(Float64, size(p))
    hv = zeros(Float64, size(p))
    R, C = size(p)
    for i = 1:R
        Z[i,1] = p[i,1]
        Z[i,R] = p[i,R]
    end
    for i = 1:C
        Z[1,i] = q[1,i]
        Z[C,i] = q[C,i]
    end
    for i = 2:(R-1)
        for j = 2:(C-1)
            h[i,j] = (p[i,j+1] - p[i,j-1]) / 2.0
            v[i,j] = (q[i+1,j] - q[i-1,j]) / 2.0
            hv[i,j] = (ϵ / 4) * (h[i,j] - v[i,j])
        end
    end
    for k = 1:iter
        copyto!(Zᵏ⁺¹, Z)
        for i = 2:(R-1)
            for j = 2:(C-1)
                Zᵏ⁺¹[i,j] = (Z[i-1,j] + Z[i+1,j] + Z[i,j-1] + Z[i,j+1]) / 4.0
                Zᵏ⁺¹[i,j] = Zᵏ⁺¹[i,j] - hv[i,j]
            end
        end
        Z = Zᵏ⁺¹
    end
    return Z
end

@doc raw"""
```
Durou(ϵ::Real = 1.0, max_iter::Real = 1000)
```
Implements the Durou and Courteille method of integrating surface normals. This
algorithm offers an iterative solution to the Poisson equation describing the
surface extending Horn and Brook's method by improving the boundary approximation
and providing good reconstructions under most conditions.
# Output
`Durou()` returns a Durou integrator which can then be called to integrate
a gradient field.
# Details
TheDurou and Courteille's method attempts to solve the Poisson equation
``\nabla^2z = \partial_up + \partial_vq`` by the discretization given below:
```math
z_{u+1,v}+z_{u,v+1}-2z_{u,v}=\frac{p_{u+1,v}+p_{u,v}}{2}+\frac{q_{u,v+1}+q_{u,v}}{2}
```
Which can be rearranged to give the iterative scheme provided by:
```math
z_{u,v}^{k+1}= \frac{z_{u+1,v}^k + z_{u,v+1}^k}{2} - \frac{p_{u+1,v}+p_{u,v}}{4} - \frac{q_{u,v+1}+q_{u,v}}{4}
```
This scheme will always converge to a solution however the rate of convergence
may depend upon the initial solution. This implementation will initialize with a
zero solution. Natural boundary conditions are imposed at the edges using the
condition ``\partial_uz-p+\partial_v-q=0``. Although faster then the Horn and Brook's
method and better at handling boundaries, it can generate a worse solution under
some conditions.
# Parameters
The function parameters are described in more detail below.
##  `Max_iter`:
An `Int` which controls the number of iterations the algorithm will run for.
the range [0,1].
## `ϵ`:
A `Real` representing the distance between pixels. This will Control how tall
the final reconstruction is relative the array grid.
# Example
The following example demonstrates the use of the `Durou` integrator.
```julia
using ShapeFromShading, Makie

# Generate synthetic gradients
p, q = synthetic_gradient(SynthSphere(38), img_size = 151)

# Create a Durou() integrator
durou = Durou(ϵ = 0.03, max_iter = 10000)

# Calculate the heightmap from the gradients
Z = durou(p, q)

# Display using Makie (Note: Makie can often take several minutes first time)
r = 0.0:0.1:4
surface(r, r, Z)
```
# Reference
[1] Y. Quéau, J. Durou and J. Aujol, "Normal Integration: A Survey", Journal of Mathematical Imaging and Vision, vol. 60, no. 4, pp. 576-593, 2017. [doi: 10.1007/s10851-017-0773-x](https://doi.org/10.1007/s10851-017-0773-x )
"""
function (scheme::Durou)(p::AbstractArray, q::AbstractArray)
    iter = copy(scheme.max_iter)
    ϵ = copy(scheme.ϵ)
    Z = zeros(Float64, size(p))
    Zᵏ⁺¹ = zeros(Float64, size(p))
    h = zeros(Float64, size(p))
    v = zeros(Float64, size(p))
    hv = zeros(Float64, size(p))
    R, C = size(p)
    for i = 1:(R-1)
        for j = 1:(C-1)
            h[i,j] = (p[i,j+1] + p[i,j]) / 2
            v[i,j] = (q[i+1,j] + q[i,j]) / 2
            hv[i,j] = (ϵ / 2) * (h[i,j] - v[i,j])
        end
    end
    for k = 1:iter
        copyto!(Zᵏ⁺¹, Z)
        for i = 1:(R-1)
            for j = 1:(C-1)
                Zᵏ⁺¹[i,j] = (Z[i+1,j] + Z[i,j+1]) / 2
                Zᵏ⁺¹[i,j] = Zᵏ⁺¹[i,j] - hv[i,j]
            end
        end
        Z = Zᵏ⁺¹
    end
    return Z
end

# Generates matricies Dᵤ⁺, Dᵤ⁻, Dᵥ⁺ and Dᵥ⁻ as per eq 11, 12, 13, 14 and 17
function gen_matrix(Dᵤ⁺, Dᵤ⁻, Dᵥ⁺, Dᵥ⁻,T,mask)
    U,V = size(mask)
    for i in CartesianIndices(mask)
        #u+
        if i[1] + 1 <= U && mask[i] == 1 && mask[i[1]+1,i[2]] == 1
            Dᵤ⁺[i[1]+(i[2]-1)*U, i[1]+(i[2]-1)*U] = -1.0
            Dᵤ⁺[i[1]+(i[2]-1)*U, i[1]+(i[2]-1)*U+1] = 1.0
        end
        #u-
        if i[1] - 1 > 0 && mask[i] == 1 && mask[i[1]-1,i[2]] == 1
            Dᵤ⁻[i[1]+(i[2]-1)*U, i[1]+(i[2]-1)*U] = 1.0
            Dᵤ⁻[i[1]+(i[2]-1)*U, i[1]+(i[2]-1)*U-1] = -1.0
        end
        #v+
        if i[2] + 1 <= V && mask[i] == 1 && mask[i[1],i[2]+1] == 1
            Dᵥ⁺[i[1]+(i[2]-1)*V, i[1]+(i[2]-1)*V] = -1.0
            Dᵥ⁺[i[1]+(i[2]-1)*V, i[1]+(i[2])*V] = 1.0
        end
        #v-
        if i[2] - 1 > 0 && mask[i] == 1 && mask[i[1],i[2]-1] == 1
            Dᵥ⁻[i[1]+(i[2]-1)*V, i[1]+(i[2]-1)*V] = 1.0
            Dᵥ⁻[i[1]+(i[2]-1)*V, i[1]+(i[2]-2)*V] = -1.0
        end
    end
end

@doc raw"""
```
Quadratic(z::AbstractArray, λ::AbstractArray = fill(10.0^-6, size(z)), mask::AbstractArray = fill(1.0, size(z)))
```
Implements the quadratic variational least squared method proposed by Aujol, Durou
and Quéau. The algorithm solves the least squares system generated from minimizing
a fidelity term ``\mathcal{F}(z) = \iint_{(u,v)\in \Omega}\Phi(||\nabla z(u,v) - g(u,v)||)dudv``
and a regularization term ``\mathcal{R}(z) = \iint_{(u,v)\in \Omega}\lambda \left[z(u,v)-z^0(u,v)\right]^2dudv``
where ``\Phi(s) = s^2``. This method is able to quickly produce good solutions on
a smooth surface and can easily handle non-rectangular domains or sub-divided
domains and can provide a good starting solution for other algorithms.
# Output
`Quadratic()` returns a Quadratic integrator which can then be called to run the
quadratic method on a gradient field.
# Details
As mentioned above the quadratic variational least squared method aims to minimize
a fidelity term ``\mathcal{F}(z) = \iint_{(u,v)\in \Omega}\Phi(||\nabla z(u,v) - g(u,v)||)dudv``
and a regulization term ``\mathcal{R}(z) = \iint_{(u,v)\in \Omega}\lambda \left[z(u,v)-z^0(u,v)\right]^2dudv``
where ``\Phi(s) = s^2``.

This leads to the minimization problem given by:
```math
min\iint_{(u,v)\in \Omega}||\nabla z(u,v)-\mathbf{g}(u,v)||^2+\lambda(u,v)\left[z(u,v)-z^0(u,v)\right]^2dudv
```
By discretizing the problem we can arrive at the functional:
```math
\begin{aligned}
E(z) = \frac{1}{2}\Big(\sum_{(u,v)\in \Omega_u^+}[\partial_u^+z_{u,v}-p_{u,v}]^2 \\+ \sum_{(u,v)\in \Omega_u^-}[\partial_u^-z_{u,v}-p_{u,v}]^2 \\+ \sum_{(u,v)\in \Omega_v^+}[\partial_v^+z_{u,v}-q_{u,v}]^2 \\+ \sum_{(u,v)\in \Omega_v^-}[\partial_v^-z_{u,v}-q_{u,v}]^2\Big) \\+ \sum_{(u,v)\in \Omega}[z_{u,v}-z^0_{u,v}]^2
\end{aligned}
```
where ``\Omega_u^+`` represents the domain where ``(u,v)\in\{(u,v)\in\Omega|(u,v)(u+1,v)\in\Omega\}`` etc..
Using this definition the discrete differences can be converted into matrix form,
where ``p_{u,v}=z_{u+1,v}-z_{u,v}`` is the forward difference in the u direction etc.
This data is then stacked into three vectors; ``\mathbf{z},\mathbf{p},\mathbf{q}\in\R^{|\Omega|}``.
Thus the matrix reresenting each of these is defined as below where m(i) is the
mapping of the ith element of this vector to its corresponding point in ``(u,v)``
and ``D_u^+`` is a ``|\Omega|\times|\Omega|`` matrix.
```math
D_u^+[i,j]=\begin{cases}
   0 &\text{if } m(i)\notin\Omega_u^+ \text{or } j \ne i \text{or } j \ne i+1\\
   -1 &\text{if } j = i \text {and } m(i)\in\Omega_u^+\\
   1 &\text{if } j = i+1 \text {and } m(i)\in\Omega_u^+
\end{cases}
```
For a 2X2 domain this looks like:
```math
D_u^+ = \begin{bmatrix}
   -1 & 1 & 0 & 0 \\
   0 & 0 & 0 & 0 \\
   0 & 0 & -1 & 1 \\
   0 & 0 & 0 & 0
\end{bmatrix}
```
The other three discrete differences matrices are similarly defined from there
definitions to be ``D_u^-``, ``D_v^+`` and ``D_v^-``. These can then be used to
redefine the minimization problem to be in the form:
```math
E(\mathbf{z})=\frac{1}{2}\left(||D_u^+\mathbf{z}-\mathbf{p}||^2+||D_u^-\mathbf{z}-\mathbf{p}||^2+||D_v^+\mathbf{z}-\mathbf{q}||^2+||D_v^-\mathbf{z}-\mathbf{q}||^2\right)+||\Lambda(\mathbf{z}-\mathbf{z}^0)||^2
```
Where ``\Lambda`` is the ``|\Omega|\times|\Omega|`` diagonal matrix containing
the values of ``\sqrt{\lambda_{u,v}}``. Using the above definitions the negative
Laplacian matrix can then be defined as:
```math
L=\frac{1}{2}[D_u^{+\top}D_u^++D_u^{-\top}D_u^-+D_v^{+\top}D_v^++D_v^{-\top}D_v^-]
```
Finally the least minimization problem can be represented in the form of a least
squares problem of the form ``A\mathbf{z}=\mathbf{b}`` where:
```math
\begin{gathered}
A=L+\Lambda^2\\
\mathbf{b}=\frac{1}{2}\left[D_u^{+\top}+D_u^{-\top}\right]\mathbf{p}+\frac{1}{2}\left[D_v^{+\top}+D_v^{-\top}\right]\mathbf{q}+\Lambda^2\mathbf{z}^0\\
=D_u\mathbf{p}+D_v\mathbf{q}+\Lambda^2\mathbf{z}^0
\end{gathered}
```
This system is then solved using a standard conjugate gradient algorithm where
the initialization has only a slight impact on the runtime and no impact on the
final solution. The algorithm provides good results on smooth surfaces but
struggles in the presence of discontinuities.
# Parameters
## `z`:
An `AbstractArray` which defines the value of ``z^0`` the initial solution and
prior to be used in the regularization term. Must be provided.
## `λ`:
An `AbstractArray` the same size as `z`, defaulting to ``10.0^{-6}`` everywhere.
This defines theregulization weight at each point. Large values will force the
algorithm to keep the solution near to ``z^0`` at that position. Can be used to
keep the solution near the initial solution or guide the solution to a certain
known value at points (i.e. known maxima and minima). This value should be set
uniformly small otherwise.
## `mask`:
An `AbstractArray` the same size as `z`, which guides the algorithm as to where
the valid domain is. Values of `1` will be in the domain ``\Omega`` while other
values will be ignored and set to ``z^0``. This can be used to
integrate over sub-domain or to segment the domain into parts. The `gen_mask()`
funtion can be used to generate a mask which will remove non-integrable regions
dramatically improving the solution under most condition at the cost of not
integrating the entire solution.
# Example
The following example demonstrates the use of the `Quadratic` integrator.
```julia
using ShapeFromShading, Makie

# Generate synthetic gradients
p, q = synthetic_gradient(Prism(75), img_size = 151)

# Create a Quadratic() integrator
quadratic = Quadratic(z=zeros(size(p)))
quadraticMasked = Quadratic(z=zeros(size(p)), mask=gen_mask(p,q,1.0)[:,:,1])

# Calculate the heightmap from the gradients
Z = quadratic(p, q)
Z2 = quadraticMasked(p, q)

# Normalize to maximum of 1 (not necessary but makes displaying easier)
Z = Z./maximum(Z)
Z2 = Z2./maximum(Z2)

# Display using Makie (Note: Makie can often take several minutes first time)
r = 0.0:0.1:4
vbox(surface(r, r, Z), surface(r, r, Z2))
```
# Reference
[1] Y. Quéau, J. Durou and J. Aujol, "Variational Methods for Normal Integration", Journal of Mathematical Imaging and Vision, vol. 60, no. 4, pp. 609-632, 2017. [doi: 10.1007/s10851-017-0777-6](https://doi.org/10.1007/s10851-017-0777-6 )
"""
function (scheme::Quadratic)(pIn::AbstractArray, qIn::AbstractArray)
    z⁰ = copy(scheme.z)
    λ = copy(scheme.λ)
    mask = copy(scheme.mask)
    z = copy(z⁰)
    mask = rotr90(mask)
    λ = rotr90(λ)
    p = rotr90(copy(pIn)).*mask
    q = rotr90(copy(qIn)).*mask
    T = first(size(z))
    index = CartesianIndices(z)
    index = reshape(index, length(index),1)

    Dᵤ⁺ = spzeros(Float64, T^2, T^2)
    Dᵤ⁻ = spzeros(Float64, T^2, T^2)
    Dᵥ⁺ = spzeros(Float64, T^2, T^2)
    Dᵥ⁻ = spzeros(Float64, T^2, T^2)
    gen_matrix(Dᵤ⁺, Dᵤ⁻, Dᵥ⁺, Dᵥ⁻, T, mask)

    # Calculate Dᵤ, Dᵥ and L as per eq 23 and 24
    Dᵤ = 0.5*(transpose(Dᵤ⁺) .+ transpose(Dᵤ⁻))
    Dᵥ = 0.5*(transpose(Dᵥ⁺) .+ transpose(Dᵥ⁻))
    L = 0.5*((transpose(Dᵤ⁺) * Dᵤ⁺) .+ (transpose(Dᵤ⁻) * Dᵤ⁻) .+ (transpose(Dᵥ⁺) * Dᵥ⁺) .+ (transpose(Dᵥ⁻) * Dᵥ⁻))

    Λ = spzeros(Float64,T^2,T^2)
    A = spzeros(Float64,T^2,T^2)
    P = zeros(Float64,T^2)
    Q = zeros(Float64,T^2)
    Z⁰ = zeros(Float64,T^2)
    Z = zeros(Float64,T^2)

    # Vectorize inputs
    for j in CartesianIndices(z)
        Λ[j[1]+(j[2]-1)*T,j[1]+(j[2]-1)*T] = sqrt(λ[j])
        P[j[1]+(j[2]-1)*T] = p[j]
        Q[j[1]+(j[2]-1)*T] = q[j]
        Z⁰[j[1]+(j[2]-1)*T] = z⁰[j]
        Z[j[1]+(j[2]-1)*T] = z[j]
    end

    # Calculate A and b as per eq 23 abd 24
    A = L .+ (Λ^2)
    b = Dᵤ*P .+ Dᵥ*Q .+ (Λ^2)*Z⁰

    # Solve system
    pl = AMGPreconditioner{RugeStuben}(A)
    cg!(Z, A, b, tol=10.0^-4, Pl = pl)

    for j in eachindex(index)
        z[index[j]] = Z[j]
    end
    z = z.*mask
    z = rotl90(z)
    return z
end

@doc raw"""
```
TotalVariation(z::AbstractArray, α::Real = 1.0, λ::AbstractArray = fill(10.0^-6, size(z)), mask::AbstractArray = fill(1.0, size(z)), max_iter::Int = 100)
```
Implements the total variational method proposed by Aujol, Durou and Quéau. The
algorithm solves the same minimization problem as th `Quadratic`method exept the
fidelity terms function ``\Phi(s)=||s||_{L_1}``. This method is able to produce
good solutions on a smooth and piecewise smooth surface and can easily handle
non-rectangular domains or sub-divided domains.
# Output
`TotalVariation()` returns a TotalVariation integrator which can then be called to run the
quadratic method on a gradient field.
# Details
As discussed above this algorithm used the same fidelity and regularization terms
as the `Quadratic` method exept the ``\Phi(s)=s^2`` term is replaced with
``\Phi(s)=||s||_{L_1}``. This leads to the minimization problem defined by:
```math
min\iint_{(u,v)\in \Omega}||\nabla z(u,v)-\mathbf{g}(u,v)||+\lambda(u,v)\left[z(u,v)-z^0(u,v)\right]^2dudv
```
By considering the four posible discreatisations of ``∇z(u,v)`` we can generate
four posible domains to consider given by; ``\Omega^{UV}=\Omega_u^U\cup\Omega_v^V, (U,V)\in\{+,-\}^2``
where ``\{+,-\}^2`` refers to all posible combinations of +,-.
Using these the following discreate functional can be generated:
```math
\begin{aligned}
E(\mathbf{z})=\frac{1}{4}\Big(&\sum_{(u,v)\in\Omega^{++}}\sqrt{[\partial_u^+z_{u,v}-p_{u,v}]^2+[\partial_v^+z_{u,v}-q_{u,v}]^2}\\
+&\sum_{(u,v)\in\Omega^{+-}}\sqrt{[\partial_u^+z_{u,v}-p_{u,v}]^2+[\partial_v^-z_{u,v}-q_{u,v}]^2}\\
+&\sum_{(u,v)\in\Omega^{-+}}\sqrt{[\partial_u^-z_{u,v}-p_{u,v}]^2+[\partial_v^+z_{u,v}-q_{u,v}]^2}\\
+&\sum_{(u,v)\in\Omega^{--}}\sqrt{[\partial_u^-z_{u,v}-p_{u,v}]^2+[\partial_v^-z_{u,v}-q_{u,v}]^2}\Big)\\
+&\sum_{(u,v)\in\Omega}\lambda_{u,v}\left[z_{u,v}-z^0_{u,v}\right]^2
\end{aligned}
```
Which simplifies to the minimization problem:
```math
\begin{gathered}
min\frac{1}{4}\sum_{(U,V)\in\{+,-\}^2}\sum_{(u,v)\in\Omega^{UV}}||\mathbf{r}_{(u,v)}^{UV}||+\sum_{(u,v)\in\Omega}\lambda_{u,v}\left[z_{u,v}-z^0_{u,v}\right]^2\\
\mathbf{r}_(u,v)^{UV}=\nabla^{UV}z_{u,v}-\mathbf{g}_{u,v}
\end{gathered}
```
This leads to the optimization scheme using an ADMM algorithm defined by:
```math
\begin{aligned}
z^{(k+1)}=&min\frac{\alpha}{8}\sum_{(U,V)\in\{+,-\}^2}\sum_{(u,v)\in\Omega^{UV}}||\nabla^{UV}z_{u,v}-(\mathbf{g}_{u,v}+\mathbf{r}_{(u,v)}^{UV^{(k)}}-\mathbf{b}_{(u,v)}^{UV^{(k)}})||\\&+\sum_{(u,v)\in\Omega}\lambda_{u,v}\left[z_{u,v}-z^0_{u,v}\right]^2\\
\mathbf{r}_{(u,v)}^{UV^{(k+1)}}=&min\frac{\alpha}{8}||\mathbf{r}-(\nabla^{UV}z_{u,v}-\mathbf{g}_{u,v}+\mathbf{b}_{(u,v)}^{UV^{(k)}})||+||\mathbf{r}||\\
\mathbf{b}_{(u,v)}^{UV^{(k+1)}}=&\mathbf{b}_{(u,v)}^{UV^{(k)}})+\nabla^{UV}z_{u,v}-\mathbf{g}_{u,v}-\mathbf{r}_{(u,v)}^{UV^{(k+1)}}
\end{aligned}
```
The z update then can be solved using the linear system defined below, where D_{u,v}^{U,V}
and ``\Lambda`` are the same at those defined in [`Quadratic`](@ref).
```math
\begin{aligned}
    A_{TV}&\mathbf{z}^{(k+1)}=b_{TV}^{(k)}\\
    A_{TV}&=\frac{\alpha}{8}\sum_{(U,V)\in\{+,-\}^2}\left[D_u^{U\top}D_u^{U}+D_v^{V\top}D_v^{V}\right] + \Lambda^2\\
    b_{TV}^{(k)}&=\frac{\alpha}{8}\sum_{(U,V)\in\{+,-\}^2}\left[D_u^{U\top}\mathbf{P}^{UV^{(k)}} + D_v^{V\top}\mathbf{Q}^{UV^{(k)}}\right]+ \Lambda^2\mathbf{z}^0\\
\end{aligned}
```
Where ``\mathbf{P}^{UV^{(k)}}, \mathbf{Q}^{UV^{(k)}}`` are the u and v components
of ``\mathbf{g}+\mathbf{r}^{UV^{(k)}}-\mathbf{b}^{UV^{(k)}}``. This can be solved
using conjugate gradient. Finally the update to ``\mathbf{r}^{UV}`` can be computed
as:
```math
\begin{aligned}
    \mathbf{r}^{UV^{(k+1)}}&=max\Big\{||\mathbf{s}_{u,v}^{UV^{(k+1)}}||-\frac{4}{\alpha},0\Big\}\frac{\mathbf{s}_{u,v}^{UV^{(k+1)}}}{||\mathbf{s}_{u,v}^{UV^{(k+1)}}||}\\
    \text{Where:}&\\
    \mathbf{s}^{UV^{(k+1)}}&=\nabla^{UV}z_{u,v}^{(k+1)}-\mathbf{g}_{u,v}+\mathbf{b}_{u,v}^{UV^{(k)}}
\end{aligned}
```
# Parameters
## `z`:
An `AbstractArray` which defines the value of ``z^0`` the initial solution and
prior to be used in the regulization term. Must be provided.
## `α`:
A `Real` with defult value of `1.0` which controls the step size. In theory this
value should have no impact on final solution but in practice larger values can
lead to worse solutions while values which are two small may lead non-convergance
in the least square update step, causing the algorithm to hang for long periods
of time. Values below `0.25` are not recomended but may work depending on domain
size and inputs.
## `λ`:
An `AbstractArray` the same size as `z` defulting to ``10.0^{-6}`` everywhere.
This defines theregulization weight at each point. Large valueas will force the
algorithm to keep the solution near to ``z^0`` at that position. Can be used to
keep the solution near the initial solution or guide the solution to a certian
known value at points (i.e. known maxima and minima). This value should be set
uniformly small otherwise.
## `mask`:
An `AbstractArray` the same size as `z`, which guides the algorithm as to where
the valid domain is. Values of `1` will be in the domain ``\Omega`` while other
values will be ignored and set to ``z^0``. This can be used to
integrate over sub-domain or to segment the domain into parts. The `gen_mask()`
funtion can be used to generate a mask which will remove non-integrable regions
dramatically improving the solution under most condition at the cost of not
integrating the entire solution.
# Example
The following example demonstrates the use of the `TotalVariation` integrator.
```julia
using ShapeFromShading, Makie

# Generate synthetic gradients
p, q = synthetic_gradient(Prism(75), img_size = 151)

# Create a TotalVariation() integrator
totalVariation = TotalVariation(z=zeros(size(p)), α=1.0)
totalVariationMasked = TotalVariation(z=zeros(size(p)), α=0.5, mask=gen_mask(p,q,1.0)[:,:,1])

# Calculate the heightmap from the gradients
Z = totalVariation(p, q)
Z2 = totalVariationMasked(p, q)

# Normalize to maximum of 1 (not necessary but makes displaying easier)
Z = Z./maximum(Z)
Z2 = Z2./maximum(Z2)

# Display using Makie (Note: Makie can often take several minutes first time)
r = 0.0:0.1:4
vbox(surface(r, r, Z), surface(r, r, Z2))
```
# Reference
[1] Y. Quéau, J. Durou and J. Aujol, "Variational Methods for Normal Integration", Journal of Mathematical Imaging and Vision, vol. 60, no. 4, pp. 609-632, 2017. [doi: 10.1007/s10851-017-0777-6](https://doi.org/10.1007/s10851-017-0777-6 )
"""
function (scheme::TotalVariation)(pIn::AbstractArray, qIn::AbstractArray)
    z⁰ = copy(scheme.z)
    λ = copy(scheme.λ)
    mask = copy(scheme.mask)
    α = copy(scheme.α)
    max_iter = copy(scheme.max_iter)
    z = copy(z⁰)
    mask = rotr90(mask)
    λ = rotr90(λ)
    p = rotr90(copy(pIn)).*mask
    q = rotr90(copy(qIn)).*mask
    T = first(size(z))
    index = CartesianIndices(z)
    index = reshape(index, length(index),1)

    Dᵤ⁺ = spzeros(Float64,T^2,T^2)
    Dᵤ⁻ = spzeros(Float64,T^2,T^2)
    Dᵥ⁺ = spzeros(Float64,T^2,T^2)
    Dᵥ⁻ = spzeros(Float64,T^2,T^2)
    gen_matrix(Dᵤ⁺, Dᵤ⁻, Dᵥ⁺, Dᵥ⁻, T, mask)
    Λ = spzeros(Float64,T^2,T^2)
    A = spzeros(Float64,T^2,T^2)
    P = zeros(Float64,T^2)
    Q = zeros(Float64,T^2)
    Z⁰ = zeros(Float64,T^2)
    Z = zeros(Float64,T^2)

    # Vectorize inputs
    for j in CartesianIndices(z)
        Λ[j[1]+(j[2]-1)*T,j[1]+(j[2]-1)*T] = sqrt(λ[j])
        P[j[1]+(j[2]-1)*T] = p[j]
        Q[j[1]+(j[2]-1)*T] = q[j]
        Z⁰[j[1]+(j[2]-1)*T] = z⁰[j]
        Z[j[1]+(j[2]-1)*T] = z[j]
    end

    # Calculate A as per eq 59
    A = (α/8)*(((transpose(Dᵤ⁺) * Dᵤ⁺) .+ (transpose(Dᵥ⁺) * Dᵥ⁺)) .+ ((transpose(Dᵤ⁻) * Dᵤ⁻) .+ (transpose(Dᵥ⁻) * Dᵥ⁻)) .+ ((transpose(Dᵤ⁻) * Dᵤ⁻) .+ (transpose(Dᵥ⁺) * Dᵥ⁺)) .+ ((transpose(Dᵤ⁺) * Dᵤ⁺) .+ (transpose(Dᵥ⁻) * Dᵥ⁻))) .+ (Λ^2)
    pl = AMGPreconditioner{RugeStuben}(A)

    # Calculate b, s, r, P and Q in each combination of {+,-} as per equations 57, 61, 62
    bᵤ⁺⁺ = zeros(Float64,T^2)
    bᵥ⁺⁺ = zeros(Float64,T^2)
    sᵤ⁺⁺ = Dᵤ⁺*Z - P + bᵤ⁺⁺
    sᵥ⁺⁺ = Dᵥ⁺*Z - Q + bᵥ⁺⁺
    rᵤ⁺⁺ = max(norm(sᵤ⁺⁺) - (4/α),0)*(sᵤ⁺⁺./(norm(sᵤ⁺⁺)+eps(Float64)))
    rᵥ⁺⁺ = max(norm(sᵥ⁺⁺) - (4/α),0)*(sᵥ⁺⁺./(norm(sᵥ⁺⁺)+eps(Float64)))
    bᵤ⁺⁺ = bᵤ⁺⁺ + Dᵤ⁺*Z - P - rᵤ⁺⁺
    bᵥ⁺⁺ = bᵥ⁺⁺ + Dᵥ⁺*Z - Q - rᵥ⁺⁺
    P⁺⁺ = P + rᵤ⁺⁺ - bᵤ⁺⁺
    Q⁺⁺ = Q + rᵥ⁺⁺ - bᵥ⁺⁺

    bᵤ⁺⁻ = zeros(Float64,T^2)
    bᵥ⁺⁻ = zeros(Float64,T^2)
    sᵤ⁺⁻ = Dᵤ⁺*Z - P + bᵤ⁺⁻
    sᵥ⁺⁻ = Dᵥ⁻*Z - Q + bᵥ⁺⁻
    rᵤ⁺⁻ = max(norm(sᵤ⁺⁻) - (4/α),0)*(sᵤ⁺⁻./(norm(sᵤ⁺⁻)+eps(Float64)))
    rᵥ⁺⁻ = max(norm(sᵥ⁺⁻) - (4/α),0)*(sᵥ⁺⁻./(norm(sᵥ⁺⁻)+eps(Float64)))
    bᵤ⁺⁻ = bᵤ⁺⁻ + Dᵤ⁺*Z - P - rᵤ⁺⁻
    bᵥ⁺⁻ = bᵥ⁺⁻ + Dᵥ⁻*Z - Q - rᵥ⁺⁻
    P⁺⁻ = P + rᵤ⁺⁻ - bᵤ⁺⁻
    Q⁺⁻ = Q + rᵥ⁺⁻ - bᵥ⁺⁻

    bᵤ⁻⁺ = zeros(Float64,T^2)
    bᵥ⁻⁺ = zeros(Float64,T^2)
    sᵤ⁻⁺ = Dᵤ⁻*Z - P + bᵤ⁻⁺
    sᵥ⁻⁺ = Dᵥ⁺*Z - Q + bᵥ⁻⁺
    rᵤ⁻⁺ = max(norm(sᵤ⁻⁺) - (4/α),0)*(sᵤ⁻⁺./(norm(sᵤ⁻⁺)+eps(Float64)))
    rᵥ⁻⁺ = max(norm(sᵥ⁻⁺) - (4/α),0)*(sᵥ⁻⁺./(norm(sᵥ⁻⁺)+eps(Float64)))
    bᵤ⁻⁺ = bᵤ⁻⁺ + Dᵤ⁻*Z - P - rᵤ⁻⁺
    bᵥ⁻⁺ = bᵥ⁻⁺ + Dᵥ⁺*Z - Q - rᵥ⁻⁺
    P⁻⁺ = P + rᵤ⁻⁺ - bᵤ⁻⁺
    Q⁻⁺ = Q + rᵥ⁻⁺ - bᵥ⁻⁺

    bᵤ⁻⁻ = zeros(Float64,T^2)
    bᵥ⁻⁻ = zeros(Float64,T^2)
    sᵤ⁻⁻ = Dᵤ⁻*Z - P + bᵤ⁻⁻
    sᵥ⁻⁻ = Dᵥ⁻*Z - Q + bᵥ⁻⁻
    rᵤ⁻⁻ = max(norm(sᵤ⁻⁻) - (4/α),0)*(sᵤ⁻⁻./(norm(sᵤ⁻⁻)+eps(Float64)))
    rᵥ⁻⁻ = max(norm(sᵥ⁻⁻) - (4/α),0)*(sᵥ⁻⁻./(norm(sᵥ⁻⁻)+eps(Float64)))
    bᵤ⁻⁻ = bᵤ⁻⁻ + Dᵤ⁻*Z - P - rᵤ⁻⁻
    bᵥ⁻⁻ = bᵥ⁻⁻ + Dᵥ⁻*Z - Q - rᵥ⁻⁻
    P⁻⁻ = P + rᵤ⁻⁻ - bᵤ⁻⁻
    Q⁻⁻ = Q + rᵥ⁻⁻ - bᵥ⁻⁻

    # Calculate bₜᵥ as per eq 60
    bₜᵥ = (transpose(Dᵤ⁺) * P⁺⁺) .+ (transpose(Dᵥ⁺) * Q⁺⁺) .+ (transpose(Dᵤ⁺) * P⁺⁻) .+ (transpose(Dᵥ⁻) * Q⁺⁻) .+ (transpose(Dᵤ⁻) * P⁻⁺) .+ (transpose(Dᵥ⁺) * Q⁻⁺) .+ (transpose(Dᵤ⁻) * P⁻⁻) .+ (transpose(Dᵥ⁻) * Q⁻⁻) .+ (Λ^2)*Z⁰
    for i = 1:max_iter
        cg!(Z, A, bₜᵥ, tol=10.0^-4, Pl = pl)

        # Recalculate b, s, r, P and Q in each combination of {+,-} as per equations 57, 61, 62
        sᵤ⁺⁺ = Dᵤ⁺*Z - P + bᵤ⁺⁺
        sᵥ⁺⁺ = Dᵥ⁺*Z - Q + bᵥ⁺⁺
        rᵤ⁺⁺ = max(norm(sᵤ⁺⁺) - (4/α),0)*(sᵤ⁺⁺./(norm(sᵤ⁺⁺)+eps(Float64)))
        rᵥ⁺⁺ = max(norm(sᵥ⁺⁺) - (4/α),0)*(sᵥ⁺⁺./(norm(sᵥ⁺⁺)+eps(Float64)))
        bᵤ⁺⁺ = bᵤ⁺⁺ + Dᵤ⁺*Z - P - rᵤ⁺⁺
        bᵥ⁺⁺ = bᵥ⁺⁺ + Dᵥ⁺*Z - Q - rᵥ⁺⁺
        P⁺⁺ = P + rᵤ⁺⁺ - bᵤ⁺⁺
        Q⁺⁺ = Q + rᵥ⁺⁺ - bᵥ⁺⁺

        sᵤ⁺⁻ = Dᵤ⁺*Z - P + bᵤ⁺⁻
        sᵥ⁺⁻ = Dᵥ⁻*Z - Q + bᵥ⁺⁻
        rᵤ⁺⁻ = max(norm(sᵤ⁺⁻) - (4/α),0)*(sᵤ⁺⁻./(norm(sᵤ⁺⁻)+eps(Float64)))
        rᵥ⁺⁻ = max(norm(sᵥ⁺⁻) - (4/α),0)*(sᵥ⁺⁻./(norm(sᵥ⁺⁻)+eps(Float64)))
        bᵤ⁺⁻ = bᵤ⁺⁻ + Dᵤ⁺*Z - P - rᵤ⁺⁻
        bᵥ⁺⁻ = bᵥ⁺⁻ + Dᵥ⁻*Z - Q - rᵥ⁺⁻
        P⁺⁻ = P + rᵤ⁺⁻ - bᵤ⁺⁻
        Q⁺⁻ = Q + rᵥ⁺⁻ - bᵥ⁺⁻

        sᵤ⁻⁺ = Dᵤ⁻*Z - P + bᵤ⁻⁺
        sᵥ⁻⁺ = Dᵥ⁺*Z - Q + bᵥ⁻⁺
        rᵤ⁻⁺ = max(norm(sᵤ⁻⁺) - (4/α),0)*(sᵤ⁻⁺./(norm(sᵤ⁻⁺)+eps(Float64)))
        rᵥ⁻⁺ = max(norm(sᵥ⁻⁺) - (4/α),0)*(sᵥ⁻⁺./(norm(sᵥ⁻⁺)+eps(Float64)))
        bᵤ⁻⁺ = bᵤ⁻⁺ + Dᵤ⁻*Z - P - rᵤ⁻⁺
        bᵥ⁻⁺ = bᵥ⁻⁺ + Dᵥ⁺*Z - Q - rᵥ⁻⁺
        P⁻⁺ = P + rᵤ⁻⁺ - bᵤ⁻⁺
        Q⁻⁺ = Q + rᵥ⁻⁺ - bᵥ⁻⁺

        sᵤ⁻⁻ = Dᵤ⁻*Z - P + bᵤ⁻⁻
        sᵥ⁻⁻ = Dᵥ⁻*Z - Q + bᵥ⁻⁻
        rᵤ⁻⁻ = max(norm(sᵤ⁻⁻) - (4/α),0)*(sᵤ⁻⁻./(norm(sᵤ⁻⁻)+eps(Float64)))
        rᵥ⁻⁻ = max(norm(sᵥ⁻⁻) - (4/α),0)*(sᵥ⁻⁻./(norm(sᵥ⁻⁻)+eps(Float64)))
        bᵤ⁻⁻ = bᵤ⁻⁻ + Dᵤ⁻*Z - P - rᵤ⁻⁻
        bᵥ⁻⁻ = bᵥ⁻⁻ + Dᵥ⁻*Z - Q - rᵥ⁻⁻
        P⁻⁻ = P + rᵤ⁻⁻ - bᵤ⁻⁻
        Q⁻⁻ = Q + rᵥ⁻⁻ - bᵥ⁻⁻

        # Recalculate bₜᵥ as per eq 60
        bₜᵥ = (transpose(Dᵤ⁺) * P⁺⁺) .+ (transpose(Dᵥ⁺) * Q⁺⁺) .+ (transpose(Dᵤ⁺) * P⁺⁻) .+ (transpose(Dᵥ⁻) * Q⁺⁻) .+ (transpose(Dᵤ⁻) * P⁻⁺) .+ (transpose(Dᵥ⁺) * Q⁻⁺) .+ (transpose(Dᵤ⁻) * P⁻⁻) .+ (transpose(Dᵥ⁻) * Q⁻⁻) .+ (Λ^2)*Z⁰
    end
    for j in eachindex(index)
        z[index[j]] = Z[j]
    end
    z = z.*mask
    z = rotl90(z)
    return z
end

@doc raw"""
```
NonConvex1(z::AbstractArray, β::Real = 0.5, λ::AbstractArray = fill(10.0^-6, size(z)), mask::AbstractArray = fill(1.0, size(z)), max_iter::Int = 100)
```
The first of two non-convex regularization methods proposed by Aujol, Durou and
Quéau. The same fidelity and regularization terms to minimize as [`Quadratic`](@ref)
are used, but the convexity of ``\Phi`` is sacrificed in order to gain better
behaviour around discontinuities and outliers.
# Output
`NonConvex1()` returns a NonConvex1 integrator which can then be called to run the
non-convex regularization method method on a gradient field.
# Details
The initial minimization problem is the same as in [`TotalVariation`](@ref)
as given below exept ``Phi`` is defined using the given function instead of the
``L_1`` norm.
```math
\begin{aligned}
&min\iint_{(u,v)\in \Omega}||\nabla z(u,v)-\mathbf{g}(u,v)||+\lambda(u,v)\left[z(u,v)-z^0(u,v)\right]^2dudv\\
&\text{where:}\\
&\Phi(s)=log(s^2+\beta^2)
\end{aligned}
```
The problem can then be discretized to form the following functional where ``D_{u,v}^{UV}``
is a ``2\times|\Omega|`` matrix formed from stacking the vectors of the finite
diferences.
```math
\begin{aligned}
E(\mathbf{z})&=\frac{1}{4}\sum_{(U,V)\in\{+,-\}^2}\sum_{(u,v)\in\Omega^{UV}}\Phi(||\nabla z_{u,v}-\mathbf{g}_{u,v}||)+\sum_{(u,v)\in\Omega}\lambda_{u,v}\left[z_{u,v}-z^0_{u,v}\right]^2\\
&=\frac{1}{4}\sum_{(U,V)\in\{+,-\}^2}\sum_{(u,v)\in\Omega^{UV}}\Phi(||D_{u,v}^{UV}\mathbf{z}-\mathbf{g}_{u,v}||)+||\Lambda\left(\mathbf{z}-\mathbf{z}^0\right)||^2\\
&=f(\mathbf{z})+g(\mathbf{z})
\end{aligned}
```
As ``f(\mathbf{z})`` is smooth but not convex and ``g(\mathbf{z})`` is convex
the iPiano algorithm is then used to iterativly solve the minimization of the
functional such that:
```math
\mathbf{z}^{(k+1)}=(I+\alpha_1\partial g)^{-1}(\mathbf{z}^{(k)}-\alpha_1\nabla f(\mathbf{z}^{(k)})+\alpha_2(\mathbf{z}^{(k)}-z^{(k+1)}))
```
where ``(I+\alpha_1\partial g)^{-1}`` is a proximal operator defined as:
```math
(I+\alpha_1\partial g)^{-1}(\mathbf{\hat{x}})=(I+2\alpha_1\Lambda^2)^{-1}(\mathbf{\hat{x}}+2\alpha_1\Lambda\mathbf{z}^0)
```
Using the definition of ``\Phi`` given above, the derivative of ``f(\mathbf{z})``
can be computed as below:
```math
\nabla f(\mathbf{z})=\frac{1}{4}\sum_{(U,V)\in\{+,-\}^2}\sum_{(u,v)\in\Omega^{UV}}\frac{D_{u,v}^{UV^\top}(D_{u,v}^{UV}\mathbf{z}-\mathbf{g}_{u,v})}{||D_{u,v}^{UV}\mathbf{z}-\mathbf{g}_{u,v}||^2+\beta^2}
```
The values of ``\alpha_1`` and ``\aplha_2`` control the step size of the algorith
and are chosen such that ``\alpha_2`` is fixed at 0.8 while ``\alpha_1`` is chosen
using the lazy backtracking method, which uses a Lipschitz constant to caculate
a suitably small step size at each step using the below relationship where ``\eta>1``
is a constant:
```math
\begin{aligned}
&\alpha_1 < 2(1-\alpha_2)/L_n\\
&\text{where:}\\
&L_k\in\{L_{k-1},\eta L_{k-1},\eta^2L_{k-1},...\}\\
&\text{such that it is minimal and satisfies:}\\
&f(x^{(k+1)})\le f(x^{(k)}) + \left\langle\nabla f(x^{(k)}), x^{(k+1)}-x^{(k)}\right\rangle+\frac{L_k}{2}||x^{(k+1)}-x^{(k)}||^2
\end{aligned}
```
This method, being non-convex, highly relies on a good initial solution and will
often only provide a minimal improvment to the solution. A bad initial solution
will produce a final solution which does not resemble the surface under most
conditions.
# Parameters
## `z`:
An `AbstractArray` which defines the value of ``z^0`` the initial solution and
prior to be used in the regulization term.
## `β`:
A `Real` which acts as a hyper-parameter to the function. Large values for β will
produce smoother functions but loose discontinuities. Smaller value will preserve
discontinuities but lead to staircassing in the solution. Defults to `0.5`.
## `λ`:
An `AbstractArray` the same size as `z` defulting to ``10.0^-6`` everywhere.
This defines theregulization weight at each point. Large valueas will force the
algorithm to keep the solution near to ``z^0`` at that position. Can be used to
keep the solution near the initial solution or guide the solution to a certian
known value at points (i.e. known maxima and minima). This value should be set
uniformly small otherwise.
## `mask`:
An `AbstractArray` the same size as `z`, which guides the algorithm as to where
the valid domain is. Values of `1` will be in the domain ``\Omega`` while other
values will be ignored and set to ``z^0``. This can be used to
integrate over sub-domain or to segment the domain into parts. The `gen_mask()`
funtion can be used to generate a mask which will remove non-integrable regions
dramatically improving the solution under most condition at the cost of not
integrating the entire solution.
# Example
The following example demonstrates the use of the `NonConvex1` integrator.
```julia
using ShapeFromShading, Makie

# Generate synthetic gradients
p, q = synthetic_gradient(Prism(75), img_size = 151)

# Create a NonConvex1() integrator
nonConvex1 = NonConvex1(z=zeros(size(p)), β=0.5)
nonConvex1Init = NonConvex1(z=Horn()(p,q), β=0.5)

# Calculate the heightmap from the gradients
Z = nonConvex1(p, q)
Z2 = nonConvex1Init(p, q)

# Normalize to maximum of 1 (not necessary but makes displaying easier)
Z = Z./maximum(Z)
Z2 = Z2./maximum(Z2)

# Display using Makie (Note: Makie can often take several minutes first time)
r = 0.0:0.1:4
vbox(surface(r, r, Z), surface(r, r, Z2))
```
# Reference
[1] Y. Quéau, J. Durou and J. Aujol, "Variational Methods for Normal Integration", Journal of Mathematical Imaging and Vision, vol. 60, no. 4, pp. 609-632, 2017. [doi: 10.1007/s10851-017-0777-6](https://doi.org/10.1007/s10851-017-0777-6 )
"""
function (scheme::NonConvex1)(pIn::AbstractArray, qIn::AbstractArray)
    z⁰ = copy(scheme.z)
    λ = copy(scheme.λ)
    mask = copy(scheme.mask)
    β = copy(scheme.β)
    max_iter = copy(scheme.max_iter)
    z = copy(z⁰)
    mask = rotr90(mask)
    λ = rotr90(λ)
    p = rotr90(copy(pIn)).*mask
    q = rotr90(copy(qIn)).*mask
    T = first(size(z))
    index = CartesianIndices(z)
    index = reshape(index, length(index),1)

    Dᵤ⁺ = spzeros(Float64,T^2,T^2)
    Dᵤ⁻ = spzeros(Float64,T^2,T^2)
    Dᵥ⁺ = spzeros(Float64,T^2,T^2)
    Dᵥ⁻ = spzeros(Float64,T^2,T^2)
    gen_matrix(Dᵤ⁺, Dᵤ⁻, Dᵥ⁺, Dᵥ⁻, T, mask)
    Λ = zeros(Float64,T^2,T^2)
    P = zeros(Float64,T^2)
    Q = zeros(Float64,T^2)
    Z⁰ = zeros(Float64,T^2)
    Z = zeros(Float64,T^2)

    # Vectorize inputs
    for j in CartesianIndices(z)
        Λ[j[1]+(j[2]-1)*T,j[1]+(j[2]-1)*T] = sqrt(λ[j])
        P[j[1]+(j[2]-1)*T] = p[j]
        Q[j[1]+(j[2]-1)*T] = q[j]
        Z⁰[j[1]+(j[2]-1)*T] = z⁰[j]
        Z[j[1]+(j[2]-1)*T] = z[j]
    end
    Λ = Diagonal(Λ)
    Zold = copy(Z)
    Ztemp = copy(Z)
    α₁ = 0.8
    α₂ = 0.8

    # Initialize Lₙ (Lipschitz constant) such that lazy backtracking starts with α₁ = α₂.
    # This is mostly arbitary and value will approach needed value over time.
    Lₙ = (18.0/α₁)*(1-α₂)
    η = 1.1
    i = 1
    while i <= max_iter && (norm(Z - Zold) > 10.0^-4.0 || i == 1)
        # Calculate f₁, ∇f₁ and previous f₁ (f₁ⁿ⁻¹ )
        ∇f₁ = ((transpose(Dᵤ⁺)*(Dᵤ⁺*Z-P))/(norm(Dᵤ⁺*Z-P)^2+β^2)) + ((transpose(Dᵥ⁺)*(Dᵥ⁺*Z-Q))/(norm(Dᵥ⁺*Z-Q)^2+β^2)) + ((transpose(Dᵤ⁻)*(Dᵤ⁻*Z-P))/(norm(Dᵤ⁻*Z-P)^2+β^2)) + ((transpose(Dᵥ⁻)*(Dᵥ⁻*Z-Q))/(norm(Dᵥ⁻*Z-Q)^2+β^2))
        f₁ⁿ⁻¹ = log(norm(Dᵤ⁺*Zold-P)^2+β^2) + log(norm(Dᵥ⁺*Zold-Q)^2+β^2) + log(norm(Dᵤ⁻*Zold-P)^2+β^2) + log(norm(Dᵥ⁻*Zold-Q)^2+β^2)
        f₁ = log(norm(Dᵤ⁺*Z-P)^2+β^2) + log(norm(Dᵥ⁺*Z-Q)^2+β^2) + log(norm(Dᵤ⁻*Z-P)^2+β^2) + log(norm(Dᵥ⁻*Z-Q)^2+β^2)

        # Calculate lazy backtracking condition as per ipiano algortihm 4
        while f₁ > f₁ⁿ⁻¹  + dot(∇f₁, Z - Zold) + (Lₙ/2)*norm(Z - Zold)
            Lₙ = Lₙ*η
        end

        # Calculate new α₁ as per ipiano algortihm 4
        α₁ = (2*(1-α₂)/Lₙ)*0.9

        # Calcualte x̂ as per eq 69 then calculate new Z as per eq 71
        x̂ = Z - α₁.*∇f₁ + α₂.*(Z-Zold)
        Ztemp = inv(I+(2*α₁).*Λ^2)*(x̂+(2*α₁).*Λ*Z⁰)
        copyto!(Zold, Z)
        copyto!(Z, Ztemp)
        i += 1
    end
    println("Converged after ", i-1, " iterations")
    for j in eachindex(index)
        z[index[j]] = Z[j]
    end
    z = rotl90(z)
    return z
end

@doc raw"""
```
NonConvex2(z::AbstractArray, γ::Real = 1.0, λ::AbstractArray = fill(10.0^-6, size(z, mask::AbstractArray = fill(1.0, size(z)), max_iter::Int = 100)
```
The second of two non-convex regularization methods proposed by Aujol, Durou and
Quéau. The same as [`NonConvex1`](@ref) exept ``\Phi`` has been replaced with a
different non-convex function.
# Output
`NonConvex2()` returns a NonConvex2 integrator which can then be called to run the
non-convex regularization method method on a gradient field.
# Details
The implimentation is the same as the one in [`TotalVariation`](@ref)
as given below exept ``Phi`` is defined as below:
```math
\Phi(s)=\frac{s^2}{s^2+\gamma^2}
```
Using the definition of ``\Phi`` given above, the derivative of ``f(\mathbf{z})``
can be computed as below:
```math
\nabla f(\mathbf{z})=\frac{1}{4}\sum_{(U,V)\in\{+,-\}^2}\sum_{(u,v)\in\Omega^{UV}}\frac{\gamma^2D_{u,v}^{UV^\top}(D_{u,v}^{UV}\mathbf{z}-\mathbf{g}_{u,v})}{()||D_{u,v}^{UV}\mathbf{z}-\mathbf{g}_{u,v}||^2+\gamma^2)^2}
```
This method, being non-convex, highly relies on a good initial solution and will
often only provide a minimal improvment to the solution. A bad initial solution
will produce a final solution which does not resemble the surface under most
conditions.
# Parameters
## `z`:
An `AbstractArray` which defines the value of ``z^0`` the initial solution and
prior to be used in the regulization term.
## `γ`:
A `Real` which acts as a hyper-parameter to the function. Large values for γ will
produce smoother functions but loose discontinuities. Smaller value will preserve
discontinuities but lead to staircassing in the solution. Defults to `1.0`.
## `λ`:
An `AbstractArray` the same size as `z` defulting to ``10.0^-6`` everywhere.
This defines theregulization weight at each point. Large valueas will force the
algorithm to keep the solution near to ``z^0`` at that position. Can be used to
keep the solution near the initial solution or guide the solution to a certian
known value at points (i.e. known maxima and minima). This value should be set
uniformly small otherwise.
## `mask`:
An `AbstractArray` the same size as `z`, which guides the algorithm as to where
the valid domain is. Values of `1` will be in the domain ``\Omega`` while other
values will be ignored and set to ``z^0``. This can be used to
integrate over sub-domain or to segment the domain into parts. The `gen_mask()`
funtion can be used to generate a mask which will remove non-integrable regions
dramatically improving the solution under most condition at the cost of not
integrating the entire solution.
# Example
The following example demonstrates the use of the `NonConvex2` integrator.
```julia
using ShapeFromShading, Makie

# Generate synthetic gradients
p, q = synthetic_gradient(Prism(75), img_size = 151)

# Create a NonConvex2() integrator
nonConvex2 = NonConvex2(z=zeros(size(p)), γ=1.0)
nonConvex2Init = NonConvex2(z=Horn()(p,q), γ=1.0)

# Calculate the heightmap from the gradients
Z = nonConvex2(p, q)
Z2 = nonConvex2Init(p, q)

# Normalize to maximum of 1 (not necessary but makes displaying easier)
Z = Z./maximum(Z)
Z2 = Z2./maximum(Z2)

# Display using Makie (Note: Makie can often take several minutes first time)
r = 0.0:0.1:4
vbox(surface(r, r, Z), surface(r, r, Z2))
```
# Reference
[1] Y. Quéau, J. Durou and J. Aujol, "Variational Methods for Normal Integration", Journal of Mathematical Imaging and Vision, vol. 60, no. 4, pp. 609-632, 2017. [doi: 10.1007/s10851-017-0777-6](https://doi.org/10.1007/s10851-017-0777-6 )
"""
function (scheme::NonConvex2)(pIn::AbstractArray, qIn::AbstractArray)
    z⁰ = copy(scheme.z)
    λ = copy(scheme.λ)
    mask = copy(scheme.mask)
    γ = copy(scheme.γ)
    max_iter = copy(scheme.max_iter)
    z = copy(z⁰)
    mask = rotr90(mask)
    λ = rotr90(λ)
    p = rotr90(copy(pIn)).*mask
    q = rotr90(copy(qIn)).*mask
    T = first(size(z))
    index = CartesianIndices(z)
    index = reshape(index, length(index),1)

    Dᵤ⁺ = spzeros(Float64,T^2,T^2)
    Dᵤ⁻ = spzeros(Float64,T^2,T^2)
    Dᵥ⁺ = spzeros(Float64,T^2,T^2)
    Dᵥ⁻ = spzeros(Float64,T^2,T^2)
    gen_matrix(Dᵤ⁺, Dᵤ⁻, Dᵥ⁺, Dᵥ⁻, T, mask)
    Λ = zeros(Float64,T^2,T^2)
    P = zeros(Float64,T^2)
    Q = zeros(Float64,T^2)
    Z⁰ = zeros(Float64,T^2)
    Z = zeros(Float64,T^2)

    # Vectorize inputs
    for j in CartesianIndices(z)
        Λ[j[1]+(j[2]-1)*T,j[1]+(j[2]-1)*T] = sqrt(λ[j])
        P[j[1]+(j[2]-1)*T] = p[j]
        Q[j[1]+(j[2]-1)*T] = q[j]
        Z⁰[j[1]+(j[2]-1)*T] = z⁰[j]
        Z[j[1]+(j[2]-1)*T] = z[j]
    end
    Λ = Diagonal(Λ)
    Zold = copy(Z)
    Ztemp = copy(Z)
    α₁ = 0.8
    α₂ = 0.8

    # Initialize Lₙ (Lipschitz constant) such that lazy backtracking starts with α₁ = α₂.
    # This is mostly arbitary and value will approach needed value over time.
    Lₙ = (18.0/α₁)*(1-α₂)
    η = 1.1
    i = 1
    while i <= max_iter && (norm(Z - Zold) > 10.0^-4.0 || i == 1)
        # Calculate f₁, ∇f₁ and previous f₁ (f₁ⁿ⁻¹ )
        ∇f₁ = (((γ^2)*transpose(Dᵤ⁺)*(Dᵤ⁺*Z-P))/(norm(Dᵤ⁺*Z-P)^2+γ^2)^2) + (((γ^2)*transpose(Dᵥ⁺)*(Dᵥ⁺*Z-Q))/(norm(Dᵥ⁺*Z-Q)^2+γ^2)^2) + (((γ^2)*transpose(Dᵤ⁻)*(Dᵤ⁻*Z-P))/(norm(Dᵤ⁻*Z-P)^2+γ^2)^2) + (((γ^2)*transpose(Dᵥ⁻)*(Dᵥ⁻*Z-Q))/(norm(Dᵥ⁻*Z-Q)^2+γ^2)^2)
        f₁ⁿ⁻¹ = (norm(Dᵤ⁺*Zold-P)^2)/(norm(Dᵤ⁺*Zold-P)^2+γ^2) + (norm(Dᵥ⁺*Zold-Q)^2)/(norm(Dᵥ⁺*Zold-Q)^2+γ^2) + (norm(Dᵤ⁻*Zold-P)^2)/(norm(Dᵤ⁻*Zold-P)^2+γ^2) + (norm(Dᵥ⁻*Zold-Q)^2)/(norm(Dᵥ⁻*Zold-Q)^2+γ^2)
        f₁ = (norm(Dᵤ⁺*Z-P)^2)/(norm(Dᵤ⁺*Z-P)^2+γ^2) + (norm(Dᵥ⁺*Z-Q)^2)/(norm(Dᵥ⁺*Z-Q)^2+γ^2) + (norm(Dᵤ⁻*Z-P)^2)/(norm(Dᵤ⁻*Z-P)^2+γ^2) + (norm(Dᵥ⁻*Z-Q)^2)/(norm(Dᵥ⁻*Z-Q)^2+γ^2)

        # Calculate lazy backtracking condition as per ipiano algortihm 4
        while f₁ > f₁ⁿ⁻¹ + dot(∇f₁, Z - Zold) + (Lₙ/2)*norm(Z - Zold)
            Lₙ = Lₙ*η
        end

        # Calculate new α₁ as per ipiano algortihm 4
        α₁ = (2*(1-α₂)/Lₙ)*0.9

        # Calcualte x̂ as per eq 69 then calculate new Z as per eq 71
        x̂ = Z - α₁.*∇f₁ + α₂.*(Z-Zold)
        Ztemp = inv(I+(2*α₁).*Λ^2)*(x̂+(2*α₁).*Λ*Z⁰)
        copyto!(Zold, Z)
        copyto!(Z, Ztemp)
        i += 1
    end
    println("Converged after ", i-1, " iterations")
    for j in eachindex(index)
        z[index[j]] = Z[j]
    end
    z = rotl90(z)
    return z
end

@doc raw"""
```
AnisotropicDiffusion(z::AbstractArray, μ::Real = 5.0, ν::Real = 10., λ::AbstractArray = fill(10.0^-6, size(z)), mask::AbstractArray = fill(1.0, size(z)), max_iter::Int = 10)
```
Defines the anisotropic diffusion method as proposed by Aujol, Durou and Quéau.
It utilizes an anisotropic diffusion like process to create a weighted version
of the least squares problem solved in the [`Quadratic`](@ref) method.
# Output
`AnisotropicDiffusion()` returns a AnisotropicDiffusion integrator which can then be called to run the
anisotropic diffusion method method on a gradient field.
# Details
The minimization problem from [`Quadratic`](@ref) is modified with the addition
weighting term ``W(u,v)`` to reach the following minimization problem.
```math
min\iint_{(u,v)\in \Omega}||W(u,v)[\nabla z(u,v)-\mathbf{g}(u,v)]||^2+\lambda(u,v)\left[z(u,v)-z^0(u,v)\right]^2dudv
```
The weighting term can be defined as below where ``\mu`` and ``\nu`` are parameters
which control the impact of gradient of the reconstruction at each iteration and
the input gradients.
```math
W(u,v)=\frac{1}{\sqrt{\left(\frac{||\nabla z(u,v)||}{\mu}\right)^2+1}}\begin{bmatrix}\frac{1}{\sqrt{1+\left(\frac{p(u,v)}{\nu}\right)^2}} &0\\ 0 &\frac{1}{\sqrt{1+\left(\frac{q(u,v)}{\nu}\right)^2}}\end{bmatrix}
```
This lends itself to the creation of the terms ``A^{UV}`` and ``B^{UV}`` which
are ``|\Omega|\times|\Omega|`` diagonal matrices, respectively,  contianing the
values of;
```math
    \frac{1}{\sqrt{1+\left(\frac{p(u,v)}{\nu}\right)^2}\sqrt{\frac{(\partial^U_uz_{u,v})^2+(\partial^V_vz_{u,v})^2}{\mu^2}+1}}
```
and
```math
    \frac{1}{\sqrt{1+\left(\frac{q(u,v)}{\nu}\right)^2}\sqrt{\frac{(\partial^U_uz_{u,v})^2+(\partial^V_vz_{u,v})^2}{\mu^2}+1}}
```
where ``(U,V)\in\{+,-\}^2``. Using these definititions the original minimization
problem can be rewriten as the following iterative scheme:
```math
\begin{aligned}
z^{(k+1)}=&min\frac{1}{4}\sum_{(U,V)\in\{+,-\}^2}\left\{||A^{UV}(D^U_u\mathbf{z}-\mathbf{p})||^2+||B^{UV}(D^V_v\mathbf{z}-\mathbf{q})||^2\right\}\\
&+||\Lambda(\mathbf{z}-\mathbf{z}^0)||
\end{aligned}
```
which is then solved using Cholesky factorization.
# Parameters
## `z`:
An `AbstractArray` which defines the value of ``z^0`` the initial solution and
prior to be used in the regulization term.
## `μ`:
A `Real` which acts as a hyper-parameter to the function. Allows for the tuning
of discontinuities. Defults to `5.0`.
## `ν`:
A `Real` which acts as a hyper-parameter to the function. Allows for the tuning
of discontinuities. Unlike `μ` is has a minimal impact on the solution and generally
does not need adjusting unless an extreme value for `μ` is used where it can help
to balance out the two terms. Defults to `10.0`.
## `λ`:
An `AbstractArray` the same size as `z` defulting to ``10.0^-6`` everywhere.
This defines theregulization weight at each point. Large valueas will force the
algorithm to keep the solution near to ``z^0`` at that position. Can be used to
keep the solution near the initial solution or guide the solution to a certian
known value at points (i.e. known maxima and minima). This value should be set
uniformly small otherwise.
## `mask`:
An `AbstractArray` the same size as `z`, which guides the algorithm as to where
the valid domain is. Values of `1` will be in the domain ``\Omega`` while other
values will be ignored and set to ``z^0``. This can be used to
integrate over sub-domain or to segment the domain into parts. The `gen_mask()`
funtion can be used to generate a mask which will remove non-integrable regions
dramatically improving the solution under most condition at the cost of not
integrating the entire solution.
# Example
The following example demonstrates the use of the `AnisotropicDiffusion` integrator.
```julia
using ShapeFromShading, Makie

# Generate synthetic gradients
p, q = synthetic_gradient(Prism(75), img_size = 151)

# Create a AnisotropicDiffusion() integrator
anisotropicDiffusion = AnisotropicDiffusion(z=zeros(size(p)), γ=1.0)
anisotropicDiffusionInit = AnisotropicDiffusion(z=Horn()(p,q), γ=1.0)

# Calculate the heightmap from the gradients
Z = anisotropicDiffusion(p, q)
Z2 = anisotropicDiffusionInit(p, q)

# Normalize to maximum of 1 (not necessary but makes displaying easier)
Z = Z./maximum(Z)
Z2 = Z2./maximum(Z2)

# Display using Makie (Note: Makie can often take several minutes first time)
r = 0.0:0.1:4
vbox(surface(r, r, Z), surface(r, r, Z2))
```
# Reference
[1] Y. Quéau, J. Durou and J. Aujol, "Variational Methods for Normal Integration", Journal of Mathematical Imaging and Vision, vol. 60, no. 4, pp. 609-632, 2017. [doi: 10.1007/s10851-017-0777-6](https://doi.org/10.1007/s10851-017-0777-6 )
"""
function (scheme::AnisotropicDiffusion)(pIn::AbstractArray, qIn::AbstractArray)
    λ = copy(scheme.λ)
    z⁰ = copy(scheme.z)
    mask = copy(scheme.mask)
    ν = copy(scheme.ν)
    μ = copy(scheme.μ)
    max_iter = copy(scheme.max_iter)
    z = copy(z⁰)
    mask = rotr90(mask)
    λ = rotr90(λ)
    p = rotr90(copy(pIn)).*mask
    q = rotr90(copy(qIn)).*mask
    T = first(size(z))
    index = CartesianIndices(z)
    index = reshape(index, length(index),1)

    Dᵤ⁺ = spzeros(Float64,T^2,T^2)
    Dᵤ⁻ = spzeros(Float64,T^2,T^2)
    Dᵥ⁺ = spzeros(Float64,T^2,T^2)
    Dᵥ⁻ = spzeros(Float64,T^2,T^2)
    gen_matrix(Dᵤ⁺, Dᵤ⁻, Dᵥ⁺, Dᵥ⁻,T,mask)
    Λ = zeros(Float64,T^2,T^2)
    P = zeros(Float64,T^2)
    Q = zeros(Float64,T^2)
    Z⁰ = zeros(Float64,T^2)
    Z = zeros(Float64,T^2)

    # Vectorize inputs
    for j in CartesianIndices(z)
        Λ[j[1]+(j[2]-1)*T,j[1]+(j[2]-1)*T] = sqrt(λ[j])
        P[j[1]+(j[2]-1)*T] = p[j]
        Q[j[1]+(j[2]-1)*T] = q[j]
        Z⁰[j[1]+(j[2]-1)*T] = z⁰[j]
        Z[j[1]+(j[2]-1)*T] = z[j]
    end
    Λ = sparse(Diagonal(Λ))

    # Calcualte A and B for each conination in {+,-} as per eq 90 and 91
    A⁺⁺ = Diagonal(1.0./(sqrt.(1.0.+(P./ν).^2) .* sqrt.(((Dᵤ⁺*Z).^2 .+ (Dᵥ⁺*Z).^2)./(μ^2) .+ 1)))
    A⁻⁺ = Diagonal(1.0./(sqrt.(1.0.+(P./ν).^2) .* sqrt.(((Dᵤ⁻*Z).^2 .+ (Dᵥ⁺*Z).^2)./(μ^2) .+ 1)))
    A⁺⁻ = Diagonal(1.0./(sqrt.(1.0.+(P./ν).^2) .* sqrt.(((Dᵤ⁺*Z).^2 .+ (Dᵥ⁻*Z).^2)./(μ^2) .+ 1)))
    A⁻⁻ = Diagonal(1.0./(sqrt.(1.0.+(P./ν).^2) .* sqrt.(((Dᵤ⁻*Z).^2 .+ (Dᵥ⁻*Z).^2)./(μ^2) .+ 1)))
    B⁺⁺ = Diagonal(1.0./(sqrt.(1.0.+(Q./ν).^2) .* sqrt.(((Dᵤ⁺*Z).^2 .+ (Dᵥ⁺*Z).^2)./(μ^2) .+ 1)))
    B⁻⁺ = Diagonal(1.0./(sqrt.(1.0.+(Q./ν).^2) .* sqrt.(((Dᵤ⁻*Z).^2 .+ (Dᵥ⁺*Z).^2)./(μ^2) .+ 1)))
    B⁺⁻ = Diagonal(1.0./(sqrt.(1.0.+(Q./ν).^2) .* sqrt.(((Dᵤ⁺*Z).^2 .+ (Dᵥ⁻*Z).^2)./(μ^2) .+ 1)))
    B⁻⁻ = Diagonal(1.0./(sqrt.(1.0.+(Q./ν).^2) .* sqrt.(((Dᵤ⁻*Z).^2 .+ (Dᵥ⁻*Z).^2)./(μ^2) .+ 1)))

    # Calculate A and b as derived from eq 92 in form Az=b
    A = (A⁺⁺*Dᵤ⁺)*transpose(A⁺⁺*Dᵤ⁺) + (B⁺⁺*Dᵥ⁺)*transpose(B⁺⁺*Dᵥ⁺) + (A⁺⁻*Dᵤ⁺)*transpose(A⁺⁻*Dᵤ⁺) + (B⁺⁻*Dᵥ⁻)*transpose(B⁺⁻*Dᵥ⁻) + (A⁻⁺*Dᵤ⁻)*transpose( A⁻⁺*Dᵤ⁻) + (B⁻⁺*Dᵥ⁺)*transpose(B⁻⁺*Dᵥ⁺) + (A⁻⁻*Dᵤ⁻)*transpose(A⁻⁻*Dᵤ⁻) + (B⁻⁻*Dᵥ⁻)*transpose(B⁻⁻*Dᵥ⁻)
    A = A/4.0 + Λ^2
    b = (transpose(A⁺⁺*Dᵤ⁺)*A⁺⁺)*P + (transpose(B⁺⁺*Dᵥ⁺)*B⁺⁺)*Q + (transpose(A⁺⁻*Dᵤ⁺)*A⁺⁻)*P + (transpose(B⁺⁻*Dᵥ⁻)*B⁺⁻)*Q + (transpose(A⁻⁺*Dᵤ⁻)*A⁻⁺)*P + (transpose(B⁻⁺*Dᵥ⁺)*B⁻⁺)*Q + (transpose(A⁻⁻*Dᵤ⁻)*A⁻⁻)*P + (transpose(B⁻⁻*Dᵥ⁻)*B⁻⁻)*Q
    b = b/4.0 + (Λ^2)*Z⁰
    A₀ = factorize(A)
    for i = 1:max_iter
        # Solve system at iteration k
        Z = A₀\b

        # Recalcualte A and B for each conination in {+,-} as per eq 90 and 91
        A⁺⁺ = Diagonal(1.0./(sqrt.(1.0.+(P./ν).^2) .* sqrt.(((Dᵤ⁺*Z).^2 .+ (Dᵥ⁺*Z).^2)./(μ^2) .+ 1)))
        A⁻⁺ = Diagonal(1.0./(sqrt.(1.0.+(P./ν).^2) .* sqrt.(((Dᵤ⁻*Z).^2 .+ (Dᵥ⁺*Z).^2)./(μ^2) .+ 1)))
        A⁺⁻ = Diagonal(1.0./(sqrt.(1.0.+(P./ν).^2) .* sqrt.(((Dᵤ⁺*Z).^2 .+ (Dᵥ⁻*Z).^2)./(μ^2) .+ 1)))
        A⁻⁻ = Diagonal(1.0./(sqrt.(1.0.+(P./ν).^2) .* sqrt.(((Dᵤ⁻*Z).^2 .+ (Dᵥ⁻*Z).^2)./(μ^2) .+ 1)))
        B⁺⁺ = Diagonal(1.0./(sqrt.(1.0.+(Q./ν).^2) .* sqrt.(((Dᵤ⁺*Z).^2 .+ (Dᵥ⁺*Z).^2)./(μ^2) .+ 1)))
        B⁻⁺ = Diagonal(1.0./(sqrt.(1.0.+(Q./ν).^2) .* sqrt.(((Dᵤ⁻*Z).^2 .+ (Dᵥ⁺*Z).^2)./(μ^2) .+ 1)))
        B⁺⁻ = Diagonal(1.0./(sqrt.(1.0.+(Q./ν).^2) .* sqrt.(((Dᵤ⁺*Z).^2 .+ (Dᵥ⁻*Z).^2)./(μ^2) .+ 1)))
        B⁻⁻ = Diagonal(1.0./(sqrt.(1.0.+(Q./ν).^2) .* sqrt.(((Dᵤ⁻*Z).^2 .+ (Dᵥ⁻*Z).^2)./(μ^2) .+ 1)))

        # Recalculate A and b as derived from eq 92 in form Az=b
        A = (A⁺⁺*Dᵤ⁺)*transpose(A⁺⁺*Dᵤ⁺) + (B⁺⁺*Dᵥ⁺)*transpose(B⁺⁺*Dᵥ⁺) + (A⁺⁻*Dᵤ⁺)*transpose(A⁺⁻*Dᵤ⁺) + (B⁺⁻*Dᵥ⁻)*transpose(B⁺⁻*Dᵥ⁻) + (A⁻⁺*Dᵤ⁻)*transpose( A⁻⁺*Dᵤ⁻) + (B⁻⁺*Dᵥ⁺)*transpose(B⁻⁺*Dᵥ⁺) + (A⁻⁻*Dᵤ⁻)*transpose(A⁻⁻*Dᵤ⁻) + (B⁻⁻*Dᵥ⁻)*transpose(B⁻⁻*Dᵥ⁻)
        A = A/4.0 + Λ^2
        b = (transpose(A⁺⁺*Dᵤ⁺)*A⁺⁺)*P + (transpose(B⁺⁺*Dᵥ⁺)*B⁺⁺)*Q + (transpose(A⁺⁻*Dᵤ⁺)*A⁺⁻)*P + (transpose(B⁺⁻*Dᵥ⁻)*B⁺⁻)*Q + (transpose(A⁻⁺*Dᵤ⁻)*A⁻⁺)*P + (transpose(B⁻⁺*Dᵥ⁺)*B⁻⁺)*Q + (transpose(A⁻⁻*Dᵤ⁻)*A⁻⁻)*P + (transpose(B⁻⁻*Dᵥ⁻)*B⁻⁻)*Q
        b = b/4.0 + (Λ^2)*Z⁰
        A₀ = factorize(A)
    end

    for j in eachindex(index)
        z[index[j]] = Z[j]
    end
    z = z.*mask
    z = rotl90(z)
    return z
end

@doc raw"""
```
MumfordShah(z::AbstractArray, μ::Real = 10.0, ϵ::Real = 0.1, λ::AbstractArray = fill(10.0^-6, size(z, mask::AbstractArray = fill(1.0, size(z)), max_iter::Int = 50)
```
Defines the Mumford-Shah method, the final method proposed by Aujol, Durou and
Quéau. It utilizes an adepted version of Mumford and Shah functional to provide
a good reconstruction in the presence of discontinuities at the const of a longer
runtime then other methods.
!!! warning
    This function can take several minutes to run on grids larger then 64X64.
# Output
`MumfordShah()` returns a MumfordShah integrator which can then be called to run the
Mumford-Shah method method on a gradient field.
# Details
This method involves modifying the minimization problem given by the Mumford and
Shah functional given below:
```math
min\;\mu\iint_{(u,v)\in\Omega\backslash K}||\nabla z(u,v)||dudv+\int_Kd\sigma\\+\lambda\iint_{(u,v)\in\Omega\backslash K}[z(u,v)-z^0(u,v)]^2dudv
```
Where ``K`` is the set of discontinuities. This can then be adapted to our problem
by utilizing a Ambrosio-Tortelli approximation to achaive the following functional
where ``\epsilon\to0``.
```math
\begin{aligned}
E(z)=\mu&\iint_{(u,v)\in\Omega}w(u,v)^2||\nabla z(u,v)-\mathbf{g}(u,v)||^2dudv\\
+&\iint_{(u,v)}\left[\epsilon||\nabla w(u,v)||^2+\frac{1}{4\epsilon}(w(u,v)-1)^2\right]dudv\\
+&\iint_{(u,v)}\lambda(u,v)\left[z(u,v)-z^0(u,v)\right]dudv
\end{aligned}
```
This leads to the following discretization of the functional where ``\mathbf{w}^{+/-}_{u/v}``
are thevecotr of weights in each forward and backward direction and the matrix
``W^{+/-}_{u/v}`` is the diagonal matrix formed from this vector.
```math
\begin{aligned}
E(\mathbf{z},\mathbf{w}^{+}_{u},\mathbf{w}^{-}_{u},\mathbf{w}^{+}_{v},\mathbf{w}^{-}_{v})&=\frac{\mu}{2}\Big(||W^+_u(D^+_u\mathbf{z}-\mathbf{p})||^2+||W^-_u(D^-_u\mathbf{z}-\mathbf{p})||^2\\
&+||W^+_v(D^+_v\mathbf{z}-\mathbf{q})||^2+||W^-_v(D^-_v\mathbf{z}-\mathbf{q})||^2\Big)\\
&+\frac{\epsilon}{2}\Big(||D^+_u\mathbf{w}^{+}_{u}||^2+||D^-_u\mathbf{w}^{-}_{u}||^2+||D^+_v\mathbf{w}^{+}_{v}||^2\\
&+||D^-_v\mathbf{w}^{-}_{v}||^2\Big)+\frac{1}{8\epsilon}\Big(||\mathbf{w}^{+}_{u}-\mathbf{1}||^2+||\mathbf{w}^{-}_{u}-\mathbf{1}||^2\\
&+||\mathbf{w}^{+}_{v}-\mathbf{1}||^2+||\mathbf{w}^{-}_{v}-\mathbf{1}||^2\Big)+||\Lambda(\mathbf{z}-\mathbf{z}^0)||^2
\end{aligned}
```
This is then solved with a conjugate gradient algorithm and an alternating optimization
scheme at each step where the updates are found using the relationships below:
```math
\begin{gathered}
\mathbf{z}^{(k+1)}=min\;E(\mathbf{z}^{(k)},\mathbf{w}^{+(k)}_{u},\mathbf{w}^{-(k)}_{u},\mathbf{w}^{+(k)}_{v},\mathbf{w}^{-(k)}_{v})\\
\mathbf{w}^{+/-(k+1)}_{u/v}=min\;E(\mathbf{z}^{(k+1},\mathbf{w}^{+(k)}_{u},\mathbf{w}^{-(k)}_{u},\mathbf{w}^{+(k)}_{v},\mathbf{w}^{-(k)}_{v})\\
\end{gathered}
```
This method can produce good results if the parameters are appropriotly tuned.
If the parameters are too large the solution will suffer heavily from staircasing
artifacts while setting it too small will result in a smooth solution. Even if the
value is chosen correctly the algorithm tends to overfit the final solution to the
discontinuities and they will extend into parts of the solution which are actually
smooth.
# Parameters
## `z`:
An `AbstractArray` which defines the value of ``z^0`` the initial solution and
prior to be used in the regulization term.
## `ϵ`:
A `Real` which acts as a hyper-parameter to the function. Controls how the final
solution will converge. Large values will lead to staircasing while small values
will over-smooth the surface. Must be relativly small to achieve convergence to a
solution. Defults to `0.1`.
## `μ`:
A `Real` which acts as a hyper-parameter to the function. This value controls the
smoothness of the final solution. Large values will lead to staircasing while
small values will lead to over-smoothed solutions. Defults to `10.0`.
## `λ`:
An `AbstractArray` the same size as `z` defulting to ``10.0^-6`` everywhere.
This defines the regulization weight at each point. Large values will force the
algorithm to keep the solution near to ``z^0`` at that position. Can be used to
keep the solution near the initial solution or guide the solution to a certian
known value at points (i.e. known maxima and minima). This value should be set
uniformly small otherwise.
## `mask`:
An `AbstractArray` the same size as `z`, which guides the algorithm as to where
the valid domain is. Values of `1` will be in the domain ``\Omega`` while other
values will be ignored and set to ``z^0``. This can be used to
integrate over sub-domain or to segment the domain into parts. The `gen_mask()`
funtion can be used to generate a mask which will remove non-integrable regions
dramatically improving the solution under most condition at the cost of not
integrating the entire solution.
# Example
The following example demonstrates the use of the `MumfordShah` integrator.
```julia
using ShapeFromShading, Makie

# Generate synthetic gradients
p, q = synthetic_gradient(Prism(75), img_size = 151)

# Create a MumfordShah() integrator
mumfordShah = MumfordShah(z=zeros(size(p)), γ=1.0)
mumfordShahInit = MumfordShah(z=Horn()(p,q), γ=1.0)

# Calculate the heightmap from the gradients
Z = mumfordShah(p, q)
Z2 = mumfordShahInit(p, q)

# Normalize to maximum of 1 (not necessary but makes displaying easier)
Z = Z./maximum(Z)
Z2 = Z2./maximum(Z2)

# Display using Makie (Note: Makie can often take several minutes first time)
r = 0.0:0.1:4
vbox(surface(r, r, Z), surface(r, r, Z2))
```
# Reference
[1] Y. Quéau, J. Durou and J. Aujol, "Variational Methods for Normal Integration", Journal of Mathematical Imaging and Vision, vol. 60, no. 4, pp. 609-632, 2017. [doi: 10.1007/s10851-017-0777-6](https://doi.org/10.1007/s10851-017-0777-6 )
"""
function (scheme::MumfordShah)(pIn::AbstractArray, qIn::AbstractArray)
    z⁰ = copy(scheme.z)
    λ = copy(scheme.λ)
    mask = copy(scheme.mask)
    ϵ = copy(scheme.ϵ)
    μ = copy(scheme.μ)
    max_iter = copy(scheme.max_iter)
    z = copy(z⁰)
    mask = rotr90(mask)
    λ = rotr90(copy(λ))
    p = rotr90(copy(pIn)).*mask
    q = rotr90(copy(qIn)).*mask
    T = first(size(z))
    index = CartesianIndices(z)
    index = reshape(index, length(index),1)

    Dᵤ⁺ = spzeros(Float64,T^2,T^2)
    Dᵤ⁻ = spzeros(Float64,T^2,T^2)
    Dᵥ⁺ = spzeros(Float64,T^2,T^2)
    Dᵥ⁻ = spzeros(Float64,T^2,T^2)
    gen_matrix(Dᵤ⁺, Dᵤ⁻, Dᵥ⁺, Dᵥ⁻,T,mask)
    Λ = zeros(Float64,T^2,T^2)
    P = zeros(Float64,T^2)
    Q = zeros(Float64,T^2)
    Z⁰ = zeros(Float64,T^2)
    Z = zeros(Float64,T^2)
    # Vectorize inputs
    for j in CartesianIndices(z)
        Λ[j[1]+(j[2]-1)*T,j[1]+(j[2]-1)*T] = sqrt(λ[j])
        P[j[1]+(j[2]-1)*T] = p[j]
        Q[j[1]+(j[2]-1)*T] = q[j]
        Z⁰[j[1]+(j[2]-1)*T] = z⁰[j]
        Z[j[1]+(j[2]-1)*T] = z[j]
    end
    Λ = sparse(Diagonal(Λ))
    wᵤ⁺ = fill(1.0, T^2)
    wᵤ⁻ = fill(1.0, T^2)
    wᵥ⁺ = fill(1.0, T^2)
    wᵥ⁻ = fill(1.0, T^2)

    Wᵤ⁺ = Diagonal(wᵤ⁺)
    Wᵤ⁻ = Diagonal(wᵤ⁻)
    Wᵥ⁺ = Diagonal(wᵥ⁺)
    Wᵥ⁻ = Diagonal(wᵥ⁻)

    A1 = μ*(transpose(Dᵤ⁺)*transpose(Wᵤ⁺)*Wᵤ⁺*Dᵤ⁺ + transpose(Dᵤ⁻)*transpose(Wᵤ⁻)*Wᵤ⁻*Dᵤ⁻ + transpose(Dᵥ⁺)*transpose(Wᵥ⁺)*Wᵥ⁺*Dᵥ⁺ + transpose(Dᵥ⁻)*transpose(Wᵥ⁻)*Wᵥ⁻*Dᵥ⁻)
    A1 = A1 + Λ^2
    b = μ*((transpose(Dᵤ⁺)*transpose(Wᵤ⁺)*Wᵤ⁺)*P + (transpose(Dᵤ⁻)*transpose(Wᵤ⁻)*Wᵤ⁻)*P + (transpose(Dᵥ⁺)*transpose(Wᵥ⁺)*Wᵥ⁺)*Q + (transpose(Dᵥ⁻)*transpose(Wᵥ⁻)*Wᵥ⁻)*Q)
    b = b + (Λ^2)*Z⁰

    b₀ =  Diagonal(fill(1.0/(4.0*ϵ), T^2))
    DDᵤ⁺ = (ϵ*transpose(Dᵤ⁺)*Dᵤ⁺+b₀)
    DDᵤ⁻ = (ϵ*transpose(Dᵤ⁻)*Dᵤ⁻+b₀)
    DDᵥ⁺ = (ϵ*transpose(Dᵥ⁺)*Dᵥ⁺+b₀)
    DDᵥ⁻ = (ϵ*transpose(Dᵥ⁻)*Dᵥ⁻+b₀)
    A2 = ((Diagonal(Dᵤ⁻*Z-P)^2)*μ) + DDᵤ⁻
    b2 = fill(1.0/(4.0*ϵ), T^2)
    for i = 1:max_iter
        cg!(Z, A1, b)

        A2 = ((Diagonal(Dᵤ⁺*Z-P)^2)*μ) + DDᵤ⁺
        cg!(wᵤ⁺, A2, b2)
        Wᵤ⁺ = Diagonal(wᵤ⁺)

        A2 = ((Diagonal(Dᵤ⁻*Z-P)^2)*μ) + DDᵤ⁻
        cg!(wᵤ⁻, A2, b2)
        Wᵤ⁻ = Diagonal(wᵤ⁻)

        A2 = ((Diagonal(Dᵥ⁺*Z-Q)^2)*μ) + DDᵥ⁺
        cg!(wᵥ⁺, A2, b2)
        Wᵥ⁺ = Diagonal(wᵥ⁺)

        A2 = ((Diagonal(Dᵥ⁻*Z-Q)^2)*μ) + DDᵥ⁻
        cg!(wᵥ⁻, A2, b2)
        Wᵥ⁻ = Diagonal(wᵥ⁻)

        A1 = μ*(transpose(Dᵤ⁺)*transpose(Wᵤ⁺)*Wᵤ⁺*Dᵤ⁺ + transpose(Dᵤ⁻)*transpose(Wᵤ⁻)*Wᵤ⁻*Dᵤ⁻ + transpose(Dᵥ⁺)*transpose(Wᵥ⁺)*Wᵥ⁺*Dᵥ⁺ + transpose(Dᵥ⁻)*transpose(Wᵥ⁻)*Wᵥ⁻*Dᵥ⁻)
        A1 = A1 + Λ^2
        b = μ*((transpose(Dᵤ⁺)*transpose(Wᵤ⁺)*Wᵤ⁺)*P + (transpose(Dᵤ⁻)*transpose(Wᵤ⁻)*Wᵤ⁻)*P + (transpose(Dᵥ⁺)*transpose(Wᵥ⁺)*Wᵥ⁺)*Q + (transpose(Dᵥ⁻)*transpose(Wᵥ⁻)*Wᵥ⁻)*Q)
        b = b + (Λ^2)*Z⁰
    end

    for j in eachindex(index)
        z[index[j]] = Z[j]
    end
    z = z.*mask
    z = rotl90(z)
    return z
end

# function (scheme::Horn)(p::AbstractArray, q::AbstractArray, μ, ν)
#     @show "new"
#     iter = copy(scheme.max_iter)
#     ϵ = copy(scheme.ϵ)
#     Z = zeros(Float64, size(p))
#     Zᵏ⁺¹ = zeros(Float64, size(p))
#     R, C = size(p)
#     wᵤ⁺ = zeros(Float64, size(p))
#     wᵤ⁻ = zeros(Float64, size(p))
#     wᵥ⁺ = zeros(Float64, size(p))
#     wᵥ⁻ = zeros(Float64, size(p))
#     for i in CartesianIndices(p)
#         u,v=i.I
#         wᵤ⁺[i] = 1.0 / (sqrt((((q[min(u+1,R),v]-q[u,v])-(p[u,min(v+1,C)]-p[u,max(v-1,1)]))/μ)^2 + 1)*sqrt(((Z[min(u+1,R),v]-Z[u,v])^2/(ν^2)+1)))
#         wᵤ⁻[i] = 1.0 / (sqrt((((q[u,v]-q[max(u-1,1),v])-(p[u,min(v+1,C)]-p[u,max(v-1,1)]))/μ)^2 + 1)*sqrt(((Z[u,v]-Z[max(u-1,1),v])^2/(ν^2)+1)))
#         wᵥ⁺[i] = 1.0 / (sqrt((((q[min(u+1,R),v]-q[max(u-1,1),v])-(p[u,min(v+1,C)]-p[u,v]))/μ)^2 + 1)*sqrt(((Z[u,min(v+1,C)]-Z[u,v])^2/(ν^2)+1)))
#         wᵥ⁻[i] = 1.0 / (sqrt((((q[min(u+1,R),v]-q[max(u-1,1),v])-(p[u,v]-p[u,max(v-1,1)]))/μ)^2 + 1)*sqrt(((Z[u,v]-Z[u,max(v-1,1)])^2/(ν^2)+1)))
#     end
#     H = zeros(Float64, size(p))
#     V = zeros(Float64, size(p))
#     hv = zeros(Float64, size(p))
#     for i = 1:R
#         Z[i,1] = p[i,1]
#         Z[i,R] = p[i,R]
#     end
#     for i = 1:C
#         Z[1,i] = q[1,i]
#         Z[C,i] = q[C,i]
#     end
#     for i = 2:(R-1)
#         for j = 2:(C-1)
#             H[i,j] = (wᵥ⁻[i,j+1]*(p[i,j+1]-p[i,j]) + wᵥ⁺[i,j-1]*(p[i,j]-p[i,j-1])) / 2
#             V[i,j] = (wᵤ⁻[i+1,j]*(q[i+1,j]-q[i,j]) + wᵤ⁺[i-1,j]*(q[i,j]-q[i-1,j])) / 2
#             # h[i,j] = (p[i,j+1] - p[i,j-1]) / 2
#             # v[i,j] = (q[i+1,j] - q[i-1,j]) / 2
#             hv[i,j] = ϵ * (H[i,j] - V[i,j])
#             # hv[i,j] = hv[i,j]
#         end
#     end
#     for k = 1:iter
#         for i in CartesianIndices(p)
#             u,v=i.I
#             wᵤ⁺[i] = 1.0 / (sqrt((((q[min(u+1,R),v]-q[u,v])-(p[u,min(v+1,C)]-p[u,max(v-1,1)]))/μ)^2 + 1)*sqrt(((Z[min(u+1,R),v]-Z[u,v])^2/(ν^2)+1)))
#             wᵤ⁻[i] = 1.0 / (sqrt((((q[u,v]-q[max(u-1,1),v])-(p[u,min(v+1,C)]-p[u,max(v-1,1)]))/μ)^2 + 1)*sqrt(((Z[u,v]-Z[max(u-1,1),v])^2/(ν^2)+1)))
#             wᵥ⁺[i] = 1.0 / (sqrt((((q[min(u+1,R),v]-q[max(u-1,1),v])-(p[u,min(v+1,C)]-p[u,v]))/μ)^2 + 1)*sqrt(((Z[u,min(v+1,C)]-Z[u,v])^2/(ν^2)+1)))
#             wᵥ⁻[i] = 1.0 / (sqrt((((q[min(u+1,R),v]-q[max(u-1,1),v])-(p[u,v]-p[u,max(v-1,1)]))/μ)^2 + 1)*sqrt(((Z[u,v]-Z[u,max(v-1,1)])^2/(ν^2)+1)))
#         end
#         copyto!(Zᵏ⁺¹, Z)
#         for i = 2:(R-1)
#             for j = 2:(C-1)
#                 H[i,j] = (wᵥ⁻[i,j+1]*(p[i,j+1]-p[i,j]) + wᵥ⁺[i,j-1]*(p[i,j]-p[i,j-1])) / 2
#                 V[i,j] = (wᵤ⁻[i+1,j]*(q[i+1,j]-q[i,j]) + wᵤ⁺[i-1,j]*(q[i,j]-q[i-1,j])) / 2
#                 hv[i,j] = ϵ * (H[i,j] - V[i,j])
#                 Zᵏ⁺¹[i,j] = (wᵤ⁺[i-1,j]*Z[i-1,j] + wᵤ⁻[i+1,j]*Z[i+1,j] + wᵥ⁺[i,j-1]*Z[i,j-1] + wᵥ⁻[i,j+1]*Z[i,j+1])
#                 Zᵏ⁺¹[i,j] = Zᵏ⁺¹[i,j] - hv[i,j]
#                 Zᵏ⁺¹[i,j] = Zᵏ⁺¹[i,j] / (wᵤ⁻[i+1,j] + wᵤ⁺[i-1,j] + wᵥ⁻[i,j+1] + wᵥ⁺[i,j-1])
#             end
#         end
#         Z = Zᵏ⁺¹
#     end
#     return Z
# end


text:
# 2bit format
# ===========
#
# Reader and writer of the 2bit file format.
#
# This file is a part of BioJulia.
# License is MIT: https://github.com/BioJulia/Bio.jl/blob/master/LICENSE.md

"""
The 2bit file format stores multiple DNA sequences (up to 4 Gbp total) as a
compact randomly-accessible format.
See https://genome.ucsc.edu/FAQ/FAQformat.html#format7 for the details.
"""
immutable TwoBit <: FileFormat end

type TwoBitParser{T<:IO} <: AbstractParser
    # input stream
    input::T

    # sequence names
    names::Vector{AbstractString}

    # file offsets
    offsets::Vector{UInt32}

    # byte-swapped or not
    swapped::Bool
end

function TwoBitParser(input::IO)
    swapped, seqcount = read_2bit_header(input)
    names, offsets = read_2bit_index(input, swapped, seqcount)
    @assert seqcount == length(names) == length(offsets)
    return TwoBitParser(input, names, offsets, swapped)
end

function Base.open(filename::AbstractString, ::Type{TwoBit})
    return TwoBitParser(open(filename))
end

function Base.close(p::TwoBitParser)
    close(p.input)
end

function Base.eof(p::TwoBitParser)
    return eof(p.input)
end

function Base.eltype(::Type{TwoBitParser})
    return SeqRecord{ReferenceSequence,Vector{UnitRange{Int}}}
end

function Base.length(p::TwoBitParser)
    return length(p.names)
end

function Base.start(p::TwoBitParser)
    return 1
end

function Base.done(p::TwoBitParser, k)
    return k > endof(p)
end

function Base.next(p::TwoBitParser, k)
    return p[k], k + 1
end

# Random access
# -------------

function Base.checkbounds(p::TwoBitParser, k::Integer)
    if 1 ≤ k ≤ endof(p)
        return true
    end
    throw(BoundsError(p, k))
end

function Base.checkbounds(p::TwoBitParser, r::Range)
    if isempty(r) || (1 ≤ first(r) && last(r) ≤ endof(p))
        return true
    end
    throw(BoundsError(p, k))
end

function Base.endof(p::TwoBitParser)
    return length(p)
end

function Base.getindex(p::TwoBitParser, k::Integer)
    checkbounds(p, k)
    seek(p.input, p.offsets[k])
    seq, mblocks = read_2bit_seq(p.input, p.swapped)
    return SeqRecord{ReferenceSequence,Vector{UnitRange{Int}}}(p.names[k], seq, mblocks)
end

function Base.getindex(p::TwoBitParser, name::AbstractString)
    k = findfirst(p.names, name)
    if k == 0
        throw(KeyError(name))
    end
    return p[k]
end

function Base.getindex(p::TwoBitParser, r::Range)
    checkbounds(p, r)
    return [p[k] for k in r]
end

function Base.getindex{S<:AbstractString}(p::TwoBitParser, names::AbstractVector{S})
    return [p[name] for name in names]
end

# Parser
# ------

function read_2bit_header(input)
    signature = read(input, UInt32)
    if signature == 0x1A412743
        swapped = false
    elseif signature == 0x4327411A
        swapped = true
    else
        error("invalid 2bit signature")
    end
    version  = read32(input, swapped)
    seqcount = read32(input, swapped)
    reserved = read32(input, swapped)
    @assert version == reserved == 0
    return swapped, seqcount
end

function read_2bit_index(input, swapped, seqcount)
    names = AbstractString[]
    offsets = UInt32[]
    for i in 1:seqcount
        namesize = read(input, UInt8)
        name = read(input, UInt8, namesize)
        offset = read32(input, swapped)
        push!(names, bytestring(name))
        push!(offsets, offset)
    end
    return names, offsets
end

function read_2bit_seq(input, swap)
    # read record header
    dnasize = read32(input, swap)
    blockcount = read32(input, swap)  # N blocks
    blockstarts = read32(input, blockcount, swap)
    blocksizes = read32(input, blockcount, swap)
    mblockcount = read32(input, swap)  # masked blocks
    mblockstarts = read32(input, mblockcount, swap)
    mblocksizes = read32(input, mblockcount, swap)
    reserved = read32(input, swap)
    @assert reserved == 0

    # read packed DNAs
    data = zeros(UInt64, cld(dnasize, 32))
    i_stop = BitIndex(dnasize, 2)
    i = BitIndex(1, 2)
    while i ≤ i_stop
        x = read(input, UInt8)
        data[index(i)] |= twobit2refseq(x) << offset(i)
        i += 8
    end

    # make an N vector
    nmask = falses(Int(dnasize))
    for j in 1:blockcount
        s = blockstarts[j]
        for k in s+1:s+blocksizes[j]
            nmask[k] = true
        end
    end

    # make masked blocks
    mblocks = UnitRange{Int}[]
    for j in 1:mblockcount
        s = mblockstarts[j]
        push!(mblocks, s+1:s+mblocksizes[j])
    end

    return ReferenceSequence(data, nmask, 1:dnasize), mblocks
end

# mapping table from .2bit DNA encoding to Bio.jl DNA encoding
const twobit2refseq_table = let
    # T: 00, C: 01, A: 10, G: 11
    f(x) = x == 0b00 ? UInt64(DNA_T) :
           x == 0b01 ? UInt64(DNA_C) :
           x == 0b10 ? UInt64(DNA_A) :
           x == 0b11 ? UInt64(DNA_G) : error()
    tcag = 0b00:0b11
    tbl = UInt64[]
    for a in tcag, b in tcag, c in tcag, d in tcag
        x::UInt64 = 0
        x |= f(a) << 0
        x |= f(b) << 2
        x |= f(c) << 4
        x |= f(d) << 6
        push!(tbl, x)
    end
    tbl
end

twobit2refseq(x::UInt8) = twobit2refseq_table[Int(x)+1]

# read 32 bits and swap bytes if necessary
function read32(input, swap)
    x = read(input, UInt32)
    if swap
        x = bswap(x)
    end
    return x
end

function read32(input, n, swap)
    xs = read(input, UInt32, n)
    if swap
        map!(bswap, xs)
    end
    return xs
end

# Writer
# ------

type TwoBitWriter{T<:IO} <: AbstractWriter
    # output stream
    output::T

    # sequence names
    names::Vector{AbstractString}

    # bit vector to check if each sequence is already written or not
    written::BitVector
end

function TwoBitWriter(output::IO, names::AbstractVector)
    writer = TwoBitWriter(
        output,
        convert(Vector{AbstractString}, names),
        falses(length(names)))
    write_header(writer)
    write_index(writer)
    return writer
end

function Base.close(writer::TwoBitWriter)
    if !all(writer.written)
        error("one or more sequences are not written")
    end
    close(writer.output)
end

function write_header(writer::TwoBitWriter)
    output = writer.output
    n = 0
    n += write(output, 0x1A412743)
    n += write(output, UInt32(0))
    n += write(output, UInt32(length(writer.names)))
    n += write(output, UInt32(0))
    return n
end

function write_index(writer::TwoBitWriter)
    output = writer.output
    n = 0
    for name in writer.names
        n += write(output, UInt8(length(name)))
        n += write(output, name)
        n += write(output, UInt32(0))
    end
    return n
end

# Update the file offset of a sequence in the index section.
function update_offset(writer::TwoBitWriter, seqname, seqoffset)
    @assert seqname ∈ writer.names
    output = writer.output
    offset = 16
    for name in writer.names
        offset += sizeof(UInt8) + length(name)
        if name == seqname
            old = position(output)
            seek(output, offset)
            write(output, UInt32(seqoffset))
            seek(output, old)
            return
        end
        offset += sizeof(UInt32)
    end
end

function Base.write(writer::TwoBitWriter, record::SeqRecord)
    i = findfirst(writer.names, record.name)
    if i == 0
        error("sequence \"", record.name, "\" doesn't exist in the writing list")
    elseif writer.written[i]
        error("sequence \"", record.name, "\" is already written")
    end

    output = writer.output
    update_offset(writer, record.name, position(output))

    n = 0
    n += write(output, UInt32(length(record.seq)))
    n += write_n_blocks(output, record.seq)
    n += write_masked_blocks(output, record.metadata)
    n += write(output, UInt32(0))  # reserved bytes
    n += write_twobit_sequence(output, record.seq)

    writer.written[i] = true
    return n
end

function make_n_blocks(seq)
    starts = UInt32[]
    sizes = UInt32[]
    i = 1
    while i ≤ endof(seq)
        nt = seq[i]
        if nt == DNA_N
            start = i - 1  # 0-based index
            push!(starts, start)
            while i ≤ endof(seq) && seq[i] == DNA_N
                i += 1
            end
            push!(sizes, (i - 1) - start)
        elseif isambiguous(nt)
            error("ambiguous nucleotide except N is not supported")
        else
            i += 1
        end
    end
    return starts, sizes
end

function write_n_blocks(output, seq)
    blockstarts, blocksizes = make_n_blocks(seq)
    @assert length(blockstarts) == length(blocksizes)
    n = 0
    n += write(output, UInt32(length(blockstarts)))
    n += write(output, blockstarts)
    n += write(output, blocksizes)
    return n
end

function write_masked_blocks(output, metadata)
    n = 0
    if isa(metadata, Vector{UnitRange{Int}})
        n += write(output, UInt32(length(metadata)))
        for mblock in metadata
            n += write(output, UInt32(first(mblock) - 1))  # 0-based
        end
        for mblock in metadata
            n += write(output, UInt32(length(mblock)))
        end
    elseif metadata === nothing
        n += write(output, UInt32(0))
    else
        error("metadata is not serializable in the 2bit file format")
    end
    return n
end

function write_twobit_sequence(output, seq)
    n = 0
    i = 4
    while i ≤ endof(seq)
        x::UInt8 = 0
        x |= nuc2twobit(seq[i-3]) << 6
        x |= nuc2twobit(seq[i-2]) << 4
        x |= nuc2twobit(seq[i-1]) << 2
        x |= nuc2twobit(seq[i-0]) << 0
        n += write(output, x)
        i += 4
    end
    r = length(seq) % 4
    if r > 0
        let x::UInt8 = 0
            i = endof(seq) - r + 1
            while i ≤ endof(seq)
                x = x << 2 | nuc2twobit(seq[i])
                i += 1
            end
            x <<= (4 - r) * 2
            n += write(output, x)
        end
    end
    return n
end

function nuc2twobit(nt::DNANucleotide)
    return (
        nt == DNA_A ? 0b10 :
        nt == DNA_C ? 0b01 :
        nt == DNA_G ? 0b11 :
        nt == DNA_T ? 0b00 :
        nt == DNA_N ? 0b00 : error())
end


text:
using DataFrames, PrettyTables
using StatsBase, StatsBase.Statistics

"""
    print_decision_strategy(diagram::InfluenceDiagram, Z::DecisionStrategy, state_probabilities::StateProbabilities; show_incompatible_states::Bool = false)

Print decision strategy.

# Arguments
- `diagram::InfluenceDiagram`: Influence diagram structure.
- `Z::DecisionStrategy`: Decision strategy structure with optimal decision strategy.
- `state_probabilities::StateProbabilities`: State probabilities structure corresponding to optimal decision strategy.
- `show_incompatible_states::Bool`: Choice to print rows also for incompatible states.

# Examples
```julia
print_decision_strategy(diagram, Z, S_probabilities)
```
"""
function print_decision_strategy(diagram::InfluenceDiagram, Z::DecisionStrategy, state_probabilities::StateProbabilities; show_incompatible_states::Bool = false)
    probs = state_probabilities.probs

    for (d, I_d, Z_d) in zip(Z.D, Z.I_d, Z.Z_d)
        s_I = vec(collect(paths(diagram.S[I_d])))
        s_d = [Z_d(s) for s in s_I]

        if !isempty(I_d)
            informations_states = [join([String(diagram.States[i][s_i]) for (i, s_i) in zip(I_d, s)], ", ") for s in s_I]
            decision_probs = [ceil(prod(probs[i][s1] for (i, s1) in zip(I_d, s))) for s in s_I]
            decisions = collect(p == 0 ? "--" : diagram.States[d][s] for (s, p) in zip(s_d, decision_probs))
            df = DataFrame(informations_states = informations_states, decisions = decisions)
            if !show_incompatible_states
                 filter!(row -> row.decisions != "--", df)
            end
            pretty_table(df, header = ["State(s) of $(join([diagram.Names[i] for i in I_d], ", "))", "Decision in $(diagram.Names[d])"], alignment=:l)
        else
            df = DataFrame(decisions = diagram.States[d][s_d])
            pretty_table(df, header = ["Decision in $(diagram.Names[d])"], alignment=:l)
        end
    end
end

"""
    print_utility_distribution(U_distribution::UtilityDistribution; util_fmt="%f", prob_fmt="%f")

Print utility distribution.

# Examples
```julia
U_distribution = UtilityDistribution(diagram, Z)
print_utility_distribution(U_distribution)
```
"""
function print_utility_distribution(U_distribution::UtilityDistribution; util_fmt="%f", prob_fmt="%f")
    df = DataFrame(Utility = U_distribution.u, Probability = U_distribution.p)
    formatters = (
        ft_printf(util_fmt, [1]),
        ft_printf(prob_fmt, [2]))
    pretty_table(df; formatters = formatters)
end

"""
    print_state_probabilities(diagram::InfluenceDiagram, state_probabilities::StateProbabilities, nodes::Vector{Name}; prob_fmt="%f")

Print state probabilities with fixed states.

# Examples
```julia
S_probabilities = StateProbabilities(diagram, Z)
print_state_probabilities(S_probabilities, ["R"])
print_state_probabilities(S_probabilities, ["A"])
```
"""
function print_state_probabilities(diagram::InfluenceDiagram, state_probabilities::StateProbabilities, nodes::Vector{Name}; prob_fmt="%f")
    node_indices = [findfirst(j -> j==node, diagram.Names) for node in nodes]
    states_list = diagram.States[node_indices]
    state_sets = unique(states_list)
    n = length(states_list)

    probs = state_probabilities.probs
    fixed = state_probabilities.fixed

    prob(p, state) = if 1≤state≤length(p) p[state] else NaN end
    fix_state(i) = if i∈keys(fixed) string(diagram.States[i][fixed[i]]) else "" end

    for state_set in state_sets
        node_indices2 = filter(i -> diagram.States[i] == state_set, node_indices)
        state_names = diagram.States[node_indices2[1]]
        states = 1:length(state_names)
        df = DataFrame()
        df[!, :Node] = diagram.Names[node_indices2]
        for state in states
            df[!, Symbol("$(state_names[state])")] = [prob(probs[i], state) for i in node_indices2]
        end
        df[!, Symbol("Fixed state")] = [fix_state(i) for i in node_indices2]
        pretty_table(df; formatters = ft_printf(prob_fmt, (first(states)+1):(last(states)+1)))
    end
end

"""
    print_statistics(U_distribution::UtilityDistribution; fmt = "%f")

Print statistics about utility distribution.
"""
function print_statistics(U_distribution::UtilityDistribution; fmt = "%f")
    u = U_distribution.u
    w = ProbabilityWeights(U_distribution.p)
    names = ["Mean", "Std", "Skewness", "Kurtosis"]
    statistics = [mean(u, w), std(u, w, corrected=false), skewness(u, w), kurtosis(u, w)]
    df = DataFrame(Name = names, Statistics = statistics)
    pretty_table(df, formatters = ft_printf(fmt, [2]))
end

"""
    print_risk_measures(U_distribution::UtilityDistribution, αs::Vector{Float64}; fmt = "%f")

Print risk measures.
"""
function print_risk_measures(U_distribution::UtilityDistribution, αs::Vector{Float64}; fmt = "%f")
    VaR = [value_at_risk(U_distribution, α) for α in αs]
    CVaR = [conditional_value_at_risk(U_distribution, α) for α in αs]
    df = DataFrame(α = αs, VaR = VaR, CVaR = CVaR)
    pretty_table(df, formatters = ft_printf(fmt))
end