#
# Conjugate gradient
#
# This is an independent implementation of:
#   W. W. Hager and H. Zhang (2006) Algorithm 851: CG_DESCENT, a
#     conjugate gradient method with guaranteed descent. ACM
#     Transactions on Mathematical Software 32: 113–137.
#
# Code comments such as "HZ, stage X" or "HZ, eqs Y" are with
# reference to a particular point in this paper.
#
# Several aspects of the following have also been incorporated:
#   W. W. Hager and H. Zhang (2012) The limited memory conjugate
#     gradient method.
#
# This paper will be denoted HZ2012 below.
#
# There are some modifications and/or extensions from what's in the
# paper (these may or may not be extensions of the cg_descent code
# that can be downloaded from Hager's site; his code has undergone
# numerous revisions since publication of the paper):
#
# cgdescent: the termination condition employs a "unit-correct"
#   expression rather than a condition on gradient
#   components---whether this is a good or bad idea will require
#   additional experience, but preliminary evidence seems to suggest
#   that it makes "reasonable" choices over a wider range of problem
#   types.
#
# linesearch: the Wolfe conditions are checked only after alpha is
#   generated either by quadratic interpolation or secant
#   interpolation, not when alpha is generated by bisection or
#   expansion. This increases the likelihood that alpha will be a
#   good approximation of the minimum.
#
# linesearch: In step I2, we multiply by psi2 only if the convexity
#   test failed, not if the function-value test failed. This
#   prevents one from going uphill further when you already know
#   you're already higher than the point at alpha=0.
#
# both: checks for Inf/NaN function values
#
# both: support maximum value of alpha (equivalently, c). This
#   facilitates using these routines for constrained minimization
#   when you can calculate the distance along the path to the
#   disallowed region. (When you can't easily calculate that
#   distance, it can still be handled by returning Inf/NaN for
#   exterior points. It's just more efficient if you know the
#   maximum, because you don't have to test values that won't
#   work.) The maximum should be specified as the largest value for
#   which a finite value will be returned.  See, e.g., limits_box
#   below.  The default value for alphamax is Inf. See alphamaxfunc
#   for cgdescent and alphamax for linesearch_hz.

macro cgtrace()
    quote
        if tracing
            dt = Dict()
            if o.extended_trace
                dt["x"] = copy(x)
                dt["g(x)"] = copy(g)
                dt["Current step size"] = alpha
            end
            g_norm = vecnorm(g, Inf)
            update!(tr,
                    iteration,
                    f_x,
                    g_norm,
                    dt,
                    o.store_trace,
                    o.show_trace,
                    o.show_every,
                    o.callback)
        end
    end
end

immutable ConjugateGradient{T} <: Optimizer
    eta::Float64
    P::T
    precondprep!::Function
    linesearch!::Function
end

function ConjugateGradient(;
                           linesearch!::Function = hz_linesearch!,
                           eta::Real = 0.4,
                           P::Any = nothing,
                           precondprep! = (P, x) -> nothing)
    ConjugateGradient{typeof(P)}(Float64(eta),
                                 P, precondprep!,
                                 linesearch!)
end

function optimize{T}(df::DifferentiableFunction,
                     initial_x::Array{T},
                     mo::ConjugateGradient,
                     o::OptimizationOptions)
    # Print header if show_trace is set
    print_header(o)

    # Maintain current state in x and previous state in x_previous
    x, x_previous = copy(initial_x), copy(initial_x)

    # Count the total number of iterations
    iteration = 0

    # Track calls to function and gradient
    f_calls, g_calls = 0, 0

    # Count number of parameters
    n = length(x)

    # Maintain current gradient in g and previous gradient in g_previous
    g, g_previous = similar(x), similar(x)

    # Maintain the preconditioned gradient in pg
    pg = similar(x)

    # The current search direction
    s = similar(x)

    # Buffers for use in line search
    x_ls, g_ls = similar(x), similar(x)

    # Intermediate value in CG calculation
    y = similar(x)
    py = similar(x)

    # Store f(x) in f_x
    f_x = df.fg!(x, g)
    @assert typeof(f_x) == T
    f_x_previous = convert(T, NaN)
    f_calls, g_calls = f_calls + 1, g_calls + 1
    copy!(g_previous, g)

    # Keep track of step-sizes
    alpha = alphainit(one(T), x, g, f_x)

    # TODO: How should this flag be set?
    mayterminate = false

    # Maintain a cache for line search results
    lsr = LineSearchResults(T)

    # Trace the history of states visited
    tr = OptimizationTrace{typeof(mo)}()
    tracing = o.store_trace || o.show_trace || o.extended_trace || o.callback != nothing
    @cgtrace

    # Output messages
    if !isfinite(f_x)
        error("Must have finite starting value")
    end
    if !all(isfinite(g))
        @show g
        @show find(!isfinite(g))
        error("Gradient must have all finite values at starting point")
    end

    # Determine the intial search direction
    #    if we don't precondition, then this is an extra superfluous copy
    #    TODO: consider allowing a reference for pg instead of a copy
    mo.precondprep!(mo.P, x)
    A_ldiv_B!(pg, mo.P, g)
    scale!(copy!(s, pg), -1)

    # Assess multiple types of convergence
    x_converged, f_converged = false, false
    g_converged = vecnorm(g, Inf) < o.g_tol

    # Iterate until convergence
    converged = g_converged
    while !converged && iteration < o.iterations
        # Increment the number of steps we've had to perform
        iteration += 1

        # Reset the search direction if it becomes corrupted
        dphi0 = vecdot(g, s)
        if dphi0 >= 0
            @simd for i in 1:n
                @inbounds s[i] = -pg[i]
            end
            dphi0 = vecdot(g, s)
            if dphi0 >= 0
                break
            end
        end

        # Refresh the line search cache
        clear!(lsr)
        @assert typeof(f_x) == T
        @assert typeof(dphi0) == T
        push!(lsr, zero(T), f_x, dphi0)

        # Pick the initial step size (HZ #I1-I2)
        alpha, mayterminate, f_update, g_update =
          alphatry(alpha, df, x, s, x_ls, g_ls, lsr)
        f_calls, g_calls = f_calls + f_update, g_calls + g_update

        # Determine the distance of movement along the search line
        alpha, f_update, g_update =
          mo.linesearch!(df, x, s, x_ls, g_ls, lsr, alpha, mayterminate)
        f_calls, g_calls = f_calls + f_update, g_calls + g_update

        # Maintain a record of previous position
        copy!(x_previous, x)

        # Update current position # x = x + alpha * s
        LinAlg.axpy!(alpha, s, x)

        # Maintain a record of the previous gradient
        copy!(g_previous, g)

        # Update the function value and gradient
        f_x_previous, f_x = f_x, df.fg!(x, g)
        f_calls, g_calls = f_calls + 1, g_calls + 1

        x_converged,
        f_converged,
        g_converged,
        converged = assess_convergence(x,
                                       x_previous,
                                       f_x,
                                       f_x_previous,
                                       g,
                                       o.x_tol,
                                       o.f_tol,
                                       o.g_tol)

        # Check sanity of function and gradient
        if !isfinite(f_x)
            error("Function must finite function values")
        end

        # Determine the next search direction using HZ's CG rule
        #  Calculate the beta factor (HZ2012)
        # -----------------
        # Comment on py: one could replace the computation of py with
        #    ydotpgprev = vecdot(y, pg)
        #    vecdot(y, py)  >>>  vecdot(y, pg) - ydotpgprev
        # but I am worried about round-off here, so instead we make an
        # extra copy, which is probably minimal overhead.
        # -----------------
        mo.precondprep!(mo.P, x)
        dPd = dot(s, mo.P, s)
        etak::T = mo.eta * vecdot(s, g_previous) / dPd
        @simd for i in 1:n
            @inbounds y[i] = g[i] - g_previous[i]
        end
        ydots = vecdot(y, s)
        copy!(py, pg)        # below, store pg - pg_previous in py
        A_ldiv_B!(pg, mo.P, g)
        @simd for i in 1:n     # py = pg - py
           @inbounds py[i] = pg[i] - py[i]
        end
        betak = (vecdot(y, pg) - vecdot(y, py) * vecdot(g, s) / ydots) / ydots
        beta = max(betak, etak)
        @simd for i in 1:n
            @inbounds s[i] = beta * s[i] - pg[i]
        end

        @cgtrace
    end

    return MultivariateOptimizationResults("Conjugate Gradient",
                                           initial_x,
                                           x,
                                           Float64(f_x),
                                           iteration,
                                           iteration == o.iterations,
                                           x_converged,
                                           o.x_tol,
                                           f_converged,
                                           o.f_tol,
                                           g_converged,
                                           o.g_tol,
                                           tr,
                                           f_calls,
                                           g_calls)
end


import StatsBase: sample

# Model selection

# Taken from https://royalsocietypublishing.org/doi/pdf/10.1098/rspa.2017.0009
"""
	AIC(k, X, Y; likelyhood = (X, Y) = sum(abs2, X-Y))

Computes the Akaike Information Criterion (AIC) given the free parameters `k` for the data `X` and its
estimate `Y` of the model. `likelyhood` can be any function of `X` and `Y`.
"""
function AIC(k::Int64, X::AbstractArray, Y::AbstractArray; likelyhood = (X,Y) -> sum(abs2, X-Y))
    @assert size(X) == size(Y) "Dimensions of trajectories should be equal !"
    return 2*k - 2*log(likelyhood(X, Y))
end

# Taken from https://royalsocietypublishing.org/doi/pdf/10.1098/rspa.2017.0009
"""
	AICC(k, X, Y; likelyhood = (X, Y) = sum(abs2, X-Y))

Computes the Akaike Information Criterion compensated for finite samples (AICC) given the free parameters `k` for the data `X` and its
estimate `Y` of the model. `likelyhood` can be any function of `X` and `Y`.
"""
function AICC(k::Int64, X::AbstractMatrix, Y::AbstractMatrix; likelyhood = (X,Y) -> sum(abs2, X-Y))
    @assert size(X) == size(Y) "Dimensions of trajectories should be equal !"
    return AIC(k, X, Y, likelyhood = likelyhood)+ 2*(k+1)*(k+2)/(size(X)[2]-k-2)
end

function AICC(k::Int64, X::AbstractVector, Y::AbstractVector; likelyhood = (X,Y) -> sum(abs2, X-Y))
    @assert size(X) == size(Y) "Dimensions of trajectories should be equal !"
    return AIC(k, X, Y, likelyhood = likelyhood)+ 2*(k+1)*(k+2)/(length(X)-k-2)
end

# Double check on that
# Taken from https://www.immagic.com/eLibrary/ARCHIVES/GENERAL/WIKIPEDI/W120607B.pdf
"""
	BIC(k, X, Y; likelyhood = (X, Y) = sum(abs2, X-Y))

Computes Bayes Information Criterion (BIC) given the free parameters `k` for the data `X` and its
estimate `Y` of the model. `likelyhood` can be any function of `X` and `Y`.
"""
function BIC(k::Int64, X::AbstractMatrix, Y::AbstractMatrix; likelyhood = (X,Y) -> sum(abs2, X-Y))
    @assert size(X) == size(Y) "Dimensions of trajectories should be equal !"
    return - 2*log(likelyhood(X, Y)) + k*log(size(X)[2])
end

function BIC(k::Int64, X::AbstractVector, Y::AbstractVector; likelyhood = (X,Y) -> sum(abs2, X-Y))
    @assert size(X) == size(Y) "Dimensions of trajectories should be equal !"
    return - 2*log(likelyhood(X, Y)) + k*log(length(X))
end

# Optimal Shrinkage for data in presence of white noise
# See D. L. Donoho and M. Gavish, "The Optimal Hard Threshold for Singular
# Values is 4/sqrt(3)", http://arxiv.org/abs/1305.5870
# Code taken from https://github.com/erichson/optht

function optimal_svht(m::Int64, n::Int64; known_noise::Bool = false)
    @assert m/n > 0
    @assert m/n <= 1

    β = m/n
    ω = (8*β) / (β+1+sqrt(β^2+14β+1))
    c = sqrt(2*(β+1)+ω)

    if known_noise
        return c
    else
        median = median_marcenko_pastur(β)
        return c / sqrt(median)
    end
end

function marcenko_pastur_density(t, lower, upper, beta)
    sqrt((upper-t).*(t-lower))./(2π*beta*t)
end

function incremental_marcenko_pastur(x, beta, gamma)
    @assert beta <= 1
    upper = (1+sqrt(beta))^2
    lower = (1-sqrt(beta))^2

    @inline marcenko_pastur(x) = begin
        if (upper-x)*(x-lower) > 0
            return marcenko_pastur_density(x, lower, upper, beta)
        else
            return zero(eltype(x))
        end
    end

    if gamma ≈ zero(eltype(gamma))
        i, ϵ = quadgk(x->(x^gamma)*marcenko_pastur(x), x, upper)
        return i
    else
        i, ϵ = quadgk(x->marcenko_pastur(x), x, upper)
        return i
    end
end

function median_marcenko_pastur(beta)
    @assert 0 < beta <= 1
    upper = (1+sqrt(beta))^2
    lower = (1-sqrt(beta))^2
    change = true
    x = ones(eltype(upper), 5)
    y = similar(x)
    while change && (upper - lower > 1e-5)
        x = range(lower, upper, length = 5)
        for (i,xi) in enumerate(x)
            y[i] = one(eltype(x)) - incremental_marcenko_pastur(xi, beta, 0)
        end
        any(y .< 0.5) ? lower = maximum(x[y .< 0.5]) : change = false
        any(y .> 0.5) ? upper = minimum(x[y .> 0.5]) : change = false
    end
    return (lower+upper)/2
end

"""
	optimal_shrinkage(X)
	optimal_shrinkage!(X)

Compute a feature reduced version of the data array `X` via tresholding the
singular values by computing the [optimal threshold for singular values](http://arxiv.org/abs/1305.5870).
"""
function optimal_shrinkage(X::AbstractArray{T, 2}) where T <: Number
    m,n = minimum(size(X)), maximum(size(X))
    U, S, V = svd(X)
    τ = optimal_svht(m,n)
    inds = S .>= τ*median(S)
    return U[:, inds]*Diagonal(S[inds])*V[:, inds]'
end

function optimal_shrinkage!(X::AbstractArray{T, 2}) where T <: Number
    m,n = minimum(size(X)), maximum(size(X))
    U, S, V = svd(X)
    τ = optimal_svht(m,n)
    inds = S .>= τ*median(S)
    X .= U[:, inds]*Diagonal(S[inds])*V[:, inds]'
    return
end

"""
	savitzky_golay(X, windowSize, polyOrder; deriv, dt, crop)

Estimate the time derivative via the savitzky_golay filter. `X` is the data matrix containing the trajectories, which is interpolated
via polynomials of order `polyOrder` over `windowSize` points repeatedly. `deriv` defines the order of the derivative, `dt`
the timestepsize. `crop` indicates if the original data should be returned cropped along the derivative approximation.
"""
function savitzky_golay(x::AbstractVector{T}, windowSize::Integer, polyOrder::Integer; deriv::Integer=0, dt::Real=1.0, crop::Bool = true) where T <: Number
	# Polynomial smoothing with the Savitzky Golay filters
	# Adapted from: https://github.com/BBN-Q/Qlab.jl/blob/master/src/SavitskyGolay.jl
	# More information: https://pdfs.semanticscholar.org/066b/7534921b308925f6616480b4d2d2557943d1.pdf
	# Requires LinearAlgebra and DSP modules loaded.

	# Some error checking
	@assert isodd(windowSize) "Window size must be an odd integer."
	@assert polyOrder < windowSize "Polynomial order must be less than window size."

	# Calculate filter coefficients
	filterCoeffs = calculate_filterCoeffs(windowSize, polyOrder, deriv, dt)

	# Pad the signal with the endpoints and convolve with filter
	halfWindow = Int(ceil((windowSize - 1)/2))
	paddedX = [x[1]*ones(halfWindow); x; x[end]*ones(halfWindow)]
	y = conv(filterCoeffs[end:-1:1], paddedX)

	if !crop
		# Return the valid midsection
		return y[2*halfWindow+1:end-2*halfWindow]
	else
		# Return cropped data. Excluding borders, where the estimation is less accurate
		return x[halfWindow+2:end-halfWindow-1], y[3*halfWindow+2:end-3*halfWindow-1]
	end
end

function savitzky_golay(x::AbstractMatrix{T}, windowSize::Integer, polyOrder::Integer; deriv::Integer=0, dt::Real=1.0, crop::Bool = true) where T <: Number
	# Polynomial smoothing with the Savitzky Golay filters
	# Adapted from: https://github.com/BBN-Q/Qlab.jl/blob/master/src/SavitskyGolay.jl
	# More information: https://pdfs.semanticscholar.org/066b/7534921b308925f6616480b4d2d2557943d1.pdf
	# Requires LinearAlgebra and DSP modules loaded.

	# Some error checking
	@assert isodd(windowSize) "Window size must be an odd integer."
	@assert polyOrder < windowSize "Polynomial order must be less than window size."

	# Calculate filter coefficients
	filterCoeffs = calculate_filterCoeffs(windowSize, polyOrder, deriv, dt)

	# Apply filter to each component
	halfWindow = Int(ceil((windowSize - 1)/2))

	if !crop
		y = similar(x)
		for (i, xi) in enumerate(eachrow(x))
			paddedX = [xi[1]*ones(halfWindow); xi; xi[end]*ones(halfWindow)]
			y₀ = conv(filterCoeffs[end:-1:1], paddedX)
			y[i,:] = y₀[2*halfWindow+1:end-2*halfWindow]
		end
		return y
	else
		cropped_x = x[:,halfWindow+2:end-halfWindow-1]
		y = similar(cropped_x)
		for (i, xi) in enumerate(eachrow(x))
			paddedX = [xi[1]*ones(halfWindow); xi; xi[end]*ones(halfWindow)]
			y₀ = conv(filterCoeffs[end:-1:1], paddedX)
			y[i,:] = y₀[3*halfWindow+2:end-3*halfWindow-1]
		end
		return cropped_x, y
	end
end

function calculate_filterCoeffs(windowSize::Integer, polyOrder::Integer, deriv::Integer, dt::Real)
	# Some error checking
	@assert isodd(windowSize) "Window size must be an odd integer."
	@assert polyOrder < windowSize "Polynomial order must be less than window size."

	# Form the design matrix A
	halfWindow = Int(ceil((windowSize - 1)/2))
	A = zeros(windowSize, polyOrder+1)
	for order = 0:polyOrder
		A[:, order+1] = (-halfWindow:halfWindow).^(order)
	end

	# Compute the required column of the inverse of A'*A
	# and calculate filter coefficients
	ei = zeros(polyOrder+1)
	ei[deriv+1] = 1.0
	inv_col = (A'*A) \ ei
	return A*inv_col * factorial(deriv) ./(dt^deriv)
end

"""
	burst_sampling(X, samplesize, n)

Randomly selects `n` bursts of data with size `samplesize` from the data `X`.
"""
@inline function burst_sampling(x::AbstractArray, samplesize::Int64, bursts::Int64)
    @assert size(x)[end] >= samplesize*bursts "Bursting impossible. Please provide more data or reduce bursts or samplesize."
    inds = sample(1:size(x)[end]-samplesize, bursts, replace = false)
    inds = sort(unique(vcat([collect(i:i+samplesize) for i in inds]...)))
    return resample(x, inds)
end

"""
	burst_sampling(X, Y, samplesize, n)

Randomly selects `n` bursts of data with size `samplesize` from the data `X` and `Y`.
"""
@inline function burst_sampling(x::AbstractArray, y::AbstractArray, samplesize::Int64, bursts::Int64)
    @assert size(x)[end] >= samplesize*bursts "Bursting impossible. Please provide more data or reduce bursts or samplesize"
    @assert size(x)[end] == size(y)[end]
    inds = sample(1:size(x)[end]-samplesize, bursts, replace = false)
    inds = sort(unique(vcat([collect(i:i+samplesize) for i in inds]...)))
    return resample(x, inds), resample(y, inds)
end

"""
	burst_sampling(X, t, period, n)

Randomly selects `n` bursts of data within a time window `period` from the data `X`. The time information
has to be provided in `t`.
"""
@inline function burst_sampling(x::AbstractArray, t::AbstractVector, period::T, bursts::Int64) where T <: AbstractFloat
    @assert period > zero(typeof(period)) "Sampling period has to be positive."
    @assert size(x)[end] == size(t)[end] "Provide consistent data."
    @assert bursts >= 1 "Number of bursts has to be positive."
    @assert t[end]-t[1]>= period*bursts "Bursting impossible. Please provide more data or reduce bursts or samplesize"
    t_ids = zero(eltype(t)) .<= t .- period  .<= t[end] .- 2*period
    samplesize = Int64(floor(period/(t[end]-t[1])*length(t)))
    inds = sample(collect(1:length(t))[t_ids], bursts, replace = false)
    inds = sort(unique(vcat([collect(i:i+samplesize) for i in inds]...)))
    return resample(x, inds), resample(t, inds)
end

"""
	subsample(X, n)

Returns the subsampled `X` with only every `n`-th entry.
"""
@inline function subsample(x::AbstractVector, frequency::Int64)
    @assert frequency > 0 "Sampling frequency has to be positive."
    return x[1:frequency:end]
end

@inline function subsample(x::AbstractArray, frequency::Int64)
    @assert frequency > 0 "Sampling frequency has to be positive."
    return x[:, 1:frequency:end]
end

"""
	subsample(X, t, dt)

Returns the subsampled `X` with a a minimum period of `dt` between two data points. `t` provides the
time information.
"""
@inline function subsample(x::AbstractArray, t::AbstractVector, period::T) where T <: AbstractFloat
    @assert period > zero(typeof(period)) "Sampling period has to be positive."
    @assert size(x)[end] == size(t)[end] "Provide consistent data."
    @assert t[end]-t[1]>= period "Subsampling impossible. Sampling period exceeds time window."
    idx = Int64[1]
    t_now = t[1]
    @inbounds for (i, t_current) in enumerate(t)
        if t_current - t_now >= period
            push!(idx, i)
            t_now = t_current
        end
    end
    return resample(x, idx), resample(t, idx)
end

@inline function resample(x::AbstractArray{T,1}, indx::AbstractArray{Int64}) where T <: Number
    @assert maximum(indx) <= length(x) "Sampling index has to be consistent with array dimensions."
    @assert minimum(indx) >= 1 "Sampling index has to be consistent with array dimensions."
    return x[indx]
end

@inline function resample(x::AbstractArray{T,2}, indx::AbstractArray{Int64}) where T <: Number
    @assert maximum(indx) <= size(x, 2) "Sampling index has to be consistent with array dimensions."
    @assert minimum(indx) >= 1 "Sampling index has to be consistent with array dimensions."
    return x[:, indx]
end


function LinearAlgebra.det(M::Matrix{<:AbstractPolynomialLike})
    m = size(M)[1]
    if m > 2
        return sum((-1)^(i-1) * M[i,1] * LinearAlgebra.det(M[1:end .!= i, 2:end]) for i in 1:m)
    else
        return M[1,1] * M[2,2] - M[2,1] * M[1,2]
    end
end


using Qaintessent
using LinearAlgebra
using SparseArrays: sparse
using Memoize

# Simple struct that represents a graph via its edges
struct Graph
    n::Int # number of vertices (indices 1,...,n)
    edges::Set{Set{Int}} # edges, represented as sets of vertices

    function Graph(n::Integer, edges::Vector{Tuple{T, T}}) where T <: Integer
        n >= 1 || throw(DomainError("n must be a positive integer"))

        # Turn into set of sets
        edge_set = Set(Set.(edges))

        # Verify that all edges are valid
        all(edge -> edge ⊆ 1:n && length(edge) == 2, edge_set) || throw(ArgumentError("Some edges have invalid endpoints"))
        new(n, edge_set)
    end
end

"""
    Phase separation gate for Max-κ-colorable subgraph QAOA mapping.
    Represents the objective function which counts the number of invalid
    edges (i.e. between vertices of the same color).
    Implemented the for one-hot encoding.

``U_{P}(\\gamma) = e^{-i \\gamma H_{P}}``
``H_{P} = \\sum_{\\{u, v\\} \\in E} \\sum_{a=1}^{\\kappa} Z_{u, a} Z_{v, a}``

Reference:\n
    Stuart Hadfield, Zhihui Wang, Bryan O'Gorman, Eleanor G. Rieffel, Davide Venturelli and Rupak Biswas\n
    From the Quantum Approximate Optimization Algorithm to a Quantum Alternating Operator Ansatz\n
    Algorithms 12.2 (2019), p.34
"""
struct MaxKColSubgraphPhaseSeparationGate <: AbstractGate
    # use a reference type (array with 1 entry) for compatibility with Flux
    γ::Vector{Float64}
    κ::Int # the number of possible colors
    graph::Graph # the underlying graph which should be colored

    function MaxKColSubgraphPhaseSeparationGate(γ::Float64, κ::Integer, graph::Graph)
        κ > 0 || throw(ArgumentError("Parameter `κ` must be a positive integer."))
        length(graph.edges) > 0 || throw(ArgumentError("Graph `graph` must have at least one edge."))
        new([γ], κ, graph)
    end
end

@memoize function max_k_col_subgraph_phase_separation_hamiltonian(graph::Graph, κ::Int)
    z = matrix(Z)

    # Implementation of Eq. (17)
    # one-hot encoding: n * κ vector. Index (a-1)*n + (b-1) corresponds to vertex a, color b.
    H_P_enc = sum(
        kron((color == a && vertex ∈ edge ? z : I(2) for vertex ∈ 1:graph.n for color ∈ 1:κ)...) # Z_{u,a} Z_{v,a}
        for a ∈ 1:κ for edge ∈ graph.edges # Σ_{(u,v) = edge ∈ E} Σ_{a=1..κ}
    )
    return Diagonal(H_P_enc)
end

function Qaintessent.matrix(g::MaxKColSubgraphPhaseSeparationGate)
    # Calculate the hamiltonian (Eq. 17)
    H_P_enc = max_k_col_subgraph_phase_separation_hamiltonian(g.graph, g.κ)

    # Implementation of one-hot phase seperator, A.3.1, p. 34
    U_P = exp(-im * g.γ[] * H_P_enc)

    U_P
end

Qaintessent.adjoint(g::MaxKColSubgraphPhaseSeparationGate) = MaxKColSubgraphPhaseSeparationGate(-g.γ[], g.κ, g.graph)

Qaintessent.sparse_matrix(g::MaxKColSubgraphPhaseSeparationGate) = sparse(matrix(g))

# Number of wires (κ * n in the one-hot encoding)
Qaintessent.num_wires(g::MaxKColSubgraphPhaseSeparationGate)::Int = g.κ * g.graph.n