from __future__ import division, print_function

import os
import struct
import sys

import numpy as np
from scipy.stats import multivariate_normal

try:
    import sounddevice as sd

    have_sounddevice = True
except:
    have_sounddevice = False

from .acoustics import mfcc


class CircularGaussianEmission:
    def __init__(self, nstates, odim=1, examples=None):
        """Initialize the Gaussian emission object"""

        # The emissions parameters
        self.K = nstates

        if examples is None:
            # Initialize to random components
            self.O = odim
            self.mu = np.random.normal(size=(self.K, self.O))
            self.Sigma = np.ones((self.K, self.O)) * 10

        else:
            # Initialize all components to the same mean and variance of the data
            self.O = examples[0].shape[1]

            X = np.concatenate(examples, axis=0)

            self.mu = np.array([np.mean(X, axis=0)] * self.K)
            centered = X - self.mu[0]
            self.Sigma = np.array([np.mean(centered**2, axis=0)] * self.K)

    def update_parameters(self, examples, gamma):
        g = np.concatenate(gamma, axis=0)
        X = np.concatenate(examples, axis=0)
        Z = g.sum(axis=0)

        for k in range(self.K):
            self.mu[k] = np.sum(X.T * g[:, k], axis=1) / Z[k]
            centered = (X - self.mu[k]) ** 2
            self.Sigma[k] = np.sum(centered.T * g[:, k], axis=1) / Z[k]

    def get_pdfs(self):
        """Return the pdf of all the emission probabilities"""
        return [
            multivariate_normal(self.mu[k], np.diag(self.Sigma[k]))
            for k in range(self.K)
        ]

    def prob_x_given_state(self, examples):
        """
        Recompute the probability of the observation given the state of the
        latent variables
        """
        distribution = [
            multivariate_normal(self.mu[k], np.diag(self.Sigma[k]))
            for k in range(self.K)
        ]
        p_x_given_z = []

        for X in examples:
            p_x_given_z.append(np.zeros((X.shape[0], self.K)))
            for k in range(self.K):
                p_x_given_z[-1][:, k] = distribution[k].pdf(X)

        return p_x_given_z


class GaussianEmission:
    def __init__(self, nstates, odim=1, examples=None):
        """Initialize the Gaussian emission object"""

        # The emissions parameters
        self.K = nstates

        if examples is None:
            # initialize to random mean unit variance
            self.O = odim
            self.mu = np.random.normal(size=(self.K, self.O))
            self.Sigma = np.random.normal(size=(self.K, self.O, self.O))
            for k in range(self.K):
                self.Sigma[k] = np.dot(self.Sigma[k].T, self.Sigma[k]) + np.eye(self.O)

        else:
            # Initialize using mean and covariance of dataset
            self.O = examples[0].shape[1]

            X = np.concatenate(examples, axis=0)

            self.mu = np.array([np.mean(X, axis=0)] * self.K)
            centered = X - self.mu[0]
            self.Sigma = np.array([np.diag(np.mean(centered**2, axis=0))] * self.K)

    def update_parameters(self, examples, gamma):
        g = np.concatenate(gamma, axis=0)
        X = np.concatenate(examples, axis=0)
        Z = g.sum(axis=0)

        for k in range(self.K):
            self.mu[k] = np.sum(X.T * g[:, k], axis=1) / Z[k]
            centered = X - self.mu[k]
            self.Sigma[k] = np.dot(centered.T * g[:, k], centered / Z[k])

    def get_pdfs(self):
        """Return the pdf of all the emission probabilities"""
        return [multivariate_normal(self.mu[k], self.Sigma[k]) for k in range(self.K)]

    def prob_x_given_state(self, examples):
        """
        Recompute the probability of the observation given the state of the
        latent variables
        """
        distribution = [
            multivariate_normal(self.mu[k], self.Sigma[k]) for k in range(self.K)
        ]
        p_x_given_z = []

        for X in examples:
            p_x_given_z.append(np.zeros((X.shape[0], self.K)))
            for k in range(self.K):
                p_x_given_z[-1][:, k] = distribution[k].pdf(X)

        return p_x_given_z


class HMM:
    """
    Hidden Markov Model with Gaussian emissions

    Attributes
    ----------
    K : int
        Number of states in the model
    O : int
        Number of dimensions of the Gaussian emission distribution
    A : ndarray
        KxK transition matrix of the Markov chain
    pi : ndarray
        K dim vector of the initial probabilities of the Markov chain
    emission : (GaussianEmission or CircularGaussianEmission)
        An instance of emission_class
    model : string, optional
        The model used for the chain, can be 'full' or 'left-right'
    leftright_jum_max : int, optional
        The number of non-zero upper diagonals in a 'left-right' model
    """

    def __init__(self, nstates, emission, model="full", leftright_jump_max=3):
        """
        Initialize a Hidden Markov Model with nstates and Gaussian observations

        nstates: int
            The number of states in the Markov chain
        emission : emission object, optional
            The emission object (CircularGaussianEmission or GaussianEmission)
        model : string, optional
            The model used for the chain, can be 'full' or 'left-right'
        leftright_jump_max : int
            The maximum jump length in the Left-Right chain model
        """

        self.K = nstates  # number of states
        self.emission = emission  # The observation parameters

        # The Markov chain parameters
        self.model = model

        self.leftright_jump_max = leftright_jump_max
        self.A = np.zeros((self.K, self.K))  # the state transition matrix
        self.pi = np.zeros((self.K))  # the initial distribution

        # Initialize the HMM parameters to some random values
        if self.model == "full":
            self.A = np.random.uniform(size=(self.K, self.K))
            self.pi = np.random.uniform(size=(self.K))

        elif self.model == "left-right":
            self.A = np.triu(
                np.tril(
                    np.random.uniform(size=(self.K, self.K)), k=self.leftright_jump_max
                )
            )
            self.A += np.diag(np.sum(self.A[:, :], axis=1) * 2)
            self.pi = np.zeros(self.K)
            self.pi[0] = 1

        # Normalize the distributions
        for row in self.A:
            row /= row.sum()
        self.pi /= self.pi.sum()

    def fit(self, examples, tol=0.1, max_iter=10, verbose=False):
        """
        Training of the HMM using the EM algorithm

        Parameters
        ----------
        examples : (list)
            A list of examples used to train the model. Each example is
            an array of feature vectors, each row is a feature vector,
            the sequence runs on axis 0
        tol : (float)
            The training stops when the progress between to steps is less than
            this number (default 0.1)
        max_iter : (int)
            Alternatively the algorithm stops when a maximum number of
            iterations is reached (default 10)
        verbose : bool, optional
            When True, prints extra information about convergence
        """

        # Make sure to normalize parameters that should be...
        for row in self.A:
            row[:] /= row.sum()
        self.pi[:] /= self.pi.sum()

        # Run the EM algorithm
        loglikelihood_old = -np.inf  # log-likelihood
        n_iter = 0
        while True:
            # Initialize new parameters value for accumulation
            loglikelihood = 0.0

            # We need to run the forward/backward algorithm for each example and
            # and combine the result to form the new estimates
            gamma = []
            xhi = []
            p_x_given_z = self.emission.prob_x_given_state(examples)

            # Expectation-step
            # -----------------

            for X, pxz in zip(examples, p_x_given_z):
                # check dimension of emission
                if X.shape[1] != self.emission.O:
                    raise ValueError(
                        "Error: Emission vectors of all examples should have the same size"
                    )

                # First compute alpha and beta using forward/backward algo
                alpha, c = self.forward(X, pxz)
                beta = self.backward(X, pxz, c)

                # Recompute the likelihood of the sequence
                # (Bishop 13.63)
                loglikelihood += np.sum(np.log(c))

                # Now the more interesting quantities
                # gamma(z_n) = p(z_n | X, theta_old)
                # xhi(z_{n-1}, z_n) = p(z_{n-1}, z_n | X, theta_old)
                gamma.append(alpha * beta)
                xhi.append(np.zeros((X.shape[0] - 1, self.K, self.K)))
                for n in range(1, X.shape[0]):
                    xhi[-1][n - 1] = (
                        np.outer(alpha[n - 1], beta[n] * pxz[n]) * self.A / c[n]
                    )

            # Maximization-step
            # ------------------

            # update the Markov Chain parameters
            self.update_parameters(examples, gamma, xhi)

            # Update the emission distribution parameters
            self.emission.update_parameters(examples, gamma)

            # Now check for convergence
            # --------------------------
            n_iter += 1
            epsilon = loglikelihood - loglikelihood_old
            if verbose:
                print(
                    "Iterations:", n_iter, "epsilon:", epsilon, "LL_new:", loglikelihood
                )

            # some checks here
            if epsilon < tol:
                if verbose:
                    print("Tolerance reached: stopping.")
                break
            if n_iter == max_iter:
                if verbose:
                    print("Maximum iterations reached: stopping.")
                break

            loglikelihood_old = loglikelihood

        # return the number of iterations performed
        return n_iter

    def update_parameters(self, examples, gamma, xhi):
        """Update the parameters of the Markov Chain"""

        X = np.concatenate(examples, axis=0)
        x = np.concatenate(xhi, axis=0)

        self.pi[:] = np.sum([g[0, :] for g in gamma], axis=0)
        self.A = x.sum(axis=0)

        # normalize to enforce distribution constraints
        self.pi /= np.sum(self.pi)
        for k in range(self.K):
            den = np.sum(self.A[k, :])
            if den < 1e-15:
                self.A[k, :] = 0.0
            else:
                self.A[k, :] /= den

    def generate(self, N):
        """Generate a random sample of length N using the model"""
        X = np.zeros((N, self.emission.O))
        distributions = self.emission.get_pdfs()

        # pick initial state
        state = np.random.choice(self.K, p=self.pi)

        # now run the chain
        for n in range(0, N):
            # produce emission vector according to current state
            X[n, :] = distributions[state].rvs()
            # pick next state
            state = np.random.choice(self.K, p=self.A[state, :])

        return X

    def loglikelihood(self, X):
        """
        Compute the log-likelihood of a sample vector using the sum-product algorithm
        """
        p_x_given_z = self.emission.prob_x_given_state([X])[0]
        alpha, c = self.forward(X, p_x_given_z)

        return np.sum(np.log(c))

    def forward(self, X, p_x_given_z):
        """The forward recursion for HMM as described in Bishop Ch. 13"""

        # initialize the alpha vector
        alpha = np.zeros((X.shape[0], self.K))
        c = np.zeros(X.shape[0])

        # initialize the recursion as
        # p(X | z_k) pi_k
        alpha[0] = p_x_given_z[0] * self.pi
        c[0] = np.sum(alpha[0])
        alpha[0] /= c[0]

        # Run the forward recursion
        for n in range(1, X.shape[0]):
            alpha[n] = p_x_given_z[n] * np.dot(self.A.T, alpha[n - 1])
            c[n] = np.sum(alpha[n])
            alpha[n] /= c[n]

        return alpha, c

    def backward(self, X, p_x_given_z, c):
        """The backward recursion for HMM as described in Bishop Ch. 13"""

        # intialize the beta vectors
        beta = np.zeros((X.shape[0], self.K))

        # initialize the recursion
        beta[-1, :] = 1

        # Run the backward recursion
        for n in range(X.shape[0] - 2, -1, -1):
            beta[n] = np.dot(self.A, p_x_given_z[n + 1] * beta[n + 1]) / c[n + 1]

        return beta

    def viterbi(self):
        x = 1
