import torch
from bottleneck import Bottleneck
torch.manual_seed(23337)

# use True to print layerwise sum for all outputs in reference code path
DEBUG = False#True

for stride, o_channel in [(1,32), (1,128), (2,32)]:
    print("testing stride ==", stride, ", in_channel == 32 , out_channel ==", o_channel)
    a_ = torch.randn(17,32,28,28)

    a = a_.cuda().half().to(memory_format=torch.channels_last).requires_grad_()
    model = Bottleneck(32,8,o_channel,stride=stride).cuda().half().to(memory_format=torch.channels_last)

    # test model
    b = model(a)
    b.mean().backward()
    d_grad = a.grad.float()
    a.grad = None
    torch.cuda.synchronize()

    if DEBUG:
        print("[DEBUG] ref dx :", d_grad.sum().item())
        # print wgrad. we don't need to reset since later cpp print before accumulation
        for i, w in enumerate(model.w_conv):
            print("[DEBUG] ref wgrad{} :".format(i+1), w.grad.sum().item())

    wgrads = []
    for w in model.w_conv:
        wgrads.append(w.grad.float())

    model.use_cudnn = True
    model.zero_grad()
    c = model(a)
    c.mean().backward()

    torch.cuda.synchronize()
    print("comparing native and channels_last:")
    print("max error fprop:", (b-c).abs().max().item(), "max elem:", b.abs().max().item())
    print("max error dgrad:", (d_grad-a.grad.float()).abs().max().item(), "max elem:", d_grad.abs().max().item())
    for i, (w, wgrad) in enumerate(zip(model.w_conv, wgrads)):
        print("max error wgrad{}:".format(i+1), (wgrad - w.grad.float()).abs().max().item(), "max elem:", wgrad.abs().max().item())

    nhwc_a = a_.permute(0,2,3,1).contiguous().cuda().half().requires_grad_()
    nhwc_model = Bottleneck(32,8,o_channel,stride=stride,explicit_nhwc=True, use_cudnn=True).cuda().half()
    for p,q in zip(model.parameters(), nhwc_model.parameters()):
        # model's storage is already in nhwc, we clone and assign to explicit nhwc model
        q.data.copy_(p.data.permute(0,2,3,1).contiguous())
    for p,q in zip(model.buffers(), nhwc_model.buffers()):
        q.data.copy_(p.data)

    d = nhwc_model(nhwc_a)
    d.mean().backward()
    torch.cuda.synchronize()

    # reset reference to cudnn channels_last permute
    #c_s = c.storage().tolist()
    #d_s = d.storage().tolist()
    #print(max([x-y for x,y in zip(c_s,d_s)]))
    c = c.contiguous(memory_format=torch.contiguous_format).permute(0,2,3,1).contiguous()
    d_grad = a.grad.float().permute(0,2,3,1).contiguous()
    wgrads = []
    for w in model.w_conv:
        wgrads.append(w.grad.float().permute(0,2,3,1).contiguous())

    torch.cuda.synchronize()
    print("comparing nhwc and channels_last:")
    print("max error fprop:", (d-c).abs().max().item(), "max elem:", c.abs().max().item())
    print("max error dgrad:", (d_grad-nhwc_a.grad.float()).abs().max().item(), "max elem:", d_grad.abs().max().item())
    for i, (w, wgrad) in enumerate(zip(nhwc_model.w_conv, wgrads)):
        print("max error wgrad{}:".format(i+1), (wgrad - w.grad.float()).abs().max().item(), "max elem:", wgrad.abs().max().item())