"""
Recover the toy dataset generated by example/generate_toy/bnmf/generate_bnmf.py
using ICM, and plot the MSE against timestamps.

We can plot the MSE, R2 and Rp as it converges, on the entire dataset.

We have I=100, J=80, K=10, and no test data.
We give flatter priors (1/10) than what was used to generate the data (1).
"""

import sys, os
project_location = os.path.dirname(__file__)+"/../../../../"
sys.path.append(project_location)

from BNMTF.code.models.nmf_icm import nmf_icm

import numpy, random, scipy, matplotlib.pyplot as plt

##########

input_folder = project_location+"BNMTF/data_toy/bnmf/"

repeats = 10

iterations = 2000
init_UV = 'random'
I, J, K = 100,80,10

minimum_TN = 0.1

alpha, beta = 1., 1. #1., 1.
lambdaU = numpy.ones((I,K))/10.
lambdaV = numpy.ones((J,K))/10.
priors = { 'alpha':alpha, 'beta':beta, 'lambdaU':lambdaU, 'lambdaV':lambdaV }

# Load in data
R = numpy.loadtxt(input_folder+"R.txt")
M = numpy.ones((I,J))

# Run the VB algorithm, <repeats> times
times_repeats = []
performances_repeats = []
for i in range(0,repeats):
    # Set all the seeds
    numpy.random.seed(0)
    random.seed(0)
    scipy.random.seed(0)
    
    # Run the classifier
    nmf = nmf_icm(R,M,K,priors) 
    nmf.initialise(init_UV)
    nmf.run(iterations,minimum_TN=minimum_TN)

    # Extract the performances and timestamps across all iterations
    times_repeats.append(nmf.all_times)
    performances_repeats.append(nmf.all_performances)

# Check whether seed worked: all performances should be the same
assert all([numpy.array_equal(performances, performances_repeats[0]) for performances in performances_repeats]), \
    "Seed went wrong - performances not the same across repeats!"

# Print out the performances, and the average times
icm_all_times_average = list(numpy.average(times_repeats, axis=0))
icm_all_performances = performances_repeats[0]
print "icm_all_times_average = %s" % icm_all_times_average
print "icm_all_performances = %s" % icm_all_performances


# Print all time plots, the average, and performance vs iterations
plt.figure()
plt.title("Performance against time")
plt.ylim(0,10)
for times in times_repeats:
    plt.plot(times, icm_all_performances['MSE'])

plt.figure()
plt.title("Performance against average time")
plt.plot(icm_all_times_average, icm_all_performances['MSE'])
plt.ylim(0,10)

plt.figure()
plt.title("Performance against iteration")
plt.plot(icm_all_performances['MSE'])
plt.ylim(0,10)