# Source: content/notes/nlp/sequence-labeling.md
# Independent CPU example; use the curriculum environment.
# See /notes/ml/#example-environment or /notes/deep-learning/#example-environment.

import itertools
import numpy as np
emission = np.array([[.2, -.1], [.1, .7], [.4, .0]])
transition = np.array([[.3, -.2], [-.1, .2]])
start = np.array([0., -np.inf])  # second label cannot begin a sequence
forward = start + emission[0]
best = forward.copy()
back = []
for token in emission[1:]:
    candidates = best[:, None] + transition
    back.append(candidates.argmax(0))
    best = candidates.max(0) + token
    forward = np.logaddexp.reduce(forward[:, None]+transition, axis=0)+token
paths = list(itertools.product(range(2), repeat=3))
scores = np.array([start[p[0]] + sum(emission[t, p[t]] for t in range(3))
                  + sum(transition[p[t-1], p[t]] for t in range(1, 3)) for p in paths])
assert np.isclose(np.logaddexp.reduce(forward), np.logaddexp.reduce(scores))
decoded = [int(best.argmax())]
for pointers in reversed(back):
    decoded.append(int(pointers[decoded[-1]]))
decoded.reverse()
assert tuple(decoded) == paths[int(scores.argmax())]
assert decoded[0] == 0
print("best path", decoded, "log partition", np.logaddexp.reduce(forward))