# -*- coding: utf-8 -*-#
'''
# Name:         GMM
# Description:  高斯混合模型实现
# Author:       super
# Date:         2020/5/10
'''

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Ellipse
from scipy.stats import multivariate_normal
plt.style.use('seaborn')

def generate_X(true_Mu, true_Var):
    '''
    生成三个高斯分布数据
    :param true_Mu: 均值
    :param true_Var: 方差
    :return:
    '''
    # 第一簇的数据
    num1, mu1, var1 = 400, true_Mu[0], true_Var[0]
    X1 = np.random.multivariate_normal(mu1, np.diag(var1), num1)
    # 第二簇的数据
    num2, mu2, var2 = 600, true_Mu[1], true_Var[1]
    X2 = np.random.multivariate_normal(mu2, np.diag(var2), num2)
    # 第三簇的数据
    num3, mu3, var3 = 1000, true_Mu[2], true_Var[2]
    X3 = np.random.multivariate_normal(mu3, np.diag(var3), num3)
    # 合并在一起
    X = np.vstack((X1, X2, X3))
    # 显示数据
    plt.figure(figsize=(10, 8))
    plt.axis([-10, 15, -5, 15])
    plt.scatter(X1[:, 0], X1[:, 1], s=5)
    plt.scatter(X2[:, 0], X2[:, 1], s=5)
    plt.scatter(X3[:, 0], X3[:, 1], s=5)
    plt.show()
    return X

# E步骤更新W，也就是第i个变量属于第m簇的概率
# 更新W
def update_W(X, Mu, Var, Pi):
    n_points, n_clusters = len(X), len(Pi)
    pdfs = np.zeros(((n_points, n_clusters)))
    for i in range(n_clusters):
        # multivariate_normal.pdf：多元正态分布的概率密度函数
        pdfs[:, i] = Pi[i] * multivariate_normal.pdf(X, Mu[i], np.diag(Var[i]))
    W = pdfs / pdfs.sum(axis=1).reshape(-1, 1)
    return W

# 根据更新的W，更新每一簇的占比
# 更新pi
def update_Pi(W):
    Pi = W.sum(axis=0) / W.sum()
    return Pi


# 计算log似然函数
def logLH(X, Pi, Mu, Var):
    n_points, n_clusters = len(X), len(Pi)
    pdfs = np.zeros(((n_points, n_clusters)))
    for i in range(n_clusters):
        pdfs[:, i] = Pi[i] * multivariate_normal.pdf(X, Mu[i], np.diag(Var[i]))
    return np.mean(np.log(pdfs.sum(axis=1)))


# 画出聚类图像
def plot_clusters(X, Mu, Var, Mu_true=None, Var_true=None):
    colors = ['b', 'g', 'r']
    n_clusters = len(Mu)
    plt.figure(figsize=(10, 8))
    plt.axis([-10, 15, -5, 15])
    plt.scatter(X[:, 0], X[:, 1], s=5)
    ax = plt.gca()
    for i in range(n_clusters):
        plot_args = {'fc': 'None', 'lw': 2, 'edgecolor': colors[i], 'ls': ':'}
        ellipse = Ellipse(Mu[i], 3 * Var[i][0], 3 * Var[i][1], **plot_args)
        ax.add_patch(ellipse)
    if (Mu_true is not None) & (Var_true is not None):
        for i in range(n_clusters):
            plot_args = {'fc': 'None', 'lw': 2, 'edgecolor': colors[i], 'alpha': 0.5}
            ellipse = Ellipse(Mu_true[i], 3 * Var_true[i][0], 3 * Var_true[i][1], **plot_args)
            ax.add_patch(ellipse)
    plt.show()

# M步根据更新的W和PI来跟新均值Mu与方差Var
# 更新Mu
def update_Mu(X, W):
    n_clusters = W.shape[1]
    Mu = np.zeros((n_clusters, 2))
    for i in range(n_clusters):
        Mu[i] = np.average(X, axis=0, weights=W[:, i])
    return Mu


# 更新Var
def update_Var(X, Mu, W):
    n_clusters = W.shape[1]
    Var = np.zeros((n_clusters, 2))
    for i in range(n_clusters):
        Var[i] = np.average((X - Mu[i]) ** 2, axis=0, weights=W[:, i])
    return Var


if __name__ == '__main__':
    # 生成数据
    true_Mu = [[0.5, 0.5], [5.5, 2.5], [1, 7]]
    true_Var = [[1, 3], [2, 2], [6, 2]]
    X = generate_X(true_Mu, true_Var)
    # 初始化
    n_clusters = 3 #聚类的个数
    n_points = len(X)
    Mu = [[0, -1], [6, 0], [0, 9]]
    Var = [[1, 1], [1, 1], [1, 1]]
    Pi = [1 / n_clusters] * 3
    W = np.ones((n_points, n_clusters)) / n_clusters #隐变量
    Pi = W.sum(axis=0) / W.sum() #每一簇的比重，可以根据W求得
    # 迭代
    loglh = []
    for i in range(5):
        plot_clusters(X, Mu, Var, true_Mu, true_Var)
        loglh.append(logLH(X, Pi, Mu, Var))
        W = update_W(X, Mu, Var, Pi)
        Pi = update_Pi(W)
        Mu = update_Mu(X, W)
        print('log-likehood:%.3f'%loglh[-1])
        Var = update_Var(X, Mu, W)

# -*- coding: utf-8 -*-#
'''
# Name:         normal_show
# Description:  显示高斯混合模型的一些信息
# Author:       super
# Date:         2020/5/10
'''

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Ellipse
from scipy.stats import multivariate_normal
plt.style.use('seaborn')

# 画出聚类图像
def plot_clusters(X, Mu_true, Var_true, ang = 0):
    colors = ['b', 'g', 'r', 'c']
    n_clusters = len(Mu_true)
    plt.figure(figsize=(10, 8))
    plt.axis([-10, 15, -5, 15])
    ax = plt.gca()
    for i in range(n_clusters):
        plt.scatter(X[i][:, 0], X[i][:, 1], s=5, c=colors[i])
        plot_args = {'fc': 'None', 'lw': 2, 'edgecolor': colors[i], 'alpha': 0.5}
        if i == 3:
            ellipse = Ellipse(Mu_true[i], 3 * Var_true[i][0], 3 * Var_true[i][1], angle=20, **plot_args)
        else:
            ellipse = Ellipse(Mu_true[i], 3 * Var_true[i][0], 3 * Var_true[i][1], angle=ang,**plot_args)
        ax.add_patch(ellipse)
    plt.show()

# 画出聚类图像
def plot_cluster(X, Mu_true, Var_true, color ,ang = 0):
    plt.figure(figsize=(10, 8))
    plt.axis([-10, 15, -5, 15])
    ax = plt.gca()
    plt.scatter(X[:, 0], X[:, 1], s=5, c=color)
    plot_args = {'fc': 'None', 'lw': 2, 'edgecolor': color, 'alpha': 0.5}
    ellipse = Ellipse(Mu_true, 3 * Var_true[0], 3 * Var_true[1], angle=ang, **plot_args)
    ax.add_patch(ellipse)
    plt.show()

if __name__ == "__main__":
    nums = [400, 600, 1000, 500]
    true_Mu = [[0.5, 0.5], [5.5, 2.5], [1, 7], [9, 4.5]]
    true_Var = [[1, 3], [2, 2], [6, 2], [1, 3]]
    var = [np.diag(true_Var[0]), np.diag(true_Var[1]),np.diag(true_Var[2]),np.array([[1,-1],[-1,3]])]
    print(np.diag(true_Var[0]))
    print(type(np.diag(true_Var[0])))
    a = np.array([[1,2],[3,4]])
    print(a)
    print(type(a))
    # 第一簇的数据
    num1, mu1, var1 = nums[0], true_Mu[0], true_Var[0]
    X1 = np.random.multivariate_normal(mu1, var[0], num1)
    plot_cluster(X1, [0.5, 0.5], [1, 3], 'b')
    # 第二簇的数据
    num2, mu2, var2 = nums[1], true_Mu[1], true_Var[1]
    X2 = np.random.multivariate_normal(mu2, var[1], num2)
    plot_cluster(X2, [5.5, 2.5], [2,2], 'g')
    # 第三簇的数据
    num3, mu3, var3 = nums[2], true_Mu[2], true_Var[2]
    X3 = np.random.multivariate_normal(mu3, var[2], num3)
    plot_cluster(X3, [1, 7], [6, 2], 'r')
    # 第四簇的数据
    num4, mu4, var4 = nums[3], true_Mu[3], true_Var[3]
    X4 = np.random.multivariate_normal(mu4, var[3], num4)
    plot_cluster(X4, [9, 4.5], [1, 3], 'c' ,25)
    X = [X1, X2, X3, X4]

    # plot_clusters(X1, X2, X3, X4, true_Mu, true_Var)
    plot_clusters(X, true_Mu, true_Var)