from load_images import *
from extract_features import *
from skimage import transform
from sklearn import svm
from sklearn.feature_selection import RFE
import os
import pickle
import random

__author__ = "Wei-Chung Huang"
__copyright__ = "Copyright 2018, The SJSU MSCS Master project"
__license__ = "MIT"
__version__ = "1.0.0"

working_path = os.path.dirname(os.path.abspath(__file__))
feature_folder_name = 'extracted_feats'
cluster_folder_name = 'clustered_res'
saved_feature_path = working_path + '/../' + feature_folder_name
saved_cluster_path = working_path + '/../' + cluster_folder_name
n_families = 25


def mask_features(feat_file_name):
    """
    Mask the original feature vector to reduced features with the mask generated from SVM-RFE
    :param feat_file_name: str
            File name of saved feature.
    :return: m-by-n ndarray
            Reduced feature matrix.
            M is the number of samples and N is the reduced length.
    """
    with open(saved_feature_path + '/f1_reduced_mask', 'rb') as fp:
        mask = pickle.load(fp)

    with open(saved_feature_path + '/' + feat_file_name, 'rb') as fp:
        features = pickle.load(fp)

    feature_mat = list()
    for i in range(len(features)):
        feature_vec = list()
        for j in range(len(mask)):
            if mask[j]:
                feature_vec.append(features[i][j])
        feature_mat.append(feature_vec)

    feature_mat = np.asarray(feature_mat)
    return feature_mat


def get_reduced_features(features, n_dims_RFE=1):
    """
    Get reduced feature vectors by using SVM-RFE with specific no. dimensions
    :param features:
    :param n_dims_RFE:
    :return:
    """
    training_set = list()
    training_labels = list()
    for family in features:
        feature_mat = features.get(family)
        for j in range(len(feature_mat)):
            training_set.append(feature_mat[j])
            training_labels.append(family)

    clf = svm.SVC(kernel='linear')
    clf_reduced = RFE(clf, n_dims_RFE, step=1)
    clf_reduced = clf_reduced.fit(training_set, training_labels)
    X_new = clf_reduced.transform(training_set)
    X_mask = clf_reduced.get_support()
    return X_new, X_mask


def fitting_scoring(features, cv=5, verbose=False, is_RFE_mode=False, n_dims_RFE=1):
    """
    Train SVM model and score it.
    :param features: dict{family name: feature matrix}
                    feature_matrix is MxN where M is number of samples
                    and N is the dimension of a feature sample
    :param cv: int (optional)
            Number of cross-validation. Default is 5-fold
    :param verbose: boolean (optional)
            Flag to dump necessary message
    :param is_RFE_mode: boolean (optional)
            If it set to True, then this function will train by using linear-SVM with reduced size.
            If set to False, the RBF kernel will be used.
    :param n_dims_RFE: int (optional)
            Number of reduced dimension is used in SVM-RFE
    :return:
            Return a average accuracy over cross-validation and a confusion matrix with n_fold-by-n-by-n
            If is_RFE_mode is set, extra parameters are return.
                X_new: reduced feature vectors.
                X_mask: boolean mask showing which dim is used or not
    """
    # N-fold cross-validation
    num_fold = cv
    accuracy = [0] * num_fold
    cm = [np.zeros([n_families, n_families], dtype=np.int32) for i in range(num_fold)]

    # label each family with an index for the purpose of generating the confusion matrix
    fam_labels = []
    for fam_name in features.keys():
        fam_labels.append(fam_name)
    fam_index_dict = get_label_index(fam_labels)

    for i in range(num_fold):
        training_set = list()
        training_labels = list()
        testing_set = list()
        testing_labels = list()
        for family in features:
            feature_mat = features.get(family)
            if verbose: print(family, "sample size:", len(feature_mat))

            fold_start = i * int(len(feature_mat) / num_fold)
            fold_end = fold_start + int(len(feature_mat) / num_fold) - 1

            # separate training and testing set
            for j in range(len(feature_mat)):
                if fold_start <= j <= fold_end:
                    testing_set.append(feature_mat[j])
                    testing_labels.append(family)
                else:
                    training_set.append(feature_mat[j])
                    training_labels.append(family)

        p_res = None
        X_new = None
        X_mask = None
        if is_RFE_mode:
            clf = svm.SVC(kernel='linear')
            clf_reduced = RFE(clf, n_dims_RFE, step=1)
            clf_reduced = clf_reduced.fit(training_set, training_labels)
            X_new = clf_reduced.transform(training_set)
            X_mask = clf_reduced.get_support()
            p_res = clf_reduced.predict(testing_set)
        else:
            clf = svm.SVC()
            clf.fit(training_set, training_labels)
            p_res = clf.predict(testing_set)

        accuracy[i] = 0
        for j in range(len(p_res)):
            if p_res[j] == testing_labels[j]:
                accuracy[i] += 1
            in_idx = fam_index_dict[testing_labels[j]]
            out_idx = fam_index_dict[p_res[j]]
            cm[i][in_idx, out_idx] += 1
        accuracy[i] = (accuracy[i] / len(p_res)) * 100

    if is_RFE_mode:
        if verbose: print('n_dims:', n_dims_RFE, accuracy)
        return np.mean(accuracy), cm, X_new, X_mask

    return np.mean(accuracy), cm


def get_label_index(labels):
    """
    Labeling family name with an index.
    :param labels: list of strings
    :return: dict with key is the label and value is the index.
    """
    fam_index_dict = dict()
    for i in range(len(labels)):
        fam_index_dict[labels[i]] = i
    return fam_index_dict


def test():
    """
    Run n-fold cross-validation on the pre-trained model
    :return:
    """
    features = {}
    for i in range(n_families):
        f1_fam_name = 'f1_fam_' + str(i) + '_features'
        feature_mat_fam1 = mask_features(f1_fam_name)
        features[i] = feature_mat_fam1

    accu, cm = fitting_scoring(features, cv=5, verbose=True)
    print('done:', accu)

    return


def train(dataset_path):
    random.seed(1)
    use_imbalance_dataset = True
    reduces_sample_size = -1

    image_list = prefetch_images(dataset_path)

    if len(image_list) <= 0:
        return

    features = {}  # dictionary<family name, list of feature vector>
    for i in range(len(image_list)):
        family = image_list[i][0]
        images = image_list[i][1]
        n_images = len(images)
        if reduces_sample_size > 0:
            n_images = min(reduces_sample_size, n_images)
        min_samples = min(80, n_images)
        max_samples = 120
        num_samples = min(max(n_images, min_samples), max_samples)
        image_idx_list = random.sample(range(n_images), num_samples)
        if use_imbalance_dataset:
            image_idx_list = range(n_images)
        print("fetching...", family)
        sample_count = 0
        for j in image_idx_list:
            image_path = dataset_path + family + '/' + images[j]
            print(image_path, sample_count + 1, "(", j, ")", "out of", len(image_idx_list))
            img = io.imread(image_path)

            # concatenate selected features
            feature_vec = get_features(img, 'h-edge')
            feature_vec += get_features(img, 'hog')
            feature_vec += get_features(img, 'mean')
            feature_vec += get_features(img, 'lbp')
            feature_vec += get_features(img, 'contrast')
            feature_vec += get_features(img, 'median')

            if family not in features:
                features[family] = list()
            l = features.get(family)
            l.append(feature_vec)
            sample_count += 1

    accu, cm = fitting_scoring(features, verbose=True)
    print('overall accuracy:', accu)


def feature_selection_RFE(features):
    """
    Enumerate all accuracies of reducing feature dimension down to 1.
    :param features: dict{family name: feature matrix}
                    feature_matrix is MxN where M is number of samples
                    and N is the dimension of a feature sample.
    :return: list
            All accuracies w.r.t. to each step when on reduced features.
    """
    feature_test = features[list(features.keys())[0]]
    n_dims = len(feature_test[0])
    overall_accu = list()
    for dim_picked in range(n_dims, 0, -1):
        accu, cm = fitting_scoring(features, cv=2, verbose=False, is_RFE_mode=True, n_dims_RFE=dim_picked)
        overall_accu.append(accu[0])

    return overall_accu
