"""

describe what is in this file

"""

# %%

import numpy as np

import random as rnd

import pandas as pd

from pathlib import Path

from scipy.stats import pearsonr

import matplotlib.pyplot as plt

from sklearn.datasets import make_sparse_spd_matrix

from move.data.preprocessing import scale

########################### Hyperparameters ####################################

PROJECT_NAME = "ffa_sim"

MODE = "linear"  # "non-linear"

SEED_1 = 1234

np.random.seed(SEED_1)

rnd.seed(SEED_1)

COV_ALPHA = 0.99

N_SAMPLES = 500

# SETTINGS = {

#     "proteomics": {

#         "features": 3000,

#         "frequencies": [0.002, 0.01, 0.02],

#         "coefficients": [500, 100, 50],

#         "phase": 0,

#         "offset": 500,

#     },

#     "metagenomics": {

#         "features": 1000,

#         "frequencies": [0.001, 0.05, 0.08],

#         "coefficients": [80, 20, 10],

#         "phase": np.pi / 2,

#         "offset": 400,

#     },

# }

SETTINGS = {

    "LP": {

        "features": 900,

        "frequencies": [0.002, 0.01, 0.02],

        "coefficients": [500, 100, 50],

        "phase": 0,

        "offset": 500,

    },

    "transcriptomics": {

        "features": 1000,

        "frequencies": [0.001, 0.05, 0.08],

        "coefficients": [80, 20, 10],

        "phase": np.pi / 2,

        "offset": 400,

    },

    # "PRS": {

    #     "features": 15,

    #     "frequencies": [0.1, 0.5, 0.8],

    #     "coefficients": [.2, .1, .05],

    #     "phase": np.pi / 2,

    #     "offset": 10,

    # },

    # "age": {

    # "features": 1,

    # "frequencies": [0.1, 0.5, 0.8],

    # "coefficients": [.1, .05, 0.1],

    # "phase": np.pi / 2,

    # "offset": 50,

    # },

}

COR_THRES = 0.02

PAIRS_OF_INTEREST = [(1, 2), (3, 4)]  # ,(77,75),(99,70),(38,2),(67,62)]

# Path to store output files

outpath = Path("./") / f"dataset_creation_outputs_{PROJECT_NAME}"

outpath.mkdir(exist_ok=True, parents=True)

################################ Functions ####################################

def get_feature_names(settings):

    all_feature_names = [

        f"{key}_{i+1}"

        for key in settings.keys()

        for i in range(settings[key]["features"])

    ]

    return all_feature_names

def create_mean_profiles(settings):

    feature_means = []

    for key in settings.keys():

        mean = settings[key]["offset"]

        for freq, coef in zip(

            settings[key]["frequencies"], settings[key]["coefficients"]

        ):

            mean += coef * (

                np.sin(

                    freq * np.arange(settings[key]["features"]) + settings[key]["phase"]

                )

                + 1

            )

        feature_means.extend(list(mean))

    return feature_means

def create_ground_truth_correlations_file(correlations):

    sort_ids = np.argsort(abs(correlations), axis=None)[::-1]  # 1D: N x C

    corr = np.take(correlations, sort_ids)  # 1D: N x C

    sig_ids = sort_ids[abs(corr) > COR_THRES]

    sig_ids = np.vstack(

        (sig_ids // len(all_feature_names), sig_ids % len(all_feature_names))

    ).T

    associations = pd.DataFrame(sig_ids, columns=["feature_a_id", "feature_b_id"])

    a_df = pd.DataFrame(dict(feature_a_name=all_feature_names))

    a_df.index.name = "feature_a_id"

    a_df.reset_index(inplace=True)

    b_df = pd.DataFrame(dict(feature_b_name=all_feature_names))

    b_df.index.name = "feature_b_id"

    b_df.reset_index(inplace=True)

    associations = associations.merge(a_df, on="feature_a_id", how="left").merge(

        b_df, on="feature_b_id", how="left"

    )

    associations["Correlation"] = corr[abs(corr) > COR_THRES]

    associations = associations[

        associations.feature_a_id > associations.feature_b_id

    ]  # Only one half of the matrix

    return associations

def plot_score_matrix(

    array, feature_names, cmap="bwr", vmin=None, vmax=None, label_step=50

):

    if vmin is None:

        vmin = np.min(array)

    elif vmax is None:

        vmax = np.max(array)

    # if ax is None:

    fig = plt.figure(figsize=(15, 15))

    plt.imshow(array, cmap=cmap, vmin=vmin, vmax=vmax)

    plt.xticks(

        np.arange(0, len(feature_names), label_step),

        feature_names[::label_step],

        fontsize=8,

        rotation=90,

    )

    plt.yticks(

        np.arange(0, len(feature_names), label_step),

        feature_names[::label_step],

        fontsize=8,

    )

    plt.tight_layout()

    # ax

    return fig

def plot_feature_profiles(dataset, feature_means):

    ## Plot profiles

    fig = plt.figure(figsize=(15, 5))

    plt.plot(

        np.arange(len(feature_means)), feature_means, lw=1, marker=".", markersize=0

    )

    for sample in dataset:

        plt.plot(

            np.arange(len(feature_means)), sample, lw=0.1, marker=".", markersize=0

        )

    plt.xlabel("Feature number")

    plt.ylabel("Count number")

    plt.title("Patient specific profiles")

    plt.tight_layout()

    return fig

def plot_feature_correlations(dataset, pairs_2_plot):

    fig = plt.figure()

    for f1, f2 in pairs_2_plot:

        plt.plot(

            dataset[:, f1],

            dataset[:, f2],

            lw=0,

            marker=".",

            markersize=1,

            label=f"{correlations[f1,f2]:.2f}",

        )

    plt.xlabel("Feature A")

    plt.ylabel("Feature B")

    plt.legend(

        loc="upper center",

        bbox_to_anchor=(0.5, -0.1),

        fancybox=True,

        shadow=True,

        ncol=5,

    )

    plt.title("Feature correlations")

    plt.tight_layout()

    return fig

def save_splitted_datasets(

    settings: dict, PROJECT_NAME, dataset, all_feature_names, n_samples, outpath

):

    # Save index file

    index = pd.DataFrame({"ID": list(np.arange(1, n_samples + 1))})

    index.to_csv(outpath / f"random.{PROJECT_NAME}.ids.txt", index=False)

    # Save continuous files

    df = pd.DataFrame(

        dataset, columns=all_feature_names, index=list(np.arange(1, n_samples + 1))

    )

    cum_feat = 0

    for key in settings.keys():

        df_feat = settings[key]["features"]

        df_cont = df.iloc[:, cum_feat : cum_feat + df_feat]

        df_cont.insert(0, "ID", np.arange(1, n_samples + 1))

        df_cont.to_csv(

            outpath / f"random.{PROJECT_NAME}.{key}.tsv", sep="\t", index=False

        )

        cum_feat += df_feat

################################## Main script ##################################

if __name__ == "__main__":

    # %%

    # Add all datasets in a single matrix:

    all_feature_names = get_feature_names(SETTINGS)

    feat_means = create_mean_profiles(SETTINGS)

    # %%

    ###### Covariance matrix definition ######

    if MODE == "linear":

        covariance_matrix = make_sparse_spd_matrix(

            dim=len(all_feature_names),

            alpha=COV_ALPHA,

            smallest_coef=-30,

            largest_coef=30,

            norm_diag=False,

            random_state=SEED_1,

        )

    elif MODE == "non-linear":

        covariance_matrix = np.identity(len(all_feature_names))

    ABS_MAX = np.max(abs(covariance_matrix))

    fig = plot_score_matrix(

        covariance_matrix, all_feature_names, vmin=-ABS_MAX, vmax=ABS_MAX

    )

    fig.savefig(outpath / f"Covariance_matrix_{PROJECT_NAME}.png")

    #    dataset = np.array(

    #        [

    #            list(np.random.multivariate_normal(feat_means, covariance_matrix))

    #            for _ in range(N_SAMPLES)

    #        ]

    #    )

    dataset = np.random.multivariate_normal(feat_means, covariance_matrix, N_SAMPLES)

    print(dataset.shape)

    # Add non-linearities

    if MODE == "non-linear":

        for i, j in PAIRS_OF_INTEREST:

            freq = np.random.choice([4, 5, 6])

            dataset[:, i] += np.sin(freq * dataset[:, j])

    scaled_dataset, _ = scale(dataset)

    # Actual correlations:

    # correlations = np.empty(np.shape(covariance_matrix))

    # for ifeat in range(len(covariance_matrix)):

    #     for jfeat in range(len(covariance_matrix)):

    #         correlations[ifeat, jfeat] = pearsonr(dataset[:, ifeat], dataset[:, jfeat])[

    #             0

    #         ]

    correlations = np.corrcoef(dataset, rowvar=False)

    fig = plot_score_matrix(correlations, all_feature_names, vmin=-1, vmax=1)

    fig.savefig(outpath / f"Correlations_{PROJECT_NAME}.png", dpi=200)

    # Sort correlations by absolute value

    associations = create_ground_truth_correlations_file(correlations)

    associations.to_csv(outpath / f"changes.{PROJECT_NAME}.txt", sep="\t", index=False)

    # Plot feature profiles per sample

    fig = plot_feature_profiles(dataset, feat_means)

    fig.savefig(outpath / "Multi-omic_profiles.png")

    ## Plot correlations

    fig = plot_feature_correlations(dataset, PAIRS_OF_INTEREST)

    fig.savefig(outpath / "Feature_correlations.png")

    fig = plot_feature_correlations(scaled_dataset, PAIRS_OF_INTEREST)

    fig.savefig(outpath / "Feature_correlations_scaled.png")

    # Write tsv files with feature values for all samples in both datasets:

    save_splitted_datasets(

        SETTINGS, PROJECT_NAME, dataset, all_feature_names, N_SAMPLES, outpath

    )