# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.15.2
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # MAP and MPE Inference
#
# **COMP9418-19T3, W09 Tutorial**
#
# - Instructor: Gustavo Batista
# - School of Computer Science and Engineering, UNSW Sydney
# - Last Update 5th Novermber at 18:00, 2020
# $$
# % macros
# \newcommand{\indep}{\perp \!\!\!\perp}
# $$

# In this week's tutorial, we will implement MPE and MAP algorithms with variable elimination. We will extend our factor representation to store MAP and MPE instantiations. We will build our code using implementations from previous tutorials, in particular, inference with variable elimination and Markov chain models.

# ## Technical prerequisites
#
# You will need certain packages installed to run this notebook.
#
# If you are using ``conda``'s default
# [full installation](https://conda.io/docs/install/full.html),
# these requirements should all be satisfied already.

# Once we have done all that, we
# import some useful modules for later use.

# +
# Make division default to floating-point, saving confusion
from __future__ import division
from __future__ import print_function

# combinatorics
from itertools import product, combinations, permutations
# ordered dictionaries are useful for keeping ordered sets of varibles
from collections import OrderedDict as odict
# table formating for screen output
from tabulate import tabulate
# math library to get access to math.inf
# -

# ## The Medical Conditional Network
#
# In this tutorial, we will implement MPE and MAP inference using variable elimination (VE). To test our implementations, we will use the same example of a Bayesian network from the lecture. Therefore, we will be able to compare our results with those obtained in the slides.
#
# The Bayesian network has the following structure:
#
# ![](./img/Simple_net.png)
#
# In summary, it describes a problem in which we have a population of males and females ($S$) that can suffer from a medical condition ($C$) that is more likely in males than in females. There are two diagnostic tests for $𝐶$, named $𝑇_1$ and $𝑇_2$. The $𝑇_2$ test is more effective on females, and both tests are equally effective on males. Finally, we have a deterministic variable indicating if tests $T_1$ and $T_2$ have the same outcome.
#
# The next cell defines the factors for each CPT and the outcome space for all five variables.

# +
factors = {
    'S': {
        'dom': ('S'), 
        'table': odict([
            (('male',), 0.55),
            (('female',), 0.45),
        ]),
    },
    
    'C': {
        'dom': ('S', 'C'), 
        'table': odict([
            (('male', 'yes'), 0.05),
            (('male', 'no'), 0.95),
            (('female', 'yes'), .01),
            (('female', 'no'), .99),
        ])
    },

    'T1': {
        'dom': ('C', 'T1'), 
        'table': odict([
            (('yes', '+ve'), 0.80),
            (('yes', '-ve'), 0.20),
            (('no', '+ve'), .20),
            (('no', '-ve'), .80),
        ])
    },
    
    'T2': {
        'dom': ('S', 'C', 'T2'), 
        'table': odict([
            (('male', 'yes', '+ve'), 0.80),
            (('male', 'yes', '-ve'), 0.20),
            (('male', 'no', '+ve'), 0.20),
            (('male', 'no', '-ve'), 0.80),
            (('female', 'yes', '+ve'), .95),
            (('female', 'yes', '-ve'), .05),
            (('female', 'no', '+ve'), .05),
            (('female', 'no', '-ve'), .95),
        ])
    },

    'A': {
        'dom': ('T1', 'T2', 'A'), 
        'table': odict([
            (('+ve', '+ve','yes'), 1),
            (('+ve', '+ve','no'), 0),
            (('+ve', '-ve','yes'), 0),
            (('+ve', '-ve','no'), 1),
            (('-ve', '+ve','yes'), 0),
            (('-ve', '+ve','no'), 1),
            (('-ve', '-ve','yes'), 1),
            (('-ve', '-ve','no'), 0),
        ])
    }, 
}
    
outcomeSpace = dict(
    S=('male','female'),
    C=('yes','no'),
    A=('yes','no'),
    T1=('+ve','-ve'),
    T2=('+ve','-ve')
)
# -

# ## The Extended Factor
#
# An essential part of answering MAP and MPE queries is recovering the instantiation with maximal probability. Remember that an MPE probability query has the following definition:
#
# $$MPE_P(\textbf{e}) = max_{\textbf{q}}⁡P(\textbf{q},\textbf{e})$$
#
# where, $\textbf{q}$ are all variables in the network not present in $\textbf{e}$. Therefore, we are looking for an instantiation $\textbf{q}$ with maximal probability. The $MPE_P$ query returns such maximal probability, i.e., a numerical value. 
#
# The MPE instantiation query returns the instantiation, that is, values associated with each variable in $\textbf{q}$. We define the MPE instantiation query as
#
# $$MPE(e) = argmax_{\textbf{q}}P(\textbf{q},\textbf{e})$$
#
# We can use similar definitions for MAP queries. Remember that the main difference between MAP and MPE is that MAP queries typically involve a subset of the network variables. Therefore, MPE queries can be seen as a particular case of MAP queries.
#
# To keep track of the most likely instantiations, we can use extended factors. The general idea is simple. We record the value of the eliminated variable in a list when computing the maximization operation. The next figure from the slides illustrates the concept.
#
# ![](./img/Ext_factors.png)
#
# In this example, we are eliminating variable $S$ through the maximization operation. In the case of the instantiation $C=yes, T_2=ve$, the two candidates are $S=male$ with probability $P(S=male,C=yes,T_2=ve)=.0220$ and $S=female$ with probability $P(S=female,C=yes,T_2=ve)=.0043$. The largest probability comes from $S=male$ and this value is recorded in the exteded field of the factor.
#
# We need an extended factor representation. While there are many possibilities, we will use something similar to our table representation. Our requirements are the following:
#
# 1. We need one separated list of each possible combination of variable values in a factor. 
# 2. We need to record the variable name as well as the value assigned to this variable, such as $A=yes$.
# 3. We need to start with an empty structure and add entries as we eliminate variables with maximization.
# 4. We need to merge two lists when we join two factors since each one will have its extended field. Notice, as each variable is eliminated only once, the two lists of variables will be disjoint.
#
# We implement the extended field by adding an `ext` key to the factor dictionary. This key has a similar structure to the `table` key. It is an `odict` object will all possible combinations of variable values in the factor domain. However, instead of storing a number (probability), we will store a dictionary.
#
# The next cell illustrates this one factor with the extended representation.

phi = {
    'dom': ('C', 'T2'), 
    'table': odict([
        (('yes', '+ve'), 0.0220),
        (('yes', '-ve'), 0.0055),
        (('no', '+ve'),  0.1045),
        (('no', '-ve'),  0.4237),
    ]),
    'ext': odict([
        (('yes', '+ve'), dict(S='male')),
        (('yes', '-ve'), dict(S='male')),
        (('no', '+ve'),  dict(S='male')),
        (('no', '-ve'),  dict(S='female')),
    ]),    
}


# Let us start our code with a new version of the `printFactor` function. The new function prints the extended field if it is present in the factor. This code is a simple modification of the previous `printFactor` sub-routine, so we will provide it to you.

def printFactor(f):
    """
    argument 
    `f`, a factor to print on screen. If the `ext` field is presented, it will be printed as an additional column
    """
    # Create a empty list that we will fill in with the probability table entries
    table = list()
    
    # Iterate over all keys and probability values in the table
    for key, item in f['table'].items():
        # Convert the tuple to a list to be able to manipulate it
        k = list(key)
        # Append the probability value to the list with key values
        k.append(item)
        
        # NEW CODE
        # if the 'ext' key is present in the factor, we apppend the ext dictionary converted to a list of items
        if 'ext' in f:
            k.append(list(f['ext'][key].items()))
        
        # Append an entire row to the table
        table.append(k)
    # dom is used as table header. We need it converted to list
    dom = list(f['dom'])
    # Append a 'Pr' to indicate the probabity column
    dom.append('Pr')
    
    # NEW CODE
    # Append a 'Ext' to indicate the extended column
    dom.append('Ext')
    print(tabulate(table,headers=dom,tablefmt='orgtbl'))


# We can test the new `printFactor` function. It works with the extended and non-extended factors.

printFactor(factors['C'])
print()
printFactor(phi)


# ## Maximize Operation
#
# Let us start with the `maximize` operation. We implemented this function in week 5 tutorial when coding the Viterbi algorithm. The next cell brigs this code. Your task is to adapt the implementation to store the extended factors.
#
# ### Exercise
#
# Adapt the `maximize` function in the next cell to update the `ext` dictionary with the name and the corresponding value of the maximized out variable. In your code, verify if the `ext` dictionary exists in the factor `f`. If it does not exist, start with an empty dictionary. 

# +
def maximize(f, var, outcomeSpace):
    """
    argument 
    `f`, factor to be maximized.
    `var`, variable to be maximized out.
    `outcomeSpace`, dictionary with the domain of each variable.
    
    Returns a new factor f' with dom(f') = dom(f) - {var}. The `ext` field is updated to store the most likely instantiation
    """    
    
    # Let's make a copy of f domain and convert it to a list. We need a list to be able to modify its elements
    new_dom = list(f['dom']) 
    new_dom.remove(var)            # Remove var from the list new_dom by calling the method remove(). 1 line
    table = list()                 # Create an empty list for table. We will fill in table from scratch. 1 line
    for entries in product(*[outcomeSpace[node] for node in new_dom]):
        m = -1;                    # Initialize the maximization variable m. 1 line

        # We need to iterate over all possible outcomes of the variable var
        for val in outcomeSpace[var]:
            # To modify the tuple entries, we will need to convert it to a list
            entriesList = list(entries)
            # We need to insert the value of var in the right position in entriesList
            entriesList.insert(f['dom'].index(var), val)

            p = prob(f, *tuple(entriesList))     # Calculate the probability of factor f for entriesList. 1 line
            m = max(m, p)                        # Maximize over all values of var by storing the max value in m. 1 line
            
        # Create a new table entry with the multiplication of p1 and p2
        table.append((entries, m))
    return {'dom': tuple(new_dom), 'table': odict(table)}

################
# Test code

def prob(factor, *entry):
    """
    argument 
    `factor`, a dictionary of domain and probability values,
    `entry`, a list of values, one for each variable in the same order as specified in the factor domain.
    
    Returns p(entry)
    """

    return factor['table'][entry]     # insert your code here, 1 line   

printFactor(factors['T2'])
print()
f = maximize(factors['T2'], 'S', outcomeSpace)
printFactor(f)
print()
printFactor(maximize(f, 'T2', outcomeSpace))


# +
# Answer

def maximize(f, var, outcomeSpace):
    """
    argument 
    `f`, factor to be maximized.
    `var`, variable to be maximized out.
    `outcomeSpace`, dictionary with the domain of each variable.
    
    Returns a new factor f' with dom(f') = dom(f) - {var}
    """    
    
    # Let's make a copy of f domain and convert it to a list. We need a list to be able to modify its elements
    new_dom = list(f['dom'])
    new_dom.remove(var)            # Remove var from the list new_dom by calling the method remove(). 1 line

    # NEW CODE: new_ext is a list of variables names and values that will compose the exteded field of the factor
    new_ext = list()
    table = list()                 # Create an empty list for table. We will fill in table from scratch. 1 line
    for entries in product(*[outcomeSpace[node] for node in new_dom]):
        m = -1;                    # Initialize the maximization variable m. 1 line

        # We need to iterate over all possible outcomes of the variable var
        for val in outcomeSpace[var]:
            # To modify the tuple entries, we will need to convert it to a list
            entriesList = list(entries)
            # We need to insert the value of var in the right position in entriesList
            entriesList.insert(f['dom'].index(var), val)

            p = prob(f, *tuple(entriesList))     # Calculate the probability of factor f for entriesList. 1 line
            
            # NEW CODE: test to see if we have a new candidate for instantiation with maximum probability
            if p > m:
                # NEW CODE: Store the new maximum probability for instantiation
                m = p
                # NEW CODE: Store the value of the variable that is candidate to be instantiation with maximum probability
                max_val = val
                # NEW CODE: if the 'ext' field in present in the factor, we make a copy and will add values to it
                if ('ext' in f):
                    max_ext = f['ext'][tuple(entriesList)].copy()
                # NEW CODE: Otherwise, we start with an empty dictionary
                else:
                    max_ext = dict()
        
        # Create a new table entry with the multiplication of p1 and p2
        table.append((entries, m))
        # NEW CODE: Add a new entry to the dictionary: var = max_val
        max_ext[var] = max_val
        # NEW CODE: Append a line to the new_ext list with the updated extended dictionary
        new_ext.append((entries, max_ext))
        # NEW CODE: return a new factor with the `ext` field as well
    return {'dom': tuple(new_dom), 'table': odict(table), 'ext': odict(new_ext)}

################
# Test code

def prob(factor, *entry):
    """
    argument 
    `factor`, a dictionary of domain and probability values,
    `entry`, a list of values, one for each variable in the same order as specified in the factor domain.
    
    Returns p(entry)
    """

    return factor['table'][entry]  

printFactor(factors['T2'])
print()
f = maximize(factors['T2'], 'S', outcomeSpace)
printFactor(f)
print()
printFactor(maximize(f, 'T2', outcomeSpace))


# -

# If your code is correct, you should see the following output:
#
# ```
# | S      | C   | T2   |   Pr |
# |--------+-----+------+------|
# | male   | yes | +ve  | 0.8  |
# | male   | yes | -ve  | 0.2  |
# | male   | no  | +ve  | 0.2  |
# | male   | no  | -ve  | 0.8  |
# | female | yes | +ve  | 0.95 |
# | female | yes | -ve  | 0.05 |
# | female | no  | +ve  | 0.05 |
# | female | no  | -ve  | 0.95 |
#
# | C   | T2   |   Pr | Ext               |
# |-----+------+------+-------------------|
# | yes | +ve  | 0.95 | [('S', 'female')] |
# | yes | -ve  | 0.2  | [('S', 'male')]   |
# | no  | +ve  | 0.2  | [('S', 'male')]   |
# | no  | -ve  | 0.95 | [('S', 'female')] |
#
# | C   |   Pr | Ext                              |
# |-----+------+----------------------------------|
# | yes | 0.95 | [('S', 'female'), ('T2', '+ve')] |
# | no  | 0.95 | [('S', 'female'), ('T2', '-ve')] |
# ```

# ## Join Operation with Extended Factors
#
# We also need to modify the `join` function to work with extended factors. If we have two factors with a non-empty `ext` dictionaries, we need to merge those dictionaries for each table entry. 
#
# ### Exercise
#
# Implement a new version of the `join` operation that works with extended factors. The new 'join' operation must work with the following scenarios:
#
# 1. Join two factors with no extended factor fields. In this case, the resulting factor has an extended field with empty dictionaries.
# 2. Join a factor that has an extended factor field `e` with another one that does not have it. In this case, the resulting factor should have a copy of the extended field `e`. 
# 3. Join two factors with existing extended factor fields. In this case, we need to merge the two dictionaries.
#
# We have started the code for you.

# +
def join(f1, f2, outcomeSpace):
    """
    argument 
    `f1`, first factor to be joined.
    `f2`, second factor to be joined.
    `outcomeSpace`, dictionary with the domain of each variable
    
    Returns a new factor with a join of f1 and f2. This version tests of the existance of `ext` field and merge them in the resulting factor
    """
    
    # First, we need to determine the domain of the new factor. It will be union of the domain in f1 and f2
    # But it is important to eliminate the repetitions
    common_vars = list(f1['dom']) + list(set(f2['dom']) - set(f1['dom']))
    
    # We will build a table from scratch, starting with an empty list. Later on, we will transform the list into a odict
    table = list()
    # ext stores a list similar to table but for the extended factor field. We start with an empty list and convert to a odict in the end
    ext = None                                                     # Initialize ext with an empty list: 1 line
    
    # Here is where the magic happens. The product iterator will generate all combinations of varible values 
    # as specified in outcomeSpace. Therefore, it will naturally respect observed values
    for entries in product(*[outcomeSpace[node] for node in common_vars]):
        
        # We need to map the entries to the domain of the factors f1 and f2
        entryDict = dict(zip(common_vars, entries))
        f1_entry = tuple((entryDict[var] for var in f1['dom']))
        f2_entry = tuple((entryDict[var] for var in f2['dom']))
     
        # Insert your code here
        p1 = prob(f1, *f1_entry)           # Use the fuction prob to calculate the probability in factor f1 for entry f1_entry 
        p2 = prob(f2, *f2_entry)           # Use the fuction prob to calculate the probability in factor f2 for entry f2_entry 

        # Test if `ext` key exists in f1. If it does exist, make a copy of the dictionary for f1_entry
        if None:                                          # Test if `ext` key exists in the f1's extended factor: 1 line
            e = None                                      # Assign to e a copy of the ext dictionary for f1_entry: 1 line
        else:
        # Otherwise, initialize e with an empty dictionary
            e = None                                      # Assign an empty dictionary to e
        # Test if `ext` key exists in f2. If it does exist, merge the ext dictionary for f2_entry with e
        if None:                                          # Test if `ext` key exists in the f2's extended factor: 1 line
            None                                          # Merge e and the ext dictionary for f2_entry. Use update: 1 line
        
        # Create a new table entry with the multiplication of p1 and p2
        table.append((entries, p1 * p2))
        # Create a new ext table entry with the dictionary e            
        ext.append(None)                                   # Append to ext a row with the entries values and the dictionary e: 1 line
        
    return {'dom': tuple(common_vars), 'table': odict(table), 'ext': None}      # Return ext as an odict object: 1 line

########################
# Test code

phi1 = {
    'dom': ('A', 'B'), 
    'table': odict([
        (('+a', '+b'), 0.3),
        (('+a', '-b'), 0.7),
        (('-a', '+b'),  0.2),
        (('-a', '-b'),  0.8),
    ]),
    'ext': odict([
        (('+a', '+b'), dict(C='+c')),
        (('+a', '-b'), dict(C='+c')),
        (('-a', '+b'),  dict(C='+c')),
        (('-a', '-b'),  dict(C='-c')),
    ]),    
}

phi2 = {
    'dom': ('B', 'D'), 
    'table': odict([
        (('+b', '+d'), 0.1),
        (('+b', '-d'), 0.9),
        (('-b', '+d'),  0.4),
        (('-b', '-d'),  0.6),
    ]),
    'ext': odict([
        (('+b', '+d'), dict(E='-e')),
        (('+b', '-d'), dict(E='+e')),
        (('-b', '+d'),  dict(E='-e')),
        (('-b', '-d'),  dict(E='+e')),
    ]),    
}

o = dict(
    A=('+a','-a'),
    B=('+b','-b'),
    C=('+c','-c'),
    D=('+d','-d'),
    E=('+e','-e')
)

printFactor(join(phi1, phi2, o))


# +
# Answer

def join(f1, f2, outcomeSpace):
    """
    argument 
    `f1`, first factor to be joined.
    `f2`, second factor to be joined.
    `outcomeSpace`, dictionary with the domain of each variable
    
    Returns a new factor with a join of f1 and f2. This version tests of the existance of `ext` field and merge them in the resulting factor
    """
    
    # First, we need to determine the domain of the new factor. It will be union of the domain in f1 and f2
    # But it is important to eliminate the repetitions
    common_vars = list(f1['dom']) + list(set(f2['dom']) - set(f1['dom']))
    
    # We will build a table from scratch, starting with an empty list. Later on, we will transform the list into a odict
    table = list()
    # ext stores a list similar to table but for the extended factor field. We start with an empty list and convert to a odict in the end
    ext = list()
    
    # Here is where the magic happens. The product iterator will generate all combinations of varible values 
    # as specified in outcomeSpace. Therefore, it will naturally respect observed values
    for entries in product(*[outcomeSpace[node] for node in common_vars]):
        
        # We need to map the entries to the domain of the factors f1 and f2
        entryDict = dict(zip(common_vars, entries))
        f1_entry = tuple((entryDict[var] for var in f1['dom']))
        f2_entry = tuple((entryDict[var] for var in f2['dom']))
     
        # Insert your code here
        p1 = prob(f1, *f1_entry)           # Use the fuction prob to calculate the probability in factor f1 for entry f1_entry 
        p2 = prob(f2, *f2_entry)           # Use the fuction prob to calculate the probability in factor f2 for entry f2_entry 

        # Test if `ext` key exists in f1. If it does exist, make a copy of the dictionary for f1_entry
        if ('ext' in f1):
            e = f1['ext'][tuple(f1_entry)].copy()
        else:
        # Otherwise, initialize e with an empty dictionary
            e = dict()
        # Test if `ext` key exists in f2. If it does exist, merge the dictionary for f2_entry with e
        if ('ext' in f2):
            e.update(f2['ext'][f2_entry])
        
        # Create a new table entry with the multiplication of p1 and p2
        table.append((entries, p1 * p2))
        # Create a new ext table entry with the dictionary e            
        ext.append((entries, e))
        
    return {'dom': tuple(common_vars), 'table': odict(table), 'ext': odict(ext)}

########################
# Test code

phi1 = {
    'dom': ('A', 'B'), 
    'table': odict([
        (('+a', '+b'), 0.3),
        (('+a', '-b'), 0.7),
        (('-a', '+b'),  0.2),
        (('-a', '-b'),  0.8),
    ]),
    'ext': odict([
        (('+a', '+b'), dict(C='+c')),
        (('+a', '-b'), dict(C='+c')),
        (('-a', '+b'),  dict(C='+c')),
        (('-a', '-b'),  dict(C='-c')),
    ]),    
}

phi2 = {
    'dom': ('B', 'D'), 
    'table': odict([
        (('+b', '+d'), 0.1),
        (('+b', '-d'), 0.9),
        (('-b', '+d'),  0.4),
        (('-b', '-d'),  0.6),
    ]),
    'ext': odict([
        (('+b', '+d'), dict(E='-e')),
        (('+b', '-d'), dict(E='+e')),
        (('-b', '+d'),  dict(E='-e')),
        (('-b', '-d'),  dict(E='+e')),
    ]),    
}

o = dict(
    A=('+a','-a'),
    B=('+b','-b'),
    C=('+c','-c'),
    D=('+d','-d'),
    E=('+e','-e')
)

printFactor(join(phi1, phi2, o))


# -

# If your code is correct, you should see the following output:
#
# ```
# | A   | B   | D   |   Pr | Ext                        |
# |-----+-----+-----+------+----------------------------|
# | +a  | +b  | +d  | 0.03 | [('C', '+c'), ('E', '-e')] |
# | +a  | +b  | -d  | 0.27 | [('C', '+c'), ('E', '+e')] |
# | +a  | -b  | +d  | 0.28 | [('C', '+c'), ('E', '-e')] |
# | +a  | -b  | -d  | 0.42 | [('C', '+c'), ('E', '+e')] |
# | -a  | +b  | +d  | 0.02 | [('C', '+c'), ('E', '-e')] |
# | -a  | +b  | -d  | 0.18 | [('C', '+c'), ('E', '+e')] |
# | -a  | -b  | +d  | 0.32 | [('C', '-c'), ('E', '-e')] |
# | -a  | -b  | -d  | 0.48 | [('C', '-c'), ('E', '+e')] |
# ```

# ## MPE Variable Elimination
#
# We can now implement the MPE VE algorithm. Let us make it in two parts. The first one is the variable elimination part and the second one the query part. Therefore, we will be able to reuse the VE implementation from week 4. 
#
# The MPE VE implementation is very similar to the `VE` function we did before. 
#
# ### Exercise
#
# Implement the `MPE_VE` function based on the `VE` function from week 6. The main difference is that we now eliminate variables maximizing out instead of summing out.

# +
def MPE_VE(factors, order, outcomeSpace):
    """
    argument 
    `factors`, a dictionary of factors, each factor is a dictionary of domain and probability values.
    `order`, a list of variable names specifying an elimination order.
    `outcomeSpace`, a dictionary with variable names and respective domains.
    Returns a dictionary with non-eliminated factors. Variables are eliminated by maximization.
    """    

    # Let's make a copy of factors, so we can freely modify it without destroying the original dictionary
    f = factors.copy()
    # We process the factor in elimination order  
    for i, var in enumerate(order):
        # This is the domain of the new factor. We use sets as it is handy to eliminate duplicate variables
        newFactorDom = set()
        # This is a list of factors that will be removed from f because they were joined with other factors
        listFactorsRemove = list()
        # This is a flag to indicate if we are processing the first factor
        first = True
        # Lets iterate over all factors
        for f_id in f.keys():
            # and select the ones that have the variable to be eliminated
            if var in f[f_id]['dom']:        
                if first:
                    # We need this code since join requires two factors, so we save the first one in fx and wait for the next
                    fx = f[f_id]
                    first = False
                else:
                    # Join fx and f[f_id] and save the result in fx
                    fx = join(fx, f[f_id], outcomeSpace)
                # f_id was joined, so we will need to eliminate it from f later. Let's save that factor id for future removal
                listFactorsRemove.append(f_id)
        # Now, we need to remove var from the domain of the new factor doing a maximization    
        fx = None                                     # Eliminate variable var from fx by doing maximization: 1 line
        # Now, we remove all factors that we joined. We do it outside the for loop since it modifies the data structure
        for f_id in listFactorsRemove:
            del f[f_id]
        # We will create a new factor with id equal a sequential number and insert it into f, so it can be used in future joins          
        f[i] = fx
    return f


#########################
# Test code

printFactor(MPE_VE(factors, ('A', 'T1', 'T2','S', 'C'), outcomeSpace)[4])


# +
## Answer

def MPE_VE(factors, order, outcomeSpace):
    """
    argument 
    `factors`, a dictionary of factors, each factor is a dictionary of domain and probability values.
    `order`, a list of variable names specifying an elimination order.
    `outcomeSpace`, a dictionary with variable names and respective domains.
    Returns a dictionary with non-eliminated factors. Variables are eliminated by maximization.
    """    

    # Let's make a copy of factors, so we can freely modify it without destroying the original dictionary
    f = factors.copy()
    # We process the factor in elimination order  
    for i, var in enumerate(order):
        # This is the domain of the new factor. We use sets as it is handy to eliminate duplicate variables
        newFactorDom = set()
        # This is a list of factors that will be removed from f because they were joined with other factors
        listFactorsRemove = list()
        # This is a flag to indicate if we are processing the first factor
        first = True
        # Lets iterate over all factors
        for f_id in f.keys():
            # and select the ones that have the variable to be eliminated
            if var in f[f_id]['dom']:        
                if first:
                    # We need this code since join requires two factors, so we save the first one in fx and wait for the next
                    fx = f[f_id]
                    first = False
                else:
                    # Join fx and f[f_id] and save the result in fx
                    fx = join(fx, f[f_id], outcomeSpace)
                # f_id was joined, so we will need to eliminate it from f later. Let's save that factor id for future removal
                listFactorsRemove.append(f_id)
        # Now, we need to remove var from the domain of the new factor doing a maximization    
        fx = maximize(fx, var, outcomeSpace)        # Eliminate variable var from fx by doing maximization: 1 line
        # Now, we remove all factors that we joined. We do it outside the for loop since it modifies the data structure
        for f_id in listFactorsRemove:
            del f[f_id]
        # We will create a new factor with id equal a sequential number and insert it into f, so it can be used in future joins          
        f[i] = fx
    return f


#########################
# Test code

printFactor(MPE_VE(factors, ('A', 'T1', 'T2','S', 'C'), outcomeSpace)[4])


# -

# If your code is correct, you should see the following output:
#
# ```
# |      Pr | Ext                                                                        |
# |---------+----------------------------------------------------------------------------|
# | 0.33858 | [('A', 'yes'), ('T1', '-ve'), ('T2', '-ve'), ('S', 'female'), ('C', 'no')] |
# ````
#
# This result matches with the one we obtained in slide 27 from lecture 14.

# ## MPE Query
#
# Now that we have MPE VE, we can easily implement MPE Query. The function `MPE_query` receives as argument a dictionary `q_evi` with the evidence information. After setting the evidence according to the entries in this dictionary, `MPE_query` calls `MPE_VE` to eliminate all variables according to a provided elimination `order`. 
#
# ### Exercise
#
# Implement the function `MPE_query`. We have provided most of the code for you.

# +
def MPE_query(factors, order, outcomeSpace, **q_evi):
    """
    argument 
    `factors`, a dictionary of factors
    `order`, a list with variable elimination order
    `outcomeSpace`, dictionary will variable domains
    `q_evi`, dictionary of evidence in the form of variables names and values
    
    Returns a new factor with MPE(e)
    """     
    
    # Let's make a copy of these structures, since we will reuse the variable names
    outSpace = outcomeSpace.copy()
    
    # First, we set the evidence 
    for var_evi, e in q_evi.items():
        outSpace = None                     # Set the evidence according to q_evi. Use the evidence function: 1 line
 
    # Call MPE_VE to eliminate all variables according to the provided elimination order
    f = None                                # Call MPE_VE to eliminate variables using maximization: 1 line
    
    # The remaining code will join all remaining factors in f into a single factor fx
    first = True
    for f_id in f.keys():
        if first:
            # We need this code since join requires two factors, so we save the first one in fx and wait for the next
            fx = f[f_id]
            first = False
        else:
            # Join fx and f[f_id] and save the result in fx
            fx = join(fx, f[f_id], outSpace)    
    
    return fx

#########################
# Test code
#########################

def evidence(var, e, outcomeSpace):
    """
    argument 
    `var`, a valid variable identifier.
    `e`, the observed value for var.
    `outcomeSpace`, dictionary with the domain of each variable
    
    Returns dictionary with a copy of outcomeSpace with var = e
    """    
    newOutcomeSpace = outcomeSpace.copy()      # Make a copy of outcomeSpace with a copy to method copy(). 1 line
    newOutcomeSpace[var] = (e,)                # Replace the domain of variable var with a tuple with a single element e. 1 line
    return newOutcomeSpace


printFactor(MPE_query(factors, ('A', 'T1', 'T2','S','C'), outcomeSpace, A='yes'))


# +
#Answer

def MPE_query(factors, order, outcomeSpace, **q_evi):
    """
    argument 
    `factors`, a dictionary of factors
    `order`, a list with variable elimination order
    `outcomeSpace`, dictionary will variable domains
    `q_evi`, dictionary of evidence in the form of variables names and values
    
    Returns a new factor with MPE(e)
    """     
    
    # Let's make a copy of these structures, since we will reuse the variable names
    outSpace = outcomeSpace.copy()
    
    # First, we set the evidence 
    for var_evi, e in q_evi.items():
        outSpace = evidence(var_evi, e, outSpace) # Set the evidence according to q_evi. Use the evidence function: 1 line
 
    # Call MPE_VE to eliminate all variables according to the provided elimination order
    f = MPE_VE(factors, order, outSpace)    # Call MPE_VE to eliminate variables using maximization: 1 line
    
    # The remaining code will join all remaining factors in f into a single factor fx
    first = True
    for f_id in f.keys():
        if first:
            # We need this code since join requires two factors, so we save the first one in fx and wait for the next
            fx = f[f_id]
            first = False
        else:
            # Join fx and f[f_id] and save the result in fx
            fx = join(fx, f[f_id], outSpace)    
    
    return fx

#########################
# Test code
#########################

def evidence(var, e, outcomeSpace):
    """
    argument 
    `var`, a valid variable identifier.
    `e`, the observed value for var.
    `outcomeSpace`, dictionary with the domain of each variable
    
    Returns dictionary with a copy of outcomeSpace with var = e
    """    
    newOutcomeSpace = outcomeSpace.copy()      # Make a copy of outcomeSpace with a copy to method copy(). 1 line
    newOutcomeSpace[var] = (e,)                # Replace the domain of variable var with a tuple with a single element e. 1 line
    return newOutcomeSpace

printFactor(MPE_query(factors, ('A', 'T1', 'T2','S','C'), outcomeSpace, A='yes'))


# -

# If your code is correct, you should see the following output:
#
# ```
# |      Pr | Ext                                                                        |
# |---------+----------------------------------------------------------------------------|
# | 0.33858 | [('A', 'yes'), ('T1', '-ve'), ('T2', '-ve'), ('S', 'female'), ('C', 'no')] |
# ````
#
# This output matches the results presented in slide 43 of lecture 14.

# ## MAP Variable Elimination
#
# Let us implement the MAP VE algorithm in the same steps we did for MPE. We start with a `MAP_VE` function that is very similar to the `MPE_VE` we just implemented.
#
# The main difference between then is that `MAP_VE` receives an additional argument `map_vars` with a list of MAP variables. To provide the correct output, `map_vars` should be the last variables of `order`. However, we will not enforce that constraint. Calling `map_vars` when `map_vars` are not the last variables in `order` has its utility since the output will be an upper bound to the MAP probability. 
#
# MAP VE should remove all variables in `order`. Differently from MPE VE, non-MAP variables are eliminated using summation while MAP variables are eliminated using maximization. 
#
# ### Exercise
#
# Implement the `MAP_VE` function. You can use the `MPE_VE` function implemented before in the tutorial. Remember to eliminate MAP and non-MAP variables with maximization and summation, respectively.
#
# We create a stub for your function. You will implement most of the code. Use `MPE_VE` as a starting point.

# +
def MAP_VE(factors, order, map_vars, outcomeSpace):
    """
    argument 
    `factors`, a dictionary of factors, each factor is a dictionary of domain and probability values,
    `order`, a list of variable names specifying an elimination order,
    `map_vars`, a list of MAP variables. Although the code does not enforce that, these variables should be the last ones in `order`    
    `outcomeSpace`, a dictionary with variable names and respective domains.
    Returns a dictionary with non-eliminated factors
    """
    
    None
    
##########################
# Test code

def marginalize(f, var, outcomeSpace):
    """
    argument 
    `f`, factor to be marginalized.
    `var`, variable to be summed out.
    `outcomeSpace`, dictionary with the domain of each variable
    
    Returns a new factor f' with dom(f') = dom(f) - {var}
    """    
    
    # Let's make a copy of f domain and convert it to a list. We need a list to be able to modify its elements
    new_dom = list(f['dom'])
    new_dom.remove(var)            # Remove var from the list new_dom by calling the method remove(). 1 line
    table = list()                 # Create an empty list for table. We will fill in table from scratch. 1 line
    for entries in product(*[outcomeSpace[node] for node in new_dom]):
        s = 0;                     # Initialize the summation variable s. 1 line

        # We need to iterate over all possible outcomes of the variable var
        for val in outcomeSpace[var]:
            # To modify the tuple entries, we will need to convert it to a list
            entriesList = list(entries)
            # We need to insert the value of var in the right position in entriesList
            entriesList.insert(f['dom'].index(var), val)
          
            p = prob(f, *tuple(entriesList))     # Calculate the probability of factor f for entriesList. 1 line
            s = s + p                            # Sum over all values of var by accumulating the sum in s. 1 line
            
        # Create a new table entry with the multiplication of p1 and p2
        table.append((entries, s))
    return {'dom': tuple(new_dom), 'table': odict(table)}

printFactor(MAP_VE(factors, ('A', 'T1', 'T2', 'S', 'C'), ('S', 'C'), outcomeSpace)[4])    


# +
## Answer

def MAP_VE(factors, order, map_vars, outcomeSpace):
    """
    argument 
    `factors`, a dictionary of factors, each factor is a dictionary of domain and probability values,
    `order`, a list of variable names specifying an elimination order,
    `map_vars`, a list of MAP variables. Although the code does not enforce that, these variables should be the last ones in `order`    
    `outcomeSpace`, a dictionary with variable names and respective domains.
    Returns a dictionary with non-eliminated factors
    """    

    # Let's make a copy of factors, so we can freely modify it without destroying the original dictionary
    f = factors.copy()
    # We process the factor in elimination order
    for i, var in enumerate(order):
        # This is the domain of the new factor. We use sets as it is handy to eliminate duplicate variables
        newFactorDom = set()
        # This is a list of factors that will be removed from f because they were joined with other factors
        listFactorsRemove = list()
        # This is a flag to indicate if we are processing the first factor
        first = True
        # Lets iterate over all factors
        for f_id in f.keys():
            # and select the ones that have the variable to be eliminated
            if var in f[f_id]['dom']:        
                if first:
                    # We need this code since join requires two factors, so we save the first one in fx and wait for the next
                    fx = f[f_id]
                    first = False
                else:
                    # Join fx and f[f_id] and save the result in fx
                    fx = join(fx, f[f_id], outcomeSpace)
                # f_id was joined, so we will need to eliminate it from f later. Let's save that factor id for future removal
                listFactorsRemove.append(f_id)
        # Now, we need to remove var from the domain of the new factor doing a marginalization    
        if (var in map_vars):
            fx = maximize(fx, var, outcomeSpace)
        else:
            fx = marginalize(fx, var, outcomeSpace)
        # Now, we remove all factors that we joined in the simulation. We do it outside the for loop since it modifies the data structure
        for f_id in listFactorsRemove:
            del f[f_id]
        # We will create a new factor with id equal a sequential number and insert it into f, so it can be used in future joins          
        f[i] = fx
    return f

##########################
# Test code

def marginalize(f, var, outcomeSpace):
    """
    argument 
    `f`, factor to be marginalized.
    `var`, variable to be summed out.
    `outcomeSpace`, dictionary with the domain of each variable
    
    Returns a new factor f' with dom(f') = dom(f) - {var}
    """    
    
    # Let's make a copy of f domain and convert it to a list. We need a list to be able to modify its elements
    new_dom = list(f['dom'])
    new_dom.remove(var)            # Remove var from the list new_dom by calling the method remove(). 1 line
    table = list()                 # Create an empty list for table. We will fill in table from scratch. 1 line
    for entries in product(*[outcomeSpace[node] for node in new_dom]):
        s = 0;                     # Initialize the summation variable s. 1 line

        # We need to iterate over all possible outcomes of the variable var
        for val in outcomeSpace[var]:
            # To modify the tuple entries, we will need to convert it to a list
            entriesList = list(entries)
            # We need to insert the value of var in the right position in entriesList
            entriesList.insert(f['dom'].index(var), val)
          
            p = prob(f, *tuple(entriesList))     # Calculate the probability of factor f for entriesList. 1 line
            s = s + p                            # Sum over all values of var by accumulating the sum in s. 1 line
            
        # Create a new table entry with the multiplication of p1 and p2
        table.append((entries, s))
    return {'dom': tuple(new_dom), 'table': odict(table)}

printFactor(MAP_VE(factors, ('A', 'T1', 'T2', 'S', 'C'), ('S', 'C'), outcomeSpace)[4])


# -

# If your output is correct, you should see the following output:
#
# ```
# |     Pr | Ext                          |
# |--------+------------------------------|
# | 0.5225 | [('S', 'male'), ('C', 'no')] |
# ```

# ## MAP Query
#
# Now that we have MAP VE, we can also implement MAP Query. The function `MAP_query`  is very similar to `MPE_query` with the difference it uses `MAP_VE` instead of `MPE_VE`. Therefore, it will eliminate MAP and non-MAP variables differently.
#
# ### Exercise
#
# Implement the function `MAP_query`. We have provided a stub for you. Use the function `MPE_query` as a starting point.

# +
def MAP_query(factors, order, outcomeSpace, map_vars, **q_evi):
    """
    argument 
    `factors`, a dictionary of factors
    `order`, a list with variable elimination order
    `outcomeSpace`, dictionary will variable domains
    `map_vars`, a list of MAP variables. Although the code does not enforce that, these variables should be the last ones in `order`    
    `q_evi`, dictionary of evidence in the form of variables names and values
    
    Returns a new factor with P(Q, e) 
    """     
    None
    
#########################
# Test code
printFactor(MAP_query(factors, ('A', 'T1', 'T2','S','C'), outcomeSpace, ('S', 'C'), A='yes'))


# +
#Answer

def MAP_query(factors, order, outcomeSpace, map_vars, **q_evi):
    """
    argument 
    `factors`, a dictionary of factors
    `order`, a list with variable elimination order
    `outcomeSpace`, dictionary will variable domains
    `map_vars`, a list of MAP variables. Although the code does not enforce that, these variables should be the last ones in `order`    
    `q_evi`, dictionary of evidence in the form of variables names and values
    
    Returns a new factor with P(Q, e) or P(Q|e)
    """     
    
    # Let's make a copy of these structures, since we will reuse the variable names
    outSpace = outcomeSpace.copy()
    
    # First, we set the evidence 
    for var_evi, e in q_evi.items():
        outSpace = evidence(var_evi, e, outSpace)
 
    f = MAP_VE(factors, order, map_vars, outSpace)
    
    first = True
    for f_id in f.keys():
        if first:
            # We need this code since join requires two factors, so we save the first one in fx and wait for the next
            fx = f[f_id]
            first = False
        else:
            # Join fx and f[f_id] and save the result in fx
            fx = join(fx, f[f_id], outSpace)    
    
    return fx

#########################
# Test code
printFactor(MAP_query(factors, ('A', 'T1', 'T2','S','C'), outcomeSpace, ('S', 'C'), A='yes'))
# -

# If your code is correct, you should see the following output:
#
# ```
# |     Pr | Ext                          |
# |--------+------------------------------|
# | 0.3553 | [('S', 'male'), ('C', 'no')] |
# ```
#
# This matches the results shown in slide 78 of lecture 12.

# That is all for today. See you next week!
