#!/bin/env python3

"""
Use this script to append summary features to deep features after
dimensionality reduction.

deep features: N feature vectors per patient.
summary features: One summary feature vector per patient.
"""

import numpy as np
import pickle
from scipy.io import loadmat, savemat
from optparse import OptionParser

try:
    # Storage for huge lists.
    import h5py
    hdf5_enabled = True
except:
    hdf5_enabled = False

def append_summary_features(deep, summary, segment_lengths = None):
    """
    deep: Multiple feature vectors per patient.
    summary: One summary feature vector per patient.
    
    Returns:
    Returns the deep feature vectors concatenated with corresponding summary vector.
    """
    
    if segment_lengths is not None:
        result = [
                        [np.concatenate((d, l, np.squeeze(s))) for d, l in zip(fvs, sl)]
                    for fvs, sl, s in zip(deep, segment_lengths, summary)
                ]
    else:
        result = [
                        [np.concatenate((d, s)) for d in fvs]
                    for fvs, s in zip(deep, summary)
                ]
            
    return result

if __name__ == '__main__':
    
    parser = OptionParser()
    parser.add_option("-d", "--deep-features", dest="deep_filename", help="Load deep feature vectors from FILE", metavar="FILE")
    parser.add_option("-s", "--summary-features", dest="summary_filename", help="Load summary feature vectors from FILE", metavar="FILE")
    parser.add_option("-l", "--segment-lengths", dest="seg_len_filename", help="Load segment lengths from FILE", metavar="FILE")
    parser.add_option("-o", "--output", dest="output_filename", help="Store dimension-reduced feature vectors in FILE", metavar="FILE")
    
    (options, args) = parser.parse_args()
    
    if options.deep_filename is None:
        print('Path to deep feature .mat file is required.')
        exit(1)
        
    if options.summary_filename is None:
        print('Path to summary feature .mat file is required.')
        exit(1)
        
    if options.output_filename is None:
        print('Output file name is required.')
        exit(1)
        
    
    segment_lengths = None
    
    # Load summary features either from mat or HDF5 file.
    if hdf5_enabled and h5py.is_hdf5(options.summary_filename):
        # Load HDF5 file
        print('Loading data from HDF5 file...')
        f = h5py.File(options.summary_filename, 'r')
        summary_fv = f['summary_features']
        if options.seg_len_filename is not None:
            f = h5py.File(options.seg_len_filename, 'r')
            segment_lengths = f['segment_lengths']
    else:
        # Load mat file
        print('Fall-back to mat.')
        summary_fv = loadmat(options.summary_filename)['summary_features'][0]
        if options.seg_len_filename is not None:
            segment_lengths = loadmat(options.seg_len_filename)['segment_lengths'][0]
        
    deep_fv = loadmat(options.deep_filename)
    deep_fv = deep_fv['features'][0]
    
    
#     print(len(deep_fv))
#     print(len(deep_fv[0]))
#     print(len(deep_fv[0][0]))
#     print(len(segment_lengths))
#     print(len(segment_lengths[0]))
#     print(len(segment_lengths[0][0]))
#     print(len(summary_fv))
#     print(summary_fv[0].shape)
    
    result = append_summary_features(deep_fv, summary_fv, segment_lengths)
    
    savemat(options.output_filename, {'features': [result]})
