import numpy as np
import pickle
from scipy.io import loadmat, savemat
from sklearn.decomposition import PCA
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
from ols import ols
from optparse import OptionParser
from util import *
from itertools import zip_longest, chain, islice

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

_patient_limit = 1000 # Maximal number of patients to use for training.
_heartbeat_limit = 20 # Maximal number of heartbeats per patient to use for training.

"""
Reduce the dimentsion of all patients feature vectors,
while preserving the nested list structure.
"""
def _reduce_dimension(dimred, patients):
    reduced_patients = [dimred.transform(p) for p in patients]
    return np.array(reduced_patients)



if __name__ == '__main__':
    
    parser = OptionParser()
    parser.add_option("-i", "--input", dest="input_filename", help="Load feature vectors from FILE", metavar="FILE")
    parser.add_option("-o", "--output", dest="output_filename", help="Store dimension-reduced feature vectors in FILE", metavar="FILE")
    parser.add_option("-n", "--numberofdimensions", dest="ndimens", help="Reduce feature vectors to N dimensions", type='int', metavar="N", default=100)
    parser.add_option("-p", "--params", dest="params_filename", help="Save trained transformation to file", metavar="FILE")
    parser.add_option("-t", "--training", action="store_true", dest="training", default=False, help="Operate in training mode.")
    parser.add_option("-l", "--labels", dest="label_filename", help="Feature vector labels.", metavar="FILE")
    parser.add_option("--algo", dest="algorithm", help="PCA or OLS", metavar="ALGO", default='pca')

    (options, args) = parser.parse_args()
    
    if options.input_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.params_filename is None:
        print('Parameter file name is required.')
        exit(1)
        
    algorithm = options.algorithm.lower()
    
    if algorithm not in ['pca', 'ols', 'lda']:
        print('Unknown algorithm: %s'%algorithm)
        exit(1)

    supervised = algorithm in ['ols', 'lda']


    infilename = options.input_filename
    outfilename = options.output_filename
    ndims = options.ndimens
    trainingmode = options.training
    
    if trainingmode and supervised and options.label_filename is None:
        print('Label file name is required.')
        exit(1)
        
    if algorithm == 'pca':   
        dimred = PCA(n_components=ndims)
    elif algorithm == 'ols':
        dimred = ols(n_components=ndims)
    elif algorithm == 'lda':
        dimred = LinearDiscriminantAnalysis(n_components=ndims)
    
    print('loading data...')
    if hdf5_enabled and h5py.is_hdf5(infilename):
        # Load HDF5 file
        print('Loading data from HDF5 file...')
        f = h5py.File(infilename, 'r')
        patients = hdf5_load_patients(f)
    else:
        # Load mat file
        matfile = loadmat(infilename)
        patients = matfile["scattered"] #[0]

    if trainingmode:
        print('training dimensionality reduction...')
        
        # Load labels if necessary.
        if options.label_filename is not None:
            targets = loadmat(options.label_filename)
            targets = targets['label_list'][0]
            targets = targets[:len(patients)]
        else:
            targets = None
            
        if supervised:
            assert targets is not None, 'A list of labels is required for supervised algorithms.'

        # Subsample dataset such that normal and abnormal are as balanced as possible.
        if _patient_limit > 0 and targets is not None:
            normal = -1
            abnormal = 1
            
            pat_indices = np.arange(len(patients))
            
            idx_normal = pat_indices[targets == normal]
            idx_abnormal = pat_indices[targets == abnormal]
            
            np.random.seed(0)
            np.random.shuffle(idx_normal)
            np.random.shuffle(idx_abnormal)
            
            def interleave(a, b):
                chained = chain.from_iterable(zip_longest(a, b))
                return filter(lambda x: x is not None, chained) # remove all None
            
            idx_interleaved = interleave(idx_normal, idx_abnormal)
            idx_balanced = np.array(list(islice(idx_interleaved, 0, _patient_limit)))
            
            assert _patient_limit == 0 or len(idx_balanced) <= _patient_limit
            
            patients_balanced = patients[idx_balanced]
            
            if supervised:
                targets_balanced = targets[idx_balanced]
            else:
                targets_balanced = None
                
            # Subsample balanced set.
            patients_ss, targets_ss = subsample_patients(patients_balanced, targets_balanced, _patient_limit)  # Use at most _patient_limit patients.
            patients_ss = subsample_heartbeats(patients_ss, _heartbeat_limit)
        
        else:
            # _patient_limit == 0
            # Subsample balanced set.
            patients_ss, targets_ss = subsample_patients(patients, targets, _patient_limit)  # Use at most _patient_limit patients.
            patients_ss = subsample_heartbeats(patients_ss, _heartbeat_limit)
            
        # Put feature vectors in suitable format.
        featurematrix, pat_ids = nested_to_matrix(patients_ss)
        
        if supervised:
            dimred.fit(featurematrix, targets_ss[pat_ids])
        else:
            dimred.fit(featurematrix)
        
        print('Save trained transformation to %s'%options.params_filename)
        pickle.dump(dimred, open(options.params_filename,'wb'))
        
        print("Fitting complete.")

  
    else:
        # Load transform from file.
        dimred = pickle.load(open(options.params_filename,'rb'))

    fv_reduced = _reduce_dimension(dimred, patients)
    
    print("Data reduction complete, exporting dimensionality reducing transformation and reduced data...")
    savemat(outfilename, {'features': [fv_reduced]})

    
