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 = 