<commit_before>#!/usr/bin/env python
# encoding: utf-8
"""
setArrangement.py

Created by <NAME> on 28 Aug 2014.
Licensed under a 3-clause BSD license.

Revision history:
    28 Aug 2014 <NAME>
      Initial version

"""

from __future__ import division
from __future__ import print_function
import itertools
from mangaLogic import checkExposure
from sdss.internal.manga.Totoro.exceptions import NoMangaExposure
from sdss.internal.manga.Totoro import TotoroDBConnection, log, config
from sdss.internal.manga.Totoro.utils import intervals
from sqlalchemy import func
from scipy.misc import factorial
import collections
import warnings
import numpy as np


def updatePlate(plate, **kwargs):
    """Finds missing exposures """

    log.debug('plate_id={0}: updating sets'.format(plate.plate_id))

    badSetStatus = checkBadSets(plate, **kwargs)

    newExposures = getNewExposures(plate)

    if len(newExposures) > 0:

        log.info('plate_id={0}: {1} new exposure(s) found with '
                 'mangaDB.Exposure.pk={0}'.format(
                     plate.plate_id, len(newExposures),
                     ', '.join([str(exp.mangadbExposure[0].pk)
                                for exp in newExposures])))

        if len(newExposures) >= \
                config['setArrangement']['forceRearrangementMinExposures']:
            log.info('plate_id={0}: more than {1} new exposures found. '
                     'Triggering a complete set rearrangement.'
                     .format(plate.plate_id,
                             config['setArrangement']
                             ['forceRearrangementMinExposures']))
            updateStatus = rearrangeSets(plate, **kwargs)
        else:
            results = []
            for exp in newExposures:
                results.append(addExposure(exp, plate))
            if any(results):
                removeOrphanSets()
                updateStatus = True
            else:
                updateStatus = False
    else:
        log.debug('plate_id={0}: no new exposures found.'.format(
                  plate.plate_id))
        updateStatus = False

    if updateStatus or badSetStatus:
        return True
    else:
        return False


def checkBadSets(plate, **kwargs):
    """Identifies bad sets and breaks them."""

    setQuality = [set.getQuality()[0] for set in plate.sets]
    if 'Bad' in setQuality:
        nBadSets = len([True for ss in setQuality if ss == 'Bad'])
        log.debug('plate_id={0}: found {1} bad sets. Fixing them.'.format(
                  plate.plate_id, nBadSets))
        plate = fixBadSets(plate, setQuality=setQuality)
        return True
    return False


def getValidSet(totoroExp, plate, setQuality=None):
    """Gets the best possible set for an exposure"""

    from sdss.internal.manga.Totoro import dbclasses

    if not totoroExp.isValid()[0]:
        return False

    if setQuality is None:
        setQuality = [set.getQuality(silent=True)[0] for set in plate.sets]

    incompleteSets = [set for nn, set in enumerate(plate.sets)
                      if setQuality[nn].lower() == 'incomplete']

    distanceToEndOfVisibilityWindow = [
        plate.getLSTRange()[1] - np.mean(set.getLSTRange())
        for set in incompleteSets]
    order = np.argsort(distanceToEndOfVisibilityWindow)
    incompleteSetsSorted = [incompleteSets[ii] for ii in order]

    mockSetQuality = []
    for set in incompleteSetsSorted:
        mockSet = dbclasses.Set.fromExposures(
            set.totoroExposures + [totoroExp], silent=True)
        mockSetQuality.append(mockSet.getQuality(silent=True)[0])
        if mockSetQuality[-1] in ['Good', 'Excellent']:
            return set

    try:
        firstIncompleteSet = mockSetQuality.index('Incomplete')
        return incompleteSetsSorted[firstIncompleteSet]
    except ValueError:
        return None


def addExposure(exp, plate):
    """Adds an exposure to a plate."""

    from sdss.internal.manga.Totoro import dbclasses

    if not isinstance(exp, dbclasses.Exposure):
        totoroExp = dbclasses.Exposure(exp, silent=True)
    else:
        totoroExp = exp

    validSet = getValidSet(exp, plate)
    if validSet is False:
        return False
    elif validSet is None:
        validSet = dbclasses.Set()
        plate.sets.append(validSet)

    validSet.totoroExposures.append(totoroExp)

    if plate.isMock:
        return True

    db = TotoroDBConnection()
    session = db.Session()

    with session.begin():
        if validSet.pk is None:
            validSet.pk = getNewSetPK()
            session.add(validSet)
            session.flush()
        totoroExp._mangaExposure.set_pk = validSet.pk
        session.add(totoroExp)
        session.flush()

    if validSet.pk is not None:
        log.info('adding mangaDB exposure pk={0} to set {1} -> {2} set.'
                 .format(totoroExp._mangaExposure.pk, validSet.pk,
                         validSet.getQuality(flag=False)[0]))

    return True


def getNewSetPK():
    """Gets the lowest available set pk."""

    db = TotoroDBConnection()
    session = db.Session()

    with session.begin(subtransactions=True):
        setPKs = session.query(db.mangaDB.Set.pk).all()

    setPKs = np.unique(setPKs)
    return [xx for xx in np.arange(1, np.max(setPKs)+1) if xx not in setPKs][0]


def getNewExposures(plate):
    """Gets unassigned exposures in a plate."""

    newExposures = []
    for exp in plate.getScienceExposures():
        if len(exp.mangadbExposure) == 0:
            warnings.warn('exposure pk={0} has no mangaDB '
                          'counterpart.'.format(exp.pk), NoMangaExposure)
            continue
        if exp.mangadbExposure[0].set_pk is not None:
            continue
        if checkExposure(exp, silent=True, flag=True)[0] is False:
            continue
        newExposures.append(exp)

    return newExposures


def rearrangeSets(plate, **kwargs):
    """Assigns a set to each exposure for a given plate. Overwrites any current
    assignment"""

    log.info('plate_id={0}: reorganising sets'.format(plate.plate_id))

    optimalArrangement = getOptimalArrangement(plate, **kwargs)

    if optimalArrangement is False:
        return False

    for exposure in optimalArrangement['invalid']:
        removeSet(exposure._mangaExposure.set_pk)

    for set in optimalArrangement['sets']:
        updateSet(set)

    removeOrphanSets()

    log.info('plate_id={0}: rearrangement was successful'.format(
             plate.plate_id))

    return True


def removeSet(set_pk, orphan=False):
    """Removes a set and sets set_pk to None for all its exposures."""

    if set_pk is None:
        return

    db = TotoroDBConnection()
    session = db.Session()

    with session.begin():
        set = session.query(db.mangaDB.Set).get(set_pk)
        if set is not None:
            if set.exposures is not None:
                for exp in set.exposures:
                    exp.set_pk = None
            session.delete(set)

    msg = 'removed orphan set pk={0}' if orphan else 'Removed set pk={0}'
    log.debug(msg.format(set_pk))

    return


def updateSet(set):
    """Updates the set_pk for the exposures in the set. Writes to the DB."""

    db = TotoroDBConnection()
    session = db.Session()

    setPKs = [exposure._mangaExposure.set_pk
              for exposure in set.totoroExposures]

    for ss in np.unique(setPKs):
        removeSet(ss)

    with session.begin():
        newSet = db.mangaDB.Set()
        newSet.pk = getNewSetPK()
        session.add(newSet)
        session.flush()
        for exposure in set.totoroExposures:
            mangaExp = exposure._mangaExposure
            mangaExp.set_pk = newSet.pk
            session.add(mangaExp)

    log.debug('created set pk={0} with mangaDB exposures pk={1}'
              .format(newSet.pk,
                      ', '.join([str(exp._mangaExposure.pk)
                                 for exp in set.totoroExposures])))
    return


def removeOrphanSets():
    """Removes sets without exposures."""

    db = TotoroDBConnection()
    session = db.Session()

    with session.begin():
        nullSets = session.query(db.mangaDB.Set).outerjoin(<commit_msg>Totoro: added subtransactions=True to a few sessions<commit_after>#!/usr/bin/env python
# encoding: utf-8
"""
setArrangement.py

Created by <NAME> on 28 Aug 2014.
Licensed under a 3-clause BSD license.

Revision history:
    28 Aug 2014 <NAME>
      Initial version

"""

from __future__ import division
from __future__ import print_function
import itertools
from mangaLogic import checkExposure
from sdss.internal.manga.Totoro.exceptions import NoMangaExposure
from sdss.internal.manga.Totoro import TotoroDBConnection, log, config
from sdss.internal.manga.Totoro.utils import intervals
from sqlalchemy import func
from scipy.misc import factorial
import collections
import warnings
import numpy as np


def updatePlate(plate, **kwargs):
    """Finds missing exposures """

    log.debug('plate_id={0}: updating sets'.format(plate.plate_id))

    badSetStatus = checkBadSets(plate, **kwargs)

    newExposures = getNewExposures(plate)

    if len(newExposures) > 0:

        log.info('plate_id={0}: {1} new exposure(s) found with '
                 'mangaDB.Exposure.pk={0}'.format(
                     plate.plate_id, len(newExposures),
                     ', '.join([str(exp.mangadbExposure[0].pk)
                                for exp in newExposures])))

        if len(newExposures) >= \
                config['setArrangement']['forceRearrangementMinExposures']:
            log.info('plate_id={0}: more than {1} new exposures found. '
                     'Triggering a complete set rearrangement.'
                     .format(plate.plate_id,
                             config['setArrangement']
                             ['forceRearrangementMinExposures']))
            updateStatus = rearrangeSets(plate, **kwargs)
        else:
            results = []
            for exp in newExposures:
                results.append(addExposure(exp, plate))
            if any(results):
                removeOrphanSets()
                updateStatus = True
            else:
                updateStatus = False
    else:
        log.debug('plate_id={0}: no new exposures found.'.format(
                  plate.plate_id))
        updateStatus = False

    if updateStatus or badSetStatus:
        return True
    else:
        return False


def checkBadSets(plate, **kwargs):
    """Identifies bad sets and breaks them."""

    setQuality = [set.getQuality()[0] for set in plate.sets]
    if 'Bad' in setQuality:
        nBadSets = len([True for ss in setQuality if ss == 'Bad'])
        log.debug('plate_id={0}: found {1} bad sets. Fixing them.'.format(
                  plate.plate_id, nBadSets))
        plate = fixBadSets(plate, setQuality=setQuality)
        return True
    return False


def getValidSet(totoroExp, plate, setQuality=None):
    """Gets the best possible set for an exposure"""

    from sdss.internal.manga.Totoro import dbclasses

    if not totoroExp.isValid()[0]:
        return False

    if setQuality is None:
        setQuality = [set.getQuality(silent=True)[0] for set in plate.sets]

    incompleteSets = [set for nn, set in enumerate(plate.sets)
                      if setQuality[nn].lower() == 'incomplete']

    distanceToEndOfVisibilityWindow = [
        plate.getLSTRange()[1] - np.mean(set.getLSTRange())
        for set in incompleteSets]
    order = np.argsort(distanceToEndOfVisibilityWindow)
    incompleteSetsSorted = [incompleteSets[ii] for ii in order]

    mockSetQuality = []
    for set in incompleteSetsSorted:
        mockSet = dbclasses.Set.fromExposures(
            set.totoroExposures + [totoroExp], silent=True)
        mockSetQuality.append(mockSet.getQuality(silent=True)[0])
        if mockSetQuality[-1] in ['Good', 'Excellent']:
            return set

    try:
        firstIncompleteSet = mockSetQuality.index('Incomplete')
        return incompleteSetsSorted[firstIncompleteSet]
    except ValueError:
        return None


def addExposure(exp, plate):
    """Adds an exposure to a plate."""

    from sdss.internal.manga.Totoro import dbclasses

    if not isinstance(exp, dbclasses.Exposure):
        totoroExp = dbclasses.Exposure(exp, silent=True)
    else:
        totoroExp = exp

    validSet = getValidSet(exp, plate)
    if validSet is False:
        return False
    elif validSet is None:
        validSet = dbclasses.Set()
        plate.sets.append(validSet)

    validSet.totoroExposures.append(totoroExp)

    if plate.isMock:
        return True

    db = TotoroDBConnection()
    session = db.Session()

    with session.begin(subtransactions=True):
        if validSet.pk is None:
            validSet.pk = getNewSetPK()
            session.add(validSet)
            session.flush()
        totoroExp._mangaExposure.set_pk = validSet.pk
        session.add(totoroExp)
        session.flush()

    if validSet.pk is not None:
        log.info('adding mangaDB exposure pk={0} to set {1} -> {2} set.'
                 .format(totoroExp._mangaExposure.pk, validSet.pk,
                         validSet.getQuality(flag=False)[0]))

    return True


def getNewSetPK():
    """Gets the lowest available set pk."""

    db = TotoroDBConnection()
    session = db.Session()

    with session.begin(subtransactions=True):
        setPKs = session.query(db.mangaDB.Set.pk).all()

    setPKs = np.unique(setPKs)
    return [xx for xx in np.arange(1, np.max(setPKs)+1) if xx not in setPKs][0]


def getNewExposures(plate):
    """Gets unassigned exposures in a plate."""

    newExposures = []
    for exp in plate.getScienceExposures():
        if len(exp.mangadbExposure) == 0:
            warnings.warn('exposure pk={0} has no mangaDB '
                          'counterpart.'.format(exp.pk), NoMangaExposure)
            continue
        if exp.mangadbExposure[0].set_pk is not None:
            continue
        if checkExposure(exp, silent=True, flag=True)[0] is False:
            continue
        newExposures.append(exp)

    return newExposures


def rearrangeSets(plate, **kwargs):
    """Assigns a set to each exposure for a given plate. Overwrites any current
    assignment"""

    log.info('plate_id={0}: reorganising sets'.format(plate.plate_id))

    optimalArrangement = getOptimalArrangement(plate, **kwargs)

    if optimalArrangement is False:
        return False

    for exposure in optimalArrangement['invalid']:
        removeSet(exposure._mangaExposure.set_pk)

    for set in optimalArrangement['sets']:
        updateSet(set)

    removeOrphanSets()

    log.info('plate_id={0}: rearrangement was successful'.format(
             plate.plate_id))

    return True


def removeSet(set_pk, orphan=False):
    """Removes a set and sets set_pk to None for all its exposures."""

    if set_pk is None:
        return

    db = TotoroDBConnection()
    session = db.Session()

    with session.begin(subtransactions=True):
        set = session.query(db.mangaDB.Set).get(set_pk)
        if set is not None:
            if set.exposures is not None:
                for exp in set.exposures:
                    exp.set_pk = None
            session.delete(set)

    msg = 'removed orphan set pk={0}' if orphan else 'Removed set pk={0}'
    log.debug(msg.format(set_pk))

    return


def updateSet(set):
    """Updates the set_pk for the exposures in the set. Writes to the DB."""

    db = TotoroDBConnection()
    session = db.Session()

    setPKs = [exposure._mangaExposure.set_pk
              for exposure in set.totoroExposures]

    for ss in np.unique(setPKs):
        removeSet(ss)

    with session.begin(subtransactions=True):
        newSet = db.mangaDB.Set()
        newSet.pk = getNewSetPK()
        session.add(newSet)
        session.flush()
        for exposure in set.totoroExposures:
            mangaExp = exposure._mangaExposure
            mangaExp.set_pk = newSet.pk
            session.add(mangaExp)

    log.debug('created set pk={0} with mangaDB exposures pk={1}'
              .format(newSet.pk,
                      ', '.join([str(exp._mangaExposure.pk)
                                 for exp in set.totoroExposures])))
    return


def removeOrphanSets():
    """Removes sets without exposures."""

    db = TotoroDBConnection()
    session = db.Session()

    with session.begin(subtransactions=True):
        nullSets = session.query(db.mangaDB.Set).outerjoin(
{'fraction_non_alphanumeric': 0.10165585006368653, 'fraction_numerical': 0.005822769454721902, 'mean_word_length': 2.397486091077684, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3064840', 'n_tokens_mistral': 5218, 'n_tokens_neox': 4815, 'n_words': 1233}