"""
"""
import re

import pandas as pd

from simdeep.config import PATH_DATA
from simdeep.config import SURVIVAL_FLAG
from simdeep.config import SEPARATOR
from simdeep.config import ENTREZ_TO_ENSG_FILE
from simdeep.config import USE_INPUT_TRANSPOSE
from simdeep.config import DEFAULTSEP
from simdeep.config import CLASSIFIER

import  numpy as np

from os.path import isfile

from scipy.stats import rankdata

from numpy import hstack

from sklearn.metrics import pairwise_distances

from sklearn.preprocessing import LabelBinarizer
from sklearn.preprocessing import RobustScaler

from collections import defaultdict

from simdeep.coxph_from_r import coxph
from simdeep.coxph_from_r import c_index

from scipy.stats import kruskal
from scipy.stats import ranksums

from os.path import isdir
from os import mkdir


################ DEBUG ################
# supposed to be None for normal usage
MAX_FEATURE = None
#######################################


class MadScaler():
    def __init__(self):
        """
        """
        pass
    def fit_transform(self, X):
        """ """
        X = np.asarray(X)

        for i in range(len(X)):
            med = np.median(X[i])
            mad = np.median(np.abs(X[i] - med))
            X[i] = (X[i] - med) / mad

        return np.nan_to_num(np.matrix(X))

class RankNorm():
    """
    """
    def __init__(self):
        """
        """
        pass

    def fit_transform(self, X):
        """ """
        X = np.asarray(X)
        shape = list(map(float, X.shape))

        for i in range(len(X)):
            X[i] = rankdata(X[i]) / shape[1]

        return np.matrix(X)

class SampleReducer():
    """
    """
    def __init__(self, perc_sample_to_keep=0.90):
        """
        """
        assert(isinstance(perc_sample_to_keep, float))
        assert(0.0 < perc_sample_to_keep < 1.0)
        self.perc_sample_to_keep = perc_sample_to_keep

    def sample_to_keep(self, datasets, index=None):
        """
        """
        nb_samples = len(datasets.values()[0][index])
        scores = np.zeros(nb_samples)
        threshold = int(nb_samples * self.perc_sample_to_keep)

        for key in datasets:
            scores_array = np.array([vector.sum() for vector in datasets[key][index]])
            scores = scores + scores_array

        scores = [(pos, score) for pos, score in enumerate(scores)]

        scores.sort(key=lambda x:x[1], reverse=True)
        to_keep = [pos for pos, score in scores[:threshold]]
        to_remove = [pos for pos, score in scores[threshold:]]

        return to_keep, to_remove

class VarianceReducer():
    """
    """
    def __init__(self, nb_features=200):
        """
        """
        self.nb_features = nb_features
        self.index_to_keep = []

    def fit(self, dataset):
        """
        """
        if self.nb_features > dataset.shape[1]:
            self.nb_features = dataset.shape[1]

        variances = [np.var(array) for array in dataset.T]
        threshold = sorted(enumerate(variances),
                           reverse=True,
                           key=lambda x:x[1],
        )
        self.index_to_keep = [pos for pos, var in threshold[:self.nb_features]]

    def transform(self, dataset):
        """
        """
        return dataset.T[self.index_to_keep].T

    def fit_transform(self, dataset):
        """
        """
        self.fit(dataset)
        return self.transform(dataset)

class CorrelationReducer():
    """
    """
    def __init__(self, distance='correlation', threshold=None):
        """
        """
        self.distance = distance
        self.dataset = None
        self.threshold = threshold

    def fit(self, dataset):
        """ """
        self.dataset = dataset

        if self.threshold:
            self.dataset[self.dataset < self.threshold] = 0

    def transform(self, dataset):
        """ """
        if self.threshold:
            dataset[dataset < self.threshold] = 0

        return 1.0 - pairwise_distances(dataset,
                                        self.dataset,
                                        self.distance)

    def fit_transform(self, dataset):
        """ """
        self.fit(dataset)
        return self.transform(dataset)


class RankCorrNorm():
    """
    """
    def __init__(self, dataset):
        """
        """
        self.dataset = dataset


def load_survival_file(f_name,
                       path_data=PATH_DATA,
                       sep=DEFAULTSEP,
                       survival_flag=SURVIVAL_FLAG):
    """ """
    if f_name in SEPARATOR:
        sep = SEPARATOR[f_name]

    survival = {}

    filename = '{0}/{1}'.format(path_data, f_name)

    if not isfile(filename):
        raise Exception('## Error wirh unexisting file: {0}'.format(filename))

    with open(filename, 'r') as f_surv:
        first_line = f_surv.readline().strip(' \n\r\t').split(sep)
        for field in survival_flag.values():
            try:
                assert(field in first_line)
            except Exception:
                raise Exception("""#### Exception with survival file {fil} header (first line): "{header}". Header does not contain the field "{field}". Please define new survival flags  using the `survival_flag` variable. Current non-valid survival_flag: {sflag}. Needs variables in header: `{header}` defined as values for the key: "patient_id" => ID of the patient, "survival" => time of the study, "event" => event after survival time """.format(
                    header=first_line,
                    field=field,
                    fil=filename,
                    sflag=survival_flag,
                    ))

        patient_id = first_line.index(survival_flag['patient_id'])
        surv_id = first_line.index(survival_flag['survival'])
        event_id = first_line.index(survival_flag['event'])

        for line in f_surv:
            line = line.strip('\n').split(sep)
            ids  = line[patient_id].strip('"')
            ndays = line[surv_id].strip('"')
            isdead = line[event_id].strip('"')

            survival[ids] = (float(ndays), float(isdead))

    return survival


def translate_index(original_ids, new_ids):
    """ """
    index1d = {ids: pos for pos, ids in enumerate(original_ids)}

    return np.asarray([index1d[sample] for sample in new_ids])


def return_intersection_indexes(ids_1, ids_2):
    """ """
    index1d = {ids: pos for pos, ids in enumerate(ids_1)}
    index2d = {ids: pos for pos, ids in enumerate(ids_2)}

    inter = set(ids_1).intersection(ids_2)

    if len(inter) == 0:
        raise(Exception("Error! No common sample index between: {0}... and {1}...".format(
            ids_1[:2], ids_2[:2])))

    index1 = np.asarray([index1d[sample] for sample in inter])
    index2 = np.asarray([index2d[sample] for sample in inter])

    return index1, index2, list(inter)


def convert_metadata_frame_to_matrix(frame):
    """ """
    lbl = LabelBinarizer()

    normed_matrix = np.zeros((frame.shape[0], 0))
    keys = []

    for key in frame.keys():
        if str(frame[key].dtype) == 'object' or str(frame[key].dtype) == 'string':
            matrix = lbl.fit_transform(frame[key].astype('string'))
            if lbl.y_type_ == "binary":
                keys += list(["{0}_{1}".format(key, lbl.classes_[lbl.pos_label])])
            else:
                keys += ["{0}_{1}".format(key, k) for k in lbl.classes_]
        else:
            rbs = RobustScaler()
            matrix = np.asarray(frame[key]).reshape((frame.shape[0], 1))
            matrix = rbs.fit_transform(matrix)
            keys.append(key)

        normed_matrix = hstack([normed_matrix, matrix])

    return pd.DataFrame(normed_matrix, columns=keys)


def load_data_from_tsv(use_transpose=USE_INPUT_TRANSPOSE, **kwargs):
    """
    """
    if use_transpose:
        return _load_data_from_tsv_transposee(**kwargs)
    else:
        return _load_data_from_tsv(**kwargs)


def _load_data_from_tsv(
        f_name,
        key,
        path_data=PATH_DATA,
        f_type=float,
        sep=DEFAULTSEP,
        nan_to_num=True):
    """ """
    f_short = key

    if f_name in SEPARATOR:
        sep = SEPARATOR[f_name]

    f_tsv = open("{0}/{1}".format(path_data, f_name))
    header = f_tsv.readline().strip(sep + '\n').split(sep)

    feature_ids = ['{0}_{1}'.format(f_short, head)
                   for head in header][:MAX_FEATURE]
    sample_ids = []
    f_matrix = []

    for line in f_tsv:
        line = line.strip(sep + '\n').split(sep)
        sample_ids.append(line[0])

        if nan_to_num:
            line[1:] = [0 if (l.isalpha() or not l) else l
                        for l in line[1:MAX_FEATURE]]

        f_matrix.append(list(map(f_type, line[1:MAX_FEATURE])))

    f_matrix = np.array(f_matrix)

    if f_matrix.shape[1] == len(feature_ids) - 1:
        feature_ids = feature_ids[1:]

    assert(f_matrix.shape[1] == len(feature_ids))
    assert(f_matrix.shape[0] == len(sample_ids))

    f_tsv.close()

    return sample_ids, feature_ids, f_matrix

def _format_sample_name(sample_ids):
    """
    """
    regex = re.compile('_1_[A-Z][A-Z]')

    sample_ids = [regex.sub('', sample.strip('"')) for sample in sample_ids]
    return sample_ids

def _load_data_from_tsv_transposee(
        f_name,
        key,
        path_data=PATH_DATA,
        f_type=float,
        sep=DEFAULTSEP,
        nan_to_num=True):
    """ """
    if f_name in SEPARATOR:
        sep = SEPARATOR[f_name]

    f_tsv = open(path_data + f_name)
    header = f_tsv.readline().strip(sep + '\n').split(sep)

    sample_ids = header[1:]

    sample_ids = _format_sample_name(sample_ids)

    feature_ids = []
    f_matrix = []

    if f_name.lower().count('entrez'):
        ensg_dict = load_entrezID_to_ensg()
        use_ensg = True
    else:
        use_ensg = False

    for line in f_tsv:
        line = line.strip(sep + '\n').split(sep)
        feature = line[0].strip('"')

        if nan_to_num:
            line[1:] = [0 if (l.isalpha() or not l) else l
                        for l in line[1:]]

        if use_ensg and feature in ensg_dict:
            features = ensg_dict[feature]
        else:
            features = [feature]

        for feature in features:
            feature_ids.append('{0}_{1}'.format(key, feature))
            f_matrix.append(list(map(f_type, line[1:])))


    f_matrix = np.array(f_matrix).T

    assert(f_matrix.shape[1] == len(feature_ids))
    assert(f_matrix.shape[0] == len(sample_ids))

    f_tsv.close()

    return sample_ids, feature_ids, f_matrix

def select_best_classif_params(clf):
    """
    select best classifier parameters based uniquely
    on test errors
    """
    arr = []

    for fold in range(clf.cv):
        arr.append(clf.cv_results_[
            'split{0}_test_score'.format(fold)])

    score = [ar.max() for ar in np.array(arr).T]
    index = score.index(max(score))

    params = clf.cv_results_['params'][index]

    clf = CLASSIFIER(**params)

    return clf, params

def load_entrezID_to_ensg():
    """
    """
    entrez_dict = {}

    for line in open(ENTREZ_TO_ENSG_FILE):
        line = line.split()
        entrez_dict[line[0]] = line[1:]

    return entrez_dict

def _process_parallel_coxph(inp):
    """
    """
    node_id, activity, isdead, nbdays, seed, metadata_mat, use_r_packages = inp
    pvalue = coxph(activity,
                   isdead,
                   nbdays,
                   seed=seed,
                   metadata_mat=metadata_mat,
                   use_r_packages=use_r_packages)

    return node_id, pvalue

def _process_parallel_cindex(inp):
    """
    """
    (node_id,
     act_ref, isdead_ref, nbdays_ref,
     act_test, isdead_test, nbdays_test, use_r_packages) = inp

    score = c_index(act_ref, isdead_ref, nbdays_ref,
                    act_test, isdead_test, nbdays_test,
                    use_r_packages=use_r_packages
                    )

    return node_id, score

def _process_parallel_feature_importance(inp):
    """
    """
    arrays = defaultdict(list)
    feature, array, labels = inp

    for label, value in zip(labels, np.array(array).reshape(-1)):
        arrays[label].append(value)
    try:
        score, pvalue = kruskal(*arrays.values())
    except Exception:
        return feature, 1.0

    return feature, pvalue

def _process_parallel_feature_importance_per_cluster(inp):
    """
    """
    arrays = defaultdict(list)
    results = []

    feature, array, labels, pval_thres = inp

    for label, value in zip(labels, np.array(array).reshape(-1)):
        arrays[label].append(value)

    for cluster in arrays:
        array = np.array(arrays[cluster])
        array_comp = np.array([a for comp in arrays for a in arrays[comp]
                      if comp != cluster])

        score, pvalue = ranksums(array, array_comp)
        median_diff = np.median(array) - np.median(array_comp)

        if pvalue < pval_thres:
            results.append((cluster, feature, median_diff, pvalue))

    return results


def _process_parallel_survival_feature_importance_per_cluster(inp):
    """
    """

    feature, array, survival, metadata_mat, pval_thres, use_r_packages = inp
    nbdays, isdead = survival.T.tolist()

    pvalue = coxph(
        array,
        isdead,
        nbdays,
        metadata_mat=metadata_mat,
        use_r_packages=use_r_packages
    )

    if not np.isnan(pvalue) and pvalue < pval_thres:
        return (feature, pvalue)
    return None, None

def save_matrix(matrix, feature_array, sample_array,
                path_folder, project_name, key='', sep='\t'):
    """
    """
    if not isdir(path_folder):
        mkdir(path_folder)

    if key:
        key = '_' + key

    f_csv = open('{0}/{1}{2}.tsv'.format(path_folder, project_name, key), 'w')

    f_csv.write(sep + sep.join(map(lambda x:x.split('_', 1)[-1], feature_array)) + '\n')

    for sample, vector in zip(sample_array, matrix):
        vector = np.asarray(vector).reshape(-1)
        f_csv.write('{0}{1}'.format(sample, sep) + sep.join(map(str, vector)) + '\n')

    print('{0}/{1}{2}.tsv saved'.format(path_folder, project_name, key))