import os
import numpy as np
import scipy.io as sio
import tifffile

from sklearn.decomposition import PCA
from sklearn.model_selection import train_test_split


#Load dataset
def loadData(name,data_path):
    if name == 'IP':
        data = sio.loadmat(os.path.join(data_path, 'Indian_pines_corrected.mat'))['indian_pines_corrected']
        labels = sio.loadmat(os.path.join(data_path, 'Indian_pines_gt.mat'))['indian_pines_gt']
    elif name == 'SA':
        data = sio.loadmat(os.path.join(data_path, 'Salinas_corrected.mat'))['salinas_corrected']
        labels = sio.loadmat(os.path.join(data_path, 'Salinas_gt.mat'))['salinas_gt']
    elif name == 'PU':
        data = sio.loadmat(os.path.join(data_path, 'PaviaU.mat'))['paviaU']
        labels = sio.loadmat(os.path.join(data_path, 'PaviaU_gt.mat'))['paviaU_gt']
    elif name == 'HU13':
        # dict_keys(['__header__', '__version__', '__globals__', 'Houston'])
        #dict_values([b'MATLAB 5.0 MAT-file, Platform: PCWIN64, Created on: Wed Jul 17 16:45:01 2019', '1.0', [], array()])
        #data = sio.loadmat(os.path.join(data_path, 'Houston.mat'))
        #labels = sio.loadmat(os.path.join(data_path,'Houston_gt.mat'))
        data = sio.loadmat(os.path.join(data_path, 'Houston.mat'))['Houston']
        labels = sio.loadmat(os.path.join(data_path,'Houston_gt.mat'))['Houston_gt']
    elif name == 'KSC':
        data = sio.loadmat(os.path.join(data_path, 'KSC.mat'))['KSC']
        labels = sio.loadmat(os.path.join(data_path,'KSC_gt.mat'))['KSC_gt']

    return data, labels



# Use tifffile pkg read the hyperspectral img.
# Load .tiff data set and converted to .mat data
def loadTifDataTomat(data_path,save_DataPath,name):
    if name=='HU13':
        totalTif=tifffile.imread(os.path.join(data_path,'2013_IEEE_GRSS_DF_Contest_CASI.tif'))
        trainTif=tifffile.imread(os.path.join(data_path,'train_roi.tif'))
        valTif=tifffile.imread(os.path.join(data_path,'val_roi.tif'))

        print(totalTif.shape,trainTif.shape,valTif.shape)
        #spectral.imshow(totalTif)
        #spectral.imshow(trainTif)
        sio.savemat(os.path.join(save_DataPath,"totalTifHouston13.mat"),{'totalTifHouston13':totalTif})
        sio.savemat(os.path.join(save_DataPath,"trainTifHouston13.mat"),{'trainTifHouston13':trainTif})
        sio.savemat(os.path.join(save_DataPath,"valTifHouston13.mat"),{'valTifHouston13':valTif})



def loadTifMat(data_path,name):
    if name=='HU13':
        data=sio.loadmat(os.path.join(data_path, 'totalTifHouston13.mat'))['totalTifHouston13']
        train=sio.loadmat(os.path.join(data_path, 'trainTifHouston13.mat'))['trainTifHouston13']
        val=sio.loadmat(os.path.join(data_path, 'valTifHouston13.mat'))['valTifHouston13']
        return data,train,val



### Using PCA for removing the spectral redundancy(冗余)
### Reduce the spectral dimension, from high-dimensional to low-dimensional.
def applyPCA(X, numComponents=75):
    newX = np.reshape(X, (-1, X.shape[2]))
    pca = PCA(n_components=numComponents, whiten=True)
    newX = pca.fit_transform(newX)
    newX = np.reshape(newX, (X.shape[0],X.shape[1], numComponents))
    return newX, pca



### Padding zeros
def padWithZeros(X, margin=2):
    newX = np.zeros((X.shape[0] + 2 * margin, X.shape[1] + 2* margin, X.shape[2]))
    x_offset = margin
    y_offset = margin
    newX[x_offset:X.shape[0] + x_offset, y_offset:X.shape[1] + y_offset, :] = X
    return newX



### Create data cube,3D-patch.
def createImageCubes(X, y, windowSize=5, removeZeroLabels = True):
    margin = int((windowSize - 1) / 2)
    zeroPaddedX = padWithZeros(X, margin=margin)
    # split patches
    patchesData = np.zeros((X.shape[0] * X.shape[1], windowSize, windowSize, X.shape[2]))
    patchesLabels = np.zeros((X.shape[0] * X.shape[1]))
    patchIndex = 0
    for r in range(margin, zeroPaddedX.shape[0] - margin):
        for c in range(margin, zeroPaddedX.shape[1] - margin):
            patch = zeroPaddedX[r - margin:r + margin + 1, c - margin:c + margin + 1]
            patchesData[patchIndex, :, :, :] = patch
            patchesLabels[patchIndex] = y[r-margin, c-margin]
            patchIndex = patchIndex + 1
    if removeZeroLabels:
        patchesData = patchesData[patchesLabels>0,:,:,:]
        patchesLabels = patchesLabels[patchesLabels>0]
        patchesLabels -= 1
    return patchesData, patchesLabels

# Dataset split.
def splitTrainTestSet(X, y, testRatio, randomState=345):
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=testRatio, random_state=randomState,
                                                        stratify=y)
    return X_train, X_test, y_train, y_test
[((65, 11, 65, 42), 'numpy.reshape', 'np.reshape', ({(65, 22, 65, 23): 'X', (65, 25, 65, 41): '(-1, X.shape[2])'}, {}), '(X, (-1, X.shape[2]))', True, 'import numpy as np\n'), ((66, 10, 66, 54), 'sklearn.decomposition.PCA', 'PCA', (), '', False, 'from sklearn.decomposition import PCA\n'), ((68, 11, 68, 67), 'numpy.reshape', 'np.reshape', ({(68, 22, 68, 26): 'newX', (68, 28, 68, 66): '(X.shape[0], X.shape[1], numComponents)'}, {}), '(newX, (X.shape[0], X.shape[1], numComponents))', True, 'import numpy as np\n'), ((75, 11, 75, 82), 'numpy.zeros', 'np.zeros', ({(75, 20, 75, 81): '(X.shape[0] + 2 * margin, X.shape[1] + 2 * margin, X.shape[2])'}, {}), '((X.shape[0] + 2 * margin, X.shape[1] + 2 * margin, X.shape[2]))', True, 'import numpy as np\n'), ((88, 18, 88, 89), 'numpy.zeros', 'np.zeros', ({(88, 27, 88, 88): '(X.shape[0] * X.shape[1], windowSize, windowSize, X.shape[2])'}, {}), '((X.shape[0] * X.shape[1], windowSize, windowSize, X.shape[2]))', True, 'import numpy as np\n'), ((89, 20, 89, 55), 'numpy.zeros', 'np.zeros', ({(89, 30, 89, 53): 'X.shape[0] * X.shape[1]'}, {}), '(X.shape[0] * X.shape[1])', True, 'import numpy as np\n'), ((105, 39, 106, 67), 'sklearn.model_selection.train_test_split', 'train_test_split', (), '', False, 'from sklearn.model_selection import train_test_split\n'), ((40, 33, 40, 93), 'os.path.join', 'os.path.join', ({(40, 46, 40, 55): 'data_path', (40, 56, 40, 92): '"""2013_IEEE_GRSS_DF_Contest_CASI.tif"""'}, {}), "(data_path, '2013_IEEE_GRSS_DF_Contest_CASI.tif')", False, 'import os\n'), ((41, 33, 41, 72), 'os.path.join', 'os.path.join', ({(41, 46, 41, 55): 'data_path', (41, 56, 41, 71): '"""train_roi.tif"""'}, {}), "(data_path, 'train_roi.tif')", False, 'import os\n'), ((42, 31, 42, 68), 'os.path.join', 'os.path.join', ({(42, 44, 42, 53): 'data_path', (42, 54, 42, 67): '"""val_roi.tif"""'}, {}), "(data_path, 'val_roi.tif')", False, 'import os\n'), ((47, 20, 47, 71), 'os.path.join', 'os.path.join', ({(47, 33, 47, 46): 'save_DataPath', (47, 47, 47, 70): '"""totalTifHouston13.mat"""'}, {}), "(save_DataPath, 'totalTifHouston13.mat')", False, 'import os\n'), ((48, 20, 48, 71), 'os.path.join', 'os.path.join', ({(48, 33, 48, 46): 'save_DataPath', (48, 47, 48, 70): '"""trainTifHouston13.mat"""'}, {}), "(save_DataPath, 'trainTifHouston13.mat')", False, 'import os\n'), ((49, 20, 49, 69), 'os.path.join', 'os.path.join', ({(49, 33, 49, 46): 'save_DataPath', (49, 47, 49, 68): '"""valTifHouston13.mat"""'}, {}), "(save_DataPath, 'valTifHouston13.mat')", False, 'import os\n'), ((13, 27, 13, 80), 'os.path.join', 'os.path.join', ({(13, 40, 13, 49): 'data_path', (13, 51, 13, 79): '"""Indian_pines_corrected.mat"""'}, {}), "(data_path, 'Indian_pines_corrected.mat')", False, 'import os\n'), ((14, 29, 14, 75), 'os.path.join', 'os.path.join', ({(14, 42, 14, 51): 'data_path', (14, 53, 14, 74): '"""Indian_pines_gt.mat"""'}, {}), "(data_path, 'Indian_pines_gt.mat')", False, 'import os\n'), ((55, 25, 55, 73), 'os.path.join', 'os.path.join', ({(55, 38, 55, 47): 'data_path', (55, 49, 55, 72): '"""totalTifHouston13.mat"""'}, {}), "(data_path, 'totalTifHouston13.mat')", False, 'import os\n'), ((56, 26, 56, 74), 'os.path.join', 'os.path.join', ({(56, 39, 56, 48): 'data_path', (56, 50, 56, 73): '"""trainTifHouston13.mat"""'}, {}), "(data_path, 'trainTifHouston13.mat')", False, 'import os\n'), ((57, 24, 57, 70), 'os.path.join', 'os.path.join', ({(57, 37, 57, 46): 'data_path', (57, 48, 57, 69): '"""valTifHouston13.mat"""'}, {}), "(data_path, 'valTifHouston13.mat')", False, 'import os\n'), ((16, 27, 16, 75), 'os.path.join', 'os.path.join', ({(16, 40, 16, 49): 'data_path', (16, 51, 16, 74): '"""Salinas_corrected.mat"""'}, {}), "(data_path, 'Salinas_corrected.mat')", False, 'import os\n'), ((17, 29, 17, 70), 'os.path.join', 'os.path.join', ({(17, 42, 17, 51): 'data_path', (17, 53, 17, 69): '"""Salinas_gt.mat"""'}, {}), "(data_path, 'Salinas_gt.mat')", False, 'import os\n'), ((19, 27, 19, 64), 'os.path.join', 'os.path.join', ({(19, 40, 19, 49): 'data_path', (19, 51, 19, 63): '"""PaviaU.mat"""'}, {}), "(data_path, 'PaviaU.mat')", False, 'import os\n'), ((20, 29, 20, 69), 'os.path.join', 'os.path.join', ({(20, 42, 20, 51): 'data_path', (20, 53, 20, 68): '"""PaviaU_gt.mat"""'}, {}), "(data_path, 'PaviaU_gt.mat')", False, 'import os\n'), ((26, 27, 26, 65), 'os.path.join', 'os.path.join', ({(26, 40, 26, 49): 'data_path', (26, 51, 26, 64): '"""Houston.mat"""'}, {}), "(data_path, 'Houston.mat')", False, 'import os\n'), ((27, 29, 27, 69), 'os.path.join', 'os.path.join', ({(27, 42, 27, 51): 'data_path', (27, 52, 27, 68): '"""Houston_gt.mat"""'}, {}), "(data_path, 'Houston_gt.mat')", False, 'import os\n'), ((29, 27, 29, 61), 'os.path.join', 'os.path.join', ({(29, 40, 29, 49): 'data_path', (29, 51, 29, 60): '"""KSC.mat"""'}, {}), "(data_path, 'KSC.mat')", False, 'import os\n'), ((30, 29, 30, 65), 'os.path.join', 'os.path.join', ({(30, 42, 30, 51): 'data_path', (30, 52, 30, 64): '"""KSC_gt.mat"""'}, {}), "(data_path, 'KSC_gt.mat')", False, 'import os\n')]