import numpy as np
from pyugm.factor import DiscreteFactor
from pyugm.factor import DiscreteBelief

# Specify the parameters (should be the same shape as the potential table would have been)
factor_parameters = np.array([['theta_0', 'theta_1'], ['theta_2', 'theta_3']])  
variables_names_and_cardinalities = [(1, 2), (2, 2)]
# Construct the factor
factor = DiscreteFactor(variables_names_and_cardinalities, parameters=factor_parameters)
print factor

# The factor still has its default potential table (all ones)
factor.data

# Create a belief based on the factor
belief = DiscreteBelief(factor)

# Potentials are filled with the exponent of the parameters.
belief.set_parameters({'theta_0': np.log(2), 'theta_1': np.log(0.2), 'theta_2': np.log(5), 'theta_3': np.log(1)})
belief.data

belief.set_parameters({'theta_0': np.log(1), 'theta_1': np.log(2), 'theta_2': np.log(1), 'theta_3': np.log(1)})
belief.data

factor_parameters = np.array([['theta_0', 'theta_1'], ['theta_0', 'theta_0']])  
factor = DiscreteFactor([(1, 2), (2, 2)], parameters=factor_parameters)
belief = DiscreteBelief(factor)
belief.set_parameters({'theta_0': np.log(3), 'theta_1': np.log(5)})
print belief.data

import pickle
import matplotlib.pyplot as plt
import matplotlib
%matplotlib inline
import seaborn

seaborn.set_style("dark")
# Load pre-discretized image. Each pixel value is an integer between 0 and 32.
image = pickle.load(open('test_image.pkl'))
plt.figure(figsize=(14, 3))
plt.subplot(1, 2, 1)
_ = plt.imshow(image, cmap=matplotlib.cm.Greys_r, interpolation='nearest')
_ = plt.title('Image')
seaborn.set_style("darkgrid")
plt.subplot(1, 2, 2)
_ = plt.hist(image.flatten(), bins=32, color=(0.5, 0.5, 0.5))
_ = plt.title('Pixel intensity histogram')

observation_template = np.array([['obs_low'] * 32,
                                 ['obs_high'] * 32])
observation_template[0, 13:17] = 'obs_high'
observation_template[1, 13:17] = 'obs_low'

I, J = image.shape
factors = []
evidence = {}

# Add observation factors
for i in xrange(I):
    for j in xrange(J):
        label_variable_name = 'label_{}_{}'.format(i, j)
        observation_variable_name = 'obs_{}_{}'.format(i,j)
        factors.append(DiscreteFactor([(label_variable_name, 2), (observation_variable_name, 32)], parameters=observation_template))
        evidence[observation_variable_name] = image[i, j]

from pyugm.model import Model
model = Model(factors)

from pyugm.infer_message import LoopyBeliefUpdateInference
from pyugm.infer_message import FloodingProtocol

order = FloodingProtocol(model, max_iterations=30)
inference = LoopyBeliefUpdateInference(model, order)

parameters = {'obs_high': 0.1, 'obs_low': -1.0}

inference.calibrate(evidence, parameters)

labels = np.zeros(image.shape)
for i in xrange(I):
    for j in xrange(J):
        variable_name = 'label_{}_{}'.format(i, j)
        label_factor = inference.get_marginals(variable_name)[0]
        labels[i, j] = label_factor.normalized_data[0] 

seaborn.set_style("dark")
plt.figure(figsize=(14, 3))
plt.subplot(1, 3, 1)
_ = plt.imshow(image, cmap=matplotlib.cm.Greys_r, interpolation='nearest')
_ = plt.title('Original image')
plt.subplot(1, 3, 2)
_ = plt.imshow(labels, cmap=matplotlib.cm.Greys, interpolation='nearest')
_ = plt.title('Inferred labels (white=foreground, black=background)')

label_template = np.array([['same', 'different'], 
                           ['different', 'same']])

evidence = {}
factors = []

# Add observation factors
for i in xrange(I):
    for j in xrange(J):
        label_variable_name = 'label_{}_{}'.format(i, j)
        observation_variable_name = 'obs_{}_{}'.format(i, j)
        factors.append(DiscreteFactor([(label_variable_name, 2), (observation_variable_name, 32)], parameters=observation_template))
        evidence[observation_variable_name] = image[i, j] 
        
# Add label factors
for i in xrange(I):
    for j in xrange(J):
        variable_name = 'label_{}_{}'.format(i, j)
        if i + 1 < I:
            neighbour_down_name = 'label_{}_{}'.format(i + 1, j)
            factors.append(DiscreteFactor([(variable_name, 2), (neighbour_down_name, 2)], parameters=label_template))

model = Model(factors)

from pyugm.infer_message import LoopyBeliefUpdateInference
from pyugm.infer_message import FloodingProtocol

# Get some feedback on how inference is converging by listening in on some of the label beliefs.
var_values = {'label_1_1': [], 'label_10_10': [], 'label_20_20': [], 'label_30_30': [], 'label_40_40': []}
changes = []
def reporter(infe, orde):
    for var in var_values.keys():
        marginal = infe.get_marginals(var)[0].data[0]
        var_values[var].append(marginal)
    change = orde.last_iteration_delta
    changes.append(change)
    energy = infe.partition_approximation()
    print '{:3} {:8.2f} {:5.2f} {:8.2f}'.format(orde.total_iterations, change, marginal, energy)
    
order = FloodingProtocol(model, max_iterations=15)
inference = LoopyBeliefUpdateInference(model, order, callback=reporter)

parameters = {'same': 2.0, 'different': -1.0, 'obs_high': 1.0, 'obs_low': -0.0}

inference.calibrate(evidence, parameters)

labels = np.zeros(image.shape)
for i in xrange(I):
    for j in xrange(J):
        variable_name = 'label_{}_{}'.format(i, j)
        label_factor = inference.get_marginals(variable_name)[0]
        labels[i, j] = label_factor.normalized_data[0] 

plt.figure(figsize=(14, 3))
plt.subplot(1, 3, 1)
_ = plt.imshow(image, cmap=matplotlib.cm.Greys_r, interpolation='nearest')
_ = plt.title('Original image')
plt.subplot(1, 3, 2)
_ = plt.imshow(labels, cmap=matplotlib.cm.Greys, interpolation='nearest')
_ = plt.title('Label beliefs \n(darker=higher background belief,\n lighter=higher foreground belief')
plt.subplot(1, 3, 3)
_ = plt.imshow(labels > 0.5, cmap=matplotlib.cm.Greys, interpolation='nearest')
_ = plt.title('Thresholded beliefs')

evidence = {}
factors = []
# Add observation factors
for i in xrange(I):
    for j in xrange(J):
        label_variable_name = 'label_{}_{}'.format(i, j)
        observation_variable_name = 'obs_{}_{}'.format(i, j)
        factors.append(DiscreteFactor([(label_variable_name, 2), (observation_variable_name, 32)], parameters=observation_template))
        evidence[observation_variable_name] = image[i, j]
        
# Add label factors
for i in xrange(I):
    for j in xrange(J):
        variable_name = 'label_{}_{}'.format(i, j)
        if i + 1 < I:
            neighbour_down_name = 'label_{}_{}'.format(i + 1, j)
            factors.append(DiscreteFactor([(variable_name, 2), (neighbour_down_name, 2)], parameters=label_template))
        if j + 1 < J:
            neighbour_right_name = 'label_{}_{}'.format(i, j + 1)
            factors.append(DiscreteFactor([(variable_name, 2), (neighbour_right_name, 2)], parameters=label_template))

model = Model(factors)

from pyugm.infer_message import LoopyBeliefUpdateInference
from pyugm.infer_message import FloodingProtocol, LoopyDistributeCollectProtocol

# Get some feedback on how inference is converging by listening in on some of the label beliefs.
var_values = {'label_1_1': [], 'label_10_10': [], 'label_20_20': [], 'label_30_30': [], 'label_40_40': []}
changes = []
partitions = []
def reporter(infe, orde):
    for var in var_values.keys():
        marginal = infe.get_marginals(var)[0].data[0]
        var_values[var].append(marginal)
    change = orde.last_iteration_delta
    changes.append(change)
    energy = infe.partition_approximation()
    partitions.append(energy)
    print '{:3} {:8.2f} {:5.2f} {:8.2f}'.format(orde.total_iterations, change, marginal, energy)
    
order = FloodingProtocol(model, max_iterations=15)
inference = LoopyBeliefUpdateInference(model, order, callback=reporter)

parameters = {'same': 0.1, 'different': -1.0, 'obs_high': 1.0, 'obs_low': -0.0}
inference.calibrate(evidence, parameters)

seaborn.set_style("darkgrid")
ax1 = plt.axes()
ax2 = ax1.twinx()
_ = ax1.plot([np.log(change) for change in changes])
seaborn.set_style('dark')
_ = ax2.plot(partitions, color='g')
ax2.set_yticks(np.linspace(ax2.get_yticks()[0], ax2.get_yticks()[-1], len(ax1.get_yticks())))
_ = ax1.set_ylabel('Log of total belief change')
_ = ax2.set_ylabel('Log of the approximation to the partition function Z')
_ = ax1.set_xlabel('Iteration number')

for key, value in var_values.items():
    plt.plot(value, label=key)
_ = plt.ylabel('Normalized belief that the variable equals 0')
_ = plt.xlabel('Iteration number')
plt.legend()

labels = np.zeros(image.shape)
for i in xrange(I):
    for j in xrange(J):
        variable_name = 'label_{}_{}'.format(i, j)
        label_factor = inference.get_marginals(variable_name)[0]
        labels[i, j] = label_factor.normalized_data[0] 

seaborn.set_style("dark")
plt.figure(figsize=(14, 3))
plt.subplot(1, 3, 1)
_ = plt.imshow(image, cmap=matplotlib.cm.Greys_r, interpolation='nearest')
_ = plt.title('Origin image')
plt.subplot(1, 3, 2)
_ = plt.imshow(labels, cmap=matplotlib.cm.Greys, interpolation='nearest')
_ = plt.title('Label beliefs \n(darker=higher background belief,\n lighter=higher foreground belief')
plt.subplot(1, 3, 3)
_ = plt.imshow(labels > 0.5, cmap=matplotlib.cm.Greys, interpolation='nearest')
_ = plt.title('Thresholded beliefs')