#!/bin/env python3

"""
This module processes the output of the Springer segmentation.
"""

import numpy as np
#from skimage.io import imshow, show
from scipy.io import loadmat, savemat
from util import load_matlab_cell, hdf5_store_patients

from scipy.signal import resample
from scipy.signal import spectrogram
from optparse import OptionParser

from numpy.random import random, normal, seed

from pyscatter.scatteringtree import scattering_tree
from pyscatter.filters.swt1 import swt1
from pyscatter.filters.swt2 import swt2
from pyscatter.pooling import pool

try:
    # Storage for huge lists.
    import h5py
    hdf5_enabled = True
except:
    hdf5_enabled = False
    
try:
    from joblib import Parallel, delayed
    from multiprocessing import cpu_count
    
    n_jobs = cpu_count()
except:
    n_jobs = 1
    

resample_to = [256,  256,  256,  512] # Target lengths for s1, systole, s2, diastole.

def _build_scattering_tree_1d():

    prop = [
        [swt1(wavelet='db1', level=3, start_level=0)],
        [swt1(wavelet='db1', level=3, start_level=0)],
        [swt1(wavelet='db1', level=3, start_level=0)],
        []
    ]

    out = [
        [],
        [None],
        [None],
        [None]
    ]
    
    relu = lambda x: np.maximum(np.zeros(np.array(x).shape),np.array(x))
#     nonlin = [np.abs]*4
    nonlin = [relu]*4
             
    poolings = [
#         pool(2, np.mean),   # Reduce the signal size by a factor 2.
#         pool(2, np.mean),
#         pool(2, np.mean),
#         pool(2, np.mean),
        pool(2, np.max),   # Reduce the signal size by a factor 2.
        pool(2, np.max),
        pool(2, np.max),
        pool(2, np.max),
    ]

    return scattering_tree(prop, out, nonlin, poolings)

def _build_scattering_tree_2d():

    prop = [
        [swt2(wavelet='db1', level=3, start_level=0)],
        [swt2(wavelet='db1', level=3, start_level=0)],
        []
    ]

    out = [
        [],
        [None],
        [None],
    ]
    # ReLU
    # relu = lambda x: np.maximum(np.zeros(np.array(x).shape),x)
    nonlin = [np.abs]*4

    poolings = [
        pool((2,2), np.mean),   # Reduce the signal size by a factor 2.
        pool((2,2), np.mean),
        pool((2,2), np.mean),
    ]

    return scattering_tree(prop, out, nonlin, poolings)

def _normalize(signal):
    mean = np.mean(signal)
#     scale = np.sum(np.abs(signal))/len(signal)+1e-3
    scale = np.std(signal)+1e-6
    signal = ((signal-mean)/scale)
    return signal

def _prepare_segments(segments):
    """Resample and normalize a list of segments.
    """

    # resample
    segments = [resample(s, r) for s, r in zip(segments, resample_to)]
    
    # normalize amplitude
    segments = [_normalize(s) for s in segments]

    return segments
    
def _gen_deep_feature_vector(heartbeats):
    """Generate a feature vectors from all heartbeats of a patient.

    beat_segments   =   [[s11, s12, s13, s14], [s21, s22, s23, s24], ...]
    s** = [x1, x2, ... ]
    
    Returns: deep_features
    """

    st_1d = _build_scattering_tree_1d()
#     st_2d = _build_scattering_tree_2d()

    deep_features = []
    # extract deep features
    for beat_segments in heartbeats:

#         lengths = [len(segment) for segment in beat_segments]

#         prepared = _prepare_segments(beat_segments)
        
        fvs = []
        #fvs.append(st_1d.scatter(prepared))
        
        concatenated = np.concatenate(beat_segments)
        concatenated = _normalize(concatenated)
        concatenated = resample(concatenated, 1024)
        fvs.append(st_1d.scatter(concatenated))
        
#         spectrogram = _spectrogram(concatenated)
#         fvs.append(st_2d.scatter(spectrogram))
        
        fv = np.concatenate(fvs)

        deep_features.append(fv)
    
    return np.array(deep_features)

def _spectrogram(x):
#     x = resample(x, 1024)
    f, t, Sxx = spectrogram(x,
                            nperseg=128,
                            noverlap=114,
                            )

    Sxx = Sxx[:64, :64]
    #assert Sxx.shape == (16, 32)
    
    return Sxx

def _gen_summary_feature_vector(heartbeats):
    """Generate a feature vectors from all heartbeats of a patient.

    beat_segments   =   [[s11, s12, s13, s14], [s21, s22, s23, s24], ...]
    s** = [x1, x2, ... ]
    
    Returns: (deep_features, summary_features)
    """

    by_phase = heartbeats.transpose()

    # Calculate statistics on the interval lenghts.
    # mean_lengths = [mean(len(s1)), mean(len(S2)), mean(len(systole)), mean(len(diastole))]
    mean_lengths = [np.mean(list(map(len, phase))) for phase in by_phase]
    stddev_lengths = [np.std(list(map(len, phase))) for phase in by_phase]
    summary_fv = np.concatenate([mean_lengths, stddev_lengths])


    return np.array(summary_fv)
    
def _scatter_all(signals):
    """
    Compute a set of feature vectors per patient.

    Expects a list:

    [
        [patient0_list_of_segments0, ...],
        [patient1_list_of_segments0, ...],
        ...
    ]

    Where patient?_list_of_segments? is of type
    [[x11, x12, x13, x14], [x21, x22, x23, x24], ...]

    signals = List of patient data.
    """
    if n_jobs > 1:
        # Run multiple jobs in parallel.
        deep = Parallel(n_jobs=n_jobs, verbose=10)(delayed(_gen_deep_feature_vector)(heartbeats) for heartbeats in signals)
    else:
        deep = [_gen_deep_feature_vector(heartbeats) for heartbeats in signals]
        
#     summary = [_gen_summary_feature_vector(heartbeats) for heartbeats in signals]
# 
#     return deep, summary
    
    segment_lengths = [_get_segment_lengths(heartbeats) for heartbeats in signals]

    return deep, segment_lengths

def _get_segment_lengths(heartbeats):
    by_phase = heartbeats.transpose()
    
    segment_lengths = [np.array(list(map(len, phase)), dtype=np.int32) for phase in by_phase]
#     segment_lengths[:] = [sl + [np.sum(sl)] for sl in segment_lengths];
    
    return segment_lengths


if __name__ == '__main__':
    # Run as standalone script.

    parser = OptionParser()
    parser.add_option("-i", "--input", dest="input_filename",
                      help="Load segmented heart beats from FILE", metavar="FILE")
    parser.add_option("-o", "--output", dest="output_filename",
                      help="Store feature vectors in FILE", metavar="FILE")
    parser.add_option("-l", "--limit", dest="limit",
                      help="Process only N patient records.", type='int', metavar="N", default=0)
    parser.add_option("--hdf5", action="store_true", dest="hdf5", default=False, help="Use hdf5 output format.")

    (options, args) = parser.parse_args()

    filename = options.input_filename

    if filename is None:
        print('Path to .mat file is required.')
        exit(1)
        
    if options.output_filename is None:
        print('Output file name is required.')
        exit(1)
        
    if options.hdf5 and not hdf5_enabled:
        print('HDF5 is not installed. Install h5py!')
        exit(1)
        
    # Load the signals.
    
    print('Loading data from %s.'%filename)
    
    patientlist = load_matlab_cell(filename, 'patientlist')
    
    if False:
        # Find average relative lengths of s1, systole, s2, diastole.
        # Typical: [0.15949056  0.22123093  0.13376002  0.48551849]
        heartbeats = np.concatenate(patientlist)
        lengths = np.array([[len(s) for s in h] for h in heartbeats])
        lengths = np.mean(lengths, 0)
        lengths /= np.sum(lengths)
        print(lengths)
    
	#signals = _generate_dummy_signals() # generate dummy test data
    
    # Optionally limit the number of processed records.
    if options.limit > 0:
        patientlist = patientlist[0:min(len(patientlist), options.limit)]

    # Process the data.
    print('Run feature extraction...')
#     deep_features, summary_features = _scatter_all(patientlist)
    deep_features, segment_lengths = _scatter_all(patientlist)
    

    """
    fvs is of shape:
    
    [
        patient[0],
        patient[1],
        ...
    ]
    
    where patient[i] is of shape:
    [
        featurevector[0],
        featurevector[1],
        ...
    ]
    
    =>
    
    fvs = [
        [featurevector[0][0], featurevector[0][1], ...],
        [featurevector[1][0], featurevector[1][1], ...],
        ...
    ]
    
    Loading the result in MATLAB:
    
    loadmat('[file].mat')
    % Variable scattered is a cell of length num_patients.
    % Each cell element contains the feature vectors.
    patient1 = scattered{1}
    patient1_fv1 = patient1(1,:)
    
    """
    
    
    # Store the feature vectors.
    if options.hdf5:
        # HDF5 will handle very large data sets.
        f = h5py.File(options.output_filename, 'w')
        hdf5_store_patients(np.array(deep_features), f)
        
# #         summary = np.array(summary_features)
# #         
# #         s = f.create_dataset('summary_features', summary.shape)
# #         s[...] = summary
#         lengths = np.array(segment_lengths, dtype=object)
#         dt = h5py.special_dtype(vlen=np.dtype('int32'))
#         dset = f.create_dataset('segment_lengths', lengths.shape, dtype=dt)
#         print(lengths.shape)
#         print(lengths.dtype)
#         print(lengths[0,1])
#         dset[...] = lengths
# #         s = f.create_dataset('segment_lengths', data=lengths)
        f.close()
        
    else:
        # Fall-back to mat format.
        print('Saving feature vectors to %s.'%options.output_filename)
#         savemat(options.output_filename, {'scattered': np.array(deep_features), 'summary_features': np.array(summary_features)})
        savemat(options.output_filename, {'scattered': np.array(deep_features), 'segment_lengths': np.array(segment_lengths)})



