#!/usr/bin/python3
"""
Python3 implementation of the correlation-based regression method
by Han et al. (https://ieeexplore.ieee.org/document/6981664)
"""

import argparse
import logging
import math
import os
import random
import sys
import json

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

from sklearn import linear_model
from sklearn.preprocessing import PolynomialFeatures
from sklearn.preprocessing import normalize
from sklearn.preprocessing import minmax_scale

from sklearn.metrics import mean_absolute_percentage_error, mean_absolute_error

import numpy as np
import scipy.stats

import logging

logfile = sys.argv[0] + ".log"
logging.basicConfig(filename=logfile, level=logging.DEBUG,
                    format='%(asctime)s <%(filename)s:%(lineno)d> [%(levelname)s] %(message)s',
                    datefmt='%Y-%m-%d %H:%M:%S')
logging.getLogger().setLevel(logging.INFO)

print("Starting...")
logging.info("=================")
logging.info(sys.argv)

parser = argparse.ArgumentParser(description='''Correlation-based regression''',
                                 epilog="Author: Sebastian Kopetzky (sebastian.kopetzky@labvantage-biomax.com)")

parser.add_argument("-d", metavar="<file>", action="store", dest="data_file", required=True,
                    help="File containing the connectivity data in desquared form [subjects x connections]")
parser.add_argument("-a", metavar="<file>", action="store", dest="attr_file", required=True,
                    help="File containing the attribute of interest [subjects x attribute]. NOTE: currently only one attribute/a vector is supported (i.e. one column)")
parser.add_argument("-s", metavar="<file>", action="store", dest="sex_file", required=False,
                    help="File containing the subjects gender or sex. NOTE: Must be in the same order like the other two files")
parser.add_argument("-ids", metavar="<file>", action="store", dest="id_file", required=False,
                    help="File containing the subject ids. If given JSON output will contain the MAE and subject ids. NOTE: Must be in the same order like the data & attribute file")

parser.add_argument('-spearman', dest="spearman", action='store_true',
                    help='Use spearman correlation instead of pearson correlation', default=False)
parser.add_argument('-p', '-poly', dest="poly", action='store_true',
                    help='Use polynomial regression for the cross-validation', default=False)
parser.add_argument('-t', metavar="<float>", dest="threshold", required=False,
                    help='Apply threshold (as absolute value) to CORRELATION values to select connections to use')
parser.add_argument('-o', metavar="<dir>", dest="output_dir", action='store', required=False,
                    help='Output dir for prediction data')
parser.add_argument("-fold", metavar="<int>", action="store", dest="fold", required=False,
                    help="Fold to use for cross-validation (default: 10)")

parser.add_argument("-plot-feature", action="store_true", dest="plot_feature", required=False,
                    help="Set true if distribution of the input feature should be plotted")
parser.add_argument("-feature", metavar="<str>", action="store", dest="feature", required=False,
                    help="Name of the investigated feature (for the plot)")

parser.add_argument('-stability', dest="stability", action='store_true',
                    help='Check stability of correlation vector when drawing random samples', default=False)


parser._optionals.title = "Arguments"
params, unknown_args = parser.parse_known_args()

connectivity_file = params.data_file
attribute_file = params.attr_file


def find_correlations(data_matrix, attribute_matrix, force_positive=False):
    """
    Get correlation values between each edge
    and the attribute(s) of interest

    n: number of subjects
    m: number of edges
    l: number of attributes

    Parameter:
      data_matrix      -- n x m numpy matrix
      attribute_matrix -- n x l numpy matrix

    Returns:
      A tuple with [0] m x l matrix with correlation values
                   [1] the p-values
    """
    (n,m) = data_matrix.shape
    l = 1 # NOTE: only supports vector (single attribute) for now

    correlations = np.zeros((m,l)) # this is C in the paper
    pvalues = np.zeros((m,l))

    for attr_idx in range(l):
        for edge_idx in range(m):
            
            # skip zero columns
            if np.sum(data_matrix[:,edge_idx]) != 0:
                # get correlation
                (r,p) = (None,None)
                if params.spearman:
                    (r,p) = scipy.stats.spearmanr(data_matrix[:,edge_idx], attribute_matrix)
                else:
                    (r,p) = scipy.stats.pearsonr(data_matrix[:,edge_idx], attribute_matrix)

                if force_positive and r < 0.0:
                    r = -1.0 * r

                correlations[edge_idx, attr_idx] = r
                pvalues[edge_idx, attr_idx] = p
            
    return (correlations, pvalues)


def compute_S(data_matrix, C):
    """
    Parameters:
        data_matrix -- the descquared connectivity matrix [n x m]
        C           -- the correlation vector (matrix) returned by find_correlations() [m x l]

    Returns:
        n x l numpy matrix
    """
    (n,m) = data_matrix.shape
    l = 1 # NOTE: attribute matrix not supported at the moment

    s = np.zeros((n,l)); # result matrix n x l

    for attr_index in range(l):

        nonnan = ~np.isnan(C[:,attr_index])
        nonzero = (C[:,attr_index] != 0 )
        # get index of values that are not NaN and not zero
        non_zero_nan = np.logical_and(nonnan, nonzero)
        # edge_number = non_zero_nan.sum()

        for subject_idx in range(n):
            s[subject_idx,attr_index]= np.matmul(data_matrix[subject_idx,non_zero_nan], C[non_zero_nan,attr_index])

    return s


def plot_hist(data, filename):
    fig, ax = plt.subplots()

    ax.hist(data, bins=25, color="lightblue", edgecolor = "black")
    ax.set_xlabel("Pearson's r")

    plt.savefig(filename)


def histogram(values1, values2, title, filepath, feature):
    fig, axs = plt.subplots(1, 2, sharey=True, tight_layout=False)
    fig.suptitle(title)

    # get min/max of all values
    all_vals = []
    for v in values1:
        all_vals.append(v)
    for v in values2:
        all_vals.append(v)
    min_val = math.floor(min(all_vals))
    max_val = math.ceil(max(all_vals))

    # We can set the number of bins with the *bins* keyword argument.
    axs[0].set_xlim(min_val, max_val)
    axs[1].set_xlim(min_val, max_val)

    axs[0].hist(values1, color='lightblue', edgecolor="black") #, bins=n_bins)
    axs[0].set_xlabel(feature + " (male)")
    axs[0].set_ylabel("Number of subjects")

    axs[1].hist(values2, color='lightblue', edgecolor="black") #, bins=n_bins)
    axs[1].set_xlabel(feature + " (female)")

    fig.savefig(filepath)


def scatterplot(x, y, title, xlabel, ylabel, filename, line=None, label=""):

    test = plt.subplot(111)
    test.clear()

    plt.scatter(x,y)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    #if title is not None:
    #    plt.title(title)
    if line is not None:
        plt.plot(line[0], line[1], color='red', label=label)
        plt.legend(bbox_to_anchor=(0, 1, 1, 0), loc="lower left")
    #fig1, ax1 = plt.subplots()

    #ax1.set_title(title)
    #if line is not None:
    #    ax1.plot(line[0], line[1])
    #ax1.scatter(x,y)
    #ax1.set_xticks([1,2])
    #ax1.set_xticklabels(xlabels)
    #ax1.set_xlabel(xlabel)
    #ax1.set_ylabel(ylabel)

    plt.savefig(filename)


def read_ids(id_file):
    """
    Read ordered HCP ids from file
    Expected format is one id per line
    """
    ids = []
    with open(id_file) as infile:
        for line in infile:
            if line.rstrip():
                ids.append(line.rstrip())
    return ids    


def run_crossvalidation(data_matrix, attr_matrix, ids=[], fold=10):
    """
    Run 10-fold cross-validation for the given data

    Parameter:
      data_matrix      -- n x m numpy matrix
      attribute_matrix -- n x l numpy matrix
      fold (int)       -- <int>-fold cross validation
    """
    (n,m) = data_matrix.shape
    test_size = round(n / fold)
    num_iterations = round(n / test_size)

    used_ids = []

    print("Size test set: {} ({} total)".format(test_size, n))
    total_errors = np.array([])
    actual_errors = np.array([]) # non-absolute errors
    total_predictions = np.array([])
    total_real = np.array([])

    for i in range(1, num_iterations + 1):
        print("  Running iteration {}/{}...".format(i, num_iterations))
        logging.info("Running {}-fold crossvalidation iteration {}/{}...".format(fold, i, num_iterations))

        # set subject start and stop index for the test set
        test_start = (i - 1) * test_size
        test_stop = test_start + test_size

        M_testing = data_matrix[test_start:test_stop,:];
        L_testing =  attr_matrix[test_start:test_stop];

        print("   Test size: " + str(L_testing.shape[0]))

        # defined training & testing sets
        M_training = np.array([]);
        L_training = np.array([]);
        if test_start > 0:
            M_training = data_matrix[:test_start,:];
            L_training = attr_matrix[:test_start];

        if test_stop < n:
            if len(M_training) == 0:
                M_training = data_matrix[test_stop:n,:]
                L_training = attr_matrix[test_stop:n]
            else:
                M_training = np.concatenate((M_training, data_matrix[test_stop:n,:]))
                L_training = np.concatenate((L_training, attr_matrix[test_stop:n]))

        # fit the model with the training data and predict the testing data
        (C_training, p_training) = find_correlations(M_training, L_training, force_positive=False)
        S_training = compute_S(M_training, C_training)

        S_testing = compute_S(M_testing, C_training) # use C obtained from the training data

        regr = linear_model.LinearRegression()
        prediction = None

        if params.poly:
            # exponential regression
            msg = "Using polynomial regression!"
            print(msg)
            logging.info(msg)
            poly = PolynomialFeatures(degree=2)
            S_training_poly = poly.fit_transform(S_training)
            S_testing_poly = poly.fit_transform(S_testing)
            regr.fit(S_training_poly, L_training)
            prediction = regr.predict(S_testing_poly)
        else:
            # linear regression
            regr.fit(S_training, L_training)
            prediction = regr.predict(S_testing)

        error = np.absolute(np.subtract(prediction, L_testing))
        total_errors = np.append(total_errors, error)
        actual_errors = np.append(actual_errors, np.subtract(prediction, L_testing))

        if params.id_file is not None:
            current_ids = ids[test_start:test_stop]
            assert(len(current_ids) == len(error))
            used_ids.extend(current_ids)

        total_predictions = np.append(total_predictions, prediction)
        total_real = np.append(total_real, L_testing)

    mean_error = np.mean(total_errors)
    std = np.std(total_errors)
    print("MAE: {:.2f} (±{:.2f})".format(mean_error, std))

    # Range-normalized MAE:
    NMAE = mean_error / (np.max(total_real) - np.min(total_real))
    print("NMAE: {}".format(NMAE))

    MAPE =  mean_absolute_percentage_error(total_real, total_predictions)
    print("MAPE:", MAPE)

    return(total_real, total_predictions, actual_errors, used_ids)


def get_sex_mask(sex_file, sex):
    """
    Get a boolean list depending on the specified sex
    The order of subjects in the sex file must be the same
    as in the data file for which the mask will be applied!

    The mask can be used to filter the data matrix by sex

    Parameters:
        sex_file (str) -- path to file containig subject sex
        sex (str)      -- sex for which to set the mask values to True ("M" or "F")

    Returns:
        mask (list) -- list of boolean values - True for the specified sex
    """
    mask = []
    with open(sex_file) as infile:
        for line in infile:
            fields = line.rstrip().split("\t")
            mask.append(fields[0] == sex)
    return mask


def normalize_matrix_rows(data):
    """
    Normalize the rows of the matrix by the maximum, i.e.
    normalization per subject by dividing all connectivity strengths
    by the maximum connectivity strength of the subject
    """
    print("Normalizing connectivity strength per subject!")
    return normalize(data, norm="max", axis=1)
    

def normalize_all_matrix_cols(data):
    """
    Normalize the columns of the matrix by the maximum, i.e.
    normalization per connection by dividing all connectivity strengths
    by the maximum connectivity strength over all subjects of the connection
    """
    print("Normalizing connectivity strength per connection!")
    return normalize(data, norm="max", axis=0)


def normalize_matrix_cols(data, num_connections):
    """
    Note: This assumes that the first num_connections columns
          of the matrix contain connectivity strength features.
          These columns will be normalized by the maximum connectivity
          strenght occurring within them. The subsequent columns are
          normalized individually by the max of each column
    """
    num_columns = data.shape[1]
    n = num_connections

    # normalize first n columns by the max of that sub-matrix
    # and subsequent columns by the max of each column
    data_norm = data[:, 0:n]
    max_strength = np.max(data_norm)
    data_norm = np.true_divide(data_norm, max_strength)

    # normalize & append the rest of the columns
    for i in range(n, num_columns):
        col = data[:, i]
        col_max = np.max(col)
        data_norm = np.c_[data_norm, np.true_divide(col, col_max)]
        
    return data_norm


def get_top_connections(C, n):
    """
    Get a map with indices of the top n connections in C (strongest correlation values)
    """
    #print("Getting {} top connections".format(n))
    connections = []

    for i,val in enumerate(C[:,0]):
        connections.append((abs(val), i))

    # sort the connection (indices) by descending absolute correlation values
    connections.sort(reverse=True, key=lambda x: x[0])

    # get top connections
    top_indices = {}
    for i in range(n):
        x = connections[i]
        top_indices[x[1]] = True

    return top_indices


def check_stability(M, L):
    """
    Check how stable the correlation vector is when drawing random samples
    """
    print("Checking stability!")

    # get base vector C
    (C,Pval) = find_correlations(M, L)
    num_subjects = M.shape[0]

    top_connections = 20
    top_indices = get_top_connections(C, top_connections)

    training_set_fraction = 0.99
    sample_size = math.floor(training_set_fraction * num_subjects)
    print("Sample size: {} (fraction: {})".format(sample_size, training_set_fraction))
    repetitions = 10

    MAE_vec = []
    reidentified_top_connections = []

    for repetition in range(repetitions):

        # create random sample
        mask = [False] * num_subjects
        sample_idx = random.sample(range(num_subjects), sample_size)
        for idx in sample_idx:
            mask[idx] = True

        M_sample = M[mask, :]
        L_sample = L[mask]

        # get correlations for sample
        (C_rep, p_rep) = find_correlations(M_sample, L_sample)

        # diff
        abs_errors = np.absolute(np.subtract(C, C_rep))
        MAE_vec.append(np.mean(abs_errors))

        # get number of reidentified top connections
        top_indices_r = get_top_connections(C_rep, top_connections)
        top_reidentified = 0
        for index in top_indices_r:
            if index in top_indices:
                top_reidentified += 1
        reidentified_top_connections.append(top_reidentified)
        print("Reidentified: {}/{} top connections".format(top_reidentified, top_connections))

    print(MAE_vec)
    print("Mean MAE ({} repetitions): {}".format(repetitions, np.mean(MAE_vec)))
    mean_reidentified = np.mean(reidentified_top_connections)
    print("Mean reidentified top connections ({} repetitions): {}/{}".format(repetitions, mean_reidentified, top_connections))

    sys.exit(0)


print("Input data file: {}".format(connectivity_file))

# n: number of subjects
# m: number of edges
# l: number of 'measures'/'attributes', e.g. IQ or age
M = np.loadtxt(connectivity_file) # n x m
L = np.loadtxt(attribute_file)    # m x l

# optional - apply input threshold
#(n,m) = M.shape
#print(n,m)
#for attr_idx in range(n):
#   for edge_idx in range(m):
#        if M[attr_idx,edge_idx] < 1.0:
#            M[attr_idx,edge_idx] = 0.0

if params.stability:
    check_stability(M, L)


fold = 10  # 10-fold crossvalidation
if params.fold is not None:
    fold = int(params.fold)

ids = []
if params.id_file is not None:
    ids = read_ids(params.id_file)


if True or params.sex_file is None:
    # run cross-validation on full data

    if len(ids):
        assert(len(ids) == M.shape[0])

    (total_real, total_predictions, total_errors, used_ids) = run_crossvalidation(M, L, ids, fold)

    img_file = params.output_dir + "/" + os.path.basename(params.data_file) + "_real_vs_pred_{}.png".format(fold)

    # write values as json for more flexible plotting/analysis later without having to redo the cross validation
    data = {'x': total_real.tolist(), 'y': total_predictions.tolist()}
    
    if len(total_errors) and len(used_ids):
        # also write absolute errors and subject IDs into the output json
        data["errors"] = list(total_errors)
        data["subjects"] = used_ids
    
    json_file = params.output_dir + "/" + os.path.basename(params.data_file) + "_TOTAL.json"
    with open(json_file, "w") as out:
        json.dump(data, out)

    slope, intercept, r, p, std_err = scipy.stats.linregress(total_real, total_predictions)
    print("Slope: {}, Δslope: {}".format(slope, slope - 1))
    print("intercept: {}".format(intercept))
    print("LinReg r={} (p={})".format(r,p))
    print("R-squared: {:.6f}".format(r**2))

    x_max = np.amax(total_real)
    x_min = np.amin(total_real)
    x = [x_min, x_max]
    y = [x_min * slope + intercept, x_max * slope + intercept] 
    scatterplot(total_real, total_predictions, "test", "Actual values", "Predicted values", img_file, (x,y), "r={:.4f}".format(r))

    (r,p) = scipy.stats.pearsonr(total_real, total_predictions)
    print("Pearson result real/predicted: {} (p={})".format(r,p))


if params.sex_file is not None:
    # split data up by sex and run cross-validation for each subgroup

    mask_male = get_sex_mask(params.sex_file, "M") # mask to get only male subjects
    mask_female = get_sex_mask(params.sex_file,"F") # mask to get only female subjects

    M_male = M[mask_male,:]
    L_male = L[mask_male]
    M_female = M[mask_female,:]
    L_female = L[mask_female]

    ids_male = []
    ids_female = []
    if len(ids):
        ids_male = [i for idx, i in enumerate(ids) if mask_male[idx]]
        ids_female = [i for idx, i in enumerate(ids) if mask_female[idx]]
        assert(len(ids_male) == M_male.shape[0])
        assert(len(ids_female) == M_female.shape[0])

    if params.plot_feature:
        hist_file = os.path.join(params.output_dir, os.path.basename(params.data_file) + "_hist_{}.png".format(params.feature))
        print("Plotting histogram...")
        histogram(L_male, L_female, "Age distribution (HCP)", hist_file, params.feature)
        sys.exit(0)

    print("MALE:")
    (total_real, total_predictions, total_errors, used_ids) = run_crossvalidation(M_male, L_male, ids_male, fold)

    # write values as json for more flexible plotting later without having to redo CV
    data = {'x': total_real.tolist(), 'y': total_predictions.tolist()}
    if len(total_errors) and len(used_ids):
        # also write absolute errors and subject IDs into the output json
        data["errors"] = list(total_errors)
        data["subjects"] = used_ids

    json_file = params.output_dir + "/" + os.path.basename(params.data_file) + "_MALE.json"
    with open(json_file, "w") as out:
        json.dump(data, out)

    img_file = params.output_dir + "/" + os.path.basename(params.data_file) + "_real_vs_pred_{}_MALE.png".format(fold)
    slope, intercept, r, p, std_err = scipy.stats.linregress(total_real, total_predictions)
    print("Slope: {}, Δslope: {}".format(slope, slope - 1))
    print("intercept: {}".format(intercept))
    print("LinReg r={} (p={})".format(r,p))
    print("R-squared: {:.6f}".format(r**2))

    x_max = np.amax(total_real)
    x_min = np.amin(total_real)
    x = [x_min, x_max]
    y = [x_min * slope + intercept, x_max * slope + intercept]
    scatterplot(total_real, total_predictions, "test", "Actual values", "Predicted values", img_file, (x,y), "r={:.4f}".format(r))

    (r,p) = scipy.stats.pearsonr(total_real, total_predictions)
    print("Pearson result real/predicted MALE: {} (p={})".format(r,p))

    print("FEMALE:")
    (total_real, total_predictions, total_errors, used_ids) = run_crossvalidation(M_female, L_female, ids_female, fold)

    # write values as json for more flexible plotting later without having to redo CV
    data = {'x': total_real.tolist(), 'y': total_predictions.tolist()}
    if len(total_errors) and len(used_ids):
        # also write absolute errors and subject IDs into the output json
        data["errors"] = list(total_errors)
        data["subjects"] = used_ids
    
    json_file = params.output_dir + "/" + os.path.basename(params.data_file) + "_FEMALE.json"
    with open(json_file, "w") as out:
        json.dump(data, out)

    img_file = params.output_dir + "/" + os.path.basename(params.data_file) + "_real_vs_pred_{}_FEMALE.png".format(fold)
    slope, intercept, r, p, std_err = scipy.stats.linregress(total_real, total_predictions)
    print("Slope: {}, Δslope: {}".format(slope, slope - 1))
    print("intercept: {}".format(intercept))
    print("LinReg r={} (p={})".format(r,p))
    print("R-squared: {:.6f}".format(r**2))

    x_max = np.amax(total_real)
    x_min = np.amin(total_real)
    x = [x_min, x_max]
    y = [x_min * slope + intercept, x_max * slope + intercept]
    scatterplot(total_real, total_predictions, "test", "Actual values", "Predicted values", img_file, (x,y), "r={:.4f}".format(r))

    (r,p) = scipy.stats.pearsonr(total_real, total_predictions)
    print("Pearson result real/predicted: {} (p={})".format(r,p))


logging.info("Finished!")