import numpy as np
import matplotlib.pyplot as plt
from scipy.spatial.distance import cdist
from mpl_toolkits.mplot3d import Axes3D
import random
from sklearn.datasets import make_spd_matrix
import math 
from sklearn.datasets.samples_generator import make_blobs

def uniformRandomSum(s, n):
    """
    Uniformly samples a combination from all possible combinations of n numbers that sum to s. Avoids
    the non-uniform result produced by scaling plain draws from a uniform distribution. 

    Args:
        s (int): the sum
        n (int): the number of component numbers
    
    Returns:
        res (numpy array): a list of n numbers that sum to s
    """
    boundaries = random.sample(range(s), n-1)
    boundaries.sort()
    boundariesMaxes = list(boundaries) + [s]
    boundariesMins = [0] + list(boundaries)
    res = [boundaryMax-boundaryMin for boundaryMax, boundaryMin in zip(boundariesMaxes, boundariesMins)]
    return(np.array(res))




def gaussianClusterParams(c, k, n, minMean, maxMean, minVar, maxVar, stretch="circle"):
    """
    genClusterParameters: given a number of clusters c, data dimensionality k, and number of observations n
                          randomly generate mean and variance matrices for each cluster as well as counts
                          of observations to draw from each distribution s.t. the counts sum to n. 

    Args:
        c (int) number of clusters
        k (int): dimensionality of data
        n (int): total number of data points
        minMean (int): min value of cluster means
        maxMean (int): max value of cluster means
        minVar (int): min value of covariance matrix
        maxVar (int): max value of covariance matrix
        stretch (str): valid inputs = {"circle", "ellipse"}. Specifies shape distributions; 
                       i.e. circle implies the off diagonal values of the covariance matrices by 0

    Returns:
        M (numpy array): c by k matrix specifying the cluster means
        V (numpy array): c by k by k matrix specifying the covariance matrix for each Gaussian
        N (numpy array): c by 1 vector specifying number of data points generated by each distribution
    """
    M = np.zeros((c, k))
    V = np.zeros((c, k, k))
    
    N = uniformRandomSum(n, c)
    
    M = np.random.rand(c, k)
    M = M*(abs(maxMean)+abs(minMean))-minMean # Scale 
    
    # Randomly draw scalars in variance range. Expand to c by k by k matrix. Multiply by identity of same shape.
    if (stretch=="circle"): 
        I = np.tile(np.identity(k), reps=(c, 1, 1))
        V = np.random.uniform(low=minVar, high=maxVar, size=c)
        V = np.tile(V, reps=(k, k, 1))
        V = I*np.swapaxes(V, 2, 0)

    # Randomly generate symmetric matrix with values in range
    elif (stretch=="ellipse"):
        V = np.zeros((c, k, k))
        for i in range(c):
            V[i, :, :] = make_spd_matrix(k)*maxVar
    
    else:
        print("Invalid stretch setting")
        return

    return M, V, N


def gaussianClusters(M, V, N, random_colors=None, plot=True):
    """
    gaussianClustersSpecific: takes in parameters specifying multivariate Gaussian distributions. Draws the specified
                number of observations from each distribution and returns the combined, shuffled data.

                c: num clusters
                k: num features
                n: number of data points
    Args:
        M (numpy array): c by k matrix specifying the cluster means
        V (numpy array): c by k by k matrix specifying the covariance matrix for each Gaussian
        N (numpy array): c by 1 vector specifying number of data points generated by each distribution

    Returns:
        X (numpy array): n by k matrix of shuffled data points
    """
    # Housekeeping
    n = np.sum(N)
    c, k = np.shape(M)
    X = np.zeros((n, k))
    
    if not random_colors: random_colors = np.random.random((c, 3))
    colors = np.zeros((n, 3))
    
    # Draw from each distribution
    index = 0
    for i in range(len(N)):
        group_n = N[i]
        nextData = np.random.multivariate_normal(M[i], V[i], N[i])
        X[index:(index+group_n)] = nextData
        colors[index:(index+group_n), :] = np.tile(random_colors[i, :], reps=(group_n, 1))
        index += group_n
    
    # Display
    if plot:
        if k==1:
            print(X.shape)
            plt.hist(X[:,0], bins=math.ceil(np.amax(X)))
            plt.show() 
        elif k==2:
            plt.scatter(X[:,0], X[:,1], c=colors, alpha=.5)
            plt.show()            
        elif k==3:
            fig = plt.figure()
            ax = fig.add_subplot(111, projection='3d')
            ax.scatter(X[:,0], X[:,1], X[:,2], c=colors)
            plt.show()
        
        else:
            print("Does not display " + str(k) + " dimensional data") 

    # Mix and return
    np.random.shuffle(X)
    return(X)


if __name__ == '__main__':
    M, V, N = gaussianClusterParams(c=3, k=2, n=100, minMean=0, maxMean=10, minVar=0, maxVar=1, stretch="circle")
    print(M)
    print(V)
    print(N)
    X = gaussianClusters(M, V, N)