"""Mixture model for collaborative filtering"""
from typing import NamedTuple, Tuple
import numpy as np
from matplotlib import pyplot as plt
from matplotlib.patches import Circle, Arc


class GaussianMixture(NamedTuple):
    """Tuple holding a gaussian mixture"""
    mu: np.ndarray  # (K, d) array - each row corresponds to a gaussian component mean
    var: np.ndarray  # (K, ) array - each row corresponds to the variance of a component
    p: np.ndarray  # (K, ) array = each row corresponds to the weight of a component


def init(X: np.ndarray, K: int,
         seed: int = 0) -> Tuple[GaussianMixture, np.ndarray]:
    """Initializes the mixture model with random points as initial
    means and uniform assignments

    Args:
        X: (n, d) array holding the data
        K: number of components
        seed: random seed

    Returns:
        mixture: the initialized gaussian mixture
        post: (n, K) array holding the soft counts
            for all components for all examples

    """
    np.random.seed(seed)
    n, _ = X.shape
    p = np.ones(K) / K

    # select K random points as initial means
    mu = X[np.random.choice(n, K, replace=False)]
    var = np.zeros(K)
    # Compute variance
    for j in range(K):
        var[j] = ((X - mu[j]) ** 2).mean()

    mixture = GaussianMixture(mu, var, p)
    post = np.ones((n, K)) / K

    return mixture, post


def plot(X: np.ndarray, mixture: GaussianMixture, post: np.ndarray,
         title: str):
    """Plots the mixture model for 2D data"""
    _, K = post.shape

    percent = post / post.sum(axis=1).reshape(-1, 1)
    fig, ax = plt.subplots()
    ax.title.set_text(title)
    ax.set_xlim((-20, 20))
    ax.set_ylim((-20, 20))
    r = 0.25
    color = ["r", "b", "k", "y", "m", "c"]
    for i, point in enumerate(X):
        theta = 0
        for j in range(K):
            offset = percent[i, j] * 360
            arc = Arc(point,
                      r,
                      r,
                      0,
                      theta,
                      theta + offset,
                      edgecolor=color[j])
            ax.add_patch(arc)
            theta += offset
    for j in range(K):
        mu = mixture.mu[j]
        sigma = np.sqrt(mixture.var[j])
        circle = Circle(mu, sigma, color=color[j], fill=False)
        ax.add_patch(circle)
        legend = "mu = ({:0.2f}, {:0.2f})\n stdv = {:0.2f}".format(
            mu[0], mu[1], sigma)
        ax.text(mu[0], mu[1], legend)
    plt.axis('equal')
    plt.show()


def rmse(X, Y):
    return np.sqrt(np.mean((X - Y) ** 2))


def bic(X: np.ndarray, mixture: GaussianMixture,
        log_likelihood: float) -> float:
    """Computes the Bayesian Information Criterion for a
    mixture of Gaussians

    Args:
        X: (n, d) array holding the data
        mixture: a mixture of spherical gaussian
        log_likelihood: the log-likelihood of the data

    Returns:
        float: the BIC for this mixture
    """
    n, _ = X.shape
    K, d = mixture.mu.shape  # len(mixture.var) and len(mixture.p) = K
    p = (d * K) + K + (K - 1)

    BIC = log_likelihood - (p * np.log(n)) / 2

    return BIC


def multi_plot(X: np.ndarray, mixtureG: GaussianMixture, mixtureK: GaussianMixture,
               postG: np.ndarray, postK: np.ndarray, titleG: str, titleK: str):
    """Plots the mixture model comparison for 2D data"""
    _, K = postG.shape
    _, K = postK.shape

    percentG = postG / postG.sum(axis=1).reshape(-1, 1)
    percentK = postK / postK.sum(axis=1).reshape(-1, 1)
    fig, ax = plt.subplots(2, sharey=True)
    ax[0].title.set_text(titleG)
    ax[1].title.set_text(titleK)
    ax[0].set_xlim((-20, 20))
    ax[1].set_xlim((-20, 20))
    ax[0].set_ylim((-20, 20))
    ax[1].set_ylim((-20, 20))
    r = 0.25
    color = ["r", "b", "k", "y", "m", "c"]
    for i, point in enumerate(X):
        thetaG = 0
        thetaK = 0
        for j in range(K):
            offsetG = percentG[i, j] * 360
            offsetK = percentK[i, j] * 360
            arcG = Arc(point,
                       r,
                       r,
                       0,
                       thetaG,
                       thetaG + offsetG,
                       edgecolor=color[j])
            arcK = Arc(point,
                       r,
                       r,
                       0,
                       thetaK,
                       thetaK + offsetK,
                       edgecolor=color[j])
            ax[0].add_patch(arcG)
            ax[1].add_patch(arcK)
            thetaG += offsetG
            thetaK += offsetK

    for j in range(K):
        muG = mixtureG.mu[j]
        muK = mixtureK.mu[j]
        sigmaG = np.sqrt(mixtureG.var[j])
        sigmaK = np.sqrt(mixtureK.var[j])
        circleG = Circle(muG, sigmaG, color=color[j], fill=False)
        circleK = Circle(muK, sigmaK, color=color[j], fill=False)
        ax[0].add_patch(circleG)
        ax[1].add_patch(circleK)
        legendG = "mu = ({:0.2f}, {:0.2f})\n stdv = {:0.2f}".format(
            muG[0], muG[1], sigmaG)
        legendK = "mu = ({:0.2f}, {:0.2f})\n stdv = {:0.2f}".format(
            muK[0], muK[1], sigmaK)
        ax[0].text(muG[0], muG[1], legendG)
        ax[1].text(muK[0], muK[1], legendK)
    plt.axis('equal')
    plt.show()
