from itertools import combinations, product
import numpy as np
import re
import pickle
import random


DATA_TARGET_DIR = "./data/"
FIG_TARGET_DIRS = ["./Figure1/","./Figure2/","./Figure5/","./Figure6/"]
SUBTYPES = ["EC","ENR1","ENR2","ES","LC","LS","WT"]
DEVELOPMENT =["P9P10","P12P13","P14P18","P30P40"]
OPTIONS = ["n","y"]
POSTFIX = ["","_sub_ipsi_contra"]
INPUT_DIR = "../SpaethBahugunaData/ProcessedData/"



Fig2_panel_name = dict({"modularity_index":"H","participation_pos":"I","module_degree_zscore":"J","local_assortativity_pos_whole":"K"})

# Uncomment the lower two lines for the first time. Once the seeds are generated, its better to comment and use the seeds generated
#========================================================================== 
#SEEDS = list(np.random.randint(1,99999999,20))
#pickle.dump(SEEDS,open(DATA_TARGET_DIR+"seeds.pickle","wb"))

SEEDS = list(pickle.load(open(DATA_TARGET_DIR+"seeds.pickle","rb")))[:5] # For full simulation run for all seeds
print(SEEDS)

#seed_plot = random.sample(list(SEEDS),1)[0]
seed_plot = SEEDS[0]

rule all:
	input:
		DATA_TARGET_DIR+"meta_data.csv",
		DATA_TARGET_DIR+"data_2d_maps.pickle",
		DATA_TARGET_DIR+"meta_data_days.csv",
		DATA_TARGET_DIR+"data_2d_maps_days.pickle",	

		expand(DATA_TARGET_DIR+"covariance_maps_norm_{seed}.pickle",seed=SEEDS),
		expand(FIG_TARGET_DIRS[0]+"corr_maps_rearranged_{st}_norm_{seed}.png",st=SUBTYPES,seed=seed_plot),
		expand(DATA_TARGET_DIR+"graph_properties_norm_{seed}.pickle",seed=SEEDS),

		expand(DATA_TARGET_DIR+"covariance_maps_days_norm_{seed}.pickle",seed=SEEDS),
		expand(DATA_TARGET_DIR+"graph_properties_days_norm_{seed}.pickle",seed=SEEDS),
		expand(FIG_TARGET_DIRS[0]+"corr_maps_rearranged_{st}_norm_{seed}.png",st=DEVELOPMENT,seed=seed_plot),
       
		expand(DATA_TARGET_DIR+"graph_properties_pandas_for_behav_{seed}.csv",seed=SEEDS),
		expand(DATA_TARGET_DIR+"graph_properties_pandas_{seed}.csv",seed=SEEDS),
		DATA_TARGET_DIR+"graph_properties_pandas_for_behav_all.csv",
		DATA_TARGET_DIR+"graph_properties_pandas_all.csv",

		expand(DATA_TARGET_DIR+"graph_properties_pandas_for_behav_days_{seed}.csv",seed=SEEDS),
		expand(DATA_TARGET_DIR+"graph_properties_pandas_days_{seed}.csv",seed=SEEDS),
		DATA_TARGET_DIR+"graph_properties_pandas_for_behav_days_all.csv",
		DATA_TARGET_DIR+"graph_properties_pandas_days_all.csv",
		expand(DATA_TARGET_DIR+"graph_properties_pandas_for_behav_sub_contra_ipsi_days_{seed}.csv",seed=SEEDS),
		expand(DATA_TARGET_DIR+"graph_properties_pandas_sub_contra_ipsi_days_{seed}.csv",seed=SEEDS),
		DATA_TARGET_DIR+"graph_properties_pandas_for_behav_sub_contra_ipsi_days_all.csv",
		DATA_TARGET_DIR+"graph_properties_pandas_sub_contra_ipsi_days_all.csv",

		expand(DATA_TARGET_DIR+"graph_properties_pandas_for_behav_sub_contra_ipsi_{seed}.csv",seed=SEEDS),
		expand(DATA_TARGET_DIR+"graph_properties_pandas_sub_contra_ipsi_{seed}.csv",seed=SEEDS),
		DATA_TARGET_DIR+"graph_properties_pandas_for_behav_sub_contra_ipsi_all.csv",
		DATA_TARGET_DIR+"graph_properties_pandas_sub_contra_ipsi_all.csv",


        expand(FIG_TARGET_DIRS[1]+"Figure2_{N}_development_n.png",N=[ Fig2_panel_name[x] for x in Fig2_panel_name.keys()] ),
        expand(DATA_TARGET_DIR+"behavior_features_pandas{pf}.csv",pf=POSTFIX),
        expand(DATA_TARGET_DIR+"graph_properties_with_behavior_pandas{pf}_all.csv",pf=POSTFIX),


        expand(FIG_TARGET_DIRS[2]+"tsne_all_subtypes_{sub_ipsi_contra}_{st}_seeds.png",sub_ipsi_contra='y',st="subtype"),

        FIG_TARGET_DIRS[2]+"Accuracy_comparison_y_subtype.png",
        FIG_TARGET_DIRS[2]+"Confusion_matrix_random_forest_y_subtype.png",

        expand(DATA_TARGET_DIR+"graph_properties_behavior_enr{pf}_all.csv",pf=POSTFIX),
        DATA_TARGET_DIR+"Predicted_actual_scatter_points_slope_n.csv",
        FIG_TARGET_DIRS[3]+"Predicted_actual_scatter_jittered_slope_n.png",

        DATA_TARGET_DIR+"Predicted_actual_scatter_points_total_distance_n.csv",
        FIG_TARGET_DIRS[3]+"Predicted_actual_scatter_jittered_total_distance_n.png",

        DATA_TARGET_DIR+"Predicted_actual_scatter_points_lcls_post_op15_n.csv",
        FIG_TARGET_DIRS[3]+"Predicted_actual_scatter_jittered_post_op15_n.png",

        expand(FIG_TARGET_DIRS[1]+"Graph_{day}.png",day=DEVELOPMENT)














rule read_data_adaptive:
	input:
		INPUT_DIR+"Adaptive_Dataset/",
	output:
		DATA_TARGET_DIR+"meta_data.csv",
		DATA_TARGET_DIR+"data_2d_maps.pickle",
	run:
		shell("python ./common/read_data_adaptive.py")


rule calc_graph_features_adaptive:
	input:
		DATA_TARGET_DIR+"meta_data.csv",
		DATA_TARGET_DIR+"data_2d_maps.pickle",
	output:
		expand(DATA_TARGET_DIR+"covariance_maps_norm_{seed}.pickle",seed=SEEDS),
		expand(FIG_TARGET_DIRS[0]+"corr_maps_rearranged_{st}_norm_{seed}.png",st=SUBTYPES,seed=seed_plot),
		expand(DATA_TARGET_DIR+"graph_properties_norm_{seed}.pickle",seed=SEEDS)
	run:
		shell("python ./common/plot_correlation_maps_calculate_graph_features_adaptive.py {seed}".format(seed=seed_plot))

rule read_data_development:
	input:
		INPUT_DIR+"Development_Dataset/",
	output:
		DATA_TARGET_DIR+"meta_data_days.csv",
		DATA_TARGET_DIR+"data_2d_maps_days.pickle"	
	run:
		shell("python ./common/read_data_development.py")


rule calc_graph_features_development:
	input:
		DATA_TARGET_DIR+"data_2d_maps_days.pickle",
		DATA_TARGET_DIR+"meta_data_days.csv"	
	output:
		expand(DATA_TARGET_DIR+"covariance_maps_days_norm_{seed}.pickle",seed=SEEDS),
		expand(DATA_TARGET_DIR+"graph_properties_days_norm_{seed}.pickle",seed=SEEDS),
		expand(FIG_TARGET_DIRS[0]+"corr_maps_rearranged_{st}_norm_{seed}.png",st=DEVELOPMENT,seed=seed_plot)

	run:
		shell("python ./common/plot_correlation_maps_calculate_graph_features_development.py {seed}".format(seed=seed_plot))



rule collate_graph_features_adaptive_dataframe:
    input:
        expand(DATA_TARGET_DIR+"graph_properties_norm_{seed}.pickle",seed=SEEDS)
    output:
        expand(DATA_TARGET_DIR+"graph_properties_pandas_for_behav_{seed}.csv",seed=SEEDS),
        expand(DATA_TARGET_DIR+"graph_properties_pandas_{seed}.csv",seed=SEEDS),
        DATA_TARGET_DIR+"graph_properties_pandas_for_behav_all.csv",
        DATA_TARGET_DIR+"graph_properties_pandas_all.csv"
    run:
        shell("python common/combine_graph_props_seeds_pandas.py subtype")



rule collate_graph_features_development_dataframe:
    input:
        expand(DATA_TARGET_DIR+"graph_properties_days_norm_{seed}.pickle",seed=SEEDS)
    output:
        expand(DATA_TARGET_DIR+"graph_properties_pandas_for_behav_days_{seed}.csv",seed=SEEDS),
        expand(DATA_TARGET_DIR+"graph_properties_pandas_days_{seed}.csv",seed=SEEDS),
        DATA_TARGET_DIR+"graph_properties_pandas_for_behav_days_all.csv",
        DATA_TARGET_DIR+"graph_properties_pandas_days_all.csv",
        DATA_TARGET_DIR+"graph_properties_pandas_days_null_all.csv"
    run:
        shell("python common/combine_graph_props_seeds_pandas.py development")


rule collate_graph_features_development_ipsi_contra_dataframe:
    input:
        expand(DATA_TARGET_DIR+"graph_properties_days_norm_{seed}.pickle",seed=SEEDS)
    output:
        expand(DATA_TARGET_DIR+"graph_properties_pandas_for_behav_sub_contra_ipsi_days_{seed}.csv",seed=SEEDS),
        expand(DATA_TARGET_DIR+"graph_properties_pandas_sub_contra_ipsi_days_{seed}.csv",seed=SEEDS),
        DATA_TARGET_DIR+"graph_properties_pandas_for_behav_sub_contra_ipsi_days_all.csv",
        DATA_TARGET_DIR+"graph_properties_pandas_sub_contra_ipsi_days_all.csv"
    run:
        shell("python common/combine_graph_props_ipsi_contra_seeds_pandas.py development")


rule collate_graph_features_adaptive_ipsi_contra_dataframe:
    input:
        expand(DATA_TARGET_DIR+"graph_properties_norm_{seed}.pickle",seed=SEEDS)
    output:
        expand(DATA_TARGET_DIR+"graph_properties_pandas_for_behav_sub_contra_ipsi_{seed}.csv",seed=SEEDS),
        expand(DATA_TARGET_DIR+"graph_properties_pandas_sub_contra_ipsi_{seed}.csv",seed=SEEDS),
        DATA_TARGET_DIR+"graph_properties_pandas_for_behav_sub_contra_ipsi_all.csv",
        DATA_TARGET_DIR+"graph_properties_pandas_sub_contra_ipsi_all.csv"
    run:
        shell("python common/combine_graph_props_ipsi_contra_seeds_pandas.py subtype")


rule Figure2_panelHIJK:
    input:
        DATA_TARGET_DIR+"graph_properties_pandas_days_all.csv",
        DATA_TARGET_DIR+"graph_properties_pandas_days_null_all.csv"
    output:
        expand(FIG_TARGET_DIRS[1]+"Figure2_{N}_development_n.png",N=[ Fig2_panel_name[x] for x in Fig2_panel_name.keys()] )
    run:
        shell("python Figure2/compare_graph_props_subtype_null.py development n")


rule Figure2_panelE:
    input:
       expand(DATA_TARGET_DIR+"graph_properties_norm_{seed}.pickle",seed=seed_plot) 
    output:
       expand(FIG_TARGET_DIRS[1]+"Graph_{day}.png",day=DEVELOPMENT)
    run:
       for d in DEVELOPMENT:
           shell("python Figure2/draw_graph_networkx.py days {st}".format(st=d)) 


rule read_data_behavior:
    input:
        "../DataSource/Catwalk_Raw_Profiles_Cuff_Sham_Ctrl.xlsx",
        DATA_TARGET_DIR+"graph_properties_pandas_for_behav_all.csv",
        DATA_TARGET_DIR+"graph_properties_pandas_for_behav_sub_contra_ipsi_all.csv"
    output:
        expand(DATA_TARGET_DIR+"behavior_features_pandas{pf}.csv",pf=POSTFIX),
        expand(DATA_TARGET_DIR+"graph_properties_with_behavior_pandas{pf}_all.csv",pf=POSTFIX)
    run:
        for op in OPTIONS:
            shell("python common/read_data_behavior.py {sub_ipsi_contra}".format(sub_ipsi_contra=op))


rule Figure5_panelD:
    input:
        DATA_TARGET_DIR+"graph_properties_with_behavior_pandas_sub_ipsi_contra_all.csv"
    output:
        expand(FIG_TARGET_DIRS[2]+"tsne_all_subtypes_{sub_ipsi_contra}_{st}_seeds.png",sub_ipsi_contra='y',st="subtype")
    run:
        shell("python Figure5/Figure5_PanelD.py y")

rule Figure5_panelE:
    input:
        DATA_TARGET_DIR+"graph_properties_with_behavior_pandas_all.csv",
        DATA_TARGET_DIR+"graph_properties_with_behavior_pandas_sub_ipsi_contra_all.csv"
    output:
        FIG_TARGET_DIRS[2]+"Accuracy_comparison_y_subtype.png",
        FIG_TARGET_DIRS[2]+"Confusion_matrix_random_forest_y_subtype.png"
    run:
        shell("python Figure5/Figure5_PanelE.py y subtype")


rule read_enrichment:
    input:
        "../DataSource/Enrichment.xlsx"
    output:
        #expand(DATA_TARGET_DIR+"graph_properties_pandas_for_behav_{pf}all.csv",pf=POSTFIX),
        expand(DATA_TARGET_DIR+"graph_properties_behavior_enr{pf}_all.csv",pf=POSTFIX)
    run:
        for op in OPTIONS:
            shell("python common/read_enrichment.py {sub_ipsi_contra}".format(sub_ipsi_contra=op))

rule Figure6_PanelA:
    input:
        DATA_TARGET_DIR+"graph_properties_behavior_enr_all.csv"
    output:
        DATA_TARGET_DIR+"Predicted_actual_scatter_points_slope_n.csv",
        FIG_TARGET_DIRS[3]+"Predicted_actual_scatter_jittered_slope_n.png"
    run:
        shell("python Figure6/Figure6_PanelA.py")


rule Figure6_PanelB:
    input:
        DATA_TARGET_DIR+"graph_properties_behavior_enr_all.csv"
    output:
        DATA_TARGET_DIR+"Predicted_actual_scatter_points_total_distance_n.csv",
        FIG_TARGET_DIRS[3]+"Predicted_actual_scatter_jittered_total_distance_n.png" 
    run:
        shell("python Figure6/Figure6_PanelB.py")

rule Figure6_PanelD:
    input:
        DATA_TARGET_DIR+"graph_properties_with_behavior_pandas_all.csv"
    output:
        DATA_TARGET_DIR+"Predicted_actual_scatter_points_lcls_post_op15_n.csv",
        FIG_TARGET_DIRS[3]+"Predicted_actual_scatter_jittered_post_op15_n.png"
    run:
        shell("python Figure6/Figure6_PanelD.py")

