import math
import torch
import torch.nn as nn
from torch.optim import SGD
from torch.optim.lr_scheduler import LambdaLR

from ding.policy import Policy
from ding.model import model_wrap
from ding.torch_utils import to_device
from ding.utils import EasyTimer


class ImageClassificationPolicy(Policy):
    config = dict(
        type='image_classification',
        on_policy=False,
    )

    def _init_learn(self):
        self._optimizer = SGD(
            self._model.parameters(),
            lr=self._cfg.learn.learning_rate,
            weight_decay=self._cfg.learn.weight_decay,
            momentum=0.9
        )
        self._timer = EasyTimer(cuda=True)

        def lr_scheduler_fn(epoch):
            if epoch <= self._cfg.learn.warmup_epoch:
                return self._cfg.learn.warmup_lr / self._cfg.learn.learning_rate
            else:
                ratio = epoch // self._cfg.learn.decay_epoch
                return math.pow(self._cfg.learn.decay_rate, ratio)

        self._lr_scheduler = LambdaLR(self._optimizer, lr_scheduler_fn)
        self._lr_scheduler.step()
        self._learn_model = model_wrap(self._model, 'base')
        self._learn_model.reset()

        self._ce_loss = nn.CrossEntropyLoss()

    def _forward_learn(self, data):
        if self._cuda:
            data = to_device(data, self._device)
        self._learn_model.train()

        with self._timer:
            img, target = data
            logit = self._learn_model.forward(img)
            loss = self._ce_loss(logit, target)
        forward_time = self._timer.value

        with self._timer:
            self._optimizer.zero_grad()
            loss.backward()
        backward_time = self._timer.value

        with self._timer:
            if self._cfg.multi_gpu:
                self.sync_gradients(self._learn_model)
        sync_time = self._timer.value
        self._optimizer.step()

        cur_lr = [param_group['lr'] for param_group in self._optimizer.param_groups]
        cur_lr = sum(cur_lr) / len(cur_lr)
        return {
            'cur_lr': cur_lr,
            'total_loss': loss.item(),
            'forward_time': forward_time,
            'backward_time': backward_time,
            'sync_time': sync_time,
        }

    def _monitor_vars_learn(self):
        return ['cur_lr', 'total_loss', 'forward_time', 'backward_time', 'sync_time']

    def _init_eval(self):
        self._eval_model = model_wrap(self._model, 'base')

    def _forward_eval(self, data):
        if self._cuda:
            data = to_device(data, self._device)
        self._eval_model.eval()
        with torch.no_grad():
            output = self._eval_model.forward(data)
        if self._cuda:
            output = to_device(output, 'cpu')
        return output

    def _init_collect(self):
        pass

    def _forward_collect(self, data):
        pass

    def _process_transition(self):
        pass

    def _get_train_sample(self):
        pass