"""Clustering plots"""

"""
1. Dendrograms:
"""
#amp.plot_clustering_dendrograms(Z_struct=Z_struct, prot_names=prot_names, labels=arr_df.index[time_dict['D21']], fig_path=None)
amp.plot_clustering_dendrograms(Z_struct=Z_struct, prot_names=['EBOV-GP'], labels=arr_df.group, fig_path=fig_path, fig_prefix='')

"""
2. Raw responses by cluster:
"""
# for p, s in zip(['EBOV-GP'], ['EBOV-GP']):
# 	amp.plot_raw_responses_by_clusters(raw_arr_df, ind_dict[p], num_clusters, clusters=clusters[p],
# 										  fig_path=fig_path, fig_prefix=s + '_Raw', fig_size=(18,11), y_lims=[0, 20000])
 
"""
3. Median responses by cluster:
"""
for p, s in zip(['EBOV-GP'], ['EBOV-GP']):
	amp.plot_median_responses_by_clusters(arr_df, ind_dict[p], num_clusters, clusters=clusters[p],
										  fig_path=fig_path, fig_prefix=s + '_Raw', fig_size=(18,11), y_lims=[0, 10000])

"""
4. Summary stat boxplots by clusters:
"""
amp.plot_summary_stat_boxplots_by_clusters(arr_df, clusters, ['EBOV-GP'], arr_summary_stats,
										   sample_inds=None, fig_path=fig_path)

"""
amp.plot_summary_stat_boxplots_by_exp_groups(raw_arr_df, arr_summary_stats, sample_inds=None, fig_path=fig_path)
    

"""Clustering plots"""

"""
1. Dendrograms:
"""
#amp.plot_clustering_dendrograms(Z_struct=Z_struct, prot_names=prot_names, labels=arr_df.index[time_dict['D21']], fig_path=None)
amp.plot_clustering_dendrograms(Z_struct=Z_struct, prot_names=['VN1203', 'Indo05'], labels=arr_df.index, fig_path=fig_path, fig_prefix='Raw_')

"""
2. Raw responses by cluster:
"""
for p, s in zip(['VN1203', 'Indo05'], ['VN1203', 'Indo05']):
    amp.plot_raw_responses_by_clusters(raw_arr_df, ind_dict[p], num_clusters, clusters=clusters[p],
                                          fig_path=fig_path, fig_prefix=s + '_Raw', fig_size=(18,11), y_lims=[0, 50000])
 
"""
3. Median responses by cluster:
"""
for p, s in zip(['VN1203', 'Indo05'], ['VN1203', 'Indo05']):
    amp.plot_median_responses_by_clusters(raw_arr_df, ind_dict[p], num_clusters, clusters=clusters[p],
                                          fig_path=fig_path, fig_prefix=s + '_Raw', fig_size=(18,11), y_lims=y_lims)

"""
4. Summary stat boxplots by clusters:
"""
amp.plot_summary_stat_boxplots_by_clusters(raw_arr_df, clusters, ['VN1203', 'Indo05'], arr_summary_stats,
                                           sample_inds=None, fig_path=fig_path)
        filename = "".join([FIG_PATH, t, "_", assay, "_boxplots_by_groups.png"])
        f.savefig(filename, dpi=200)

        filename = "".join([FIG_PATH, t, "_", assay, "_boxplots_by_groups.eps"])    
        f.savefig(filename, dpi=1000)


# Figure 4, 5 and 6 - clustering dendrograms, median responses and summary stats of WT vs. Obese for each group:
for a in ['Vac', 'AS03']:
    curr_inds = group_inds['Ob_post_' + a].append(group_inds['WT_post_' + a])
    amp.plot_clustering_dendrograms(Z_struct=Z_struct[a], prot_names=['SHA_ha'], labels=arr_df.loc[curr_inds].group, fig_prefix=a + '_', fig_path=FIG_PATH)
    amp.plot_median_responses_by_clusters(arr_df=arr_df.loc[curr_inds], antigen_inds=ind_dict['SHA_ha'], num_clusters=4,
                                          clusters=clusters[a]['SHA_ha'], y_lims=[0, 25000], fig_prefix=a, fig_path=FIG_PATH)
    #amp.plot_raw_responses_by_clusters(arr_df=arr_df.loc[curr_inds], antigen_inds=ind_dict['SHA_ha'], num_clusters=4, clusters=clusters[a]['SHA_ha'])
    amp.plot_summary_stat_boxplots_by_clusters(arr_df=arr_df.loc[curr_inds], clusters=clusters[a], prot_names=['SHA_ha'], 
                                               arr_summary_stats=['H7_breadth', 'H7_mag'], fig_prefix = a + '_', fig_path=FIG_PATH)
    


# Export all post data to excel:
# now generate dataframes for sig antigens from proteins of interest:
writer = pd.ExcelWriter(FIG_PATH + 'H7H9_WT_OB_raw_data.xlsx')
for p, p_name in zip(prot_names, prot_strs):
    curr_df = arr_df.loc[time_dict['Post']][ind_dict[p]]
    
    curr_df.to_excel(writer, sheet_name='post_' + p_name)
writer.save()



# # plot correlation matrix with dendrogram overlayed: