# See the preprint arXiv:1902.03505, ("Universal Optimal Configurations for
# the p-Frame Potentials" by Chen, Gonzalez, Goodman, Kang, and Okoudjou)
# for details and acknowledgements.

# This program tests the conjecture concerning lifted equiangular tight frames
# and the p-frame potential using a very basic gradient descent procedure. We
# use this program just as a guide. This code was written in connection with
# the MAPS REU, which was supported by the NSF.

# For more information, see the web page
#     https://www.math.umd.edu/~okoudjou/.............................

# Note: This program requires Sage; run this in Sage with Python.

from sage.all import *
from sage.matrix.matrix import *

import random

print '\nDone importing\n'

class Vector:
    def __init__(self, vector):
        # Expects a vector in R^N in the following form:
        # [ [x_1], [x_2], ..., [x_N]]
        assert len(vector) > 0 and all((len(vector[i]) == 1 for i in range(len(vector))))
        self.vec = vector
    
    def __str__(self):
        return 'Vector:\n' + str(self.vec)
    
    def __getitem__(self, row):
        # Gets the component of this vector in a particular row.
        return self.vec[row][0]
    
    def __setitem__(self, row, value):
        # Sets the row component of this vector to a particular value.
        self.vec[row][0] = value
    
    def dimension(self):
        # Returns the dimension in which this vector lives.
        return len(self.vec)
    
    def dot(self, v):
        # Returns the dot product of this vector and v.
        assert self.dimension() == v.dimension()
        return sum(self[r] * v[r] for r in range(self.dimension()))
    
    def length(self):
        # The length of this vector.
        return sqrt(self.dot(self))
    
    def add(self, v):
        # Returns a new Vector that is the sum of this one and v.
        assert self.dimension() == v.dimension()
        return Vector([ [self[r]+v[r]] for r in range(self.dimension())])
    
    def subtract(self, v):
        # Returns a new Vector with v subtracted from this Vector.
        return self.add(v.scaled(-1))
    
    def scaled(self, c):
        # Returns a new Vector that is a scaled version of this vector
        return Vector([ [c*self[r]] for r in range(self.dimension())])
    
    def normalized(self):
        # Returns a new Vector that is the normalized version of this one,
        # unless this is the zero vector. If this Vector is the zero vector,
        # then it reutrns another zero Vector.
        l = self.length()
        if l == 0:
            print 'Error: Vector of length 0 cannot be normalized. Returning self.duplicate(), instead.'
            return self.duplciate()
        return self.scaled(1. / self.length())
    
    def duplicate(self):
        # Duplicates a vector.
        return Vector([[self[r]] for r in range(self.dimension())])


class Frame:
    def __init__(self, fr):
        # For a frame of M vectors in R^N, expects an array of the form
        # [[x_{11}, ..., x_{M1}],
        #  [x_{12}, ..., x_{M2}],
        #  ...,
        #  [x_{1N}, ..., x_{MN}]]
        # where each column (for example, entries of the form x_{1j}) is a frame
        # vector.
        self.frame = fr
    
    def __str__(self):
        return 'Frame:\n' + str(matrix(self.frame))
    
    def __getitem__(self, (row, col)):
        # Gets the row-column entry of the frame (that is, the row-column entry
        # of the synthesis matrix).
        return self.frame[row][col]
    
    def __setitem__(self, (row, col), value):
        # Sets the row-column entry of the frame (as above, the row-column entry
        # of the synthesis matrix).
        self.frame[row][col] = value
    
    def get_column(self, col):
        # Returns the col frame vector as a Vector.
        return Vector([[self[r, col]] for r in range(self.vector_dimension())])
    
    def set_column(self, col, vector):
        # Sets the col frame vector to be the parameter vector, which is assumed
        # to be a Vector object.
        for r in range(self.vector_dimension()):
            self[r, col] = vector[r]
    
    # Computes the synthesis matrix, frame operator, and Grammian matrices
    # for this frame as Sage matrix objects.
    def synthesis(self):
        return matrix(self.frame)
    def frameoperator(self):
        m = self.synthesis()
        return m*transpose(m)
    def grammian(self):
        m = self.synthesis()
        return transpose(m)*m
    def S(self): return self.frameoperator()
    def G(self): return self.grammian()
    
    def vector_dimension(self):
        # Returns the dimension of the space in which the frame vectors live.
        return len(self.frame)
    
    def num_vectors(self):
        # Returns the number of frame vectors in this frame.
        return len(self.frame[0])
    
    def FP(self, p):
        # Finds p-Frame Potential of this frame for finite p.
        fp = 0
        p = float(p)
        numvecs = self.num_vectors()
        for col in range(numvecs):
            v1 = self.get_column(col)
            # Compute <v1, v1> separately so that it is not added twice, as the
            # other terms are in the loop below.
            fp += abs(v1.dot(v1))**p
            for col2 in range(col + 1, numvecs):
                v2 = self.get_column(col2)
                inprod = abs(v1.dot(v2))**p
                # Double the inner product to count both <v1, v2> and <v2, v1>
                # at the same time.
                fp += 2 * inprod
        return fp
    
    def normalize(self):
        # Normalizes all of the columns in this frame.
        for col in range(self.num_vectors()):
            self.set_column(col, self.get_column(col).normalized())
    
    def duplicate(self):
        # Duplicates this frame.
        m = [[self[r,c] for c in range(self.num_vectors())] for r in range(self.vector_dimension())]
        return Frame(m)





def random_normalized_frame(M, N, tol=1.e-10):
    # Generate a normalized frame of M vectors in R^N. The variable tol is used
    # when checking whether the generated set of M vectors spans R^N.
    spans = False
    while not spans:
        m = [[random.random() for col in range(M)] for row in range(N)]
        f = Frame(m)
        f.normalize()
        
        # Now check that this spans the space:
        spans = True
        for i in range(N):
            # See if the standard basis vector is projected onto any of
            # the frame vectors. First, build the standard basis vector:
            e_i = Vector([ [1 if i==j else 0] for j in range(N) ])
            nonzero_dot_product = False
            for j in range(M):
                if abs(e_i.dot(f.get_column(j))) > tol:
                    nonzero_dot_product = True
                    break
            if not nonzero_dot_product:
                # The basis vector does not get projected onto any frame vector
                # so this does not span
                spans = False
                break
    return f


def normalized_decrease_fp_one_step(Phi, p, step_size, derivative_step=1.e-8):
    # Numerically minimizes the frame potential of the frame Phi by shifting the
    # frame vectors according to the gradient of the p-frame potential. It will
    # return a normalized frame.
    # Notes:
    # - Phi must be a normalized frame.
    # - p is the value at which the p-Frame Potential is calculated.
    # - step_size is the amount by which to shift vectors.
    # - derivative_step is the step size used to approximate partial
    #   derivatives.
    M, N = Phi.num_vectors(), Phi.vector_dimension()
    
    # The array derivatives contains the partial derivatives of the p-FP with
    # respect to changes in the frame Phi. The i,j entry is the partial
    # derivative with respect to the i,j entry of the frame.
    derivatives = []
    # Set up empty (zero) entries in the array, one for each entry in Phi.
    derivatives = [[0 for col in range(M)] for row in range(N)]
    
    # Now the program computes the partial derivatives.
    fp_orig = Phi.FP(p)
    for column in range(M):
        original_vector = Phi.get_column(column)
        for row in range(N):
            changed = original_vector.duplicate()
            # Modify the component:
            changed[row] += derivative_step
            changed = changed.normalized()
            # Put the modified column in the original frame; it is reset below.
            Phi.set_column(column, changed)
            derivatives[row][column] = (Phi.FP(p) - fp_orig) / derivative_step
        # Reset the column:
        Phi.set_column(column, original_vector)
    
    # updatedPhi will be the new frame. First, set up an empty (with zeros)
    # frame of the needed size.
    updatedPhi = [[0 for col in range(M)] for row in range(N)]
    updatedPhi = Frame(updatedPhi)
    # Now compute the columns of the new frame:
    for col in range(M):
        # The gradient vector for this column
        gradientVector = Vector([ [derivatives[r][col]] for r in range(N) ])
        column = Phi.get_column(col)
        # The vector perpendicular contains the portion of the gradient vector
        # that is perpendicular to the current column.
        perpendicular = gradientVector.subtract( column.scaled( column.dot( gradientVector ) ) )
        # Scale perpendicular to have length step_size. This vector is then used
        # to shift the current column.
        if perpendicular.length() != 0:
            perpendicular = perpendicular.scaled(step_size / perpendicular.length())
        updatedPhi.set_column(col, column.subtract(perpendicular))
    updatedPhi.normalize()
    return updatedPhi


def normalized_decrease(Phi, p, num_iterations, step_sizes=(0.1, 1.e-2, 1.e-3, 1.e-4, 1.e-5, 1.e-6)):
    # Applies the function normalized_decrease_fp_one_step iteratively using
    # multiple step sizes to numerically minimize the frame potential.
    # Note, as above:
    # - Phi must be a normalized frame
    # - p is the value p for the p-Frame Potential
    # - num_iterations is the number of times to apply the function
    #   normalized_decrease_fp_one_step for each step size in step_sizes.
    # - step_sizes (not needed) is a list containing the values to use for
    #   step_size in the function normalized_decrease_fp_one_step.
    attempts = []
    for step in step_sizes:
        for i in range(num_iterations):
            Phi = normalized_decrease_fp_one_step(Phi, p, step)
            attempts.append((Phi, Phi.FP(p)))
        # Choose the frame with the minimum frame potential that it found
        Phi = min(attempts, key=lambda x: x[1])[0]
    return Phi





# The following two functions compute the values p_k and the p-Frame Potential
# of the L_k^N frames.
def p_k(k):
    return (log(k + 2) - log(k)) / (log(k + 1) - log(k))

def FPp_LkN(k, N, p):
    p = float(p)
    return N + 1. + (k + 1.) * k * (1. / k)**p





def test_conjecture_for_p_k(path, N, k, num_iterations):
    # Tests the conjecture that the L_k^N is the minimizer between p_{k-1}
    # and p_k for frames in R^N. Saves output to a file with a name of the form
    #   L_k^N Conjecture-num_iterations.txt
    # such as
    #   L_2^3 Conjecture-500.txt
    # to the location specified by the path parameter. The parameter
    # num_iterations is used when calling normalized_decrease.
    
    out_file = open(path + 'L_' + str(k) + '^' + str(N) + ' Conjecture-' + str(num_iterations) + '.txt', 'w')
    # Write the output to a file and print it to the window.
    def write(output):
        out_file.write(output)
        print output,
    
    interval = ()
    if k == 1:   interval = (0.1, p_k(1))
    elif k == N: interval = (p_k(N-1), 2)
    else:        interval = (p_k(k-1), p_k(k))
    
    # How many times to test the lower and upper bounds of the interval, along
    # with how many random values of p to test between the lower and upper
    # bounds.
    frames_lower_bound = 5
    frames_upper_bound = 5
    frames_random_p = 40
    
    num_frames = frames_lower_bound + frames_upper_bound + frames_random_p
    
    write('Testing L_' + str(k) + '^' + str(N) + ' Conjecture\n')
    write('num_iterations = ' + str(num_iterations) + '\n')
    write('frames_lower_bound = ' + str(frames_lower_bound) + '\n')
    write('frames_upper_bound = ' + str(frames_upper_bound) + '\n')
    write('frames_random_p    = ' + str(frames_random_p) + '\n')
    out_file.flush()
    
    # Counts the number of frames with a lower frame potential:
    counter = 0
    
    for frame_number in range(num_frames):
        write('\n')
        # Test a random value of p
        p = random.random() * (interval[1] - interval[0]) + interval[0]
        # Unless testing one of the boundaries:
        if frame_number < frames_lower_bound:
            p = interval[0]
        elif frame_number < frames_lower_bound + frames_upper_bound:
            p = interval[1]
        write('Test #' + str(frame_number + 1) + '\nTesting at p = ' + str(p) + '\n')
        Phi = random_normalized_frame(N+1, N)
        write('Generated random frame\n' + str(Phi) + '\nWith Grammian\n' + str(Phi.G()) + '\n')
        out_file.flush()
        Phi = normalized_decrease(Phi, float(p), num_iterations)
        
        write('\nMinimum p-Frame Potential found was with frame:\n' + str(Phi) + '\nWith Grammian\n' + str(Phi.G()) + '\n')
        phi_fp = float(Phi.FP(p))
        write('And with p-Frame Potential (p = ' + str(p) + ') equal to ' + str(phi_fp) + '\n')
        
        fp_lkn = FPp_LkN(k, N, p)
        write('Comparing to FPp(LkN) = ' + str(fp_lkn) + ' = ' + str(float(fp_lkn)) + '\n')
        write('Difference: ' + str(phi_fp - fp_lkn) + ' = ' + str(float(phi_fp - fp_lkn)) + '\n')
        if (phi_fp < fp_lkn):
            write('!'*100 + '\n')
            write('Found Lower Frame Potential\n')
            write('!'*100 + '\n')
            counter += 1
        write('-'*100 + '\n')
        out_file.flush()
    write('Found ' + str(counter) + ' frames with a lower frame potential.')
    out_file.close()





# This path tells the program where to save the output text files (include a
# final / after the path).
path = ''

# To run, call the testing function. Below is an example to test for 4 vectors
# in R^3 (by comparing to ONB+/L_1^3, L_2^3, and ETF/L_3^3) each time using 500
# iterations.
test_conjecture_for_p_k(path, 3, 1, 500)
test_conjecture_for_p_k(path, 3, 2, 500)
test_conjecture_for_p_k(path, 3, 3, 500)