import numpy as np
from sklearn.mixture import GaussianMixture, BayesianGaussianMixture
import matplotlib.pyplot as plt
from math import sin, cos, atan2, pi

npts = 100

def rot(theta):
    return np.array([[cos(theta), -sin(theta)], [sin(theta), cos(theta)]])

print("Type 'q' in the graphics window to proceed to the next trial.")

while True:
    sigmas = 1 + np.random.rand(3,2)*5
    phis = np.random.rand(3)*2*pi

    d1 = np.random.randn(npts,2)
    d1[:,0] *= sigmas[0,0]
    d1[:,1] *= sigmas[0,1]
    d1 = d1.dot(rot(phis[0]))
    d1[:,0] += 10 + np.random.randint(-10,10)
    d1[:,1] += 5 + np.random.randint(-10,10)

    d2 = np.random.randn(npts,2) * sigmas[2]
    d2 = d2.dot(rot(phis[1]))
    d2[:,0] += 20 + np.random.randint(-10,10)
    d2[:,1] += 30 + np.random.randint(-10,10)

    d3 = np.random.randn(npts,2) * sigmas[1]
    d3 = d3.dot(rot(phis[2]))
    d3[:,0] += 30 + np.random.randint(-10,10)
    d3[:,1] += -5 + np.random.randint(-10,10)

    data = np.concatenate((d1,d2,d3))

    plt.plot(data[:,0],data[:,1],'.')

    gmm = GaussianMixture(n_components=3)
    #gmm = BayesianGaussianMixture(n_components=7)
    gmm.fit(data)

    means = gmm.means_
    covariances = gmm.covariances_

    circle = np.array([[cos(x*pi/20), sin(x*pi/20)] for x in range(41)])
    for i in range(len(means)):
        pts = circle + 0
        maxvar = max(covariances[i,0,0], covariances[i,1,1])
        minvar = min(covariances[i,0,0], covariances[i,1,1])
        pts[:,0] *= 2*maxvar**0.5
        pts[:,1] *= 2*minvar**0.5
        val, vec = np.linalg.eig(covariances[i,:,:])
        principal_eig = vec[:,0] if abs(val[0]) > abs(val[1]) else vec[:,1]
        theta = -atan2(principal_eig[1], principal_eig[0])
        pts = pts.dot(rot(theta))
        pts[:,0] += means[i,0]
        pts[:,1] += means[i,1]
        plt.plot(pts[:,0], pts[:,1], 'y')
    #plt.gca().set_aspect('equal')

    plt.show()