import numpy as np
import pickle
from optparse import OptionParser
from scipy.io import loadmat, savemat
from sklearn import cross_validation
from sklearn.svm import SVC, LinearSVC, l1_min_c
from sklearn import linear_model
from copy import deepcopy

from meta_classifier import *
from util import *

try:
    from joblib import Parallel, delayed
    from multiprocessing import cpu_count
    
    n_jobs = cpu_count()
except:
    n_jobs = 1


_patient_subsamples = 2000  # Maximal number of patients to use for parameter sweep
_heartbeat_subsamples = 10  # Maximal number of heartbeats per patient to use for parameter sweep


def _F1_score(expected, predicted):
    """
    Returns: (F1_score, sensibility, specificity)
    """
    
    assert len(expected) == len(predicted)
    
    
#     wa1=0.8602
#     wa2=0.1398
#     wn1=0.9252
#     wn2=0.0748
    
    Aa1 = np.sum(np.logical_and(expected == 1, predicted == 1))
    Aa2Aq2 = np.sum(np.logical_and(expected == 2, predicted >= 0))
    A1 = np.sum(expected == 1)
    A2 = np.sum(expected == 2)
    
    Nn1 = np.sum(np.logical_and(expected == -1, predicted == -1))
    Nn2Nq2 = np.sum(np.logical_and(expected == -2, predicted <= 0))
    N1 = np.sum(expected == -1)
    N2 = np.sum(expected == -2)
    
    # Compute weights based on the ground truth
    wa1 = A1/(A1 + A2)
    wa2 = A2/(A1 + A2)
    wn1 = N1/(N1 + N2)
    wn2 = N2/(N1 + N2)
    
    Se = wa1*Aa1/A1 + wa2*Aa2Aq2/A2
    Sp = wn1*Nn1/N1 + wn2*Nn2Nq2/N2
    
#     w_unsure = 0.5
#     
#     normal = -1
#     abnormal = 1
#     uncertain = 0
# 
#     N = predicted[expected == normal]      # This predictions are expected to be -1.
#     A = predicted[expected == abnormal]    # This predictions are expected to be +1.
#     
#     Nn = np.sum(N == normal)
#     Nq = np.sum(N == uncertain)
#     Na = np.sum(N == abnormal)
#     
#     An = np.sum(A == normal)
#     Aq = np.sum(A == uncertain)
#     Aa = np.sum(A == abnormal)
# 
#     #Sensibility
#     if len(A) != 0:
#         Se = (Aa + w_unsure*Aq)/len(A)
#     else:
#         Se = 0.5
# 
#     if len(N) != 0:
#         Sp = (Nn + w_unsure*Nq)/len(N)
#     else:
#         Sp = 0.5

    return ((Se+Sp)/2, Se, Sp)


"""
 performs stratified k fold cross validation, SVM training, SVM prediction
 input: fvs and targets
    where fvs is of shape:    
    [
        patient[0],
        patient[1],
        ...
    ]
    
    where patient[i] is of shape:
    [
        featurevector[0],
        featurevector[1],
        ...
    ]

    and targets is of shape:
    [
        target_patient[0],
        target_patient[0],
        ...
    ]

output: the overall F1_score of physionet
"""

def __single_score(meta_classifier, fvs, targets, train_index, test_index):
    """Calculate F1 score.
    Used inside cross validation.
    """ 
    train_patients  = fvs[train_index]
    # Reduce unsure labels: Don't need the classifier to distinguish between normal/abnormal unsure recordigs!
    train_targets   = _get_train_targets(targets[train_index])
    
    test_patients   = fvs[test_index]
    test_targets    = targets[test_index]

    meta_classifier.fit(train_patients, train_targets)

    expected = test_targets
    predicted = meta_classifier.predict(test_patients)
    score = _F1_score(expected,predicted)
    
    return np.array(score)
        
def _get_cross_validation_score(meta_classifier, fvs, targets, n_folds, parallelize=False):
    """
    Perform cross validation on fvs and targets using given classifier.
    
    Returns: The average (F1_score, sensibility, specificity)
    """
    
    
    skf = cross_validation.StratifiedKFold(targets, n_folds=n_folds, shuffle=True, random_state=0)
    
    score = np.zeros(3)
     
    
    if parallelize and (n_jobs > 1):
        # Run multiple jobs in parallel.
        scores = Parallel(n_jobs=n_jobs, verbose=10)(
            delayed(__single_score)(deepcopy(meta_classifier), fvs, targets, train_index, test_index) for train_index, test_index in skf
        )
    else:
        scores = [__single_score(deepcopy(meta_classifier), fvs, targets, train_index, test_index) for train_index, test_index in skf]
        
    scores = np.array(scores)
    score = np.mean(scores, axis = 0)
    
    #print('crossvalidation: %s => score = %f'%(str(classifier), score))

    return tuple(score)
    
    
    
def _c_gamma_sweep(classifier, fvs, targets, C_space, gamma_space, n_folds=3):
    """Fit C and gamma parameters for best F1 score.
    
    Parameters
    ----------
    
    classifier      : The classifier to be used. This must be a constructor function: Classifier: (C, gamma) -> Classifier(C, gamma)
    fvs             : Feature vectors.
    targets         : Labels of the feature vectors.
    C_space         : All values for C to try.
    gamma_space     : All values for gamma to try.
    n_folds         : Default = 3.
    
    """
    
    assert(len(fvs) == len(targets))
    
    print("C_space     : [%10.6f ... %10.6f]"%(C_space[0], C_space[-1]))
    print("gamma_space : [%10.6f ... %10.6f]"%(gamma_space[0], gamma_space[-1]))
    
    best_score = 0
    best_params = (C_space[0], gamma_space[0])
    
    grid_C, grid_gamma = np.meshgrid(C_space, gamma_space)
    grid_C = grid_C.flatten()
    grid_gamma = grid_gamma.flatten()
    
    parameters = zip(grid_C, grid_gamma)
   
    if n_jobs > 1:
        # Run multiple jobs in parallel.
        scores = Parallel(n_jobs=n_jobs, verbose=10)(
            delayed(_get_cross_validation_score)(classifier(C, gamma), fvs, targets, n_folds) for C,gamma in parameters
        )
    else:
        scores = [_get_cross_validation_score(classifier(C, gamma), fvs, targets, n_folds) for C,gamma in parameters]
    
    for (s, sens, sel),C,gamma in zip(scores, grid_C, grid_gamma):
        # Find the best parameters so far.
        #print('C = %.4f, gamma = %.4f:  score = %.6f'%(C, gamma, s))
        
        if s > best_score:
            best_score = s
            best_params = (C,gamma)
                
    return (best_params, best_score)


def _iterative_c_gamma_sweep(classifier, fvs, targets, C_log_range=(-4,6), gamma_log_range=(-8,1)):
    """Perform 2D parameter sweep with adaptive grid size.
    
    Parameters
    ----------
    
    classifier      : The classifier to be used. This must be a constructor function: Classifier: (C, gamma) -> Classifier(C, gamma)
    fvs             : Feature vectors.
    targets         : Labels of the feature vectors.
    C_log_range     : Start and end exponent of C space.
    gamma_log_range : Start and end exponent of gamma space.
   
    
    Returns: (best_C, best_gamma)
    
    """
    
    assert len(targets) == len(fvs)

    grid_size = 5
    
    C_start, C_end = C_log_range
    C_space = np.logspace(C_start, C_end, num=grid_size, base=2)
    
    gamma_start, gamma_end = gamma_log_range
    
    gamma_space = np.logspace(gamma_start, gamma_end, num=grid_size, base=2)
    
    for i in range(0, 3):
        C_space = np.logspace(C_start, C_end, num=grid_size+2, base=2)
        gamma_space = np.logspace(gamma_start, gamma_end, num=grid_size+2, base=2)
        
        fvs_subsampled, targets_subsampled = subsample_patients(fvs, targets, _patient_subsamples)  # Use at most _patient_subsamples patients.
        fvs_subsampled = subsample_heartbeats(fvs_subsampled, _heartbeat_subsamples)  # Used in parameter sweep only
    
        # Sweep over C_space x gamma_space.
        (best_params, best_score) = _c_gamma_sweep(classifier, fvs_subsampled, targets_subsampled, C_space[1:-1], gamma_space[1:-1])
        
        print('Best parameters sofar: C = %.4f, gamma = %.4f, score = %.6f'%(best_params + (best_score,)))
        
        C, gamma = best_params
        idx_C = np.argmax(C_space == C)
        idx_gamma = np.argmax(gamma_space == gamma)
        
        # Calculate borders of new grid.
        C_start, C_end = np.log2(C_space[[idx_C-1,idx_C+1]])
        gamma_start, gamma_end = np.log2(gamma_space[[idx_gamma-1,idx_gamma+1]])

    print('Best parameters: C = %.4f, gamma = %.4f, score = %.6f, sensitivity = %.6f'%(best_params + (best_score,)))
    
    return best_params
    
def _default_svc(C, gamma, probability):
    """
    Shortcut for creating a SVC instance.
    Used to centralize parameters such as cache_size.
    
    Parameters
    ----------
    C       : SVM parameter.
    gamma   : SVM parameter.
    probability: Enable probability estimation.
    """
    
    svc = SVC(
            C=C, gamma=gamma,
            cache_size=2000, 
            class_weight='balanced', 
            random_state = 0,
            probability=probability
            )
            
    return svc

def _default_logisticRegression(C, gamma):
    """
    Shortcut for creating a RBF_LogReg instance.
    """
    
    return RBF_LogReg(C, gamma)


class RBF_LogReg:
    """ Wrapper for chaining RBFSampler and LogisticRegression.
    
    Parameters
    ----------
    
    C               : C parameter of LogisticRegression.
    gamma           : gamma parameter of RBFSampler.
    probability     : Ignored.
    rbf_factor      : Enlarge dimension by this factor. Default = 5.
    
    """
    
    def __init__(self, C, gamma, probability = False, rbf_factor = 5):
        self.C = C
        self.gamma = gamma
        self._rbf_factor = rbf_factor
        self._rbf = None
        self._logReg = None
        
    def fit(self, fvs, targets):
        n_components = self._rbf_factor*len(fvs[0])  # take a multiple of the input dimension as output dimension
        
        rbf = RBFSampler(gamma=self.gamma, n_components=n_components, random_state=1)
        rbf.fit(fvs)
        
        self._rbf = rbf
        
        logReg = linear_model.LogisticRegression(
                            C=self.C,
                            penalty='l1',
                            tol=1e-6,
                            class_weight='balanced',
                            random_state = 0
                        )
        
        
        logReg.fit(rbf.transform(fvs), targets)
                        
        self._logReg = logReg
        
    def predict(self, fvs):
        assert self._logReg is not None, "Classifier is not yet trained."
        assert self._rbf is not None, "Classifier is not yet trained."
        
        return self._logReg.predict(self._rbf.transform(fvs))
        
    def predict_proba(self, fvs):
        assert self._logReg is not None, "Classifier is not yet trained."
        assert self._rbf is not None, "Classifier is not yet trained."
        
        return self._logReg.predict_proba(self._rbf.transform(fvs))

    @property
    def classes_(self):
        assert self._logReg is not None, "Classifier is not yet trained."
        return self._logReg.classes_


def _default_classifier(C, gamma=1, probability=False):
    
    # majority vote
    _default_meta_classifier, probability = (MajorityClf, False)
    
    # maximum likelihood
    #_default_meta_classifier, probability = (MaximumLikelihoodClf, True)
    
    # SVM
    clf = _default_svc(C, gamma, probability)
    
    # logistic regression
    #clf = _default_logisticRegression(C, gamma)
    
    return _default_meta_classifier(clf)

"""
Reduce class-wise unsure label to single unsure label
Currently does not consider unsure classification to get baseline
"""
def _get_train_targets(targets):
    return (-1.0)*(targets < 0) + (targets > 0)

#     train_targets = targets.copy()
#     train_targets[np.abs(train_targets) > 1] = 0
#     return train_targets


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

    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 predicted labels in FILE.", metavar="FILE")
    parser.add_option("-p", "--param", dest="parameter_filename",
                        help="Store trained SVM in FILE", metavar="FILE")
    parser.add_option("-l", "--limit", dest="limit",
                        help="Process only N patient records.", type='int', metavar="N")
    parser.add_option("--labels", dest="label_filename",
                        help="File containing one label per line.", metavar="FILE")
    parser.add_option("-t", "--training", action="store_true", dest="training", default=False,
                        help="Operate in training mode.")
                      
    (options, args) = parser.parse_args()

    if options.input_filename is None:
        print('Path to .mat file is required.')
        exit(1)
    
    if options.parameter_filename is None:
        print('Parameter file name is required.')
        exit(1)
        
    if options.training:
        if options.label_filename is None:
            print('Path to label file must be given.')
            exit(1)
        
    # Load the feature vectors and the targets to the patients
    
    print('Loading data from %s.'%options.input_filename)
    
    matfile = loadmat(options.input_filename)
    
    fvs = matfile['features'][0]
    
    if options.training:
        labels = loadmat(options.label_filename)
        targets = labels['label_list'][0]
        targets = targets[0:len(fvs)]
        
            
        print('Fitting classifier...')
        
        classifier = _default_classifier
        
        # C_range = (-4, 6)
#         gamma_range = (-6, 4)
        C_range = np.logspace(-1, 2, num=6, base=10)
        gamma_range = np.logspace(-2, 1, num=6, base=10)
        
        enable_parameters_sweep = True
        
        if enable_parameters_sweep:
#             params = _iterative_c_gamma_sweep(classifier, fvs, targets, C_range, gamma_range)
            (params, best_score) = _c_gamma_sweep(classifier, fvs, targets, C_range, gamma_range, n_folds=3)
            print('Best score for 3-fold cross validation: %f'%best_score)
            print('Corresponding parameters: C = %.4f, gamma = %.4f'%params)
        else:
            params = (29.628, 0.25) # SVM
            #params = (1.83, 1) # LogisticRegression
            
        print('Fitting classifier with whole dataset...')
        meta_classifier = classifier(*params)
        meta_classifier.fit(fvs, _get_train_targets(targets))
        
        if type(meta_classifier) is MaximumLikelihoodClf:
            print("ML decision boundaries: %s"%str(meta_classifier._bounds))
        
        # Store trained svm or logisticRegression to file.
        print('Saving classifier parameters to %s.'%options.parameter_filename)
        pickle.dump(meta_classifier, open(options.parameter_filename, 'wb'))
        
        # TEMP HACK: CHECK PREDICTION FOR VALIDATION SET
        print("Evaluating score for validation set...")
        validx = loadmat('valindices.mat')
        validation_indices = validx['valindices'][0]
        fvs_val = fvs[validation_indices] #[fvs[i] for i in validation_indices]
        targets_val = targets[validation_indices] #[targets[i] for i in validation_indices]
        pred_val = meta_classifier.predict(fvs_val)
        score_val = _F1_score(targets_val, pred_val)
        print("test set = validation set:\tscore = %f, sensitivity = %f, specificity = %f"%score_val)
        savemat('validpred.mat', {'validpred': np.array(pred_val)})
        ###
        
        # Inspect prediction.
        prediction = meta_classifier.predict(fvs)
        pred_normal = np.sum(prediction == -1)
        pred_abnormal = np.sum(prediction == 1)
        pred_uncertain = np.sum(prediction == 0)
        
        eff_normal = np.sum(targets == -1)
        eff_abnormal = np.sum(targets == 1)
        
        print('\t\tPrediction\tEffective')
        print('Normal:\t\t%4d\t\t%4d'%(pred_normal, eff_normal))
        print('Abnormal:\t%4d\t\t%4d'%(pred_abnormal, eff_abnormal))
        print('Uncertain:\t%4d'%pred_uncertain)
        
        score = _F1_score(targets, prediction)
        print("test set = training set:\tscore = %f, sensitivity = %f, specificity = %f"%score)
        
        cv_score = _get_cross_validation_score(deepcopy(meta_classifier), fvs, targets, n_folds = 5, parallelize = True)
        print("cross validation:\t\tscore = %f, sensitivity = %f, specificity = %f"%cv_score)
        
        
    else:
        # Prediction mode.
        
        # Load classifier from file.
        meta_classifier = pickle.load(open(options.parameter_filename, 'rb'))
        
        predicted = meta_classifier.predict(fvs)
        
        if options.output_filename is not None:
            # Store prediction in file.
            print('Writing predicted values to %s...'%options.output_filename)
            fp = open(options.output_filename, 'w')
            for p in predicted:
                fp.write('%d\n'%p)
            fp.close()
                
        else:
            print(predicted)

