import os
import tkinter as tk
from tkinter import filedialog
import matplotlib.pyplot as plt
import openpyxl as py
import seaborn as sns

#  Extract differential gene expression data from a provided file generated by edgeR or other RNAseq DGE package
#  (Combined logFC.xlsx) and generate a gene expression heatmap for a provided list of genes (selected by user)
#  at each condition (script assumes 4 conditions)


path = 'C:/Users/Carl/Box Sync/Lab - Zhao/Non-model yeast/FOH RNAseq/YPDvsCN5/heatmaps'
os.chdir(path)

root = tk.Tk()
root.withdraw()  # Keeps the coming file dialog from staying forever

file_path = filedialog.askopenfilename()  # Select list of genes to create heatmap

geneListSheet = py.load_workbook(file_path)
geneList = geneListSheet.active

logFC_listsheet = py.load_workbook('Combined logFC.xlsx')
logFClist = logFC_listsheet.active  # Contains fold-change of gene expression for each gene and condition pair

pathway = []

for symbol in geneList:  # Pull the expression data for each gene symbol listed in geneList from the compiled
    # differential gene expression data and save it to pathway variable
    for geneExp in logFClist:
        if symbol[0].value == geneExp[0].value:
            load = [0, 0, 0, 0]  # Updated with fold-change for each condition for current gene
            for (idx, DGE) in enumerate(geneExp[1:]):
                load[idx] = geneExp[idx + 1].value
            pathway.append(load)

y_labels = []

for name in geneList:  # Retrieve symbol for each gene for y axis labels
    y_labels.append(name[0].value)

plt.figure(figsize=(3, (len(pathway) + 2) / 5.2))
hm = sns.heatmap(pathway, vmin=-2, vmax=2, center=0, cmap='bwr', cbar=False, annot=True,
                 yticklabels=y_labels, xticklabels=('Comparison 1', 'Comparison 2', 'Comparison 3', 'Comparison 4'))
hm.figure.subplots_adjust(left=0.2)
title = 'Plot title'
hm.set_title(title)
saver = hm.get_figure()
saver.savefig(title + '.png')
