import numpy as np

def data_loader(density_load_path, displace_load_path):
    # Load datasets
    global_density = np.load(density_load_path)
    global_displace = np.load(displace_load_path)
    global_displace = global_displace.reshape(181,61,2)
    global_displace = np.dstack((global_displace[:,:,0].T, global_displace[:,:,1].T))
    print(global_displace.shape)
    print(global_density.shape)
    
    m=5
    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__':
    data_loader()