# -*- coding: utf-8 -*-
# Copyright (c) 2014-2017, NVIDIA CORPORATION.  All rights reserved.

import os

import flask
from flask_wtf import FlaskForm
import wtforms
from wtforms import validators

from digits.config import config_value
from digits.device_query import get_device, get_nvml_info
from digits import utils
from digits.utils import sizeof_fmt
from digits.utils.forms import validate_required_iff
from digits import frameworks
from flask_babel import lazy_gettext as _

class ModelForm(FlaskForm):

    # Methods

    def selection_exists_in_choices(form, field):
        found = False
        for choice in field.choices:
            if choice[0] == field.data:
                found = True
        if not found:
            raise validators.ValidationError(_("Selected job doesn't exist. Maybe it was deleted by another user."))

    def validate_NetParameter(form, field):
        fw = frameworks.get_framework_by_id(form['framework'].data)
        try:
            # below function raises a BadNetworkException in case of validation error
            fw.validate_network(field.data)
        except frameworks.errors.BadNetworkError as e:
            raise validators.ValidationError(_('Bad network: %(message)s', message=e.message))

    def validate_file_exists(form, field):
        from_client = bool(form.python_layer_from_client.data)

        filename = ''
        if not from_client and field.type == 'StringField':
            filename = field.data

        if filename == '':
            return

        if not os.path.isfile(filename):
            raise validators.ValidationError(_('Server side file, %(filename)s, does not exist.', filename=filename))

    def validate_py_ext(form, field):
        from_client = bool(form.python_layer_from_client.data)

        filename = ''
        if from_client and field.type == 'FileField':
            filename = flask.request.files[field.name].filename
        elif not from_client and field.type == 'StringField':
            filename = field.data

        if filename == '':
            return

        (root, ext) = os.path.splitext(filename)
        if ext != '.py' and ext != '.pyc':
            raise validators.ValidationError(_('Python file, %(filename)s, needs .py or .pyc extension.',
                                               filename=filename))

    # Fields

    # The options for this get set in the view (since they are dynamic)
    dataset = utils.forms.SelectField(
        _('Select Dataset'),
        choices=[],
        tooltip=_("Choose the dataset to use for this model.")
    )

    python_layer_from_client = utils.forms.BooleanField(
        _('Use client-side file'),
        default=False,
    )

    python_layer_client_file = utils.forms.FileField(
        _('Client-side file'),
        validators=[
            validate_py_ext
        ],
        tooltip=_("Choose a Python file on the client containing layer definitions.")
    )
    python_layer_server_file = utils.forms.StringField(
        _('Server-side file'),
        validators=[
            validate_file_exists,
            validate_py_ext
        ],
        tooltip=_("Choose a Python file on the server containing layer definitions.")
    )

    train_epochs = utils.forms.IntegerField(
        _('Training epochs'),
        validators=[
            validators.NumberRange(min=1)
        ],
        default=30,
        tooltip=_("How many passes through the training data?")
    )

    snapshot_interval = utils.forms.FloatField(
        _('Snapshot interval (in epochs)'),
        default=1,
        validators=[
            validators.NumberRange(min=0),
        ],
        tooltip=_("How many epochs of training between taking a snapshot?")
    )

    val_interval = utils.forms.FloatField(
        _('Validation interval (in epochs)'),
        default=1,
        validators=[
            validators.NumberRange(min=0)
        ],
        tooltip=_("How many epochs of training between running through one pass of the validation data?")
    )

    traces_interval = utils.forms.IntegerField(
        _('Tracing Interval (in steps)'),
        validators=[
            validators.NumberRange(min=0)
        ],
        default=0,
        tooltip=_("Generation of a timeline trace every few steps")
    )

    random_seed = utils.forms.IntegerField(
        _('Random seed'),
        validators=[
            validators.NumberRange(min=0),
            validators.Optional(),
        ],
        tooltip=_('If you provide a random seed, then back-to-back runs with '
                  'the same model and dataset should give identical results.')
    )

    batch_size = utils.forms.MultiIntegerField(
        _('Batch size'),
        default=100,
        validators=[
            utils.forms.MultiNumberRange(min=1),
            utils.forms.MultiOptional(),
        ],
        tooltip=_("How many images to process at once. If blank, values are used from the network definition.")
    )

    batch_accumulation = utils.forms.IntegerField(
        _('Batch Accumulation'),
        validators=[
            validators.NumberRange(min=1),
            validators.Optional(),
        ],
        tooltip=_("Accumulate gradients over multiple batches (useful when you "
                  "need a bigger batch size for training but it doesn't fit in memory).")
    )

    # Solver types

    solver_type = utils.forms.SelectField(
        _('Solver type'),
        choices=[
            ('SGD', _('SGD (Stochastic Gradient Descent)')),
            ('MOMENTUM', _('Momentum')),
            ('NESTEROV', _("NAG (Nesterov's accelerated gradient)")),
            ('ADAGRAD', _('AdaGrad (Adaptive Gradient)')),
            ('ADAGRADDA', _('AdaGradDA (AdaGrad Dual Averaging)')),
            ('ADADELTA', _('AdaDelta')),
            ('ADAM', _('Adam (Adaptive Moment Estimation)')),
            ('RMSPROP', _('RMSprop')),
            ('FTRL', _('FTRL (Follow-The-Regularized-Leader)')),
        ],
        default='SGD',
        tooltip=_("What type of solver will be used?"),
    )

    def validate_solver_type(form, field):
        fw = frameworks.get_framework_by_id(form.framework)
        if fw is not None:
            if not fw.supports_solver_type(field.data):
                raise validators.ValidationError(
                    _('Solver type not supported by this framework'))

    # Additional settings specific to selected solver

    rms_decay = utils.forms.FloatField(
        _('RMS decay value'),
        default=0.99,
        validators=[
            validators.NumberRange(min=0),
        ],
        tooltip=_("If the gradient updates results in oscillations the gradient is reduced "
                  "by times 1-rms_decay. Otherwise it will be increased by rms_decay.")
    )

    # Learning rate

    learning_rate = utils.forms.MultiFloatField(
        _('Base Learning Rate'),
        default=0.01,
        validators=[
            utils.forms.MultiNumberRange(min=0),
        ],
        tooltip=_("Affects how quickly the network learns. If you are getting "
                 "NaN for your loss, you probably need to lower this value.")
    )

    lr_policy = wtforms.SelectField(
        _('Policy'),
        choices=[
            ('fixed', _('Fixed')),
            ('step', _('Step Down')),
            ('multistep', _('Step Down (arbitrary steps)')),
            ('exp', _('Exponential Decay')),
            ('inv', _('Inverse Decay')),
            ('poly', _('Polynomial Decay')),
            ('sigmoid', _('Sigmoid Decay')),
        ],
        default='step'
    )

    lr_step_size = wtforms.FloatField(_('Step Size'), default=33)
    lr_step_gamma = wtforms.FloatField(_('Gamma'), default=0.1)
    lr_multistep_values = wtforms.StringField(_('Step Values'), default="50,85")

    def validate_lr_multistep_values(form, field):
        if form.lr_policy.data == 'multistep':
            for value in field.data.split(','):
                try:
                    float(value)
                except ValueError:
                    raise validators.ValidationError(_('invalid value'))

    lr_multistep_gamma = wtforms.FloatField(_('Gamma'), default=0.5)
    lr_exp_gamma = wtforms.FloatField(_('Gamma'), default=0.95)
    lr_inv_gamma = wtforms.FloatField(_('Gamma'), default=0.1)
    lr_inv_power = wtforms.FloatField(_('Power'), default=0.5)
    lr_poly_power = wtforms.FloatField(_('Power'), default=3)
    lr_sigmoid_step = wtforms.FloatField(_('Step'), default=50)
    lr_sigmoid_gamma = wtforms.FloatField(_('Gamma'), default=0.1)

    # Network

    # Use a SelectField instead of a HiddenField so that the default value
    # is used when nothing is provided (through the REST API)
    method = wtforms.SelectField(
        _('Network type'),
        choices=[
            ('standard', _('Standard network')),
            ('previous', _('Previous network')),
            ('pretrained', _('Pretrained network')),
            ('custom', _('Custom network')),
        ],
        default='standard',
    )

    # framework - hidden field, set by Javascript to the selected framework ID
    framework = wtforms.HiddenField(
        _('framework'),
        validators=[
            validators.AnyOf(
                [fw.get_id() for fw in frameworks.get_frameworks()],
                message=_('The framework you choose is not currently supported.')
            )
        ],
        default=frameworks.get_frameworks()[0].get_id()
    )

    # The options for this get set in the view (since they are dependent on the data type)
    standard_networks = wtforms.RadioField(
        _('Standard Networks'),
        validators=[
            validate_required_iff(method='standard'),
        ],
    )

    previous_networks = wtforms.RadioField(
        _('Previous Networks'),
        choices=[],
        validators=[
            validate_required_iff(method='previous'),
            selection_exists_in_choices,
        ],
    )

    pretrained_networks = wtforms.RadioField(
        _('Pretrained Networks'),
        choices=[],
        validators=[
            validate_required_iff(method='pretrained'),
            selection_exists_in_choices,
        ],
    )

    custom_network = utils.forms.TextAreaField(
        _('Custom Network'),
        validators=[
            validate_required_iff(method='custom'),
            validate_NetParameter,
        ],
    )

    custom_network_snapshot = utils.forms.TextField(
        _('Pretrained model(s)'),
        tooltip=_("Paths to pretrained model files, separated by '%(pathsep)s'. "
                  "Only edit this field if you understand how fine-tuning "
                  "works in caffe or torch.", pathsep=os.path.pathsep)
    )

    def validate_custom_network_snapshot(form, field):
        pass
#       if form.method.data == 'custom':
#           for filename in field.data.strip().split(os.path.pathsep):
#               if filename and not os.path.lexists(filename):
#                   raise validators.ValidationError('File "%s" does not exist' % filename)

    # Select one of several GPUs
    select_gpu = wtforms.RadioField(
        _('Select which GPU you would like to use'),
        choices=[('next', 'Next available')] + [(
            index,
            '#%s - %s (%s memory)' % (
                index,
                get_device(index).name,
                sizeof_fmt(
                    get_nvml_info(index)['memory']['total']
                    if get_nvml_info(index) and 'memory' in get_nvml_info(index)
                    else get_device(index).totalGlobalMem)
            ),
        ) for index in config_value('gpu_list').split(',') if index],
        default='next',
    )

    # Select N of several GPUs
    select_gpus = utils.forms.SelectMultipleField(
        _('Select which GPU[s] you would like to use'),
        choices=[(
            index,
            '#%s - %s (%s memory)' % (
                index,
                get_device(index).name,
                sizeof_fmt(
                    get_nvml_info(index)['memory']['total']
                    if get_nvml_info(index) and 'memory' in get_nvml_info(index)
                    else get_device(index).totalGlobalMem)
            ),
        ) for index in config_value('gpu_list').split(',') if index],
        tooltip=_("The job won't start until all of the chosen GPUs are available.")
    )

    # XXX For testing
    # The Flask test framework can't handle SelectMultipleFields correctly
    select_gpus_list = wtforms.StringField(_('Select which GPU[s] you would like to use (comma separated)'))

    def validate_select_gpus(form, field):
        if form.select_gpus_list.data:
            field.data = form.select_gpus_list.data.split(',')

    # Use next available N GPUs
    select_gpu_count = wtforms.IntegerField(_('Use this many GPUs (next available)'),
                                            validators=[
                                                validators.NumberRange(min=1, max=len(
                                                    config_value('gpu_list').split(',')))
                                            ],
                                            default=1,
                                            )

    def validate_select_gpu_count(form, field):
        if field.data is None:
            if form.select_gpus.data:
                # Make this field optional
                field.errors[:] = []
                raise validators.StopValidation()

    model_name = utils.forms.StringField(_('Model Name'),
                                         validators=[
                                             validators.DataRequired()
                                         ],
                                         tooltip=_("An identifier, later used to refer to this model in the Application.")
                                         )

    group_name = utils.forms.StringField(_('Group Name'),
                                         tooltip=_("An optional group name for organization on the main page.")
                                         )

    # allows shuffling data during training (for frameworks that support this, as indicated by
    # their Framework.can_shuffle_data() method)
    shuffle = utils.forms.BooleanField(_('Shuffle Train Data'),
                                       default=True,
                                       tooltip=_('For every epoch, shuffle the data before training.')
                                       )

    steps = utils.forms.IntegerField("训练总步长",
                                     default=4000,
                                     validators=[
                                         validators.NumberRange(min=1)
                                     ],
                                     tooltip="本次训练总步长数（迭代次数）")

    iter_store_step = utils.forms.IntegerField("步长间隔",
                                     default=1000,
                                     validators=[
                                         validators.NumberRange(min=1)
                                     ],
                                     tooltip="要间隔多少个步长来进行快照保存")

    train_batch_size = utils.forms.IntegerField("批处理大小",
                                     default=100,
                                     validators=[
                                         validators.NumberRange(min=1)
                                     ],
                                     tooltip="一次处理多少图片，默认为100")

    # bottleneck_dir = utils.forms.StringField("瓶颈值目录",
    #                                  tooltip="计算出每个图片的瓶颈值并存储于此目录下")
