import bnpy

from matplotlib import pylab
%pylab inline

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

import BarsK10V900
Data = BarsK10V900.get_data(nDocTotal=1000, nWordsPerDoc=100)
Data.name = 'BarsK10V900'

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

bnpy.viz.BarsViz.plotExampleBarsDocs(Data, nDocToPlot=16, **imshowArgs);

hmodel, RInfo = bnpy.run(Data, 'HDPTopicModel', 'Mult', 'moVB',
                         moves='merge,delete',
                         mergePerLap=25, mergePairSelection='corr', 
                         mergeNumStuckBeforeQuit=5, mergeStartLap=3,
                         dtargetMaxSize=100,
                         K=20, alpha=0.1, lam=0.1, initname='kmeansplusplus',
                         nLap=50, printEvery=5, nCoordAscentItersLP=25,
                         nTask=2, jobname='demomergedelete-Kinit=20-kmeans')

hmodel, RInfo = bnpy.run(Data, 'HDPTopicModel', 'Mult', 'moVB',
                         moves='merge,delete',
                         mergePerLap=25, mergePairSelection='corr', 
                         mergeNumStuckBeforeQuit=5, mergeStartLap=3,
                         dtargetMaxSize=100,
                         K=40, alpha=0.1, lam=0.1, initname='kmeansplusplus',
                         nLap=50, printEvery=5, nCoordAscentItersLP=25,
                         nTask=2, jobname='demomergedelete-Kinit=40-kmeans')

bnpy.viz.PlotELBO.plotJobsThatMatchKeywords('BarsK10V900/demomergedelete-*');
pylab.ylabel('objective score', fontsize=20);
pylab.xlabel('num laps thru dataset', fontsize=20);
pylab.legend(loc='lower right');

bnpy.viz.PlotK.plotJobsThatMatchKeywords('BarsK10V900/demomergedelete-*', yvar='K');
pylab.ylim([0, 50]); pylab.ylabel('num topics K', fontsize=20);
pylab.xlabel('num laps thru dataset', fontsize=20);

bnpy.viz.PlotComps.plotCompsForTask('BarsK10V900/demomergedelete-Kinit=20-kmeans/1/', **imshowArgs)

bnpy.viz.PlotComps.plotCompsForTask('BarsK10V900/demomergedelete-Kinit=40-kmeans/1/', **imshowArgs)