from __future__ import print_function, division
import torch
import matplotlib.pyplot as plt
import warnings
import random
import numpy as np
warnings.filterwarnings("ignore")
plt.ion()


#Initiate parameters in model 
def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv2d') != -1:
        torch.nn.init.xavier_uniform_(m.weight.data)
        torch.nn.init.constant_(m.bias.data, 0.0)
    elif classname.find('ConvTranspose2d') != -1:
        torch.nn.init.xavier_uniform_(m.weight.data)
        torch.nn.init.constant_(m.bias.data, 0.0)
    if classname.find('Conv3d') != -1:
        torch.nn.init.xavier_uniform_(m.weight.data)
        torch.nn.init.constant_(m.bias.data, 0.0)
    elif classname.find('ConvTranspose3d') != -1:
        torch.nn.init.xavier_uniform_(m.weight.data)
        torch.nn.init.constant_(m.bias.data, 0.0)
    elif classname.find('Linear') != -1:
        torch.nn.init.xavier_uniform_(m.weight.data)
        torch.nn.init.constant_(m.bias.data, 0.0)

            


def SetupSeed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.deterministic = True

def ExtractPatch(lf, H, W, patchSize):
    indx=random.randrange(0,H-patchSize,8)
    indy=random.randrange(0,W-patchSize,8)
    indc=random.randint(0,2)

    lfPatch=lf[:,:,indc:indc+1,
                   indx:indx+patchSize,
                   indy:indy+patchSize]
    return lfPatch #[u v c x y] 
    



    
