# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import os
import logging
import sys
import re

import mxnet as mx
from common import TemporaryDirectory
from mxnet import np
from mxnet.gluon import nn, loss
from mxnet.gluon.contrib.estimator import estimator, event_handler
from mxnet.gluon.contrib.estimator.event_handler import LoggingHandler
from mxnet.gluon.contrib.estimator.event_handler import ValidationHandler
from mxnet.gluon.contrib.estimator import EpochEnd
from mxnet.gluon.data.dataset import Dataset
try:
    from StringIO import StringIO
except ImportError:
    from io import StringIO

mx.npx.reset_np()

class AxisArrayDataset(Dataset):
    def __init__(self, * args):
        self._length = len(args[1])
        self._data = []
        for _, data in enumerate(args):
            self._data.append(data)

    def __getitem__(self, idx):
        if len(self._data) == 1:
            return self._data[idx][0]
        else:
            return tuple(data[:, idx] for data in self._data)

    def __len__(self):
        return self._length

class Handler(EpochEnd):
    def __init__(self):
        pass

    def epoch_end(self, estimator, *args, **kwargs):
        estimator.run_test_handler = True

def _get_test_network(net=nn.Sequential()):
    net.add(nn.Dense(128, activation='relu', flatten=False),
            nn.Dense(64, activation='relu'),
            nn.Dense(10, activation='relu'))
    return net


def _get_test_data(in_size=32):
    data = np.ones((in_size, 100))
    label = np.zeros((in_size, 1))
    data_arr = mx.gluon.data.dataset.ArrayDataset(data, label)
    return mx.gluon.data.DataLoader(data_arr, batch_size=8)

def _get_batch_axis_test_data(in_size=32):
    data = np.ones((100, in_size))
    label = np.zeros((1, in_size))
    data_arr = AxisArrayDataset(data, label)
    return mx.gluon.data.DataLoader(data_arr, batch_size=8)

@mx.util.use_np
def test_checkpoint_handler():
    with TemporaryDirectory() as tmpdir:
        model_prefix = 'test_epoch'
        file_path = os.path.join(tmpdir, model_prefix)
        test_data = _get_test_data()

        net = _get_test_network()
        ce_loss = loss.SoftmaxCrossEntropyLoss()
        acc = mx.gluon.metric.Accuracy()
        est = estimator.Estimator(net, loss=ce_loss, train_metrics=acc)
        checkpoint_handler = event_handler.CheckpointHandler(model_dir=tmpdir,
                                                             model_prefix=model_prefix,
                                                             monitor=acc,
                                                             save_best=True,
                                                             epoch_period=1)
        est.fit(test_data, event_handlers=[checkpoint_handler], epochs=1)
        assert checkpoint_handler.current_epoch == 1
        assert checkpoint_handler.current_batch == 4
        assert os.path.isfile(file_path + '-best.params')
        assert os.path.isfile(file_path + '-best.states')
        assert os.path.isfile(file_path + '-epoch0batch4.params')
        assert os.path.isfile(file_path + '-epoch0batch4.states')

        model_prefix = 'test_batch'
        file_path = os.path.join(tmpdir, model_prefix)
        net = _get_test_network(nn.HybridSequential())
        net.hybridize()
        est = estimator.Estimator(net, loss=ce_loss, train_metrics=acc)
        checkpoint_handler = event_handler.CheckpointHandler(model_dir=tmpdir,
                                                             model_prefix=model_prefix,
                                                             epoch_period=None,
                                                             batch_period=2,
                                                             max_checkpoints=2)
        est.fit(test_data, event_handlers=[checkpoint_handler], batches=10)
        assert checkpoint_handler.current_batch == 10
        assert checkpoint_handler.current_epoch == 3
        assert not os.path.isfile(file_path + 'best.params')
        assert not os.path.isfile(file_path + 'best.states')
        assert not os.path.isfile(file_path + '-epoch0batch0.params')
        assert not os.path.isfile(file_path + '-epoch0batch0.states')
        assert os.path.isfile(file_path + '-symbol.json')
        assert os.path.isfile(file_path + '-epoch1batch7.params')
        assert os.path.isfile(file_path + '-epoch1batch7.states')
        assert os.path.isfile(file_path + '-epoch2batch9.params')
        assert os.path.isfile(file_path + '-epoch2batch9.states')

@mx.util.use_np
def test_resume_checkpoint():
    with TemporaryDirectory() as tmpdir:
        model_prefix = 'test_net'
        file_path = os.path.join(tmpdir, model_prefix)
        test_data = _get_test_data()

        net = _get_test_network()
        ce_loss = loss.SoftmaxCrossEntropyLoss()
        acc = mx.gluon.metric.Accuracy()
        est = estimator.Estimator(net, loss=ce_loss, train_metrics=acc)
        checkpoint_handler = event_handler.CheckpointHandler(model_dir=tmpdir,
                                                             model_prefix=model_prefix,
                                                             monitor=acc,
                                                             max_checkpoints=1)
        est.fit(test_data, event_handlers=[checkpoint_handler], epochs=2)
        assert os.path.isfile(file_path + '-epoch1batch8.params')
        assert os.path.isfile(file_path + '-epoch1batch8.states')
        checkpoint_handler = event_handler.CheckpointHandler(model_dir=tmpdir,
                                                             model_prefix=model_prefix,
                                                             monitor=acc,
                                                             max_checkpoints=1,
                                                             resume_from_checkpoint=True)
        est.fit(test_data, event_handlers=[checkpoint_handler], epochs=5)
        # should only continue to train 3 epochs and last checkpoint file is epoch4
        assert est.max_epoch == 3
        assert os.path.isfile(file_path + '-epoch4batch20.states')


@mx.util.use_np
def test_early_stopping():
    test_data = _get_test_data()

    net = _get_test_network()
    ce_loss = loss.SoftmaxCrossEntropyLoss()
    acc = mx.gluon.metric.Accuracy()
    est = estimator.Estimator(net, loss=ce_loss, train_metrics=acc)
    early_stopping = event_handler.EarlyStoppingHandler(monitor=acc,
                                                        patience=0,
                                                        mode='min')
    est.fit(test_data, event_handlers=[early_stopping], epochs=5)
    assert early_stopping.current_epoch == 2
    assert early_stopping.stopped_epoch == 1

    early_stopping = event_handler.EarlyStoppingHandler(monitor=acc,
                                                        patience=2,
                                                        mode='auto')
    est.fit(test_data, event_handlers=[early_stopping], epochs=1)
    assert early_stopping.current_epoch == 1


@mx.util.use_np
def test_logging():
    with TemporaryDirectory() as tmpdir:
        test_data = _get_test_data()
        file_name = 'test_log'
        output_dir = os.path.join(tmpdir, file_name)

        net = _get_test_network()
        ce_loss = loss.SoftmaxCrossEntropyLoss()
        acc = mx.gluon.metric.Accuracy()
        est = estimator.Estimator(net, loss=ce_loss, train_metrics=acc)

        est.logger.addHandler(logging.FileHandler(output_dir))

        train_metrics = est.train_metrics
        val_metrics = est.val_metrics
        logging_handler = event_handler.LoggingHandler(metrics=train_metrics)
        est.fit(test_data, event_handlers=[logging_handler], epochs=3)
        assert logging_handler.batch_index == 0
        assert logging_handler.current_epoch == 3
        assert os.path.isfile(output_dir)
        del est  # Clean up estimator and logger before deleting tmpdir


@mx.util.use_np
def test_custom_handler():
    class CustomStopHandler(event_handler.TrainBegin,
                            event_handler.BatchEnd,
                            event_handler.EpochEnd):
        def __init__(self, batch_stop=None, epoch_stop=None):
            self.batch_stop = batch_stop
            self.epoch_stop = epoch_stop
            self.num_batch = 0
            self.num_epoch = 0
            self.stop_training = False

        def train_begin(self, estimator, *args, **kwargs):
            self.num_batch = 0
            self.num_epoch = 0

        def batch_end(self, estimator, *args, **kwargs):
            self.num_batch += 1
            if self.num_batch == self.batch_stop:
                self.stop_training = True
            return self.stop_training

        def epoch_end(self, estimator, *args, **kwargs):
            self.num_epoch += 1
            if self.num_epoch == self.epoch_stop:
                self.stop_training = True
            return self.stop_training

    # total data size is 32, batch size is 8
    # 4 batch per epoch
    test_data = _get_test_data()
    net = _get_test_network()
    ce_loss = loss.SoftmaxCrossEntropyLoss()
    acc = mx.gluon.metric.Accuracy()
    est = estimator.Estimator(net, loss=ce_loss, train_metrics=acc)
    custom_handler = CustomStopHandler(3, 2)
    est.fit(test_data, event_handlers=[custom_handler], epochs=3)
    assert custom_handler.num_batch == 3
    assert custom_handler.num_epoch == 1
    custom_handler = CustomStopHandler(100, 5)
    est.fit(test_data, event_handlers=[custom_handler], epochs=10)
    assert custom_handler.num_batch == 5 * 4
    assert custom_handler.num_epoch == 5

@mx.util.use_np
def test_logging_interval():
    ''' test different options for logging handler '''
    ''' test case #1: log interval is 1 '''
    batch_size = 8
    data_size = 100
    old_stdout = sys.stdout
    sys.stdout = mystdout = StringIO()
    log_interval = 1
    net = _get_test_network()
    dataloader = _get_test_data(in_size=data_size)
    num_epochs = 1
    ce_loss = loss.SoftmaxCrossEntropyLoss()
    acc = mx.gluon.metric.Accuracy()
    logging = LoggingHandler(metrics=[acc], log_interval=log_interval)
    est = estimator.Estimator(net=net,
                              loss=ce_loss,
                              train_metrics=acc)

    est.fit(train_data=dataloader,
            epochs=num_epochs,
            event_handlers=[logging])

    sys.stdout = old_stdout
    log_info_list = mystdout.getvalue().splitlines()
    info_len = 0
    for info in log_info_list:
        match = re.match(
            r'(\[Epoch \d+\]\[Batch \d+\]\[Samples \d+\] time\/interval: \d+.\d+s' +
            r' training accuracy: \d+.\d+)', info)
        if match:
            info_len += 1

    assert(info_len == int(data_size/batch_size/log_interval) + 1)
    ''' test case #2: log interval is 5 '''
    old_stdout = sys.stdout
    sys.stdout = mystdout = StringIO()
    acc = mx.gluon.metric.Accuracy()
    log_interval = 5
    logging = LoggingHandler(metrics=[acc], log_interval=log_interval)
    est = estimator.Estimator(net=net,
                              loss=ce_loss,
                              train_metrics=acc)
    est.fit(train_data=dataloader,
            epochs=num_epochs,
            event_handlers=[logging])
    sys.stdout = old_stdout
    log_info_list = mystdout.getvalue().splitlines()
    info_len = 0
    for info in log_info_list:
        match = re.match(
            r'(\[Epoch \d+\]\[Batch \d+\]\[Samples \d+\] time\/interval: \d+.\d+s' +
            r' training accuracy: \d+.\d+)', info)
        if match:
            info_len += 1

    assert(info_len == int(data_size/batch_size/log_interval) + 1)

@mx.util.use_np
def test_validation_handler_batch_axis():
    # test case #1: test batch_axis=0
    test_data = _get_test_data()
    net = _get_test_network()
    ce_loss = loss.SoftmaxCrossEntropyLoss()
    acc = mx.gluon.metric.Accuracy()
    est = estimator.Estimator(net, loss=ce_loss, train_metrics=acc)
    est.fit(test_data, epochs=3)

    #test case #2: test batch_axis=1
    test_data = _get_batch_axis_test_data()
    val_data = _get_batch_axis_test_data(in_size=30)
    est = estimator.Estimator(net, loss=ce_loss, train_metrics=acc)
    est.fit(test_data, val_data=val_data,
            epochs=3, batch_axis=1)

@mx.util.use_np
def test_validation_handler():
    test_data = _get_test_data()

    net = _get_test_network()
    ce_loss = loss.SoftmaxCrossEntropyLoss()
    acc = mx.gluon.metric.Accuracy()
    est = estimator.Estimator(net, loss=ce_loss, train_metrics=acc)
    val_handler = ValidationHandler(val_data=test_data,
                                    eval_fn=est.evaluate,
                                    event_handlers=Handler())

    est.fit(train_data=test_data, val_data=test_data,
            event_handlers=[val_handler], epochs=2)
    assert est.run_test_handler == True