import numpy as np
import scipy.stats as st
import matplotlib.pyplot as plt
import time
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt

np.random.seed(42)

### Point 1 ###

n = 10


# Generate order statistic
def generate_ord_stat_sort(n):
    X = st.uniform.rvs(size=n)
    return np.sort(X)


### Point 3 ###
def generate_ord_stat_nosort(n):
    X = np.empty(n)

    x = 1
    for i in range(0, n):
        # Sample from X_{(j)}|X_{(j+1)}=x by inversion
        U = st.uniform.rvs()
        X[n - i - 1] = x * U ** (1 / (n - i))
        x = X[n - i - 1]

    return X


# Vectorize previous function
def generate_ord_stat_nosort_vec(n):
    U = st.uniform.rvs(size=n)
    X = U ** (1 / (np.arange(1, n + 1)))
    X = np.cumprod(X[::-1])[::-1]

    return X


# Compare procedures

n_grid = 10 ** np.arange(3, 8)
time_sort = []
time_nosort = []

for n in n_grid:
    # Compute time for sorting-based sampling
    start = time.time()
    X_ord = generate_ord_stat_sort(n)
    end = time.time()
    time_sort.append(end - start)

    # Compute time for inversion-based sampling
    start = time.time()
    X_ord = generate_ord_stat_nosort_vec(n)
    end = time.time()
    time_nosort.append(end - start)

plt.plot(n_grid, time_sort, label="sort")
plt.plot(n_grid, time_nosort, label="no sort")
# plt.xscale('log')
# plt.yscale('log')
plt.legend()
plt.show()


### Point 4 ###
def generate_uniform_on_unit_simplex(num_samps, n):
    X = np.empty((num_samps, n))

    for i in range(num_samps):
        # Generate order statistics
        X_ord = generate_ord_stat_nosort_vec(n)
        X_old = X_ord.copy()
        X_ord[1:] = X_ord[1:] - X_ord[: n - 1]

        X[i, :] = X_ord

    return X


n = 3
num_samps = 1000
X = generate_uniform_on_unit_simplex(num_samps, n)

fig = plt.figure()
ax = fig.add_subplot(111, projection="3d")

# Defines the values
x = X[:, 0]
y = X[:, 1]
z = X[:, 2]

ax.scatter(x, y, z, c="r", marker="o")

ax.set_xlabel("X")
ax.set_ylabel("Y")
ax.set_zlabel("Z")
plt.show()