lab6/Lab_6_HMM.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lab 6: Hidden Markov Model\n",
    "\n",
    "\n",
    "In this lab we will look into Hidden Markov Models (HMM) to model sequential data. HMMs use Markov Chains to compute probabilities for a sequence of observed events. Moreover, HMMs are based on the Markov assumption, which states that the present state $z_n$ is sufficient to predict the future $z_{n+1}$, so the past $z_{0:n-1}$ can be discarded.\n",
    "\n",
    "<img src=\"hmm.png\" width=\"400\">\n",
    "\n",
    "Often, we are in a situation where the states we are interested in are hidden, we cannot observe them directly. As we will see in Exercise 2, the part-of-speech (POS) tags in a text are hidden and we can only observe the words. From these words we have to infer the tags. Similarly in Exercises 3 where a robot needs to be localised, we cannot observe the robot's position but rather the measurements of its sensors. Both the tags and the robot position are called hidden variables because they are not observed.\n",
    "An HMM is specified by the following components:\n",
    "\n",
    "* A set of $N$ states.\n",
    "\n",
    "* A transition probability matrix $A$ where each element $a_{ij}$ represents the probability of moving from state $i$ to state $j$, s.t.  $ \\sum^{N}_{j=1} a_{ij} = 1$ $ \\forall i$\n",
    "\n",
    "* An emission probability distribution, the probabilities of observations $x_n$ being generated from a state $z_n$\n",
    "\n",
    "* An initial probability distribution over states. $\\pi_n$ is the probability that the Markov chain will start in state $n$. Also, $\\sum^{N}_{n=1} \\pi_n = 1 $\n",
    "\n",
    "\n",
    "### Additional packages (if using your own machine)\n",
    "\n",
    "For this lab we need `hmmlearn`, `nltk`, `ipywidgets`, so **if you are using your own machine** which you can install with conda:\n",
    "\n",
    "```python\n",
    "conda install -c conda-forge hmmlearn\n",
    "conda install -c anaconda nltk\n",
    "conda install -c conda-forge ipywidgets\n",
    "```\n",
    "\n",
    "Let's import the libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from hmmlearn import hmm\n",
    "from IPython import display\n",
    "\n",
    "import nltk\n",
    "from nltk.corpus import brown\n",
    "from nltk.probability import ConditionalFreqDist,ConditionalProbDist,MLEProbDist\n",
    "import itertools\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "\n",
    "# interactive display\n",
    "import ipywidgets as widgets\n",
    "%config InlineBackend.figure_format = 'retina'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1) HMM with 2 states\n",
    "We will start with a warm up exercise where you have to implement an HMM with 2 states and then analyse how the transition probability and emission noise influence the output. Your task is to create a function `hmm_model` that has three arguments:\n",
    "1. `switch_prob`: The probability of switching to the other state.\n",
    "2. `noise_level`: The variance of the emission distribution, i.e. the measurement/observation noise.\n",
    "3. `startprob`: The probabilities of starting in each state (`shape=(2,1)`).\n",
    "\n",
    "Your function should create a [hmm.GaussianHMM](https://hmmlearn.readthedocs.io/en/latest/api.html#hmmlearn.hmm.GaussianHMM) model with two states (hidden states) and `covariance_type=\"full\"`. Modify the following model attributes:\n",
    "* `startprob_`: Using the `startprob` argument.\n",
    "* `transmat_`: Implement a transition matrix with probability of transition for both states being $p_{transition} = $ `switch_prob`.\n",
    "* `means_`: Add mean $\\mu_1 = 1$ and $\\mu_2 = -1$ for each state\n",
    "* `covars_`: Using the the `noise_level` argument. (note: the shape should be `(2, 1, 1)`, i.e. two $1\\times1$ covariance matrices)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_hmm(model, states, observations):\n",
    "    \"\"\"Plots HMM states and observations for 1d states and observations.\n",
    "\n",
    "    Args:\n",
    "    model (hmmlearn model):               hmmlearn model used to get state means.\n",
    "    states (numpy array of floats):       Samples of the states.\n",
    "    observations (numpy array of floats): Samples of the states.\n",
    "    \"\"\"\n",
    "\n",
    "    nsteps = states.size\n",
    "    fig, ax1 = plt.subplots(figsize=(8,4))\n",
    "    states_forplot = list(map(lambda s: model.means_[s], states))\n",
    "    ax1.step(np.arange(nstep), states_forplot, \"--\", where=\"mid\", alpha=1.0, c=\"green\")\n",
    "    ax1.set_xlabel(\"Time\")\n",
    "    ax1.set_ylabel(\"Latent State\", c=\"green\")\n",
    "    ax1.set_yticks([-1, 1])\n",
    "    ax1.set_yticklabels([\"State 1\", \"State 0\"])\n",
    "\n",
    "    ax2 = ax1.twinx()\n",
    "    ax2.plot(np.arange(nstep), observations.flatten(), c=\"blue\")\n",
    "    ax2.set_ylabel(\"Observations\", c=\"blue\")\n",
    "    ax1.set_ylim(ax2.get_ylim())\n",
    "    plt.show(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# write your code here\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, execute the cell below to run your `hmm_model` and visualise the output with the provided `plot_hmm` function. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "np.random.seed(101)\n",
    "nstep = 50\n",
    "model = hmm_model(switch_prob=0.1, noise_level=1e-8, startprob=np.array([1.0, 0.0]))\n",
    "observations, states = model.sample(nstep)\n",
    "plot_hmm(model, states, observations)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the cell below and experiment with the interactive widget window. Try changing the value of `switch_prob` and `noise_level`, how does that changes the output?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#@title\n",
    "\n",
    "#@markdown Make sure you execute this cell to enable the widget!\n",
    "np.random.seed(101)\n",
    "nstep = 50\n",
    "\n",
    "@widgets.interact\n",
    "def plot(switch_prob=(0., 1, .01), log10_noise_level=(-8., 1., .01)):\n",
    "    model = hmm_model(switch_prob=switch_prob,\n",
    "                    noise_level=10.**log10_noise_level,\n",
    "                    startprob=[1.0, 0.0])\n",
    "\n",
    "    observations, states = model.sample(nstep)\n",
    "    observations = observations.flatten()\n",
    "    plot_hmm(model, states, observations)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2) HMM Part-of-Speech Tagging\n",
    "\n",
    "Part-of-speech (POS) tagging enables the extraction of meaningful information about words in a sentence and their relation to neighbouring words. Knowing whether a word is a noun or a verb provides us information about their most likely neighboring words (such as nouns are preceded by determiners and adjectives, verbs by nouns) and syntactic structure (nouns are generally part of noun phrases), making POS tagging a key aspect of parsing. Parts of speech are useful features for labeling named entities like people or organisations in information extraction, or for coreference resolution. A word’s part of speech can even play a role in speech recognition or synthesis, e.g., the word content is pronounced CONtent when it is a noun and conTENT when it is an adjective.\n",
    "\n",
    "Part-of-speech tagging is the process of assigning a POS tag to each word token in a given text. The input to a tagging algorithm is a sequence of tokenised words and the output is a sequence of tags, one per token. Particularly, words are ambiguous as they have more than one possible POS and the goal is to find the correct tag for each situation.\n",
    "For example, \"book\" can be a verb (\"book that flight\") or a noun (\"hand me that book\"). The goal of POS-tagging is to resolve these ambiguities, choosing the proper tag for the context.\n",
    "\n",
    "\n",
    "In this section we introduce the use of the Hidden Markov Model for part-of-speech tagging. The HMM can be used as a sequence tagger to assign a tag or class label to each unit in an observed sequence, thus mapping a sequence of observations to a sequence of labels. A HMM is a probabilistic sequence model: given a sequence of units (words, letters, morphemes, sentences), it computes a probability distribution over possible sequences of labels and chooses the best label sequence."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1) Brown corpus dataset\n",
    "The Brown corpus is a common dataset in Natural Language Processing (NLP), it is available from the NLTK library and it supports POS tagged information. Run the cell below to download the dataset and the universal tagset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Brown dataset\n",
    "nltk.download('brown')\n",
    "nltk.download('universal_tagset')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First we need to get the dataset and split it into training and testing.\n",
    "Run the following cells to achieve that and to prepare the dataset in the correct format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nltk_data = list(brown.tagged_sents(tagset='universal'))\n",
    "train_set,test_set = train_test_split(nltk_data,\n",
    "                                      train_size=0.80,\n",
    "                                      test_size=0.20,\n",
    "                                      random_state=101\n",
    "                                      )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create list of train and test tagged words\n",
    "train_tagged_words = [ tup for sent in train_set for tup in sent ]\n",
    "test_tagged_words = [ tup for sent in test_set for tup in sent ]\n",
    "print(f'Number of tagged words in the training set: {len(train_tagged_words)}')\n",
    "print(f'Number of tagged words in the test set: {len(test_tagged_words)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's explore the different types of tags by running the next cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tags = {tag for word, tag in train_tagged_words}\n",
    "print(f'Number of possible tags: {len(tags)}')\n",
    "print(f'Possible tags: {tags}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next cell, explores the structure of the dataset in terms of words and sentences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Words example: {}'.format(train_tagged_words[0:5]))\n",
    "print('Sentence example: {}'.format(train_set[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2) POS tagging model\n",
    "\n",
    "Two of the main components of HMMs are the transition model and the emission model. You already got a taste of how the two models operate in the exercise 1. Your task now is two implement both models for POS tagging model.\n",
    "\n",
    "Specifically, the transition model $P(tag_{t+1}|tag_t)$ will estimate the probability of the next tag given the current tag while the emission model $P(word|tag_t)$ will estimate the probability of observing a word given the current tag.\n",
    "\n",
    "Note: the provided hints link to the `ntlk.probability` functions, however you are free to use any other libraries such as `numpy`.\n",
    "\n",
    "#### 2.2.1) Emission Model\n",
    "Given the tagged words, the main steps are:\n",
    "1. Create a list of tuples ($tag, word$) from the input `words` which will be utilised in step 2. (careful: `words` contains tuples of ($word, tag$))\n",
    "2. Calculate the frequency for each word given the corresponding tag. (hint: [ConditionalFreqDist](https://www.nltk.org/api/nltk.html?highlight=conditionalfreqdist#nltk.probability.ConditionalFreqDist)) \n",
    "3. Calculate the conditional probability distribution given the frequency distribution of step 2. (hint: [ConditionalProbDist](https://www.nltk.org/api/nltk.html?highlight=conditionalprobdist#nltk.probability.ConditionalProbDist) with `probdist_factory`=`MLEProbDist`)\n",
    "\n",
    "Implement these steps in the cell below in a function `emission_model` that has one argument, `words`, containing a list of `(word, tag)` pairs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# write your code here\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.2.2) Transition Model\n",
    "Given the tagged sentences, the main steps are:\n",
    "1. Create chain of tuples ($tag_{t}, tag_{t+1}$) to be passed to `ConditionalFreqDist()`. To achieve this efficiently:\n",
    "    * Create a generator expression (similar to list comprehension but replacing `[]` with `()` which iterates over each sentence and for each sentence creates a pair of ($tag_{t}, tag_{t+1}$) (if you are curious about generator expressions, you can read [this](https://wiki.python.org/moin/Generators) and [this](https://realpython.com/introduction-to-python-generators/))\n",
    "    * Create a chain from the ouput of the generator (hint: [itertools.chain.from_iterable()](https://docs.python.org/3.4/library/itertools.html#itertools.chain.from_iterable))\n",
    "2. Calculate the frequency of the next tag given the previous tag. (hint: [ConditionalFreqDist](https://www.nltk.org/api/nltk.html))\n",
    "3. Calculate the conditional probability distribution of the next tag given the previous tag. (hint: [ConditionalProbDist](https://www.nltk.org/api/nltk.html))\n",
    "\n",
    "Implement these steps in the cell below in a function `transition_model` that has one argument, `sentences`, containing the dataset of sentences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# write your code here\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Utilising the previously created models, the cell below trains the emission and transition models, the former using tagged words and the later on tagged sentences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "emission_p = emission_model(train_tagged_words)\n",
    "transition_p = transition_model(train_set)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Based on the Brown corpus, the trained HMM gives the following probabilities:\n",
    "* The probability of observing the world `city` given the tag `NOUN`.\n",
    "* The probability of the tag `VERB` given the tag `NOUN`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p_city_NOUN = emission_p['NOUN'].prob('city')\n",
    "p_VERB_NOUN = transition_p['NOUN'].prob('VERB')\n",
    "print('p(city|NOUN) = ', p_city_NOUN)\n",
    "print('p(VERB|NOUN) = ', p_VERB_NOUN)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3) Viterbi algorithm\n",
    "\n",
    "The goal of the Viterbi algorithm is to find the most likely sequence of hidden states for a given sequence of observations. Specifically, we will utilise the Viterbi algorithm to find the tags of a sequence of untagged sentences, known as $decoding$. \n",
    "The benefit of this algorithm comes from its ability to efficiently determine the most probable path amongst the exponentially many possibilities. Doing so, it can reduce the complexity from $O(N^T)$ to $O(NT)$ where $N$ is the total number of words and $T$ is the total number of tags.\n",
    "\n",
    "Your task is to implement Viterbi algorithm and perform HMM decoding.\n",
    "Complete the missing `to-do` lines.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def viterbi(observed_seq, states, start_p, transition_p, emit_p):\n",
    "    eps = 0.00000001\n",
    "    \n",
    "    # Initialise the list, V, that will contain a dictionary for each element in the sequence. Each dictionary \n",
    "    # has the possible states (tags) as keys and its values are a dictionary with the keys:\n",
    "    # 'prob': the probability of the sequence until this point\n",
    "    # 'prev': the previous state used for backtracking.\n",
    "    V = [dict()] \n",
    "    for state in states:\n",
    "        # Compute the probability of each possible state for the first token in the sequence.\n",
    "        V[0][state] = {\"prob\": start_p[state] * emit_p[state].prob(observed_seq[0]), \"prev\": None}\n",
    "\n",
    "    # Run Viterbi for t > 0\n",
    "    for t in range(1, len(observed_seq)):\n",
    "        V.append({})\n",
    "        for state in states:\n",
    "            # In the forward pass, for each element of the observed_seq we want to store the maximum \n",
    "            # probability for the each state as well as the previous state. This probability is the product of\n",
    "            # the maximum transition probability and the emission probability.\n",
    "            \n",
    "            # First we calculate the max_transition_prob by going through all the combinations between states.\n",
    "            # The next line just initialises the max_transition_prob:\n",
    "            max_transition_prob = -1\n",
    "\n",
    "            for prev_state in states:\n",
    "                # Next, we calculate each possible transition probability that results in tag_t = state.\n",
    "                # The transition probability is given by the product between\n",
    "                # the prob of the sequence until t-1 where tag_{t-1} = prev_state,\n",
    "                # and the transition probability from prev_state to state.\n",
    "                #to-do\n",
    "                \n",
    "                #update the maximum transition_prob accordingly\n",
    "                if transition_prob > max_transition_prob:\n",
    "                    max_transition_prob = transition_prob\n",
    "                    prev_state_selected = prev_state\n",
    "                    \n",
    "            # Calculate the max_prob given by the product of the maximum transition probability and the \n",
    "            # emission probability . Remember to add eps to each probability before multiplying them together\n",
    "            #to-do\n",
    "            \n",
    "            V[t][state] = {\"prob\": max_prob, \"prev\": prev_state_selected}\n",
    "\n",
    "    most_likely_seq = []\n",
    "    max_prob = 0.0\n",
    "    previous = None\n",
    "    # Get most probable final state before backtracking\n",
    "    for state, data in V[-1].items():\n",
    "        if data[\"prob\"] > max_prob:\n",
    "            max_prob = data[\"prob\"]\n",
    "            best_state = state\n",
    "    most_likely_seq.append(best_state)\n",
    "    previous = best_state\n",
    "\n",
    "    # Backtrack until the first observation\n",
    "    for t in range(len(V) - 2, -1, -1):\n",
    "        # add the most likely tag to the most_likely_seq. Remember: we are backtracking hence we need to \n",
    "        # add before the last tag (hint: most_likely_seq.insert() )\n",
    "        # to-do\n",
    "        \n",
    "        previous = V[t + 1][previous][\"prev\"]\n",
    "\n",
    "    # print ('The steps of states are ' + ' '.join(most_likely_seq) + ' with highest probability of %s' % max_prob)\n",
    "\n",
    "    return list(zip(observed_seq, most_likely_seq))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the cell below to get `n=10` random sentences as test data (untagged words)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random, time\n",
    "# Let's test our Viterbi algorithm on a few sample sentences of test dataset\n",
    "random.seed(1234)      #define a random seed to get same sentences when run multiple times\n",
    "\n",
    "n = 10\n",
    "# n = 11468 # whole test set\n",
    "# choose random n numbers\n",
    "rndom = [random.randint(1,len(test_set)) for x in range(n)]\n",
    "\n",
    "# list of 10 sentences on which we test the model\n",
    "test_run = [test_set[i] for i in rndom]\n",
    " \n",
    "# list of tagged words\n",
    "test_run_base = [tup for sent in test_run for tup in sent]\n",
    " \n",
    "# list of untagged words\n",
    "test_tagged_words = [tup[0] for sent in test_run for tup in sent]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the cell below to check your implementation of the Viterbi algorithm on the test set and its POS tagging accuracy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "possible_tags = list(set([pair[1] for pair in train_tagged_words]))\n",
    "start_pr = {}\n",
    "for tag in possible_tags:\n",
    "    start_pr[tag] = 1.0/len(possible_tags)\n",
    "    \n",
    "start = time.time()\n",
    "predicted_tags = viterbi(test_tagged_words, possible_tags, start_pr, transition_p, emission_p)\n",
    "end = time.time()\n",
    "difference = end-start\n",
    "\n",
    "print(\"Time taken in seconds: \", difference)\n",
    "\n",
    "# accuracy \n",
    "predicted_correct = [i for i, j in zip(predicted_tags, test_run_base) if i == j]  \n",
    "\n",
    "accuracy = len(predicted_correct)/len(predicted_tags)\n",
    "print('Viterbi Algorithm Accuracy: ', accuracy*100)\n",
    "\n",
    "# accuracy ~= 97 for only 10 sentences (136 words)\n",
    "# accuracy ~= 94.239 in 45 secs for all test data (234100 words)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's print the print the first 10 words and their predicted tags."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predicted_tags[0:10]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "as well as the mispredicted tags:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, j in zip(predicted_tags, test_run_base):\n",
    "    if i != j:\n",
    "        print(i,j)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "How well does the model predict? Can you make any observations regarding the wrongly-predicted tags?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3) Robot localisation -- Optional\n",
    "\n",
    "Another application of HMM application is robot localisation, which is the focus of this exercise. \n",
    "\n",
    "We have a robot which is deployed to Mars to collect some samples. However, during the landing phase, the connection got temporarily lost and we are not sure where exactly the robot is. Luckily, the map of the area is available and the robot is provided with a sensor that can detect rocky surface.\n",
    "\n",
    "The aim is to use an HMM to localise the robot where the hidden variable $z$ is the position of the robot and the observation $x$ is given by the sensor measurement. \n",
    "The figure at the top of the lab sheet, taken from Bishop, shows the relation between hidden variables $z$ with the observable variables $x$.  \n",
    "\n",
    "We will achieve this through the discrete case of the Bayes Filter which consists of mainly two steps:\n",
    "prediction step and measurement step. Your task will be to implement these two steps which is given by the following pseudocode.\n",
    "\n",
    "First step:\n",
    "$$ \\overline{p}_{k,t} = \\sum_i p(z_k | u_t, z_i) p_{i, t-1} $$\n",
    "where:\n",
    "* $\\overline{p}_{k,t}$: the predicted belief of the new state $z_k$ at the current time $t$ based on the action (control $u_t$). \n",
    "* $p_{i, t-1}$: the belief of state $z_i$ at the previous time step $t_1$\n",
    "* $p(z_k | u_t, z_i)$ : the transition probability of moving from state $z_i$ to state $z_k$ when taking action $u_t$.\n",
    "\n",
    "In the second step, the knowledge from the measurement model $p(x_t | z_k)$ is combined with the predicted belief $\\overline{p}_{k,t}$ to obtain the posterior belief $p_{i, t-1}$ of each state. This posterior belief becomes the current belief for the next time step.\n",
    "\n",
    "$$ p_{i, t} = \\eta p(x_t | z_k) \\overline{p}_{k,t}  $$\n",
    "$p(x_t | z_k)$ is also known as the likelihood, i.e. how likely it is to see an observation given the robot is in location $z_k$. We need to get the likelihood before implementing the belief. Moreover $\\eta$ is the normalising term such that the probability adds to 1."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We provide the implementation of both World and Robot class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class World():\n",
    "    def __init__(self,start=(0,0),N=15):\n",
    "        self.world = np.zeros((N,N))\n",
    "        self.world[0,:] = 1\n",
    "        self.world[-1,:] = 1\n",
    "        self.world[:,0] = 1\n",
    "        self.world[:,-1] = 1\n",
    "        self.world[5, 6:12] = 1\n",
    "        self.world[6:11, 9] = 1\n",
    "        self.world[7:10, 6] = 1\n",
    "        self.world[11, 4:10] = 1\n",
    "        self.N = N\n",
    "    \n",
    "    def get_world(self):\n",
    "        return self.world\n",
    "    def set_world(self,N):\n",
    "        self.world = np.zeros((N,N))\n",
    "    def show_world(self):\n",
    "        fig, ax = plt.subplots(figsize=(5,5))\n",
    "        ax.grid(which='major', axis='both', linestyle='-', color='k', linewidth=1)\n",
    "        ax.set_xticks(np.arange(-.5, self.N, 1));\n",
    "        ax.set_yticks(np.arange(-.5, self.N, 1));\n",
    "#         ax.axis('off')\n",
    "        im = ax.imshow(self.world)\n",
    "        plt.colorbar(im)\n",
    "        plt.show()\n",
    "        \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The robot has 4 possible actions: up, down, left and right. Each action will transition the robot in the new state with probability `p_move`. At each state the robot can measure whether the surface is rocky or normal by calling `measurement()` with the sensor's noise value of `p_sense`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Robot:\n",
    "    def __init__(self, p_move, p_sense, world):\n",
    "        self.location = np.random.randint(0,world.shape[0],2)\n",
    "        self.p_move = p_move\n",
    "        self.p_sense = p_sense\n",
    "        self.world =  world\n",
    "        self.N = self.world.shape[0]\n",
    "        self.possibleActions = ['U', 'D', 'L', 'R']\n",
    "        self.tracking_error = []\n",
    "        self.homing_error = []\n",
    "\n",
    "    # The robot transition are noisy with prob_move to follow the action and (1-p_move) of staying\n",
    "    # in the same location\n",
    "    def transition(self,action):\n",
    "        if (np.random.rand() < self.p_move):\n",
    "            if(action == 'U' ):\n",
    "                self.location =  (max(0,self.location[0]-1), self.location[1])\n",
    "            elif(action == 'R'):\n",
    "                self.location = (self.location[0], min(self.N-1,self.location[1]+1))\n",
    "            elif(action == 'D'):\n",
    "                self.location = (min(self.N-1,self.location[0]+1), self.location[1])\n",
    "            elif(action == 'L'):\n",
    "                self.location = (self.location[0], max(0,self.location[1]-1))\n",
    "            \n",
    "        return self.location\n",
    "\n",
    "    # Measurement model which returns 1 when the surface is rocky and 0 otherwise.\n",
    "    # As in transition model, the sensor can be noisy with p_sense of returning the correct value.\n",
    "    def measurement(self):\n",
    "        if (np.random.rand() < self.p_sense):\n",
    "            return self.world[int(self.location[0]),int(self.location[1])]\n",
    "        else:\n",
    "            if self.world[int(self.location[0]),int(self.location[1])] == 0:\n",
    "                return 1\n",
    "            else:\n",
    "                return 0\n",
    "\n",
    "    # transition model of the robot assuming no noise\n",
    "    def transition_no_noise(self,location,action):\n",
    "        if(action == 'U' ):\n",
    "            location =  (max(0,location[0]-1), location[1])\n",
    "        elif(action == 'R'):\n",
    "            location = (location[0], min(self.N-1,location[1]+1))\n",
    "        elif(action == 'D'):\n",
    "            location = (min(self.N-1,location[0]+1), location[1])\n",
    "        elif(action == 'L'):\n",
    "            location = (location[0], max(0,location[1]-1))\n",
    "            \n",
    "        return location\n",
    "\n",
    "    def visualise_true_position(self):\n",
    "        x = self.world.copy()\n",
    "        x[int(self.location[0]),int(self.location[1])] = 0.5\n",
    "#         x[int(self.home[0]),int(self.home[1])] = 0.7\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The map is discrete and represented through grids where each grid indicates a position. \n",
    "As displayed below, the yellow corresponds to the rocky surface while the purple corresponds to the normal surface. The sensor measurement will output a value of 1 for the rocky surface and a value of 0 for the normal surface."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "world = World(15)\n",
    "world.show_world()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the following two cells to initialise the robot class and show the starting position."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p_move = .9  # original value 0.9\n",
    "p_sense = .8 # original value 0.8\n",
    "N = 15\n",
    "# Initialise belief in position - uniform distribution over world\n",
    "state_belief = np.ones((N,N))/(N*N)\n",
    "\n",
    "robot = Robot(p_move,p_sense,world.get_world())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(robot.visualise_true_position());"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next cell provides helper functions to visualise the true robot position, the predicted belief, the likelihood and the posterior belief."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# helper function to visualise \n",
    "def show(robot, predicted_state_belief, likelihood, state_belief):\n",
    "    plt.clf()\n",
    "    plt.subplot(2,4,1)\n",
    "    plt.cla()\n",
    "    plt.imshow(robot.visualise_true_position())\n",
    "    plt.title('True position in world')\n",
    "    \n",
    "    plt.subplot(2,4,2)\n",
    "    plt.cla()\n",
    "    plt.imshow(predicted_state_belief)\n",
    "    plt.title('Predicted belief given action')\n",
    "    \n",
    "    plt.subplot(2,4,3)\n",
    "    plt.cla()\n",
    "    plt.imshow(likelihood)\n",
    "    plt.title('Likelihood')\n",
    "    \n",
    "    plt.subplot(2,4,4)\n",
    "    plt.cla()\n",
    "    plt.imshow(state_belief)\n",
    "    plt.title('Posterior belief given measurement')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next cell contains the main algorithm and your task is to fill the missing lines `to-do` (hint: check the provided pseudocode above)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Run the discrete Bayes filter\n",
    "plt.figure(figsize=(15,9))\n",
    "\n",
    "N = 15\n",
    "steps = 100\n",
    "# Create list of possible states in world (x,y coordinates)\n",
    "possible_states = np.array(np.meshgrid(np.linspace(0,N-1,N),np.linspace(0,N-1,N))).reshape(2,-1).T\n",
    "\n",
    "for step in range(steps):\n",
    "    \n",
    "    # Sample a random action\n",
    "    a = np.random.choice(['U','D','L','R'])\n",
    "    \n",
    "    # Move the robot according to the action and sense the measurement\n",
    "    state = robot.transition(a)\n",
    "    measurement = robot.measurement()\n",
    "    \n",
    "    # Make prediction about state belief given action\n",
    "    predicted_state_belief = np.zeros((N,N))\n",
    "    \n",
    "    for state in possible_states.astype(int):\n",
    "        \n",
    "        new_state = robot.transition_no_noise(state,a)\n",
    "        \n",
    "        # to-do\n",
    "        # to-do\n",
    "   \n",
    "    # Evaluate sensor belief for each possible state\n",
    "    likelihood = np.zeros((N,N))\n",
    "    for state in possible_states.astype(int):\n",
    "        \n",
    "        match = measurement == world.get_world()[state[0],state[1]]\n",
    "        # to-do\n",
    "    \n",
    "    # Compute posterior belief over possible states\n",
    "    # to-do\n",
    "    \n",
    "    # normalising\n",
    "    state_belief = (state_belief+1e-15)\n",
    "    state_belief = state_belief/np.sum(state_belief) \n",
    "    \n",
    "    # Make a function to display\n",
    "    show(robot, predicted_state_belief, likelihood, state_belief)\n",
    "\n",
    "    display.clear_output(wait=True)\n",
    "    display.display(plt.gcf())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Change the `p_move` and `p_sense` values and re-run the experiment with different values. How these probabilities affect the performance?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Wrap-up"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Congratulations, you have reached the end of this lab. Let's summarise what we have learnt:\n",
    "* Understand the role of the HMM components \n",
    "* Key applications of HMM such as POS tagger\n",
    "* Perform POS tagging through the transition model and the emission model.\n",
    "* Implemented the Viterbi algorithm and the discrete Bayes filter algorithm.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "* COMS30035 Machine Learning\n",
    "* Bishop - Pattern Recognition and Machine Learning: Chapter 13 - Sequential Data\n",
    "* [Probabilistic Robotics](http://www.probabilistic-robotics.org) \n",
    "* University of Edinburgh's Foundations of Natural Language Processing (FNLP) and Decision Making in Robots and Autonomous Agents (DMR) courses\n",
    "* [Hussain Mujtaba](https://www.mygreatlearning.com/blog/pos-tagging/)\n",
    "* [Neuromatch Academy 2020](https://www.neuromatchacademy.org)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The End :)\n",
    "\n",
    "<video controls src=\"gif.mp4\" width=\"200\">"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}