import os
import time
import numpy as np
import megengine.distributed as dist
import megengine as mge
import megengine.functional as F
from megengine.autodiff import GradManager
from edit.core.hook.evaluation import psnr, ssim
from edit.utils import imwrite, tensor2img, bgr2ycbcr, img_multi_padding, img_de_multi_padding, ensemble_forward, ensemble_back
from edit.utils import img_multi_padding, img_de_multi_padding, flow_to_image
from ..base import BaseModel
from ..builder import build_backbone, build_loss
from ..registry import MODELS
from tqdm import tqdm

def get_bilinear(image):
    B,T,C,h,w = image.shape
    image = image.reshape(-1, C,h,w)
    return F.nn.interpolate(image, scale_factor=4).reshape(B,T,C,4*h, 4*w)

def train_generator_batch(image, label, *, gm, netG, netloss):
    B,T,_,h,w = image.shape
    biup = get_bilinear(image)
    netG.train()
    with gm:
        forward_hiddens = []
        backward_hiddens = []
        res = []
        hidden = F.zeros((2*B, netG.hidden_channels, h, w))
        for i in range(T):
            now_frame = F.concat([image[:, i, ...], image[:, T-i-1, ...]], axis=0)
            if i==0:
                flow = netG.flownet(now_frame, now_frame)
            else:
                ref = F.concat([image[:, i-1, ...], image[:, T-i, ...]], axis=0)
                flow = netG.flownet(now_frame, ref)
            hidden = netG(hidden, flow, now_frame)
            forward_hiddens.append(hidden[0:B, ...])
            backward_hiddens.append(hidden[B:2*B, ...])
        for i in range(T):
            res.append(netG.do_upsample(forward_hiddens[i], backward_hiddens[T-i-1]))
        res = F.stack(res, axis = 1) # [B,T,3,H,W]
        loss = netloss(res+biup, label)
        gm.backward(loss)
        if dist.is_distributed():
            loss = dist.functional.all_reduce_sum(loss) / dist.get_world_size()
    return loss

def test_generator_batch(image, *, netG):
    # image: [1,100,3,180,320]
    B,T,_,h,w = image.shape
    biup = get_bilinear(image)
    netG.eval()
    forward_hiddens = []
    backward_hiddens = []
    res = []
    hidden = F.zeros((2*B, netG.hidden_channels, h, w))
    for i in range(T):
        now_frame = F.concat([image[:, i, ...], image[:, T-i-1, ...]], axis=0)
        if i==0:
            flow = netG.flownet(now_frame, now_frame)
        else:
            ref = F.concat([image[:, i-1, ...], image[:, T-i, ...]], axis=0)
            flow = netG.flownet(now_frame, ref)
        hidden = netG(hidden, flow, now_frame)
        forward_hiddens.append(hidden[0:B, ...])
        backward_hiddens.append(hidden[B:2*B, ...])
    for i in range(T):
        res.append(netG.do_upsample(forward_hiddens[i], backward_hiddens[T-i-1]))
    res = F.stack(res, axis = 1) # [B,T,3,H,W]
    return res + biup

epoch_dict = {}

def adjust_learning_rate(optimizer, epoch):
    if epoch>=8 and epoch % 1 == 0  and epoch_dict.get(epoch, None) is None:
        epoch_dict[epoch] = True
        for param_group in optimizer.param_groups:
            param_group["lr"] = param_group["lr"] * 0.8
        print("adjust lr! , now lr: {}".format(param_group["lr"]))

# TODO 可以再写一个父类，抽象一些公共方法，当大于1个模型时，代码重复了，如getimdid和test step
@MODELS.register_module()
class BidirectionalRestorer(BaseModel):
    allowed_metrics = {'PSNR': psnr, 'SSIM': ssim}

    def __init__(self, generator, pixel_loss, train_cfg=None, eval_cfg=None, pretrained=None, Fidelity_loss=None):
        super(BidirectionalRestorer, self).__init__()

        self.train_cfg = train_cfg
        self.eval_cfg = eval_cfg
        # generator
        self.generator = build_backbone(generator)
        # loss
        self.pixel_loss = build_loss(pixel_loss)

        if Fidelity_loss:
            self.Fidelity_loss = build_loss(Fidelity_loss)
        else:
            self.Fidelity_loss = None

        # load pretrained
        self.init_weights(pretrained)

    def init_weights(self, pretrained=None):
        self.generator.init_weights(pretrained)

    def train_step(self, batchdata, now_epoch, now_iter):
        LR_tensor = mge.tensor(batchdata['lq'], dtype="float32")
        HR_tensor = 
mge.tensor(batchdata['gt'], dtype="float32")