@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


"""
    is_graphical(degs::Vector{Int})

Check whether the degree sequence `degs` is graphical, according to
[Erdös-Gallai condition](http://mathworld.wolfram.com/GraphicSequence.html).

Time complexity: O(length(degs)^2)
"""
function is_graphical(degs::Vector{Int})
    iseven(sum(degs)) || return false
    n = length(degs)
    cumdeg = cumsum(degs)
    for r=1:n-1
        cond = cumdeg[r] <= r*(r-1) + sum(i->min(r,degs[i]), r+1:n)
        cond || return false
    end
    return true
end

"""
    has_self_loops(g)

Returns true if `g` has any self loops.
"""
has_self_loops(g::AGraphOrDiGraph) = any(v->has_edge(g, v, v), vertices(g))

"""
    num_self_loops(g)

Returns the number of self loops in `g`.
"""
num_self_loops(g::AGraphOrDiGraph) = count(v->has_edge(g, v, v), vertices(g))


using Random, Plots, Distributions, Measures; pyplot()

N = 5*10^3
alpha = 0.05
confLevel = 1 - alpha
z = quantile(Normal(),1-alpha/2)

function randCI(n,p)
    sample = rand(n) .< p
    pHat = sum(sample)/n
    serr = sqrt(pHat*(1-pHat)/n)
    (pHat - z*serr, pHat + z*serr)
end
cover(p,ci) = ci[1] <= p && p <= ci[2]

pGrid = 0.1:0.01:0.9
nGrid = 5:1:50
errs = zeros(length(nGrid),length(pGrid))

for i in 1:length(nGrid)
    for j in 1:length(pGrid)
        Random.seed!(0)
        n, p = nGrid[i], pGrid[j]
        coverageRatio = sum([cover(p,randCI(n,p)) for _ in 1:N])/N
        errs[i,j] = confLevel - coverageRatio
    end
end

default(xlabel = "p", ylabel = "n",
    xticks =([1:5:length(pGrid);], minimum(pGrid):0.05:maximum(pGrid)),
    yticks =([1:5:length(nGrid);], minimum(nGrid):5:maximum(nGrid)))

p1 = heatmap(errs, c=cgrad([:white, :black]))
p2 = heatmap(abs.(errs) .<= 0.04, legend = false, c=cgrad([:black, :white]))
plot(p1,p2, size = (1000,400), margin = 5mm)


## Exercise 5-2
## The function time returns the current Greenwich Mean Time in seconds since “the epoch”, which is an arbitrary time used as a reference point. On UNIX systems, the epoch is 1 January 1970.

## julia> time()
## 1.554236060271472e9

## Write a script that reads the current time and converts it to a time of day in hours, minutes, and seconds, plus the number of days since the epoch.
println("Ans: ")
using Printf

curr_time = time()

seconds = curr_time % 60
h = 60 * 60
minutes = ((curr_time - seconds) % (h))/h
d = 60 * 60 * 24
hours = ((curr_time - (seconds + minutes)) % (d)) / d
days = (curr_time - (seconds + minutes + hours)) / (60 * 60 * 24)

@printf("%d Days, %d Hours, %d Minutes, %d Seconds\n", days, hours, minutes, seconds)

println("End.")


#' ---
#' title: Intro to Julia
#' ---

#' **Originally Contributed by**: Juan Pablo Vielma

#' ## Introduction

#' Since JuMP is embedded in Julia, knowing some basic Julia is important
#' for learning JuMP. This notebook is designed to provide a minimalist
#' crash course in the basics of Julia. You can find resources that provide
#' a more comprehensive introduction to Julia [here](https://julialang.org/learning/).

#' ### How to Print
#' In Julia, we usually use println() to print

println("Hello, World!")

#' ### Basic Data Types
#' Integers

typeof(1 + -2)

#' Floating point numbers

typeof(1.2 - 2.3)

#' There are also some cool things like an irrational representation of π. To make π
#' (and most other greek letters), type \pi and then press [TAB].

π

#+

typeof(π)

#' Julia has native support for complex numbers

typeof(2 + 3im)

#' Double quotes are used for strings

typeof("This is Julia")

#' Unicode is fine in strings

typeof("π is about 3.1415")

#'Julia symbols provide a way to make human readable unique identifiers.

:my_id
typeof(:my_id)

#' ### Arithmetic and Equality Testing

#' Julia is great for math

1 + 1

#'Even math involving complex numbers

(2 + 1im) * (1 - 2im)

#' We can also write things like the following using √ (\sqrt)

sin(2π/3) == √3/2

#' Wait. What???

sin(2π/3) - √3/2

#' Let's try again using ≈ (\approx).

sin(2π/3) ≈ √3/2

#' Note that this time we used ≈ instead of ==. That is because computers don't use
#' real numbers. They use a discrete representation called floating point. If you aren't
#' careful, this can throw up all manner of issues. For example:

1 + 1e-16 == 1

#' It even turns out that floating point numbers aren't associative!

(1 + 1e-16) - 1e-16 == 1 + (1e-16 - 1e-16)

#' ### Vectors, Matrices and Arrays
#' Similar to Matlab, Julia has native support for vectors, matrices and tensors; all of which are represented by arrays of different dimensions.
#' Vectors are constructed by comma-separated elements surrounded by square brackets:

b = [5, 6]

#' Matrices can by constructed with spaces separating the columns, and semicolons separating the rows:

A = [1 2; 3 4]

#' We can do linear algebra:

x = A \ b

#+

A * x

#+

A * x == b

#' Note that when multiplying vectors and matrices, dimensions matter. For example, you can't multiply a vector by a vector:
#+ tangle = false

b * b

#' But multiplying transposes works:

@show b' * b
@show b * b';

#' ### Tuples
#' Julia makes extensive use of a simple data structure called Tuples.  Tuples are immutable collections of values.
#' For example,

t = ("hello", 1.2, :foo)

#+

typeof(t)

#' Tuples can be accessed by index, similar to arrays,

t[2]

#' And can be "unpacked" like so,

a, b, c = t
b

#' The values can also be given names, which is a convenient way of making light-weight data structures.

t = (word="hello", num=1.2, sym=:foo)

#' Then values can be accessed using a dot syntax,

t.word

#' ### Dictionaries
#' Similar to Python, Julia has native support for dictionaries.  Dictionaries provide a very generic way of mapping keys to values.  For example, a map of integers to strings,

d1 = Dict(1 => "A", 2 => "B", 4 => "D")

#' Looking up a values uses the bracket syntax,

d1[2]

#' Dictionaries support non-integer keys and can mix data types,

Dict("A" => 1, "B" => 2.5, "D" => 2 - 3im)

#' Dictionaries can be nested

d2 = Dict("A" => 1, "B" => 2, "D" => Dict(:foo => 3, :bar => 4))

#+

d2["B"]

#+

d2["D"][:foo]

#' ### For-Each Loops
#' Julia has native support for for-each style loops with the syntax `for <value> in <collection> end`.

for i in 1:5
    println(i)
end

#+

for i in [1.2, 2.3, 3.4, 4.5, 5.6]
    println(i)
end

#' This for-each loop also works with dictionaries.

for (key, value) in Dict("A" => 1, "B" => 2.5, "D" => 2 - 3im)
    println("$key: $value")
end

#' Note that in contrast to vector languages like Matlab and R, loops do not result in a significant performance degradation in Julia.

#' ### Control Flow
#' Julia control flow is similar to Matlab, using the keywords `if-elseif-else-end`, and the logical operators `||` and `&&` for *or* and *and* respectively.

i = 10
for i in 0:3:15
    if i < 5
        println("$(i) is less than 5")
    elseif i < 10
        println("$(i) is less than 10")
    else
        if i == 10
            println("the value is 10")
        else
            println("$(i) is bigger than 10")
        end
    end
end

#' ### Comprehensions
#' Similar to languages like Haskell and Python, Julia supports the use of simple loops in the construction of arrays and dictionaries, called comprehenions.
#'
#' A list of increasing integers,

[i for i in 1:5]

#' Matrices can be built by including multiple indices,

[i*j for i in 1:5, j in 5:10]

#' Conditional statements can be used to filter out some values,

[i for i in 1:10 if i%2 == 1]

#' A similar syntax can be used for building dictionaries

Dict("$i" => i for i in 1:10 if i%2 == 1)

#' ### Functions
#' A simple function is defined as follows,

function print_hello()
    println("hello")
end
print_hello()

#' Arguments can be added to a function,

function print_it(x)
    println(x)
end
print_it("hello")
print_it(1.234)
print_it(:my_id)

#' Optional keyword arguments are also possible

function print_it(x; prefix="value:")
    println("$(prefix) $x")
end
print_it(1.234)
print_it(1.234, prefix="val:")

#' The keyword `return` is used to specify the return values of a function.

function mult(x; y=2.0)
    return x * y
end
mult(4.0)

#+

mult(4.0, y=5.0)

#' ### Other notes on types
#' Usually, specifing types is not required to use Julia.  However, it can be helpful to understand the basics of Julia types for debugging.
#' For example this list has a type of `Array{Int64,1}` indicating that it is a one dimensional array of integer values.

[1, 5, -2, 7]

#' In this example, the decimal values lead to a one dimensional array of floating point values, i.e. `Array{Float64,1}`.  Notice that the integer `7` is promoted to a `Float64`, because all elements in the array need share a common type.

[1.0, 5.2, -2.1, 7]

#' ### Mutable vs immutable objects
#' Some types in Julia are *mutable*, which means you can change the values inside them. A good example is an array. You can modify the contents of an array without having to make a new array.
#' In contrast, types like `Float64` are *immutable*. You can't modify the contents of a `Float64`.
#' This is something to be aware of when passing types into functions. For example:

function mutability_example(mutable_type::Vector{Int}, immutable_type::Int)
    mutable_type[1] += 1
    immutable_type += 1
    return
end

mutable_type = [1, 2, 3]
immutable_type = 1

mutability_example(mutable_type, immutable_type)

println("mutable_type: $(mutable_type)")
println("immutable_type: $(immutable_type)")

#' Because `Vector{Int}` is a mutable type, modifying the variable inside the function changed the value outside of the function. In constrast, the change to `immutable_type` didn't modify the value outside the function.
#' You can check mutability with the `isimmutable` function.

@show isimmutable([1, 2, 3])
@show isimmutable(1);

#' ### Using Packages and the Package Manager
#' No matter how wonderful Julia's base language is, at some point you will want to use an extension package.  Some of these are built-in, for example random number generation is available in the `Random` package in the standard library. These packages are loaded with the commands `using` and `import`.

using Random
Random.seed!(33);

[rand() for i in 1:10]

#' The Package Manager is used to install packages that are not part of Julia's standard library.
#' For example the following can be used to install JuMP,
#+ tangle = false

using Pkg
Pkg.add("JuMP")

#' For a complete list of registed Julia packages see the package listing at https://pkg.julialang.org/.
#' From time to you may wish to use a Julia package that is not registered.  In this case a git repository URL can be used to install the package.
#+ tangle = false

using Pkg
Pkg.add("https://github.com/user-name/MyPackage.jl.git")

#' Note that for clarity this example uses the package manager `Pkg`.  Julia includes an interactive package manager that can be accessed using `]`.  [This video](https://youtu.be/76KL8aSz0Sg) gives an overview of using the interactive package manager environment.
#' The state of installed packages can also be saved in two files: `Project.toml` and `Manifest.toml`. If these files are stored in the same directory than a notebook, the state of the packages can be recovered by running
#+ tangle = false
import Pkg
Pkg.activate(@__DIR__)
Pkg.instantiate()

#' ### HELP!
#' Julia includes a help mode that can be accessed using `?`.  Entering any object (e.g. function, type, struct, ...) into the help mode will show its documentation, if any is available.

#' ### Some Common Gotchas

#' #### MethodError
#' A common error in Julia is `MethodError`, which indicates that the function is not defined for the given value.  For example, by default the `ceil` function is not defined for complex numbers.  The "closest candidates" list suggest some Julia types that the function is defined for.
#+ tangle = false

ceil(1.2 + 2.3im)


# poissonreg.jl
#
# Code for poisson regression.

using Distributions;

# (A, b, x) = GeneratePoissonData(m::Int64 = 200, n::Int64 = 20)
#
# Generates an m-by-n matrix A of measurement vectors a_i = A[i, :], where each
# a_i is uniform on the unit sphere. Chooses x uniformly from the sphere of
# radius sqrt(n), and draws b_i as
#
#   b_i ~ Poisson(exp(a_i' * x))
function GeneratePoissonData(m::Int64 = 200, n::Int64 = 20;
                             generate_gaussian::Bool = false,
                             condition_number::Float64 = 1.0)
  A = randn(m, n);
  if (!generate_gaussian)
    A = A ./ sqrt.(repeat(sum(A.^2, dims=2), outer = (1, n)));
  end
  if (condition_number > 1.0)
    # Make A have condition number given
    condition_vec = Vector(LinRange(1.0, condition_number, n));
    condition_vec = condition_vec / norm(condition_vec);
    A = A .* repeat(condition_vec', outer = (m, 1));
  end
  x = randn(n);
  x = sqrt(n) * x / norm(x);
  b = zeros(Int64, m);
  for ii = 1:m
    R = Poisson(exp(dot(A[ii, :], x)));
    b[ii] = rand(R);
  end
  return (A, b, x);
end

# PoissonObj(A::Matrix{Float64}, b::Vector{Int64}, x::Vector{Float64})
#
# Computes the (average) log loss for the poisson regression problem,
# where we recall that the log-likelihood for b ~ Poisson(lambda) is
#
#  p(b | lambda) = exp(-lambda) * lambda^b / b!
#
# For a data vector a, the parameter lambda = exp(a' * x), so we have log loss
# (ignoring the factorial term)
#
# l(x; (a, b)) = -b * dot(a, x) + exp(dot(a, x)).
function PoissonObj(A::Matrix{Float64}, b::Vector{Int64}, x::Vector{Float64})
  mm = size(A, 1);
  objective = -sum(b .* (A * x)) + sum(exp.(A*x)); # no need to add b!
  return objective / mm;
end

# g = FullPoissonGradient(A, b, x)
#
# Sets g to be the gradient of the poisson objective and returns it.
function FullPoissonGradient(A::Matrix{Float64}, b::Vector{Int64},
                             x::Vector{Float64})
  mm = size(A, 1);
  grad = (-A' * b + A' * exp.(A * x)) / mm;
  return grad;
end

# x = ProximalUpdatePoisson(a::Vector{Float64}, b::Float64,
#                              x_init::Vector, stepsize::Float64; accuracy)
#
# Sets x to minimize
#
#  -b * a' * x + exp(x' * a) + norm(x - x_init)^2 / (2 * stepsize).
#
# Returns the minimizing x. The accuracy parameter governs the
# accuracy of the solution: we recommend a line search.
function ProximalUpdatePoisson(a::Vector{Float64}, b::Int64,
                               x_init::Vector{Float64}, stepsize::Float64;
                               eps_accuracy::Float64 = 1e-10)
  # Replace this code to return the correct update
  return x_init;
end

# x = TruncatedUpdatePoisson(a::Vector{Float64}, b::Int64,
#                            x_init::Vector{Float64}, stepsize::Float64)
#
# Let l(x; (a, b)) = -b * a' * x + exp(a' * x). Sets x to minimize the
# lower-bounded approximation to F at x_init, that is, for
#
#  l_approx(x) = max{ l(x_init) + l'(x_init) * (x - x_init)),
#                     inf_x l(x; (a, b)) }
#
# sets x to minimize
#
#  l_approx(x) + norm(x - x_init)^2 / (2 * stepsize).
#
# Returns the minimizing x.
function TruncatedUpdatePoisson(a::Vector{Float64}, b::Int64,
                                x_init::Vector{Float64}, stepsize::Float64)
  # Replace this code to return the correct update
  return x_init;
end

# x = SGUpdatePoisson(a::Vector{Float64}, b::Int64,
#                     x_init::Vector{Float64}, stepsize::Float64)
#
# Computes the stochastic gradient update for Poisson regression.
function SGUpdatePoisson(a::Vector{Float64}, b::Int64,
                         x_init::Vector{Float64}, stepsize::Float64)
  # Replace this code to return the correct update
  return x_init;
end

Machine Learning Summer School 2019 (London, UK)/tutorials/optimization/poissonreg.jl


# граница xэмминга про упаковку шарами
function hammingND(n, d)
  t = div(d - 1, 2)
  binoms = [binomial(n, i) for i in 0:t]
  println("binoms = $binoms")
  z = sum(binoms)
  println("z = $z")
  M = 2^n / z
  return (M, log2(M))
end

function hammingNK(n, k)
  t = 0
  b = 2^(n - k)
  while sum([binomial(n, i) for i in 0:t]) ≤ b
    t += 1
  end
  return t - 1
end

function varshamovGuilbertND(n, d)
  k = 0
  binomSum = sum([binomial(n - 1, i) for i in 0:d-2])
  println("binomSum = $binomSum")

  while 2^(n - k) > binomSum
    k += 1
  end

  return k - 1
end

function varshamovGuilbertNK(n, k)
  d = 2
  b = 2^(n - k)
  while b > sum([binomial(n - 1, i) for i in 0:d-2])
    d += 1
  end
  return d - 1
end

function graismer(k, d)
  n = 0
  for i = 0:k-1
    n += ceil(Int64, d / 2^i)
  end
  return n
end

function graismerD(n, k)
  # n > sum(d/2^i) from 0 to k
  d = 0
  while true
    sums = sum([ceil(Int64, d / 2^i) for i in 0:k-1])
    if n > sums
      d += 1
    else
      break
    end
  end

  return d - 1
end

n, k, d = 21, 13, 14

# наxождение k по n и d
#
# M, klog = hammingND(n, d)
# println("граница Хэмминга: M ≤ $M, k ≈ $(klog)")

# println("==============================")

# mink = varshamovGuilbertND(n, d)
# println("граница Варшамова-Гилберта: k ≥ $mink")

# println("==============================")
# println(graismer(9, d))
# println("###################################")

# наxождение d по n и k
#
t = hammingNK(n, k)
println("граница Хэмминга: t ≤ $t")

println("==============================")

mind = varshamovGuilbertNK(n, k)
println("граница Варшамова-Гилберта: d ≥ $mind")

println("==============================")
println("граница Грайсмера: $(graismerD(n, k))")
println("###################################")


### ntt.jl
#
# Copyright (C) 2019 Jakub Wronowski.
#
# Maintainer: Jakub Wronowski <jakubwro@users.noreply.github.com>
# Keywords: number theoretic transform
#
# This file is a part of NumberTheoreticTransforms.jl.
#
# License is MIT.
#
### Commentary:
#
# This file contains implementation of general Number Theoretic Transform.
#
### Code:

export ntt, intt

"""
Constraints on NTT params to ensure that inverse can be computed
"""
function validate(N, g, q)
    @assert mod(q - 1, N) == 0
    @assert powermod(g, N, q) == 1
    @assert !(1 in powermod.(g, 2:N-1, q)) # this may be redundant
    @assert gcd(N, q) == 1
end

"""
    ntt(x::Array{T,1}, g, q) -> Array{T,1}

The [Number Theoretic Transform](https://en.wikipedia.org/wiki/Discrete_Fourier_transform_(general)#Number-theoretic_transform)
transforms data in a similar fashion to DFT, but instead complex roots of unity
it uses integer roots with all operation defined in a finite field (modulo an
integer number).

`ntt` function implements Number Theoretic Transform directly from the formula,
so it is flexible about choosing transformation params but lacks performance.

``\\bar{x}_k = \\sum_{n=1}^N{x_n g^{(n-1)(k-1)} } \\mod q``

There is also a few constraints on choosing parameters and input length to
ensure that inverse exists and equals to the original input:
-   ``g`` must ``N``-th root of one in modulo ``q`` arithmetic
-   ``q-1`` mod ``N`` must be equal zero
-   ``q`` must be grater than maximum element present in transformed array

To find parameter set you may use fint-ntt script.

The arguments of `ntt` function are

-   `x`: input data, its elements must be smaller than q
-   `g`: transform power base, must have inversion modulo q
-   `q`: defines modulo arithmetic
"""
function ntt(x::Array{T,1}, g::T, q::T) where {T<:Integer}
    N = length(x)
    #TODO: more validation of p,q, decompose it to struct
    #TODO: create transform object that validates input in the constructor

    validate(N, g, q)
    (lo, hi) = extrema(x)
    @assert lo >= 0
    @assert hi <= q-1

    t = [powermod(g, n * k, q) for n in 0:N-1, k in 0:N-1]
    #TODO: make result of ntt a struct that will hold infrmation about g
    return mod.(t * x, q)
end

# TODO: change implementation to handle any number of dimensions
# it can be done calling N-1 dimensional transform in a loop

function ntt(x::Array{T,2}, g::T, q::T) where {T<:Integer}
    N, M = size(x)
    @assert N == M #TODO: make it work for N != M (need different g for each dim)
    y = zeros(T, size(x))

    for n in 1:N
        y[n, :] = ntt(x[n, :], g, q)
    end

    for m in 1:M
        y[:, m] = ntt(y[:, m], g, q)
    end

    return y
end

"""
    intt(y::Array{T,1}, g, q) -> Array{T,1}

Inverse Number Theoretic Transform implementation directly from the formula.

``x_k = N^{-1} \\sum_{n=1}^N{\\bar{x}_n g^{-(n-1)(k-1)} } \\mod q``

The same input parameters constraints as for `ntt` function must be applied
"""
function intt(y::Array{T,1}, g::T, q::T) where {T<:Integer}
    N = length(y)

    validate(N, g, q)

    inv_g = invmod(g, q)
    inv_N = invmod(N, q)
    t = [powermod(inv_g, l * k, q) for k in 0:N-1, l in 0:N-1]
    return mod.(inv_N * t * y, q)
end

function intt(y::Array{T,2}, g::T, q::T) where {T<:Integer}
    N, M = size(y)
    x = zeros(T, size(y))

    for m in 1:M
        x[:, m] = intt(y[:, m], g, q)
    end

    for n in 1:N
        x[n, :] = intt(x[n, :], g, q)
    end

    return x
end

# those implementations are incorrect but transform results were interesing anyway
# TODO: check if those calculations are useful

# function ntt(g::T, q::T, x::Array{T,N}) where {T <: Integer,N}
#     return reshape(ntt(g, q, reshape(x, length(x))), size(x))
# end

# function intt(g::T, q::T, y::Array{T,N}) where {T <: Integer,N}
#     return reshape(intt(g, q, reshape(y, length(y))), size(y))
# end


"""
Area scaling factor for a function f with partial derivatives fx, fy
    * fx::Function
    * fy::Function

Returns:
    * [Function] : Area scaling factor at x.
"""
function af(fx, fy)
    return x -> norm(fx(x)×fy(x))
end

function af(f)
    # Find partial derivatives
    fx(x) = ForwardDiff.derivative(z -> f([z, x[2]]), x[1])
    fy(x) = ForwardDiff.derivative(z -> f([x[1], z]), x[2])

    # Compute area scaling factor
    return x -> norm(fx(x)×fy(x))
end

"""
Derivative of area scaling factor for a function f with derivatives
    * fx::Function
    * fy::Function
    * fxx::Function
    * fxy::Function
    * fyy::Function

Returns:
    * [Function] : Derivative of area scaling factor at x.
"""
function Daf(af, fx, fy, fxx, fxy, fyy)
    return x -> [
        (fxx(x) × fy(x) - fxy(x) × fx(x)) ⋅ (fx(x) × fy(x)) / af(x)
        (fxy(x) × fy(x) - fyy(x) × fx(x)) ⋅ (fx(x) × fy(x)) / af(x)
    ]
end

"""
Get q-map and derivative of a given function
    * f::Function

Returns:
    * q::Function, Dq::Function
"""
function Qmap(f)
    # Find partial derivatives
    fx(x) = ForwardDiff.derivative(z -> f([z, x[2]]), x[1])
    fy(x) = ForwardDiff.derivative(z -> f([x[1], z]), x[2])

    # Compute area scaling factor
    a = af(fx, fy)

    # Return q-map and derivative
    q(x) = √a(x) * f(x)
    Dq(x) = ForwardDiff.jacobian(q, x)
    return q, Dq
end

"""
Get q-map and derivative of a given function with derivatives
    * f::Function
    * fx::Function
    * fy::Function

Returns:
    * q::Function, Dq::Function
"""
function Qmap(f, fx, fy)
    a = af(fx, fy)

    q(x) = √a(x) * f(x)
    Dq(x) = ForwardDiff.jacobian(q, x)
    return q, Dq
end

"""
Get q-map and derivative of a given function with derivatives
    * f::Function
    * fx::Function
    * fy::Function
    * fxx::Function
    * fxy::Function
    * fyy::Function

Returns:
    * q::Function, Dq::Function
"""
function Qmap(f, fx, fy, fxx, fxy, fyy)
    a = af(fx, fy)
    Da = Daf(a, fx, fy, fxx, fxy, fyy)

    Df(x) = [fx(x) fy(x)]
    q(x) = √a(x) * f(x)
    Dq(x) = 1. / (2. * √a(x)) * f(x) * Da(x)' + √a(x) * Df(x)
    return q, Dq
end

"""
Compute the centroid of the parametric surface.
    * f::Function
    * af:: Function

Returns:
    * [Array{Float}] : Vector value of surface centre
"""
function surface_centroid(f, af)
    return hcubature(x -> f(x)*af(x), [0, 0], [1, 1])[1] / hcubature(af, [0, 0], [1, 1])[1]
end

"""
Translate function so that centroid coincides with the origin.
    * f::Function
    * af:: Function

Returns:
    * [Function] : Translated Surface
"""
function centered(f, af)
    c = surface_centroid(f, af)
    return x -> f(x) - c
end

"""
Compute the parametric surface area.
    * f::Function
    * af:: Function

Returns:
    * [Float] : Surface Area
"""
function surface_area(f, af)
    return hcubature(af, [0, 0], [1, 1])[1]
end

"""
Scale surface to unit area.
    * f::Function
    * af:: Function

Returns:
    * [Function] : Scaled Surface
"""
function scale_surface(f, af)
    s = surface_area(f, af)
    return x -> f(x) / s
end

"""
Translate and scale surface
    * f::Function

Returns:
    * [Function] : Scaled Surface
"""
function normalize_surface(f)
    # Find partial derivatives
    fx(x) = ForwardDiff.derivative(z -> f([z, x[2]]), x[1])
    fy(x) = ForwardDiff.derivative(z -> f([x[1], z]), x[2])

    # Compute area scaling factor
    a = af(fx, fy)

    c = surface_centroid(f, a)
    s = surface_area(f, a)
    return x -> (f(x) - c) / s
end

"""
Translate and scale surface
    * f::Function
    * af:: Function

Returns:
    * [Function] : Scaled Surface
"""
function normalize_surface(f, af)
    c = surface_centroid(f, af)
    s = surface_area(f, af)
    return x -> (f(x) - c) / s
end


using Plots; plotlyjs()
using FFTW

# Construct a signal
f = 1
fs = 5
ts = 1 / fs
l = 2^10 + 1
t = (-(l - 1) ÷ 2 : (l - 1) ÷ 2) * ts
x = sinc.(f * t)

# Construct the frequency spectrum
xf = fft(x)
ff = (0 : 1 / (l - 1) : 1) * fs

# Signal energy
Δt = ts
Δf = fs / (l - 1)
Et = sum(abs.(x).^2) * Δt
Ef = sum(abs.(xf).^2) * Δf / (fs^2)
@show Et, Ef

# Plots
plt = plot(layout=2)
plot!(t, x, subplot=1)
plot!(ff, abs.(xf), subplot=2)
display(plt)


# This file is pat of RLinearAlgebra.jl
# 1. Specifies type
# 2. Implements rsubsolve function
# 3. Exports Type

#using LinearAlgebra

"""
    LinSysVecRowProjFO <: LinSysVecRowProjection

A mutable structure that represents the standard row projection method
with full orthogonalization against all previous projections. Equivalently, this type
represents a solver for incrementally constructed matrix sketches.

See McCormick, S. F. "The methods of Kaczmarz and row orthogonalization for solving linear
    equations and least squares problems in Hilbert space." Indiana University Mathematics
    Journal 26.6 (1977): 1137-1150. https://www.jstor.org/stable/24891603

See Patel, Vivak, Mohammad Jahangoshahi, and Daniel Adrian Maldonado. "An Implicit
    Representation and Iterative Solution of Randomly Sketched Linear Systems."
    SIAM Journal on Matrix Analysis and Applications (2021) 42:2, 800-831.
    https://doi.org/10.1137/19M1259481

# Fields

- `S::Union{Matrix{Float64}, Nothing}` is a matrix used for orthogonalizing against
    all previous search directions.

Calling `LinSysVecRowProjFO()` defaults to `LinSysVecRowProjFO(nothing)`.
"""
mutable struct LinSysVecRowProjFO <: LinSysVecRowProjection
    S::Union{Matrix{Float64}, Nothing}
end
LinSysVecRowProjFO() = LinSysVecRowProjFO(nothing)

# Common rsubsolve interface for linear systems
function rsubsolve!(
    type::LinSysVecRowProjFO,
    x::AbstractVector,
    samp::Tuple{U,V} where {U<:AbstractVector,V<:Real},
    iter::Int64,
)
    # samp[1] is vector in the row space of the coefficient matrix
    # samp[2] is a scalar corresponding to samp[1]

    # Allocate space for orthogonalization matrix S
    if iter == 1
        d = length(x)
        type.S = diagm(ones(Float64, d))
    end

    #Compute orthogonal component of q using S (projection matrix)
    u = type.S * samp[1]

    if dot(u,u) < eps()^2 * length(x)
        return nothing
    end

    #Update Iterate
    res = samp[2] - dot(samp[1], x)
    γ = dot(u, samp[1])
    x .= x + u * ( res / γ )

    #Update Projection matrix
    type.S .= (I - (u/γ)*samp[1]')*type.S

    return nothing
end