# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed 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 time
from collections import deque, defaultdict
import pickle
import shutil

import numpy as np
import paddle
import paddle.nn.functional as F
from paddleseg.utils import TimeAverager, calculate_eta, resume, logger, train_profiler

from .val import evaluate


def visual_in_traning(log_writer, vis_dict, step):
    """
    Visual in vdl

    Args:
        log_writer (LogWriter): The log writer of vdl.
        vis_dict (dict): Dict of tensor. The shape of thesor is (C, H, W)
    """
    for key, value in vis_dict.items():
        value_shape = value.shape
        if value_shape[0] not in [1, 3]:
            value = value[0]
            value = value.unsqueeze(0)
        value = paddle.transpose(value, (1, 2, 0))
        min_v = paddle.min(value)
        max_v = paddle.max(value)
        if (min_v > 0) and (max_v < 1):
            value = value * 255
        elif (min_v < 0 and min_v >= -1) and (max_v <= 1):
            value = (1 + value) / 2 * 255
        else:
            value = (value - min_v) / (max_v - min_v) * 255

        value = value.astype('uint8')
        value = value.numpy()
        log_writer.add_image(tag=key, img=value, step=step)


def save_best(best_model_dir, metrics_data, iter):
    with open(os.path.join(best_model_dir, 'best_metrics.txt'), 'w') as f:
        for key, value in metrics_data.items():
            line = key + ' ' + str(value) + '\n'
            f.write(line)
        f.write('iter' + ' ' + str(iter) + '\n')


def get_best(best_file, metrics, resume_model=None):
    '''Get best metrics and iter from file'''
    best_metrics_data = {}
    if os.path.exists(best_file) and (resume_model is not None):
        values = []
        with open(best_file, 'r') as f:
            lines = f.readlines()
            for line in lines:
                line = line.strip()
                key, value = line.split(' ')
                best_metrics_data[key] = eval(value)
                if key == 'iter':
                    best_iter = eval(value)
    else:
        for key in metrics:
            best_metrics_data[key] = np.inf
        best_iter = -1
    return best_metrics_data, best_iter


def train(model,
          train_dataset,
          val_dataset=None,
          optimizer=None,
          save_dir='output',
          iters=10000,
          batch_size=2,
          resume_model=None,
          save_interval=1000,
          log_iters=10,
          log_image_iters=1000,
          num_workers=0,
          use_vdl=False,
          losses=None,
          keep_checkpoint_max=5,
          eval_begin_iters=None,
          metrics='sad',
          precision='fp32',
          amp_level='O1',
          profiler_options=None):
    """
    Launch training.
    Args:
        model（nn.Layer): A matting model.
        train_dataset (paddle.io.Dataset): Used to read and process training datasets.
        val_dataset (paddle.io.Dataset, optional): Used to read and process validation datasets.
        optimizer (paddle.optimizer.Optimizer): The optimizer.
        save_dir (str, optional): The directory for saving the model snapshot. Default: 'output'.
        iters (int, optional): How may iters to train the model. Defualt: 10000.
        batch_size (int, optional): Mini batch size of one gpu or cpu. Default: 2.
        resume_model (str, optional): The path of resume model.
        save_interval (int, optional): How many iters to save a model snapshot once during training. Default: 1000.
        log_iters (int, optional): Display logging information at every log_iters. Default: 10.
        log_image_iters (int, optional): Log image to vdl. Default: 1000.
        num_workers (int, optional): Num workers for data loader. Default: 0.
        use_vdl (bool, optional): Whether to record the data to VisualDL during training. Default: False.
        losses (dict, optional): A dict of loss, refer to the loss function of the model for details. Default: None.
        keep_checkpoint_max (int, optional): Maximum number of checkpoints to save. Default: 5.
        eval_begin_iters (int): The iters begin evaluation. It will evaluate at iters/2 if it is None. Defalust: None.
        metrics(str|list, optional): The metrics to evaluate, it may be the combination of ("sad", "mse", "grad", "conn"). 
        precision (str, optional): Use AMP if precision='fp16'. If precision='fp32', the training is normal.
        amp_level (str, optional): Auto mixed precision level. Accepted values are “O1” and “O2”: O1 represent mixed precision, 
            the input data type of each operator will be casted by white_list and black_list; O2 represent Pure fp16, all operators 
            parameters and input data will be casted to fp16, except operators in black_list, don’t support fp16 kernel and batchnorm. Default is O1(amp)
        profiler_options (str, optional): The option of train profiler.
    """
    model.train()
    nranks = paddle.distributed.ParallelEnv().nranks
    local_rank = paddle.distributed.ParallelEnv().local_rank

    start_iter = 0
    if resume_model is not None:
        start_iter = resume(model, optimizer, resume_model)

    if not os.path.isdir(save_dir):
        if os.path.exists(save_dir):
            os.remove(save_dir)
        os.makedirs(save_dir)

    # Use amp
    if precision == 'fp16':
        logger.info('use AMP to train. AMP level = {}'.format(amp_level))
        scaler = paddle.amp.GradScaler(init_loss_scaling=1024)
        if amp_level == 'O2':
            model, optimizer = paddle.amp.decorate(
                models=model,
                optimizers=optimizer,
                level='O2',
                save_dtype='float32')

    if nranks > 1:
        # Initialize parallel environment if not done.
        if not paddle.distributed.parallel.parallel_helper._is_parallel_ctx_initialized(
        ):
            paddle.distributed.init_parallel_env()
            ddp_model = paddle.DataParallel(model)
        else:
            ddp_model = paddle.DataParallel(model)

    batch_sampler = paddle.io.DistributedBatchSampler(
        train_dataset, batch_size=batch_size, shuffle=True, drop_last=True)

    loader = paddle.io.DataLoader(
        train_dataset,
        batch_sampler=batch_sampler,
        num_workers=num_workers,
        return_list=True, )

    if use_vdl:
        from visualdl import LogWriter
        log_writer = LogWriter(save_dir)

    if isinstance(metrics, str):
        metrics = [metrics]
    elif not isinstance(metrics, list):
        metrics = ['sad']
    best_metrics_data, best_iter = get_best(
        os.path.join(save_dir, 'best_model', 'best_metrics.txt'),
        metrics,
        resume_model=resume_model)
    avg_loss = defaultdict(float)
    iters_per_epoch = len(batch_sampler)
    reader_cost_averager = TimeAverager()
    batch_cost_averager = TimeAverager()
    save_models = deque()
    batch_start = time.time()

    iter = start_iter
    while iter < iters:
        for data in loader:
            iter += 1
            if iter > iters:
                break
            reader_cost_averager.record(time.time() - batch_start)

            if precision == 'fp16':
                with paddle.amp.auto_cast(
                        level=amp_level,
                        enable=True,
                        custom_white_list={
                            "elementwise_add", "batch_norm", "sync_batch_norm"
                        },
                        custom_black_list={'bilinear_interp_v2', 'pad3d'}):
                    logit_dict, loss_dict = ddp_model(
                        data) if nranks > 1 else model(data)

                scaled = scaler.scale(loss_dict['all'])  # scale the loss
                scaled.backward()  # do backward
                scaler.minimize(optimizer, scaled)  # update parameters
            else:
                logit_dict, loss_dict = ddp_model(
                    data) if nranks > 1 else model(data)
                loss_dict['all'].backward()
                optimizer.step()

            lr = optimizer.get_lr()
            if isinstance(optimizer._learning_rate,
                          paddle.optimizer.lr.LRScheduler):
                optimizer._learning_rate.step()

            train_profiler.add_profiler_step(profiler_options)

            model.clear_gradients()

            for key, value in loss_dict.items():
                avg_loss[key] += float(value)
            batch_cost_averager.record(
                time.time() - batch_start, num_samples=batch_size)

            if (iter) % log_iters == 0 and local_rank == 0:
                for key, value in avg_loss.items():
                    avg_loss[key] = value / log_iters
                remain_iters = iters - iter
                avg_train_batch_cost = batch_cost_averager.get_average()
                avg_train_reader_cost = reader_cost_averager.get_average()
                eta = calculate_eta(remain_iters, avg_train_batch_cost)
                # loss info
                loss_str = ' ' * 26 + '\t[LOSSES]'
                loss_str = loss_str
                for key, value in avg_loss.items():
                    if key != 'all':
                        loss_str = loss_str + ' ' + key + '={:.4f}'.format(
                            value)
                logger.info(
                    "[TRAIN] epoch={}, iter={}/{}, loss={:.4f}, lr={:.6f}, batch_cost={:.4f}, reader_cost={:.5f}, ips={:.4f} samples/sec | ETA {}\n{}\n"
                    .format((iter - 1) // iters_per_epoch + 1, iter, iters,
                            avg_loss['all'], lr, avg_train_batch_cost,
                            avg_train_reader_cost,
                            batch_cost_averager.get_ips_average(
                            ), eta, loss_str))
                if use_vdl:
                    for key, value in avg_loss.items():
                        log_tag = 'Train/' + key
                        log_writer.add_scalar(log_tag, value, iter)

                    log_writer.add_scalar('Train/lr', lr, iter)
                    log_writer.add_scalar('Train/batch_cost',
                                          avg_train_batch_cost, iter)
                    log_writer.add_scalar('Train/reader_cost',
                                          avg_train_reader_cost, iter)
                    if iter % log_image_iters == 0:
                        vis_dict = {}
                        # ground truth
                        vis_dict['ground truth/img'] = data['img'][0]
                        for key in data['gt_fields']:
                            key = key[0]
                            vis_dict['/'.join(['ground truth', key])] = data[
                                key][0]
                        # predict
                        for key, value in logit_dict.items():
                            vis_dict['/'.join(['predict', key])] = logit_dict[
                                key][0]
                        visual_in_traning(
                            log_writer=log_writer, vis_dict=vis_dict, step=iter)

                for key in avg_loss.keys():
                    avg_loss[key] = 0.
                reader_cost_averager.reset()
                batch_cost_averager.reset()

            # save model
            if (iter % save_interval == 0 or iter == iters) and local_rank == 0:
                current_save_dir = os.path.join(save_dir,
                                                "iter_{}".format(iter))
                if not os.path.isdir(current_save_dir):
                    os.makedirs(current_save_dir)
                paddle.save(model.state_dict(),
                            os.path.join(current_save_dir, 'model.pdparams'))
                paddle.save(optimizer.state_dict(),
                            os.path.join(current_save_dir, 'model.pdopt'))
                save_models.append(current_save_dir)
                if len(save_models) > keep_checkpoint_max > 0:
                    model_to_remove = save_models.popleft()
                    shutil.rmtree(model_to_remove)

            # eval model
            if eval_begin_iters is None:
                eval_begin_iters = iters // 2
            if (iter % save_interval == 0 or iter == iters) and (
                    val_dataset is not None
            ) and local_rank == 0 and iter >= eval_begin_iters:
                num_workers = 1 if num_workers > 0 else 0
                metrics_data = evaluate(
                    model,
                    val_dataset,
                    num_workers=1,
                    print_detail=True,
                    save_results=False,
                    metrics=metrics,
                    precision=precision,
                    amp_level=amp_level)
                model.train()

            # save best model and add evaluation results to vdl
            if (iter % save_interval == 0 or iter == iters) and local_rank == 0:
                if val_dataset is not None and iter >= eval_begin_iters:
                    if metrics_data[metrics[0]] < best_metrics_data[metrics[0]]:
                        best_iter = iter
                        best_metrics_data = metrics_data.copy()
                        best_model_dir = os.path.join(save_dir, "best_model")
                        paddle.save(
                            model.state_dict(),
                            os.path.join(best_model_dir, 'model.pdparams'))
                        save_best(best_model_dir, best_metrics_data, iter)

                    show_list = []
                    for key, value in best_metrics_data.items():
                        show_list.append((key, value))
                    log_str = '[EVAL] The model with the best validation {} ({:.4f}) was saved at iter {}.'.format(
                        show_list[0][0], show_list[0][1], best_iter)
                    if len(show_list) > 1:
                        log_str += " While"
                        for i in range(1, len(show_list)):
                            log_str = log_str + ' {}: {:.4f},'.format(
                                show_list[i][0], show_list[i][1])
                        log_str = log_str[:-1]
                    logger.info(log_str)

                    if use_vdl:
                        for key, value in metrics_data.items():
                            log_writer.add_scalar('Evaluate/' + key, value,
                                                  iter)

            batch_start = time.time()

    # Sleep for half a second to let dataloader release resources.
    time.sleep(0.5)
    if use_vdl:
        log_writer.close()