# encoding: utf-8
"""
@author: zeming li
@contact: zengarden2009@gmail.com
"""
import torch
import math
from torch.nn import Module
import resnet_mbn as resnet


class M2Teacher(Module):
    def __init__(self, cfg):
        super(M2Teacher, self).__init__()
        self.device = torch.device(cfg.MODEL.DEVICE)

        self.param_momentum = cfg.MODEL.M2T.PARAM_MOMENTUM
        self.total_iters = cfg.SOLVER.LR_SCHEDULER.MAX_ITER * cfg.SOLVER.BATCH_SUBDIVISIONS
        self.current_train_iter = 0

        self.student_encoder = resnet.resnet50(
            low_dim=256, width=1, hidden_dim=4096, MLP="byol", CLS=False, bn="customized", predictor=True
        )
        self.teacher_encoder = resnet.resnet50(
            low_dim=256, width=1, hidden_dim=4096, MLP="byol", CLS=False, bn="mbn", predictor=False
        )
        for p in self.teacher_encoder.parameters():
            p.requires_grad = False

        self.momentum_update(m=0)
        for m in self.teacher_encoder.modules():
            if isinstance(m, resnet.MomentumBatchNorm1d) or isinstance(m, resnet.MomentumBatchNorm2d):
                m.total_iters = self.total_iters

        self.to(self.device)

    @torch.no_grad()
    def momentum_update(self, m):
        for p1, p2 in zip(self.student_encoder.parameters(), self.teacher_encoder.parameters()):
            # p2.data.mul_(m).add_(1 - m, p1.detach().data)
            p2.data = m * p2.data + (1.0 - m) * p1.detach().data

    def get_param_momentum(self):
        return 1.0 - (1.0 - self.param_momentum) * (
            (math.cos(math.pi * self.current_train_iter / self.total_iters) + 1) * 0.5
        )

    def forward(self, batched_inputs, update_param=True):
        if update_param:
            current_param_momentum = self.get_param_momentum()
            self.momentum_update(current_param_momentum)

        x1 = torch.stack([bi['q']["image"] for bi in batched_inputs]).to(self.device)
        x2 = torch.stack([bi['k']["image"] for bi in batched_inputs]).to(self.device)

        q1 = self.student_encoder(x1)
        q2 = self.student_encoder(x2)

        with torch.no_grad():
            k1 = self.teacher_encoder(x2)
            k2 = self.teacher_encoder(x1)
        con_loss = (4 - 2 * ((q1 * k1).sum(dim=-1, keepdim=True) + (q2 * k2).sum(dim=-1, keepdim=True))).mean()

        self.current_train_iter += 1
        if self.training:
            return {"loss": con_loss}

    def preprocess_image(self, batched_inputs):
        """
        Normalize, pad and batch the input images.
        """
        images = [x.float().to(self.device) for x in batched_inputs]
        images = [self.normalizer(x) for x in images]
        images = ImageList.from_tensors(images, self.size_divisibility)

        return images

