from __future__ import print_function
import json
import numpy as np


class HMM:

    def __init__(self, pi, A, B, obs_dict, state_dict):
        """
        - pi: (1*num_state) A numpy array of initial probabilities. pi[i] = P(Z_1 = s_i)
        - A: (num_state*num_state) A numpy array of transition probabilities. A[i, j] = P(Z_t = s_j|Z_{t-1} = s_i)
        - B: (num_state*num_obs_symbol) A numpy array of observation probabilities. B[i, k] = P(X_t = o_k| Z_t = s_i)
        - obs_dict: A dictionary mapping each observation symbol to its index 
        - state_dict: A dictionary mapping each state to its index
        """
        self.pi = pi
        self.A = A
        self.B = B
        self.obs_dict = obs_dict
        self.state_dict = state_dict

    def forward(self, Osequence):
        """
        Inputs:
        - Osequence: (1*L) A numpy array of observation sequence with length L

        Returns:
        - alpha: (num_state*L) A numpy array where alpha[i, t-1] = P(Z_t = s_i, X_{1:t}=x_{1:t})
                 (note that this is alpha[i, t-1] instead of alpha[i, t])
        """
        S = len(self.pi)
        L = len(Osequence)
        O = self.find_item(Osequence)
        alpha = np.zeros([S, L])
        ######################################################
        # TODO: compute and return the forward messages alpha
        ######################################################
        for t in range(L):
            if t == 0:
                Ot=self.obs_dict[Osequence[t]]
                alpha[:,t]=self.pi*self.B[:,Ot]
                
            else:
                Ot=self.obs_dict[Osequence[t]]
                sum1=np.dot(alpha[:,t-1],self.A)
                alpha[:,t]=sum1*self.B[:,Ot]

        return alpha;

    def backward(self, Osequence):
        """
        Inputs:
        - Osequence: (1*L) A numpy array of observation sequence with length L

        Returns:
        - beta: (num_state*L) A numpy array where beta[i, t-1] = P(X_{t+1:T}=x_{t+1:T} | Z_t = s_i)
                    (note that this is beta[i, t-1] instead of beta[i, t])
        """
        S = len(self.pi)
        L = len(Osequence)
        O = self.find_item(Osequence)
        beta = np.zeros([S, L])
        #######################################################
        # TODO: compute and return the backward messages beta
        #######################################################
        for i in range(S):
            beta[i,L-1]=1
        for t in range(L-2,-1,-1):
            Ot=self.obs_dict[Osequence[t+1]]
            beta[:,t]=np.dot(self.A,self.B[:,Ot]*beta[:,t+1])
            
        return beta
            

    def sequence_prob(self, Osequence):
        """
        Inputs:
        - Osequence: (1*L) A numpy array of observation sequence with length L

        Returns:
        - prob: A float number of P(X_{1:T}=x_{1:T})
        """
        
        #####################################################
        # TODO: compute and return prob = P(X_{1:T}=x_{1:T})
        #   using the forward/backward messages
        #####################################################
        L=len(Osequence)
        alpha=self.forward(Osequence)
        prob=0
        S = len(self.pi)
        for i in range(S):
            prob+=alpha[i][L-1];
        return prob
        
        

    def posterior_prob(self, Osequence):
        """
        Inputs:
        - Osequence: (1*L) A numpy array of observation sequence with length L

        Returns:
        - gamma: (num_state*L) A numpy array where gamma[i, t-1] = P(Z_t = s_i | X_{1:T}=x_{1:T})
		           (note that this is gamma[i, t-1] instead of gamma[i, t])
        """
        ######################################################################
        # TODO: compute and return gamma using the forward/backward messages
        ######################################################################
        alpha=self.forward(Osequence)
        beta=self.backward(Osequence)
        prob=self.sequence_prob(Osequence)
        gamma=alpha*beta/prob
        return gamma

    
    def likelihood_prob(self, Osequence):
        """
        Inputs:
        - Osequence: (1*L) A numpy array of observation sequence with length L

        Returns:
        - prob: (num_state*num_state*(L-1)) A numpy array where prob[i, j, t-1] = 
                    P(Z_t = s_i, Z_{t+1} = s_j | X_{1:T}=x_{1:T})
        """
        S = len(self.pi)
        L = len(Osequence)
        prob = np.zeros([S, S, L - 1])
        #####################################################################
        # TODO: compute and return prob using the forward/backward messages
        #####################################################################
        alpha=self.forward(Osequence)
        beta=self.backward(Osequence)
        p=self.sequence_prob(Osequence)
        for s in range(S):
            pp=prob[s]
            for t in range(L-1):
                Ot=self.obs_dict[Osequence[t+1]]
                pp[:,t]=alpha[s][t]*self.A[s]*self.B[:,Ot]*beta[:,t+1]/p
        return prob
        

    def viterbi(self, Osequence):
        """
        Inputs:
        - Osequence: (1*L) A numpy array of observation sequence with length L

        Returns:
        - path: A List of the most likely hidden states (return actual states instead of their indices;
                    you might find the given function self.find_key useful)
        """
        path = []
        ################################################################################
        # TODO: implement the Viterbi algorithm and return the most likely state path
        ################################################################################
        
        S = len(self.pi)
        L = len(Osequence)
        delta=np.zeros([S,L])
        bigpath=np.zeros([S,L],dtype=int)
        delta[:,0]=self.pi*self.B[:,self.obs_dict[Osequence[0]]]
        for i in range(1,L):
            delta[:,i]=self.B[:,self.obs_dict[Osequence[i]]]*(self.A*delta[:,i-1].reshape(S,1)).max(axis=0)
            bigpath[:,i]=np.argmax(self.A*delta[:,i-1].reshape(S,1),axis=0)
        tpath=np.argmax(delta[:,L-1])
        path.append(tpath)
        for i in range(L-1,0,-1):
            state=bigpath[:,i]
            tpath=state[tpath]
            path.append(tpath)
        path.reverse()
        for i in range(L):
            path[i]=list(self.state_dict.keys())[list(self.state_dict.values()).index(path[i])]
        
        return path
        

    #DO NOT MODIFY CODE BELOW
    def find_key(self, obs_dict, idx):
        for item in obs_dict:
            if obs_dict[item] == idx:
                return item

    def find_item(self, Osequence):
        O = []
        for item in Osequence:
            O.append(self.obs_dict[item])
        return O
