doSaveEPS = 1; # If set to 1, will save each plots to .eps for LaTeX publication.

import numpy as np

import glob; # for simple check that saved runs exist
import os;
import bnpy; # bnpy package for learning and plotting BNP models
import SpeakerDiar # Module for the dataset

from bnpy.viz.PlotUtil import pylab
%pylab inline

bnpy.viz.PlotUtil.ConfigPylabDefaults(pylab);

import Helper_SpeakerDiar as Helper
reload(Helper);
reload(Helper.PlotUtil);

jobName = 'nips2015'
pathPattern = "%s/%s/%s" % (os.environ['BNPYOUTDIR'], 'SpeakerDiar*', jobName + '*')
dirList = glob.glob(pathPattern)
if len(dirList) == 0:
    raise ValueError("STOP! You have not run the expected experiments yet!")
print "Success! Found %d directories for jobs under the name '%s' for SpeakerDiar" % (len(dirList), jobName)

Jdict = Helper.setUp()

for key in Jdict.keys():
    print key

scores = Helper.plotScatterComparison(Jdict['delmerge Sticky=100 K=25'],
                                      Jdict['sampler Sticky=100 K=25'],
                                      nTask=10, pylab=pylab)
if doSaveEPS:
    pylab.savefig('SpeakerDiar_CompareFinalHamming.eps', bbox_inches='tight');

Helper.makeLegendInOwnFigure(Jdict, names=['sampler', 'memo', 'delmerge:delete,merge', 'birth:birth,delete,merge'])
if doSaveEPS:
    pylab.savefig('SpeakerDiar_LegendForTracePlots.eps', bbox_inches='tight');
    # Trim the bounding box a bit
    !sed -i -e 's/BoundingBox: 171 275 440 516/BoundingBox: 205 390 425 505/' SpeakerDiar_LegendForTracePlots.eps

n = '11'
Helper.PlotUtil.plotHammingDist(Jdict, n=n, loc=None, xscale='log', xvar='times');
if doSaveEPS:
    pylab.savefig('SpeakerDiar11_HammingVsTime.eps', bbox_inches='tight');

n = '11'
Helper.PlotUtil.plotELBO(Jdict, n=n, loc=None, xscale='log', xvar='times');
if doSaveEPS:
    pylab.savefig('SpeakerDiar11_ELBOVsTime.eps', bbox_inches='tight');

Helper.PlotUtil.plotStateSeq(Jdict['birth Sticky=100 K=25'], taskids='.best', sequences=['11']);

n = '21'
Helper.PlotUtil.plotHammingDist(Jdict, n=n, loc=None, xscale='log', xvar='times');
#pylab.ylim([-.01, 0.4]);
#pylab.gca().set_yticks([0, 0.1, 0.2, 0.3]);
if doSaveEPS:
    pylab.savefig('SpeakerDiar21_HammingVsTime.eps', bbox_inches='tight');

n = '21'
Helper.PlotUtil.plotELBO(Jdict, n=n, loc=None, xscale='log', xvar='times');
pylab.gca().set_yticks([-2.55, -2.5, -2.45, -2.4]);
pylab.ylim([-2.56, -2.39]);
if doSaveEPS:
    pylab.savefig('SpeakerDiar21_ELBOVsTime.eps', bbox_inches='tight');

Helper.PlotUtil.plotStateSeq(Jdict['delmerge Sticky=100 K=25'], taskids='.best', sequences=['21']);

n = '16'
Helper.PlotUtil.plotHammingDist(Jdict, n=n, loc=None, xscale='log', xvar='times');
if doSaveEPS:
    pylab.savefig('SpeakerDiar16_HammingVsTime.eps', bbox_inches='tight');

n = '16'
Helper.PlotUtil.plotELBO(Jdict, n=n, loc=None, xscale='log', xvar='times');
pylab.gca().set_yticks([-2.7, -2.65, -2.6, -2.55]);
pylab.ylim([-2.71, -2.54]);
if doSaveEPS:
    pylab.savefig('SpeakerDiar16_ELBOVsTime.eps', bbox_inches='tight');

Helper.PlotUtil.plotStateSeq(Jdict['delmerge Sticky=100 K=25'], taskids='.best', sequences=['16']);