<commit_before>#!/usr/bin/env python

# The MIT License (MIT)
#
# Copyright (c) 2016 <NAME>, National Institutes of Health / NINDS
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# Top level for EM supervoxel classifier that learns merge / no merge classification based on some feature set.
# Featured Region Adjacency Graph (dpFRAG.py) creates the graph of neighboring supervoxels and their features.
#
# There is no iteration in this procedure, typically referred to as "flat learning" in the context of merging
#   supervoxels.
# After investigating flat learning, added a simple iterative procedure that agglomerates some percentage of the most
#   confident mergers and continues for set number of iterations.


import numpy as np
import time
import argparse
import os
#import sys
from io import BytesIO

from configobj import ConfigObj, flatten_errors
from validate import Validator, ValidateError
import dill
from io import StringIO
from sklearn.preprocessing import scale
from scipy import ndimage as nd

from sklearn.neighbors import KNeighborsClassifier
from sklearn.svm import SVC
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier, AdaBoostClassifier
from sklearn.naive_bayes import GaussianNB
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
from sklearn.discriminant_analysis import QuadraticDiscriminantAnalysis
from sklearn.linear_model import LogisticRegression

# xxx - how to make imports as optional? make plotting class?
from matplotlib import pylab as pl
from matplotlib import pyplot as plt
#import matplotlib as mpl
#from matplotlib import colors
#from cycler import cycler

from dpFRAG import dpFRAG
from metrics import pixel_error_fscore

class dpSupervoxelClassifier():

    # Constants
    LIST_ARGS = ['test_chunks', 'label_subgroups', 'label_subgroups_out', 'iterate_merge_perc',
        'thresholds', 'threshold_subgroups']
    n_jobs = 8
    #n_jobs = 1

    def __init__(self, args):

        # save command line arguments from argparse, see definitions in main or run with --help
        for k, v in vars(args).items():
            # do not override any values that are already set as a method of allowing inherited classes to specify
            if hasattr(self,k): continue
            if type(v) is list and k not in self.LIST_ARGS:
                if len(v)==1:
                    setattr(self,k,v[0])  # save single element lists as first element
                elif type(v[0]) is int:   # convert the sizes and offsets to numpy arrays
                    setattr(self,k,np.array(v,dtype=np.int32))
                else:
                    setattr(self,k,v)   # store other list types as usual (floats)
            else:
                setattr(self,k,v)

        # save the command line argument dict as a string
        out = StringIO(); print( vars(args), file=out )
        self.arg_str = out.getvalue(); out.close()

        print('dpSupervoxelClassifier: config file ''%s''' % self.cfgfile)
        # retrieve / save options from ini files, see definitions dpSupervoxelClassifier.ini
        opts = dpSupervoxelClassifier.get_options(self.cfgfile)
        d = vars(self)
        for k, v in opts.items():
            # do not import if have a "True" (non-empty) value from command line
            if k in d and d[k]: continue
            if type(v) is list and k not in self.LIST_ARGS:
                if len(v)==1:
                    setattr(self,k,v[0])  # save single element lists as first element
                elif len(v)>0 and type(v[0]) is int:   # convert the sizes and offsets to numpy arrays
                    setattr(self,k,np.array(v,dtype=np.int32))
                else:
                    setattr(self,k,v)   # store other list types as usual (floats, empties)
            else:
                setattr(self,k,v)

        # save the command line argument dict as a string
        out = StringIO(); print( opts, file=out )
        self.ini_str = out.getvalue(); out.close()

        # Options / Inits

        # these are so standard cubeIter inputs can be used from command line to override from .ini
        if (self.chunk >= 0).all():
            self.chunk_range_beg = self.chunk
            self.chunk_range_end = []
            self.offset_list = self.offset
            self.size_list = self.size

        self.doplots = (self.show_plots or self.export_plots)

        # default is to have sklearn calculate the priors
        self.priors = None
        if self.merge_prior > 0: self.priors = np.array([1-self.merge_prior, self.merge_prior],dtype=np.double)

        # keep the list of merge percentages, use original to store current iteration value
        self.iterate_merge_perc_list = self.iterate_merge_perc

        # initialize for "chunkrange" or "chunklist" mode if these parameters are not empty.
        # this code was modified from that in the em data parser for cuda-convnets2.
        self.chunk_range_beg = self.chunk_range_beg.reshape(-1,3); self.use_chunk_range = False
        self.nchunk_list = self.chunk_range_beg.shape[0]
        self.nchunks = self.nchunk_list

        if len(self.chunk_range_end) > 0:
            # "chunkrange" mode, chunks are selected based on defined beginning and end of ranges in X,Y,Z
            # range is open ended (python-style, end is not included in range)
            self.chunk_range_end = self.chunk_range_end.reshape(-1,3);
            assert( self.chunk_range_end.shape[0] == self.nchunk_list )
            self.chunk_range_rng = self.chunk_range_end - self.chunk_range_beg
            assert( (self.chunk_range_rng >= 0).all() )     # some bad ranges
            self.chunk_range_size = self.chunk_range_rng.prod(axis=1)
            self.chunk_range_cumsize = np.concatenate((np.zeros((1,),dtype=self.chunk_range_size.dtype),
                self.chunk_range_size.cumsum()))
            self.chunk_range_nchunks = self.chunk_range_cumsize[-1]
            self.use_chunk_range = True; self.nchunks = self.chunk_range_nchunks

        # offsets are either per chunk or per range, depending on above mode (whether chunk_range_end empty or not)
        self.offset_list = self.offset_list.reshape(-1,3)   # list must have multiple of 3 elements for volumes
        if self.offset_list.shape[0] > 1:
            assert( self.offset_list.shape[0] == self.nchunk_list )
        else:
            self.offset_list = np.ones_like(self.chunk_range_beg)*self.offset_list

        # sizes are either per chunk or per range, depending on above mode (whether chunk_range_end empty or not).
        self.size_list = self.size_list.reshape(-1,3)   # list must have multiple of 3 elements for volumes
        if self.size_list.shape[0] > 1:
            assert( self.size_list.shape[0] == self.nchunk_list )
        else:
            self.size_list = np.ones_like(self.chunk_range_beg)*self.size_list

        # print out all initialized variables in verbose mode
        if self.dpSupervoxelClassifier_verbose:
            # print out info for chunklist / chunkrange modes so that input data is logged
            print(('Using %d ' % self.nchunk_list) + ('ranges' if self.use_chunk_range else 'chunks') + ':')
            fh = BytesIO()
            if self.use_chunk_range:
                np.savetxt(fh, np.concatenate((np.arange(self.nchunk_list).reshape((self.nchunk_list,1)),
                    self.chunk_range_beg, self.chunk_range_end, self.chunk_range_size.reshape((self.nchunk_list,1)),
                    self.size_list, self.offset_list), axis=1),
                    fmt='\t(%d) range %d %d %d to %d %d %d (%d chunks), size %d %d %d, offset %d %d %d',
                    delimiter='', newline='\n', header='', footer='', comments='')
            else:
                np.savetxt(fh, np.concatenate((np.arange(self.nchunk_list).reshape((self.nchunk_list,1)),
                    self.chunk_range_beg, self.size_list, self.offset_list), axis=1),
                    fmt='\t(%d) chunk %d %d %d, size %d %d %d, offset %d %d %d',
                    delimiter='', newline='\n', header='', footer='', comments='')
            cstr = fh.getvalue(); fh.close(); print(cstr.decode('UTF-8'))
            print('Test chunks: %s' % ' '.join([str(x) for x in self.test_chunks]))

            #print('dpSupervoxelClassifier, verbose mode:\n'); print(vars(self))

        # inits for iterative prior mode
        self.iterative_mode = (self.iterate_count > 0)
        self.iterative_frag = [None] * self.nchunks
        self.iterative_mode_count = 0

        if self.iterative_mode:
            # expand the list out to the number of iterations by repeating the last merge percentage
            tmp = np.zeros((self.iterate_count,), dtype=np.double)
            tmp[:len(self.iterate_merge_perc_list)] = self.iterate_merge_perc_list
            tmp[len(self.iterate_merge_perc_list):] = self.iterate_merge_perc_list[-1]
            self.iterate_merge_perc_list = tmp

        # xxx - clean this up at some point, basically hack for legacy metric comparison scripts
        if len(self.threshold_subgroups) == 0:
            if self.iterative_mode:
                self.threshold_subgroups = np.arange(1,self.iterate_count+1,dtype=np.double)
            else:
                self.threshold_subgroups = np.sort(self.thresholds)[::-1]
        else:
            self.threshold_subgroups = np.sort(self.threshold_subgroups)[::-1]
            if self.iterative_mode:
                assert( self.iterate_count == len(self.threshold_subgroups) )
            else:
                assert( len(self.thresholds) == len(self.threshold_subgroups) )

        # dpFRAG features are now selectable as sets
        #self.FEATURES, self.FEATURES_NAMES, self.nfeatures = dpFRAG.make_FEATURES()
        d = dpFRAG.make_features(self.feature_set, self.has_ECS)
        for k in ['features','features_names','nfeatures']: setattr(self,k,d[k])

    def train(self):

        if self.dpSupervoxelClassifier_verbose: print('\nTRAIN')

        if self.trainin and (not self.classifierin or self.doplots):
            if self.dpSupervoxelClassifier_verbose:
                print('Loading training data')
            with open(self.trainin, 'rb') as f: data = dill.load(f)
            target = data['target']; fdata = data['data']
            ntargets = target.size; nfeatures = fdata.shape[1]
            assert( nfeatures == self.nfeatures )

        elif not self.classifierin:
            #dict_keys(['feature_names', 'DESCR', 'target_names', 'target', 'data'])
            nalloc = self.nchunks*self.nalloc_per_chunk
            nfeatures = self.nfeatures
            target = np.zeros((nalloc,), dtype=np.int64)
            fdata = np.zeros((nalloc,nfeatures), dtype=np.double)

            # accumulate training data from all training chunks
            cnt_targets = 0; ntargets = np.zeros((self.nchunks,),dtype=np.int64)
            for chunk in range(self.nchunks):
                cchunk, chunk_list_index, chunk_range_index = self.get_chunk_inds(chunk)
                offset = self.offset_list[chunk_list_index,:]; size = self.size_list[chunk_list_index,:]

                if chunk_list_index in self.test_chunks: continue
                print('Appending training data for chunk %d,%d,%d' % tuple(cchunk.tolist()))

                if self.iterative_mode:
                    if self.iterative_frag[chunk] is None:
                        frag = dpFRAG.makeBothFRAG(self.labelfile, cchunk, size, offset,
                            [self.probfile, self.probaugfile], [self.rawfile, self.rawaugfile],
                            self.raw_dataset, self.gtfile, self.outfile, self.label_subgroups, ['training','thr'],
                            progressBar=self.progress_bar, feature_set=self.feature_set, has_ECS=self.has_ECS,
                            verbose=self.dpSupervoxelClassifier_verbose)
                        frag.isTraining = True; self.iterative_frag[chunk] = frag
                    else:
                        frag = self.iterative_frag[chunk]
                else:
                    frag = dpFRAG.makeTrainingFRAG(self.labelfile, cchunk, size, offset,
                        [self.probfile, self.probaugfile], [self.rawfile, self.rawaugfile],
                        self.raw_dataset, self.gtfile, self.label_subgroups, feature_set=self.feature_set,
                        has_ECS=self.has_ECS,
                        progressBar=self.progress_bar, verbose=self.dpSupervoxelClassifier_verbose)
                frag.createFRAG(update = self.iterative_mode)
                #frag.createFRAG(update = False)
                data = frag.createDataset()
                ntargets[chunk] = data['target'].shape[0]
                target[cnt_targets:cnt_targets+ntargets[chunk]] = data['target']
                fdata[cnt_targets:cnt_targets+ntargets[chunk],:] = data['data']
                cnt_targets += ntargets[chunk]
            target = target[:cnt_targets]; fdata = fdata[:cnt_targets,:]

            if self.trainout:
                #dict_keys(['feature_names', 'DESCR', 'target_names', 'target', 'data'])
                descr = 'Training data from dpFRAG.py with command line:\n' + self.arg_str
                descr = ('With ini file "%s":\n' % (self.cfgfile,)) + self.ini_str
                data['data'] = fdata; data['target'] = target; data['DESCR'] = descr
                with open(self.trainout, 'wb') as f: dill.dump(data, f)

        if not self.classifierin or self.doplots:
            # everyone wants to be norml
            sdata = scale(fdata)   # normalize for the classifiers

        if self.classifierin:
            if self.dpSupervoxelClassifier_verbose:
                print('\nLoading classifier:'); t = time.time()

            with open(self.classifierin, 'rb') as f: d = dill.load(f)
            self.clf = d['classifier'];
        else:
            if self.dpSupervoxelClassifier_verbose:
                print('\nTraining classifier %s with %d examples and %d features:' % (self.classifier,
                    cnt_targets, nfeatures)); t = time.time()

            # train a classifier
            if self.classifier == 'lda':
                self.clf = LinearDiscriminantAnalysis(solver='svd', store_covariance=False, priors=self.priors)
                #self.clf = LinearDiscriminantAnalysis(solver='eigen', store_covariance=True, priors=self.priors)
            elif self.classifier == 'qda':
                self.clf = QuadraticDiscriminantAnalysis(priors=self.priors)
            elif self.classifier == 'rf':
                # the gala parameters
                #self.clf = RandomForestClassifier(n_estimators=100, criterion='entropy', max_depth=20,
                #    bootstrap=False, random_state=None)
                #self.clf = RandomForestClassifier(n_estimators=5*nfeatures,n_jobs=self.n_jobs,max_depth=10)
                self.clf = RandomForestClassifier(n_estimators=256,n_jobs=self.n_jobs,max_depth=16)
            elif self.classifier == 'svm':
                self.clf = SVC(kernel='rbf',probability=True,cache_size=2000)
            elif self.classifier == 'nb':
                self.clf = GaussianNB()
            elif self.classifier == 'kn':
                self.clf = KNeighborsClassifier(n_neighbors=10,n_jobs=self.n_jobs)
            elif self.classifier == 'dc':
                self.clf = DecisionTreeClassifier(max_depth=10)
            elif self.classifier == 'ada':
                self.clf = AdaBoostClassifier()
            elif self.classifier == 'lr':
                self.clf = LogisticRegression(penalty='l2',dual=False,solver='sag',n_jobs=self.n_jobs)
            else:
                assert(False)   # i never try anything, i just do it

            # train to the normalized data and merge or no merge targets
            self.clf.fit(sdata, target)

            if self.classifierout:
                with open(self.classifierout, 'wb') as f: dill.dump({'classifier':self.clf}, f)

        if self.dpSupervoxelClassifier_verbose:
            print('\tdone in %.4f s' % (time.time() - t))

        # do the agglomeration to use as input to next iteration for iterative mode
        thr = -1
        if self.iterative_mode:
            cnt_targets = 0
            for chunk in range(self.nchunks):
                if self.iterative_frag[chunk] is None or not self.iterative_frag[chunk].isTraining: continue

                # get the feature data for current training chunk only
                cdata = sdata[cnt_targets:cnt_targets+ntargets[chunk],:]
                cnt_targets += ntargets[chunk]

                # merge based on current classifier
                self.iterative_frag[chunk].subgroups_out[-1] = '%.8f' \
                    % self.threshold_subgroups[self.iterative_mode_count]
                clf_predict, thr = self.get_merge_predict_thr(cdata)
                self.iterative_frag[chunk].agglomerate(clf_predict)

                # make the next training iteration load from the current agglomerated supervoxels
                self.iterative_frag[chunk].srcfile = self.iterative_frag[chunk].outfile
                self.iterative_frag[chunk].subgroups = self.iterative_frag[chunk].subgroups_out

        if self.doplots:
            return self.createPlots(target,sdata,self.clf,self.export_plots,
                name=self.classifier + '_train_' + '_'.join([str(x) for x in self.test_chunks]) + \
                '_iter_' + str(self.iterative_mode_count), thr=thr, plot_features=self.plot_features)

    def test(self):
        if self.dpSupervoxelClassifier_verbose: print('\nTEST')

        for chunk in range(self.nchunks):
            cchunk, chunk_list_index, chunk_range_index = self.get_chunk_inds(chunk)

            if chunk_list_index not in self.test_chunks: continue
            data,sdata,thr = self._test(chunk, cchunk, chunk_list_index, chunk_range_index)

        # xxx - this only plots the LAST test chunk (currently this is only used for leave-one-out cross-validations)
        if self.doplots:
            assert(len(self.test_chunks) == 1)  # exporting plots/test data only works for one test chunk
            return self.createPlots(data['target'],sdata,self.clf,self.export_plots,
                name=self.classifier + '_test_' + '_'.join([str(x) for x in self.test_chunks]) + \
                '_iter_' + str(self.iterative_mode_count), thr=thr, plot_features=self.plot_features)

    def _test(self, ichunk, cchunk, chunk_list_index, chunk_range_index):
        print('Exporting testing data for chunk %d,%d,%d' % tuple(cchunk.tolist()))
        offset = self.offset_list[chunk_list_index,:]; size = self.size_list[chunk_list_index,:]

        FRAG = None
        if len(self.test_chunks) == 1 and self.testin:
            if self.dpSupervoxelClassifier_verbose:
                print('Loading testing data')
            with open(self.testin, 'rb') as f: data = dill.load(f)
            FRAG = data['FRAG']; data = data['data']

        frag = None; subgroups_out= list(self.label_subgroups_out)
        if self.iterative_mode:
            if self.iterative_frag[ichunk] is None: subgroups_out += ['thr']
            else: frag = self.iterative_frag[ichunk]

        if frag is None:
            if self.doplots:
                frag = dpFRAG.makeBothFRAG(self.labelfile, cchunk, size, offset,
                    [self.probfile, self.probaugfile], [self.rawfile, self.rawaugfile],
                    self.raw_dataset, self.gtfile, self.outfile, self.label_subgroups, subgroups_out,
                    G=FRAG, progressBar=self.progress_bar, feature_set=self.feature_set, has_ECS=self.has_ECS,
                    verbose=self.dpSupervoxelClassifier_verbose)
            else:
                frag = dpFRAG.makeTestingFRAG(self.labelfile, cchunk, size, offset,
                    [self.probfile, self.probaugfile], [self.rawfile, self.rawaugfile],
                    self.raw_dataset, self.outfile, self.label_subgroups, subgroups_out, G=FRAG,
                    progressBar=self.progress_bar, feature_set=self.feature_set, has_ECS=self.has_ECS,
                    verbose=self.dpSupervoxelClassifier_verbose)

        if self.iterative_mode and self.iterative_frag[ichunk] is None:
            frag.isTraining = False; self.iterative_frag[ichunk] = frag

        if not (len(self.test_chunks) == 1 and self.testin):
            frag.createFRAG(update = self.iterative_mode)
            #frag.createFRAG(update = False)
            data = frag.createDataset(train=self.doplots)

            if self.testout:
                if self.dpSupervoxelClassifier_verbose:
                    print('Dumping testing data')
                descr = 'Testing data from dpFRAG.py with command line:\n' + self.arg_str
                descr = ('With ini file "%s":\n' % (self.cfgfile,)) + self.ini_str
                data['DESCR'] = descr
                with open(self.testout, 'wb') as f: dill.dump({'data':data,'FRAG':frag.FRAG}, f)

        sdata = scale(data['data'])     # normalize for the classifiers

        thr = -1
        if self.iterative_mode:
            # merge based on current classifier
            frag.subgroups_out[-1] = '%.8f' % self.threshold_subgroups[self.iterative_mode_count]
            clf_predict, thr = self.get_merge_predict_thr(sdata)
            frag.agglomerate(clf_predict)

            # make the next training iteration load from the current agglomerated supervoxels
            self.iterative_frag[ichunk].srcfile = self.iterative_frag[ichunk].outfile
            self.iterative_frag[ichunk].subgroups = self.iterative_frag[ichunk].subgroups_out
        else:
            try:
                # predict merge or not on testing cube and write outputs at specified probability thresholds
                frag.threshold_agglomerate(self.clf.predict_proba(sdata), self.thresholds, self.threshold_subgroups)
                # there's an issue here if there are no mergers left, would be better to just copy the current
                #   agglomeration, xxx - deal with this later. handled this explicitly in other locations.
                #except AttributeError:
            except:
                # if the classifier doesn't do probabilities just export single prediction
                frag.subgroups_out += ['single_' + self.classifier]
                frag.agglomerate(self.clf.predict(sdata))

        return data,sdata,thr

    def iterative_classify(self):
        assert(self.iterative_mode)

        if self.classifierin:
            if self.dpSupervoxelClassifier_verbose:
                print('\nLoading iterative classifiers:'); t = time.time()

            with open(self.classifierin, 'rb') as f: d = dill.load(f)
            self.clfs = d['classifiers'];
            print('\tdone in %.4f s' % (time.time() - t))
        else:
            self.clfs = [None] * self.iterate_count

        # save the ROC-style metrics for later analysis / plotting
        train_metrics = [None] * self.iterate_count; test_metrics = [None] * self.iterate_count

        # each iteration loop trains on the current supervoxels and
        #   then performs a merge (test) based with a normal sklearn predict based on a small merge prior.
        self.classifierin = ''; self.trainin = ''; self.trainout = ''; self.testin = ''; self.testout = ''
        if self.test_only:
            print('Test-only iterative mode with merge prior %.4f' % (self.merge_prior,))
            for chunk in range(self.nchunks):
                cchunk, chunk_list_index, chunk_range_index = self.get_chunk_inds(chunk)
                #if chunk_list_index not in self.test_chunks: continue

                for i in range(self.iterate_count):
                    self.iterate_merge_perc = self.iterate_merge_perc_list[i]; self.iterative_mode_count = i
                    print('Iteration %d, merge perc %.4f' % (self.iterative_mode_count+1,self.iterate_merge_perc))

                    assert( self.clfs[i] is not None) # classifierin must be specified for test only
                    self.clf = self.clfs[i]
                    self._test(chunk, cchunk, chunk_list_index, chunk_range_index)

                # let the previous chunk get garbage collected to save memory.
                # that is essentially the purpose of this special mode (iteration loop as inner loop).
                self.iterative_frag[chunk] = None
        else:
            print('Normal iterative mode with merge prior %.4f' % (self.merge_prior,))
            for i in range(self.iterate_count):
                self.iterate_merge_perc = self.iterate_merge_perc_list[i]; self.iterative_mode_count = i
                print('Iteration %d, merge perc %.4f' % (self.iterative_mode_count+1,self.iterate_merge_perc))

                if self.clfs[i] is None:
                    train_metrics[i] = self.train(); self.clfs[i] = self.clf
                else:
                    self.clf = self.clfs[i]
                test_metrics[i] = self.test()

                if self.classifierout:
                    with open(self.classifierout, 'wb') as f:
                        dill.dump({'classifiers':self.clfs,'train_metrics':train_metrics,'test_metrics':test_metrics},f)

    def get_merge_predict_thr(self,data):
        thr = -1; #ntargets = data.shape[0]
        # without merge percentage specified just use the classifier predict().
        clf_predict = self.clf.predict(data)
        # use the merge percentage to merge only this percentage of the most confident mergers.
        # do not use any potential mergers that are against the classifier predict().
        if self.iterate_merge_perc > 0:
            clf_probs = self.clf.predict_proba(data)
            # xxx - figure out how to handle no mergers left more gracefully, just stop at this iteration?
            assert(clf_probs.ndim > 1 and clf_probs.shape[1] > 1) # ran out of mergers
            # use merge percentage as percentage of estimated yes merges remaining
            nmerge = (clf_predict==1).sum(dtype=np.int64); clf_probs = clf_probs[:,1]
            if nmerge > 0:
                thr = np.sort(clf_probs)[::-1][int(nmerge*self.iterate_merge_perc)]
                clf_predict = np.logical_and(clf_predict, clf_probs >= thr)
        return clf_predict, thr

    def get_chunk_inds(self, chunk):
        if self.use_chunk_range:
            chunk_list_index = np.nonzero(chunk >= self.chunk_range_cumsize)[0][-1]
            chunk_range_index = chunk - self.chunk_range_cumsize[chunk_list_index]
            cchunk = np.unravel_index(chunk_range_index, self.chunk_range_rng[chunk_list_index,:]) \
                + self.chunk_range_beg[chunk_list_index,:]
        else:
            chunk_list_index = chunk; cchunk = self.chunk_range_beg[chunk,:]; chunk_range_index = None
        return cchunk, chunk_list_index, chunk_range_index

    # this method assumes binary classification
    # xxx - plotting code got really messy as per usual, clean this up<commit_msg>allow chunk range sizes bigger than knossos cube size<commit_after>#!/usr/bin/env python

# The MIT License (MIT)
#
# Copyright (c) 2016 <NAME>, National Institutes of Health / NINDS
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# Top level for EM supervoxel classifier that learns merge / no merge classification based on some feature set.
# Featured Region Adjacency Graph (dpFRAG.py) creates the graph of neighboring supervoxels and their features.
#
# There is no iteration in this procedure, typically referred to as "flat learning" in the context of merging
#   supervoxels.
# After investigating flat learning, added a simple iterative procedure that agglomerates some percentage of the most
#   confident mergers and continues for set number of iterations.


import numpy as np
import time
import argparse
import os
#import sys
from io import BytesIO

from configobj import ConfigObj, flatten_errors
from validate import Validator, ValidateError
import dill
from io import StringIO
from sklearn.preprocessing import scale
from scipy import ndimage as nd

from sklearn.neighbors import KNeighborsClassifier
from sklearn.svm import SVC
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier, AdaBoostClassifier
from sklearn.naive_bayes import GaussianNB
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
from sklearn.discriminant_analysis import QuadraticDiscriminantAnalysis
from sklearn.linear_model import LogisticRegression

# xxx - how to make imports as optional? make plotting class?
from matplotlib import pylab as pl
from matplotlib import pyplot as plt
#import matplotlib as mpl
#from matplotlib import colors
#from cycler import cycler

from dpLoadh5 import dpLoadh5
from dpFRAG import dpFRAG
from metrics import pixel_error_fscore

class dpSupervoxelClassifier():

    # Constants
    LIST_ARGS = ['test_chunks', 'label_subgroups', 'label_subgroups_out', 'iterate_merge_perc',
        'thresholds', 'threshold_subgroups']
    n_jobs = 8
    #n_jobs = 1

    def __init__(self, args):

        # save command line arguments from argparse, see definitions in main or run with --help
        for k, v in vars(args).items():
            # do not override any values that are already set as a method of allowing inherited classes to specify
            if hasattr(self,k): continue
            if type(v) is list and k not in self.LIST_ARGS:
                if len(v)==1:
                    setattr(self,k,v[0])  # save single element lists as first element
                elif type(v[0]) is int:   # convert the sizes and offsets to numpy arrays
                    setattr(self,k,np.array(v,dtype=np.int32))
                else:
                    setattr(self,k,v)   # store other list types as usual (floats)
            else:
                setattr(self,k,v)

        # save the command line argument dict as a string
        out = StringIO(); print( vars(args), file=out )
        self.arg_str = out.getvalue(); out.close()

        print('dpSupervoxelClassifier: config file ''%s''' % self.cfgfile)
        # retrieve / save options from ini files, see definitions dpSupervoxelClassifier.ini
        opts = dpSupervoxelClassifier.get_options(self.cfgfile)
        d = vars(self)
        for k, v in opts.items():
            # do not import if have a "True" (non-empty) value from command line
            if k in d and d[k]: continue
            if type(v) is list and k not in self.LIST_ARGS:
                if len(v)==1:
                    setattr(self,k,v[0])  # save single element lists as first element
                elif len(v)>0 and type(v[0]) is int:   # convert the sizes and offsets to numpy arrays
                    setattr(self,k,np.array(v,dtype=np.int32))
                else:
                    setattr(self,k,v)   # store other list types as usual (floats, empties)
            else:
                setattr(self,k,v)

        # save the command line argument dict as a string
        out = StringIO(); print( opts, file=out )
        self.ini_str = out.getvalue(); out.close()

        # Options / Inits

        # these are so standard cubeIter inputs can be used from command line to override from .ini
        if (self.chunk >= 0).all():
            self.chunk_range_beg = self.chunk
            self.chunk_range_end = []
            self.offset_list = self.offset
            self.size_list = self.size

        self.doplots = (self.show_plots or self.export_plots)

        # default is to have sklearn calculate the priors
        self.priors = None
        if self.merge_prior > 0: self.priors = np.array([1-self.merge_prior, self.merge_prior],dtype=np.double)

        # keep the list of merge percentages, use original to store current iteration value
        self.iterate_merge_perc_list = self.iterate_merge_perc

        # initialize for "chunkrange" or "chunklist" mode if these parameters are not empty.
        # this code was modified from that in the em data parser for cuda-convnets2.
        self.chunk_range_beg = self.chunk_range_beg.reshape(-1,3); self.use_chunk_range = False
        self.nchunk_list = self.chunk_range_beg.shape[0]
        self.nchunks = self.nchunk_list

        # offsets are either per chunk or per range, depending on above mode (whether chunk_range_end empty or not)
        self.offset_list = self.offset_list.reshape(-1,3)   # list must have multiple of 3 elements for volumes
        if self.offset_list.shape[0] > 1:
            assert( self.offset_list.shape[0] == self.nchunk_list )
        else:
            self.offset_list = np.ones_like(self.chunk_range_beg)*self.offset_list

        # sizes are either per chunk or per range, depending on above mode (whether chunk_range_end empty or not).
        self.size_list = self.size_list.reshape(-1,3)   # list must have multiple of 3 elements for volumes
        if self.size_list.shape[0] > 1:
            assert( self.size_list.shape[0] == self.nchunk_list )
        else:
            self.size_list = np.ones_like(self.chunk_range_beg)*self.size_list

        if len(self.chunk_range_end) > 0:
            # "chunkrange" mode, chunks are selected based on defined beginning and end of ranges in X,Y,Z
            # range is open ended (python-style, end is not included in range)
            self.chunk_range_end = self.chunk_range_end.reshape(-1,3);
            assert( self.chunk_range_end.shape[0] == self.nchunk_list )
            self.chunk_range_rng = self.chunk_range_end - self.chunk_range_beg
            assert( (self.chunk_range_rng >= 0).all() )     # some bad ranges

            # this allows for iterating in chunk_range mode with sizes greater than chunksize
            # just read the h5info for the rawfile to get the chunksize
            rawh5 = dpLoadh5.readInith5(self.rawfile, self.raw_dataset, [0,0,0], [0,0,0], [128,128,128], 'uint8')
            self.chunksize = rawh5.chunksize
            self.chunkscale = np.ones_like(self.size_list)
            # only allow this if the sizes are a multiple of chunksize
            if (self.size_list % self.chunksize == 0).all():
                scale = self.size_list // self.chunksize
                # only allow this if the ranges are a multiple of the sizes
                if (self.chunk_range_rng % scale == 0).all():
                    self.chunkscale = scale
            self.chunk_range_rng //= self.chunkscale

            self.chunk_range_size = self.chunk_range_rng.prod(axis=1)
            self.chunk_range_cumsize = np.concatenate((np.zeros((1,),dtype=self.chunk_range_size.dtype),
                self.chunk_range_size.cumsum()))
            self.chunk_range_nchunks = self.chunk_range_cumsize[-1]
            self.use_chunk_range = True; self.nchunks = self.chunk_range_nchunks

        # print out all initialized variables in verbose mode
        if self.dpSupervoxelClassifier_verbose:
            # print out info for chunklist / chunkrange modes so that input data is logged
            print(('Using %d ' % self.nchunk_list) + ('ranges' if self.use_chunk_range else 'chunks') + ':')
            fh = BytesIO()
            if self.use_chunk_range:
                np.savetxt(fh, np.concatenate((np.arange(self.nchunk_list).reshape((self.nchunk_list,1)),
                    self.chunk_range_beg, self.chunk_range_end, self.chunk_range_size.reshape((self.nchunk_list,1)),
                    self.size_list, self.offset_list), axis=1),
                    fmt='\t(%d) range %d %d %d to %d %d %d (%d chunks), size %d %d %d, offset %d %d %d',
                    delimiter='', newline='\n', header='', footer='', comments='')
            else:
                np.savetxt(fh, np.concatenate((np.arange(self.nchunk_list).reshape((self.nchunk_list,1)),
                    self.chunk_range_beg, self.size_list, self.offset_list), axis=1),
                    fmt='\t(%d) chunk %d %d %d, size %d %d %d, offset %d %d %d',
                    delimiter='', newline='\n', header='', footer='', comments='')
            cstr = fh.getvalue(); fh.close(); print(cstr.decode('UTF-8'))
            print('Test chunks: %s' % ' '.join([str(x) for x in self.test_chunks]))

            #print('dpSupervoxelClassifier, verbose mode:\n'); print(vars(self))

        # inits for iterative prior mode
        self.iterative_mode = (self.iterate_count > 0)
        self.iterative_frag = [None] * self.nchunks
        self.iterative_mode_count = 0

        if self.iterative_mode:
            # expand the list out to the number of iterations by repeating the last merge percentage
            tmp = np.zeros((self.iterate_count,), dtype=np.double)
            tmp[:len(self.iterate_merge_perc_list)] = self.iterate_merge_perc_list
            tmp[len(self.iterate_merge_perc_list):] = self.iterate_merge_perc_list[-1]
            self.iterate_merge_perc_list = tmp

        # xxx - clean this up at some point, basically hack for legacy metric comparison scripts
        if len(self.threshold_subgroups) == 0:
            if self.iterative_mode:
                self.threshold_subgroups = np.arange(1,self.iterate_count+1,dtype=np.double)
            else:
                self.threshold_subgroups = np.sort(self.thresholds)[::-1]
        else:
            self.threshold_subgroups = np.sort(self.threshold_subgroups)[::-1]
            if self.iterative_mode:
                assert( self.iterate_count == len(self.threshold_subgroups) )
            else:
                assert( len(self.thresholds) == len(self.threshold_subgroups) )

        # dpFRAG features are now selectable as sets
        #self.FEATURES, self.FEATURES_NAMES, self.nfeatures = dpFRAG.make_FEATURES()
        d = dpFRAG.make_features(self.feature_set, self.has_ECS)
        for k in ['features','features_names','nfeatures']: setattr(self,k,d[k])

    def train(self):

        if self.dpSupervoxelClassifier_verbose: print('\nTRAIN')

        if self.trainin and (not self.classifierin or self.doplots):
            if self.dpSupervoxelClassifier_verbose:
                print('Loading training data')
            with open(self.trainin, 'rb') as f: data = dill.load(f)
            target = data['target']; fdata = data['data']
            ntargets = target.size; nfeatures = fdata.shape[1]
            assert( nfeatures == self.nfeatures )

        elif not self.classifierin:
            #dict_keys(['feature_names', 'DESCR', 'target_names', 'target', 'data'])
            nalloc = self.nchunks*self.nalloc_per_chunk
            nfeatures = self.nfeatures
            target = np.zeros((nalloc,), dtype=np.int64)
            fdata = np.zeros((nalloc,nfeatures), dtype=np.double)

            # accumulate training data from all training chunks
            cnt_targets = 0; ntargets = np.zeros((self.nchunks,),dtype=np.int64)
            for chunk in range(self.nchunks):
                cchunk, chunk_list_index, chunk_range_index = self.get_chunk_inds(chunk)
                offset = self.offset_list[chunk_list_index,:]; size = self.size_list[chunk_list_index,:]

                if chunk_list_index in self.test_chunks: continue
                print('Appending training data for chunk %d,%d,%d' % tuple(cchunk.tolist()))

                if self.iterative_mode:
                    if self.iterative_frag[chunk] is None:
                        frag = dpFRAG.makeBothFRAG(self.labelfile, cchunk, size, offset,
                            [self.probfile, self.probaugfile], [self.rawfile, self.rawaugfile],
                            self.raw_dataset, self.gtfile, self.outfile, self.label_subgroups, ['training','thr'],
                            progressBar=self.progress_bar, feature_set=self.feature_set, has_ECS=self.has_ECS,
                            verbose=self.dpSupervoxelClassifier_verbose)
                        frag.isTraining = True; self.iterative_frag[chunk] = frag
                    else:
                        frag = self.iterative_frag[chunk]
                else:
                    frag = dpFRAG.makeTrainingFRAG(self.labelfile, cchunk, size, offset,
                        [self.probfile, self.probaugfile], [self.rawfile, self.rawaugfile],
                        self.raw_dataset, self.gtfile, self.label_subgroups, feature_set=self.feature_set,
                        has_ECS=self.has_ECS,
                        progressBar=self.progress_bar, verbose=self.dpSupervoxelClassifier_verbose)
                frag.createFRAG(update = self.iterative_mode)
                #frag.createFRAG(update = False)
                data = frag.createDataset()
                ntargets[chunk] = data['target'].shape[0]
                target[cnt_targets:cnt_targets+ntargets[chunk]] = data['target']
                fdata[cnt_targets:cnt_targets+ntargets[chunk],:] = data['data']
                cnt_targets += ntargets[chunk]
            target = target[:cnt_targets]; fdata = fdata[:cnt_targets,:]

            if self.trainout:
                #dict_keys(['feature_names', 'DESCR', 'target_names', 'target', 'data'])
                descr = 'Training data from dpFRAG.py with command line:\n' + self.arg_str
                descr = ('With ini file "%s":\n' % (self.cfgfile,)) + self.ini_str
                data['data'] = fdata; data['target'] = target; data['DESCR'] = descr
                with open(self.trainout, 'wb') as f: dill.dump(data, f)

        if not self.classifierin or self.doplots:
            # everyone wants to be norml
            sdata = scale(fdata)   # normalize for the classifiers

        if self.classifierin:
            if self.dpSupervoxelClassifier_verbose:
                print('\nLoading classifier:'); t = time.time()

            with open(self.classifierin, 'rb') as f: d = dill.load(f)
            self.clf = d['classifier'];
        else:
            if self.dpSupervoxelClassifier_verbose:
                print('\nTraining classifier %s with %d examples and %d features:' % (self.classifier,
                    cnt_targets, nfeatures)); t = time.time()

            # train a classifier
            if self.classifier == 'lda':
                self.clf = LinearDiscriminantAnalysis(solver='svd', store_covariance=False, priors=self.priors)
                #self.clf = LinearDiscriminantAnalysis(solver='eigen', store_covariance=True, priors=self.priors)
            elif self.classifier == 'qda':
                self.clf = QuadraticDiscriminantAnalysis(priors=self.priors)
            elif self.classifier == 'rf':
                # the gala parameters
                #self.clf = RandomForestClassifier(n_estimators=100, criterion='entropy', max_depth=20,
                #    bootstrap=False, random_state=None)
                #self.clf = RandomForestClassifier(n_estimators=5*nfeatures,n_jobs=self.n_jobs,max_depth=10)
                self.clf = RandomForestClassifier(n_estimators=256,n_jobs=self.n_jobs,max_depth=16)
            elif self.classifier == 'svm':
                self.clf = SVC(kernel='rbf',probability=True,cache_size=2000)
            elif self.classifier == 'nb':
                self.clf = GaussianNB()
            elif self.classifier == 'kn':
                self.clf = KNeighborsClassifier(n_neighbors=10,n_jobs=self.n_jobs)
            elif self.classifier == 'dc':
                self.clf = DecisionTreeClassifier(max_depth=10)
            elif self.classifier == 'ada':
                self.clf = AdaBoostClassifier()
            elif self.classifier == 'lr':
                self.clf = LogisticRegression(penalty='l2',dual=False,solver='sag',n_jobs=self.n_jobs)
            else:
                assert(False)   # i never try anything, i just do it

            # train to the normalized data and merge or no merge targets
            self.clf.fit(sdata, target)

            if self.classifierout:
                with open(self.classifierout, 'wb') as f: dill.dump({'classifier':self.clf}, f)

        if self.dpSupervoxelClassifier_verbose:
            print('\tdone in %.4f s' % (time.time() - t))

        # do the agglomeration to use as input to next iteration for iterative mode
        thr = -1
        if self.iterative_mode:
            cnt_targets = 0
            for chunk in range(self.nchunks):
                if self.iterative_frag[chunk] is None or not self.iterative_frag[chunk].isTraining: continue

                # get the feature data for current training chunk only
                cdata = sdata[cnt_targets:cnt_targets+ntargets[chunk],:]
                cnt_targets += ntargets[chunk]

                # merge based on current classifier
                self.iterative_frag[chunk].subgroups_out[-1] = '%.8f' \
                    % self.threshold_subgroups[self.iterative_mode_count]
                clf_predict, thr = self.get_merge_predict_thr(cdata)
                self.iterative_frag[chunk].agglomerate(clf_predict)

                # make the next training iteration load from the current agglomerated supervoxels
                self.iterative_frag[chunk].srcfile = self.iterative_frag[chunk].outfile
                self.iterative_frag[chunk].subgroups = self.iterative_frag[chunk].subgroups_out

        if self.doplots:
            return self.createPlots(target,sdata,self.clf,self.export_plots,
                name=self.classifier + '_train_' + '_'.join([str(x) for x in self.test_chunks]) + \
                '_iter_' + str(self.iterative_mode_count), thr=thr, plot_features=self.plot_features)

    def test(self):
        if self.dpSupervoxelClassifier_verbose: print('\nTEST')

        for chunk in range(self.nchunks):
            cchunk, chunk_list_index, chunk_range_index = self.get_chunk_inds(chunk)

            if chunk_list_index not in self.test_chunks: continue
            data,sdata,thr = self._test(chunk, cchunk, chunk_list_index, chunk_range_index)

        # xxx - this only plots the LAST test chunk (currently this is only used for leave-one-out cross-validations)
        if self.doplots:
            assert(len(self.test_chunks) == 1)  # exporting plots/test data only works for one test chunk
            return self.createPlots(data['target'],sdata,self.clf,self.export_plots,
                name=self.classifier + '_test_' + '_'.join([str(x) for x in self.test_chunks]) + \
                '_iter_' + str(self.iterative_mode_count), thr=thr, plot_features=self.plot_features)

    def _test(self, ichunk, cchunk, chunk_list_index, chunk_range_index):
        print('Exporting testing data for chunk %d,%d,%d' % tuple(cchunk.tolist()))
        offset = self.offset_list[chunk_list_index,:]; size = self.size_list[chunk_list_index,:]

        FRAG = None
        if len(self.test_chunks) == 1 and self.testin:
            if self.dpSupervoxelClassifier_verbose:
                print('Loading testing data')
            with open(self.testin, 'rb') as f: data = dill.load(f)
            FRAG = data['FRAG']; data = data['data']

        frag = None; subgroups_out= list(self.label_subgroups_out)
        if self.iterative_mode:
            if self.iterative_frag[ichunk] is None: subgroups_out += ['thr']
            else: frag = self.iterative_frag[ichunk]

        if frag is None:
            if self.doplots:
                frag = dpFRAG.makeBothFRAG(self.labelfile, cchunk, size, offset,
                    [self.probfile, self.probaugfile], [self.rawfile, self.rawaugfile],
                    self.raw_dataset, self.gtfile, self.outfile, self.label_subgroups, subgroups_out,
                    G=FRAG, progressBar=self.progress_bar, feature_set=self.feature_set, has_ECS=self.has_ECS,
                    verbose=self.dpSupervoxelClassifier_verbose)
            else:
                frag = dpFRAG.makeTestingFRAG(self.labelfile, cchunk, size, offset,
                    [self.probfile, self.probaugfile], [self.rawfile, self.rawaugfile],
                    self.raw_dataset, self.outfile, self.label_subgroups, subgroups_out, G=FRAG,
                    progressBar=self.progress_bar, feature_set=self.feature_set, has_ECS=self.has_ECS,
                    verbose=self.dpSupervoxelClassifier_verbose)

        if self.iterative_mode and self.iterative_frag[ichunk] is None:
            frag.isTraining = False; self.iterative_frag[ichunk] = frag

        if not (len(self.test_chunks) == 1 and self.testin):
            frag.createFRAG(update = self.iterative_mode)
            #frag.createFRAG(update = False)
            data = frag.createDataset(train=self.doplots)

            if self.testout:
                if self.dpSupervoxelClassifier_verbose:
                    print('Dumping testing data')
                descr = 'Testing data from dpFRAG.py with command line:\n' + self.arg_str
                descr = ('With ini file "%s":\n' % (self.cfgfile,)) + self.ini_str
                data['DESCR'] = descr
                with open(self.testout, 'wb') as f: dill.dump({'data':data,'FRAG':frag.FRAG}, f)

        sdata = scale(data['data'])     # normalize for the classifiers

        thr = -1
        if self.iterative_mode:
            # merge based on current classifier
            frag.subgroups_out[-1] = '%.8f' % self.threshold_subgroups[self.iterative_mode_count]
            clf_predict, thr = self.get_merge_predict_thr(sdata)
            frag.agglomerate(clf_predict)

            # make the next training iteration load from the current agglomerated supervoxels
            self.iterative_frag[ichunk].srcfile = self.iterative_frag[ichunk].outfile
            self.iterative_frag[ichunk].subgroups = self.iterative_frag[ichunk].subgroups_out
        else:
            try:
                # predict merge or not on testing cube and write outputs at specified probability thresholds
                frag.threshold_agglomerate(self.clf.predict_proba(sdata), self.thresholds, self.threshold_subgroups)
                # there's an issue here if there are no mergers left, would be better to just copy the current
                #   agglomeration, xxx - deal with this later. handled this explicitly in other locations.
                #except AttributeError:
            except:
                # if the classifier doesn't do probabilities just export single prediction
                frag.subgroups_out += ['single_' + self.classifier]
                frag.agglomerate(self.clf.predict(sdata))

        return data,sdata,thr

    def iterative_classify(self):
        assert(self.iterative_mode)

        if self.classifierin:
            if self.dpSupervoxelClassifier_verbose:
                print('\nLoading iterative classifiers:'); t = time.time()

            with open(self.classifierin, 'rb') as f: d = dill.load(f)
            self.clfs = d['classifiers'];
            print('\tdone in %.4f s' % (time.time() - t))
        else:
            self.clfs = [None] * self.iterate_count

        # save the ROC-style metrics for later analysis / plotting
        train_metrics = [None] * self.iterate_count; test_metrics = [None] * self.iterate_count

        # each iteration loop trains on the current supervoxels and
        #   then performs a merge (test) based with a normal sklearn predict based on a small merge prior.
        self.classifierin = ''; self.trainin = ''; self.trainout = ''; self.testin = ''; self.testout = ''
        if self.test_only:
            print('Test-only iterative mode with merge prior %.4f' % (self.merge_prior,))
            for chunk in range(self.nchunks):
                cchunk, chunk_list_index, chunk_range_index = self.get_chunk_inds(chunk)
                #if chunk_list_index not in self.test_chunks: continue

                for i in range(self.iterate_count):
                    self.iterate_merge_perc = self.iterate_merge_perc_list[i]; self.iterative_mode_count = i
                    print('Iteration %d, merge perc %.4f' % (self.iterative_mode_count+1,self.iterate_merge_perc))

                    assert( self.clfs[i] is not None) # classifierin must be specified for test only
                    self.clf = self.clfs[i]
                    self._test(chunk, cchunk, chunk_list_index, chunk_range_index)

                # let the previous chunk get garbage collected to save memory.
                # that is essentially the purpose of this special mode (iteration loop as inner loop).
                self.iterative_frag[chunk] = None
        else:
            print('Normal iterative mode with merge prior %.4f' % (self.merge_prior,))
            for i in range(self.iterate_count):
                self.iterate_merge_perc = self.iterate_merge_perc_list[i]; self.iterative_mode_count = i
                print('Iteration %d, merge perc %.4f' % (self.iterative_mode_count+1,self.iterate_merge_perc))

                if self.clfs[i] is None:
                    train_metrics[i] = self.train(); self.clfs[i] = self.clf
                else:
                    self.clf = self.clfs[i]
                test_metrics[i] = self.test()

                if self.classifierout:
                    with open(self.classifierout, 'wb') as f:
                        dill.dump({'classifiers':self.clfs,'train_metrics':train_metrics,'test_metrics':test_metrics},f)

    def get_merge_predict_thr(self,data):
        thr = -1; #ntargets = data.shape[0]
        # without merge percentage specified just use the classifier predict().
        clf_predict = self.clf.predict(data)
        # use the merge percentage to merge only this percentage of the most confident mergers.
        # do not use any potential mergers that are against the classifier predict().
        if self.iterate_merge_perc > 0:
            clf_probs = self.clf.predict_proba(data)
            # xxx - figure out how to handle no mergers left more gracefully, just stop at this iteration?
            assert(clf_probs.ndim > 1 and clf_probs.shape[1] > 1) # ran out of mergers
            # use merge percentage as percentage of estimated yes merges remaining
            nmerge = (clf_predict==1).sum(dtype=np.int64); clf_probs = clf_probs[:,1]
            if nmerge > 0:
                thr = np.sort(clf_probs)[::-1][int(nmerge*self.iterate_merge_perc)]
                clf_predict = np.logical_and(clf_predict, clf_probs >= thr)
        return clf_predict, thr

    def get_chunk_inds(self, chunk):
        if self.use_chunk_range:
            chunk_list_index = np.nonzero(chunk >= self.chunk_range_cumsize)[0][-1]
            chunk_range_index = chunk - self.chunk_range_cumsize[chunk_list_index]
            cchunk = np.unravel_index(chunk_range_index, self.chunk_range_rng[chunk_list_index,:]) \
                *self.chunkscale[chunk_list_index,:] + self.chunk_range_beg[chunk_list_index,:]
        else:
            chunk_list_index = chunk; cchunk = self.chunk_range_beg[chunk,:]; chunk_range_index = None
        return cchunk, chunk_list_index, chunk_range_index

    # this method assumes binary classification
    # xxx - plotting code got really messy as per usual, clean this up
{'fraction_non_alphanumeric': 0.09484713700532692, 'fraction_numerical': 0.004623688050491377, 'mean_word_length': 2.4212678936605316, 'pattern_counts': {'":': 4, '<': 5, '<?xml version=': 0, '>': 33, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6945305', 'n_tokens_mistral': 17302, 'n_tokens_neox': 16343, 'n_words': 4316}