<commit_before>import uuid
import logging
import operator
import itertools

from contexttimer import Timer

from PIL import Image

from django.db.models import Q
from django.db import (
    transaction,
    IntegrityError,
)
from django.utils import timezone

from huey.djhuey import (
    db_task,
    periodic_task,
    crontab,
)

from mozy.apps.mosaic.models import (
    NormalizedSourceImage,
    SourceImageTile,
    MosaicImage,
)
from mozy.apps.mosaic.utils import (
    decompose_an_image,
    convert_image_to_django_file,
    extract_pixel_data_from_image,
    cast_image_data_to_numpy_array,
)
from mozy.apps.mosaic.backends import (
    get_mosaic_backend,
)


logger = logging.getLogger(__file__)


@db_task()
def create_source_image_tiles(source_image_pk):
    source_image = NormalizedSourceImage.objects.get(pk=source_image_pk)

    if source_image.image.file.closed:
        source_image.image.file.open()

    tile_data = decompose_an_image(
        Image.open(source_image.image.file),
        tile_size=source_image.tile_size,
    )

    with Timer() as timer:
        for box_coords, tile_image in tile_data.items():
            already_exists = SourceImageTile.objects.filter(
                main_image=source_image,
                upper_left_x=box_coords[0],
                upper_left_y=box_coords[1],
            ).exists()
            if already_exists:
                continue
            tile = SourceImageTile(
                main_image=source_image,
                upper_left_x=box_coords[0],
                upper_left_y=box_coords[1],
                tile_data=extract_pixel_data_from_image(tile_image),
            )
            try:
                with transaction.atomic():
                    tile.tile_image.save(
                        "{0}.png".format(str(uuid.uuid4())),
                        convert_image_to_django_file(tile_image),
                        save=True
                    )
            except IntegrityError:
                pass
            tile_image.close()
    logger.info(
        "Took %s to create source image tiles for NormalizedSourceImage: %s",
        timer.elapsed,
        source_image_pk,
    )
    # go ahead and trigger tile matching
    queue_source_image_tiles_for_matching()


MATCH_BATCH_SIZE = 100


@periodic_task(crontab(minute='*/5'))
def queue_source_image_tiles_for_matching():
    tiles_needing_matching = SourceImageTile.objects.filter(
        Q(status=SourceImageTile.STATUS_PENDING) |
        Q(
            status=SourceImageTile.STATUS_QUEUED,
            updated_at__lte=SourceImageTile.get_errored_datetime(),
        ) |
        Q(
            status=SourceImageTile.STATUS_MATCHING,
            updated_at__lte=SourceImageTile.get_errored_datetime(),
        )
    ).distinct().order_by('main_image', 'pk')

    tile_data = tuple(tiles_needing_matching.values_list('main_image', 'pk'))

    with transaction.atomic():
        tiles_needing_matching.select_for_update().update(
            status=SourceImageTile.STATUS_QUEUED,
            updated_at=timezone.now(),
        )

    key_fn = operator.itemgetter(0)

    for _, tile_pks_group in itertools.groupby(tile_data, key_fn):
        # extract the tile pks
        tile_pks = zip(*tile_pks_group)[1]

        for i in range(0, len(tile_pks), MATCH_BATCH_SIZE):
            match_souce_image_tiles(tile_pks[i:i + MATCH_BATCH_SIZE])
    logger.info(
        "Enqueued %s SourceImageTile instances for matching", len(tile_data)
    )


@db_task()
def match_souce_image_tiles(source_image_tile_pks):
    with Timer() as timer:
        try:
            source_image = NormalizedSourceImage.objects.get(all_tiles__pk=source_image_tile_pks[0])
        except:
            logger.info("%s", repr(source_image_tile_pks))
            raise

        matcher = get_mosaic_backend()(source_image)

        tile_data_array = tuple((
            (tile_pk, cast_image_data_to_numpy_array(tile_data))
            for tile_pk, tile_data
            in SourceImageTile.objects.filter(pk__in=source_image_tile_pks).values_list(
                'pk', 'tile_data',
            )
        ))

        match_data = matcher(tile_data_array)

        for tile_pk, stock_id, match_similarity in match_data:<commit_msg>remove debug<commit_after>import uuid
import logging
import operator
import itertools

from contexttimer import Timer

from PIL import Image

from django.db.models import Q
from django.db import (
    transaction,
    IntegrityError,
)
from django.utils import timezone

from huey.djhuey import (
    db_task,
    periodic_task,
    crontab,
)

from mozy.apps.mosaic.models import (
    NormalizedSourceImage,
    SourceImageTile,
    MosaicImage,
)
from mozy.apps.mosaic.utils import (
    decompose_an_image,
    convert_image_to_django_file,
    extract_pixel_data_from_image,
    cast_image_data_to_numpy_array,
)
from mozy.apps.mosaic.backends import (
    get_mosaic_backend,
)


logger = logging.getLogger(__file__)


@db_task()
def create_source_image_tiles(source_image_pk):
    source_image = NormalizedSourceImage.objects.get(pk=source_image_pk)

    if source_image.image.file.closed:
        source_image.image.file.open()

    tile_data = decompose_an_image(
        Image.open(source_image.image.file),
        tile_size=source_image.tile_size,
    )

    with Timer() as timer:
        for box_coords, tile_image in tile_data.items():
            already_exists = SourceImageTile.objects.filter(
                main_image=source_image,
                upper_left_x=box_coords[0],
                upper_left_y=box_coords[1],
            ).exists()
            if already_exists:
                continue
            tile = SourceImageTile(
                main_image=source_image,
                upper_left_x=box_coords[0],
                upper_left_y=box_coords[1],
                tile_data=extract_pixel_data_from_image(tile_image),
            )
            try:
                with transaction.atomic():
                    tile.tile_image.save(
                        "{0}.png".format(str(uuid.uuid4())),
                        convert_image_to_django_file(tile_image),
                        save=True
                    )
            except IntegrityError:
                pass
            tile_image.close()
    logger.info(
        "Took %s to create source image tiles for NormalizedSourceImage: %s",
        timer.elapsed,
        source_image_pk,
    )
    # go ahead and trigger tile matching
    queue_source_image_tiles_for_matching()


MATCH_BATCH_SIZE = 100


@periodic_task(crontab(minute='*/5'))
def queue_source_image_tiles_for_matching():
    tiles_needing_matching = SourceImageTile.objects.filter(
        Q(status=SourceImageTile.STATUS_PENDING) |
        Q(
            status=SourceImageTile.STATUS_QUEUED,
            updated_at__lte=SourceImageTile.get_errored_datetime(),
        ) |
        Q(
            status=SourceImageTile.STATUS_MATCHING,
            updated_at__lte=SourceImageTile.get_errored_datetime(),
        )
    ).distinct().order_by('main_image', 'pk')

    tile_data = tuple(tiles_needing_matching.values_list('main_image', 'pk'))

    with transaction.atomic():
        tiles_needing_matching.select_for_update().update(
            status=SourceImageTile.STATUS_QUEUED,
            updated_at=timezone.now(),
        )

    key_fn = operator.itemgetter(0)

    for _, tile_pks_group in itertools.groupby(tile_data, key_fn):
        # extract the tile pks
        tile_pks = zip(*tile_pks_group)[1]

        for i in range(0, len(tile_pks), MATCH_BATCH_SIZE):
            match_souce_image_tiles(tile_pks[i:i + MATCH_BATCH_SIZE])
    logger.info(
        "Enqueued %s SourceImageTile instances for matching", len(tile_data)
    )


@db_task()
def match_souce_image_tiles(source_image_tile_pks):
    with Timer() as timer:
        source_image = NormalizedSourceImage.objects.get(all_tiles__pk=source_image_tile_pks[0])

        matcher = get_mosaic_backend()(source_image)

        tile_data_array = tuple((
            (tile_pk, cast_image_data_to_numpy_array(tile_data))
            for tile_pk, tile_data
            in SourceImageTile.objects.filter(pk__in=source_image_tile_pks).values_list(
                'pk', 'tile_data',
            )
        ))

        match_data = matcher(tile_data_array)

        for tile_pk, stock_id, match_similarity in match_data:
{'fraction_non_alphanumeric': 0.08330353950661423, 'fraction_numerical': 0.003336908592539626, 'mean_word_length': 2.5201342281879193, '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': '87708', 'n_tokens_mistral': 2707, 'n_tokens_neox': 2605, 'n_words': 413}