import numpy as np


def interior_point_repeat(n, K, r):
    m = n / r
    assert m % 1 == 0, "r does not divide n"
    m = int(m)

    b = 1. / r - (K - 1) ** 0.5 / (r * (r - 1) ** 0.5)
    a = 1. - (r - 1) * b

    U0 = (a - b) * np.eye(r) + np.full((r, r), b)
    U0 = np.repeat(U0, m, axis=0) / (m ** 0.5)

    P, S, Q0 = np.linalg.svd(U0, full_matrices=False)
    if P[0, 0] < 0:
        P = -P
        Q0 = -Q0
    V0 = P[:, 1:] @ np.diag(S[1:])

    return V0, Q0


def interior_point_rotate(n, K, r, seed=None):
    # Random generator
    rng = np.random.default_rng(seed=seed)

    # Cycle through points on the standard basis
    U0 = np.zeros((n, r))
    U0[range(n), [m % r for m in range(n)]] = 1
    rng.shuffle(U0)

    # Construct rotation to put the standard simplical cone from its axis
    # at [1,1,...,1] to [1,0,...,0]
    Q0 = np.hstack((np.ones(shape=(r, 1)),
                    rng.normal(size=(r, r - 1))))
    Q0, _ = np.linalg.qr(Q0)
    Q0 = Q0 * np.sign(Q0[0, 0])

    # Shrink the other components towards the axis
    V0 = U0.dot(Q0) / np.sqrt(n / r)
    V0[:, 1:] -= V0[:, 1:].mean(axis=0)
    V0[:, 1:] /= np.linalg.norm(V0[:, 1:], 'fro')
    V0[:, 1:] *= np.sqrt(K - 1)

    # Rotate back
    V0 = V0[:, 1:]
    Q0 = Q0.T

    return V0, Q0
