import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.mixture import GaussianMixture as GMM
from matplotlib.patches import Ellipse

fig = plt.figure(1,(10,8),dpi = 400)
# 用于将两个列表转为array的数据形式
def list_to_array(list1, list2):
    array1 = [[list1[i], list2[i]] for i in range(len(list1))]
    array1 = np.array(array1)
    return array1


# 给定的位置和协方差画一个椭圆
def draw_ellipse(position, covariance, ax=None, **kwargs):
    ax = ax or plt.gca()
    # 将协方差转换为主轴
    if covariance.shape == (2, 2):
        U, s, Vt = np.linalg.svd(covariance)
        angle = np.degrees(np.arctan2(U[1, 0], U[0, 0]))
        width, height = 2 * np.sqrt(s)
    else:
        angle = 0
        width, height = 2 * np.sqrt(covariance)

    # 画出椭圆
    for nsig in range(1, 4):
        ax.add_patch(Ellipse(position, nsig * width, nsig * height,
                             angle, **kwargs))


# 画图
def plot_gmm(gmm, X, label=True, ax=None):
    ax = ax or plt.gca()
    labels = gmm.fit(X).predict(X)
    if label:
        ax.scatter(X[:, 0], X[:, 1], c=labels, s=4, cmap='viridis', zorder=2)
    else:
        ax.scatter(X[:, 0], X[:, 1], s=4, zorder=2)
    ax.axis('equal')
    w_factor = 0.2 / gmm.weights_.max()
    for pos, covar, w in zip(gmm.means_, gmm.covariances_, gmm.weights_):
        draw_ellipse(pos, covar, alpha=w * w_factor)

def GMM_xy(d1):
    print(d1)
    xdata, ydata = list(d1['longitude']), list(d1['latitude'])
    datas = list_to_array(xdata, ydata)
    plt.scatter(datas[:, 0], datas[:, 1], 1)
    gmm10 = GMM(n_components=3, covariance_type='full', random_state=0)
    plot_gmm(gmm10, datas, label=False)
    plt.show()