using SumOfSquares
using DynamicPolynomials
using LinearAlgebra
using Oscar

using MosekTools
model = SOSModel(MosekTools.Optimizer)

d=3

@polyvar p[1:d] q[1:d]

x = [q;p]
n=length(x)

@polyvar z[1:n]

H(x) = 0.5*(x[4]^2+x[5]^2+x[6]^2+x[1]^2+x[2]^2+x[3]^2) + (x[1]^2*x[2]^2 + x[2]^2*x[3]^2 + x[3]^2*x[1]^2)

J = [zeros(d,d) I; -I zeros(d,d)]
f = J*differentiate(H(x),x)
Df = differentiate(f,x)

Φ = dot(z, Df*z)
l = Df*z - Φ*z

function make_z_symmetries(n, symmetries)
    T = eltype(symmetries[1])
    expandedsymmetries = Vector{Matrix{Int}}()
    
    signsymmetry = [I zeros(Int, n,n); zeros(Int, n,n) -I]

    push!(expandedsymmetries, signsymmetry)


    for symmetry in symmetries
        @assert(n == size(symmetry,1))
        @assert(n == size(symmetry,2))

        expandedsymmetry = zeros(T, 2n, 2n)
        expandedsymmetry[1:n,1:n] = symmetry
        expandedsymmetry[n+1:end,n+1:end] = symmetry

        push!(expandedsymmetries, expandedsymmetry)
    end
      
    return expandedsymmetries
end

symmetry_generators = Vector{Matrix{Int}}([
    [[-1 0 0 0 0 0;
      0 1 0 0 0 0;
      0 0 1 0 0 0;
      0 0 0 -1 0 0;
      0 0 0 0 1 0;
      0 0 0 0 0 1]];
    [[0 1 0 0 0 0;
      1 0 0 0 0 0;
      0 0 1 0 0 0;
      0 0 0 0 1 0;
      0 0 0 1 0 0;
      0 0 0 0 0 1]];
    [[0 0 1 0 0 0;
      0 1 0 0 0 0;
      1 0 0 0 0 0;
      0 0 0 0 0 1;
      0 0 0 0 1 0;
      0 0 0 1 0 0]];
    ])

symmetries = make_z_symmetries(n,symmetry_generators) # expand the symmetries to the augmented system

K,κ = cyclotomic_field(4,"κ")
G = matrix_group([matrix(K, symm) for symm in symmetries])
IR = invariant_ring(G)

import GroupsCore
struct GroupWrapper <: GroupsCore.Group
    G
end
struct GroupWrapperElement <: GroupsCore.GroupElement
    e
end

G = GroupWrapper(G)

Base.one(G::GroupWrapper) = GroupWrapperElement(Base.one(G.G))
Base.eltype(G::GroupWrapper) = typeof(one(G))
Base.IteratorSize(::GroupWrapper) = Base.HasLength()
function Base.iterate(G::GroupWrapper)
    result = Base.iterate(G.G)
    if isnothing(result)
        return nothing
    else
        return GroupWrapperElement(result[1]), result[2]
    end
end

function Base.iterate(G::GroupWrapper, prev)
    result = Base.iterate(G.G, prev)
    if isnothing(result)
        return nothing
    else
        return GroupWrapperElement(result[1]), result[2]
    end
end
GroupsCore.order(::Type{T}, G::GroupWrapper) where {T} = convert(T, Oscar.order(G.G))
GroupsCore.gens(G::GroupWrapper) = GroupWrapperElement.(Oscar.gens(G.G))
Base.parent(e::GroupWrapperElement) = GroupWrapper(Base.parent(e.e))
Base.:(==)(a::GroupWrapperElement, b::GroupWrapperElement) = a.e == b.e
Base.inv(e::GroupWrapperElement) = GroupWrapperElement(Base.inv(e.e))
function Base.:*(a::GroupWrapperElement, b::GroupWrapperElement)
    return GroupWrapperElement(a.e*b.e)
end
Base.copy(e::GroupWrapperElement) = GroupWrapperElement(e.e)
function Base.deepcopy_internal(e::GroupWrapperElement, id::IdDict)
    haskey(id, e) && return id[e]
    haskey(id, e.e) && return GroupWrapperElement(id[e.e])
    id[e] = GroupWrapperElement(Base.deepcopy_internal(e.e, id))
    return id[e]
end
import SymbolicWedderburn
struct GroupAction <: Symmetry.OnMonomials end
SymbolicWedderburn.coeff_type(::GroupAction) = Int
function SymbolicWedderburn.action(::GroupAction, e::GroupWrapperElement, mono::AbstractMonomial)
    mono([x;z]=>map_coeff.(Matrix(matrix(e.e)))*[x;z])
end

map_coeff = c -> iszero(c) ? 0 : evalpoly(cispi(2/4), Int.(collect(Oscar.coefficients(parent(defining_polynomial(K))(c)))))


function make_basis_of_degree(d)
    b = Oscar.basis(IR,d)
    basis = Vector{Polynomial{DynamicPolynomials.Commutative{DynamicPolynomials.CreationOrder}, Graded{LexOrder}, Int}}()

    polyvars = [x;z]

    for f in b
        ff = 0*polyvars[1] 
        for (c, e) in zip(Oscar.coefficients(f), Oscar.exponents(f))
            ff += map_coeff(c)*prod(polyvars[i]^e[i] for i in 1:length(polyvars))
        end
        push!(basis, ff)
    end
    return basis
end

function makebasis(maxdeg, mindeg)
    return reduce(vcat, [make_basis_of_degree(j) for j=mindeg:maxdeg])
end

bound = @variable(model)

basisV    = makebasis(2, 1)
Vc = @variable(model,[1:length(basisV)])
V = sum(Vc.*basisV)
dVdx = differentiate(V,x)
dVdz = differentiate(V,z)

basisrho = makebasis(4, 0)
rho = sum(@variable(model,[1:length(basisrho)]).*basisrho)



# finally, this is the polynomial we're interested in:
P = bound - Φ - dot(f, dVdx) - dot(l, dVdz) - rho*(1-dot(z,z))


# Check it's invariant under G
for g in G
    display(SymbolicWedderburn.action(GroupAction(),g,P)-P)
end

con_ref=@constraint(model, P  >= 0, symmetry = Symmetry.Pattern(G, GroupAction()))

@objective(model, Min, bound)

optimize!(model)

display(solution_summary(model))
display(value(bound))