import_data.py

import numpy as np
import pandas as pd
import random



## all samples MUST include at least one view.
def import_incomplete_handwritten():
    npz = np.load('./data/Handwritten_Missing/data_with_missingviews.npz', allow_pickle=True)

    X_set    = npz['X_set'].tolist()
    Y_onehot = npz['Y_onehot']

    M        = len(X_set)

    ### Construct Mask Vector to indicate available (m=1) or missing (m=0) values
    Mask     = np.ones([np.shape(X_set[0])[0], M])
    for m_idx in range(M):
        Mask[np.isnan(X_set[m_idx]).all(axis=1), m_idx] = 0
        X_set[m_idx][Mask[:, m_idx] == 0] = np.mean(X_set[m_idx][Mask[:, m_idx] == 1], axis=0)
    
    return X_set, Y_onehot, Mask



def import_dataset_TCGA(year=1):
    filename = '/media/vdslab/Genomics/TCGA/dataset/FINAL/cleaned/incomplete_multi_view_pca_{}yr.npz'.format(int(year))
    npz      = np.load(filename)

    Mask  = npz['m']
    M     = np.shape(Mask)[1]

    X_set = {}
    for m in range(M):
        tmp = npz['x{}'.format(m+1)]
        tmp[np.isnan(tmp[:, 0]), :] = np.nanmean(tmp, axis=0)
        X_set[m] = tmp

    Y     = npz['y']

    X_set_incomp = {}
    X_set_comp   = {}
    for m in range(M):
        X_set_comp[m]    = X_set[m][np.sum(Mask, axis=1) == 4]
        X_set_incomp[m]  = X_set[m][np.sum(Mask, axis=1) != 4]

    Y_comp    = Y[np.sum(Mask, axis=1) == 4]
    Y_incomp  = Y[np.sum(Mask, axis=1) != 4]

    Mask_comp    = Mask[np.sum(Mask, axis=1) == 4]
    Mask_incomp  = Mask[np.sum(Mask, axis=1) != 4]


    Y_onehot_incomp = np.zeros([np.shape(Y_incomp)[0], 2])
    Y_onehot_comp   = np.zeros([np.shape(Y_comp)[0], 2])

    Y_onehot_incomp[Y_incomp == 0, 0] = 1
    Y_onehot_incomp[Y_incomp == 1, 1] = 1

    Y_onehot_comp[Y_comp == 0, 0] = 1
    Y_onehot_comp[Y_comp == 1, 1] = 1
    
    return X_set_comp, Y_onehot_comp, Mask_comp, X_set_incomp, Y_onehot_incomp, Mask_incomp