import sys
import numpy as np
import pdb


class DependencyDecoder():
    '''
    Dependency decoder class
    '''
    def __init__(self):
        self.verbose = False

    def parse_marginals_nonproj(self, scores):
        '''
        Compute marginals and the log-partition function using the matrix-tree theorem
        '''
        nr, nc = np.shape(scores)
        if nr != nc:
            raise ValueError("scores must be a squared matrix with nw+1 rows")
            return []

        nw = nr - 1;

        s = np.matrix(scores)
        lap = np.matrix(np.zeros((nw+1, nw+1)))
        for m in range(1, nw + 1):
            d = 0.0
            for h in range(0, nw + 1):
                if m != h:
                    d += np.exp(s[h,m])
                    lap[h,m] = -np.exp(s[h,m])
            lap[m,m] = d
        r = lap[0,1:]
        minor = lap[1:,1:]

        #logZ = np.linalg.slogdet(minor)[1]
        logZ = np.log(np.linalg.det(minor))
        invmin = np.linalg.inv(minor)
        marginals = np.zeros((nw+1, nw+1))
        for m in range(1, nw + 1):
            marginals[0,m] = np.exp(s[0,m]) * invmin[m-1,m-1]
            for h in range(1, nw + 1):
                if m != h:
                    marginals[h,m] = np.exp(s[h,m]) * (invmin[m-1,m-1] - invmin[m-1,h-1])

        return marginals, logZ


    def parse_proj(self, scores):
        '''
        Parse using Eisner's algorithm.
        '''

        ###########################
        # Solution to Exercise 4.3.6 
        nr, nc = np.shape(scores)
        if nr != nc:
            raise ValueError("scores must be a squared matrix with nw+1 rows")
            return []

        N = nr - 1 # Number of words (excluding root).
        
        # Initialize CKY table.
        complete = np.zeros([N+1, N+1, 2]) # s, t, direction (right=1). 
        incomplete = np.zeros([N+1, N+1, 2]) # s, t, direction (right=1). 
        complete_backtrack = -np.ones([N+1, N+1, 2], dtype=int) # s, t, direction (right=1). 
        incomplete_backtrack = -np.ones([N+1, N+1, 2], dtype=int) # s, t, direction (right=1).

        incomplete[0, :, 0] -= np.inf
        
        # Loop from smaller items to larger items.
        for k in xrange(1,N+1):
            for s in xrange(N-k+1):
                t = s+k
                
                # First, create incomplete items.
                # left tree
                incomplete_vals0 = complete[s, s:t, 1] + complete[(s+1):(t+1), t, 0] + scores[t, s]
                incomplete[s, t, 0] = np.max(incomplete_vals0)
                incomplete_backtrack[s, t, 0] = s + np.argmax(incomplete_vals0)
                # right tree
                incomplete_vals1 = complete[s, s:t, 1] + complete[(s+1):(t+1), t, 0] + scores[s, t]
                incomplete[s, t, 1] = np.max(incomplete_vals1)
                incomplete_backtrack[s, t, 1] = s + np.argmax(incomplete_vals1)

                # Second, create complete items.
                # left tree
                complete_vals0 = complete[s, s:t, 0] + incomplete[s:t, t, 0]
                complete[s, t, 0] = np.max(complete_vals0)
                complete_backtrack[s, t, 0] = s + np.argmax(complete_vals0)
                # right tree
                complete_vals1 = incomplete[s, (s+1):(t+1), 1] + complete[(s+1):(t+1), t, 1]
                complete[s, t, 1] = np.max(complete_vals1)
                complete_backtrack[s, t, 1] = s + 1 + np.argmax(complete_vals1)
                
        value = complete[0][N][1]
        heads = -np.ones(N+1, dtype=int)
        self.backtrack_eisner(incomplete_backtrack, complete_backtrack, 0, N, 1, 1, heads)

        value_proj = 0.0
        for m in xrange(1,N+1):
            h = heads[m]
            value_proj += scores[h,m]
        
        return heads

        # End of solution to Exercise 4.3.6 
        ###########################

    def backtrack_eisner(self, incomplete_backtrack, complete_backtrack, s, t, direction, complete, heads):
        '''
        Backtracking step in Eisner's algorithm.
        - incomplete_backtrack is a (NW+1)-by-(NW+1) numpy array indexed by a start position,
        an end position, and a direction flag (0 means left, 1 means right). This array contains
        the arg-maxes of each step in the Eisner algorithm when building *incomplete* spans.
        - complete_backtrack is a (NW+1)-by-(NW+1) numpy array indexed by a start position,
        an end position, and a direction flag (0 means left, 1 means right). This array contains
        the arg-maxes of each step in the Eisner algorithm when building *complete* spans.
        - s is the current start of the span
        - t is the current end of the span
        - direction is 0 (left attachment) or 1 (right attachment)
        - complete is 1 if the current span is complete, and 0 otherwise
        - heads is a (NW+1)-sized numpy array of integers which is a placeholder for storing the 
        head of each word.
        '''
        if s == t:
            return
        if complete:
            r = complete_backtrack[s][t][direction]
            if direction == 0:
                self.backtrack_eisner(incomplete_backtrack, complete_backtrack, s, r, 0, 1, heads)
                self.backtrack_eisner(incomplete_backtrack, complete_backtrack, r, t, 0, 0, heads)
                return
            else:
                self.backtrack_eisner(incomplete_backtrack, complete_backtrack, s, r, 1, 0, heads)
                self.backtrack_eisner(incomplete_backtrack, complete_backtrack, r, t, 1, 1, heads)
                return
        else:
            r = incomplete_backtrack[s][t][direction]
            if direction == 0:
                heads[s] = t
                self.backtrack_eisner(incomplete_backtrack, complete_backtrack, s, r, 1, 1, heads)
                self.backtrack_eisner(incomplete_backtrack, complete_backtrack, r+1, t, 0, 1, heads)
                return
            else:
                heads[t] = s
                self.backtrack_eisner(incomplete_backtrack, complete_backtrack, s, r, 1, 1, heads)
                self.backtrack_eisner(incomplete_backtrack, complete_backtrack, r+1, t, 0, 1, heads)
                return

                
    def parse_nonproj(self, scores):
        '''
        Parse using Chu-Liu-Edmonds algorithm.
        '''
        nr, nc = np.shape(scores)
        if nr != nc:
            raise ValueError("scores must be a squared matrix with nw+1 rows")
            return []

        nw = nr - 1;

        curr_nodes = np.ones(nw+1, int)
        reps = []
        oldI = -np.ones((nw+1, nw+1), int)
        oldO = -np.ones((nw+1, nw+1), int)
        for i in range(0, nw+1): 
            reps.append({i : 0})
            for j in range(0, nw+1):
                oldI[i,j] = i
                oldO[i,j] = j
                if i==j or j==0:
                    continue

        if self.verbose:
            print "Starting C-L-E...\n"

        scores_copy = scores.copy()
        final_edges = self.chu_liu_edmonds(scores_copy, curr_nodes, oldI, oldO, {}, reps)
        heads = np.zeros(nw+1, int)
        heads[0] = -1
        for key in final_edges.keys():
            ch = key
            pr = final_edges[key]
            heads[ch] = pr

        return heads


    def chu_liu_edmonds(self, scores, curr_nodes, oldI, oldO, final_edges, reps):
        '''
        Chu-Liu-Edmonds algorithm
        '''

        # need to construct for each node list of nodes they represent (here only!)
        nw = np.size(curr_nodes) - 1;

        # create best graph
        par = -np.ones(nw+1, int)
        for m in range(1, nw+1):
            # only interested in current nodes
            if 0 == curr_nodes[m]: 
                continue
            max_score = scores[0,m]
            par[m] = 0
            for h in range(nw+1):
                if m == h:
                    continue
                if 0 == curr_nodes[h]:
                    continue
                if scores[h,m] > max_score:
                    max_score = scores[h,m]
                    par[m] = h

        if self.verbose:
            print "After init\n"
            for m in range(0, nw+1):
                if 0 < curr_nodes[m]:
                    print "{0}|{1} ".format(par[m],m)
            print "\n"


        # find a cycle
        cycles = []
        added = np.zeros(nw+1, int)
        for m in range(0, nw+1):
            if np.size(cycles) > 0:
                break
            if added[m] or 0 == curr_nodes[m]:
                continue
            added[m] = 1
            cycle = {m : 0}
            l = m
            while True:
                if par[l] == -1:
                    added[l] = 1
                    break
                if par[l] in cycle:
                    cycle = {}
                    lorg = par[l]
                    cycle[lorg] = par[lorg]
                    added[lorg] = 1
                    l1 = par[lorg]
                    while l1 != lorg:
                        cycle[l1] = par[l1]
                        added[l1] = True
                        l1 = par[l1]
                    cycles.append(cycle)
                    break
                cycle[l] = 0
                l = par[l]
                if added[l] and (l not in cycle):
                    break
                added[l] = 1

        # get all edges and return them
        if np.size(cycles) == 0:
            for m in range(0, nw+1):
                if 0 == curr_nodes[m]:
                    continue
                if par[m] != -1:
                    pr = oldI[par[m], m]
                    ch = oldO[par[m], m]
                    final_edges[ch] = pr
                else:
                    final_edges[0] = -1
            return final_edges

        max_cyc = 0
        wh_cyc = 0
        for cycle in cycles:
            if np.size(cycle.keys()) > max_cyc:
                max_cyc = np.size(cycle.keys())
                wh_cyc = cycle

        cycle = wh_cyc
        cyc_nodes = cycle.keys()
        rep = cyc_nodes[0]
        
        if self.verbose:
            print "Found Cycle\n"
            for node in cyc_nodes:
                print "{0} ".format(node)
            print "\n"

        cyc_weight = 0.0
        for node in cyc_nodes:
            cyc_weight += scores[par[node], node] 

        for i in range(0, nw+1):