<commit_before>import logging
from abc import abstractmethod

import ee

from aoi import Aoi
from dates import parse_date, day_of_year, add_years, to_millis, milis_per_day
from image_spec import ImageSpec
from mosaic import Mosaic


class MosaicSpec(ImageSpec):
    def __init__(self, spec):
        super(MosaicSpec, self).__init__()
        recipe = spec['recipe']
        model = recipe['model']
        dates = model['dates']
        bands = spec.get('bands', {})
        self.spec = spec
        self.aoi = Aoi.create(model['aoi'])
        self.target_date = parse_date(dates['targetDate'])
        self.season_start = parse_date(dates['seasonStart'])
        self.season_end = parse_date(dates['seasonEnd'])
        self.years_before = int(dates['yearsBefore'])
        self.years_after = int(dates['yearsAfter'])
        self.target_day = day_of_year(parse_date(dates['targetDate']))
        self.from_date = add_years(parse_date(dates['seasonStart']), - int(dates['yearsBefore']))
        self.to_date = add_years(parse_date(dates['seasonEnd']), int(dates['yearsAfter']))
        self.target_day_weight = self._filter('DAY_OF_YEAR')
        self.shadow_tolerance = 1 - self._filter('SHADOW')
        self.haze_tolerance = 1 - self._filter('HAZE')
        self.greenness_weight = self._filter('NDVI')
        self.bands = bands.get('selection', [])
        self.median_composite = model['compositeOptions']['compose'] == 'MEDIAN'
        self.mask_clouds = 'CLOUDS' in model['compositeOptions']['mask']
        self.mask_snow = 'SNOW' in model['compositeOptions']['mask']
        self.brdf_correct = 'BRDF' in model['compositeOptions']['corrections']
        self.surface_reflectance = 'SR' in model['compositeOptions']['corrections']
        self.pan_sharpen = bool(bands.get('panSharpen', False))

    def _viz_params(self):
        viz_by_band = _sr_viz_by_bands if self.surface_reflectance else _toa_viz_by_bands
        return viz_by_band[', '.join(self.bands)]({
            'from_days_since_epoch': to_millis(self.from_date) / milis_per_day,
            'to_days_since_epoch': to_millis(self.to_date) / milis_per_day
        })

    def _ee_image(self):
        logging.info('Creating mosaic of ' + str(self))
        return Mosaic(self).create(self._data_sets())

    def _date_filter(self):
        def filter(year_diff):
            return ee.Filter.date(
                add_years(self.season_start, year_diff),
                add_years(self.season_end, year_diff)
            )

        filters = ee.Filter.Or(
            [filter(0)]
            + [filter(-(i + 1)) for i in range(0, self.years_before)]
            + [filter(i + 1) for i in range(0, self.years_after)]
        )
        return filters

    @abstractmethod
    def _data_sets(self):
        raise AssertionError('Method in subclass expected to have been invoked')

    def _filter(self, type):
        filters_of_type = [filter for filter in self.spec['recipe']['model']['compositeOptions']['filters'] if
                           filter['type'] == type]
        filter = filters_of_type[0] if len(filters_of_type) else None
        return _to_float(filter, 'percentile', 0) / 100


_toa_viz_by_bands = {
    'red, green, blue': lambda params: {
        'bands': 'red,green,blue',
        'min': '200, 400, 600',
        'max': '2400, 2200, 2400',
        'gamma': 1.2},
    'nir, red, green': lambda params: {
        'bands': 'nir,red,green',
        'min': '500, 200, 400',
        'max': '5000, 2400, 2200'},
    'nir, swir1, red': lambda params: {
        'bands': 'nir, swir1, red',
        'min': 0,
        'max': 5000,
        'gamma': 1.5},
    'swir2, nir, red': lambda params: {
        'bands': 'swir2,nir,red',
        'min': '0, 500, 200',
        'max': '1800, 6000, 3500'},
    'swir2, swir1, red': lambda params: {
        'bands': 'swir2,swir1,red',
        'min': '0, 500, 200',
        'max': '1800, 3000, 2400'},
    'swir2, nir, green': lambda params: {
        'bands': 'swir2,nir,green',<commit_msg>Specified bands without spaces in viz params.<commit_after>import logging
from abc import abstractmethod

import ee

from aoi import Aoi
from dates import parse_date, day_of_year, add_years, to_millis, milis_per_day
from image_spec import ImageSpec
from mosaic import Mosaic


class MosaicSpec(ImageSpec):
    def __init__(self, spec):
        super(MosaicSpec, self).__init__()
        recipe = spec['recipe']
        model = recipe['model']
        dates = model['dates']
        bands = spec.get('bands', {})
        self.spec = spec
        self.aoi = Aoi.create(model['aoi'])
        self.target_date = parse_date(dates['targetDate'])
        self.season_start = parse_date(dates['seasonStart'])
        self.season_end = parse_date(dates['seasonEnd'])
        self.years_before = int(dates['yearsBefore'])
        self.years_after = int(dates['yearsAfter'])
        self.target_day = day_of_year(parse_date(dates['targetDate']))
        self.from_date = add_years(parse_date(dates['seasonStart']), - int(dates['yearsBefore']))
        self.to_date = add_years(parse_date(dates['seasonEnd']), int(dates['yearsAfter']))
        self.target_day_weight = self._filter('DAY_OF_YEAR')
        self.shadow_tolerance = 1 - self._filter('SHADOW')
        self.haze_tolerance = 1 - self._filter('HAZE')
        self.greenness_weight = self._filter('NDVI')
        self.bands = bands.get('selection', [])
        self.median_composite = model['compositeOptions']['compose'] == 'MEDIAN'
        self.mask_clouds = 'CLOUDS' in model['compositeOptions']['mask']
        self.mask_snow = 'SNOW' in model['compositeOptions']['mask']
        self.brdf_correct = 'BRDF' in model['compositeOptions']['corrections']
        self.surface_reflectance = 'SR' in model['compositeOptions']['corrections']
        self.pan_sharpen = bool(bands.get('panSharpen', False))

    def _viz_params(self):
        viz_by_band = _sr_viz_by_bands if self.surface_reflectance else _toa_viz_by_bands
        return viz_by_band[', '.join(self.bands)]({
            'from_days_since_epoch': to_millis(self.from_date) / milis_per_day,
            'to_days_since_epoch': to_millis(self.to_date) / milis_per_day
        })

    def _ee_image(self):
        logging.info('Creating mosaic of ' + str(self))
        return Mosaic(self).create(self._data_sets())

    def _date_filter(self):
        def filter(year_diff):
            return ee.Filter.date(
                add_years(self.season_start, year_diff),
                add_years(self.season_end, year_diff)
            )

        filters = ee.Filter.Or(
            [filter(0)]
            + [filter(-(i + 1)) for i in range(0, self.years_before)]
            + [filter(i + 1) for i in range(0, self.years_after)]
        )
        return filters

    @abstractmethod
    def _data_sets(self):
        raise AssertionError('Method in subclass expected to have been invoked')

    def _filter(self, type):
        filters_of_type = [filter for filter in self.spec['recipe']['model']['compositeOptions']['filters'] if
                           filter['type'] == type]
        filter = filters_of_type[0] if len(filters_of_type) else None
        return _to_float(filter, 'percentile', 0) / 100


_toa_viz_by_bands = {
    'red, green, blue': lambda params: {
        'bands': 'red,green,blue',
        'min': '200, 400, 600',
        'max': '2400, 2200, 2400',
        'gamma': 1.2},
    'nir, red, green': lambda params: {
        'bands': 'nir,red,green',
        'min': '500, 200, 400',
        'max': '5000, 2400, 2200'},
    'nir, swir1, red': lambda params: {
        'bands': 'nir,swir1,red',
        'min': 0,
        'max': 5000,
        'gamma': 1.5},
    'swir2, nir, red': lambda params: {
        'bands': 'swir2,nir,red',
        'min': '0, 500, 200',
        'max': '1800, 6000, 3500'},
    'swir2, swir1, red': lambda params: {
        'bands': 'swir2,swir1,red',
        'min': '0, 500, 200',
        'max': '1800, 3000, 2400'},
    'swir2, nir, green': lambda params: {
        'bands': 'swir2,nir,green',
{'fraction_non_alphanumeric': 0.14510923535253228, 'fraction_numerical': 0.027557100297914597, 'mean_word_length': 2.8366666666666664, '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': '2494183', 'n_tokens_mistral': 2863, 'n_tokens_neox': 2631, 'n_words': 549}