import numpy as np
from sklearn import cross_validation
from sklearn.metrics import confusion_matrix
from sklearn import metrics
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib import colors
import multiclass_metrics



class Validation:
    """Cross validation (using stratified k fold) and reporting
    Validates the classifier (clf) using cross validation and 
    creates a report with confidence intervals (+/- 2 std)
    Args:
    clf (sklearn classifier): the model to validate
    data (numpy array): data, shape (n_samples, m_features)
    targets (numpy array): targets, shape - n_samples
    target_names: unique label names
    n_folds (int): number of folds in cross validation
            (str): 'loocv' to run leave one out cross validation
            (str): 'own', user provides own folds. 
    folds: (list): list should contain dictionaries
                   with train and test indices:
                    [{"test_ind": [], "train_ind": []},
                    {"test_ind": [], "train_ind": []}]

    score: "mean" or "combined " - averaging fold scores vs combined folds 
    Attributes:
    report (string):
        report with precision, recall, f1 and confidence intervals
    y_true (numpy array) : 
        true test fold targets concatenated to the 1d array 
    y_pred (numpy array):
        predicted test fold targets concatenated to the 1d array

    folds: fold indices generated by stratitified k fold
    TODO: rewrite printing
    Notes: 
    1)  Mean(scores) vs Score(concatenation) in cross validation discussion:
        http://stats.stackexchange.com/questions/34611/meanscores-vs-scoreconcatenation-in-cross-validation
    """
    def __init__(self, clf, data, targets, target_names, 
                       n_folds=5, shuffle=False, 
                       random_state=None, score="mean", folds=None):


        self.target_names = target_names
        self.n_folds = n_folds if (type(n_folds) is int) or (n_folds=='loocv') \
                               else len(folds)
        self.shuffle = shuffle
        self.random_state = random_state
        self.folds = folds
        self.fold_metrics = []


        self.y_true = np.array([])
        self.y_pred = np.array([])


        if (type(n_folds) is int) or (n_folds=='own'):
            #array with precision, recall, f1 and support for every class and fold
            self._prfs_cv = np.zeros((4, len(target_names), self.n_folds))
            self.f1_macro_train = np.zeros((self.n_folds, ))
            self.f1_macro_test = np.zeros((self.n_folds, ))
            self._skf(clf, data, targets, folds)
            self.report = self._make_report(score)

        elif n_folds == 'loocv':
            self._loocv(clf, data, targets)
            self.report = metrics.classification_report(self.y_true, self.y_pred)
            self.report += "\n" + "f1_macro: "
            self.report += "%.3f" % self.metrics.macro_fscore
            self.report += "\n" + "f1_micro: "
            self.report += "%.3f" % self.metrics.micro_fscore

        else:
            raise ValueError("Wrong n_fods");

        self.metrics = multiclass_metrics.MulticlassMetrics( \
                                                    self.y_true,
                                                    self.y_pred,
                                                    self.target_names)



    def _skf(self, clf, data, targets, folds=None):
        if folds is None:
            skf = cross_validation.StratifiedKFold(targets, 
                                                   n_folds=self.n_folds,
                                                   shuffle=self.shuffle,
                                                   random_state=self.random_state)
            self.folds = []
            save_folds = True
        else:
            skf = [(item["train_ind"], item["test_ind"]) for item in folds]
            save_folds = False

        for it, (train_idx, test_idx) in enumerate(skf):
            #split the data
            data_train, data_test = data[train_idx], data[test_idx]
            targets_train, targets_test = targets[train_idx], targets[test_idx]
            #fit the classifier
            clf.fit(data_train, targets_train)
            targets_train_pred = clf.predict(data_train)
        
        
            self.f1_macro_train[it] = metrics.f1_score(targets_train,
                                                       targets_train_pred, 
                                                       pos_label=None,
                                                       average='macro')

            #predict
            targets_pred = clf.predict(data_test)

            self._prfs_cv[:,:,it] = metrics.precision_recall_fscore_support(targets_test,
                                                                      targets_pred)
    
            self.f1_macro_test[it] = metrics.f1_score(targets_test,
                                                      targets_pred,
                                                      pos_label=None,
                                                      average='macro')

            self.y_true = np.concatenate((self.y_true, targets_test))
            self.y_pred = np.concatenate((self.y_pred, targets_pred))
            #save fold ind
            if save_folds:
                self.folds.append({"id": it, "train_ind": train_idx, \
                                         "test_ind": test_idx} )

            fmetrics = multiclass_metrics.MulticlassMetrics(targets_test,
                                                            targets_pred,
                                                            self.target_names)
            self.fold_metrics.append({"id": it, "metrics": fmetrics})



    def _loocv(self, clf, data, targets):
        y_pred = []
        y_true = []
        loo = cross_validation.LeaveOneOut(len(targets))
        for train_index, test_index in loo:
            X_train, X_test = data[train_index], data[test_index]
            y_train, y_test = targets[train_index], targets[test_index]
            clf.fit(X_train, y_train)
            y_pred.append(clf.predict(X_test))
            y_true.append(y_test)

        self.y_true = np.array([y[0] for y in y_true])
        self.y_pred = np.array([y[0] for y in y_pred])
        



    def _make_report(self, score):

        #compute combined scores
        comb_prfs = metrics.precision_recall_fscore_support(self.y_true,
                                                            self.y_pred)

        # compute mean scores
        means = self._prfs_cv.mean(axis=2)[:3,:]
        stds = self._prfs_cv.std(axis=2)[:3,:]
        supports = self._prfs_cv[3].sum(axis=1)

        report = str()
        report += "%46s" % "precision"
        report += "%16s" % "recall"
        report += "%19s" % "f1 score"
        report += "%12s" % "support\n"

        for i in range(means.shape[1]):
            row = "%30s" % self.target_names[i]
            for j in range(3):
                if( score == "mean"):
                    row += "%7.2f" % means[j,i]
                else:
                    row += "%7.2f" % comb_prfs[j][i]
                row += " (+/-%0.2f)" % (2*stds[j,i])
            row += " %5d" % supports[i]
            report += row + "\n"

        report += "\nf1_macro train:         "  + " %7.3f" % self.f1_macro_train.mean()
        report += " (+/-%0.3f)" % (2*self.f1_macro_train.std())
        report += "\nf1_macro test:          "  + " %7.3f" %  self.f1_macro_test.mean() 
        report += " (+/-%0.3f)" % (2*self.f1_macro_test.std()) + "\n"
        report += "overall test accuracy: " + " %7.3f" % (metrics.accuracy_score(
                                                     self.y_true,
                                                     self.y_pred))
        return report



        