import hashlib
import os
import pickle
import time
import csv
import gzip
import numpy
import pandas
import matplotlib as mpl
import matplotlib.pyplot as plt
import scipy
import pprint
import sklearn.cluster
import mpl_toolkits.mplot3d
import matplotlib.patches

def get_matrix(filename):
    return pandas.DataFrame(scipy.io.mmread(filename).toarray())

def get_feature_ids(filename):
    return [row[0] for row in csv.reader(gzip.open(filename, mode="rt"), delimiter="\t")]

def get_gene_name(filename):
    return [row[1] for row in csv.reader(gzip.open(filename, mode="rt"), delimiter="\t")]

def get_feature_type(filename):
    return [row[2] for row in csv.reader(gzip.open(filename, mode="rt"), delimiter="\t")]

def get_barcodes(filename):
    return [row[0] for row in csv.reader(gzip.open(filename, mode="rt"), delimiter="\t")]

data = dict()

def get_all(ID):
    if ID in data:
        return data[ID]
    dirname = "/home/jwlee/Spermatogenesis/result/" + ID + "/outs/filtered_feature_bc_matrix"
    matrix_dir = os.path.join(dirname, "matrix.mtx.gz")
    features_path = os.path.join(dirname, "features.tsv.gz")
    barcodes_path = os.path.join(dirname, "barcodes.tsv.gz")

    data[ID] = {"matrix": get_matrix(matrix_dir), "feature_ids": get_feature_ids(features_path), "gene_name": get_gene_name(features_path), "feature_type": get_feature_type(features_path), "barcodes": get_barcodes(barcodes_path)}

    return data[ID]

now = time.strftime("%m%d%H%M%S")
figure_directory = "/home/jwlee/Spermatogenesis/figures/"
IDs = ["NS_SW1", "NS_SW2", "NS_SW3", "NS_SW4"]

def select_highly_variable_genes(raw_data, show=True, datum_point=95):
    a = raw_data.mean(axis=1).to_numpy()
    b = raw_data.var(axis=1).to_numpy()
    data = pandas.DataFrame.from_dict({"means": a, "cvs": numpy.divide(b, a)})

    data = data.loc[(data["cvs"] > 0) & (data["means"] > 0)]

    selected = data.loc[(data["cvs"] >= numpy.percentile(data["cvs"], datum_point)) & (data["means"] >= numpy.percentile(data["means"], datum_point))]
    unselected = data.loc[(data["cvs"] < numpy.percentile(data["cvs"], datum_point)) | (data["means"] < numpy.percentile(data["means"], datum_point))]

    raw_data = raw_data.iloc[selected.index]
    print("Gene & Cell:", raw_data.shape)

    if show:
        mpl.use("Agg")
        mpl.rcParams.update({"font.size": 30})

        plt.figure()
        plt.scatter(numpy.log(selected["means"]), numpy.log(selected["cvs"]), c="blue", alpha=0.6, label="Selected")
        plt.scatter(numpy.log(unselected["means"]), numpy.log(unselected["cvs"]), c="red", alpha=0.6, label="Unselected")

        plt.grid(True)
        plt.title(str(selected.shape[0]) + " Genes: " + str(100 - datum_point) + "%")
        plt.xlabel("log(means)")
        plt.ylabel("log(CV)")
        plt.legend()

        fig = plt.gcf()
        fig.set_size_inches(24, 18)
        fig.savefig(figure_directory + "HighlyVariableGene_" + now + ".png")
        plt.close()

    return raw_data

def get_whole_data(genes=None):
    def make_md5(data):
        if data is None:
            return hashlib.md5("".encode("utf-8")).hexdigest()
        else:
            return hashlib.md5(str(sorted(genes)).encode("utf-8")).hexdigest()

    if os.path.exists(make_md5(genes) + ".data"):
        with open(make_md5(genes) + ".data", "rb") as f:
            return pickle.load(f)

    if genes is not None and "ref" in genes:
        data = get_matrix("/home/jwlee/Spermatogenesis/result/ref/outs/filtered_feature_bc_matrix/matrix.mtx.gz")
        print(data)

        data = sklearn.decomposition.PCA(random_state=0, n_components=data.shape[1]).fit_transform(numpy.swapaxes(data.values, 0, 1))
        print("PCA data:", data)
        print("Cell & Gene-like:", len(data), len(data[0]))

        data = numpy.swapaxes(sklearn.manifold.TSNE(n_components=2, random_state=0).fit_transform(data), 0, 1)

        projection = dict()
        projection["Barcode"] = get_barcodes("/home/jwlee/Spermatogenesis/result/ref/outs/filtered_feature_bc_matrix/barcodes.tsv.gz")
        projection["std_TSNE-1"] = scipy.stats.zscore(data[0])
        projection["std_TSNE-2"] = scipy.stats.zscore(data[1])

        with open(make_md5(genes) + ".data", "wb") as f:
            pickle.dump(projection, f)

        return projection

    data = get_matrix("/home/jwlee/Spermatogenesis/result/aggr/outs/filtered_feature_bc_matrix/matrix.mtx.gz")

    if genes is None:
        data = select_highly_variable_genes(data)
    else:
        data["gene"] = get_gene_name("/home/jwlee/Spermatogenesis/result/aggr/outs/filtered_feature_bc_matrix/features.tsv.gz")
        data = data[data["gene"].isin(genes)]
        del data["gene"]

    data = sklearn.decomposition.PCA(random_state=0, n_components="mle").fit_transform(numpy.swapaxes(data.values, 0, 1))
    print("PCA data: ", data)
    print("Cell & Gene-like:", len(data), len(data[0]))

    data = numpy.swapaxes(sklearn.manifold.TSNE(n_components=2, random_state=0).fit_transform(data), 1, 0)

    projection = dict()
    projection["Barcode"] = numpy.array(get_barcodes("/home/jwlee/Spermatogenesis/result/aggr/outs/filtered_feature_bc_matrix/barcodes.tsv.gz"))
    projection["std_TSNE-1"] = scipy.stats.zscore(data[0])
    projection["std_TSNE-2"] = scipy.stats.zscore(data[1])

    projection = pandas.DataFrame.from_dict(projection)

    with open(make_md5(genes) + ".data", "wb") as f:
        pickle.dump(projection, f)

    return projection

def draw_all():
    mpl.use("Agg")
    mpl.rcParams.update({"font.size": 30})

    projection = get_whole_data()

    plt.figure()
    plt.scatter(projection["std_TSNE-1"], projection["std_TSNE-2"], alpha=0.6)

    plt.grid(True)
    plt.title("Total")
    plt.xlabel("Standardized TSNE-1")
    plt.ylabel("Standardized TSNE-2")

    fig = plt.gcf()
    fig.set_size_inches(24, 18)
    fig.savefig(figure_directory + "total" + "_" + now + ".png")
    plt.close()

def get_real_barcodes(ID):
    projection = pandas.read_csv("/home/jwlee/Spermatogenesis/result/" + ID + "/outs/analysis/tsne/2_components/projection.csv", header=0)

    return [barcode[:-1] + ID[-1] for barcode in projection["Barcode"]]

def get_data_from_id(ID, genes=None):
    if ID == "ref":
        return get_whole_data(genes=["ref"])

    projection = get_whole_data(genes)
    return projection[numpy.isin(projection["Barcode"], get_real_barcodes(ID))]

def get_data_from_id_3d(ID, genes=None):
    projection = get_whole_data_3d(genes)
    return projection[numpy.isin(projection["Barcode"], get_real_barcodes(ID))]

def draw_all_with_color():
    mpl.use("Agg")
    mpl.rcParams.update({"font.size": 30})

    plt.figure()
    for ID in IDs:
        projection = get_data_from_id(ID)
        plt.scatter(projection["std_TSNE-1"], projection["std_TSNE-2"], alpha=0.6, label=ID)

    plt.grid(True)
    plt.title("Total")
    plt.xlabel("Standardized TSNE-1")
    plt.ylabel("Standardized TSNE-2")
    plt.legend()

    fig = plt.gcf()
    fig.set_size_inches(24, 18)
    fig.savefig(figure_directory + "total_" + now + ".png")
    plt.close()

def draw_tSNE(ID, genes=None):
    mpl.use("Agg")
    mpl.rcParams.update({"font.size": 30})

    whole_projection = get_whole_data(genes)

    wanted = whole_projection[numpy.isin(whole_projection["Barcode"], get_real_barcodes(ID))]
    unwanted = whole_projection[numpy.invert(numpy.isin(whole_projection["Barcode"], get_real_barcodes(ID)))]

    plt.figure()
    plt.scatter(unwanted["std_TSNE-1"], unwanted["std_TSNE-2"], c="tab:gray", alpha=0.6)
    plt.scatter(wanted["std_TSNE-1"], wanted["std_TSNE-2"], c="tab:blue", alpha=1)

    plt.grid(True)
    plt.title(ID)
    plt.xlabel("Standardized TSNE-1")
    plt.ylabel("Standardized TSNE-2")

    fig = plt.gcf()
    fig.set_size_inches(24, 18)
    fig.savefig(figure_directory + ID + "_" + now + ".png")
    plt.close()

def make_cluster_dict(cells):
    cells = cells.tolist()
    given = dict()
    for i in range(max(cells) + 1):
        given[i] = list(filter(lambda x: cells[x] == i, list(range(len(cells)))))
    return given

def clustering_Spectral_with_num(ID, num_groups):
    mpl.use("Agg")
    mpl.rcParams.update({"font.size": 30})

    projection = get_data_from_id(ID)

    projection["group"] = sklearn.cluster.SpectralClustering(n_clusters=num_groups, random_state=0, n_jobs=-1).fit_predict(projection[["std_TSNE-1", "std_TSNE-2"]].values)

    group = make_cluster_dict(projection["group"])
    data = [group[i] for i in group]
    cluster_centers = [numpy.mean([projection.loc[d, "std_TSNE-1"], projection.loc[d, "std_TSNE-2"]], axis=1) for d in data]

    plt.figure()
    plt.scatter(projection["std_TSNE-1"], projection["std_TSNE-2"], c=projection["group"])
    plt.scatter([elem[0] for elem in cluster_centers], [elem[1] for elem in cluster_centers], c="k", marker="X")
    for i, loc in enumerate(cluster_centers):
        plt.text(loc[0] + 0.05, loc[1], str(i), fontsize=30, bbox=dict(color="white", alpha=0.8))

    plt.grid(True)
    plt.title("Spectral: " + str(num_groups))
    plt.xlabel("Standardized TSNE-1")
    plt.ylabel("Standardized TSNE-2")

    fig = plt.gcf()
    fig.set_size_inches(24, 18)
    fig.savefig(figure_directory + "Spectral_" + ID + "_" + str(num_groups) + "_" + now + ".png")
    plt.close()

    return (group, cluster_centers)

def clustering_Kmeans_with_num(ID, num_groups):
    mpl.use("Agg")
    mpl.rcParams.update({"font.size": 30})

    projection = get_data_from_id(ID)

    kmeans = sklearn.cluster.KMeans(n_clusters=num_groups, random_state=0, n_jobs=-1).fit(projection[["std_TSNE-1", "std_TSNE-2"]].values)

    projection["group"] = kmeans.fit_predict(projection[["std_TSNE-1", "std_TSNE-2"]].values)

    plt.figure()
    plt.scatter(projection["std_TSNE-1"], projection["std_TSNE-2"], c=projection["group"])
    plt.scatter([elem[0] for elem in kmeans.cluster_centers_], [elem[1] for elem in kmeans.cluster_centers_], c="k", marker="X", s=500)
    for i, loc in enumerate(kmeans.cluster_centers_):
        plt.text(loc[0] + 0.05, loc[1], str(i), fontsize=30, bbox=dict(color="white", alpha=0.8))

    plt.grid(True)
    plt.title("KMeans: " + str(num_groups))
    plt.xlabel("Standardized TSNE-1")
    plt.ylabel("Standardized TSNE-2")

    fig = plt.gcf()
    fig.set_size_inches(24, 18)
    fig.savefig(figure_directory + "KMeans_" + ID + "_" + str(num_groups) + "_" + now + ".png")
    plt.close()

    return (make_cluster_dict(projection["group"]), kmeans.cluster_centers_)

def clustering_Kmeans(ID, num=10):
    return [clustering_Kmeans_with_num(ID, i) for i in range(2, num + 1)]

def clustering_Kmeans_with_num_3d(ID, num_groups):
    def change_str(num):
        color = ["#7e1e9c", "#15b01a", "#0343df", "#ff81c0", "#653700", "#e50000", "#95d0fc", "#029386", "#f97306", "#96f97b", "c20078", "#ffff14", "#75bbfd", "#929591", "#89fe05", "#bf77f6", "#9a0eea", "#033500"]
        return color[num % len(color)]

    mpl.use("Agg")
    mpl.rcParams.update({"font.size": 30})

    projection = get_data_from_id_3d(ID)

    kmeans = sklearn.cluster.KMeans(n_clusters=num_groups, random_state=0, n_jobs=-1).fit(projection[["std_TSNE-1", "std_TSNE-2", "std_TSNE-3"]].values)

    projection["group"] = kmeans.fit_predict(projection[["std_TSNE-1", "std_TSNE-2", "std_TSNE-3"]].values)

    fig = plt.figure()
    ax = mpl_toolkits.mplot3d.Axes3D(fig, elev=45, azim=135)
    ax.scatter(projection["std_TSNE-1"], projection["std_TSNE-2"], projection["std_TSNE-3"], c=projection["group"])
    ax.scatter([elem[0] for elem in kmeans.cluster_centers_], [elem[1] for elem in kmeans.cluster_centers_], [elem[2] for elem in kmeans.cluster_centers_], c="k", marker="X", s=500)

    ax.set_title("KMeans: " + str(num_groups))
    ax.set_xlabel("Standardized TSNE-1")
    ax.set_ylabel("Standardized TSNE-2")
    ax.set_zlabel("Standardized TSNE-3")

    fig = plt.gcf()
    fig.set_size_inches(24, 18)
    fig.savefig(figure_directory + "KMeans3D_" + ID + "_" + str(num_groups) + "_" + now + ".png")
    plt.close()

    if not os.path.exists("Kmeans_" + ID + "_" + str(num_groups) + ".data"):
        with open("KMeans_" + ID + "_" + str(num_groups) + ".data", "w") as f:
            f.write("x,y,z,c\n")
            for x, y, z, c in zip(projection["std_TSNE-1"], projection["std_TSNE-2"], projection["std_TSNE-3"], list(map(lambda x: change_str(x), projection["group"]))):
                f.write(str(x) + "," + str(y) + "," + str(z) + "," + c + "\n")

    return (make_cluster_dict(projection["group"]), kmeans.cluster_centers_)

def gene_in_cells(ID, cell_numbers=None):
    all_data = get_all(ID)
    all_data["matrix"].index = all_data["gene_name"]

    if cell_numbers is None:
        return all_data["matrix"]

    data = all_data["matrix"].copy()

    data.drop(all_data["matrix"].columns[list(filter(lambda x: x not in cell_numbers, list(range(all_data["matrix"].shape[1]))))], axis=1, inplace=True)

    return data

def gene_sum_in_cells(ID, cell_numbers=None, num_gene=None):
    data = gene_in_cells(ID, cell_numbers).sum(axis=1).sort_values(ascending=False)
    data = data[data > 0]

    return data if num_gene is None else data[:num_gene]

def gene_mean_in_cells(ID, cell_numbers=None, num_gene=100, text=True):
    data = gene_in_cells(ID, cell_numbers).mean(axis=1).sort_values(ascending=False)
    data = data[data > 0]

    return data if num_gene is None else data[:num_gene]

def check_valid_function(cluster_function):
    allowed_functions = [clustering_Kmeans_with_num, clustering_Spectral_with_num]
    if cluster_function not in allowed_functions:
        print("cluster_function must be in", allowed_functions)
        return False
    else:
        return True

def check_valid_function_3d(cluster_function):
    allowed_functions = [clustering_Kmeans_with_num_3d]
    if cluster_function not in allowed_functions:
        print("cluster_function must be in", allowed_functions)
        return False
    else:
        return True

def stacked_bar_gene_sum(ID, cluster_function, num_groups=10, num_gene=5):
    if not check_valid_function(cluster_function):
        return
    cluster_group, cluster_centers = cluster_function(ID, num_groups)

    gene_list = numpy.swapaxes([list(gene_sum_in_cells(ID, cluster_group[i], num_gene)) for i in cluster_group], 0, 1)
    gene_name = numpy.swapaxes([list(gene_sum_in_cells(ID, cluster_group[i], num_gene).index) for i in cluster_group], 0, 1)

    pprint.pprint(gene_list)
    pprint.pprint(gene_name)

    mpl.use("Agg")
    mpl.rcParams.update({"font.size": 30})

    plt.figure()
    plt.bar(numpy.arange(num_groups), gene_list[0], 0.35)
    for i in range(1, num_gene):
        plt.bar(numpy.arange(num_groups), gene_list[i], 0.35, bottom=numpy.sum(numpy.array([gene_list[j] for j in range(i)]), axis=0))

    gene_tick = numpy.amax(numpy.sum(gene_list, axis=0)) / 5 / num_gene
    for i in range(num_groups):
        for j in range(num_gene):
            plt.text(i + 0.05, (j + 1) * gene_tick, gene_name[j][i], fontsize=10, bbox=dict(color="white", alpha=0.3))

    plt.grid(True)
    plt.title("Stacked Bar " + ID + " with " + str(num_gene) + " Gene")
    plt.xlabel("Group")
    plt.ylabel("# of Genes")
    plt.xticks(numpy.arange(num_groups), list(range(num_groups)))

    fig = plt.gcf()
    fig.set_size_inches(24, 18)
    fig.savefig(figure_directory + "StackedBar_" + ID + "_" + str(num_groups) + "_" + str(num_gene) + "_" + now + ".png")
    plt.close()

def stacked_bar_gene_mean(ID, cluster_function, num_groups=10, num_gene=5):
    if not check_valid_function(cluster_function):
        return
    cluster_group, cluster_centers = cluster_function(ID, num_groups)

    gene_list = numpy.swapaxes([list(gene_mean_in_cells(ID, cluster_group[i], num_gene)) for i in cluster_group], 0, 1)
    gene_name = numpy.swapaxes([list(gene_mean_in_cells(ID, cluster_group[i], num_gene).index) for i in cluster_group], 0, 1)

    pprint.pprint(gene_list)
    pprint.pprint(gene_name)

    mpl.use("Agg")
    mpl.rcParams.update({"font.size": 30})

    plt.figure()
    plt.bar(numpy.arange(num_groups), gene_list[0], 0.35)
    for i in range(1, num_gene):
        plt.bar(numpy.arange(num_groups), gene_list[i], 0.35, bottom=numpy.sum(numpy.array([gene_list[j] for j in range(i)]), axis=0))

    gene_tick = numpy.amax(numpy.sum(gene_list, axis=0)) / 5 / num_gene
    for i in range(num_groups):
        for j in range(num_gene):
            plt.text(i + 0.05, (j + 1) * gene_tick, gene_name[j][i], fontsize=10, bbox=dict(color="white", alpha=0.3))

    plt.grid(True)
    plt.title("Stacked Bar " + ID + " with " + str(num_gene) + " Genes")
    plt.xlabel("Group")
    plt.ylabel("# of Gene")
    plt.xticks(numpy.arange(num_groups), list(range(num_groups)))

    fig = plt.gcf()
    fig.set_size_inches(24, 18)
    fig.savefig(figure_directory + "StackedBar_" + ID + "_" + str(num_groups) + "_" + str(num_gene) + "_" + now + ".png")
    plt.close()

def sort_index(gene_list):
    group_order = [tuple(list(scipy.stats.rankdata(data)) + [i]) for i, data in enumerate(gene_list)]

    group_order.sort()

    group_order = [list(elem)[-1] for elem in group_order]
    answer = [[i for i in gene_list[j]] for j in group_order]

    return (group_order, answer)

def heatmap_sum_top(ID, cluster_function, num_groups=10, num_gene=None, show_text=True):
    if not check_valid_function(cluster_function) and not check_valid_function_3d(cluster_function):
        return

    cluster_group, cluster_centers = cluster_function(ID, num_groups)

    gene_name = list(gene_sum_in_cells(ID).index)
    if num_gene is not None:
        gene_name = gene_name[:num_gene]
    gene_name = sorted(gene_name)

    group_order, gene_list = sort_index([gene_sum_in_cells(ID, cluster_group[i]) for i in cluster_group])
    for i, data in enumerate(gene_list):
        data.drop(labels=list(filter(lambda x: x not in gene_name, list(data.index))), inplace=True)
        data.sort_index(inplace=True)
        gene_list[i] = scipy.stats.zscore(data.tolist())

    pprint.pprint(gene_name)
    pprint.pprint(gene_list)

    mpl.use("Agg")
    mpl.rcParams.update({"font.size": 30})

    plt.figure()
    plt.imshow(gene_list)

    plt.title("HeatMap _ " + ID + "_" + str(num_gene) + " Genes")
    plt.xlabel("Genes")
    plt.ylabel("Groups")
    plt.xticks(numpy.arange(len(gene_name)), gene_name, fontsize=10, rotation=90)
    plt.yticks(numpy.arange(len(group_order)), group_order, fontsize=10)

    threshold = numpy.amax([numpy.amax(i) for i in gene_list]) / 2
    for i in range(len(gene_name)):
        for j in range(num_groups):
            if show_text:
                plt.text(j, i, str(gene_list[i][j]), color="white" if gene_list[i][j] < threshold else 'black', fontsize=10)

    fig = plt.gcf()
    fig.set_size_inches(max(24, len(gene_name) * 0.5), 18)
    fig.savefig(figure_directory + "HeatMap_" + ID + "_" + str(num_groups) + "_" + str(len(gene_name)) + "_" + now + ".png")
    plt.close()

    return (cluster_group, group_order, cluster_centers)

def heatmap_mean_top(ID, cluster_function, num_groups=10, num_gene=10, show_text=False):
    if not check_valid_function(cluster_function) and not check_valid_function_3d(cluster_function):
        return

    cluster_group, cluster_centers = cluster_function(ID, num_groups)

    gene_name = list(gene_mean_in_cells(ID).index)
    if num_gene is not None:
        gene_name = gene_name[:num_gene]
    gene_name = sorted(gene_name)

    gene_list = [gene_mean_in_cells(ID, cluster_group[i]).sort_index() for i in cluster_group]
    for i, data in enumerate(gene_list):
        data = data.add(pandas.Series(0, index=gene_name), fill_value=0)
        data.drop(labels=list(filter(lambda x: x not in gene_name, list(data.index))), inplace=True)
        data.sort_index(inplace=True)
        gene_list[i] = scipy.stats.zscore(data.tolist())

    mpl.use("Agg")
    mpl.rcParams.update({"font.size": 30})

    plt.figure()
    plt.imshow(gene_list)

    plt.title("HeatMap_" + ID + "_" + str(num_gene) + " Genes")
    plt.xlabel("Genes")
    plt.ylabel("Groups")
    plt.xticks(numpy.arange(len(gene_name)), gene_name, fontsize=10, rotation=90)
    plt.yticks(numpy.arange(num_groups), list(range(num_groups)), fontsize=10)

    threshold = numpy.amax([numpy.amax(i) for i in gene_list]) / 2
    for i in range(len(gene_name)):
        for j in range(num_groups):
            if show_text:
                plt.text(j, i, str(gene_list[i][j]), color="white" if gene_list[i][j] < threshold else 'black', fontsize=10)

    fig = plt.gcf()
    fig.set_size_inches(max(24, len(gene_name) * 0.5), 18)
    fig.savefig(figure_directory + "HeatMap_" + ID + "_" + str(num_groups) + "_" + str(len(gene_name)) + "_" + now + ".png")
    plt.close()

    return (cluster_group, list(range(num_groups)), cluster_centers)

def find_marker_gene(ID, cluster_function, num_groups=10):
    if not check_valid_function(cluster_function):
        return

    cluster_group, cluster_centers = cluster_function(ID, num_groups)

    whole_cells = gene_in_cells(ID)
    marker_gene = list()
    for i in cluster_group:
        selected_gene = list()
        data = gene_in_cells(ID, cell_numbers=cluster_group[i])
        for row in list(data.index):
            value = scipy.stats.ttest_ind(list(data.loc[row]), list(whole_cells.loc[row]))
            if value[0] > 0 and value[1] < 0.05:
                selected_gene.append((value[1], row))
        selected_gene.sort()
        print(selected_gene[:10])
        marker_gene.append(tuple(selected_gene[i][1] for i in range(3)))

    return marker_gene

gene_1 = ["Grfa1", "Zbtb16", "Nanos3", "Nanos2", ",Sohlh1", "Neurog3", "Piwil4", "Lin28a", "Utf1", "Kit", "Uchl1", "Dmrt1", "Sohlh2", "Dazl", "Stra8", "Scml2", "Rpa2", "Rad51", "Rhox13", "Dmc1", "Melob", "Sycp1", "Sycp3", "Ccnb1ip1", "Hormad1", "Piwil2", "Piwil1", "Atr", "Mybl1", "Dyx1c1", "Msh3", "Ccnb1", "Spo11", "Ldha", "Ldhc", "Cetn4", "Tekt1", "Acr", "Ssxb1", "Ssxb2", "Acrv1", "Catsper3", "Catsper1", "Saxo1", "Hsfy2", "Txndc8", "Tnp1", "Tnp2", "Tmod4", "Gapdhs", "Car2", "Prm2", "Prm1", "Prm3", "Pgk2", "Wt1", "Sox9", "Cyp11a1", "Nr5a1", "Star", "Hsd3b1", "Clu", "Cyp17a1", "Gata4", "Acta2"]
gene_2 = ["Id4", "Gfra1", "Zbtb16", "Stra8", "Rhox13", "Sycp3", "Dmc1", "Piwil1", "Pgk2", "Acr", "Gapdhs", "Prm1"]

def heatmap_given_genes(ID, cluster_function, gene_name=gene_1, num_groups=10):
    if not check_valid_function(cluster_function) and not check_valid_function_3d(cluster_function):
        return

    cluster_group, cluster_centers = cluster_function(ID, num_groups)

    raw_gene_list = [gene_mean_in_cells(ID, cluster_group[i]) for i in cluster_group]
    gene_list = [[None for i in gene_name] for j in raw_gene_list]
    for i, data in enumerate(raw_gene_list):
        for j, gene in enumerate(gene_name):
            gene_list[i][j] = float(data.loc[gene]) if (gene in list(data.index)) else 0.0

    for i, data in enumerate(gene_list):
        if numpy.unique(data).size > 1:
            gene_list[i] = scipy.stats.zscore(data)
        else:
            gene_list[i] = [0 for gene in data]

    group_order, gene_list = sort_index(gene_list)
    pprint.pprint(gene_list)

    mpl.use("Agg")
    mpl.rcParams.update({"font.size": 30})

    plt.figure()
    plt.imshow(gene_list)

    plt.title("Heatmap_" + ID + "_" + str(len(gene_name)) + " Genes")
    plt.xlabel("Genes")
    plt.ylabel("Groups")
    plt.xticks(numpy.arange(len(gene_name)), gene_name, fontsize=10, rotation=90)
    plt.yticks(numpy.arange(num_groups), group_order, fontsize=10)

    fig = plt.gcf()
    fig.set_size_inches(max(24, 0.5 * len(gene_name)), max(18, 0.2 * num_groups))
    fig.savefig(figure_directory + "HeatMap_" + ID + "_" + str(num_groups) + "_" + str(len(gene_name)) + "_" + now + ".png")
    plt.close()

    return (cluster_group, group_order, cluster_centers)

def pseudotime(ID, cluster_function, num_groups=10, select_gene=True):
    if not check_valid_function(cluster_function):
        return

    if select_gene:
        cluster_group, group_order, cluster_centers = heatmap_given_genes(ID, cluster_function, num_groups=num_groups)
    else:
        cluster_group, group_order, cluster_centers = heatmap_mean_top(ID, cluster_function, show_text=False)
    projection = get_data_from_id(ID)

    mpl.use("Agg")
    mpl.rcParams.update({"font.size": 30})

    plt.figure()
    for i in cluster_group:
        plt.scatter(projection["std_TSNE-1"].iloc[cluster_group[i]], projection["std_TSNE-2"].iloc[cluster_group[i]], c=["C" + str(i % 10) for _ in range(projection["std_TSNE-1"].iloc[cluster_group[i]].size)])
    for i in range(1, len(cluster_centers)):
        plt.arrow(cluster_centers[group_order[i - 1]][0], cluster_centers[group_order[i - 1]][1], 0.8 * (cluster_centers[group_order[i]][0] - cluster_centers[group_order[i - 1]][0]), 0.8 * (cluster_centers[group_order[i]][1] - cluster_centers[group_order[i - 1]][1]), width=0.05, edgecolor=None, linestyle=":")

    plt.grid(True)
    plt.title("Ordering Groups")
    plt.xlabel("Standardized TSNE-1")
    plt.ylabel("Standardized TSNE-2")

    fig = plt.gcf()
    fig.set_size_inches(24, 18)
    fig.savefig(figure_directory + "Arrow_" + ID + "_" + str(num_groups) + "_" + now + ".png")
    plt.close()

def pseudotime_3d(ID, cluster_function, num_groups=10, select_gene=True):
    class Arrow3D(matplotlib.patches.FancyArrowPatch):
        def __init__(self, xs, ys, zs, *args, **kwargs):
            matplotlib.patches.FancyArrowPatch.__init__(self, (0, 0), (0, 0), *args, **kwargs)
            self._verts3d = xs, ys, zs

        def draw(self, renderer):
            xs3d, ys3d, zs3d = self._verts3d
            xs, ys, zs = mpl_toolkits.mplot3d.proj3d.proj_transform(xs3d, ys3d, zs3d, renderer.M)
            self.set_positions((xs[0], ys[0]), (xs[1], ys[1]))
            matplotlib.patches.FancyArrowPatch.draw(self, renderer)

    if not check_valid_function_3d(cluster_function):
        return

    if select_gene:
        cluster_group, group_order, cluster_centers = heatmap_given_genes(ID, cluster_function, num_groups=num_groups)
    else:
        cluster_group, group_order, cluster_centers = heatmap_mean_top(ID, cluster_function, num_groups=num_groups)
    projection = get_data_from_id_3d(ID)

    mpl.use("Agg")
    mpl.rcParams.update({"font.size": 30})

    fig = plt.figure()
    ax = mpl_toolkits.mplot3d.Axes3D(fig, elev=45, azim=135)

    for i in cluster_group:
        ax.scatter(projection["std_TSNE-1"].iloc[cluster_group[i]], projection["std_TSNE-2"].iloc[cluster_group[i]], projection["std_TSNE-3"].iloc[cluster_group[i]], c=["C" + str(i % 10) for _ in range(projection["std_TSNE-1"].iloc[cluster_group[i]].size)])
    for i in range(1, len(cluster_centers)):
        ax.add_artist(Arrow3D([cluster_centers[group_order[i - 1]][0], cluster_centers[group_order[i]][0]], [cluster_centers[group_order[i - 1]][1], cluster_centers[group_order[i]][1]], [cluster_centers[group_order[i - 1]][2], cluster_centers[group_order[i]][2]], mutation_scale=20, lw=3, arrowstyle="-|>", color="k"))

    ax.set_title("Ordering Groups in 3D")
    ax.set_xlabel("Standardized TSNE-1")
    ax.set_ylabel("Standardized TSNE-2")
    ax.set_zlabel("Standardized TSNE-3")

    fig = plt.gcf()
    fig.set_size_inches(24, 18)
    fig.savefig(figure_directory + "Arraw3D_" + ID + "_" + str(num_groups) + "_" + now + ".png")
    plt.close()

def bar_given_genes(ID, cluster_function, gene_name=gene_1, num_groups=10):
    if not check_valid_function(cluster_function):
        return

    cluster_group, cluster_centers = cluster_function(ID, num_groups)

    raw_gene_list = [gene_mean_in_cells(ID, cluster_group[i]) for i in cluster_group]
    gene_list = [[None for i in gene_name] for j in cluster_group]
    for i, data in enumerate(raw_gene_list):
        for j, gene in enumerate(gene_name):
            gene_list[i][j] = float(data.loc[gene]) if (gene in list(data.index)) else 0.0

    for i, data in enumerate(gene_list):
        if numpy.unique(data).size > 1:
            gene_list[i] = scipy.stats.zscore(data)
        else:
            gene_list[i] = [0 for _ in data]

    mpl.use("Agg")
    mpl.rcParams.update({"font.size": 30})

    plt.figure()
    fig, ax = plt.subplots(num_groups)

    for i, data in enumerate(gene_list):
        for j, high in enumerate(data):
            ax[i].bar(j, high, color="C" + str(j % 10), edgecolor="k", label=gene_name[j])

    plt.setp(ax, xticks=list(range(len(gene_name))), xticklabels=gene_name)

    fig = plt.gcf()
    fig.set_size_inches(max(24, 2.5 * len(gene_name)), max(18, 4 * num_groups))
    fig.savefig(figure_directory + "Bar_graph_" + ID + "_" + str(num_groups) + "_" + str(len(gene_name)) + "_" + now + ".png")
    plt.close()

def get_common_genes(ID, cluster_function, num_groups=10):
    if not check_valid_function:
        return

    cluster_group, cluster_centers = cluster_function(ID, num_groups)

    gene_list = [list(gene_mean_in_cells(ID, cluster_group[i]).index) for i in cluster_group]

    common_gene = set(gene_list[0])
    for gene in gene_list[1:]:
        if not common_gene:
            return common_gene
        common_gene = common_gene & set(gene)

    pprint.pprint(common_gene)
    print(len(common_gene))

    return common_gene

def scatter_given_genes(ID, genes=gene_1):
    def change_scale(gene_expression):
        minimum, maximum = min(gene_expression), max(gene_expression)

        return list(map(lambda x: x if x > 0.1 else 0.1, list(map(lambda x: (x - minimum) / (maximum - minimum), gene_expression))))

    data_1 = get_data_from_id(ID, genes)
    data_2 = get_all(ID)

    for gene in genes:
        try:
            number_gene = data_2["gene_name"].index(gene)
        except ValueError:
            print(gene, "is not here")
            continue

        gene_expression = change_scale(data_2["matrix"].iloc[number_gene].values)

        mpl.use("Agg")
        mpl.rcParams.update({"font.size": 30})

        plt.figure()
        for x, y, alpha in zip(data_1["std_TSNE-1"], data_1["std_TSNE-2"], gene_expression):
            plt.scatter(x, y, c='b', alpha=alpha)

        plt.grid(True)
        plt.title(ID + "_" + gene)
        plt.xlabel("Standardized TSNE-1")
        plt.ylabel("Standardized TSNE-2")

        fig = plt.gcf()
        fig.set_size_inches(24, 18)
        fig.savefig(figure_directory + "Scatter_" + ID + "_" + gene + "_" + now + ".png")
        plt.close()

        print(gene, "Done!!")

def get_whole_data_3d(genes=None):
    def make_md5(data):
        if data is None:
            return hashlib.md5("3d".encode("utf-8")).hexdigest()
        else:
            return hashlib.md5(("3d" + str(sorted(data))).encode("utf-8")).hexdigest()

    if os.path.exists(make_md5(genes) + ".data"):
        with open(make_md5(genes) + ".data", "rb") as f:
            return pickle.load(f)

    data = get_matrix("/home/jwlee/Spermatogenesis/result/aggr/outs/filtered_feature_bc_matrix/matrix.mtx.gz")

    if genes is None:
        data = select_highly_variable_genes(data)
    else:
        data["gene"] = get_gene_name("/home/jwlee/Spermatogenesis/result/aggr/outs/filtered_feature_bc_matrix/features.tsv.gz")
        data = data[data["gene"].isin(genes)]
        del data["gene"]

    data = sklearn.decomposition.PCA(random_state=0, n_components="mle").fit_transform(numpy.swapaxes(data.values, 0, 1))
    print("PCA data: ", data)
    print("Cell & Gene-like:", len(data), len(data[0]))

    data = numpy.swapaxes(sklearn.manifold.TSNE(n_components=3, random_state=0).fit_transform(data), 0, 1)

    projection = dict()
    projection["Barcode"] = numpy.array(get_barcodes("/home/jwlee/Spermatogenesis/result/aggr/outs/filtered_feature_bc_matrix/barcodes.tsv.gz"))
    projection["std_TSNE-1"] = scipy.stats.zscore(data[0])
    projection["std_TSNE-2"] = scipy.stats.zscore(data[1])
    projection["std_TSNE-3"] = scipy.stats.zscore(data[2])

    projection = pandas.DataFrame.from_dict(projection)

    with open(make_md5(genes) + ".data", "wb") as f:
        pickle.dump(projection, f)

    return projection

def draw_tSNE_3d(ID, genes=None):
    mpl.use("Agg")
    mpl.rcParams.update({"font.size": 30})

    whole_projection = get_whole_data_3d(genes)

    if not os.path.exists("whole.data"):
        with open("whole.data", "w") as f:
            f.write("x,y,z\n")
            for x, y, z in zip(whole_projection["std_TSNE-1"], whole_projection["std_TSNE-2"], whole_projection["std_TSNE-3"]):
                f.write(str(x) + "," + str(y) + "," + str(z) + "\n")

    wanted = whole_projection[numpy.isin(whole_projection["Barcode"], get_real_barcodes(ID))]
    unwanted = whole_projection[numpy.invert(numpy.isin(whole_projection["Barcode"], get_real_barcodes(ID)))]

    if not os.path.exists(ID + ".data"):
        with open(ID + ".data", "w") as f:
            f.write("x,y,z\n")
            for x, y, z in zip(wanted["std_TSNE-1"], wanted["std_TSNE-2"], wanted["std_TSNE-3"]):
                f.write(str(x) + "," + str(y) + "," + str(z) + "\n")

    fig = plt.figure()
    ax = mpl_toolkits.mplot3d.Axes3D(fig, elev=45, azim=135)

    ax.scatter(unwanted["std_TSNE-1"], unwanted["std_TSNE-2"], unwanted["std_TSNE-3"], c="tab:gray", alpha=0.6)
    ax.scatter(wanted["std_TSNE-1"], wanted["std_TSNE-2"], wanted["std_TSNE-3"], c="tab:blue", alpha=1)

    ax.set_xlabel("Standardized TSNE-1")
    ax.set_ylabel("Standardized TSNE-2")
    ax.set_zlabel("Standardized TSNE-3")
    ax.set_title(ID)

    fig = plt.gcf()
    fig.set_size_inches(24, 18)
    fig.savefig(figure_directory + ID + "_3D_" + now + ".png")
    plt.close()

if __name__ == "__main__":
    for ID in IDs:
        pseudotime_3d(ID, clustering_Kmeans_with_num_3d, select_gene=False)
    for _ in range(5):
        print("\a")


from numpy.core.numeric import array_equal
from numpy.lib import math
import scipy.ndimage
import scipy.stats
from scipy import signal
import matplotlib.pyplot as plt
import numpy as np
np.set_printoptions(suppress=True)

def gauss1d_a(pointCount: int, std: float):
    return

def gauss1d_b(pointCount: int, std: float):
    xs = range(pointCount)
    xs = [x - (pointCount - 1)/2 for x in xs]
    ys = np.exp(-np.square(xs)/(2 * std * std))
    return ys

def show1d():
    a = gauss1d_a(25, 3)
    b = gauss1d_b(25, 3)
    np.testing.assert_array_almost_equal(a, b)
    print(", ".join([str(x) for x in a]))
    plt.plot(a)
    plt.plot(b)
    plt.show()

def show2d():
    gkern1d = gauss1d_b(5, 2)
    gkern2d = np.outer(gkern1d, gkern1d)
    #plt.imshow(gkern2d, interpolation='none', cmap='Greys_r')
    #plt.show()
    print(gkern2d)

if __name__ == "__main__":
    show2d()


<gh_stars>0
import numpy as np
import scipy.stats as ss
import tensorflow as tf
from tensorflow.keras import backend as K

rating_weights = K.expand_dims(tf.constant(np.arange(1, 11), dtype='float32'), -1)

def earth_movers_distance(y_true, y_pred):
    cdf_true = K.cumsum(y_true, axis=-1)
    cdf_pred = K.cumsum(y_pred, axis=-1)
    emd = K.sqrt(K.mean(K.square(cdf_true - cdf_pred), axis=-1))
    return K.mean(emd)

def pearson_correlation(y_true, y_pred):
    # y_true = y_true * rating_weights
    # y_pred = y_pred * rating_weights
    #
    # xm = y_true - K.mean(y_true)
    # ym = y_pred - K.mean(y_pred)
    # print(xm, ym)
    # pearson_correlation = K.sum(xm * ym) / K.sqrt(K.sum(K.square(xm) * K.square(ym)))
    # print(pearson_correlation)
    # return K.square(pearson_correlation)  # is is actually R-squared from regression
    means_true = y_true - K.mean(y_true)
    means_pred = y_pred - K.mean(y_pred)

    # normalizing stage - setting a 1 variance
    means_true = K.l2_normalize(means_true, axis=0)
    means_pred = K.l2_normalize(means_pred, axis=0)

    # final result
    pearson_correlation = K.sum(means_true * means_pred)

    return pearson_correlation

def pearson_correlation_ava(y_true, y_pred):
    x = K.cumsum(K.dot(y_true, rating_weights))
    y = K.cumsum(K.dot(y_pred, rating_weights))
    return pearson_correlation(x, y)

def spearman_correlation(y_true, y_pred):
    return ss.spearmanr(y_true.numpy(), y_pred.numpy())[0]

def two_class_quality(y_true, y_pred):
    x = K.dot(y_true, rating_weights)
    y = K.dot(y_pred, rating_weights)
    score = K.equal(tf.floor(x / 5), tf.floor(y / 5))
    return K.mean(score)

def two_class_quality_acc(y_true, y_pred):
    x = np.array(y_true)
    y = np.array(y_pred)
    score = np.equal(np.floor(x / 5), np.floor(y / 5))
    return np.mean(score)

def mean_abs_percentage(y_true, y_pred):
    abs_diff = K.abs(y_pred - y_true) / y_true
    return K.mean(1 - abs_diff)

def mean_abs_percentage_acc(y_true, y_pred):
    x, y = np.array(y_true), np.array(y_pred)
    abs_diff = np.abs(x - y) / x
    return np.mean(1 - abs_diff)

def mean_abs_percentage_ava(y_true, y_pred):
    x = K.dot(y_true, rating_weights)
    y = K.dot(y_pred, rating_weights)
    return mean_abs_percentage(x, y)

if __name__ == '__main__':
    a = np.array([1.62, 1.83, 1.89, 1.55, 1.74, 1.6, 1.6, 1.72, 1.54, 1.82])
    b = np.array([57.15, 91.69, 95.27, 56.16, 78.52, 66.09, 63.71, 79.58, 50.22, 93.39])

    print(a)
    print(b)
    print(pearson_correlation(a, b))
    print(ss.pearsonr(a, b))
    print(spearman_correlation(a, b))
    print(ss.spearmanr(a, b))

    a = np.array([1.62, 4.83, 5.89, 8.55, 8.74, 6.6, ])
    b = np.array([2.62, 3.83, 1.89, 6.55, 5.74, 4.6, ])
    print(two_class_quality(a, b))
    print(mean_abs_percentage(a, b))


# import pyedflib
import numpy as np
from scipy import signal as sg
import argparse
import sys
import json
# import  matplotlib.pyplotmatplot  as plt
from pprint import pprint

import pandas as pd

class Notch():
    Q = 0
    f0 = 0
    def __init__(self,f0=60,Q=50):
        self.f0=f0
        self.Q=Q

    def argparse(self):
        parser = argparse.ArgumentParser()
        parser.add_argument('-i','--archivo',help='Ingrese el nombre del archivo .edf a utilizar',type = str)
        parser.add_argument('-fo','--fo',help='Frecuencia que se desea filtrar. Por defecto fo = 60',type = float)
        parser.add_argument('-Q','--Q',help='Factor de calidad del filtro. Por defecto Q = 50',type = int)
        parser.add_argument('-e','--edf',help='Nombre y dirección del archivo .edf de salida',type = str)
        parsedargs = parser.parse_args()
        arc = parsedargs.archivo
        output = parsedargs.edf
        if (parsedargs.fo != None):
            if (parsedargs.fo> 0):
                self.f0 = parsedargs.fo
        if (parsedargs.Q != None):
            if (parsedargs.Q>0):
                self.Q = parsedargs.Q
        return arc,output

    # def read_edf(self,nameEdf):
    #     '''
    #     Descripción: Se encarga de leer el archivo .edf
    #     Entradas: - nameEdf: nombre del archivo .edf
    #     Salidas: - in_signal: Matriz de Canales X Tiempo
    #              - fs: Frecuencia de muestro
    #              - headers: Etiquetas del archivo .edf
    #     '''
    #     edf = pyedflib.EdfReader(nameEdf)
    #     headers = edf.getSignalHeaders()
    #     nch  = edf.signals_in_file
    #     nsig = edf.getNSamples()[0]
    #     fs = edf.getSampleFrequency(0)
    #     in_signal = np.zeros((nch,nsig))
    #     for x in range(nch):
    #         in_signal[x,:] = edf.readSignal(x)
    #     edf._close()
    #     del edf
    #     return  in_signal,fs,headers

    def filt(self,in_signal,fs):
        '''
        Descripción: Se encarga de filtrar los datos del EEG
        Entradas: - in_signal: Matriz de Canales X Tiempo
                  - fs: Frecuencia de muestro
        Salidas: - out_signal: EEG filtrado (Matriz de CanalesXTiempo)
        '''
        w0 = self.f0/(fs/2)
        num,den = sg.iirnotch(w0,self.Q)
        out_signal = np.zeros((len(in_signal),len(in_signal[0])))
        for i in range(0,len(in_signal)):
            out_signal[i]=sg.filtfilt(num,den,in_signal[i])
        return out_signal,num,den

    # def write_edf(self,in_signal,headers,nameEdf):
    #     '''
    #     Descripción: Se encarga de escribir los datos del nuevo EEG
    #     Entradas: - headers: etiquetas del .edf
    #               - in_signal: Matriz de Canales X Tiempo
    #               - nameEdf : Nombre con el que se desea guardar el nuevo .edf
    #     '''
    #     edf = pyedflib.EdfWriter(nameEdf,len(in_signal),file_type=pyedflib.FILETYPE_EDFPLUS)
    #     edf_info = []
    #     edf_signal = []
    #     for i in range (len(in_signal)):
    #         channel_info={'label':headers[i]['label'],'dimension':headers[i]['dimension'],'sample_rate':headers[i]['sample_rate'],'physical_max':headers[i]['physical_max'] , 'physical_min': headers[i]['physical_min'], 'digital_max': headers[i]['digital_max'], 'digital_min': headers[i]['digital_min'], 'transducer':headers[i]['transducer'] , 'prefilter':headers[i]['prefilter']+',notch '+str(self.f0)+'Hz'}
    #         edf_info.append(channel_info)
    #         edf_signal.append(in_signal[i])
    #     edf.setSignalHeaders(edf_info)
    #     edf.writeSamples(edf_signal)
    #     edf.close()
    #     del edf

#Read data from stdin
def read_in():
    lines = sys.stdin.readlines()
    #Since our input would only be having one line, parse our JSON data from that
    return json.loads(lines[0])

if __name__ == '__main__':

    notch1 = Notch()
    # argparse input mode
     # print ("start of notch")
    # arc,output = notch1.argparse()
    # signal , fs ,headers= notch1.read_edf(arc)
    # filtered_signal,num,den = notch1.filt(signal[:,232250:234750],fs)
    # print("size of output",filtered_signal.shape)
    # print(vals)
    # print("size of input",in_signal.shape)
    # fig,subplt=plt.subplots(3,1,figsize=(8,5))
    # subplt[0].plot(t,inp[9][ni:nf])
    # subplt[0].title.set_text('Señal original')
    # subplt[0].grid()
    #notch1.write_edf(filtered_signal,headers,output)

    # python-shell input mode
    inSignals=read_in()
    nch=len(inSignals)
    nSamples = len(inSignals[0]['data'])
    fs=inSignals[0]['samplefrequency']
    # print(nch,nSamples)
    in_signal = np.zeros((nch,nSamples))

    # print(len(inSignals))
    # print(len(inSignals[0]['data']))
    currentCh=0
    for item in inSignals:
        for subitem in item['data']:
            subitem.pop('time', None)
        df = pd.DataFrame(item['data'])
        in_signal[currentCh,:]=np.array(df.values).transpose()
        # print (in_signal[currentCh,:],currentCh)
        currentCh = currentCh +1

    # python-shell execute mode
    filtered_signal,num,den = notch1.filt(in_signal,fs)

    # python-shell output mode
    response={}
    response['channels']=[]
    currentCh=0
    for channel in inSignals:

        channelObj={}
        channelObj['id']=channel['id']
        channelObj['label']=channel['label']
        channelObj['samples']=channel['samples']
        channelObj['physicalMaximum']=channel['physicalMaximum']
        channelObj['physicalMinimum']=channel['physicalMinimum']
        channelObj['digitalMaximum']=channel['digitalMaximum']
        channelObj['digitalMinimum']=channel['digitalMinimum']
        channelObj['samplefrequency']= channel['samplefrequency']
        channelObj['data']=[]
        currentD=0
        for subitem in channel['data']:
            d={}
            # d['value']=float(subitem['value'])
            d['value']=float(filtered_signal[currentCh,currentD])
            channelObj['data'].append(d)
            currentD=currentD+1
        response['channels'].append(channelObj)
        currentCh=currentCh+1
        # print(channelObj['id'])
    print (json.dumps(response))


import numpy as np
from scipy import sparse
import scipy.sparse.linalg as la
import time
from random import randint
from operators import *
from spectrum_analysis import *
from myio import *
from hamiltonians import *
from ipr import *
import sys

def simulation(dim_loc, L, n_dis, data, Hfunc, Kfunc, Zfunc=None, time_set=None):
    dim=dim_loc**L

    # Setting cycle
    if time_set is not None:
        steps=len(time_set)
        Z=np.zeros((n_dis, L, steps), dtype=complex)
        Znew=np.zeros((n_dis, L, steps), dtype=complex)
    spectral_matrix= np.zeros((n_dis, 7))

    # Disorder cycle
    for counter in range(n_dis):
        start = time.time()

        H = Hfunc(**data)
        kick = Kfunc(**data)
        U_F = np.dot(la.expm(-1j*H),kick)
        eigval, eigvec = np.linalg.eig(U_F)

        # Spectral properties
        spectrum = np.angle(eigval)
        gaps = gap(spectrum)
        shifted_gaps = shifted_gap(spectrum, dim_loc)
        shifted_gaps_2 = shifted_gap(spectrum, 2)
        log10_gaps = np.log10(gap(spectrum))
        log10_shifted_gaps = np.log10(shifted_gap(spectrum, dim_loc))
        log10_shifted_gaps_2 = np.log10(shifted_gap(spectrum, 2))
        r = ratio(spectrum)
        spectral_matrix[counter]=np.array([np.mean(gaps), np.mean(shifted_gaps),\
         np.mean(shifted_gaps_2), np.mean(log10_gaps), np.mean(log10_shifted_gaps),\
         np.mean(log10_shifted_gaps_2), r ])

        if time_set is not None:
            #Initial state
            initial_state = np.zeros(dim)
            initial_state[randint(0, dim-1)] = 1

            final_state = evolve(time_set, initial_state, eigvec, eigval)

            for i in range(L):
             Z[counter, i], Znew[counter, i] = Zfunc(initial_state, final_state, i, time_set, **data)

        elapsed = time.time()-start
        print('size', L, '\tdisorder realization', counter,'\ttime elapsed', elapsed)

    if time_set is not None:
        Z_mean=np.mean(Z, axis=(0,1))
        Z_var=np.var(Z, axis=(0,1))
        Znew_mean=np.mean(Znew, axis=(0,1))
        Znew_var=np.var(Znew, axis=(0,1))
    else:
        Z_mean=0
        Z_var=0
        Znew_mean=0
        Znew_var=0
    spectral_data=np.mean(spectral_matrix, axis=0)
    spectral_data_var=np.var(spectral_matrix, axis=0) #not really the variance!!!!
    return Z_mean, Z_var, Znew_mean, Znew_var, spectral_data, spectral_data_var

def IPR(dim_loc, L, n_dis, datavec, Hfunc, Kfunc):
    dim=dim_loc**L
    IPR_step_vec=np.zeros(len(datavec))
    IPR_tot_vec=np.zeros(len(datavec))
    IPR_sum_step_vec=np.zeros(len(datavec))
    IPR_sum_tot_vec=np.zeros(len(datavec))

    # Disorder cycle
    for counter in range(n_dis):
        start = time.time()
        JZZ_array = datavec[0]['JZZ']*(1/2+np.random.rand(L))
        hZ_array = datavec[0]['hZ']*(np.random.rand(L))
        hX_array = datavec[0]['hX']*(np.random.rand(L))
        kick = Kfunc(**datavec[0])
        for dcount, data in enumerate(datavec):
            H = Hfunc(**data, JZZ_array=JZZ_array, hZ_array=hZ_array, hX_array=hX_array)
            U_F = np.dot(la.expm(-1j*H),kick)
            eigval, eigvec = np.linalg.eig(U_F)
            if dcount==0:
                eigvec0=eigvec
                eigvec_old=eigvec
            else:
                eigvec, IPR_step =rearrange(eigvec,eigvec_old)
                IPR_step_vec[dcount]+=IPR_step
                IPR_tot_vec[dcount]+=IPR_func(eigvec, eigvec0)
                IPR_sum_step_vec[dcount]+=sum_IPR(eigvec, eigvec_old)
                IPR_sum_tot_vec[dcount]+=sum_IPR(eigvec, eigvec0)
                eigvec_old=eigvec

        elapsed = time.time()-start
        print('size', L, '\tdisorder realization', counter,'\ttime elapsed', elapsed)

    return IPR_step_vec/n_dis, IPR_tot_vec/n_dis, IPR_sum_step_vec/n_dis, IPR_sum_tot_vec/n_dis


<gh_stars>0
"""
author: @nimrobotics
description: calculates the functional connectivity between regions and plots them
"""

from multiprocessing import Condition
import numpy as np
import scipy.io
import glob
import os
import pandas as pd
# from fnirslib.activation import getPeakActivation
# from fnirslib.plots import plotData
import sys
sys.path.append('../fnirslib')
from plots import plotData
from activation import getPeakActivation, getMeanActivation

if __name__ == '__main__':
    in_dir = "./procDataAct/" #directory of the data
    out_dir = './actData/' #directory to save the plots
    subdirs = [x[0] for x in os.walk(in_dir)][1:] #get all the subdirectories
    threshold=0.4

    females = ['SAI01', 'SAI02', 'SAI04', 'SAI09', 'SAI10', 'SAI14', 'SAI19', 'SAI22', 'SAI25', 'SAI26', 'SAI27',
                'SAI28', 'SAI29', 'SAI30', 'SAI31', 'SAI32', 'SAI33', 'SAI35', 'SAI36', 'SAI38', 'SAI39', 'SAI40']
    males = ['SAI03', 'SAI05', 'SAI06', 'SAI07', 'SAI08', 'SAI11', 'SAI12', 'SAI13', 'SAI15', 'SAI16', 'SAI17',
                'SAI18', 'SAI20', 'SAI21', 'SAI23', 'SAI24', 'SAI34', 'SAI37']

    # initialize a pd df
    actDF = pd.DataFrame(columns=['ID', 'sex', 'condition', 'C1', 'C2', 'C3', 'C4', 'C5', 'C6', 'C7', 'C8', 'C9', 'C10', 'C11', 'C12', 'C13', 'C14', 'C15', 'C16', 'C17', 'C18', 'C19', 'C20', 'C21', 'C22', 'C23', 'C24', 'C25', 'C26', 'C27', 'C28', 'C29', 'C30', 'C31', 'C32', 'C33', 'C34', 'C35', 'C36', 'C37', 'C38', 'C39', 'C40', 'C41', 'C42', 'C43', 'C44', 'C45', 'C46'])

    for dir in subdirs:
        condition = dir.split('/')[-1]
        files = glob.glob(dir+'/*.mat') # get all the files in the directory
        # avgActivation = np.zeros(np.
        for file in files:
            print('Processing file: ', file)
            ID = file.split('/')[-1].split('.')[0][1:]
            if 'SAI'+ID in females:
                sex = 'F'
            if 'SAI'+ID in males:
                sex = 'M'

            print('\nProcessing file: {}, condition: {}'.format(ID, condition))
            data = scipy.io.loadmat(file) #load data from the directory
            data = data['pdata'] #get the data from the dictionary
            data = getPeakActivation(data, interval=8)
            # data = getMeanActivation(data)
            print('data shape: ', data.shape)
            print('data: ', data[0])
            actDF.loc[len(actDF)] = [ID, sex, condition] + list(data)

    # save df to csv
    actDF.to_csv(out_dir+'actDF.csv', index=False)


<filename>vislab/datasets/imagenet.py
"""
ImageNet classification and detection challenges.

Everything loaded from files, and images distributed with dataset.
"""
import os
import pandas as pd
import glob
import scipy.io
import networkx as nx
import numpy as np
import multiprocessing
import vislab
from vislab.datasets.pascal import load_annotation_files

class ImagenetGraph(object):
    """
    Represents the ImageNet structure, loaded from .mat files provided
    in the ILSVRC2013_devkit.

    Download devkit from [1] and untar into devkit_dirname.
    Then download meta_10k from [2] and place into devkit_dirname/data.

    [1]: http://imagenet.stanford.edu/image/ilsvrc2013/ILSVRC2013_devkit.tgz
    [2]: https://dl.dropboxusercontent.com/u/44891/research/meta_10k.mat
    """
    def __init__(self, metafile, type='1k'):
        """
        Parameters
        ----------
        type: string
            In ['1k', '10k', 'det'].
        """
        data = scipy.io.loadmat(metafile)['synsets']
        if not type == '10k':
            data = data[0]

        g = nx.DiGraph()

        # First pass: add nodes.
        wnids = []
        for node in data:
            if type == '10k':
                node = node[0]

            wnid = str(node[1][0])
            wnids.append(wnid)
            g.add_node(wnid, {'words': node[2][0]})

        # Second pass: add edges.
        for i, node in enumerate(data):
            if type == '10k':
                node = node[0]

            if type == 'det':
                children = node[4].flatten()
            else:
                children = node[5][0]

            # Children are IDs from the original metafile, which is 1-indexed.
            for child in children:
                g.add_edge(wnids[i], wnids[child - 1])

        self.g = g

    def node_name(self, wnid):
        word = self.g.node[wnid]['words'].split(',')[0]
        return '{} ({})'.format(word, wnid)

    def get_all_successors(self, wnid):
        children = self.g.successors(wnid)
        all_children = list(children)
        for child in children:
            all_children += self.get_all_successors(child)
        return all_children

    def get_leaf_nodes(self, wnids):
        return [
            wnid for wnid in wnids
            if not self.g.successors(wnid)
        ]

def load_imagenet_detection(year='2013', force=False, args=None):
    """
    TODO: currently only loads val split.
    TODO: current hard-coded to be 2013 split.

    Load all the annotations, including object bounding boxes.
    Loads XML data in args['num_workers'] threads using joblib.Parallel.

    Warning: this takes a few minutes to load from scratch!
    """
    if args is None:
        args = {'num_workers': multiprocessing.cpu_count()}

    cache_filename = \
        vislab.config['paths']['shared_data'] + \
        '/ilsvrc{}_dfs.h5'.format(year)
    if not force and os.path.exists(cache_filename):
        images_df = pd.read_hdf(cache_filename, 'images_df')
        objects_df = pd.read_hdf(cache_filename, 'objects_df')
        return images_df, objects_df

    # Load all annotation file data (should take < 30 s).
    # TODO: concat the dataframes here
    splits = ['val']
    for split in splits:
        annotation_filenames = glob.glob('{}/DET_bbox_{}/*.xml'.format(
            vislab.config['paths']['ILSVRC{}'.format(year)], split))
        images_df, objects_df = load_annotation_files(
            annotation_filenames, args['num_workers'])
        images_df['_split'] = split

    # Make sure that all labels are either True or False.
    images_df = images_df.fillna(False)

    # Propagate split info to objects_df
    objects_df['split'] = np.repeat(
        images_df['_split'].values,
        images_df['_num_objects'].values.astype(int)
    )

    images_df.to_hdf(cache_filename, 'images_df', mode='w')
    objects_df.to_hdf(cache_filename, 'objects_df', mode='a')
    return images_df, objects_df