import time

import numpy as np


class DNB:
    """
    Class representing Dynamic Naive Bayes.
    """

    def __init__(self, debug=False):
        self.states_prior = None
        self.states_list = None
        self.A = None
        self.B = None
        self.debug = debug

    def mle(self, df, state_col, features=None):
        t = time.process_time()
        """ Fitting dynamics in the DNB """
        self._dynamic_mle(df[state_col])
        """ Fitting observable variables """
        self.B = {}
        for st in self.states_list:
            self._features_mle(df[df[state_col]==st].drop([state_col],axis=1),st)
        if self.debug:
            elapsed_time = time.process_time() - t
            print("MLE finished in %d seconds." % elapsed_time)

    def _dynamic_mle(self, df):
        states_vec = df.as_matrix()
        self.states_list = np.unique(states_vec)
        states_nr = len(self.states_list)
        self.A = np.zeros((states_nr, states_nr))
        self.states_prior = np.zeros(states_nr)
        self.states_prior[np.searchsorted(self.states_list, states_vec[0])] += 1
        for i in range(1, len(states_vec)):
            self.A[np.searchsorted(self.states_list, states_vec[i - 1]), np.searchsorted(self.states_list,
                                                                                         states_vec[i])] += 1
            self.states_prior[np.searchsorted(self.states_list, states_vec[i])] += 1
        self.states_prior = self.states_prior / self.states_prior.sum()
        for i in range(states_nr):
            self.A[i] = self.A[i] / self.A[i].sum()

    def _features_mle(self, df, state):
        """simplified_version"""
        import scipy.stats as st
        features = list(df.columns.values)
        for f in features:
            params = st.norm.fit(df[f])
            arg = params[:-2]
            loc = params[-2]
            scale = params[-1]
            if self.debug:
                print("%s, %s, %s"%(str(arg),str(loc),str(scale)))
            self.B[(state, f)] = params

@@ -193,32 +193,66 @@
 tes_list = None%0A
+        self.features_list = None%0A
         self.A =
@@ -1753,24 +1753,29 @@
         
+self.
 features
  = list(
@@ -1766,16 +1766,21 @@
 features
+_list
  = list(
@@ -1815,24 +1815,29 @@
 or f in 
+self.
 features
 :%0A      
@@ -1828,16 +1828,21 @@
 features
+_list
 :%0A      
@@ -2092,15 +2092,637 @@
 , f)%5D = 
-param
+%5Bst.norm%5D + params%0A%0A    def _emission_prob(self,state,data):%0A        prob = 1%0A        for f in self.features_list:%0A            dist = self.B%5B(state,f)%5D%5B0%5D%0A            arg = self.B%5B(state,f)%5D%5B1:-2%5D%0A            loc = self.B%5B(state,f)%5D%5B-2%5D%0A            scale = self.B%5B(state,f)%5D%5B-1%5D%0A            prob *= dist.pdf(data.iloc%5B0%5D%5Bf%5D, loc=loc, scale=scale, *arg)%0A        return prob%0A%0A    def _transition_prob(self,state1, state2):%0A        return self.A%5Bnp.searchsorted(self.states_list, state1), np.searchsorted(self.states_list, state2)%5D%0A%0A    def viterbi(self, data):%0A        pass%0A%0A    def seq_probability(self, data, path):%0A        pas
 s%0A

add transition and emission probabilities