import datetime
import os
import sys
import time

import paddle
from paddle import nn
import paddlevision

import presets
import utils

import numpy as np
import random

apex = None

import numpy as np
from reprod_log import ReprodLogger


def train_one_epoch(
        model,
        criterion,
        optimizer,
        data_loader,
        device,
        epoch,
        print_freq, ):
    model.train()
    # training log
    train_reader_cost = 0.0
    train_run_cost = 0.0
    total_samples = 0
    acc1 = 0.0
    acc5 = 0.0
    reader_start = time.time()
    batch_past = 0

    for batch_idx, (image, target) in enumerate(data_loader):
        train_reader_cost += time.time() - reader_start
        train_start = time.time()
        output = model(image)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        optimizer.clear_grad()
        train_run_cost += time.time() - train_start
        acc = utils.accuracy(output, target, topk=(1, 5))
        acc1 += acc[0].item()
        acc5 += acc[1].item()
        total_samples += image.shape[0]
        batch_past += 1

        if batch_idx > 0 and batch_idx % print_freq == 0:
            msg = "[Epoch {}, iter: {}] top1: {:.5f}, top5: {:.5f}, lr: {:.5f}, loss: {:.5f}, avg_reader_cost: {:.5f} sec, avg_batch_cost: {:.5f} sec, avg_samples: {}, avg_ips: {:.5f} images/sec.".format(
                epoch, batch_idx, acc1 / batch_past, acc5 / batch_past,
                optimizer.get_lr(),
                loss.item(), train_reader_cost / batch_past,
                (train_reader_cost + train_run_cost) / batch_past,
                total_samples / batch_past,
                total_samples / (train_reader_cost + train_run_cost))
            if paddle.distributed.get_rank() <= 0:
                print(msg)
                sys.stdout.flush()
            train_reader_cost = 0.0
            train_run_cost = 0.0
            total_samples = 0
            acc1 = 0.0
            acc5 = 0.0
            batch_past = 0

        reader_start = time.time()


def evaluate(model, criterion, data_loader, device, print_freq=100):
    model.eval()
    metric_logger = utils.MetricLogger(delimiter="  ")
    header = 'Test:'
    with paddle.no_grad():
        for image, target in metric_logger.log_every(data_loader, print_freq,
                                                     header):
            output = model(image)
            loss = criterion(output, target)

            acc1, acc5 = utils.accuracy(output, target, topk=(1, 5))
            # FIXME need to take into account that the datasets
            # could have been padded in distributed setup
            batch_size = image.shape[0]
            metric_logger.update(loss=loss.item())
            metric_logger.meters['acc1'].update(acc1.item(), n=batch_size)
            metric_logger.meters['acc5'].update(acc5.item(), n=batch_size)

    # gather the stats from all processes
    metric_logger.synchronize_between_processes()

    print(' * Acc@1 {top1.global_avg:.3f} Acc@5 {top5.global_avg:.3f}'.format(
        top1=metric_logger.acc1, top5=metric_logger.acc5))
    return metric_logger.acc1.global_avg


def load_data(traindir, valdir, args):
    # Data loading code
    print("Loading data")
    resize_size, crop_size = (342, 299) if args.model == 'inception_v3' else (
        256, 224)

    print("Loading training data")
    st = time.time()
    auto_augment_policy = getattr(args, "auto_augment", None)
    random_erase_prob = getattr(args, "random_erase", 0.0)
    dataset = paddlevision.datasets.ImageFolder(
        traindir,
        presets.ClassificationPresetTrain(
            crop_size=crop_size,
            auto_augment_policy=auto_augment_policy,
            random_erase_prob=random_erase_prob))

    print("Took", time.time() - st)

    print("Loading validation data")
    dataset_test = paddlevision.datasets.ImageFolder(
        valdir,
        presets.ClassificationPresetEval(
            crop_size=crop_size, resize_size=resize_size))

    print("Creating data loaders")
    train_sampler = paddle.io.DistributedBatchSampler(
        dataset=dataset,
        batch_size=args.batch_size,
        shuffle=True,
        drop_last=False)

    test_sampler = paddle.io.SequenceSampler(dataset_test)

    return dataset, dataset_test, train_sampler, test_sampler


def main(args):
    if args.output_dir:
        utils.mkdir(args.output_dir)

    print(args)

    device = paddle.set_device(args.device)

    # multi cards
    if paddle.distributed.get_world_size() > 1:
        paddle.distributed.init_parallel_env()

    train_dir = os.path.join(args.data_path, 'train')
    val_dir = os.path.join(args.data_path, 'val')
    dataset, dataset_test, train_sampler, test_sampler = load_data(
        train_dir, val_dir, args)
    train_batch_sampler = train_sampler
    data_loader = paddle.io.DataLoader(
        dataset=dataset,
        num_workers=args.workers,
        return_list=True,
        batch_sampler=train_batch_sampler)
    test_batch_sampler = paddle.io.BatchSampler(
        sampler=test_sampler, batch_size=args.batch_size)