import numpy as np
import bnpy

from matplotlib import pylab
%pylab inline

imshowArgs = dict(interpolation='nearest', 
                  cmap='bone_r', 
                  vmin=0.0, 
                  vmax=10./900,
                  )

import MixBarsK10V900
Data = MixBarsK10V900.get_data()
Data.name = 'MixBarsK10V900'
bnpy.viz.BarsViz.plotExampleBarsDocs(Data)

bnpy.viz.BarsViz.showTopicsAsSquareImages(MixBarsK10V900.Defaults['topics'], **imshowArgs);

hmodel, RInfo = bnpy.run(Data, 'FiniteMixtureModel', 'Mult', 'EM', 
                         jobname='true-K-randexamples', K=10, nLap=100, nTask=5,
                         initname='randexamples')

bnpy.viz.PlotComps.plotCompsForJob('MixBarsK10V900/true-K-randexamples', **imshowArgs)

hmodel, RInfo = bnpy.run(Data, 'FiniteMixtureModel', 'Mult', 'EM', 
                         jobname='true-K-kmeans', K=10, nLap=500, nTask=5,
                         initname='kmeansplusplus')

bnpy.viz.PlotELBO.plotJobsThatMatchKeywords('MixBarsK10V900/true-K-*')
pylab.ylim([-10, -5])
pylab.legend(loc='lower right');

for K in [5, 10, 20, 50]: 
    hmodel, RInfo = bnpy.run(Data, 'FiniteMixtureModel', 'Mult', 'EM', jobname='em-K=%d'%(K),
                             K=K, nLap=100, doWriteStdOut=False, nTask = 1)
    print 'K=%3d | final ev % .5f' % (K, RInfo['evBound'])

bnpy.viz.PlotELBO.plotJobsThatMatchKeywords('MixBarsK10V900/em-K=*')
pylab.ylim([-10, -5])