<commit_before>import os

from progressbar import ProgressBar

from cnntools.descriptor_aggregator import DescriptorAggregator
from cnntools.descstore import DescriptorStoreHdf5
from cnntools.models import CaffeCNNSnapshot
from cnntools.redis_aggregator import patch_interrupt_signal, detach_patch_interrupt_signal
from cnntools.tasks import compute_cnn_features_gpu_task
from cnntools.common_utils import iter_batch, progress_bar_widgets


def get_snapshot_id(caffe_cnn, snapshot_id):
    if snapshot_id:
        snapshot = CaffeCNNSnapshot.objects.get(id=snapshot_id)
        if snapshot.training_run.net.id != caffe_cnn.id:
            raise ValueError(
                'Snapshot\'s corresponding network id doesn\'t match given'
                ' network id!'
            )
        return snapshot_id

    # Get last snapshot
    snapshot = CaffeCNNSnapshot.objects.\
        filter(training_run__net__id=caffe_cnn.id).\
        order_by('-id')[0]

    return snapshot.id


def get_slug(caffe_cnn, snapshot_id, slug_extra):
    snapshot_id = get_snapshot_id(caffe_cnn, snapshot_id)
    slug = 'netid-{}:snapshot-{}'.format(caffe_cnn.netid, snapshot_id)
    if slug_extra:
        slug += slug_extra

    return slug, snapshot_id


def get_redis_key_core(slug, feature_name_list):
    return 'desc:slug-%s:fetname%s:batchid#' % (
        slug,
        '-'.join(feature_name_list),
    )


def get_redis_key(slug, feature_name_list, batchid_counter):
    return '%s%s' % (
        get_redis_key_core(slug, feature_name_list),
        batchid_counter,
    )


def get_task_id(slug, feature_name_list):
    return '%stask' % (
        get_redis_key_core(slug, feature_name_list),
    )


def get_descstore_dirname(netid, feature_name):
    dirname = '{}-{}'.format(netid, feature_name)
    # clean up
    invalid_chars = ['/', '(', ')', '[', ']', '{', '}', ' ', ':']

    for c in invalid_chars:
        dirname = dirname.replace(c, '-')

    return dirname


def get_descstore_filename(netid, feature_name):
    dirname = get_descstore_dirname(netid, feature_name)
    return '{}.hdf5'.format(dirname)


def get_descstore_filepaths(desc_rootpath, feature_name_list, caffe_cnn,
                            snapshot_id, slug_extra=''):
    slug, snapshot_id = get_slug(caffe_cnn, snapshot_id, slug_extra)
    return [
        os.path.join(desc_rootpath, get_descstore_filename(slug, feature_name))
        for feature_name in feature_name_list
    ]


def dispatch_feature_comp(
    desc_rootpath,
    item_type,
    item_ids,
    node_batchsize,
    feature_name_list,
    num_dims_list,
    fetcomp_func,
    fetcomp_kwargs,
    slug,
    verbose=False,
):
    if verbose:
        print 'Dispatching feature computation for incomplete {} models'.format(item_type)

    # TODO: Use DescriptorAggregator instead, so we don't have to assemble the filename...
    completed_ids = set()
    for feature_name in feature_name_list:
        filename = get_descstore_filename(slug, feature_name)
        if os.path.exists(os.path.join(desc_rootpath, filename)):
            src_store = DescriptorStoreHdf5(
                path=os.path.join(desc_rootpath, filename),
                readonly=True,
                verbose=verbose,
            )
            if verbose:
                print 'Loaded completed descriptors for {}.'.format(filename)
            # Get already completed ids
            current_ids = set(src_store.ids[...])
            if completed_ids:
                completed_ids.intersection_update(current_ids)
            else:
                completed_ids = current_ids

    if verbose:
        print 'Getting ids for incomplete items...'
    todo_ids = item_ids.difference(completed_ids)
    if verbose:
        print '{} ids completed already'.format(len(completed_ids))
        print '{} ids to do...'.format(len(todo_ids))

    task_id = get_task_id(slug, feature_name_list)

    if verbose:
        pbar = ProgressBar(widgets=progress_bar_widgets(), maxval=len(todo_ids))
        pbar_counter = 0
        pbar.start()
    batchid_counter = 0
    batch = []

    for batch in iter_batch(todo_ids, node_batchsize):
        # Dispatch job
        batch_id = get_redis_key(
            slug,
            feature_name_list,
            batchid_counter,
        )
        fetcomp_func.delay(
            item_type,
            task_id,
            batch_id,
            batch,
            feature_name_list,
            fetcomp_kwargs,
        )

        batchid_counter += 1
        if verbose:
            pbar_counter += len(batch)
            pbar.update(pbar_counter)

    if verbose:
        pbar.finish()


def aggregate_feature_comp(
    desc_rootpath,
    item_type,
    item_ids,
    feature_name_list,
    num_dims_list,
    aggr_batchsize,
    slug,
    handle_interrupt_signal=True,
    verbose=False,
):
    if verbose:
        print 'Aggregating feature computation for {} models'.format(item_type)
    if handle_interrupt_signal:
        patch_interrupt_signal()

    dirname_list = [
        get_descstore_dirname(slug, feature_name)
        for feature_name in feature_name_list
    ]

    aggregator = DescriptorAggregator(
        feature_name_list, dirname_list, num_dims_list, verbose=verbose
    )
    aggregator.load(desc_rootpath, readonly=False)
    task_id = get_task_id(slug, feature_name_list)
    ret = aggregator.run(item_ids, task_id, aggr_batchsize=aggr_batchsize)

    if handle_interrupt_signal:
        detach_patch_interrupt_signal()

    return ret


def retrieve_features(
    desc_rootpath,
    item_ids,
    feature_name_list,
    slug,
    verbose=False,
):
    features = {}
    for feature_name in feature_name_list:
        filename = get_descstore_filename(slug, feature_name)
        src_store = DescriptorStoreHdf5(
            path=os.path.join(desc_rootpath, filename),
            readonly=True,
            verbose=verbose,
        )
        if item_ids is None:
            item_ids = src_store.ids[...]

        features[feature_name] = src_store.block_get(
            item_ids, show_progress=verbose
        )
        del src_store

    return features


def compute_features(
    desc_rootpath,
    item_type,
    item_ids,
    node_batchsize,
    aggr_batchsize,
    feature_name_list,
    num_dims_list,
    fetcomp_func,
    fetcomp_kwargs,
    slug,
    handle_interrupt_signal=True,
    verbose=False,
):
    """<commit_msg>Returning item_ids if it was None<commit_after>import os

from progressbar import ProgressBar

from cnntools.descriptor_aggregator import DescriptorAggregator
from cnntools.descstore import DescriptorStoreHdf5
from cnntools.models import CaffeCNNSnapshot
from cnntools.redis_aggregator import patch_interrupt_signal, detach_patch_interrupt_signal
from cnntools.tasks import compute_cnn_features_gpu_task
from cnntools.common_utils import iter_batch, progress_bar_widgets


def get_snapshot_id(caffe_cnn, snapshot_id):
    if snapshot_id:
        snapshot = CaffeCNNSnapshot.objects.get(id=snapshot_id)
        if snapshot.training_run.net.id != caffe_cnn.id:
            raise ValueError(
                'Snapshot\'s corresponding network id doesn\'t match given'
                ' network id!'
            )
        return snapshot_id

    # Get last snapshot
    snapshot = CaffeCNNSnapshot.objects.\
        filter(training_run__net__id=caffe_cnn.id).\
        order_by('-id')[0]

    return snapshot.id


def get_slug(caffe_cnn, snapshot_id, slug_extra):
    snapshot_id = get_snapshot_id(caffe_cnn, snapshot_id)
    slug = 'netid-{}:snapshot-{}'.format(caffe_cnn.netid, snapshot_id)
    if slug_extra:
        slug += slug_extra

    return slug, snapshot_id


def get_redis_key_core(slug, feature_name_list):
    return 'desc:slug-%s:fetname%s:batchid#' % (
        slug,
        '-'.join(feature_name_list),
    )


def get_redis_key(slug, feature_name_list, batchid_counter):
    return '%s%s' % (
        get_redis_key_core(slug, feature_name_list),
        batchid_counter,
    )


def get_task_id(slug, feature_name_list):
    return '%stask' % (
        get_redis_key_core(slug, feature_name_list),
    )


def get_descstore_dirname(netid, feature_name):
    dirname = '{}-{}'.format(netid, feature_name)
    # clean up
    invalid_chars = ['/', '(', ')', '[', ']', '{', '}', ' ', ':']

    for c in invalid_chars:
        dirname = dirname.replace(c, '-')

    return dirname


def get_descstore_filename(netid, feature_name):
    dirname = get_descstore_dirname(netid, feature_name)
    return '{}.hdf5'.format(dirname)


def get_descstore_filepaths(desc_rootpath, feature_name_list, caffe_cnn,
                            snapshot_id, slug_extra=''):
    slug, snapshot_id = get_slug(caffe_cnn, snapshot_id, slug_extra)
    return [
        os.path.join(desc_rootpath, get_descstore_filename(slug, feature_name))
        for feature_name in feature_name_list
    ]


def dispatch_feature_comp(
    desc_rootpath,
    item_type,
    item_ids,
    node_batchsize,
    feature_name_list,
    num_dims_list,
    fetcomp_func,
    fetcomp_kwargs,
    slug,
    verbose=False,
):
    if verbose:
        print 'Dispatching feature computation for incomplete {} models'.format(item_type)

    # TODO: Use DescriptorAggregator instead, so we don't have to assemble the filename...
    completed_ids = set()
    for feature_name in feature_name_list:
        filename = get_descstore_filename(slug, feature_name)
        if os.path.exists(os.path.join(desc_rootpath, filename)):
            src_store = DescriptorStoreHdf5(
                path=os.path.join(desc_rootpath, filename),
                readonly=True,
                verbose=verbose,
            )
            if verbose:
                print 'Loaded completed descriptors for {}.'.format(filename)
            # Get already completed ids
            current_ids = set(src_store.ids[...])
            if completed_ids:
                completed_ids.intersection_update(current_ids)
            else:
                completed_ids = current_ids

    if verbose:
        print 'Getting ids for incomplete items...'
    todo_ids = item_ids.difference(completed_ids)
    if verbose:
        print '{} ids completed already'.format(len(completed_ids))
        print '{} ids to do...'.format(len(todo_ids))

    task_id = get_task_id(slug, feature_name_list)

    if verbose:
        pbar = ProgressBar(widgets=progress_bar_widgets(), maxval=len(todo_ids))
        pbar_counter = 0
        pbar.start()
    batchid_counter = 0
    batch = []

    for batch in iter_batch(todo_ids, node_batchsize):
        # Dispatch job
        batch_id = get_redis_key(
            slug,
            feature_name_list,
            batchid_counter,
        )
        fetcomp_func.delay(
            item_type,
            task_id,
            batch_id,
            batch,
            feature_name_list,
            fetcomp_kwargs,
        )

        batchid_counter += 1
        if verbose:
            pbar_counter += len(batch)
            pbar.update(pbar_counter)

    if verbose:
        pbar.finish()


def aggregate_feature_comp(
    desc_rootpath,
    item_type,
    item_ids,
    feature_name_list,
    num_dims_list,
    aggr_batchsize,
    slug,
    handle_interrupt_signal=True,
    verbose=False,
):
    if verbose:
        print 'Aggregating feature computation for {} models'.format(item_type)
    if handle_interrupt_signal:
        patch_interrupt_signal()

    dirname_list = [
        get_descstore_dirname(slug, feature_name)
        for feature_name in feature_name_list
    ]

    aggregator = DescriptorAggregator(
        feature_name_list, dirname_list, num_dims_list, verbose=verbose
    )
    aggregator.load(desc_rootpath, readonly=False)
    task_id = get_task_id(slug, feature_name_list)
    ret = aggregator.run(item_ids, task_id, aggr_batchsize=aggr_batchsize)

    if handle_interrupt_signal:
        detach_patch_interrupt_signal()

    return ret


def retrieve_features(
    desc_rootpath,
    item_ids,
    feature_name_list,
    slug,
    verbose=False,
):
    ret_item_ids = item_ids is None
    features = {}
    for feature_name in feature_name_list:
        filename = get_descstore_filename(slug, feature_name)
        src_store = DescriptorStoreHdf5(
            path=os.path.join(desc_rootpath, filename),
            readonly=True,
            verbose=verbose,
        )
        # Will be called only once
        if item_ids is None:
            item_ids = src_store.ids[...]

        features[feature_name] = src_store.block_get(
            item_ids, show_progress=verbose
        )
        del src_store

    if ret_item_ids:
        return item_ids, features
    else:
        return features


def compute_features(
    desc_rootpath,
    item_type,
    item_ids,
    node_batchsize,
    aggr_batchsize,
    feature_name_list,
    num_dims_list,
    fetcomp_func,
    fetcomp_kwargs,
    slug,
    handle_interrupt_signal=True,
    verbose=False,
):
    """
{'fraction_non_alphanumeric': 0.08465812625405156, 'fraction_numerical': 0.0012347584503781448, 'mean_word_length': 2.65970064953403, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2421934', 'n_tokens_mistral': 4191, 'n_tokens_neox': 4068, 'n_words': 780}