import copy
import tempfile
from collections import OrderedDict

import pytest
import torch
import torch.nn as nn
from mmcv import Config
from torch.utils.data import Dataset

from mmaction.apis import train_model
from mmaction.datasets import DATASETS


@DATASETS.register_module()
class ExampleDataset(Dataset):

    def __init__(self, test_mode=False):
        self.test_mode = test_mode

    @staticmethod
    def evaluate(results, logger=None):
        eval_results = OrderedDict()
        eval_results['acc'] = 1
        return eval_results

    def __getitem__(self, idx):
        results = dict(imgs=torch.tensor([1]))
        return results

    def __len__(self):
        return 1


class ExampleModel(nn.Module):

    def __init__(self):
        super().__init__()
        self.test_cfg = None
        self.conv1 = nn.Conv2d(3, 8, kernel_size=1)
        self.norm1 = nn.BatchNorm1d(2)

    def forward(self, imgs, return_loss=False):
        self.norm1(torch.rand(3, 2).cuda())
        losses = dict()
        losses['test_loss'] = torch.tensor([0.5], requires_grad=True)
        return losses

    def train_step(self, data_batch, optimizer, **kwargs):
        imgs = data_batch['imgs']
        losses = self.forward(imgs, True)
        loss = torch.tensor([0.5], requires_grad=True)
        outputs = dict(loss=loss, log_vars=losses, num_samples=3)
        return outputs

    def val_step(self, data_batch, optimizer, **kwargs):
        imgs = data_batch['imgs']
        self.forward(imgs, False)
        outputs = dict(results=0.5)
        return outputs


@pytest.mark.skipif(
    not torch.cuda.is_available(), reason='requires CUDA support')
def test_train_model():
    model = ExampleModel()
    dataset = ExampleDataset()
    datasets = [ExampleDataset(), ExampleDataset()]
    _cfg = dict(
        seed=0,
        gpus=1,
        gpu_ids=[0],
        resume_from=None,
        load_from=None,
        workflow=[('train', 1)],
        total_epochs=5,
        evaluation=dict(interval=1, save_best='acc'),
        data=dict(
            videos_per_gpu=1,
            workers_per_gpu=0,
            val=dict(type='ExampleDataset')),
        optimizer=dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0001),
        optimizer_config=dict(grad_clip=dict(max_norm=40, norm_type=2)),
        lr_config=dict(policy='step', step=[40, 80]),
        omnisource=False,
        precise_bn=False,
        checkpoint_config=dict(interval=1),
        log_level='INFO',
        log_config=dict(interval=20, hooks=[dict(type='TextLoggerHook')]))

    with tempfile.TemporaryDirectory() as tmpdir:
        # normal train
        cfg = copy.deepcopy(_cfg)
        cfg['work_dir'] = tmpdir
        config = Config(cfg)
        train_model(model, dataset, config)

    with tempfile.TemporaryDirectory() as tmpdir:
        # train with validation
        cfg = copy.deepcopy(_cfg)
        cfg['work_dir'] = tmpdir
        config = Config(cfg)
        train_model(model, dataset, config, validate=True)

    with tempfile.TemporaryDirectory() as tmpdir:
        # train with Fp16OptimizerHook
        cfg = copy.deepcopy(_cfg)
        cfg['work_dir'] = tmpdir
        cfg['fp16'] = dict(loss_scale=512.)
        config = Config(cfg)
        model.fp16_enabled = None
        train_model(model, dataset, config)

    with tempfile.TemporaryDirectory() as tmpdir:
        cfg = copy.deepcopy(_cfg)
        cfg['work_dir'] = tmpdir
        cfg['omnisource'] = True
        config = Config(cfg)
        train_model(model, datasets, config)

    with tempfile.TemporaryDirectory() as tmpdir:
        # train with precise_bn on
        cfg = copy.deepcopy(_cfg)
        cfg['work_dir'] = tmpdir
        cfg['workflow'] = [('train', 1), ('val', 1)]
        cfg['data'] = dict(
            videos_per_gpu=1,
            workers_per_gpu=0,
            train=dict(type='ExampleDataset'),
            val=dict(type='ExampleDataset'))
        cfg['precise_bn'] = dict(num_iters=1, interval=1)
        config = Config(cfg)
        train_model(model, datasets, config)
