# Copyright (c) 2014-2017, NVIDIA CORPORATION.  All rights reserved.
from __future__ import absolute_import

import os.path
import re
import sys

import digits
from digits import utils
from digits.task import Task
from digits.utils import subclass, override

# NOTE: Increment this every time the pickled object
PICKLE_VERSION = 1


@subclass
class ParseS3Task(Task):
    """Parses a folder into textfiles"""

    def __init__(self, s3_endpoint_url, s3_bucket, s3_path, s3_accesskey, s3_secretkey, **kwargs):
        """
        Arguments:
        s3_endpoint_url -- endpoint url
        s3_bucket -- bucekt name
        s3_path -- a path to images
        s3_accesskey -- accesskey
        s3_secretkey -- secretkey

        Keyword arguments:
        percent_val -- percent of images used in the validation set
        percent_test -- percent of images used in the test set
        min_per_category -- minimum number of images per category
        max_per_category -- maximum number of images per category
        """
        # Take keyword arguments out of kwargs
        percent_val = kwargs.pop('percent_val', None)
        percent_test = kwargs.pop('percent_test', None)
        self.min_per_category = kwargs.pop('min_per_category', 2)
        self.max_per_category = kwargs.pop('max_per_category', None)

        super(ParseS3Task, self).__init__(**kwargs)
        self.pickver_task_parsefolder = PICKLE_VERSION

        self.s3_endpoint_url = s3_endpoint_url
        self.s3_bucket = s3_bucket
        self.s3_path = s3_path
        self.s3_accesskey = s3_accesskey
        self.s3_secretkey = s3_secretkey

        if percent_val is None:
            self.percent_val = 0
        else:
            pct = float(percent_val)
            if pct < 0:
                pct = 0
            elif pct > 100:
                raise ValueError('percent_val must not exceed 100')
            self.percent_val = pct

        if percent_test is None:
            self.percent_test = 0
        else:
            pct = float(percent_test)
            if pct < 0:
                pct = 0
            elif pct > 100:
                raise ValueError('percent_test must not exceed 100')
            self.percent_test = pct

        if percent_val is not None and percent_test is not None and percent_val + percent_test > 100:
            raise ValueError('the sum of percent_val and percent_test must not exceed 100')

        self.train_file = utils.constants.TRAIN_FILE
        self.val_file = utils.constants.VAL_FILE
        self.labels_file = utils.constants.LABELS_FILE

        # Results

        self.train_count = None
        self.val_count = None
        self.test_count = None
        self.label_count = None

    def __getstate__(self):
        state = super(ParseS3Task, self).__getstate__()
        return state

    def __setstate__(self, state):
        super(ParseS3Task, self).__setstate__(state)

    @override
    def name(self):
        sets = []
        if (self.percent_val + self.percent_test) < 100:
            sets.append('train')
        if self.percent_val > 0:
            sets.append('val')
        if self.percent_test > 0:
            sets.append('test')

        return 'Parse Folder (%s)' % ('/'.join(sets))

    @override
    def html_id(self):
        sets = []
        if (self.percent_val + self.percent_test) < 100:
            sets.append('train')
        if self.percent_val > 0:
            sets.append('val')
        if self.percent_test > 0:
            sets.append('test')

        return 'task-parse-folder-%s' % ('-'.join(sets))

    @override
    def offer_resources(self, resources):
        key = 'parse_folder_task_pool'
        if key not in resources:
            return None
        for resource in resources[key]:
            if resource.remaining() >= 1:
                return {key: [(resource.identifier, 1)]}
        return None

    @override
    def task_arguments(self, resources, env):
        args = [sys.executable, os.path.join(
            os.path.dirname(os.path.abspath(digits.__file__)),
            'tools', 'parse_s3.py'),
            self.s3_endpoint_url,
            self.s3_bucket,
            self.s3_path,
            self.s3_accesskey,
            self.s3_secretkey,
            self.path(utils.constants.LABELS_FILE),
            '--min=%s' % self.min_per_category,
        ]

        if (self.percent_val + self.percent_test) < 100:
            args.append('--train_file=%s' % self.path(utils.constants.TRAIN_FILE))
        if self.percent_val > 0:
            args.append('--val_file=%s' % self.path(utils.constants.VAL_FILE))
            args.append('--percent_val=%s' % self.percent_val)
        if self.percent_test > 0:
            args.append('--test_file=%s' % self.path(utils.constants.TEST_FILE))
            args.append('--percent_test=%s' % self.percent_test)
        if self.max_per_category is not None:
            args.append('--max=%s' % self.max_per_category)

        return args

    @override
    def process_output(self, line):
        timestamp, level, message = self.preprocess_output_digits(line)
        if not message:
            return False

        # progress
        match = re.match(r'Progress: ([-+]?[0-9]*\.?[0-9]+(e[-+]?[0-9]+)?)', message)
        if match:
            self.progress = float(match.group(1))
            self.emit_progress_update()
            return True

        # totals
        match = re.match(r'Found (\d+) images in (\d+) categories', message)
        if match:
            self.label_count = int(match.group(2))
            return True

        # splits
        match = re.match(r'Selected (\d+) for (\w+)', message)
        if match:
            if match.group(2).startswith('training'):
                self.train_count = int(match.group(1))
            elif match.group(2).startswith('validation'):
                self.val_count = int(match.group(1))
            elif match.group(2).startswith('test'):
                self.test_count = int(match.group(1))
            return True

        if level == 'warning':
            self.logger.warning('%s: %s' % (self.name(), message))
            return True
        if level in ['error', 'critical']:
            self.logger.error('%s: %s' % (self.name(), message))
            self.exception = message
            return True

        return True
