import numpy as np
import pandas as pd
import time
import scipy
from IPython.display import display, HTML


letters = ['e','t','a','i','n','o','s','h','r','d']
feature_params = np.fromfile("./model/feature-params.txt", sep=" ")
feature_params = feature_params.reshape((len(letters), 321))
wf = feature_params

transition_params = np.fromfile("./model/transition-params.txt", sep=" ")
transition_params = transition_params.reshape((len(letters), len(letters)))
wt = transition_params

train_words = np.genfromtxt("./data/train_words.txt", dtype='str')
test_words = np.genfromtxt("./data/test_words.txt", dtype='str')

def train_img_feat(i):
    r = np.fromfile("./data/train_img"+str(i+1)+".txt", sep=" ")
    r = r.reshape(len(train_words[i]), 321)
    return r

def test_img_feat(i):
    r = np.fromfile("./data/test_img"+str(i+1)+".txt", sep=" ")
    r = r.reshape(len(test_words[i]), 321)
    return r

def psi_y_given_x(y, x):
    return np.dot(wf[letters.index(y)], x.T)

def psi_y1_y2(y1, y2):
    return wt[letters.index(y1), letters.index(y2)]

def messages(x, wf_d, wt_d):
    m = dict()
    n=len(letters)
    m[len(x)+1,len(x)] = np.zeros(n)
    m[0,1] = np.zeros(n)

    for t in range(len(x), 1, -1):
        s = t-1
        m[t, s] = np.zeros(n)

        for xs in range(n):
            lse = []
            for xt in range(n):
                lse.append(  np.dot(wf_d[xt], x[t-1].T) + wt_d[xs][xt] +  m[t+1,t][xt] )
            m [t,s][xs] = scipy.misc.logsumexp(lse)


    for t in range(1, len(x) , 1):
        s = t+1
        m[t,s] = np.zeros(n)
        for xs in range(n):
            lse = []
            for xt in range(n):
                lse.append( np.dot(wf_d[xt], x[t-1].T) + wt_d[xs][xt] +  m[t-1,t][xt] )
            m [t,s][xs] = scipy.misc.logsumexp(lse)

    return m

def get_Z(t, x, wf_d, wt_d):
    m = messages(x, wf_d, wt_d)
    sum = 0
    lse = []
    for xt in range(len(letters)):
        lse.append(np.dot(wf_d[xt], x[t-1].T) + m[t-1,t][xt] + m[t+1,t][xt])
    return scipy.misc.logsumexp(lse)

def marginals(x,wf_d,wt_d):
    m = messages(x,wf_d,wt_d)
    marg = dict()
    Z = get_Z(1, x, wf_d, wt_d)
    for pos in range(1,len(x)+1):
        marg[pos] = np.zeros(len(letters))
        for k in range(len(letters)):
            marg[pos][k] = np.exp(np.dot(wf_d[k], x[pos-1].T) + m[pos-1,pos][k] + m[pos+1,pos][k] - Z)
    pairwise_marg=dict()

    for i in range(1,len(x)):
        j = i+1
        pairwise_marg[i,j] = dict()
        for k1 in range(len(letters)):
            pairwise_marg[i,j][k1] = np.zeros(len(letters))
            for k2 in range(len(letters)):
                pairwise_marg[i,j][k1][k2] = np.exp( np.dot(wf_d[k1], x[i-1].T) + np.dot(wf_d[k2], x[j-1].T) + m[i-1,i][k1] + m[j+1,j][k2] + wt_d[k1][k2] - Z )
    return marg, pairwise_marg

def der_feat(d,count):
    reshaped = d.reshape((321+len(letters), len(letters)))
    wf_d = reshaped[0:321][:].T
    wt_d = reshaped[321:][:]


def der(d,count):
    reshaped = d.reshape((321+len(letters), len(letters)))
    wf_d = reshaped[0:321][:].T
    wt_d = reshaped[321:][:]

    dt = np.zeros((len(letters),len(letters)))
    df = np.zeros((len(letters), 321))
    '''
    for c in range(len(letters)):
        for c_ in range(len(letters)):
            sum = 0
            for i in range(count):
                x = train_img_feat(i)
                mg, pmg = marginals(x,wf_d,wt_d)
                for j in range(len(x)-1):
                    if (train_words[i][j] == letters[c] and train_words[i][j+1] == letters[c_]):
                        sum += 1
                    sum -= pmg[j+1,j+2][c][c_]
            dt[c][c_] = -sum/count
    '''
    for i in range(count):
        x = train_img_feat(i)
        mg, pmg = marginals(x,wf_d,wt_d)
        for c in range(len(letters)):
            for c_ in range(len(letters)):
                for j in range(len(x)-1):
                    if (train_words[i][j] == letters[c] and train_words[i][j+1] == letters[c_]):
                        dt[c][c_] += 1
                    dt[c][c_] -= pmg[j+1,j+2][c][c_]
            for f in range(321):
                for j in range(len(x)):
                    if (train_words[i][j] == letters[c]):
                        df[c][f] += x[j][f]
                    df[c][f] -= mg[j+1][c] * x[j][f]
    #return df/-count, dt/-count
    return np.ravel(np.vstack((df.T/-count,dt/-count)))
    '''


    for c in range(len(letters)):
        for f in range(321):
            sum = 0
            for i in range(count):
                x = train_img_feat(i)
                mg, pmg = marginals(x,wf_d,wt_d)
                for j in range(len(x)):
                    if (train_words[i][j] == letters[c]):
                        sum += x[j][f]
                    sum -= mg[j+1][c] * x[j][f]
            df[c][f] = -sum/count
    return df
    return dt
    '''


def avg_log_likelihood(d,count):
    reshaped = d.reshape((321+len(letters), len(letters)))
    wf_d = reshaped[0:321][:].T
    wt_d = reshaped[321:][:]
    sum = 0
    for i in range(count):
        x = train_img_feat(i)
        y = train_words[i]
        Z = get_Z(1,x, wf_d, wt_d)
        psi_1 = 0
        for j in range(len(y)):
            psi_1 += np.dot(wf_d[letters.index(y[j])], x[j].T)
        psi_2 = 0
        for j in range(len(y)-1):
            psi_2 += wt_d[letters.index(y[j])][letters.index(y[j+1])]
        sum += ( (psi_1 + psi_2 - Z))
    return -sum/count

d=np.ravel(np.vstack((wf.T,wt)))
#t1=time.time()
#print(der(d,50))
#t2=time.time()

#print(t2-t1)
#display(pd.DataFrame(der_feat(d,50)))
#print(der_trans(d,50))
#print(der_trans(d,"e","e",50))
#print(avg_log_likelihood(d,50))
#for i in range(50):
#    x = train_img_feat(i)
df, dt = der_trans(d, 50)
display(pd.DataFrame(df))
display(pd.DataFrame(dt))