import time
import torch
import torch.distributed as dist
from utils import AverageMeter, calculate_accuracy


def val_epoch(epoch,
              data_loader,
              video_model,
              criterions,
              device,
              logger,
              tb_writer=None,
              distributed=False):
    print('validation at epoch {}'.format(epoch))

    video_model.eval()

    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()
    accuracies = AverageMeter()

    criterion = criterions['cls_criterion']
    end_time = time.time()
    acc_num = 0

    with torch.no_grad():
        for i, batch in enumerate(data_loader):
            inputs = batch['clip']
            targets = batch['target']

            data_time.update(time.time() - end_time)
            inputs = inputs.to(device,non_blocking = True)
            targets = targets.to(device, non_blocking=True)
            
            _,_,outputs = video_model(inputs) 
            loss = criterion(outputs, targets)
            acc = calculate_accuracy(outputs, targets)
            acc_num += acc*targets.shape[0]

            losses.update(loss.item(), inputs.size(0))
            accuracies.update(acc, inputs.size(0))

            batch_time.update(time.time() - end_time)
            end_time = time.time()

            print('Epoch: [{0}][{1}/{2}]\t'
                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                  'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
                  'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                  'Acc {acc.val:.3f} ({acc.avg:.3f})'.
                  format(epoch, i + 1, len(data_loader),
                         batch_time=batch_time,
                         data_time=data_time,
                         loss=losses,
                         acc=accuracies), flush=True)

    if distributed:
        loss_sum = torch.tensor([losses.sum], dtype=torch.float32, device=device)
        loss_count = torch.tensor([losses.count], dtype=torch.float32, device=device)
        acc_sum = torch.tensor([accuracies.sum], dtype=torch.float32, device=device)
        acc_count = torch.tensor([accuracies.count], dtype=torch.float32, device=device)

        dist.all_reduce(loss_sum, op=dist.ReduceOp.SUM)
        dist.all_reduce(loss_count, op=dist.ReduceOp.SUM)
        dist.all_reduce(acc_sum, op=dist.ReduceOp.SUM)
        dist.all_reduce(acc_count, op=dist.ReduceOp.SUM)

        losses.avg = loss_sum.item() / loss_count.item()
        accuracies.avg = acc_sum.item() / acc_count.item()

    if logger is not None:
        logger.log({'epoch': epoch, 'loss': losses.avg, 'acc': accuracies.avg, 'acc_num': acc_num})

    if tb_writer is not None:
        tb_writer.add_scalar('val/loss', losses.avg, epoch)
        tb_writer.add_scalar('val/acc', accuracies.avg, epoch)
        tb_writer.add_scalar('val/acc_num', acc_num, epoch)

    return losses.avg, accuracies.avg
