import numpy as np
import os

def data_loader(opt):
    # Load datasets
    
    # './datasets/train/180_60/u_OR_xPhys/mod1.npy' as default
    density_load_path = os.path.join(opt.dataroot, opt.phase, str(opt.nelx)+'_'+str(opt.nely), 'xPhys', opt.mod+'.npy')
    displace_load_path = os.path.join(opt.dataroot, opt.phase, str(opt.nelx)+'_'+str(opt.nely), 'u', opt.mod+'.npy')

    global_density = np.load(density_load_path) # (nely , nelx)
    global_displace = np.load(displace_load_path) # ( (nely+1)*(nelx+1)*2 , 1 )
    
    global_displace = global_displace.reshape(opt.nelx+1, opt.nely+1, 2)
    global_displace = np.dstack((global_displace[:,:,0].T, global_displace[:,:,1].T)) # -> ( (nelx+1), (nelx+1), 2 )
    
    m=opt.ms_ratio
    N=(m+1)**2
    global_nely=global_density.shape[0]
    global_nelx=global_density.shape[1]
    coarse_nely = int(global_nely/m)
    coarse_nelx = int(global_nelx/m)
    
    # Generate coarse mesh density
    coarse_density=np.zeros(shape=(coarse_nely*coarse_nelx,m*m))
    for ely in range(coarse_nely):
        for elx in range(coarse_nelx):
            coarse_density[elx + ely * m] = global_density[ely * m : (ely + 1) * m, elx * m : (elx + 1) * m].flatten()
    print(coarse_density.shape)
    
    # Generate coarse mesh displacement
    coarse_displace=np.zeros(shape=(coarse_nely*coarse_nelx,4,2))
    for ely in range(coarse_nely):
        for elx in range(coarse_nelx):
            coarse_displace[elx + ely * m][0] = global_displace[ely * m, elx * m, :]
            coarse_displace[elx + ely * m][1] = global_displace[ely * m, (elx+1) * m, :]
            coarse_displace[elx + ely * m][2] = global_displace[(ely+1) * m, elx * m, :]
            coarse_displace[elx + ely * m][3] = global_displace[(ely+1) * m, (elx+1) * m, :]
    print(coarse_displace.shape)
    
    # Generate fine mesh displacement
    fine_displace=np.zeros(shape=(coarse_nely*coarse_nelx, ((m+1)**2) * 2))
    for ely in range(coarse_nely):
        for elx in range(coarse_nelx):
            fine_displace[elx + ely * m] = global_displace[ely*m : (ely+1)*m+1, elx*m : (elx+1)*m+1, :].flatten()
    print(fine_displace.shape)
    
        
    return global_density, global_displace, coarse_density, coarse_displace, fine_displace
    
if __name__=='__main__':
    from options.train_options import TrainOptions
    opt = TrainOptions().parse()
    data_loader(opt)