
import numpy as np
from scipy.io import loadmat, savemat

try:
    import h5py
    hdf5_enabled = True
except:
    hdf5_enabled = False

# Unpack nested MATLAB cell structure.
def _unpack_cells(mat):
    #mat = [np.array([np.array([np.array(s[0]) for s in h[0]]) for h in p[0]]) for p in mat[0]]
    mat = [np.array([[np.array(s[0]) for s in h[0]] for h in p[0]]) for p in mat[0]]
    return np.array(mat)

"""
Load nested cell from MATLAB.
"""
def load_matlab_cell(path, var_name):
    mat = loadmat(path)
    arr = mat[var_name]
    arr = _unpack_cells(arr)
    return arr

def nested_to_matrix(patients):
    """
    Convert a triple nested list into a matrix.
    `matrix_to_nested()` does the opposite.
    
    Returns:
    (m, ids)
    m = [
        fv1,
        fv2,
        ...
    ]
    
    ids: Tells where a row in m has been located in the nested structure.
    """
    num_fvs = list(map(len, patients))
    patient_ids = np.repeat(np.arange(len(patients)), num_fvs)
    
    fv_matrix = np.concatenate(patients)
    
    assert(len(patient_ids) == len(fv_matrix))
    
    return (fv_matrix, patient_ids)
    
def nested_map(function, struct):
    """
    Applies `function` on each element of d.

    d = A list of lists.
    """
    
    return np.array([
        np.array([function(d) for d in l])
        for l in struct
    ])

def matrix_to_nested(matrix, ids):
    """
    Inverse function of `nested_to_matrix()`.
    """
    
    assert(len(ids) == len(matrix))
    
    unique_ids = np.sort(np.unique(ids))

    # Restore nested structure.
    nested = np.array([matrix[ids == i] for i in unique_ids])
    
    return nested
    
def subsample_heartbeats(patients, num_subsamples):
    """
    For each patient pick at most 'num_subsamples' heartbeats at random.
    """
    
    if num_subsamples <= 0:
        return patients
    
    fvs_ss = []
    np.random.seed(0)
    for heartbeats in patients:
        idx = np.arange(len(heartbeats))
        np.random.shuffle(idx)
        heartbeats = heartbeats[idx[0:num_subsamples]]
        fvs_ss.append(heartbeats)
        
    return np.array(fvs_ss)
    
def subsample_patients(fvs, targets, num_subsamples):
    """
    From all patients and corresponding labels pick at most 'num_subsamples' at random.
    
    Returns: (new_fvs, new_targets)
    """
    if num_subsamples <= 0:
        return (fvs, targets)

    # Random subsampling.
    np.random.seed(0)
    idx = np.arange(len(fvs))
    np.random.shuffle(idx)
    fvs_ss = fvs[idx[0:num_subsamples]]
    
    if targets is not None:
        targets_ss = targets[idx[0:num_subsamples]]
    else:
        targets_ss = None
    
    return (fvs_ss, targets_ss)

def hdf5_store_patients(patients, hdf5file):
    """
    patients: Nested array of feature vectors.
    hdf5file: Object created with h5py.File(..., 'w') or a path name.
    """
    
    if type(hdf5file) is str:
        hdf5file = h5py.File(hdf5file, 'w')
    
    fv_matrix, patient_ids = nested_to_matrix(patients)
    
    hdf5patients = hdf5file.create_dataset('fv_matrix', fv_matrix.shape)
    hdf5patients[...] = fv_matrix

    hdf5patient_ids = hdf5file.create_dataset('patient_ids', patient_ids.shape)
    hdf5patient_ids[...] = patient_ids

def hdf5_load_patients(hdf5file):
    """
    hdf5file: Object created with h5py.File(..., 'r') or a path name.
    
    Returns: Nested array of feature vectors.
    """
    
    if type(hdf5file) is str:
        hdf5file = h5py.File(hdf5file, 'r')
    
    patient_ids = np.array(hdf5file['patient_ids'])
    fv_matrix = np.array(hdf5file['fv_matrix'])

    patients = matrix_to_nested(fv_matrix, patient_ids)
    
    return patients
    
    
    
    
