2DGaussianVol

from numpy.linalg import inv as np_inv
from numpy.linalg import svd, det, cholesky
# from numpy.linalg import slogdet as np_slogdet
from numpy import array, trace, log, diag
from numpy import sqrt as np_sqrt
from scipy.linalg import sqrtm
from scipy.optimize import bisect, brenth, minimize_scalar, LinearConstraint, minimize
import numpy as np
import sys
from fpylll import *
from fpylll.algorithms.bkz2 import BKZReduction

ROUNDING_FACTOR = 2**64

def round_matrix_to_rational(M):
    A = matrix(ZZ, (ROUNDING_FACTOR * matrix(M)).apply_map(round))
    return matrix(QQ, A / ROUNDING_FACTOR)

def projection_matrix(A):
    """
    Construct the projection matrix orthogonally to Span(V)
    """
    S = A * A.T
    return A.T * S.inverse() * A

def square_root_inverse_degen(S, B=None, assume_full_rank=False):
    """ Compute the determinant of a symmetric matrix
    sigma (m x m) restricted to the span of the full-rank
    rectangular (k x m, k <= m) matrix V
    """
    
    if assume_full_rank:
        P = identity_matrix(S.ncols())

    elif not assume_full_rank and B is None:
        # Get an orthogonal basis for the Span of B
        V = S.echelon_form()
        V = V[:V.rank()]
        P = projection_matrix(V)

    else:
        P = projection_matrix(B)

    # make S non-degenerated by adding the complement of span(B)
    C = identity_matrix(S.ncols()) - P
    # Take matrix sqrt via SVD, then inverse
    # S = adjust_eigs(S)
    
    u, s, vh = svd(array(S + C, dtype=float))
    L_inv = np_inv(vh) @ np_inv(np_sqrt(diag(s))) @ np_inv(u)
    # L_inv = np_inv(sqrtm(array(S + C, dtype=float)))
    
    L_inv = np_inv(cholesky(array(S + C, dtype=float))).T
    L_inv = round_matrix_to_rational(L_inv)
    L = L_inv.inverse()


    # scipy outputs complex numbers, even for real valued matrices. Cast to real before rational.
    #L = round_matrix_to_rational(u @ np_sqrt(diag(s)) @ vh)
  
    return L, L_inv

def create_2d_plot(mie, direction, a, b, list_points):
    from copy import deepcopy
    base_mie = MIE(identity_matrix(2), vector([0, 0]))
    p = base_mie.plot2d() + mie.plot2d()
    for samp_point in list_points:
        if samp_point[1] == 0:
            p += point(samp_point[0], rgbcolor=hue(0))
        else:
            p += point(samp_point[0], rgbcolor=hue(0.6))

    return p

# input a coefficient matrix for an "intuitive" ellipsoid and get back the one
# scaled for the paper
def build_sigma_matrix(mat):
    return mat.inverse() ^ 2

# From the papers:
# intuitive form of ellipse: E = {c + Au : u in S^2}
# ellipoid norm form:        E = {x in R^n : <X(x-c), x-c> <= 1}
# Here, Sigma = X, and since sqrt(X^-1) = A, it follows that
# A = sqrt({Sigma}^-1)

class MIE:
    # right now only checking for positive definite matrix, but we
    # technically should verify poitive semidefinite
    def __init__(self, S, mu):
        if not S.is_positive_definite():
            print("ERROR: must input a positive definite matrix")
            return
        self.S = S
        self.mu = mu
                
    def dim(self):
        return len(self.mu)

    # WARNING: this is done assuming that self.S is in the intiuitve form
    def plot2d(self):
        ellipse_angle = atan2(self.S.column(1)[1], self.S.column(1)[0])
        
        return ellipse(list(self.mu), self.S.column(1).norm(), self.S.column(0).norm(),
                       angle = ellipse_angle)
    
    # NOTE: in the toolkit everything is done with rows instead of columns
    # go about this assuming direction is a unit vector from now on,
    # this matches the definition of what one expects when working with a
    # "direction" vector

    # but the user need not worry about this, we'll normalize it
    def integrate_parallel_cuts_hint(self, direction, a, b):
        # the meaning of the signs here is kind of superfluous, this is just
        # to determine where everything is relative to the center of the
        # ellipsoid
        # 
        # a and b are distances from the center of the ellipsoid to the
        # two hyperplanes in the hint, along the direction of direction

        # now a and b are vectors representing the center of the ellipse
        # to the hyperplanes
        direction = direction / direction.norm()
        a = a * direction
        b = b * direction

        # there are problems if the direction is not in the column space
        # right now just error out
        if not (a in self.S.column_space() and b in self.S.column_space()):
            #return InvalidHint("direction not in column space of Sigma")
            print("a or b along direction not in column space of Sigma")
            return

        # A as above := sqrt_inv_mat
        (sqrt_mat, sqrt_inv_mat) = square_root_inverse_degen(self.S)

        # Step 1, subtract
        # before: E = mu + (sqrt_inv_mat)B_n
        # after: E = (sqrt_inv_mat)B_n
        # note that this applies to everyone inside the ball, and since
        # we only care about a and b right now, only apply to a and b
        # however this step shouldn't matter I think because a and b
        # were defined with repect to the center
        # a -= self.mu
        # b -= self.mu

        # make direction actual direction (scaled) and treat a and b as vectors from this

        # this is to accommodate for step 2 in our drawing
        # Step 2: stretch the ellipsoid into ball
        # before: E = (sqrt_inv_mat)B_n
        # after: E = B_n
        # to get back to a ball for easy rotations, we must "stretch"
        # the ellipsoid back into a ball, which is done by
        # multiplying by sqrt_mat

        a_scaled = sqrt_mat * a
        b_scaled = sqrt_mat * b

        
        # this is to apply a Householder transformation
        e = zero_vector(self.dim())
        e[0] = 1

        refl = (e - direction / norm(direction)) / 2
        
        # this step might cause issues, ellipsoids are over the reals,
        # but when we multiply it by the lattice, this thing has to be rational
        # a solution is to round to the nearest rational with some precision,
        # but this can cause problems when done over and over again
        # for a single hint, this _should_ be fine, still keep note of this
        # though
        # since we are rotating and rotating back, some things should cancel out
        # intuitively, but we might have to figure this out later
        # if estimating, we need not worry about this
        G = AffineGroup(self.dim(), RR) # could be rationals, check back with this later

        if refl == zero_vector(self.dim()):
            refl_mat = G(identity_matrix(self.dim()), zero_vector(self.dim()))
        else:
            refl_mat = G.reflection(refl)

        # Step 3: rotate the ball such that a_scaled and b_scaled are aligned
        # with the first coordinate
        # before: E = B_n
        # after: E = B_n (rotated in some way)
        a_scaled_rot = refl_mat.A() * a_scaled + refl_mat.b()
        b_scaled_rot = refl_mat.A() * b_scaled + refl_mat.b()

        # now we have a unit ball with the first coordinate aligned, make sure
        # a and b are witihin this ball (since they are scalar multiples of
        # e_1, we only have to check the first coordinate)

        alpha = a_scaled_rot[0]
        beta = b_scaled_rot[0]

        alpha, beta = min(alpha, beta), max(alpha, beta)


        # in the future we might want to make it so that we assume the extreme
        # hyperplane is the tangent plane of the hypersphere
        if abs(alpha) > 1 or abs(beta) > 1:
            #raise InvalidHint("alpha or beta is too big to be useful")
            print("alpha or beta are too big to be useful")

        # this is a and b in the paper, changed names to avoid confusion
        matrix_first = 0
        matrix_rest = 0
        tau = 0
        n = self.dim()
        left_condition = 4 * n * (1 - alpha) * (1 + alpha)
        right_condition = (n + 1) * (n + 1) * (beta - alpha) * (beta + alpha)

        if alpha == -beta:
            tau = 0
            matrix_first = beta
            matrix_rest = 1
        elif left_condition < right_condition:
            tau = 0.5 * (alpha + sqrt(alpha * alpha + left_condition / pow(n + 1, 2)))
            matrix_first = tau - alpha
            matrix_rest = sqrt(matrix_first * (matrix_first + n * tau))
        else: # (left_condition >= right_condition)
            denom = 2 * (sqrt((1 - alpha) * (1 + alpha)) - sqrt((1 - beta) * (1 + beta)))
            tau = 0.5 * (beta + alpha)
            matrix_first = 0.5 * (beta - alpha)
            matrix_rest = sqrt(matrix_first ** 2 + pow((beta ** 2 - alpha ** 2) / denom, 2))



        # this is to build up the diagonal matrix as in the paper
        z = zero_vector(RR, n)
        z[0] = matrix_first
        for ind in range(1, n):
            z[ind] = matrix_rest
        A = diagonal_matrix(z)
        c = zero_vector(RR, n)
        c[0] = tau

        self.S = A
        self.mu = c
def prob_cuts(lbound, rbound, trials):
    m = MIE(identity_matrix(RR, 2), vector(RR, [0, 0]))
    m.integrate_parallel_cuts_hint(vector(RR, [1, 0]), lbound, rbound)
    mean = [0, 0]
    cov = [[1, 0], [0, 1]]
    a = m.S[0][0]
    b = m.S[1][1]
    # In MIE if (point.x/a)^2 + (point.y/b)^2 <= 1
    list_points = []
    inMIE = 0
    inEllipse = 0
    for ind in range(trials):
        point = np.random.multivariate_normal(mean, cov, int(1))
        if (((point[0][0]-m.mu[0])/a)^2 + ((point[0][1]-m.mu[1])/b)^2) <= 1:
            list_points.append([point,1])
            inMIE += 1
        elif ((point[0][0])^2 + (point[0][1])^2) <= 1 and point[0][0] <= rbound and point[0][0] >= lbound:
            list_points.append([point,0])
            inEllipse += 1
    print(f'Vol Lost = {1-N(inMIE/(inMIE+ inEllipse))}')
    show(create_2d_plot(m, vector(RR, [1, 0]), lbound, rbound, list_points)) 
prob_cuts(0.7, 1,60000)