#!/usr/bin/env python
# coding: utf-8

# We want to compute $A_4^*$. In other words, we are looking for the coefficients of a convex polynomial
# $$p(x, y) = \sum_{i=0}^{2d} c_i x^i y^{2d-i}$$
# that maximizes
# $$\frac{c_4}{70}.$$
# and satisfies
# $p(1, 0) + p(0, 1) = 2$, or equivalently $c_0 + c_8 = 1$.
# 
# In the following we perform a series of symmetry reduction techniques that will allow us to solve this problem analytically.
# 
# 

# In[1]:


# Some helpful functions

# The following function creates an N x N symmetric symbolic matrix
def symb_matrix(N, base_name='Q_%d%d', sym=False):
    var_swapper = lambda i,j: (i, j)
    if sym:
        var_swapper = lambda i,j: (min(i,j), max(i,j))
        
    Q_varnames = ','.join([base_name % var_swapper( i,j) for (i,j) in cartesian_product([range(N),range(N)])])
    Q_vars = var(Q_varnames) 
    Q = matrix(SR,N,N,Q_vars)
    
    return Q


# the following function gives a full rank description of
# matrices `Q` that satisfy the equations in `eqn`
def eliminate_var_matrix(Q, eqn, name, additional_vars=[]):
    all_vars = list(Q.variables())
    all_vars += additional_vars
    all_vars = map(SR, all_vars)

    sols = solve(map(lambda eqn_i: SR(eqn_i) == 0, eqn), all_vars)
    sols = sols[0]
    
    # write Q as sum ri * Fi, where Fi are constant matrices
    r_F = [(term.rhs(), jacobian(Q, term.lhs())[0,0]) for term in sols]
    Fr = sum([ri * Fi for ri, Fi in r_F])
    
    # substitue the variables ri with their full rank descirption
    var_to_sub = sum([list(term.rhs().variables()) for term in sols], [])
    var_sub = {v: var(name+str(i)) for i,v in enumerate(var_to_sub)}
    return Fr.subs(var_sub)


# ## Reducing the number of coefficients in $p$
# 
# Note that if $p$ is such a polynomial, then so is $p(-x, -y)$, so we
# can assume that $p$ is even (i.e., $c_i = 0$ for $i=1,3,5,7$). Furthermore, if $p$ is a solution, so is
# $p(y, x)$, so we can assume that $p$ is symmetric (i.e., $c_i = c_{8-i}$). All in all, without
# loss of generality we assume
# 
# $$p(x, y) = x^8 + y^8 +  \alpha x^4y^4 + \beta (x^2y^6 + x^6y^2) \text{ for some $\alpha, \beta \in \mathbb R$}.$$
# 

# In[2]:


get_ipython().run_line_magic('display', 'latex')
d = 4
R.<a, b, x,y,u,v> = QQ['a,b,x,y,u,v']
vars = [x,y,u,v]
q = x^8 + y^8 + a*x^4*y^4 + b*(x^2*y^6 + x^6*y^2)
q


# ## Formulating the problem as an SDP
# Because of Theorem 3.6, imposing the condition on the convexity of $p$ we can formulated as a search problem for an $8 \times 8$ positive semidefinite matrix $Q$ that satisfies
#  $${\mathbf u}^T\nabla^2q({\mathbf x}){\mathbf u} = z^TQz \quad \forall {\mathbf x},{\mathbf u} \in \mathbb R^2$$
#  where
# $$z^T({\mathbf x}, {\mathbf u}) = (u x^{2d}, ux^{2d-1}y, \ldots, v y^{2d}, v x^{2d}, v x^{2d-1}y, \ldots, v y^{2d}).$$
# 
# Once we have $Q$, we can recover the polynomial $q$ via Euler's identity as follows:
# 
# $$q(\mathbf x) = \frac1{8 \cdot 7} z(\mathbf x, \mathbf x)^T Q z(\mathbf x, \mathbf x)$$
# 
# 
# Therefore, the objective function ca be written exclusively in terms of the matrix $Q$, and we can rewrite our optimization problem (*) completely in terms of the  matrix $Q$ as follows:
# 
# $$\max_{Q \succeq 0} Q \text{ s.t. } Q \in \mathcal L$$
# where $\mathcal L$ is the following linear space
# $$\{Q \; | \; \exists q = x^8 + y^8 +  \alpha x^4y^4 + \beta (x^2y^6 + x^6y^2) \text{ s.t. } {\mathbf u}^T\nabla^2q({\mathbf x}){\mathbf u} = z^TQz.\}$$

# In[3]:


Q = symb_matrix(8, 'Q_%d%d', sym=True)
Q


# In[4]:


# vector of monomial
z = vector((x^3*u - y^3*v, x*y^2*u - x^2*y*v,
     x^3*u + y^3*v, x*y^2*u + x^2*y*v, 
     x^2*y*u - x*y^2*v, y^3*u - x^3*v,
     x^2*y*u + x*y^2*v, y^3*u + x^3*v))

# hessian of q
Hq = jacobian(jacobian(q, (x, y)), (x, y))

# Q is in L if the following polynomial has all its coefficients equal to 0
diff = z * Q * z - vector([u,v]) * Hq * vector([u, v])
eqn_L = QQ[Q.variables() + (a,b)][x,y,u,v](diff).coefficients()


Q_L = eliminate_var_matrix(Q, eqn_L, 'w', [a])
Q_L


# In[5]:


# write the objective function in terms of the new paramterization
q_from_Q = (z*Q_L*z).subs({u:x, v:y}) / (2*d*(2*d-1))
q_from_Q =  QQ[Q_L.variables()][x,y](q_from_Q)
objective = q_from_Q.coefficient({xi: 4 for xi in q_from_Q.variables()}) / binomial(2*d, d)
objective = objective.base_ring()(objective)
objective


# 
# 
# ## Symmetry reduction techniques applied to $Q$
# 
# Let $\rho: \mathbb R^4 \mapsto \mathbb R^4$ be any linear transformation of the variables $(x,y,u,v)$ that leaves the coefficient of $x^4y^4$ in $p$ invariant and leaves polynomial $q(x,y,u,v) := {\mathbf u}^T\nabla^2q({\mathbf x}){\mathbf u}$ invariant, i.e.
# $$q(x,y,u,v) = q \circ \rho(x,y,u,v).$$
# The function $\rho$ naturally acts linearly on the vector $z(\mathbf x, \mathbf y)$. If we call the $8 \times 8$ matrix of this linear operator $M_\rho$, then 
# 
# $$Q \in \mathcal L \iff M_\rho^T Q M_\rho \in \mathcal L.$$
#  
#  
# Let $G$ be a group of such transformations $\rho$, then $Q$ is an optimal solution for (*) if and only if
# 
# $$Q^* = \frac1{|G|} \sum_{\rho \in G} M_\rho^T Q M_\rho$$
# 
# is also a solution. 
# 
# We will take $G$ to be the group generated by the following transformations
# 
# \begin{align*}
#   (x,y,u,v)   & \mapsto   (y,x,v,u)\\
#    \cdot \quad & \mapsto  (-x,y,-u,v)\\
#    \cdot \quad & \mapsto  (-x,-y,u,v)\\
#    \cdot \quad & \mapsto  (-x,y,u,-v).
# \end{align*}
# 
# The matrix $Q^*$ will have a very structured sparsity pattern.
# 

# To make our computations exact, we work over the ring $R := \mathbb Q[x,y,u,v]$ of polynomials in the variables $x,y,u,v$ with rational coefficients.

# In[6]:


symmetries =[[y,x,v,u],
             [-x,y,-u,v],
             [-x,-y,u,v],
             [-x,y,u,-v],
             [x, -y, u, -v],
            [x, -y, -u, v]]
symmetries_rho = map(lambda sym_i: matrix(QQ, jacobian(sym_i, vars)), symmetries)
symmetries_rho


# In[7]:


G = gap.Group(symmetries_rho)
G = gap.Elements(G).sage()
G = [matrix(QQ, g_elem_i) for g_elem_i in G]
print("The group G has %d elements"% len(G))


# In[8]:


# action of G on the vector of monomials z
rho_acting_on_z = [z.subs({vi: subi for vi,subi in zip(vars, g*vector(vars))}) for g in G]

# matrix representation of that action
def compute_M_rho(rho_z):
    return matrix(QQ,
      [ [1 if zj == zi else -1 if  zj == -zi else 0 for zi in z]
        for zj in rho_z])

M_rho = map(compute_M_rho, rho_acting_on_z)
print("First two elements of M_rho")
M_rho[:2]


# In[9]:


# Thanks to the symmetry, Q^* is block diagonal
Q_star = sum(M.T *Q_L*M for M in M_rho) / len(M_rho)
print("Blocks on the diagonal of Q:")
[Q_star[i:i+2, i:i+2] for i in range(0, 8, 2)]


# The matrix $Q^*$ has therefore the form
# $Q^* = diag(Q_1, Q_2, Q_3, Q_4)$,
# where the $Q_i$ are $2 \times 2$ symmetric matrices. $Q^* \succeq 0$ if and only if $Q_1,\ldots,Q_4 \succeq 0$

# # Dual 
# 
# Let us now derive the dual of the previous problem

# In[10]:


#DQ is the dual variable
DQi = [symb_matrix(2, 'D'+str(i)+'_%d%d', sym=True) for i in range(4)]
DQ = block_diagonal_matrix(DQi)
DQ


# Let us compute the lagrangian
# 
# $$L(Q, D) = \langle Q, D \rangle + \text{objective function}$$

# In[11]:


hessian = lambda pp: jacobian(jacobian(pp, (x,y)), (x,y))
dot = lambda u,v: vector(u) * vector(v)
mdot = lambda A, B: dot(A.list(), B.list())


lagrangian = mdot(DQ, Q_star) + objective
var_primal = Q_star.variables()
var_dual = DQ.variables()
lagrangian = QQ[var_dual][var_primal] (lagrangian)
lagrangian


# In[12]:


# Find a full rank description of the dual varialbe DQ
lagrang_coeffs = jacobian(lagrangian, var_primal)[0]
lagrang_coeffs = list(lagrang_coeffs)
DQ_star = eliminate_var_matrix(DQ, lagrang_coeffs, 'z')
DQ_star


# # use KKT conditions to transform the SDP to a system of polynomial equations

# In[13]:


R = QQ[DQ_star.variables() + Q_star.variables()]

# KKT equations
KKT_eqn = (DQ_star * Q_star).list()
KKT_eqn = list(set(KKT_eqn))

# ideal generated by the KKT equations
I = map(lambda p: R(p), KKT_eqn)
I = R*I


# In[14]:


print("Ideal I is generated by %d equations in %d variables" % (len(I.gens()), len(R.gens())))


# In[15]:


# eliminate all variables except those present in the objective function
I_a = I.elimination_ideal([v for v in R.gens() if SR(v) not in objective.variables()])
I_a


# In[16]:


print("The ideal I_a has %d generator" % len(I_a.gens()))


# In[17]:


var('t')
p = SR(I_a.gens()[0]).subs({SR(objective.variables()[0]): 70*t})
factor(p)


# In[18]:


minimal_poly = p.factor_list()[0][0]
minimal_poly /= minimal_poly.coefficient(t^3)
minimal_poly


# In[19]:


minimal_poly.roots()


# In[20]:


minimal_poly.roots(ring=RDF)